#
# Build regressor for numeric fields
# Use time and categorical fields when available 
#
import pandas as pd
import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt
import pickle
import os.path

from sklearn.metrics import mean_squared_error
from sklearn.metrics.regression import r2_score
from sklearn.model_selection import StratifiedKFold

from include.dataset_fnames import generate_station_data_fname, generate_data_fname, generate_response_data_fname
from include.feature_lists import numeric_features, numeric_missing_features_list, numeric_features_to_work_on, categoric_features

from xgboost import XGBRegressor
from catboost import CatBoostRegressor
from sklearn.ensemble import RandomForestRegressor
from sklearn.neighbors import KNeighborsRegressor

from include.regressor_fnames import generate_regressor_fname, save_regressor, load_regressor

from include import SimpleMeanRegressor, SimpleMedianRegressor, SimpleModeRegressor

split_count = 3

def xgb_regressor_wrapper(X_train, y_train):
    xgb_regressor = XGBRegressor(objective='reg:linear', n_estimators=1000, learning_rate=0.01, base_score=0.005, n_jobs=4)
    xgb_regressor.fit(X=X_train, y=y_train) #, eval_set=[(X_test, y_test)], verbose=True)
    return xgb_regressor

def cb_regressor_wrapper_i1500(X_train, y_train):
    cb_regressor = CatBoostRegressor(iterations=1500, thread_count=8)    
    cb_regressor.fit(X=X_train, y=y_train) #, verbose=True) 
    return cb_regressor

def cb_regressor_wrapper_i1000(X_train, y_train):
    cb_regressor = CatBoostRegressor(iterations=1000, thread_count=8)    
    cb_regressor.fit(X=X_train, y=y_train) #, verbose=True) 
    return cb_regressor

def cb_regressor_wrapper_i1000_lr001(X_train, y_train):
    cb_regressor = CatBoostRegressor(iterations=1000, learning_rate=0.01, thread_count=8)    
    cb_regressor.fit(X=X_train, y=y_train) #, verbose=True) 
    return cb_regressor

def cb_regressor_wrapper_i500(X_train, y_train):
    cb_regressor = CatBoostRegressor(iterations=500, thread_count=8)    
    cb_regressor.fit(X=X_train, y=y_train) #, verbose=True) 
    return cb_regressor

def cb_regressor_wrapper_i500_lr001(X_train, y_train):
    cb_regressor = CatBoostRegressor(iterations=500, learning_rate=0.01, thread_count=8)    
    cb_regressor.fit(X=X_train, y=y_train) #, verbose=True) 
    return cb_regressor

def cb_regressor_wrapper_i500_d5(X_train, y_train):
    cb_regressor = CatBoostRegressor(iterations=500, depth=5, thread_count=8)    
    cb_regressor.fit(X=X_train, y=y_train) #, verbose=True) 
    return cb_regressor

def rf_regressor_wrapper(X_train, y_train):
    
    df_train = pd.DataFrame(data=X_train)
    df_train['y'] = y_train

    df_train.dropna(how='any', inplace=True)
    
    ytrain = df_train['y'].values
    del df_train['y'] 
    Xtrain = df_train.values
    
    rf_regressor = RandomForestRegressor(n_estimators=1000, max_features='sqrt', max_depth=6, n_jobs=2)    
    rf_regressor.fit(X=Xtrain, y=ytrain) #, verbose=True) 
    return rf_regressor

def knn_regressor_wrapper_15(X_train, y_train):
    knn_regressor = KNeighborsRegressor(n_neighbors=15)    
    knn_regressor.fit(X=X_train, y=y_train) #, verbose=True) 
    return knn_regressor

def knn_regressor_wrapper_7(X_train, y_train):
    knn_regressor = KNeighborsRegressor(n_neighbors=7)    
    knn_regressor.fit(X=X_train, y=y_train) #, verbose=True) 
    return knn_regressor

def knn_regressor_wrapper_9(X_train, y_train):
    knn_regressor = KNeighborsRegressor(n_neighbors=9)    
    knn_regressor.fit(X=X_train, y=y_train) #, verbose=True) 
    return knn_regressor

def knn_regressor_wrapper_13(X_train, y_train):
    knn_regressor = KNeighborsRegressor(n_neighbors=13)    
    knn_regressor.fit(X=X_train, y=y_train) #, verbose=True) 
    return knn_regressor

def mean_regressor_wrapper(X_train, y_train):
    mean_regressor = SimpleMeanRegressor.SimpleMeanRegressor(rtype='mean')    
    mean_regressor.fit(X=None, y=y_train) #, verbose=True) 
    return mean_regressor

