from sklearn.base import BaseEstimator, TransformerMixin
import utils
import abc
import pandas as pd
from constants import variable_type,column_names,NO_UNITS,ALL
import logger




class safe_unstacker(BaseEstimator,TransformerMixin):

    def __init__(self, *levels):
        self.levels = levels

    def fit(self, x, y=None):
        return self

    def transform(self, df):
        return safe_unstack(df,self.levels)

def safe_unstack(df,levels):
    subindex = 'subindex'
    #add subindex to facilitate unstacking
    df = utils.add_subindex(df,subindex)

    #unstack!
    df_unstacked = df.unstack(levels)

    #drop "value" level, which is derivative from value column that is being unstacked against
    df_unstacked.columns = df_unstacked.columns.droplevel(0)

    # Drop subindex
    df_unstacked.index = df_unstacked.index.droplevel(subindex)

    df_unstacked.dropna(axis=1,inplace=True,how='all')
    return df_unstacked

class add_level(BaseEstimator,TransformerMixin):
        def __init__(self,level_val,level_name,axis=0):
            self.level_val = level_val
            self.level_name = level_name
            self.axis = axis

        def fit(self, x, y=None):
            return self

        def transform(self, df):
            return utils.add_same_val_index_level(df,self.level_val,self.level_name,self.axis)

class column_standardizer(BaseEstimator,TransformerMixin):

    def __init__(self,data_dict,ureg,convert_units=True):
        self.data_dict = data_dict
        self.ureg = ureg
        self.convert_units=convert_units

    def fit(self, x, y=None):
        return self

    def transform(self, df):
        df = df.copy()
        col_cnt = df.columns.size
        if col_cnt == 0: return df
        names = ['component','status','variable_type','units','description']
        tuples=[]
        for col_ix in range(0,col_cnt):
            col = df.iloc[:,col_ix]
            new_col,new_name = self.standardize(col)
            df.iloc[:,col_ix] = new_col
            tuples.append(map(str,new_name))
        df.columns = pd.MultiIndex.from_tuples(tuples,names=names)
        df.sort_index(axis=1, inplace=True)
        return df

    def standardize(self,col):
        old_col_name = col.name
        guess_component = old_col_name[0]
        units = old_col_name[-2]
        desc = old_col_name[-1]
        dtype = col.dtype
        defs = self.data_dict.tables.definitions
        defs = defs[defs.component == guess_component]
        best_def = None
        for ix,row in defs.iterrows():
            def_units = row['units']
            if can_convert(def_units,units,self.ureg):
                best_def = row
                break

        if (best_def is None) and (dtype != pd.np.object):
            status = 'unknown'
            var_type = variable_type.QUANTITATIVE
        elif (best_def is None) or ((best_def['variable_type'] == variable_type.QUANTITATIVE) & (dtype == pd.np.object)):
            status = 'unknown'
            var_type = variable_type.NOMINAL
            if units != NO_UNITS:
                desc = utils.append_to_description(desc,units)
                units = NO_UNITS
        else:
            status = 'known'
            var_type = best_def['variable_type']
            new_units = best_def['units']
            if new_units != units:
                if not self.ureg.same_units(units,new_units) and self.convert_units:
                    col = self.ureg.convert_units(units,new_units,col)
                desc = utils.append_to_description(str(desc),units)
                units = new_units



        return (col,(guess_component,status,var_type,units,desc))

def can_convert(unit1,unit2,med_ureg):
    if (unit1 == unit2): return True
    if (NO_UNITS in [unit1,unit2]): return False
    return med_ureg.same_dimensionality(unit1,unit2)

class oob_value_remover(BaseEstimator,TransformerMixin):
    def __init__(self,data_dict):
        self.data_dict = data_dict

    def fit(self, x, y=None):
        return self

    def transform(self, df):
        logger.log('Drop OOB data | {}'.format(df.shape),new_level=True)
        df = df.copy()
        idx = pd.IndexSlice
        df = df.sort_index(axis=1).sort_index()
        for component in df.columns.get_level_values('component').unique().tolist():
            component_defs = self.data_dict.defs_for_component(component)
            for units in df[component].columns.get_level_values(column_names.UNITS).unique().tolist():
                df_slice = df.loc[:,idx[component,:,:,units,:]]
                logger.log('{}, {}, {}'.format(component,units,df_slice.count().sum()))
                matching_defs = component_defs[(component_defs.units == units)]
                if matching_defs.empty: continue
                def_row = matching_defs.iloc[0]
                lower = def_row['lower']
                upper = def_row['upper']
                df.loc[:,idx[component,:,:,units,:]] = remove_oob_values(df_slice,lower,upper)
        df.dropna(how='all',inplace=True,axis=1)
        logger.end_log_level()
        return df

