import pandas as pd
import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt
import pickle
import os
import math
import glob

from sklearn.metrics import mean_squared_error

from include.dataset_fnames import generate_station_data_fname, generate_data_fname, generate_response_data_fname, train_categorical_onehot_filename, train_station_date_filename

from include.feature_lists import numeric_features, numeric_missing_features_list, numeric_features_to_work_on, categoric_features

import theano
import theano.tensor as T

import keras
from keras import backend as K
from keras.layers import Input, Dense, Dropout, BatchNormalization
from keras.models import Model, Sequential
from keras.callbacks import Callback
#from keras.losses import mean_squared_error
from keras.optimizers import adam, adadelta
from keras.regularizers import L1L2

from BoschNNModel1 import BoschNNModel, BoschNNInput, BoschNNOutput, BoschNNCallback

from sklearn.model_selection import KFold
from keras.callbacks import History, EarlyStopping

import random

from datetime import datetime

cv_splits = 3
top_epoch_count = 30

def batch_count_csv(dirname):
    fnames = glob.glob1(dirname, "train_numeric_*")
    return len(fnames)

def batch_count_npz(dirname):
    fnames = glob.glob1(dirname, "train*.npz")
    return len(fnames)

def bosch_data_generator_from_csv(dirname, indices, shuffle=False, use_categoric_features=False, use_date_features=False):
    fname = generate_data_fname(sample_type='train', data_type='numeric')
    numeric_columns = pd.read_csv(fname, nrows=2).columns
    numeric_fnames = sorted(glob.glob1(dirname, "train_numeric_*"))

    
    if use_categoric_features:
        fname = train_categorical_onehot_filename
        categoric_columns = pd.read_csv(fname, nrows=2).columns
        categoric_fnames = sorted(glob.glob1(dirname, "train_categorical_*"))

    if use_date_features:
        fname = train_station_date_filename
        station_date_columns = pd.read_csv(fname, nrows=2).columns
        station_date_fnames = sorted(glob.glob1(dirname, "train_station_date_*"))
        
    if not shuffle:
        random.shuffle(indices)
    
    for list_index in range(len(indices)):
        t0 = datetime.now()
        numeric_fname = os.path.join(dirname, numeric_fnames[indices[list_index]])
        numeric_df = pd.read_csv(numeric_fname, names=numeric_columns, index_col='Id')
        del numeric_df['Response']
#         print datetime.now() - t0
        
        if use_categoric_features:
            t0 = datetime.now()
            categoric_fname = os.path.join(dirname, categoric_fnames[indices[list_index]])
            categoric_df = pd.read_csv(categoric_fname, names=categoric_columns, index_col='Id')
            numeric_df = numeric_df.join(categoric_df, how='inner')
#             print datetime.now() - t0
            del categoric_df

        if use_date_features:
            t0 = datetime.now()
            station_date_fname = os.path.join(dirname, station_date_fnames[indices[list_index]])
            station_date_df = pd.read_csv(station_date_fname, names=station_date_columns, index_col='Id')
            station_date_df = station_date_df / 1719.0 # Normalization
            numeric_df = numeric_df.join(station_date_df, how='inner')
#             print datetime.now() - t0
            del station_date_df
 
        yield numeric_df.values, numeric_df.values, list_index+1
        del numeric_df

def bosch_data_generator_from_npz(dirname, indices, shuffle=False, use_categoric_features=False, use_date_features=False):
    assert (use_categoric_features == True), "Compressed numeric only values not implemented."
# TODO:      
    assert (use_date_features == False), "Compressed date values not implemented."
    
    fname = generate_data_fname(sample_type='train', data_type='numeric')
    numeric_columns = pd.read_csv(fname, nrows=2).columns

    
    if use_categoric_features:
        fname = train_categorical_onehot_filename
        categoric_columns = pd.read_csv(fname, nrows=2).columns

    compressed_fnames = sorted(glob.glob1(dirname, "train_numeric+categoric_*.npz"))
        
    if not shuffle:
        random.shuffle(indices)
    
    for list_index in range(len(indices)):
        t0 = datetime.now()
        compressed_fname = os.path.join(dirname, compressed_fnames[indices[list_index]])
        dataz = np.load(compressed_fname)
        data = dataz['data']
#         print "Loaded in", datetime.now() - t0
         
        yield data, data, list_index+1
        del dataz
        del data
        