def median_regressor_wrapper(X_train, y_train):
    median_regressor = SimpleMedianRegressor.SimpleMedianRegressor(rtype='median')    
    median_regressor.fit(X=None, y=y_train) #, verbose=True) 
    return median_regressor

def mode_regressor_wrapper(X_train, y_train):
    mode_regressor = SimpleModeRegressor.SimpleModeRegressor(rtype='mode')    
    mode_regressor.fit(X=None, y=y_train) #, verbose=True) 
    return mode_regressor

regressor_list = {
#     'XGB_Regressor': xgb_regressor_wrapper,
    'CB_Regressor_i1500': cb_regressor_wrapper_i1500,
#     'CB_Regressor_i1000': cb_regressor_wrapper_i1000,
#     'CB_Regressor_i1000_lr001': cb_regressor_wrapper_i1000_lr001,
#     'CB_Regressor_i500': cb_regressor_wrapper_i500,
#     'CB_Regressor_i500_lr001': cb_regressor_wrapper_i500_lr001,
#     'CB_Regressor_i500_d5': cb_regressor_wrapper_i500_d5,
#     'RF_Regressor': rf_regressor_wrapper,
#     'KNN_Regressor_13': knn_regressor_wrapper_13,
#     'KNN_Regressor_15': knn_regressor_wrapper_15,
#     'KNN_Regressor_9': knn_regressor_wrapper_9,
#     'Mean_Regressor': mean_regressor_wrapper,
#     'Median_Regressor': median_regressor_wrapper,
#     'Mode_Regressor': mode_regressor_wrapper,
}

prev_station = {
    'L0S00' : '', 'L0S01' : 'L0S00', 'L0S02' : 'L0S01', 'L0S03' : 'L0S02', 'L0S04' : 'L0S03', 'L0S05' : 'L0S04', 'L0S06' : 'L0S05', 
    'L0S07' : 'L0S06', 'L0S08' : 'L0S07', 'L0S09' : 'L0S08', 'L0S10' : 'L0S09', 'L0S11' : 'L0S10', 'L0S12' : 'L0S11', 'L0S13' : 'L0S12', 
    'L0S14' : 'L0S13', 'L0S15' : 'L0S14', 'L0S16' : 'L0S15', 'L0S17' : 'L0S16', 'L0S18' : 'L0S17', 'L0S19' : 'L0S18', 'L0S20' : 'L0S19', 
    'L0S21' : 'L0S20', 'L0S22' : 'L0S21', 'L0S23' : 'L0S22', 'L1S24' : 'L0S23', 'L1S25' : 'L1S24', 'L2S26' : 'L1S25', 'L2S27' : 'L2S26', 
    'L2S28' : 'L2S27', 'L3S29' : 'L2S28', 'L3S30' : 'L3S29', 'L3S31' : 'L3S30', 'L3S32' : 'L3S31', 'L3S33' : 'L3S32', 'L3S34' : 'L3S33', 
    'L3S35' : 'L3S34', 'L3S36' : 'L3S35', 'L3S37' : 'L3S36', 'L3S38' : 'L3S37', 'L3S39' : 'L3S38', 'L3S40' : 'L3S39', 'L3S41' : 'L3S40', 
    'L3S42' : 'L3S41', 'L3S43' : 'L3S42', 'L3S44' : 'L3S43', 'L3S45' : 'L3S44', 'L3S46' : 'L3S45', 'L3S47' : 'L3S46', 'L3S48' : 'L3S47', 
    'L3S49' : 'L3S48', 'L3S50' : 'L3S49', 'L3S51' : 'L3S50'
    }

def load_station_df(station_id, use_time=False, use_null_tracker=False, use_category=False, use_prev_station=False, use_product_id=False, use_response=False):
   
    fname = generate_station_data_fname(station_id, sample_type='train', data_type='numeric', use_product=False, allow_nan_values=False)
    print fname
    station_features = numeric_features[station_id]    
    features = ['Id'] + station_features
    if use_time:
        features = features + ['time']
        
    station_df = pd.read_csv(fname, usecols=features, index_col=['Id'])

    if use_null_tracker:
        for feature in station_features:
            new_column = feature + '_isnull'
            null_indices = station_df[station_df[feature].isnull()].index
            station_df[new_column] = 0
            station_df.loc[null_indices, new_column] = 1

    if use_product_id:
        fname2 = generate_station_data_fname(station_id, sample_type='train', data_type='numeric', use_product=True, allow_nan_values=False)
        print fname2
        features = ['Id'] + ['product']
        station_df2 = pd.read_csv(fname2, usecols=features, index_col=['Id'], dtype=object)
    
        station_df['product'] = station_df2['product']