def remove_oob_values(data,lower,upper):
    oob_mask = (data < lower) | (data > upper)
    return data[~oob_mask]




class split_dtype(BaseEstimator,TransformerMixin):

    def fit(self, x, y=None):
        return self

    def transform(self, df):
        if df.empty: return df
        df_numeric  = df.apply(pd.to_numeric,errors='coerce')
        is_string = pd.isnull(df_numeric) & ~pd.isnull(df)

        df_string = df[is_string].dropna(how='all')
        tuples = [(col_name[0],NO_UNITS,utils.append_to_description(*map(str,col_name[3:0:-1]))) for col_name in df_string.columns]
        df_string.columns = pd.MultiIndex.from_tuples(tuples,names = df_string.columns.names)
        df_string = utils.add_same_val_index_level(df_string,level_val='string',level_name='dtype',axis=1)

        df_numeric = df_numeric.dropna(how='all')
        df_numeric = utils.add_same_val_index_level(df_numeric,level_val='number',level_name='dtype',axis=1)

        df_joined = df_numeric.join(df_string,how='outer')
        del df_string,df_numeric

        df_joined.columns = df_joined.columns.droplevel('dtype')
        df_joined.dropna(how='all',inplace=True,axis=1)
        return df_joined


class combine_like_cols(BaseEstimator,TransformerMixin):
    def fit(self, df, y=None, **fit_params):
        logger.log('FIT Combine like columns {}'.format(df.shape),new_level=True)

        self.columns_to_combine = {}
        groupby_cols = list(df.columns.names)
        groupby_cols.remove(column_names.DESCRIPTION)
        grouped = df.groupby(level=groupby_cols,axis=1)

        column_list = []
        df_out=None
        for index,group in grouped:
            index
            logger.log(index)
            if index[2] == variable_type.NOMINAL: continue

            ordered_cols = group[group.count().sort_values(ascending=False).index.tolist()].columns.tolist()
            self.columns_to_combine[index] = ordered_cols

        logger.end_log_level()
        return self

    def transform(self, df):
        logger.log('TRANSFORM Combine like columns {}'.format(df.shape),new_level=True)

        column_list = []
        for index,columns in self.columns_to_combine.iteritems():
            logger.log(index)
            df_list=[]
            for col_name in columns:
                if col_name not in df.columns:
                    df[col_name] = pd.np.nan
                col = df[col_name].dropna()
                col.name = index + (ALL,)
                df_list.append(col)

            df_combined = pd.concat(df_list).to_frame()

            # Here we will drop all duplicate values; since we sort the max col first,
            # BEFORE we loop and combine, we will be prioritizing all values from the max value
            # column. Although this may be a change in style from previous, it is easy, and will
            # most of the time be RIGHT.
            duplicates_to_drop = df_combined.index.duplicated(keep='first')
            df_combined = df_combined.loc[~duplicates_to_drop]

            #drop the combined columns
            df.drop(columns,axis=1,inplace=True)

            #join the combined column back to the DF
            df = df.join(df_combined,how='outer')

        df.columns.names = df.columns.names
        df.sort_index(inplace=True)
        df.sort_index(inplace=True,axis=1)

        logger.end_log_level()


        return df

class flatten_index(BaseEstimator,TransformerMixin):
        def __init__(self,axis=0,suffix=None):
            self.axis=axis
            self.suffix=suffix

        def fit(self, x, y=None):
            return self

        def transform(self, df):
            df = utils.flatten_index(df,axis=self.axis,suffix=self.suffix)
            return df


"""
Deal with categorical data
"""
class standardize_categories(BaseEstimator,TransformerMixin):

    def __init__(self,data_dict,category_map,use_numeric=True):
        self.data_dict = data_dict
        self.category_map = category_map
        self.use_numeric = use_numeric

    def fit(self, x, y=None):
        return self

    def transform(self, df):
        for component in utils.get_components(df):
            cat_map = self.category_map.get(component,None)
            if cat_map is None: continue
            df_slice = df.loc[:,[component]]
            categorical_mask = df_slice.columns.get_level_values('variable_type').isin([variable_type.NOMINAL,variable_type.ORDINAL])
            df_categories = self.data_dict.tables.categories
            to_replace = cat_map.keys()
            col = 'val_numeric' if self.use_numeric else 'val_text'
            values = [df_categories.loc[cat_ix,col] for cat_ix in cat_map.values()]

            df_slice.loc[:,categorical_mask] = df_slice.loc[:,categorical_mask].replace(to_replace=to_replace,value=values)
            if not self.use_numeric:
                to_replace = [df_categories.loc[cat_ix,'val_numeric'] for cat_ix in cat_map.values()]
                df_slice.loc[:,categorical_mask] = df_slice.loc[:,categorical_mask].replace(to_replace=to_replace,value=values)
            df.loc[:,[component]] = df_slice
        return df