def load_anomaly_data(use_categoric_features, use_date_features):
    fname = generate_data_fname(sample_type='train', data_type='numeric')
    numeric_columns = pd.read_csv(fname, nrows=2).columns
    numeric_fname = "d:/Kaggle_ws/Bosch/input/train_numeric_headless_1.csv"

    numeric_df = pd.read_csv(numeric_fname, names=numeric_columns, index_col='Id')
    del numeric_df['Response']
        
    if use_categoric_features:
        fname = train_categorical_onehot_filename
        categoric_columns = pd.read_csv(fname, nrows=2).columns
        categoric_fname = "d:/Kaggle_ws/Bosch/input/train_categorical_onehot_headless_1.csv"
        categoric_df = pd.read_csv(categoric_fname, names=categoric_columns, index_col='Id')
        numeric_df = numeric_df.join(categoric_df, how='inner')
        del categoric_df

    if use_date_features:
        fname = train_station_date_filename
        station_date_columns = pd.read_csv(fname, nrows=2).columns
        station_date_fname = "d:/Kaggle_ws/Bosch/input/train_station_date_headless_1.csv"
        station_date_df = pd.read_csv(station_date_fname, names=station_date_columns, index_col='Id')
        station_date_df = station_date_df / 1719.0 # Normalization
        numeric_df = numeric_df.join(station_date_df, how='inner')
        del station_date_df
    
    return numeric_df.values

def my_mean_squared_error(y_true, y_pred):
    mask = T.invert(T.isnan(y_true))
    
    y_true = y_true[mask.nonzero()]
    y_pred = y_pred[mask.nonzero()]
 
    return K.mean(K.square(y_pred - y_true), axis=-1)

# [968, 484, 32, 16, 10, 16, 32, 484, 968]
def load_model(layer_list):
    print "Building network with layers", layer_list
    cols = layer_list[0]

    model = Sequential() # was BoschNNModel()
    
# Create Input Layer
    inputLayer = BoschNNInput(units=layer_list[1], input_shape=(cols,), name="first", activation='relu')

# Add layers to model
    model.add(inputLayer)

    for i in range(2, len(layer_list) -1):
        model.add(BatchNormalization())
        model.add(Dense(layer_list[i], name='hidden_'+str(i), activation='relu', bias_regularizer=L1L2(l1=0.01)))

# # Add Dropout Layer
#     model.add(Dropout(name="Dropout", rate=0.5))
# Create Output Layer
#     outputLayer = BoschNNOutput(model, cols, name='last', activation='linear')
    outputLayer = Dense(cols, name='last', activation='linear')
    model.add(outputLayer)

    return model


def train_autoencoder_model(dirname, model_name, hidden_layers, use_categoric_features=False, use_date_features=False, use_compressed=False):


    fname = generate_data_fname(sample_type='train', data_type='numeric')
    numeric_columns = pd.read_csv(fname, nrows=2).columns
    cols = len(numeric_columns) -1 -1 # one for 'Id' one for 'Respose'

    if use_categoric_features:
        fname = train_categorical_onehot_filename
        categoric_columns = pd.read_csv(fname, nrows=2).columns
        cols = cols + len(categoric_columns) -1 # one for 'Id' 

    if use_date_features:
        fname = train_station_date_filename
        station_date_columns = pd.read_csv(fname, nrows=2).columns
        cols = cols + len(station_date_columns) -1 # one for 'Id' 
    
    anomaly_df = load_anomaly_data(use_categoric_features, use_date_features)

    layer_list = [cols] + eval(hidden_layers) + [cols]
# Load model from somewhere
    boschAutoencoderModel = load_model(layer_list)
# Compile model
    boschAutoencoderModel.compile(optimizer='adam', loss=my_mean_squared_error)
 
# this is what we have in the model
    boschAutoencoderModel.summary()
      
# Initialize Callback for weight processing    
#     boschNNCallback = BoschNNCallback()
      
#     boschModel.fit(X, y, epochs=10, batch_size=1, shuffle=False, verbose=True, callbacks=[boschNNCallback])
      
#     boschAutoencoderModel.fit(X_train, X_train,
#                     epochs=5,
#                     batch_size=1,
#                     shuffle=False,
# #                     callbacks=[boschNNCallback],
#                     validation_data=(X_test, X_test))

    batch_count = -1 # No of files to load at every CV
    if not use_compressed:
        batch_count = batch_count_csv(dirname)
    else:
        batch_count = batch_count_npz(dirname)
    print batch_count

    kf = KFold(n_splits=cv_splits)
    
    earlystoping = EarlyStopping(monitor='val_loss', min_delta=1e-6, patience=3, verbose=1, mode='auto')
    
    final_mse = 0.0
    best_cv_mse = 99999
    cv_count = 1
    for train_indices, test_indices in kf.split(range(batch_count)):
        print "----------------------------------------------------------------------"
        print "CV:", cv_count, "/", cv_splits
        

