#
# 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 datetime import datetime

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, 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=6)
    xgb_regressor.fit(X=X_train, y=y_train) #, eval_set=[(X_test, y_test)], verbose=True)
    return xgb_regressor

def cb_regressor_wrapper(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 rf_regressor_wrapper(X_train, y_train):
    rf_regressor = RandomForestRegressor(n_estimators=1500, max_features='sqrt', max_depth=6, n_jobs=-1)    
    rf_regressor.fit(X=X_train, y=y_train) #, verbose=True) 
    return rf_regressor

regressor_list = {
#     'XGB_Regressor': xgb_regressor_wrapper,
    'CB_Regressor': cb_regressor_wrapper,
#     'RF_Regressor': rf_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, sample_type='train', 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=sample_type, 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=sample_type, 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=sample_type, 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=sample_type, 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)


    response_df = None
    if (sample_type=='train'):
        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=42)
    
#     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,
            
        if r2score > best_r2_score:
#             print "new r2score"
            best_r2_score = r2score
            best_regressor = regressor
            best_regressor_name = name
    
#     print
    
    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 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 save_df(station_id, station_df, sample_type, iter_no):
    fname = generate_station_data_fname(station_id, sample_type, 'numeric', iter_no=iter_no)
    station_df.to_csv(fname)

    
def regressor_for_missing_features(station_id, iter_count=10):
    print "Station:", station_id
    train_missing_features = get_missing_features(station_id, 'train')
    test_missing_features = get_missing_features(station_id, 'test')
    
    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, sample_type='train', use_time=True, use_null_tracker=False, use_category=True, use_response=False, use_product_id=False)
    full_features = station_df.columns
# predict: Load test dataset
    station_test_df, _  = load_station_df(station_id, sample_type='test', use_time=True, use_null_tracker=False, use_category=True, use_response=False, use_product_id=False)

#     print "Station: %s, # of features: %d" % (station_id, len(station_df.columns))

    
    station_null_indices = {}
    test_null_indices = {}
    for feature in station_df_missing_features:
        feature_null_indices = station_df[station_df[feature].isnull()].index
        station_null_indices[feature] = feature_null_indices
# predict: Initialize missing indices
        feature_null_indices = station_test_df[station_test_df[feature].isnull()].index
        test_null_indices[feature] = feature_null_indices

    print "Train:", station_null_indices
    print "Test:", test_null_indices
    
    temp_df = station_df.copy()
# predict: don't know if we need this
    temp_test_df = station_test_df.copy()
    
    # Fill in initial values for missing values
    stds = station_df.std()
    means = station_df.mean()
    
    for feature in station_df_missing_features:
        std, mean = stds[feature], means[feature]
        print feature, std, mean 
        station_df.loc[station_null_indices[feature], feature] = np.random.normal(mean, std, len(station_null_indices[feature])) 
# predict: Initialize missing values
#         stationtest_df.loc[test_null_indices[feature], feature] = np.random.normal(mean, std, len(test_null_indices[feature])) 
    
    feature_scores = {}
    
    print "Initial Scores:"
    for feature in station_df_missing_features: 
        print station_id, ";", feature, ";",
        
        train_features = list((set(full_features) - set([feature])))
        train_indices = station_df[station_df[feature].notnull()].index
        test_indices = station_df[station_df[feature].isnull()].index

        X_train = station_df.loc[train_indices, train_features]
        y_train = station_df.loc[train_indices, feature]
        
        X_test = station_df.loc[test_indices, train_features] # feature == null
        
        response = response_df.loc[X_train.index]
        response = response['Response']
        
        best_score, best_regressor = regressor_for_station_feature(station_id=station_id, feature=feature, 
                                                                   X=X_train.values, y=y_train.values, 
                                                                   response=response.values)

        print "%10f; %s" % (best_score, str(type(best_regressor)))
        save_regressor(station_id, feature, best_regressor, 0)
#         break
        
        feature_scores[feature] = best_score

    print "Iteration Scores:"
    for iter in range(iter_count):
        print "Iteration:", iter+1
        t0 = datetime.now()
        for feature in sorted(feature_scores, key=feature_scores.get, reverse=True): #station_df_missing_features: 

            print station_id, ";", feature, ";",
            regressor_fname = generate_regressor_fname(station_id, feature, iter+1)
            if os.path.exists(regressor_fname):
                print "Regressor exists. Continue."
                continue
            
            train_features = list((set(full_features) - set([feature])))
            train_indices = station_df[temp_df[feature].notnull()].index
            test_indices = station_df[temp_df[feature].isnull()].index

            X_train = station_df.loc[train_indices, train_features]
            y_train = station_df.loc[train_indices, feature]
            
            X_test = station_df.loc[test_indices, train_features] # feature == null
            
            response = response_df.loc[X_train.index]
            response = response['Response']
            
            best_score, best_regressor = regressor_for_station_feature(station_id=station_id, feature=feature, 
                                                                       X=X_train.values, y=y_train.values, 
                                                                       response=response.values)

            print "%10f; %s" % (best_score, str(type(best_regressor)))
    #         break
            save_regressor(station_id, feature, best_regressor, iter+1)
            
            feature_scores[feature] = best_score

            if (len(X_test) > 0):
                y_predict = best_regressor.predict(X_test.values)
                station_df.loc[test_indices, feature] = y_predict
            else:
                print ("Station %s has no missing values in train dataset!", station_id)
                
            
# predict: use regressor to update station_test_df            
            test_indices = station_test_df[temp_test_df[feature].isnull()].index
            if (len(test_indices) > 0):
        		X_test = station_test_df.loc[test_indices, train_features] # feature == null
        		y_predict = best_regressor.predict(X_test.values)
        		station_test_df.loc[test_indices, feature] = y_predict
            else:
                print ("Station %s has no missing values in test dataset!", station_id)
            
 
        save_df(station_id, station_df, 'train', iter+1)    