from auto_vtna.VTNA_functions import VTNA_orders, on_click
import numpy as np
import pandas as pd
import itertools
from functools import partial
import matplotlib.pyplot as plt
import matplotlib.ticker as tkr
import warnings

class Automatic_VTNA:
    def __init__(self,data,VTNA_selection,order_range=[-1.5,2.5],resolution=7,deg=5, \
                 fit_metric='RMSE',iterations=7,constraint='monotonic',score_interval=0.15,\
                fixed_order_species=None,initial_mesh_denser=True):
        # Initialise the class:
        self.data=data.copy()
        self.VTNA_selection=VTNA_selection.copy()
        self.order_range=order_range
        self.resolution=resolution
        self.deg=deg
        self.fit_metric=fit_metric
        self.iterations=iterations
        self.constraint=constraint
        self.score_interval=score_interval
        self.fixed_order_species=fixed_order_species
        self.initial_mesh_denser=initial_mesh_denser
        # Perform automatic VTNA calculation:
        self.results, self.best_orders, self.interval, self.best_slope = VTNA_orders(self.data,self.VTNA_selection, \
        self.order_range,self.resolution,self.deg,self.fit_metric,self.iterations,self.constraint, \
        self.score_interval,self.fixed_order_species,self.initial_mesh_denser)

    def identify_error_intervals(self,score_interval=0.15):
        """
        Generates a Pandas dataframe containing the lowest and highest order value which is found in the \
        list of order value combinations that give an overlay score within "score_interval" of the optimal overlay \
        score. The score_interval can be set to any value above 0. 
        """
        # Get the column names except the rightmost one
        col_names = self.results.columns[:-1]
        # Create an empty list to store lists of column values
        combined_values = [[] for _ in range(len(self.results))]
        # Iterate through each row of the DataFrame
        for index, row in self.results.iterrows():
            # Iterate through each column except the rightmost one
            for col_name in col_names:
                # Append the value of each column to the corresponding list
                combined_values[index].append(row[col_name])
        # Convert the list of lists into a numpy 2D array
        order_combinations = np.array(combined_values)
        # Extract the rightmost column as a numpy array
        FQs = self.results.iloc[:, -1].values
        # Obtain the minimum point order combination (minium FQ for variance, RMSE and SE, max FQ for R2).
        minFQ=min(FQs)
        if self.fit_metric=='R2':
            minFQ=max(FQs)
        # Identify order combinations with overlay scores up to "score_interval" higher than that of the best order:
        if self.fit_metric=='R2':
            good_order_combos=[order_combinations[i] for i in range(len(order_combinations)) if FQs[i]>minFQ*(1-score_interval)]
        else: 
            good_order_combos=[order_combinations[i] for i in range(len(order_combinations)) if FQs[i]<minFQ*(1+score_interval)]
        # Identify the smallest and largest order values in the good order combinations list for each normalised species.
        largest_values=[max(column) for column in zip(*good_order_combos)]
        smallest_values=[min(column) for column in zip(*good_order_combos)]
        intervals=list(zip(smallest_values, largest_values))
        intervals=np.array(list(zip(smallest_values, largest_values)))
        # Define a list of the normalised reaction species that haven't been fixed. 
        if type(self.fixed_order_species)!=list:
            fixed_order_list=[self.fixed_order_species]
        else:
            fixed_order_list=self.fixed_order_species
        # Define the list of normalised species and generate the overlay score interval dictionary. 
        normalised_species=[i for i in list(self.VTNA_selection['normalised_species'].keys()) if i not in fixed_order_list]
        interval_dict={'Normalised species':normalised_species,f'Lower order limit {score_interval*100}%':intervals[:,0],f'Upper order limit {score_interval*100}%':intervals[:,1]}
        return pd.DataFrame(interval_dict)

    def plot_orders_vs_overlay(self,points=False,size_scaler=1,size_scaler2=1,title=None,zoom=False,grid=False,color_scaler=1,
    fixed_cbar_resolution=None,interval=False,y_unit='M',title_fontsize=12,overlay_score_range_max=None,
    overlay_score_range_min=None,show_legend_main=True,show_legend_popup=True,decimals_in_colorbar=3,annotate_optimum=True,
    specified_score_interval=None):
        """
        Creates a graph or contour plot visualising the overlay score versus order value matrix \
        for one and two normalised reaction species respectively. If the overlay score versus \
        order value matrix involves 3 or more normalised species, no plot is created. The graph \
        or contour plots can be clicked upon to generate the VTNA overlay plot corresponding to \
        the order values of the click. 
        Args:
            points (bol): Set to True to visualise the order value data points for which overlay \
            scores have been determined. 
            size_scaler (int or float): Factor by which the size of the order value versus overlay \
            score plot is adjusted. 
            size_scaler2 (int or float): Factor by which the size of the VTNA overlay plot generated \
            by clicking on the original plot is adjusted. 
            title (str): The title of the order value versus overlay score plot.
            zoom (int or float): Width of order values around the optimial order point that will be \
            included in the order value versus overlay score plot.
            grid (bol): Set to True to include a grid in the order value versus overlay score plot.
            color_scaler (int or float, ∈<0,1]): Number between 1 and 0 which determines the factor by \
            which the contour plot color bar range is truncated. NB: This will give white sections that \
            are off the contour plot color map.
            resolution_contour (float): Difference in overlay score between separate color shades in \
            the contour plot. 
            interval (bol): Set to True to show each order combination datapoint within the interval_score \
            times the best overlay score value away from the optimal order value point. 
            y_unit (str): The concentration unit of the output concentration profiles. Gets included in the \
            yaxis title. 
            title_fontsize (int): Fontsize of the title of the order value versus overlay score plot.
            overlay_score_range_max (None or float): The maximum overlay score value to be shown in the plot \
            either as a contour plot color bar (2 NS) or the y axis of an ordinary graph (1 NS).
            overlay_score_range_min (None or float): The minimum overlay score value to be shown in the plot \
            either as a contour plot color bar (2 NS) or the y axis of an ordinary graph (1 NS).
            show_legend_main (bol): Set to True to show the legend of the overlay score versus order plot. 
            show_legend_popup (bol): Set to True to show the legend of the VTNA overlay plots generated by \
            mouseclicks on the overlay score versus order plot. 
            decimals_in_colorbar (int): The number of decimals to be shown in the colorbar of an overlay score \
            versus order contour plot. 
            annotate_optimum (bol): Set to True to show the order values of the optimum point of an overlay score \
            versus order contour plot. 

        """
        #if color_scaler>1 or color_scaler<=0:
        #    raise ValueError('color_scaler should be set to ∈<0,1]')
        # Find the species corresponding to the x or x and y axis:
        axis_species=[i.replace('order in ','') for i in list(self.results.columns)[:-1]]
        output_species=self.VTNA_selection['output_species']
        analysis_columns=self.results.columns
        no_normalised_species=len(analysis_columns)-1
        # Create string to include information about reaction species with fixed order value(s) during the 
        # automatic VTNA calculation.
        fixed_order_string=''
        if self.fixed_order_species!=None:
            if type(self.fixed_order_species)==str:
                fixed_order_string=fixed_order_string+f"\nFixed order for {self.fixed_order_species}: {self.VTNA_selection['normalised_species'][self.fixed_order_species]}"
            else:
                for i in self.fixed_order_species:
                    fixed_order_string=fixed_order_string+f"\nFixed order for {i}:{self.VTNA_selection['normalised_species'][i]}."
        if interval==True:
            # Identify the score interval from the calculation or from the specified_score_interval argument.
            score_interval=self.score_interval
            if specified_score_interval!=None:
                if float(specified_score_interval)>0:
                    score_interval=float(specified_score_interval)
                else:
                    warnings.warn("The specified score interval must be a number above 0. The calculation value will be used instead.")
        if no_normalised_species>2:
            return ValueError('The overlay score : order value matrix can only be visualized for 1 or 2 varying normalised reaction species.')
        if no_normalised_species==1:
            x_best=self.best_orders[0]
            # Define x values as orders and y values as fit metric values.
            x_s=np.array(self.results[analysis_columns[0]])
            y_s=np.array(self.results[analysis_columns[1]])
            # Sort the order values from smallest to largest to get one continuous plt line. 
            xy=list(itertools.zip_longest(x_s,y_s))
            xy.sort()
            # Redefine sorted x and y values
            y=[y[1] for y in xy]
            x=[x[0] for x in xy]
            # Find the finest orders interval used to find minimum point
            resolution_orders=str(format(x_s[-1]-x_s[-2],".12f"))
            # Find the index for first non 0 digit in resolution
            for i in range(len(resolution_orders)):
                if resolution_orders[i]!='.' and resolution_orders[i]!='0':
                    break
            # Increase i by 1 if first character in order is '-'
            if str(x_best)[0]=='-':
                i=i+1
            # Create figure with size scaled by size_scaler
            fig, ax = plt.subplots(figsize=(6*size_scaler,4*size_scaler))
            # Plot the fit metric values obtained at each order.
            ax.plot(x,y,marker='o',linestyle='-',markersize=5)
            # Plot the minimum point and insert label information for legend.
            minimum_info=f'Best order: {f"%.{i-1}f" % round(x_best,i-1)}\nUncertainty: {f"%.{i}f" % round(float(resolution_orders[:(i+2)]),i)}\nMinimum {self.fit_metric}: {f"%.{i+2}f" % round(min(y),i+2)}'
            optimum_value=min(y)
            if self.fit_metric=='R2':
                minimum_info=f'Best order: {f"%.{i-1}f" % round(x_best,i-1)}\nUncertainty: {f"%.{i}f" % round(float(resolution_orders[:(i+2)]),i)}\nMaximum {self.fit_metric}: {f"%.{i+2}f" % round(max(y),i+2)}'
                optimum_value=max(y)
            minimum_info=minimum_info+fixed_order_string
            plt.plot(x_best,optimum_value,color='orange',marker='o',label=minimum_info)
            # Define the x and y axis labels. 
            ax.set_xlabel(f'{analysis_columns[0]}',fontsize=12)
            ax.set_ylabel(f'Overlay score: {self.fit_metric}',fontsize=12)
            # Change the x- and y-axis limits in case zoom has been applied. 
            if zoom!=False:
                x1=x_best-zoom/2
                x2=x_best+zoom/2
                plt.xlim(x1,x2)
                y2=[i for i in y if x[y.index(i)]<x2 and x[y.index(i)]>x1]
                y1=min(y2)
                y2=max(y2)
                leeway=(y2-y1)*0.1
                ax.set_ylim(y1-leeway,y2+leeway)
            # Plot a horizontal line showing the datapoints within the interval_score*best overlay score range of the minimum point.
            if interval==True:
                interval_limit=min(y)+min(y)*score_interval
                if self.fit_metric=='R2':
                    interval_limit=max(y)-max(y)*score_interval
                ax.plot([self.order_range[0],self.order_range[1]],[interval_limit,interval_limit],linestyle='--',label=f"{float(score_interval)*100}% from optimal point.")
            # Set the figure title. 
            if title==None:
                fig.suptitle(f'VTNA overlay score versus orders for the formation of {output_species}',fontsize=title_fontsize,y=0.95)
            else:
                fig.suptitle(title,fontsize=title_fontsize,y=0.95)
            # Apply legend and grid if selected.
            if show_legend_main:
                ax.legend()
            if grid==True:
                plt.grid()
            # Define the ylimits if overlay score range max or min have been defined. 
            if overlay_score_range_max!=None:
                ax.set_ylim(top=overlay_score_range_max)
            if overlay_score_range_min!=None:
                ax.set_ylim(bottom=overlay_score_range_min)
        # Code for contour plot to illustrate the overlay score across order values if there are 2 normalised species. 
        if no_normalised_species==2:
            # Obtain the order value arrays x_s and y_s and the overlay scores z_s.
            z_s=np.array(self.results[analysis_columns[2]])
            x_s=np.array(self.results[analysis_columns[0]])
            y_s=np.array(self.results[analysis_columns[1]])
            # Define the x and y values of the minimum point (or maximum for R2).
            x_best=self.best_orders[0]
            y_best=self.best_orders[1]
            # Create figure with size scaled by size_scaler.
            fig, ax = plt.subplots(figsize=(6*size_scaler,5*size_scaler))
            # Identify the max and min fit metric values for setting up color grid for contour plot.
            min_z=min(z_s)
            max_z=max(z_s)
            # Find the finest orders interval used to find minimum point.
            resolution_orders=str(format(y_s[-1]-y_s[-2],".12f"))
            # Find the index for first non 0 digit in resolution_orders.
            for i in range(len(resolution_orders)):
                if resolution_orders[i]!='.' and resolution_orders[i]!='0':
                    break
            # Define q and w to determine number of digits to show for each minimum point order.
            q,w=i,i
            # Increase q and w by 1 if first character in respective order is '-'
            if str(x_best)[0]=='-':
                q=1+i
            if str(y_best)[-1]=='-':
                w=1+i
            if points==True:
                # Plot the data points used to generate the countour plot.
                ax.plot(x_s,y_s,'or',markersize=0.5)
            # Plot the datapoints with an overlay score off the minimum point (or max for R2) score.
            
            if interval==True:
                if self.fit_metric!='R2':
                    x_interval=[x_s[i] for i in range(len(x_s)) if z_s[i]<(score_interval+1)*min(z_s)]
                    y_interval=[y_s[i] for i in range(len(y_s)) if z_s[i]<(score_interval+1)*min(z_s)]
                else:
                    x_interval=[x_s[i] for i in range(len(x_s)) if z_s[i]>(-score_interval+1)*min(z_s)]
                    y_interval=[y_s[i] for i in range(len(y_s)) if z_s[i]>(-score_interval+1)*min(z_s)]
                # Plot the identified interval datapoints in green.
                ax.plot(x_interval,y_interval,markersize=0.7,linestyle='',marker='x',color='b',  
                        label=f'Orders within {score_interval*100}% of optimal O.S.',zorder=2)
            # Print axis labels with species_order from data to assign normalised species.
            ax.set_xlabel(f'{analysis_columns[0]}')
            ax.set_ylabel(f'{analysis_columns[1]}')
            # Plot the data point corresponding to the best VTNA alignment and use label to provide information.
            if self.fit_metric=='R2':
                ax.plot(x_best,y_best,marker='x',color='k',linestyle='',label = f'Best orders: ({round(x_best, q-1):.{int(q-1)}f},{round(y_best, w-1):.{int(w-1)}f})\nUncertainty: {resolution_orders[:(i+2)]}\nMaximum {self.fit_metric}: {round(max(z_s), i+2):.{int(i+2)}f}{fixed_order_string}')
            else:
                ax.plot(x_best,y_best,marker='x',color='k',linestyle='',label=f'Best orders: ({round(x_best, q-1):.{int(q-1)}f},{round(y_best, w-1):.{int(w-1)}f})\nUncertainty: {resolution_orders[:(i+2)]}\nMinimum {self.fit_metric}: {round(min(z_s), i+2):.{int(i+2)}f}{fixed_order_string}')
            # If zoom argument applied, apply xlim and ylim around the minimum point and find the lowest z_value in that interval.
            min_z=min(z_s)
            max_z=max(z_s)
            z_range=max_z-min_z
            max_z=max(z_s)-z_range*(1-color_scaler)
            # If the overlay score range max or min arguments have been defined, prepare the relevant variables.
            cbar_fixing=False
            if overlay_score_range_max!=None:
                cbar_fixing=True
                max_z=overlay_score_range_max
            if overlay_score_range_min!=None:
                cbar_fixing=True
                min_z=overlay_score_range_min
            if zoom!=False:
                if cbar_fixing==True:
                    warnings.warn("overlay_score_range_max and overlay_score_range_min is ignored when zoom isn't set to False.")
                # Define x and y limits to zoom in closer to the minimum point
                min_z=min(z_s)
                ax.set_xlim(x_best-zoom/2,x_best+zoom/2)
                ax.set_ylim(y_best-zoom/2,y_best+zoom/2)
                # Obtain the z_s values in the new domain.
                z_s2=[i for i in z_s if x_s[list(z_s).index(i)]<(x_best+zoom/2) and x_s[list(z_s).index(i)]>(x_best-zoom/2) \
                       and y_s[list(z_s).index(i)]<(y_best+zoom/2) and y_s[list(z_s).index(i)]>(y_best-zoom/2)] 
                # Define the maximum fit metric value in this interval and the difference between this and the minimum value
                max_z=max(z_s2)
                z_range=max_z-min_z
                # Use color_scaler to define to adjust the maximum value of the colorbar
                max_z=max_z-z_range*(1-color_scaler)
            # Define the order value range of the plot
            delta_order=max(x_s)-min(x_s)
            # Define appropriate colorbar values if R2 is used as fit metric
            if self.fit_metric=='R2':
                max_z=max(z_s)
                max_z=max_z+(1-max_z)*0.2
                min_z=min_z+z_range*(1-color_scaler)
            if fixed_cbar_resolution!=None:
                resolution_cbar=float(fixed_cbar_resolution)
            else:
                resolution_cbar=abs((max_z-min_z)/30)
            clev = np.arange(min_z,max_z,resolution_cbar)
            # Create contour plot.
            CS = ax.tricontourf(x_s, y_s, z_s, clev,cmap=plt.cm.viridis_r)
            if self.fit_metric=='R2':
                CS = ax.tricontourf(x_s, y_s, z_s, clev,cmap=plt.cm.viridis)
            cbar=fig.colorbar(CS,orientation='vertical',label=f'Overlay score: {self.fit_metric}',ax=ax,format=tkr.FormatStrFormatter(f'%.{decimals_in_colorbar}g'))
            legend_loc='best'
            if zoom==False:
                best_value_x=x_best+0.045
                best_value_y=y_best+0.045
                if best_value_y-min(y_s)<best_value_y-max(y_s):
                    legend_loc='upper '
                else: 
                    legend_loc='lower '
                if best_value_x-min(x_s)<best_value_x-max(x_s):
                    legend_loc=legend_loc+'right'
                else:
                    legend_loc=legend_loc+'left'
            if show_legend_main:
                ax.legend(loc=legend_loc)
            # Define the title of the figure.
            if title==None:
                fig.suptitle(f'VTNA overlay score versus orders for the formation of {output_species}',fontsize=title_fontsize,y=0.95)
            else:
                fig.suptitle(title,fontsize=title_fontsize,y=0.95)
            # Annotate the minimum point to show its order values. 
            # Adjust position according to the "zoom" value and delta_order if zoom isn't set to False
            if annotate_optimum:
                if zoom!=False:
                    ax.annotate(f'{f"%.{q}f" % round(x_best,q)}, {f"%.{w}f" % round(y_best,w)}',(x_best+0.045*zoom/delta_order, y_best+0.045*zoom/delta_order))
                if zoom==False:
                    ax.annotate(f'{f"%.{q}f" % round(x_best,q)}, {f"%.{w}f" % round(y_best,w)}',(best_value_x,best_value_y))
        # Create a new click function with inputs.
        on_click_with_inputs = partial(on_click,data=self.data,VTNA_selection=self.VTNA_selection,y_unit=y_unit,\
        size_scaler=size_scaler2,fit_metric=self.fit_metric,deg=self.deg,constraint=self.constraint,fixed_order_species=self.fixed_order_species,
        axis_species=axis_species,show_legend=show_legend_popup)
        # Bind the on_click function to the mouse click event.
        cid = fig.canvas.mpl_connect('button_press_event', on_click_with_inputs)
        # Show the plot.
        plt.tight_layout()
        if grid==True:
            plt.grid()
        plt.show()
        return fig