#     print station_df.head()
    
    if use_category:
        fname2 = generate_station_data_fname(station_id, sample_type='train', data_type='categorical', use_product=False, allow_nan_values=False)
        print fname2
        station_features = categoric_features[station_id]    
        features = ['Id'] + station_features
        station_df2 = pd.read_csv(fname2, usecols=features, index_col=['Id'], dtype=object)
#         print station_df2.head()

        indices = station_df.index
        station_df = pd.concat([station_df, station_df2], axis=1)
        station_df = station_df.loc[indices]
        
        station_df = pd.get_dummies(station_df)

    if use_prev_station:
        if station_id != 'L0S00': 
            prev_id = prev_station[station_id]
    
            fname2 = generate_station_data_fname(prev_id, sample_type='train', data_type='numeric', use_product=False, allow_nan_values=False)
            print fname2
            station_features = numeric_features[prev_id]    
            features = ['Id'] + station_features
            station_df2 = pd.read_csv(fname2, usecols=features, index_col=['Id'], dtype=object)
            
            indices = station_df.index
            station_df = pd.concat([station_df, station_df2], axis=1)
            station_df = station_df.loc[indices]
            
#     print station_df.head()
#     print station_df.info()
    
#     exit(-1)


    fname = generate_response_data_fname()
    field_list = ['Id'] + ['Response']
    response_df = pd.read_csv(fname, usecols=field_list, index_col=['Id'], dtype='int')
    indices = station_df.index
    response_df = response_df.loc[indices]
    if use_response:
        station_df = pd.concat([station_df, response_df], axis=1)
#         station_df = station_df.loc[indices]
        
#     print response_df['Response'].values


    return station_df, response_df



def regressor_for_station_feature(station_id, feature, X, y, response):
    assert len(X) == len(response)
  
    skf = StratifiedKFold(n_splits=split_count, shuffle=False, random_state=False)
    
#     split_idx = len(y) / 3
#     
#     X_train = X[:-split_idx,:]
#     y_train = y[:-split_idx]
#     
#     X_test = X[-split_idx:,:]
#     y_test = y[-split_idx:]


#     fname = generate_station_data_fname(station_id, sample_type='test', data_type='numeric', allow_nan_values=False)
#     print fname
#     station_features = numeric_features[station_id]    
#     features = ['Id'] + station_features
#     station_df = pd.read_csv(fname, usecols=features, index_col=['Id'], dtype='float32')
#     dataset_df = station_df[station_df[feature].notnull()]
#     dataset_y_df = dataset_df[feature]
#     del dataset_df[feature]
#     dataset_X_df = dataset_df
# 
#     X_new = dataset_X_df.values
#     y_new = dataset_y_df.values

    best_regressor = None
    best_regressor_name = None
    best_r2_score = -9999
    for name in regressor_list:
#         print name,
        regressor_wrapper = regressor_list.get(name)
        r2score = 0
        for train_index, test_index in skf.split(X, response):
            X_train, y_train = X[train_index], y[train_index]
            X_test, y_test = X[test_index], y[test_index]

            regressor = regressor_wrapper(X_train, y_train)
            
            y_predict = regressor.predict(X_test)
#             score = ((y_test - y_predict) ** 2).sum()
#             print r2_score(y_test, y_predict),
            r2score += 1.0 * r2_score(y_test, y_predict) / split_count
    #         print r2score,
       
        print name, r2score, 
             
        if r2score > best_r2_score:
#             print "new r2score"
            best_r2_score = r2score
            best_regressor = regressor
            best_regressor_name = name
    
#     print
    
    print best_regressor_name,
    
    return best_r2_score, best_regressor
#         y_predict = regressor.predict(X_new)
#         print "Test Score: ", ((y_new - y_predict) ** 2).sum()
#         print "Test r2_score:", r2_score(y_new, y_predict)
    
    
    
# def regressor_for_station(station_id):
#     fname = generate_station_data_fname(station_id, sample_type='train', data_type='numeric', allow_nan_values=False)
# #     print fname
#     station_features = numeric_features[station_id]    
#     features = ['Id'] + station_features
#     station_df = pd.read_csv(fname, usecols=features, index_col=['Id'], dtype='float32')
# 
#     print "Station: %s, # of features: %d" % (station_id, len(station_df.columns))
#     
#     for feature in station_features:
# #         feature = 'L0_S1_F28'
#         print feature
#         dataset_df = station_df[station_df[feature].notnull()]
#         dataset_y_df = dataset_df[feature]
#         del dataset_df[feature]
#         dataset_X_df = dataset_df
#         
#         best_score, best_regressor = regressor_for_station_feature(station_id=station_id, feature=feature, 
#                                                                    X=dataset_X_df.values, y=dataset_y_df.values, 
#                                                                    response=response)
#         fname = generate_regressor_fname(station_id, feature)
#         print "Best score: %f, Saving %s in %s" % (best_score, str(type(best_regressor)), fname)
#         save_regressor(station_id, feature, best_regressor)
# #         break

