"""
Target Problem:
---------------
* Predict the evolution of brain connectivity over time.

Proposed Solution (Machine Learning Pipeline):
----------------------------------------------
* SelectKBest Algorithm -> Support Vector Regression

Input to Proposed Solution:
---------------------------
* Directories of training data and labels(measurements at two different timepoints), and testing data in csv file format.
* These data should be stored in n x m pattern in csv file format.

  Typical Example:
  ----------------
  n x m samples in train_t0 csv file (n number of samples, m number of features)
  n x m samples in train_t1 csv file (n number of samples, m number of features)
  k x m samples in test_t0 csv file (k number of samples, m number of features)

* These data set files are ready by load_data() function.

Output of Proposed Solution:
----------------------------
* Predictions generated by learning model for testing set
* They are stored in "submission.csv" file.

Code Owner:
-----------
* Copyright © Team 7. All rights reserved.
* Copyright © Istanbul Technical University, Learning From Data Fall 2021. All rights reserved.
"""

import numpy as np
import pandas as pd
from sklearn.feature_selection import SelectKBest,mutual_info_regression
from sklearn.model_selection import KFold
from sklearn.metrics import mean_squared_error as mse
from sklearn.metrics import mean_absolute_error as mae
from scipy.stats import pearsonr as pr
from sklearn.svm import SVR
from sklearn.tree import DecisionTreeRegressor
import math
from scipy.stats import pearsonr
import statistics
import random as r
r.seed(1)
np.random.seed(1)

def load_data(directories):

    """
    The method reads train and test data from their dataset files.
    Then, it splits train data into features and labels.
    Parameters
    ----------
    train_file: directory of the file in which train data set is located
    test_file: directory of the file in which test data set is located
    """

    train_t0 = pd.read_csv(directories[0])
    train_t1 = pd.read_csv(directories[1])
    test_t0 = pd.read_csv(directories[2])
    return train_t0, train_t1, test_t0

def preprocessing(x_tra, y_tra):

    """
    This method performs SelectKBest filtering algorithm to select top k
    important features.

    Parameters
    ----------
    x_tra: features of training data
    y_tra: labels of training data
    """

    FSmodel = SelectKBest(mutual_info_regression, k=130)
    FSmodel.fit(x_tra, y_tra)
    return FSmodel, FSmodel.transform(x_tra)

class MultioutputLearner:
    """
    A custom learner class that supports multiple target outputs.
    """
    def __init__(self,X,Y):
        """
        Constructor method for the custom class.

        Parameters
        ----------
        X: train dataset
        y: labels of train dataset
        """
        self.X = np.array(X).copy()
        self.Y = np.array(Y).copy()
        #array for saving learning models, includes 595 learned model, one for every dimension of y(label)
        self.listofLR = []
        #array for saving dimensionality reduction techniques, includes 595 learned model, one for every dimension of y(label)
        self.listofFS = []
    
    #function for training the dataset
    #It takes the multioutput class and desired number of features as input
    def train(self,n_features):
        """
        Constructor method for the custom class.

        Parameters
        ----------
        X: train dataset
        y: labels of train dataset
        """
        for i in range(self.Y.shape[1]): #iterating through 595 brain connectivity t1 values 
            y = np.array(self.Y[:,i]).copy()
            epsilon = 1e-100
            for i in range(y.shape[0]):
                if(y[i] > 1):
                    y[i] = 1
                if(y[i] < 0):
                    y[i] = 0
                if((1-y[i]) < epsilon): #to avoid overflow in model
                    y[i] = 1e+10
                elif(y[i] < epsilon): #to avoid overflow in model
                    y[i] = -1e+10
                else:
                    temp = y[i]
                    y[i] = np.log(temp/(1-temp)) #inverse sigmoid function

            temp_y = y.copy() #preprocessing the t1 value
            FSModel, new_X = preprocessing(self.X, temp_y)
            temp_model = SVR().fit(new_X,temp_y) #creating the model
            self.listofLR.append(temp_model) #adding learned model to list
            self.listofFS.append(FSModel) #adding dimensionality reduction technique to list

    #function for testing the dataset
    #It takes the multioutput class and testset of the dataset as input and returns the predicted t1 values
    def test(self,testset):
        testoutput = [] #array for saving predicted t1 values
        testset = np.array(testset).copy() #changing the testset's type to numpy array
        for i in range(len(self.listofLR)): #traversing through learned models
            new_test = self.listofFS[i].transform(testset) #applying dimensionality reduction to testset
            prediction = self.listofLR[i].predict(new_test).copy() #predicting the t1 values
            testoutput.append(prediction) #adding predictions to array
        # returning the transpose of predicted values, since appending prediction adds predictions one after other
        # but we should add predictions side by sides, hence transpose will provide us the desired format
        y = np.array(testoutput)
        for i in range(y.shape[0]):
            for j in range(y.shape[1]):
                y[i][j] = 1/(1 + np.exp(-y[i][j])) #sigmoid function
        return np.transpose(y) 