# Rewind model!
# Load model from somewhere
        boschAutoencoderModel = load_model(layer_list)
        
# Compile model
#         optimizer = zdam(lr=0.001, beta_1=0.9, beta_2=0.999, epsilon=1e-08, decay=0.001)
#         boschAutoencoderModel.compile(optimizer='adadelta', loss=my_mean_squared_error)
        optimizer = adadelta(lr=1.0)
        boschAutoencoderModel.compile(optimizer=optimizer, loss=my_mean_squared_error)

# Go 3 times on each batch
        for top_epoch, _ in enumerate(range(top_epoch_count),1):
            print "Top epoch:", top_epoch, "of", top_epoch_count
            if not use_compressed:
                generator = bosch_data_generator_from_csv(dirname, train_indices, shuffle=True, use_categoric_features=use_categoric_features, use_date_features=use_date_features)
            else:
                generator = bosch_data_generator_from_npz(dirname, train_indices, shuffle=True, use_categoric_features=use_categoric_features, use_date_features=use_date_features)
                
            history = History()
            for X, y, file_index in generator:
                print "Set", file_index, "of", len(train_indices)
                print "LR:", boschAutoencoderModel.optimizer.lr.get_value()
                boschAutoencoderModel.fit(X, y, 
                                        batch_size=4092,
                                        epochs=3, 
#                                         verbose=2,
#                                         validation_split=0.33,
                                        shuffle=True,
#                                         callbacks=[history, earlystoping]
                                        callbacks=[history]
                                        )
               
                print cv_count, top_epoch, file_index
                print history.history
                del X
                del y
#                 new_lr = np.float32(0.97 * boschAutoencoderModel.optimizer.lr.get_value())
#                 boschAutoencoderModel.optimizer.lr.set_value(new_lr)

            print "Validating top_epoch", top_epoch,
            if not use_compressed:
                validation_generator = bosch_data_generator_from_csv(dirname, test_indices, use_categoric_features=use_categoric_features, use_date_features=use_date_features)
            else:
                validation_generator = bosch_data_generator_from_npz(dirname, test_indices, use_categoric_features=use_categoric_features, use_date_features=use_date_features)
                
            for X, y, _ in validation_generator:
                predictions = boschAutoencoderModel.predict(X)
                
                y_pred = predictions[np.invert(np.isnan(X)).nonzero()]
                y_true = y[np.invert(np.isnan(X)).nonzero()]
     
                mse_ = mean_squared_error(y_true, y_pred)
                print "mse_:", mse_

        fname_model = os.path.join(model_name + "_CV_" + str(cv_count) + ".h5")
        print "Saving model to", fname_model, "."
        boschAutoencoderModel.save(fname_model)

        print "Validating..."
        cv_mse = 0.0
        total = 0.0
        if not use_compressed:
            validation_generator = bosch_data_generator_from_csv(dirname, test_indices, use_categoric_features=use_categoric_features, use_date_features=use_date_features)
        else:
            validation_generator = bosch_data_generator_from_npz(dirname, test_indices, use_categoric_features=use_categoric_features, use_date_features=use_date_features)
            
        for X, y, _ in validation_generator:
            predictions = boschAutoencoderModel.predict(X)
            
            y_pred = predictions[np.invert(np.isnan(X)).nonzero()]
            y_true = y[np.invert(np.isnan(X)).nonzero()]
 
            counts = y_true.shape[0]
            mse_ = mean_squared_error(y_true, y_pred)
            cv_mse += counts * mse_ 
            total += counts