def regressor_check_for_station_feature(station_id, feature, X, y, response):
    split_idx = len(y) / 3
    
    X_train = X[:-split_idx,:]
    y_train = y[:-split_idx]
    
    X_test = X[-split_idx:,:]
    y_test = y[-split_idx:]

#     regressor = load_regressor(station_id, feature)
    fname = generate_regressor_fname(station_id, feature)
    try:
        regressor = pickle.load(open(fname, "rb" ))
    except Exception as e:
        return -999, None
    
    if isinstance(regressor, XGBRegressor):
        regressor = XGBRegressor()
        r = pickle.load(open(fname, "rb" ))
        regressor._Booster = r._Booster
        regressor.set_params(**r.get_xgb_params())

#     try:
#         y_predict = regressor.predict(X_test)
#     except Exception as e:
#         y_predict = regressor.predict(X_test[:,:-1])
    y_predict = regressor.predict(X_test)

    score = ((y_test - y_predict) ** 2).sum()
    r2score = r2_score(y_test, y_predict)

    return r2score, regressor
    
    
    
def regressor_check_for_station(station_id):
    print "Station:", station_id
    train_missing_features = get_missing_features(station_id, 'train')
    test_missing_features = get_missing_features(station_id, 'test')
    
#     print train_missing_features
#     print test_missing_features
    
    station_df_missing_features = sorted(list(set(train_missing_features).union(set(test_missing_features))))

    if len(station_df_missing_features) == 0:
        print "Nothing to do. Skip."
        return

    station_df, response_df = load_station_df(station_id, use_time=True, use_category=True, use_response=True, use_product_id=False)
    
    for feature in station_df_missing_features:
        print station_id, ";", feature, ";",
        dataset_df = station_df[station_df[feature].notnull()]
        
        dataset_y_df = dataset_df[feature]
        del dataset_df[feature]
        dataset_X_df = dataset_df
        response = response_df.loc[dataset_X_df.index]
        
        best_score, best_regressor = regressor_check_for_station_feature(station_id=station_id, feature=feature, 
                                                                         X=dataset_X_df.values, y=dataset_y_df.values, 
                                                                         response=response.values)
        print "%10f; %s" % (best_score, str(type(best_regressor)))
#         break


def get_missing_features(station_id, sample_type='train'):
    fname = generate_station_data_fname(station_id=station_id, sample_type=sample_type, data_type='numeric', allow_nan_values=False)   
    features = ['Id'] + numeric_features[station_id]
    
    station_df = pd.read_csv(fname, usecols=features, index_col=['Id'], dtype=np.float32)

    station_df_null_features = []
    for feature in station_df.columns:
        if station_df[feature].isnull().any():
            station_df_null_features.append(feature)
    
    del station_df
    
    return station_df_null_features

def regressor_for_missing_features(station_id, args):
    print "Station:", station_id
    print args
    
    train_missing_features = get_missing_features(station_id, 'train')
    test_missing_features = get_missing_features(station_id, 'test')
    
#     print train_missing_features
#     print test_missing_features
    
    station_df_missing_features = sorted(list(set(train_missing_features).union(set(test_missing_features))))

    if len(station_df_missing_features) == 0:
        print "Nothing to do. Skip."
        return

#     station_df, response_df = load_station_df(station_id, use_time=True, use_null_tracker=True, use_category=True, use_response=True, use_product_id=False)
    station_df, response_df = load_station_df(station_id, **args)

#     print "Station: %s, # of features: %d" % (station_id, len(station_df.columns))
    
    for feature in station_df_missing_features:
        t0 = datetime.now()

        print station_id, ";", feature, ";",
        regressor_fname = generate_regressor_fname(station_id, feature)
        if os.path.exists(regressor_fname):
            print "Regressor exists. Continue."
            continue
        
