from sklearn.linear_model import LinearRegression
from sklearn.ensemble import RandomForestRegressor, GradientBoostingRegressor
from sklearn.cross_decomposition import PLSRegression
from sklearn.preprocessing import PolynomialFeatures
from sklearn.svm import SVR

from sklearn.linear_model import LogisticRegression
from sklearn.tree import DecisionTreeClassifier
from sklearn.ensemble import RandomForestClassifier
from sklearn.neighbors import KNeighborsClassifier
from sklearn.svm import SVC

from xgboost import XGBRegressor

from sklearn.feature_selection import RFE
from sklearn.feature_selection import SequentialFeatureSelector
from sklearn.feature_selection import SelectKBest
from sklearn.feature_selection import SelectFromModel

from sklearn.decomposition import PCA
from sklearn.decomposition import SparsePCA
from sklearn.decomposition import KernelPCA
from sklearn.discriminant_analysis import LinearDiscriminantAnalysis
from sklearn.discriminant_analysis import QuadraticDiscriminantAnalysis


from sklearn.model_selection import train_test_split

import tensorflow as tf
from tensorflow.keras.layers import Dense, Input
from tensorflow.keras import Sequential


from automl.process import Format, Preprocess
from base.utils.servant import BaseFormat

import pmdarima as pm
from pmdarima import model_selection
from prophet import Prophet

import gc

