# --*-- co


class lazy_opls_da:

    
    import os
    from glob import glob
    import pandas as pd
    import numpy as np
    import random
    from metbit import opls_da

    from .utility import project_name_generator
    from .utility import lazypair
    

    def __init__(self, data: pd.DataFrame, groups: list, working_dir: str, n_components: int = 2, scaling: str = 'pareto', 
                    estimator: str = 'opls', kfold: int = 3, random_state: int = 94, auto_ncomp: bool = True,  
                    permutation: bool = True, 
                    VIP: bool = True, 
                    linear_regression: bool = True) -> None:

        import os
        from glob import glob
        import pandas as pd
        import numpy as np
        import random
        from metbit import opls_da

        from .utility import project_name_generator
        from .utility import lazypair

        """
        This function takes in a dataframe and a list of y values and returns the project_name model.
        Parameters
        ----------
        data: pandas dataframe
            The dataframe to be used.
        y: list
            The list of y values.

        n_components: int
            The number of components to use.
        lazy_opls_da(data, y, n_components).fit()
        """


        self.groups = groups
        self.n_components = n_components
        self.working_dir = working_dir

        self.random_state = random_state        
        self.estimator = estimator
        self.scale = scaling
        self.kfold = kfold
        self.auto_ncomp = auto_ncomp
        

        data['Class'] = groups
        self.data = data

        self.permutation = permutation
        if permutation == True:
            self.n_permutataion = int(input('Enter the number of permutation: '))
            self.n_jobs = int(input('Enter the number of jobs: '))
        else:
            pass
        
        self.VIP = VIP
        if VIP == True:
            self.VIP_threshold = float(input('Enter the VIP threshold: '))

        self.linear_regression = linear_regression
        if linear_regression == True:
            self.FC_threshold = float(input('Enter the fold change threshold: '))
            self.p_val_threshold = float(input('Enter the p-value threshold: '))

        """
        This function takes in a dataframe and a list of y values and returns the project_name model.
        Parameters
        ----------
        data: pandas dataframe
            The dataframe to be used.
        y: list
            The list of y values.
        n_components: int
            The number of components to use.
        lazy_opls_da(data, y, n_components).fit()
        """
        
        project_name = project_name_generator()
        

        #Remove last / from working_dir
        if working_dir[-1] == '/':
            working_dir = working_dir[:-1]
        else:
            working_dir = working_dir

        #Replace \ with / for windows
        working_dir = working_dir.replace('\\', '/')


        if os.path.exists(working_dir + '/' + project_name + '/element'):
            print('Directory already exist')
        else:
            folder_name_plot = ['loading_plot', 's_plot', 'score_plot']
            folder_name_data = []
            if permutation == True:
                folder_name_plot.append('hist_plot')
            else:
                pass
            if VIP == True:
                folder_name_plot.append('VIP_score_plot')
                folder_name_data.append('VIP_scores')
            else:
                pass
            if linear_regression == True:
                folder_name_plot.append('Volcano_plot')
                folder_name_data.append('Lingress_data')
            else:
                pass



            os.makedirs(working_dir+'/' + project_name + '/element', exist_ok=True)
            for i in folder_name_plot:
                os.makedirs(working_dir+'/' + project_name + '/element/plots/' + i, exist_ok=True)
            for i in folder_name_data:
                os.makedirs(working_dir+'/' + project_name + '/element/data/' + i, exist_ok=True)

            os.makedirs(working_dir+'/' + project_name + '/main')

        #Create dictionary to store the path
        dir = glob(working_dir + '/' + project_name + '/element/*/*/')

        #Create dictionary to store the path
        path = {}
        for i in dir:
            path[i.split('/')[-2]] = i


        self.path = path

        #Print summary model as table text format
        Summary = f"""
        Project Name: {project_name}
        Number of groups: {len(data['Class'].unique())}
        Number of samples: {len(data)}
        Number of features: {len(data.columns) - 1}
        Number of components: {n_components}
        Estimator: {estimator}
        Scaling: {scaling}
        Kfold: {kfold}
        Random state: {random_state}
        Auto ncomp: {auto_ncomp}
        Working directory: {working_dir}
        Permutation: {permutation}
        VIP: {VIP}
        Linear regression: {linear_regression}
        """

        return print(Summary)



    def fit(self, marker_color: dict = None) -> None:

        from .metbit import opls_da
        from lingress import lin_regression
        from .utility import lazypair

        from glob import glob
        import pandas as pd
        import numpy as np
        import random



     
        data = self.data
        n_components = self.n_components
        path = self.path
        scale = self.scale

        marker_color = marker_color


        #Create object attribute
        lazy = lazypair(data, 'Class')
        data_list = lazy.get_dataset()
        name_save = lazy.get_name()

        for i in range(len(data_list)):
            
            df = data_list[i]
            name = name_save[i]

            X = df.drop('Class', axis=1)
            y = df['Class']
            feature_names = X.columns
            # Check if feature names can be converted to float
            try:
                feature_names = feature_names(float).tolist()
            except:
                feature_names = feature_names.tolist()

            #OPLS-DA
            oplsda_mod = opls_da(X=X, y=y, features_name = feature_names, n_components=n_components, scale=scale, 
                                    estimator=self.estimator, kfold=self.kfold, random_state=self.random_state, 
                                    auto_ncomp = self.auto_ncomp)
            oplsda_mod.fit()


            if len(df) <= 100:
                marker_pca_size = 35
            else:
                marker_pca_size = 16
            #Score plot
            oplsda_mod.plot_oplsda_scores(color_dict=marker_color, marker_size=marker_pca_size).write_html(path['score_plot'] + name + '_score_plot.html')
            oplsda_mod.plot_oplsda_scores(color_dict=marker_color, marker_size=marker_pca_size).write_image(path['score_plot'] + name + '_score_plot.png')

            if len(feature_names) >= 100:
                marker_load_size = 5
            else:
                marker_load_size = 20
            #Loading plot
            oplsda_mod.plot_loading(xaxis_title='Features', marker_size=marker_load_size).write_html(path['loading_plot'] + name + '_loading_plot.html')
            oplsda_mod.plot_loading(xaxis_title='Features', marker_size=marker_load_size).write_image(path['loading_plot'] + name + '_loading_plot.png')

            #S plot
            oplsda_mod.plot_s_scores(marker_size=marker_load_size).write_html(path['s_plot'] + name + '_s_plot.html')
            oplsda_mod.plot_s_scores(marker_size=marker_load_size).write_image(path['s_plot'] + name + '_s_plot.png')

            #VIP score plot
            if self.VIP == True:
                oplsda_mod.vip_scores()
                oplsda_mod.get_vip_scores().to_csv(path['VIP_scores'] + name + '_VIP_scores.csv')
                oplsda_mod.vip_plot(threshold=self.VIP_threshold, marker_size=marker_load_size).write_html(path['VIP_score_plot'] + name + '_VIP_score_plot.html')
                oplsda_mod.vip_plot(threshold=self.VIP_threshold, marker_size=marker_load_size).write_image(path['VIP_score_plot'] + name + '_VIP_score_plot.png')
            else:
                pass

            #Permutation test
            if self.permutation == True:
                oplsda_mod.permutation_test(n_permutations=self.n_permutataion, n_jobs=self.n_jobs)
                oplsda_mod.plot_hist().write_html(path['hist_plot'] + name + '_hist_plot.html')
                oplsda_mod.plot_hist().write_image(path['hist_plot'] + name + '_hist_plot.png')
            else:
                pass

            #Linear regression
            if self.linear_regression == True:
                lin_ = lin_regression(x=X, target=y, label=y, features_name=feature_names)
                lin_.create_dataset()
                lin_.fit_model(adj_method='fdr_bh')
                lin_.volcano_plot(fc_cut_off=self.FC_threshold, p_val_cut_off=self.p_val_threshold).write_html(path['Volcano_plot'] + name + '_Volcano_plot.html')
                lin_.volcano_plot(fc_cut_off=self.FC_threshold, p_val_cut_off=self.p_val_threshold).write_image(path['Volcano_plot'] + name + '_Volcano_plot.png')
                lin_.report().to_csv(path['Lingress_data'] + name + '_Lingress_data.csv', index=False)
            else:
                pass
        
        return print('Model has been fitted successfully')