class split_bad_categories(BaseEstimator,TransformerMixin):

    def __init__(self,data_dict,use_numeric=True):
        self.data_dict = data_dict
        self.use_numeric = use_numeric

    def fit(self, x, y=None):
        return self

    def transform(self, df):
        for component in utils.get_components(df):
            df_categories = self.data_dict.get_categories(component)
            if df_categories is None: continue
            df_slice = df.loc[:,[component]]
            col = 'val_numeric' if self.use_numeric else 'val_text'
            valid_values = df_categories.loc[:,col]

            categorical_mask = df_slice.columns.get_level_values('variable_type').isin([variable_type.NOMINAL,variable_type.ORDINAL])
            categorical_slice = df_slice.loc[:,categorical_mask]

            df_valid_mask  = categorical_slice.apply(lambda x: x.isin(valid_values))

            df_slice.loc[:,categorical_mask] = categorical_slice[df_valid_mask]
            df.loc[:,[component]] = df_slice

            df_invalid = categorical_slice[~df_valid_mask]
            df_invalid.columns = utils.set_level_to_same_val(df_invalid.columns,'status','unknown')
            df_invalid.columns = utils.set_level_to_same_val(df_invalid.columns,'variable_type',variable_type.NOMINAL)
            df = df.join(df_invalid,how='outer')
            del df_invalid
        df.dropna(how='all',inplace=True,axis=1)
        return df

class nominal_to_onehot(BaseEstimator,TransformerMixin):

    def fit(self, x, y=None):
        return self

    def transform(self, df):
        if df.empty: return df

        logger.log('Nominal to OneHot',new_level=True)
        nominal_cols = df.columns.get_level_values('variable_type') == variable_type.NOMINAL

        for col_name in df.loc[:,nominal_cols]:
            column = df[col_name]
            df.drop(col_name,axis=1,inplace=True)
            df_dummies = pd.get_dummies(column)
            if df_dummies.empty: continue
            dummy_col_names = [col_name[:-1] + ('{}_{}'.format(col_name[-1],text),) for text in df_dummies.columns]
            df_dummies.columns = pd.MultiIndex.from_tuples(dummy_col_names,names=df.columns.names)
            df = df.join(df_dummies,how='outer')
        logger.end_log_level()
        return df


"""
Duplicate index aggregators
"""

class same_index_aggregator(BaseEstimator,TransformerMixin):

    def __init__(self,agg_func):
        self.agg_func = agg_func

    def fit(self, x, y=None):
        return self

    def transform(self, df):

        duplicated = df.index.duplicated(keep=False)

        df_safe = df[~duplicated]
        df_duplicated = df[duplicated]

        df_fixed = df_duplicated.groupby(level=df_duplicated.index.names).agg(lambda x:self.agg_func(x))

        df_no_dups = pd.concat([df_safe,df_fixed])
        df_no_dups.sort_index(inplace=True)
        return df_no_dups

"""
Fill NA
"""

class NaNFiller(BaseEstimator,TransformerMixin):

    def fit(self, X, y, **fit_params):
        self.fill_vals = self.get_fill_vals(X, y, **fit_params)
        return self

    def transform(self,df):
        return df.apply(lambda col: col.fillna(self.fill_vals[col.name]))

    def get_fill_vals(self, X, y, **fit_params):
        return pd.Series(np.NaN,index=X.columns)

class FillerZero(NaNFiller):

    def get_fill_vals(self, X, y, **fit_params):
        return pd.Series(0,index=X.columns)

class FillerMean(NaNFiller):

    def get_fill_vals(self, X, y, **fit_params):
        return X.mean()

class FillerMode(NaNFiller):

    def get_fill_vals(self, X, y, **fit_params):
        return X.mode().iloc[0]


class do_nothing(BaseEstimator,TransformerMixin):

    def fit(self, x, y=None):
        return self

    def transform(self, df):
        return df

class GroupbyAndFFill(BaseEstimator,TransformerMixin):
        def __init__(self,level=None,by=None):
            self.level=level
            self.by=by

        def fit(self, x, y=None):
            return self

        def transform(self, df):
            return df.groupby(level=self.level,by=self.by).ffill()