class Create:
    
    def create_regression_model(self, algorithm, parameters = {}):
        """
        Create a regression model

        Parameters
        ----------


        algorithm : {"LR", "PLS", "RFR", "SVR", "PR"}
            algorithm abbreviation

            LR: LinearRegression
            PLS: PartialLeastSquaresRegressor
            RFR: RandomForestRegressor
            SVR: SupportVectorRegressor
            PR: PolinomialRegressor

            Note: Using "PR" would return a tuple including the regressor and the polynomial features object. Later, X_test should be transformed using 
            polinomial features object and predictions must be done on the X_test_polynomial object.

            X_test_polynomial = polynomial.transform(X_test)
            y_pred = regressor.predict(X_test_polynomial)

        parameters: dictionary
            Various parameters that could be defined in different choice of algorithms
        """
        """
        if algorithm == 'PR':

            regressor = LinearRegression()

            polynomial = PolynomialFeatures(degree = parameters['degree'])
            X_polynomial = polynomial.fit_transform(data[0])

            regressor.fit(X_polynomial, data[1])

            return (regressor, polynomial)
         """

        if algorithm == 'LR' : regressor = LinearRegression(**parameters)

        if algorithm == 'PLS': regressor = PLSRegression(**parameters)

        if algorithm == 'RFR' : regressor = RandomForestRegressor(**parameters)

        if algorithm == 'SVR' : regressor = SVR(**parameters)
        
        if algorithm == 'GBR' : regressor = GradientBoostingRegressor(**parameters)
        
        if algorithm == 'XGB' : regressor = XGBRegressor(**parameters)


        return regressor
    
    def create_classification_model(self, algorithm, parameters = {}):
        """
        Create a classification model

        Parameters
        ----------
    
        algorithm : {"LR", "DTC", "RFC", "SVC", "KNN", "LDA", "QDA"}
            algorithm abbreviation

            LR: LogisticRegression
            DTC: DecisionTreeClassifier
            RFC: RandomForestClassifier
            SVC: SupportVectorClassifier
            KNN: KNeighborClassifier
            LDA: LinearDisriminantAnalysis
            QDA: QuadriticDiscriminantAnalysis
            

        parameters: dictionary
            Various parameters that could be defined in different choice of algorithms
        """

        if algorithm == 'LR' : classifier = LogisticRegression(**parameters)

        if algorithm == 'DTC': classifier = DecisionTreeClassifier(**parameters)

        if algorithm == 'RFC' : classifier = RandomForestClassifier(**parameters)

        if algorithm == 'SVC' : classifier = SVC(**parameters)
        
        if algorithm == 'KNN' : classifier = KNeighborsClassifier(**parameters)
        
        if algorithm == 'LDA':model = LinearDiscriminantAnalysis(**parameters)

        if algorithm == 'QDA':model = QuadraticDiscriminantAnalysis(**parameters)


        return classifier
    
    def create_feature_selection_model(self, estimator_model, selector_algorithm, selector_parameters = {}):
        """
       Create a feature selection model

       Parameters
       ----------
        parameters: dictionary
                Various parameters that could be defined in different choice of algorithms

        algorithm:{'SFM','RFE','SKB','SFS'}
             algorithm abbreviation

        SFM: SelectFromModel
        RFE: RecursiveFeatureElimination
        SKB: SelectKBest
        SFS: SequentialFeatureSelection


        """
        
        if selector_algorithm == 'SFM' : feature_selector = SelectFromModel(estimator = estimator_model, **selector_parameters)

        if selector_algorithm == 'RFE': feature_selector = RFE(estimator = estimator_model, **selector_parameters)

        if selector_algorithm == 'SKB': feature_selector = SelectKBest(**selector_parameters)

        if selector_algorithm == 'SFS': feature_selector = SequentialFeatureSelector(estimator = estimator_model, **selector_parameters)
        
        return feature_selector
    
    def create_dimensionality_reduction_model(self,algorithm,parameters={}):
        """
        Create a dimensionality reduction model

        Parameters
        ----------
        parameters: dictionary
                Various parameters that could be defined in different choice of algorithms

        algorithm:{'PCA','SPCA','LDA'}
             algorithm abbreviation

        PCA: PrincipalComponentAnalysis
        SPCA: SparsePCA
        KPCA: KernelDiscriminantAnlaysis


        """

        if algorithm == 'PCA': model = PCA(**parameters)

        if algorithm == 'SPCA': model = SparsePCA(**parameters)

        if algorithm == 'KPCA':model = KernelPCA(**parameters)
        
        if algorithm == 'AE': model = self.create_autoencoder_model(parameters['encoder_layers'], parameters['decoder_layers'], parameters['compile_parameters'])

        return model
    
    def create_deep_learning_model(self, layers, compile_parameters={}):
        """
        Create and compile a deep learning model
        
        Parameters
        ----------
        layers: list
            List of neural network layers
            
        compile_parameters: dictionary
            Parameters for compile method in dictionary format
        """
                
        network = Sequential(layers)
        network.compile(**compile_parameters)

        return network
    
    def create_autoencoder_model(self, encoder_layers, decoder_layers, compile_parameters={}):
        """
        Create and compile a autoencoder model. Returns Autoencoder and Encoder in tuple format
        
        Parameters
        ----------
        encoder_layers: list
            List of encoder layers
            
        decoder_layers: list
            List of decoder layers
            
        compile_parameters: dictionary
            Parameters for compile method in dictionary format
        """
        encoder = Sequential(encoder_layers)
        decoder = Sequential(decoder_layers)
        autoencoder = Sequential([encoder,decoder])
        autoencoder.compile(**compile_parameters)
        
        return (autoencoder, encoder)
    
    def create_time_series_model(self, algorithm, data, parameters={}):
        
        if algorithm == 'AAR' : ts_model = pm.auto_arima(data, **parameters)

        if algorithm == 'PRH': ts_model = Prophet(**parameters)
        
        return ts_model
    