#             
#             mask = np.invert(np.isnan(X))
#             predictions = boschAutoencoderModel.predict(X)
#             y_pred = predictions[mask.nonzero()]
#             
#             del X
#             del predictions
# 
#             y_true = y[mask.nonzero()]
# 
#             del y
#             del mask
#             
#             counts = y_true.shape[0]
#             mse_ = mean_squared_error(y_true, y_pred)
#             cv_mse += counts * mse_ 
#             total += counts
#             
#             del y_true
#             del y_pred

        
        cv_mse = cv_mse / total    
        print "CV MSE:", cv_mse
        if cv_mse < best_cv_mse:
            best_cv_mse = cv_mse

            fname_model = os.path.join(model_name + "_best_model" + ".h5")
            print "New best CV MSE. Saving model to", fname_model, "."
            boschAutoencoderModel.save(fname_model)

            predictions = boschAutoencoderModel.predict(anomaly_df)
            y_pred = predictions[np.invert(np.isnan(anomaly_df)).nonzero()]
            y_true = anomaly_df[np.invert(np.isnan(anomaly_df)).nonzero()]
            a_mse_ = mean_squared_error(y_true, y_pred)
            print "Anomaly error:", a_mse_
            
        final_mse +=  cv_mse
        cv_count += 1
    
    print "Final MSE:", final_mse / cv_splits 
    return

def load_encoder_model(model_fname, bottleneck):
    print "Loading", model_fname,
    ae_model = keras.models.load_model(
        model_fname, 
        custom_objects={"BoschNNInput":BoschNNInput, "my_mean_squared_error":my_mean_squared_error})
 
    bottleneck_layer = 'hidden_' + str(bottleneck)
    encoder_model = Model(inputs=ae_model.input, outputs=ae_model.get_layer(bottleneck_layer).output)

    print "Done."
    
    return encoder_model

def load_autoencoder_model(model_fname):
    print "Loading", model_fname,
    ae_model = keras.models.load_model(
        model_fname, 
        custom_objects={"BoschNNInput":BoschNNInput, "my_mean_squared_error":my_mean_squared_error})
 
    print "Done."
    
    return ae_model

def evaluate_model(model_fname, use_categoric_features=False):

    boschAutoencoderModel = keras.models.load_model(
        model_fname, 
        custom_objects={"BoschNNInput":BoschNNInput, "my_mean_squared_error":my_mean_squared_error})
 
    boschAutoencoderModel.summary()

    anomaly_df = load_anomaly_data(use_categoric_features)

    predictions = boschAutoencoderModel.predict(anomaly_df)

    mask = np.isnan(anomaly_df)
    predictions = boschAutoencoderModel.predict(anomaly_df)

    y_pred = predictions[np.invert(mask).nonzero()]
    y_true = anomaly_df[np.invert(mask).nonzero()]
    a_mse_ = mean_squared_error(y_true, y_pred)

# Is it?
#     predictions = boschAutoencoderModel.predict(anomaly_df)
#     predictions[mask] = 0
#     anomaly_df[mask] = 0
#     a_mse_ = mean_squared_error(anomaly_df, predictions)
    
    print "Anomaly error:", a_mse_

    return

def evaluate2(model_fname, use_categoric_features=False):

    anomaly_df = load_anomaly_data(use_categoric_features)

    boschAutoencoderModel = keras.models.load_model(
        model_fname, 
        custom_objects={"BoschNNInput":BoschNNInput, "my_mean_squared_error":my_mean_squared_error})
 
    boschAutoencoderModel.summary()


    print "Anomaly error:", a_mse_

    return

def predict_with_ae_model(dirname, model_fname, use_categoric_features=False, use_date_features=False, use_compressed=False):

    ae_model = load_autoencoder_model(model_fname) 
 
    ae_model.summary()

    batch_count = batch_count_csv(dirname)
    generator = bosch_data_generator_from_csv(dirname, range(batch_count), use_categoric_features=use_categoric_features, use_date_features=use_date_features)
  
    error_list = []
    for X, y, file_index in generator:
        print "Set", file_index, "of", batch_count, 
        for row_index in range(X.shape[0]):
            Xr = np.reshape(X[row_index], (1, X.shape[1]))
            yr = np.reshape(y[row_index], (1, X.shape[1]))
            mask = np.isnan(Xr)
            predictions = ae_model.predict(Xr)
            y_pred = predictions[np.invert(mask).nonzero()]
            y_true = yr[np.invert(mask).nonzero()]
            if len(y_true) > 0: # Check for all input is NaN!
                a_mse_ = mean_squared_error(y_true, y_pred)
            else:
                a_mse = 0 # I am not sure
            error_list.append(a_mse_)
            if ((row_index % 1000) == 0): print "*", 
        print
           
    np.save('errors_normal.npy', np.asmatrix(error_list, dtype='float32'))
    pd.Series(data=error_list).to_csv('errors_normal.csv')    
    
    X = load_anomaly_data(use_categoric_features, use_date_features)
    y = X
    
    error_list = []
    for row_index in range(X.shape[0]):
        Xr = np.reshape(X[row_index], (1, X.shape[1]))
        yr = np.reshape(y[row_index], (1, X.shape[1]))
        mask = np.isnan(Xr)
        predictions = ae_model.predict(Xr)
        y_pred = predictions[np.invert(mask).nonzero()]
        y_true = yr[np.invert(mask).nonzero()]
        if len(y_true) > 0: # Check for all input is NaN!
            a_mse_ = mean_squared_error(y_true, y_pred)
        else:
            a_mse = 0 # I am not sure
        error_list.append(a_mse_)
        if ((row_index % 1000) == 0): print "*", 
    print
        
    np.save('errors_anomaly.npy', np.asmatrix(error_list, dtype='float32'))
    pd.Series(data=error_list).to_csv('errors_anomaly.csv')    
   
    return

