import numpy as np
import matplotlib.pyplot as plt
from IPython.display import display, clear_output
import time
import pandas as pd
import math
import sys
import os
import matplotlib
import matplotlib.gridspec as gridspec
import matplotlib.ticker as mtick
import types
import time
import warnings
import gc

from IPython import get_ipython
get_ipython().run_line_magic('matplotlib','inline')

warnings.filterwarnings("ignore")

class traintorch:
    def __init__(self,n_custom_plots=0,figsize=(15,20),table_metrics=[],
                 top_rows=1,top_cols=2,plot_width=4,plot_height=4,nrows=2,ncols=1,
                main_grid_hspace=0.5,main_grid_wspace=0.5,window=100,custom_window=[]):

            
        self.top_rows=top_rows
        self.top_cols=top_cols
        self.plot_width=plot_width
        self.plot_height=plot_height
        self.nrows=nrows
        self.ncols=ncols
        self.top_axes=[]
        self.main_grid_hspace=main_grid_hspace
        self.main_grid_wspace=main_grid_wspace
        self.n_splits=22
        self.counter=0
        self.window=window
        self.custom_metrics=[]

        self.main_results=pd.DataFrame()
        self.custom_data=None
        self.total_plots=0
        self.n_custom_plots=n_custom_plots
        self.figsize=figsize
        #this needs to be fixed
        if(custom_window==[] and n_custom_plots>0):
            self.custom_window=[self.window]*n_custom_plots
        else:
            if(isinstance(custom_window, list)):
                self.custom_window=custom_window
            else:
                self.custom_window=[custom_window]

        self._avg_axes=[]
        


        class plot:
            def __init__(self,parent,):
                self.parent=parent
                self.main_grid_hspace=self.parent.main_grid_hspace
                self.main_grid_wspace=self.parent.main_grid_wspace
    
            def chunks_df(self,l, n,method='r'):
                if(method=='c'):
                    for i in range(0, l.shape[1], n):
                        yield l.iloc[:,i:i + n]
                elif(method=='r'):
                    for i in range(0, l.shape[0], n):
                        yield l.iloc[i:i + n,:]
            def tail(self,):
                main_results=pd.DataFrame()
                for metric in self.parent.custom_metrics:
                    temp=metric.window().iloc[-1:,:]
                    temp.reset_index(drop=True, inplace=True)
                    main_results=pd.concat([main_results, temp], axis=1,sort=False)
                self.parent.main_results=main_results

            def create(self,custom_metrics=[]): 
                if(any([item.updated for item in custom_metrics])):
                    
                    self.parent.custom_metrics=custom_metrics
                    if(len(self.parent.custom_metrics)!=self.parent.n_custom_plots):
                        raise Exception ("Data provided does not match the number of custom plots")
                    else:
                        self.parent.total_plots=self.parent.n_custom_plots

                    if(self.parent.counter==0):
                        if((self.parent.top_rows)*(self.parent.top_cols)<self.parent.total_plots):
                            warnings.warn("Total number of plots does not match the number of rows, number of rows has been increased to \
                                          account for the difference")
                            while(((self.parent.top_rows)*(self.parent.top_cols)<self.parent.total_plots)):
                                self.parent.top_rows+=1                       


                    n_splits= self.parent.n_splits
    #                 len(cm_df_overall.index)//2


                    if(self.parent.counter==0):
                        self.parent.figure = plt.figure(figsize=self.parent.figsize)
                        self.parent.main_grid = gridspec.GridSpec(self.parent.nrows,
                                                                  self.parent.ncols,
                                                                  hspace=self.main_grid_hspace,
                                                                  wspace=self.main_grid_wspace)
                        self.parent.top_cell = self.parent.main_grid[0,0:] 
                        self.parent.bottom_cell = self.parent.main_grid[1,0:]
                        self.parent.inner_grid_top = gridspec.GridSpecFromSubplotSpec(
                                                    (self.parent.top_rows*self.parent.plot_width)+(self.parent.top_rows-1),
                                                    (self.parent.top_cols*self.parent.plot_height)+(self.parent.top_cols-1), 
                                                    self.parent.top_cell,hspace=1
                                                    )
                        self.parent.bottom_cell = self.parent.main_grid[1,0:]
                        self.parent.inner_grid_bottom = gridspec.GridSpecFromSubplotSpec(n_splits,3, self.parent.bottom_cell)


                        top_axes=[]
                        for j,k in enumerate(range(0,self.parent.top_rows*self.parent.plot_width,self.parent.plot_width)):
                            for l,m in enumerate(range(0,self.parent.top_cols*self.parent.plot_height,self.parent.plot_height)):
                                temp=plt.subplot(self.parent.inner_grid_top[k+j:k+self.parent.plot_height+j,
                                                                            l+m:(l-1)+m+self.parent.plot_width])
                                top_axes.append(temp)
                        self.parent.top_axes=top_axes
                        self.parent._avg_axes=[]
                        for i,item in enumerate(self.parent.top_axes):
                            if(i<self.parent.total_plots):
                                ax_avg=item.twiny().twinx()
                                self.parent._avg_axes.append(ax_avg)

        #                 self.parent.middle_cell = self.parent.main_grid[1,0:]
        #                 self.parent.inner_grid_middle = gridspec.GridSpecFromSubplotSpec(n_splits,3, self.parent.middle_cell)
        #                 middle_axes=[]

        
                    #adds the main plots
                    for i in range(0,self.parent.n_custom_plots):
                        try:
                            if(self.parent.custom_metrics[i].updated):
                                custom_data=self.parent.custom_metrics[i].window()
                                if(custom_data.empty):
                                    custom_data=pd.DataFrame([0,0,0,0],columns=['No Data Available Yet'])
                                #This can be optimized later
                                
                                top_axes[i].clear()
                                top_axes[i].plot(custom_data.iloc[-1*self.parent.custom_metrics[i].w_size:,:])
                                top_axes[i].legend(self.parent.custom_metrics[i].window().columns)
                                top_axes[i].set_title(self.parent.custom_metrics[i].name)
                                self.parent._avg_axes[i].clear()
                                avg=self.parent.custom_metrics[i].means
                                self.parent._avg_axes[i].plot(avg,linestyle='--',alpha=0.6)
                        except Exception as error:
                            print(error)
                            pass

                    # Adds table 
                    split_df=[] 
                    bottom_axes=[]
                    self.tail()
                    if(not self.parent.main_results.empty):
                        for i,item in enumerate(self.chunks_df(self.parent.main_results.round(6).T,
                                                               math.ceil(len(self.parent.main_results.columns)/2),'r')):
                            if((i+1)%2==0 and (i>0)):
                                loc='bottom right'
                            else:
                                loc='bottom'
                            if(not item.empty):
                                split_df.append(item)

                                temp=plt.subplot(self.parent.inner_grid_bottom[i])
                                temp.axis('off')
                                temp.axis('tight')
                                table=temp.table(cellText=item.values,
                                            loc=loc,
                                           cellLoc='left',
                                            rowLabels=item.index,
#                                             colLabels=item.columns,
                                            )
                                bottom_axes.append(table)
                                table.scale(1, 2)



                        self.parent.split_df=split_df
                        self.parent.bottom_axes=bottom_axes             


                    #moves legends to the bottom of the plots
                    for i,item in enumerate(self.parent.top_axes[:self.parent.n_custom_plots]):
                        box = item.get_position()
                        item.set_position([box.x0, box.y0 + box.height * 0.1,box.width, box.height * 0.9])
                        lines=item.get_lines()
                        if(custom_data.empty):
                            custom_data=pd.DataFrame([0,0,0,0],columns=['No Data Available Yet'])

                        lines[0].get_xydata()
                        item.legend(self.parent.custom_metrics[i].window().columns,loc='upper center', 
                                    bbox_to_anchor=(0.5, -0.1),fancybox=True, shadow=True, ncol=5)                        

                        item.grid()
                        #add scientific formatting
    #                     item.yaxis.set_major_formatter(mtick.FormatStrFormatter('%.2e'))
                        item.set_ylabel('')
                        item.set_xlabel('')  

            
                    #Removes empty plots 
                    for i,ax in enumerate(top_axes):
                        lines=ax.get_lines()
                        try:
                            lines[0].get_xydata()
                        except Exception as error:
                            self.parent.top_axes[i].axis('off')

                    self.get_confusion=False
                    plt.show()