class GroupbyAndBFill(BaseEstimator,TransformerMixin):
        def __init__(self,level=None,by=None):
            self.level=level
            self.by=by

        def fit(self, x, y=None):
            return self

        def transform(self, df):
            return df.groupby(level=self.level,by=self.by).bfill()


"""
filtering
"""


class column_filter(BaseEstimator,TransformerMixin):

    def fit(self, df, y=None, **fit_params):
        logger.log('*fit* Filter columns ({}) {}'.format(self.__class__.__name__, df.shape).format(self.__class__),new_level=True)
        if df.empty:
            self.cols_to_keep = []
        else:
            self.cols_to_keep = self.get_columns_to_keep(df, y, **fit_params)
        logger.end_log_level()
        return self

    def transform(self, df):
        logger.log('*transform* Filter columns ({}) {}'.format(self.__class__.__name__, df.shape))
        df_out = None
        if df.empty or len(self.cols_to_keep) == 0: df_out = df.drop(df.columns,axis=1)
        else: df_out = df.loc[:,self.cols_to_keep]
        logger.log(end_prev=True)
        return df_out

    def get_columns_to_keep(self,df, y=None, **fit_params):
        return df.columns

class DataSpecFilter(column_filter):

    def __init__(self,data_specs):
        self.data_specs = data_specs

    def get_columns_to_keep(self, df, y=None, **fit_params):

        df_cols = pd.DataFrame(map(list,df.columns.tolist()),columns=df.columns.names)

        mask = utils.complex_row_mask(df_cols,self.data_specs)

        return [tuple(x) for x in df_cols[mask].to_records(index=False)]

class max_col_only(column_filter):
    def get_columns_to_keep(self, df, y=None, **fit_params):
        self.max_col =  df.apply(utils.smart_count).sort_values().index.tolist()[-1]
        return [self.max_col]


class remove_small_columns(column_filter):

    def __init__(self,threshold):
        self.threshold = threshold

    def get_columns_to_keep(self, df, y=None, **fit_params):
        return df.loc[:,df.apply(utils.smart_count) > self.threshold].columns


class multislice_filter(column_filter):

    def __init__(self,slice_dict_list):
        self.slice_dict_list = slice_dict_list

    def get_columns_to_keep(self,df, y=None, **fit_params):

        cols = []
        for slice_dict in self.slice_dict_list:
            levels = slice_dict.keys()
            vals = slice_dict.values()
            cols += df.xs(vals,level=levels,axis=1,drop_level=False).columns.tolist()


        return cols

class DataNeedsFilter(multislice_filter):

    def __init__(self,data_needs):
        comp_dict = {}
        for dn in data_needs:
            component = dn[0]
            units = dn[1]
            units_list = comp_dict.get(component,[])
            units_list.append(units)

            comp_dict[component] = units_list

        slice_dict_list = []
        for component,units_list in comp_dict.iteritems():
            if ALL in units_list:
                slice_dict_list.append({column_names.COMPONENT: component})
                continue
            for unit in units_list:
                slice_dict_list.append({
                            column_names.COMPONENT: component,
                            column_names.UNITS : units
                        })
        super(DataNeedsFilter,self).__init__(slice_dict_list)

class func_filter(column_filter):

    def __init__(self,filter_func):
        self.filter_func = filter_func

    def get_columns_to_keep(self,df, y=None, **fit_params):
        return df.loc[:,df.apply(self.filter_func)].columns


class record_threshold(func_filter):

    def __init__(self,threshold):
        self.threshold = threshold
        filter_func = lambda col: col.dropna().index.get_level_values(column_names.ID).unique().size > self.threshold
        super(record_threshold,self).__init__(filter_func)


class drop_all_nan_cols(func_filter):

    def __init__(self):
        filter_func = lambda col: ~pd.isnull(col).all()
        super(drop_all_nan_cols,self).__init__(filter_func)


class known_col_only(func_filter):

    def __init__(self):
        filter_func = lambda col: col.name[1] == 'known'
        super(known_col_only,self).__init__(filter_func)

class filter_to_component(func_filter):
    def __init__(self,components):
        self.components = components
        filter_func = lambda col: col.name[0] in self.components
        super(filter_to_component,self).__init__(filter_func)

class filter_var_type(func_filter):

    def __init__(self,var_types):
        self.var_types =var_types
        filter_func = lambda col: col.name[2] in self.var_types
        super(filter_var_type,self).__init__(filter_func)

class summable_only(func_filter):

    def __init__(self,ureg,ignore_component_list):
        self.ureg = ureg
        self.ignore_component_list = ignore_component_list
        filter_func = lambda col:summable_only_filter(col,self.ureg,self.ignore_component_list)
        super(summable_only,self).__init__(filter_func)