def predict_with_encoder_model(dirname, model_fname, bottleneck, use_categoric_features=False, use_date_features=False, use_compressed=False):

    encoder_model = load_encoder_model(model_fname, bottleneck=bottleneck) 
    encoder_model.summary()

    output_feature_count = encoder_model.layers[-1].output_shape[1]
    
  
    col_names = ['col_'+str(i) for i in range(output_feature_count)]
    
    normal_df = pd.DataFrame(columns=col_names, dtype='float32')
 
    batch_count = batch_count_csv(dirname)
    if (not use_compressed):
        generator = bosch_data_generator_from_csv(dirname, range(batch_count), use_categoric_features=use_categoric_features, use_date_features=use_date_features)
    else:
        generator = bosch_data_generator_from_npz(dirname, range(batch_count), use_categoric_features=use_categoric_features, use_date_features=use_date_features)
 
    for X, _, file_index in generator:
        print "Set", file_index, "of", batch_count 
        predictions = encoder_model.predict(X)
        normal_df = normal_df.append(pd.DataFrame(data=predictions, columns=col_names, dtype='float32'))
    
    normal_df['Response'] = 0    
 
    X = load_anomaly_data(use_categoric_features, use_date_features)

    predictions = encoder_model.predict(X)
    anomaly_df = pd.DataFrame(data=predictions, columns=col_names, dtype='float32')
    anomaly_df['Response'] = 1
    
    data_df = normal_df.append(anomaly_df)

    data_df.to_csv("encoder_output.csv", index=False)
        
    return

if __name__ == '__main__':
# Experiment 1 NO
#     predict_with_encoder_model("bs30000", "experiment_1_no/BAE_model_cols_4-8_numeric_only_best_model.h5", 6, use_categoric_features=False, use_compressed=False)

# Experiment 2 NC 
#     predict_with_encoder_model("bs30000", "experiment_2_nc/BAE_model_cols_4-8_best_model.h5", 6, use_categoric_features=True, use_compressed=True)

# Experiment 3 NC 
#     predict_with_ae_model("bs30000", "experiment_5_nd_2-4/BAE_model_cols_2-4_numeric+date_best_model.h5", use_categoric_features=False, use_date_features=True, use_compressed=False)
#     predict_with_encoder_model("bs30000", "experiment_5_nd_2-4/BAE_model_cols_2-4_numeric+date_best_model.h5", 6, use_categoric_features=False, use_date_features=True, use_compressed=False)
#     exit (-1)

# Experiment 6 ND 
#     predict_with_ae_model("bs30000", "experiment_8_nc_4-8-16/BAE_model_cols_4-8-16_numeric+date_best_model.h5", use_categoric_features=True, use_date_features=True, use_compressed=False)
#     predict_with_encoder_model("bs30000", "experiment_8_nc_4-8-16/BAE_model_cols_4-8-16_numeric+date_best_model.h5", 4, use_categoric_features=True, use_date_features=True, use_compressed=False)
#     exit (-1)

    train_autoencoder_model('bs30000', 
                             'BAE_model_cols_4-8-16_numeric+date', 
#                             '[cols/4, cols/8, 64, 32, 16, 6, 16, 32, 64, cols/8, cols/4]', # 15 M
#                            '[cols/2, cols/4, 64, 32, 16, 6, 16, 32, 64, cols/4, cols/2]', # 30M 
                            '[cols/4, cols/8, cols/16, 32, cols/16, cols/8, cols/4]',
                             use_categoric_features=True,
                             use_date_features=True,
                             use_compressed=False)
#   fit2:  train_autoencoder_model('bs50000', '[cols/2, cols/4, cols/8, 64, 32, 16, 6, 16, 32, 64, cols/8, cols/4, cols/2]')