#         feature = 'L0_S0_F20'
        dataset_df = station_df[station_df[feature].notnull()]
        
        dataset_y_df = dataset_df[feature]
        del dataset_df[feature]
        dataset_X_df = dataset_df
        response = response_df.loc[dataset_X_df.index]
        response = response['Response']
        
        best_score, best_regressor = regressor_for_station_feature(station_id=station_id, feature=feature, 
                                                                   X=dataset_X_df.values, y=dataset_y_df.values, 
                                                                   response=response.values)
        print "%10f; %s" % (best_score, str(type(best_regressor)))
#         break
#         save_regressor(station_id, feature, best_regressor)
        t1 = datetime.now()
        print feature, "in", t1 - t0
        
    del station_df
        
def regressor_for_missing_features_to_work_on(station_id, missing_features_to_work_on):
    print "Station:", station_id
    
    station_df_missing_features = missing_features_to_work_on[station_id]

    if len(station_df_missing_features) == 0:
        print "Nothing to do. Skip."
        return

    station_df, response_df = load_station_df(station_id, use_time=True, use_category=True, use_response=True, use_product_id=False)

#     print "Station: %s, # of features: %d" % (station_id, len(station_df.columns))
    
    for feature in station_df_missing_features:
        print station_id, ";", feature, ";",
        regressor_fname = generate_regressor_fname(station_id, feature)
        if os.path.exists(regressor_fname):
            print "Regressor exists. Continue."
            continue
        
#         feature = 'L0_S0_F20'
        dataset_df = station_df[station_df[feature].notnull()]
        
        dataset_y_df = dataset_df[feature]
        del dataset_df[feature]
        dataset_X_df = dataset_df
        response = response_df.loc[dataset_X_df.index]
        
        best_score, best_regressor = regressor_for_station_feature(station_id=station_id, feature=feature, 
                                                                   X=dataset_X_df.values, y=dataset_y_df.values, 
                                                                   response=response.values)
        print "%10f; %s" % (best_score, str(type(best_regressor)))
        save_regressor(station_id, feature, best_regressor)
        
    del station_df
        
          
# for station_id in sorted(numeric_features):
#     station_id = 'L0S01'
#     regressor_check_for_station(station_id)
# #     break
# exit(-1)

additional_features_none = {
    'use_time': False, 
    'use_null_tracker' : False, 
    'use_category': False, 
    'use_response': False, 
    'use_product_id': False,
}

additional_features_time = {
    'use_time': True, 
    'use_null_tracker' : False, 
    'use_category': False, 
    'use_response': False, 
    'use_product_id': False,
}

additional_features_null = {
    'use_time': False, 
    'use_null_tracker' : True, 
    'use_category': False, 
    'use_response': False, 
    'use_product_id': False,
}

additional_features_cat = {
    'use_time': False, 
    'use_null_tracker' : False, 
    'use_category': True, 
    'use_response': False, 
    'use_product_id': False,
}

additional_features_prod = {
    'use_time': False, 
    'use_null_tracker' : False, 
    'use_category': False, 
    'use_response': False, 
    'use_product_id': True,
}


additional_features_time_null = {
    'use_time': True, 
    'use_null_tracker' : True, 
    'use_category': False, 
    'use_response': False, 
    'use_product_id': False,
}

additional_features_time_null_cat = {
    'use_time': True, 
    'use_null_tracker' : True, 
    'use_category': True, 
    'use_response': False, 
    'use_product_id': False,
}

additional_features_time_null_cat_prod = {
    'use_time': True, 
    'use_null_tracker' : True, 
    'use_category': True, 
    'use_response': False, 
    'use_product_id': True,
}

additional_features_reponse = {
    'use_time': False, 
    'use_null_tracker' : False, 
    'use_category': False, 
    'use_response': True, 
    'use_product_id': False,
}

additional_features_time_null_cat_prod_resp = {
    'use_time': True, 
    'use_null_tracker' : True, 
    'use_category': True, 
    'use_response': True, 
    'use_product_id': True,
}

additional_features_time_cat = {
    'use_time': True, 
    'use_null_tracker' : False, 
    'use_category': True, 
    'use_response': False, 
    'use_product_id': False,
}

from datetime import datetime

if __name__ == '__main__':
    for station_id in sorted(numeric_features)[0:24]:
#         station_id = 'L1S24'
#         regressor_check_for_station(station_id)
        # regressor_for_missing_features(station_id, additional_features_none)
        # regressor_for_missing_features(station_id, additional_features_time)
        # regressor_for_missing_features(station_id, additional_features_time_null)
        regressor_for_missing_features(station_id, additional_features_time_cat)
        
#         break
# or
#     for station_id in sorted(numeric_features_to_work_on):
#         regressor_for_missing_features_to_work_on(station_id, numeric_features_to_work_on)

     