def summable_only_filter(col,ureg,ignore_component_list):
    is_summable_unit = lambda col: (col.name[-2] != NO_UNITS) and (ureg.is_volume(str(col.name[-2])) or ureg.is_mass(str(col.name[-2])))
    should_ignore_component = lambda col: (col.name[0] in ignore_component_list)
    return lambda col: is_summable_unit(col.name) and not should_ignore_component(col.name)

class DropNaN(BaseEstimator,TransformerMixin):

    def __init__(self,axis=0,how='any',thresh=None):
        self.axis=axis
        self.how=how
        self.thresh=thresh

    def fit(self, df, y=None):
        return self

    def transform(self, df):
        return df.dropna(axis=self.axis,how=self.how,thresh=self.thresh)

class filter_ids(BaseEstimator,TransformerMixin):

    def __init__(self,print_loss=False,ids=None):
        self.print_loss = print_loss
        self.ids = ids

    def fit(self, x, y=None, **fit_params):
        if self.ids is None:
            ids = fit_params.get('ids',None)
            if (ids is None) and (y is not None):
                ids = y.index.get_level_values(column_names.ID).unique().tolist()
            self.ids = ids
        return self

    def transform(self, df):
        if self.ids is not None:
            out_df = df.loc[df.index.get_level_values(column_names.ID).isin(self.ids)]
        else: out_df = df
        if self.print_loss:
            print 'Data Loss:',utils.data_loss(df,out_df)
        return out_df

class more_than_n_component(BaseEstimator,TransformerMixin):

    def __init__(self,n,component):
        self.n = n
        self.component = component

    def fit(self, df, y=None):
        return self

    def transform(self, df):
        if df.empty: return df.drop(df.index)
        good_ids = df.loc[:,[self.component]].dropna(how='all').groupby(level=column_names.ID).count() > self.n
        good_ids = good_ids.loc[good_ids.iloc[:,0]].index.unique().tolist()
        return df.loc[df.index.get_level_values(column_names.ID).isin(good_ids)]

"""
Simple Data Manipulation
"""

class TimeShifter(TransformerMixin,BaseEstimator):

    def __init__(self,datetime_level,shift='infer',n=1):
        self.shift=shift
        self.datetime_level = datetime_level
        self.n=n

    def fit(self, X, y=None, **fit_params):
        return self

    def transform(self, df):
        shift = self.shift
        if shift == 'infer':
            infer_freq = lambda grp: grp.index.get_level_values(self.datetime_level).inferred_freq
            inferred_freqs = df.groupby(level=column_names.ID).apply(infer_freq)
            shift = inferred_freqs.value_counts().sort_values().index[-1]
        df = df.reset_index(level=self.datetime_level)
        df.loc[:,self.datetime_level] = df.loc[:,self.datetime_level] + self.n*pd.Timedelta(shift)
        df.set_index(self.datetime_level,append=True,inplace=True)
        return df

class RowShifter(TransformerMixin,BaseEstimator):

    def __init__(self,n):
        self.n=n

    def fit(self, X, y=None, **fit_params):
        return self

    def transform(self, df):
        return df.shift(self.n)

class Replacer(TransformerMixin,BaseEstimator):

    def __init__(self,to_replace=None, value=None, regex=False, method='pad'):
        self.to_replace = to_replace
        self.value = value
        self.regex=regex
        self.method = method

    def fit(self, X, y=None, **fit_params):
        return self

    def transform(self, df):
        return df.replace(
                    to_replace=self.to_replace,
                    value=self.value,
                    regex=self.regex,
                    method=self.method
                )
class Delta(TransformerMixin,BaseEstimator):

    def fit(self, X, y=None, **fit_params):
        return self

    def transform(self, df):

        df_last = df.ffill().dropna(how='any')
        df_last = utils.add_same_val_index_level(df_last,'last','temp',axis=1)


        df_next = df.shift(-1).dropna(how='any')
        df_next = utils.add_same_val_index_level(df_next,'next','temp',axis=1)

        df_all = df_last.join(df_next,how='inner')
        return df_all.loc[:,'next'] - df_all.loc[:,'last']

class ToGroupby(TransformerMixin,BaseEstimator):

    def __init__(self, by=None, axis=0, level=None, as_index=True):
        self.by=by
        self.axis=axis
        self.level=level
        self.as_index = as_index

    def fit(self, X, y=None, **fit_params):
        return self

    def transform(self, df):
        return df.groupby(by=self.by, axis=self.axis, level=self.level, as_index=self.as_index)