def train(x_tra, y_tra):

    """
    The method creates a learning model and trains it by using training data.
    Parameters
    ----------
    x_tra: features of training data
    y_tra: labels of training data
    """

    #adding train datasets of t0 and t1 values side by side
    All_train = np.append( np.delete(np.array(x_tra).copy(), [0], axis=1), np.delete(np.array(y_tra).copy(), [0], axis=1) ,axis=1)
    #creating k folds from train datasets
    K_fold = KFold(n_splits=5, shuffle = True, random_state = 1)
    #k-fold cross validation
    models = []
    for train_index, test_index in K_fold.split(All_train):
        train_fold = np.array((All_train[train_index])) #getting train fold
        fold_t0 = train_fold[:,0:int(train_fold.shape[1]/2)] #getting features of train fold
        fold_t1 = train_fold[:,int(train_fold.shape[1]/2):train_fold.shape[1]] #getting labels of train fold
        model = MultioutputLearner(fold_t0,fold_t1) #creating the multioutputlearner class from train fold
        model.train(130) #training the fold
        models.append(model)
    return models

def predict(x_tst, model):

    """
    The method predicts labels for testing data samples by using trained learning model.
    Parameters
    ----------
    x_tst: features of testing data
    model: trained learning model
    """

    # getting testset without ids
    test = np.delete(np.array(x_tst).copy(), [0], axis=1)
    #creating an empty Test_result array in order to save predictions
    Test_result = np.zeros((test.shape[0],x_tst.shape[1]))
    for model_ in model:
        prediction = model_.test(test)
        Test_result += prediction
    return Test_result/5

def write_output(filename,predictions):
    
    """
    Writes the outputted predictions to a file.
    Parameters
    ----------
    filename: file path for saving file
    predictions: numpy.ndarray object containing predictions
    """

    output = pd.DataFrame({"Predicted": predictions.flatten()})
    output.index.name = "ID"
    output.to_csv(filename)

def cv(x_tra, y_tra):

    """
    The method performs 5-fold cross-validation on training data and returns the error rates.
    Parameters
    ----------
    x_tra: features of training data
    y_tra: labels of training data
    """

    # adding train datasets of t0 and t1 values side by side
    All_train = np.append( np.delete(np.array(x_tra).copy(), [0], axis=1), np.delete(np.array(y_tra).copy(), [0], axis=1) ,axis=1)
    #creating k folds from train datasets
    K_fold = KFold(n_splits=5, shuffle = True, random_state = 1)
    #for saving all predictions
    All_prediction = np.delete(np.array(x_tra).copy(), [0], axis=1)
    #k-fold cross validation
    msqr = []
    pearson = []
    mabs = []
    i = 1
    for train_index,test_index in K_fold.split(All_train):
        train_fold = np.array((All_train[train_index]))
        fold_t0 = train_fold[:,0:int(train_fold.shape[1]/2)] #features of train fold
        fold_t1 = train_fold[:,int(train_fold.shape[1]/2):train_fold.shape[1]] #labels of train fold

        test_fold = np.array((All_train[test_index]))
        test_fold_t0 = test_fold[:,0:int(test_fold.shape[1]/2)]#getting features of test fold
        test_fold_t1 = test_fold[:,int(test_fold.shape[1]/2):test_fold.shape[1]] #getting labels of test fold
        
        model = MultioutputLearner(fold_t0,fold_t1) 
        model.train(130) #training the train fold
        prediction = model.test(test_fold_t0) #testing the test fold

        msqr.append(mse(prediction.flatten(), test_fold_t1.flatten()))
        mabs.append(mae(prediction.flatten(), test_fold_t1.flatten()))
        pearson.append(pr(prediction.flatten(), test_fold_t1.flatten()))
        All_prediction[test_index] = prediction #saving prediction
        i+=1
    print("MAE:")
    for i,stat in enumerate(mabs):
        print(i+1, stat)
    print("PCC:")
    for i,stat in enumerate(pearson):
        print(i+1, stat)
    print(statistics.mean(mabs), statistics.stdev(mabs))
    print(statistics.mean(i[0] for i in pearson), statistics.stdev(i[0] for i in pearson))
    print(statistics.mean(msqr), statistics.stdev(msqr))
    #the mean squared error between predicted train_t1 values and real train_t1 values
    total_error = mse(np.delete(np.array(y_tra).copy(), [0], axis=1),All_prediction)
    #the pearson correlation between predicted train_t1 values and real train_t1 values
    pears_err,p = pearsonr(All_prediction.flatten(),np.delete(np.array(y_tra).copy(), [0], axis=1).flatten())
    print("Mean Squared Error: ", total_error)
    print("Pearson Correlation: ",pears_err," ~ p-value: ",p)
    return total_error #mean squared error

filenames = ["train_t0.csv", "train_t1.csv", "test_t0.csv"]
x_tra, y_tra, x_tst = load_data(filenames)
cv(x_tra, y_tra)
model = train(x_tra, y_tra)
predictions = predict(x_tst, model)
write_output("submission.csv", predictions)