class Include(Create, Format, Preprocess, BaseFormat):

    def include_feature_selection(self, model, key, model_object, X_train, X_test, y_train):
        """
        Apply feature selection on train and test values
        """
        selector_model = model[1] 
        selector_name = selector_model['feature_selection'][0]
        selector_algorithm = self.format_algorithm_string(selector_name)
        selector_parameters = selector_model['feature_selection'][1]

        print(f'Selecting features with {selector_name} for {key}')
        feature_selector = self.create_feature_selection_model(model_object, selector_algorithm, selector_parameters)

        X_train_selected = feature_selector.fit_transform(X_train, y_train)
        X_test_selected = feature_selector.transform(X_test)
        
        return X_train_selected, X_test_selected, selector_name
    
    def include_dimensionality_reduction(self, model, key, X_train, X_test, y_train):
        """
        Apply dimensionality reduction on train and test values
        """
        dimensionality_reduction_model = model[1] 
        dimensionality_reduction_name = dimensionality_reduction_model['dimensionality_reduction'][0]
        dimensionality_reduction_algorithm = self.format_algorithm_string(dimensionality_reduction_name)
        dimensionality_reduction_parameters = dimensionality_reduction_model['dimensionality_reduction'][1]
        

        print(f'Reducing dimensions with {dimensionality_reduction_name} for {key}')
        
        dimensionality_reducer = self.create_dimensionality_reduction_model(dimensionality_reduction_algorithm, dimensionality_reduction_parameters)
        
        if dimensionality_reduction_algorithm == 'AE':
            autoencoder = dimensionality_reducer[0]
            encoder = dimensionality_reducer[1]
            autoencoder.fit(X_train, X_train, **dimensionality_reduction_parameters['fit_parameters'])
            X_train_reduced = encoder.predict(X_train)
            X_test_reduced = encoder.predict(X_test)
            
            return X_train_reduced, X_test_reduced, dimensionality_reduction_name
            
        X_train_reduced = dimensionality_reducer.fit_transform(X_train)
        X_test_reduced = dimensionality_reducer.transform(X_test)
        
        return X_train_reduced, X_test_reduced, dimensionality_reduction_name
    
    
