# Path: alphaware/base/factor_transformer.py
# class FactorTransformer(BaseEstimator, TransformerMixin):
#     def __init__(self, copy=True, out_container=False):
#         self.copy = copy
#         self.df_mapper = None
#         self.out_container = out_container
# 
#     @expect_types(out=bool)
#     def set_out_container(self, out):
#         self.out_container = out
# 
#     @preprocess(factor_container=ensure_factor_container)
#     def fit(self, factor_container, **kwargs):
#         self.df_mapper = self._build_mapper(factor_container)
#         return self
# 
#     def _build_mapper(self, factor_container):
#         raise NotImplementedError
# 
#     @preprocess(factor_container=ensure_factor_container)
#     def transform(self, factor_container):
#         if self.copy:
#             factor_container = copy.deepcopy(factor_container)
#         tiaocang_date = factor_container.tiaocang_date
#         calc_data = [self.df_mapper.fit_transform(factor_container.data.loc[date_]) for date_ in
#                      tiaocang_date]
#         calc_data_agg = np.vstack(calc_data)
#         factor_container.data = pd.DataFrame(calc_data_agg,
#                                              index=factor_container.data.index,
#                                              columns=factor_container.data.columns)
#         if self.out_container:
#             return factor_container
#         else:
#             return factor_container.data
# 
# Path: alphaware/enums.py
# class NAStrategy(StrEnum):
#     MEAN = 'mean'
#     MEDIAN = 'median'
#     MOST_FREQ = 'most_frequent'
#     CUSTOM = 'custom'
# 
# class FactorType(StrEnum):
#     INDUSTY_CODE = 'industry_code'
#     ALPHA_FACTOR = 'alpha_factor'
#     RETURN = 'return'
#     PRICE = 'price'
#     ALPHA_FACTOR_MV = 'alpha_factor_mv'
#     SCORE = 'score'
#     FWD_RETURN = 'fwd_return'
#     WEIGHT = 'weight'


import numpy as np
import pandas as pd
from argcheck import expect_element
from sklearn.preprocessing import Imputer
from sklearn.utils.validation import check_is_fitted
from sklearn_pandas import (CategoricalImputer,
                            DataFrameMapper)
from ..base import FactorTransformer
from ..enums import (NAStrategy,
                     FactorType)
    Compute the boolean mask X == missing_values.
    """
    if value == 'NaN' or value is None or (isinstance(value, float) and np.isnan(value)):
        return pd.isnull(X)
    else:
        return X == value


class ExtCategoricalImputer(CategoricalImputer):
    """
    ref: https://github.com/pandas-dev/sklearn-pandas/blob/master/sklearn_pandas/categorical_imputer.py
    """

    @expect_element(strategy=(NAStrategy.CUSTOM, NAStrategy.MOST_FREQ))
    def __init__(self, missing_value='NaN', copy=True, custom_value=None, strategy=NAStrategy.MOST_FREQ):
        super(ExtCategoricalImputer, self).__init__(missing_values=missing_value, copy=copy)
        self.missing_values = missing_value
        self.copy = copy
        self.strategy = strategy
        self.custom_value = custom_value
        self.fill_ = None

    def fit(self, X, y=None):
        if self.strategy == NAStrategy.MOST_FREQ:
            return super(ExtCategoricalImputer, self).fit(X, y)
        elif self.strategy == NAStrategy.CUSTOM:
            self.fill_ = self.custom_value
            return self

    def transform(self, X):

# Path: alphaware/base/factor_transformer.py
# class FactorTransformer(BaseEstimator, TransformerMixin):
#     def __init__(self, copy=True, out_container=False):
#         self.copy = copy
#         self.df_mapper = None
#         self.out_container = out_container
# 
#     @expect_types(out=bool)
#     def set_out_container(self, out):
#         self.out_container = out
# 
#     @preprocess(factor_container=ensure_factor_container)
#     def fit(self, factor_container, **kwargs):
#         self.df_mapper = self._build_mapper(factor_container)
#         return self
# 
#     def _build_mapper(self, factor_container):
#         raise NotImplementedError
# 
#     @preprocess(factor_container=ensure_factor_container)
#     def transform(self, factor_container):
#         if self.copy:
#             factor_container = copy.deepcopy(factor_container)
#         tiaocang_date = factor_container.tiaocang_date
#         calc_data = [self.df_mapper.fit_transform(factor_container.data.loc[date_]) for date_ in
#                      tiaocang_date]
#         calc_data_agg = np.vstack(calc_data)
#         factor_container.data = pd.DataFrame(calc_data_agg,
#                                              index=factor_container.data.index,
#                                              columns=factor_container.data.columns)
#         if self.out_container:
#             return factor_container
#         else:
#             return factor_container.data
# 
# Path: alphaware/enums.py
# class NAStrategy(StrEnum):
#     MEAN = 'mean'
#     MEDIAN = 'median'
#     MOST_FREQ = 'most_frequent'
#     CUSTOM = 'custom'
# 
# class FactorType(StrEnum):
#     INDUSTY_CODE = 'industry_code'
#     ALPHA_FACTOR = 'alpha_factor'
#     RETURN = 'return'
#     PRICE = 'price'
#     ALPHA_FACTOR_MV = 'alpha_factor_mv'
#     SCORE = 'score'
#     FWD_RETURN = 'fwd_return'
#     WEIGHT = 'weight'


# Path: alphaware/preprocess/imputer.py
import numpy as np
import pandas as pd
from argcheck import expect_element
from sklearn.preprocessing import Imputer
from sklearn.utils.validation import check_is_fitted
from sklearn_pandas import (CategoricalImputer,
                            DataFrameMapper)
from ..base import FactorTransformer
from ..enums import (NAStrategy,
                     FactorType)

    Compute the boolean mask X == missing_values.
    """
    if value == 'NaN' or value is None or (isinstance(value, float) and np.isnan(value)):
        return pd.isnull(X)
    else:
        return X == value


class ExtCategoricalImputer(CategoricalImputer):
    """
    ref: https://github.com/pandas-dev/sklearn-pandas/blob/master/sklearn_pandas/categorical_imputer.py
    """

    @expect_element(strategy=(NAStrategy.CUSTOM, NAStrategy.MOST_FREQ))
    def __init__(self, missing_value='NaN', copy=True, custom_value=None, strategy=NAStrategy.MOST_FREQ):
        super(ExtCategoricalImputer, self).__init__(missing_values=missing_value, copy=copy)
        self.missing_values = missing_value
        self.copy = copy
        self.strategy = strategy
        self.custom_value = custom_value
        self.fill_ = None

    def fit(self, X, y=None):
        if self.strategy == NAStrategy.MOST_FREQ:
            return super(ExtCategoricalImputer, self).fit(X, y)
        elif self.strategy == NAStrategy.CUSTOM:
            self.fill_ = self.custom_value
            return self

    def transform(self, X):

        if self.strategy == NAStrategy.MOST_FREQ: