import pandas as pd
from tabulate import tabulate
import plotly.graph_objects as go
from plotly.subplots import make_subplots
import os
import cobra
import sys
import numpy as np

from IPython.display import display
from IPython import get_ipython

from .configuration import Config

from typing import Union

class PAMValidator(object):
    RESULT_DIR = os.path.join(os.path.split(os.getcwd())[0], 'Results')
    OUT_FILE = os.path.join(RESULT_DIR, 'flux_rates_vs_glc.xls')
    GLUCOSE_EXCHANGE_RXNID = Config.GLUCOSE_EXCHANGE_RXNID
    OXYGEN_UPTAKE_RXNID = Config.OXYGEN_UPTAKE_RXNID
    BIOMASS_REACTION = Config.BIOMASS_REACTION
    ACETATE_EXCRETION_RXNID = Config.ACETATE_EXCRETION_RXNID
    CO2_EXHANGE_RXNID = Config.CO2_EXHANGE_RXNID
    PHYS_RXN_IDS = Config.PHYS_RXN_IDS
    MW_GLC = 180.15588 #g/mol
    MW_ACETATE = 59.04
    MW_CO2 = 44.009
    GRADIENT_MAX = 11 #mmol/gdw/h
    GRADIENT_STEP = 0.5 #mmol/gdw/h
    GRADIENT_MIN = 0 #mmol/gdw/h

    def __init__(self, model, phys_file, configuration=Config):
        self.GLUCOSE_EXCHANGE_RXNID = configuration.GLUCOSE_EXCHANGE_RXNID
        self.OXYGEN_UPTAKE_RXNID = configuration.OXYGEN_UPTAKE_RXNID
        self.BIOMASS_REACTION = configuration.BIOMASS_REACTION
        self.ACETATE_EXCRETION_RXNID = configuration.ACETATE_EXCRETION_RXNID
        self.CO2_EXHANGE_RXNID = configuration.CO2_EXHANGE_RXNID
        self.PHYS_RXN_IDS = configuration.PHYS_RXN_IDS

        self.model = model
        self.physiology_data = pd.read_excel(phys_file, sheet_name='Fluxes')
        self.yield_data = pd.read_excel(phys_file, sheet_name='Yields')

        self.results = pd.DataFrame(columns=['Y_biomass', 'Y_acetate', 'Y_co2']+self.PHYS_RXN_IDS)
        self._additional_rxns = []

    @staticmethod
    def check_kernel_type():
        try:
            ipy_str = str(type(get_ipython()))
            if 'zmqshell' in ipy_str:
                # jupyter
                return True
            if 'terminal' in ipy_str:
                # ipython
                return False
        except:
            # terminal
            return False

    def add_rxn_to_save(self, rxn_list: Union[str, list]):
        if isinstance(rxn_list, str):
            rxn_list = [rxn_list]

        self.PHYS_RXN_IDS += rxn_list
        for rxn in rxn_list:
            self.results[rxn] = 0
        self._additional_rxns.append(rxn_list)

    def validate_literature(self):
        self.parse_data()
        self.run_simulations()
        # for pretty printing: check if the code is ran from a jupyter notebook or not
        jupyter = self.check_kernel_type()

        if jupyter:
            pd.set_option('display.max_rows', None)
            display(self.results)
        else:
            print(tabulate(self.results, headers='keys', tablefmt='psql'))
        self.plot_mu()

    def validate_range(self, c_uptake_rxn: str = GLUCOSE_EXCHANGE_RXNID, pfba: bool = False, show: bool = True, save: bool = False):
            # reset result dataframe
            self.results = pd.DataFrame(columns=['Y_biomass', 'Y_acetate', 'Y_co2'] + self.PHYS_RXN_IDS)
            try:
                self.parse_data()
            except:
                pass
            self.run_simulations_gradient(c_uptake_rxn=c_uptake_rxn, pfba=pfba)
            jupyter = self.check_kernel_type()

            if jupyter:
                pd.set_option('display.max_rows', None)
                display(self.results)
            else:
                print(tabulate(self.results, headers='keys', tablefmt='psql'))
            if show:
                self.plot_vs_mu(c_uptake_rxn)
                if self._additional_rxns != []:
                    self.custom_plot(self._additional_rxns, c_uptake_rxn=c_uptake_rxn)
            if save:
                # make result directory if it not already exists
                if not os.path.exists(self.RESULT_DIR):
                    os.mkdir(self.RESULT_DIR)
                self.results.to_excel(self.OUT_FILE, engine='openpyxl')


    def parse_data(self):
        self.physiology_data = pd.melt(self.physiology_data, id_vars='Reaction_ID', var_name='Reference',
                                       value_vars=['Rijsewijk_2011', 'Nanchen_2006_1', 'Nanchen_2006_2', 'Nanchen_2006_3', 'Nanchen_2006_4'])

    def run_simulations(self):
        i=0
        print('Running simulations to validate the model...\n')
        #physiology dataset
        for ref in pd.unique(self.physiology_data.Reference):
            #set glucose uptake rate
            df = self.physiology_data[self.physiology_data.Reference == ref]
            bound = df[df.Reaction_ID == self.GLUCOSE_EXCHANGE_RXNID].value.iloc[0]
            self.set_glc_uptake_bounds(bound)
            self.model.optimize()
            if self.model.solver.status == 'optimal':
                self.store_results()
        print('Still validating...\n')
        for ref in pd.unique(self.yield_data.Reference):
            df = self.yield_data[self.yield_data.Reference == ref]
            for glc_flux in pd.unique(df.Glucose_IN):
                self.set_glc_uptake_bounds(glc_flux)
                self.model.optimize()
                if self.model.solver.status == 'optimal':
                    self.store_results()
                if i == 2:
                    cobra.core.solution.get_solution(self.model).fluxes.to_csv(
                        os.path.join(os.path.split(os.getcwd())[0], 'Results', 'test_out_mu.tsv'), sep='\t')
                i+=1

    def run_simulations_gradient(self, pfba: bool = False, c_uptake_rxn: str = GLUCOSE_EXCHANGE_RXNID):
        #run simulations for increasing glucose consumption rates
        print('Running simulations to validate the model...\n')
        self.model.reset_objective()
        for flux in np.arange(self.GRADIENT_MIN, self.GRADIENT_MAX, self.GRADIENT_STEP):
            self.model.change_reaction_bounds(c_uptake_rxn, -flux, -flux)
            #if glucose is not the substrate turn off glucose uptake reaction
            if c_uptake_rxn != self.GLUCOSE_EXCHANGE_RXNID:
                self.set_glc_uptake_bounds(0)
            if pfba: self.model.reset_objective()
            self.model.optimize()
            if self.model.solver.status == 'optimal':
                if pfba: self.model.pfba(fraction_of_optimum=0.1, proteins=True, reactions=False)
                self.store_results(c_uptake_rxn=c_uptake_rxn)
            print('Done running simulations\n')

    def run_simulations_glc_o2_gradient(self, oxygen_gradient: list, params_to_save:Union[str, list] = 'R_TranslationalProteinSector'):
        """
        Function to run simulations of different oxygen gradients for a range of growth rates
        This will simulate growth for the entire range of glucose concentrations for each oxygen uptake rate as given by the input.
        :param
            oxygen_gradient: list
                List of upperbounds for the oxygen uptake reaction to loop over
            params_to_save: optional, string or list
                which parameter(s) to save for further analysis (default: translational protein sector constraint)
        :return:
            results: list of dataframes
                saves the growth rate, glucose uptake rate and the user defined parameters for each oxygen uptake rate
                separate dataframes
        """
        print('Running simulations of glucose and oxygen gradients...\n')
        result = list()
        if isinstance(params_to_save,str): params_to_save = [params_to_save]
        rxns_to_save = [self.BIOMASS_REACTION, self.GLUCOSE_EXCHANGE_RXNID, self.OXYGEN_UPTAKE_RXNID] + params_to_save
        for o2 in oxygen_gradient:
            df = pd.DataFrame(columns=rxns_to_save)
            self.set_rxn_uptake_bounds(self.OXYGEN_UPTAKE_RXNID, o2)
            print(self.model.reactions.get_by_id(self.OXYGEN_UPTAKE_RXNID).bounds)
            for glc_flux in range(0, self.GRADIENT_MAX, self.GRADIENT_STEP):
                self.set_glc_uptake_bounds(glc_flux/10)
                self.model.reset_objective()
                self.model.optimize()
                if self.model.solver.status == 'optimal':
                    df = self.store_manual_results(rxns_to_save, df)
            result += [df]
        return result

    def run_simulations_ups(self, ups_gradient: list, params_to_save:Union[str, list] = 'R_TranslationalProteinSector'):
        """
        Function to run simulations with increasing unused enzyme sectors proportions for a range of growth rates
        This will simulate growth for the entire range of glucose concentrations for a range of fractions of ups_0 as
        given by the input.
        :param:
            ups_gradient: list
                List of upperbounds for the oxygen uptake reaction to loop over
            params_to_save: optional, string or list
                which parameter(s) to save for further analysis (default: translational protein sector constraint)
        :return:
            results: list of dataframes
                saves the growth rate, glucose uptake rate and the user defined parameters for each oxygen uptake rate
                separate dataframes
        """
        print('Running simulations of different unused protein intercepts and glucose gradients...\n')
        result = list()
        if isinstance(params_to_save,str): params_to_save = [params_to_save]
        rxns_to_save = [self.BIOMASS_REACTION, self.GLUCOSE_EXCHANGE_RXNID, self.OXYGEN_UPTAKE_RXNID] + params_to_save
        unused_protein_sector = self.model.sectors.get_by_id('UnusedProteinSector')
        ups_0 = unused_protein_sector.ups_0[0]
        for perc in ups_gradient:
            df = pd.DataFrame(columns=rxns_to_save)
            unused_protein_sector.ups_0[0] = ups_0*perc
            # self.set_rxn_uptake_bounds(self.OXYGEN_UPTAKE_RXNID, perc)
            for glc_flux in range(0, self.GRADIENT_MAX, self.GRADIENT_STEP):
                self.set_glc_uptake_bounds(glc_flux/10)
                self.model.reset_objective()
                self.model.optimize()
                if self.model.solver.status == 'optimal':
                    df = self.store_manual_results(rxns_to_save, df)
            result += [df]
        return result


    def set_glc_uptake_bounds(self, bound:int):
        rxn_glc_in = self.model.reactions.get_by_id(self.GLUCOSE_EXCHANGE_RXNID)
        #check if model is reversible
        if self.GLUCOSE_EXCHANGE_RXNID[-1] == 'b':
            rxn_glc_in._lower_bound, rxn_glc_in._upper_bound = bound, bound
        else:
            rxn_glc_in._lower_bound, rxn_glc_in._upper_bound = -bound, -bound
        rxn_glc_in.update_variable_bounds()

    def set_rxn_uptake_bounds(self, rxn_id:str, bound:int,
                              upper:bool = True, lower:bool = False):
        rxn = self.model.reactions.get_by_id(rxn_id)
        if rxn_id[-1] == 'b':
            if upper: rxn._upper_bound = bound
            if lower: rxn._lower_bound = bound
        else:
            if upper: rxn._lower_bound = -bound
            if lower: rxn._upper_bound = -bound
        rxn.update_variable_bounds()

    def store_results(self, c_uptake_rxn: str = GLUCOSE_EXCHANGE_RXNID):
        result = {}
        c_upt_flux = self.model.reactions.get_by_id(c_uptake_rxn).flux
        #store the relevant flux values
        if c_uptake_rxn != self.GLUCOSE_EXCHANGE_RXNID:
            for i, rxn in enumerate(self.PHYS_RXN_IDS):
                if rxn == self.GLUCOSE_EXCHANGE_RXNID:
                    self.PHYS_RXN_IDS[i] = c_uptake_rxn
                    self.results = pd.DataFrame(columns= self.PHYS_RXN_IDS)
                    # self.results = pd.DataFrame(columns=[ 'Y_biomass', 'Y_acetate', 'Y_co2']+self.PHYS_RXN_IDS)
        for id in self.PHYS_RXN_IDS:
            #getting reaction or variable
            try:result[id]=self.model.reactions.get_by_id(id).flux
            except:result[id]=self.model.variables[id].primal
            #calculate yields based on glucose for exchange reactions
            if ('EX' in id or 'BIOMASS' in id) and 'glc' not in id and '_o2' not in id and abs(c_upt_flux)> 0:
                if 'BIOMASS' in id:
                    id_ext = 'biomass'
                    MW = 1
                if 'ac' in id:
                    id_ext = 'acetate'
                    MW = 1#self.MW_ACETATE * 1e-3 #[mmol/mol]
                if '_co2' in id:
                    id_ext = 'co2'
                    MW = 1#self.MW_CO2 * 1e-3 #[mmol/mol]
                yield_id = 'Y_'+ id_ext
                result[yield_id]= (self.model.reactions.get_by_id(id).flux * MW) / (c_upt_flux * self.MW_GLC * 1e-3)
        #appending the results to the dataframe
        result_series = pd.Series(result)
        #sort the series to match the format of the output dataframe
        result_series_sorted = result_series[self.results.columns.to_list()]
        #add the resulting row to the result dataframe
        self.results.loc[len(self.results.index)] = result_series_sorted.to_list()

    def store_manual_results(self, vars_to_save:list, df:pd.DataFrame):
        row = dict()
        for id in vars_to_save:
            try:
                #for variables
                row[id] = self.model.variables[id].primal
            except:
                #for reactions:
                row[id] = self.model.reactions.get_by_id(id).flux
        # appending the results to the dataframe
        df = pd.concat([df, pd.DataFrame({key: pd.Series(value) for key, value in row.items()})], ignore_index=True)
        return df

    def plot_mu(self):
        print('Plotting the results against reference values. \n')
        # Initialize figure with subplots
        fig = make_subplots(
            rows=2, cols=2,
            specs=[[{"type": "scatter"}, {"type": "scatter"}], [{"type": "scatter"}, {"type": "scatter"}]],
        subplot_titles= ('Growth rate', 'Biomass yield', 'CO2 excretion', 'CO2 yield', 'Acetate excretion', 'Acetate yield'))
        result = self.results.__deepcopy__()
        result = result.merge(self.yield_data, left_on=self.GLUCOSE_EXCHANGE_RXNID, right_on='Glucose_IN')
        #print(tabulate(result, headers='keys', tablefmt='psql'))
        # growth rate
        fig.add_trace(go.Scatter(x=result[self.BIOMASS_REACTION], y=result['mu'], mode='markers', showlegend=False, marker=dict(color='blue')), row=1, col=1)
        fig.add_trace(go.Scatter(x=[0,0.5,1.4],y=[0,0.5,1.4], mode='lines', showlegend=False,
                                 line=dict(color='black', dash='dash', width=1)),row=1, col=1)
        fig.update_xaxes(title_text='Simulated growth rate [1/h]',row=1, col=1)
        fig.update_yaxes(title_text='Experimentally determined growth rate [1/h]',row=1, col=1)

        # biomass yield
        fig.add_trace(go.Scatter(x=result['Y_biomass'], y=result['Yield_biomass'], mode='markers',showlegend=False, marker=dict(color='blue')), row=1, col=2)
        fig.add_trace(go.Scatter(x=[0, 0.5, 0.6], y=[0, 0.5, 0.6], mode='lines',
                                 showlegend=False,
                                 line=dict(color='black', dash='dash', width =1)), row=1, col=2)
        fig.update_xaxes(title_text='Simulated biomass yield [g_x/g_glc]', row=1, col=2)
        fig.update_yaxes(title_text='Experimentally determined biomass yield [g_x/g_glc]', row=1, col=2)

        # CO2 yield
        fig.add_trace(go.Scatter(x=result['Y_co2'], y=result['Yield_co2'], mode='markers', showlegend=False, marker=dict(color='blue')), row=2, col=2)
        fig.add_trace(go.Scatter(x=[0, 15, 30], y=[0, 15, 30], mode='lines', showlegend=False,
                                 line=dict(color='black', dash='dash', width =1)), row=2, col=2)
        fig.update_xaxes(title_text='Simulated CO2 yield [mmol_x/g_glc]', row=2, col=2)
        fig.update_yaxes(title_text='Experimentally determined CO2 yield [mmol_co2/g_glc]', row=2, col=2)

        # CO2 out
        fig.add_trace(go.Scatter(x=result[self.CO2_EXHANGE_RXNID], y=result['CO2_OUT'], mode='markers', showlegend=False,
                                 marker=dict(color='blue')), row=2, col=1)
        fig.add_trace(go.Scatter(x=[0, 15, 30], y=[0, 15, 30], mode='lines', showlegend=False,
                                 line=dict(color='black', dash='dash', width=1)), row=2, col=1)
        fig.update_xaxes(title_text='Simulated CO2 exchange rate [mmol/g_cdw/h]', row=2, col=1)
        fig.update_yaxes(title_text='Experimentally determined CO2 exchange rate [mmol/g_cdw/h]', row=2, col=1)
        #fig.write_image('Results/Images/mu_co2_plots.png') #depends on the kaleido or orca package

        fig.show()
        #fig.write_image('here.png')

    def plot_vs_mu(self, c_uptake_rxn:str = GLUCOSE_EXCHANGE_RXNID):
        print('Plotting the results against reference values. \n')
        # Initialize figure with subplots
        fig = make_subplots(
            rows=2, cols=2,
            specs=[[{"type": "scatter"}, {"type": "scatter"}], [{"type": "scatter"}, {"type": "scatter"}]],
        subplot_titles= ('Glucose consumption', 'Acetate production', 'O2 consumption','CO2 production'))

        self.results[c_uptake_rxn] = abs(self.results[c_uptake_rxn])
        self.results[configuration.OXYGEN_UPTAKE_RXNID] = abs(self.results[configuration.OXYGEN_UPTAKE_RXNID])
        self.yield_data[c_uptake_rxn] = abs(self.yield_data[c_uptake_rxn])
        self.yield_data[self.OXYGEN_UPTAKE_RXNID] = abs(self.yield_data[self.OXYGEN_UPTAKE_RXNID])
        self.yield_data[self.BIOMASS_REACTION] = abs(self.yield_data[self.BIOMASS_REACTION])

        #glucose consumption
        fig.add_trace(go.Scatter(x=self.results[c_uptake_rxn], y=self.results[self.BIOMASS_REACTION],
                                 showlegend=False,  mode='lines', marker=dict(color='blue')), row=1, col=1)
        fig.add_trace(go.Scatter(x=self.yield_data[c_uptake_rxn],
                                 y=self.yield_data[self.BIOMASS_REACTION], mode='markers', showlegend=False,
                                 line=dict(color='black', dash='dash', width=1)),row=1, col=1)
        fig.update_yaxes(title_text='Growth rate [1/h]',row=1, col=1)
        fig.update_xaxes(title_text='Glucose consumption [mmol/gdw/h]',row=1, col=1)

        #acetate excretion
        fig.add_trace(go.Scatter(x=self.results[c_uptake_rxn], y=self.results[
            self.ACETATE_EXCRETION_RXNID],
                                 showlegend= False, mode='lines', marker=dict(color='blue')), row=1, col=2)
        fig.add_trace(go.Scatter(x=self.yield_data[c_uptake_rxn],
                                 y=self.yield_data[self.ACETATE_EXCRETION_RXNID],
                                 mode='markers', showlegend=False,
                                 line=dict(color='black', dash='dash', width=1)), row=1, col=2)
        fig.update_xaxes(title_text='Glucose consumption [mmol/gdw/h]', row=1, col=2)
        fig.update_yaxes(title_text='Acetate production rate [mmol/gdw/h]', row=1, col=2)

        # O2 uptake
        fig.add_trace(
            go.Scatter(x=self.results[c_uptake_rxn], y=self.results[configuration.OXYGEN_UPTAKE_RXNID], mode='lines',
                       showlegend=False, marker=dict(color='blue')), row=2, col=1)
        fig.add_trace(go.Scatter(x=self.yield_data[c_uptake_rxn],
                                 y=self.yield_data[self.OXYGEN_UPTAKE_RXNID],
                                 mode='markers', showlegend=False,
                                 line=dict(color='black', dash='dash', width=1)), row=2, col=1)
        fig.update_xaxes(title_text='Glucose consumption [mmol/gdw/h]', row=2, col=1)
        fig.update_yaxes(title_text='O2 consumption rate [mmol/gdw/h]', row=2, col=1)

        #CO2 excretion
        fig.add_trace(go.Scatter(x=self.results[c_uptake_rxn], y=self.results[
            configuration.CO2_EXHANGE_RXNID], mode='lines',
                                 showlegend=False, marker=dict(color='blue')), row=2, col=2)
        fig.add_trace(go.Scatter(x=self.yield_data[c_uptake_rxn],
                                 y=self.yield_data[self.CO2_EXHANGE_RXNID],
                                 mode='markers', showlegend=False,
                                 line=dict(color='black', dash='dash', width=1)), row=2, col=2)
        fig.update_xaxes(title_text='Glucose consumption [mmol/gdw/h]', row=2, col=2)
        fig.update_yaxes(title_text='CO2 production rate [mmol/gdw/h]', row=2, col=2)

        fig.show()

    def custom_plot(self, rxn_ids:list, valid_dataframe: pd.DataFrame = None, xaxis:str = None, c_uptake_rxn:str = GLUCOSE_EXCHANGE_RXNID):
        """
        Function to plot the results of custom reactions
        Parameters
        ----------
        rxn_ids: list of rxn_ids (str)
            reaction identifiers of the reactions to be plotted
        valid_dataframe: pandas.DataFrame, optional
            A dataframe with experimental data to validate the results with. The columns should be
             the same as the rxn_id of the reaction to be plotted and the reaction which should be
             plotted on the x-axis (by default the glucose exchange reaction 'EX_glc__D_e_b').
             If the dataframe is not provided, only the simulation results will be plotted.
        xaxis: str, optional
            the reaction identifier of the reaction which should be plotted on the x-axis
            (default: 'EX_glc__D_e_b')

        Returns
        -------
        prints scatter plots of the model simulations vs experimental datapoints (if provided)
        """
        if xaxis is None: xaxis = c_uptake_rxn
        fig = make_subplots(
            rows=1, cols=len(rxn_ids),
            specs=[[{"type": "scatter"}]*len(rxn_ids)],
            subplot_titles=tuple(rxn_id[0] for rxn_id in rxn_ids))
        for i, rxn in enumerate(rxn_ids):
            fig.add_trace(
                go.Scatter(x=self.results[xaxis], y=self.results[rxn[0]], mode='lines',
                           showlegend=False, marker=dict(color='blue')), row=1, col=i+1)
            if valid_dataframe is not None:
                fig.add_trace(go.Scatter(x=valid_dataframe[xaxis],
                                        y=valid_dataframe[rxn],
                                        mode='markers', showlegend=False,
                                        line=dict(color='black', dash='dash', width=1)), row=1, col=i+1)
            fig.update_xaxes(title_text=xaxis+' [mmol/gdw/h]', row=1, col=i+1)
            fig.update_yaxes(title_text=rxn[0] +' [mmol/gdw/h]', row=1, col=i+1)

            fig.show()