class Build(Include):
    
    def __init__(self, test_dict, feature_selection=False, dimensionality_reduction=False):
        
        self.test_dict = test_dict
        self.test_dict['predictions'] = {}
        self.test_dict['models'] = {}
        self.test_dict['X_test'] = {}
        self.test_dict['y_test'] = {}
        self.test_dict['results'] = {}
        self.feature_selection = feature_selection
        self.dimensionality_reduction = dimensionality_reduction
        
    def build_regression_models(self, models_list, dependent_variable):
        """
        Build multiple regression models on multiple data.
        
        Parameters
        ----------
        models_list: list
            List of regression models and parameters.
            
        dependent_variable: string
            Name of the variable that model is going to predict
        """      
        for key, data in self.test_dict['data'].items():
            for model in models_list:
                
                X_train, X_test, y_train, y_test = self.preprocess_test_data(data, dependent_variable)
                
                model_name, parameters = self.format_models_list(model[0])
                algorithm = self.format_algorithm_string(model_name)
                
                # Dict names
                label = key+model_name
                
                regressor = self.create_regression_model(algorithm , parameters)
                
                if self.feature_selection: 
                    X_train_selected, X_test_selected, selector_name = self.include_feature_selection(model, key, regressor, X_train, X_test, y_train)
                    X_train = X_train_selected
                    X_test = X_test_selected
                    label += selector_name

                if self.dimensionality_reduction:
                    X_train_reduced, X_test_reduced, reducer_name = self.include_dimensionality_reduction(model, key, X_train, X_test, y_train)
                    X_train = X_train_reduced
                    X_test = X_test_reduced
                    label += reducer_name

                print(f'Training regression model {model_name} for {key}')
                regressor.fit(X_train, y_train)
                print(f'Training done!')
                predictions = regressor.predict(X_test)
                print()
                
                self.test_dict['models'][label] = regressor
                self.test_dict['predictions'][label] = predictions
                self.test_dict['X_test'][key] = X_test
                self.test_dict['y_test'][key] = y_test
        
        
    def build_classification_models(self, models_list, dependent_variable):
        """
        Build multiple classification models on multiple data.
        
        Parameters
        ----------
        models_list: list
            List of classification models and parameters.
            
        dependent_variable: string
            Name of the variable that model is going to predict
        """       
        for key, data in self.test_dict['data'].items():
            for model in models_list:

                X_train, X_test, y_train, y_test = self.preprocess_test_data(data, dependent_variable)

                model_name, parameters = self.format_models_list(model[0])
                algorithm = self.format_algorithm_string(model_name)
                
                # Dict names
                label = key+model_name
                
                classifier = self.create_classification_model(algorithm , parameters)
                
                if self.feature_selection: 
                    X_train_selected, X_test_selected, selector_name = self.include_feature_selection(model, key, classifier, X_train, X_test, y_train)
                    X_train = X_train_selected
                    X_test = X_test_selected
                    label += selector_name
                    
                if self.dimensionality_reduction:
                    X_train_reduced, X_test_reduced, reducer_name = self.include_dimensionality_reduction(model, key, X_train, X_test, y_train)
                    X_train = X_train_reduced
                    X_test = X_test_reduced
                    label += reducer_name

                print(f'Training classification model {model_name} for {key}')
                classifier.fit(X_train, y_train)
                print(f'Training done!')
                predictions = classifier.predict(X_test)
                print()

                self.test_dict['models'][label] = classifier
                self.test_dict['predictions'][label] = predictions
                self.test_dict['X_test'][key] = X_test
                self.test_dict['y_test'][key] = y_test
                
    def build_deep_learning_models(self, models_list, dependent_variable, compile_parameters={}, fit_parameters={}):
        """
        Build multiple regression models on multiple data.
        
        Parameters
        ----------
        models_list: list
            List of deep learning models and parameters.
            
        dependent_variable: string
            Name of the variable that model is going to predict
        
        compile_parameters: dictionary
            Parameters that are going to be used for compiling the neural network.
        
        fit_parameters: dictionary
            Parameters that are going to be used for training the neural network
        """   
        for key, data in self.test_dict['data'].items():
            
        
            for model in models_list:
                
                X_train, X_test, y_train, y_test = self.preprocess_test_data(data, dependent_variable)
                
                model_name, layers, compile_parameters, fit_parameters = self.format_models_list(model[0])
                kind = self.format_algorithm_string(model_name)
                
                # Dict names
                label = key+model_name
                
                #Create NN 
                layers = self.format_input_layer(layers, X_train)
                nn = self.create_deep_learning_model(layers, compile_parameters)
                                
                if self.feature_selection: 
                    X_train_selected, X_test_selected, selector_name = self.include_feature_selection(model, key, nn, X_train, X_test, y_train)
                    X_train = X_train_selected
                    X_test = X_test_selected
                    label += selector_name

                if self.dimensionality_reduction:
                    X_train_reduced, X_test_reduced, reducer_name = self.include_dimensionality_reduction(model, key, X_train, X_test, y_train)
                    X_train = X_train_reduced
                    X_test = X_test_reduced
                    label += reducer_name
                
                print(f'Training deeplearning model {model_name} for {key}')
                result = nn.fit(X_train, y_train, **fit_parameters)
                print(f'Training done!')
                predictions = nn.predict(X_test)
                print()

                self.test_dict['models'][label] = nn
                self.test_dict['predictions'][label] = predictions
                self.test_dict['X_test'][key] = X_test
                self.test_dict['y_test'][key] = y_test
                self.test_dict['results'][label] = result
        
             
    def build_time_series_models(self, models_list):
        """
        Build multiple time series models on multiple data.
        
        Parameters
        ----------
        models_list: list
            List of time series models and parameters.
        """   
        for key, data in self.test_dict['data'].items():
            for model in models_list:

                model_name, layers, parameters, test_size = self.format_models_list(model[0])
                algorithm = self.format_algorithm_string(model_name) 
                
                #Change columns names if algorithm is Prophet
                if algorithm == 'PRH':
                    data = self.format_prophet_data(data)
                    
                train, test= self.preprocess_time_series_data(data, test_size)
    
                # Dict names
                label = key+model_name
                
                print(f'Creating time series model {model_name} for {key}')
                ts_model = self.create_time_series_model(algorithm, data , parameters)
                print('Model created!')
                ts_model.fit(train)
                
                if algorithm == 'PRH':
                    #Drop y values from test dataframe to make future dataframe for Prophet.
                    predictions = ts_model.predict(test.drop('y', axis=1))
                    predictions = self.format_prophet_predictions
                    
                    #Format the test and train tables so index is datetime column
                    train, test = self.format_prophet_index(train, test)
                
                else:
                    predictions = ts_model.predict(n_periods = test.shape[0])

                self.test_dict['models'][label] = ts_model
                self.test_dict['predictions'][label] = predictions
                #Because there is no X - y split in time series, train is labeled as X_test and test is labeled as y_test.
                self.test_dict['X_test'][key] = train
                self.test_dict['y_test'][key] = test