#                     for item in self.parent.top_axes:
#                         item.clear()
#                     for item in self.parent._avg_axes:
#                         item.clear()
                    for item in custom_metrics:
                        if(item.updated):
                            item.updated=False
#                     plt.close(self.parent.figure)
                    clear_output(wait=True)
                    gc.collect()
        self._plot=plot(self,)
        self.create=self._plot.create


class metric:
    def __init__(self,name=None,w_size=10):
        self.name=name
        self.__kwargs=None
        self.counter=0
        self.keys=[]
        self.updated=False
        self.w_size=w_size
        self.last_chunk=pd.DataFrame()
        self.means=[]
        if(not name):
            raise Exception('please provide a name for this metric.')

    def update(self,**kwargs):
        self.updated=True
        self.counter+=1
        self.__kwargs=kwargs
        for key in self.__kwargs.keys():
            if(key in self.__dict__ ):
                self.__dict__[key].append(self.__kwargs[key])
            else:
                self.keys.append(key)
                self.__dict__[key]=[self.__kwargs[key]]
        if(len(self.__dict__[key])%self.w_size==0):
            self.chunk()
            self.chunk_mean()
            for key in self.keys:
                self.__dict__[key]=[]
                gc.collect()
    def frame(self,):
        _dict={}
        for k,v in self.__dict__.items():
            if(k in self.__dict__['keys']):
                _dict[k]=v
        _data=pd.DataFrame(_dict)
        if('x' in _data.columns):
            _data.set_index('x',inplace=True)
        gc.collect()
        return(_data)
    def window(self,):
        _data=pd.concat([self.last_chunk,self.frame()]).reset_index(drop=True)[(-1*self.w_size):]
        _data.index=list(range(1,self.counter+1,1))[(-1*self.w_size):]
        if('x' in _data.columns):
            _data.set_index('x',inplace=True)
        gc.collect()
        return _data

    def chunk(self,method='r'):
        _data=self.frame()
        n=self.w_size
        _output=[]
        if(method=='c'):
            for i in range(0, _data.shape[1], n):
                _output.append(_data.iloc[:,i:i + n])
        elif(method=='r'):
            for i in range(0, _data.shape[0], n):
                _output.append(_data.iloc[i:i + n,:])
        self.last_chunk=_output[0]
        gc.collect()
    def chunk_mean(self,):
        self.means.append(self.last_chunk.mean(axis=0))
        gc.collect()
    def shape(self,):
        gc.collect()
        return self.frame().shape
    
    def __getitem__(self,key):
        gc.collect()
        if(key in self.keys):
            _data=self.frame()
            return(_data[key])
        else:
            return None
    def __len__(self,):
        gc.collect()
        return(len(self.keys))
    def __call__(self,):
        gc.collect()
        return self.frame().iloc[(-1*self.w_size):,:]
