import numpy as np, pandas as pd
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
import os, copy

"""
ic modules provides a class (icCMS) handling ic operation based on a pandas dataframe
The ic DataFrame is based on geom which contains several information on the ecal
geometry.
The variable ic.geom (pandas dataframe) contains all the ecal mapping (electronic, geometry...)
"""


#######################################################################################
### ieta range definition, this can be customized
### the present definition matches the one used for ic combination
def ieta_range(x):
    if   np.abs(x) <   0.5 : return -100
    elif np.abs(x) <  10.5 : return 0
    elif np.abs(x) <  25.5 : return 1
    elif np.abs(x) <  45.5 : return 2
    elif np.abs(x) <  65.5 : return 3
    elif np.abs(x) <  85.5 : return 5
    elif np.abs(x) <  95.5 : return 6
    elif np.abs(x) < 105.5 : return 7
    elif np.abs(x) < 117.5 : return 8
    elif np.abs(x) < 125.5 : return 9
    return -1

#######################################################################################
### ecal eb module defintion along ieta
def eta_module(x):
    imod = +99
    if   np.abs(x) <   0.5 : return -999
    elif np.abs(x) <   5.5 : imod = 0.5
    elif np.abs(x) <  25.5 : imod = 1
    elif np.abs(x) <  45.5 : imod = 2
    elif np.abs(x) <  65.5 : imod = 3
    elif np.abs(x) <  85.5 : imod = 4
    return imod * np.sign(x)

#######################################################################################
### The main geometry pandas DataFrame
geomDefFile = os.path.dirname(__file__) + '/data/ecalMapping.txt'
geom  = pd.read_csv(geomDefFile,sep='\t')

### add some variables
geom['eta'] = geom.groupby('ieta')['eta'].transform('mean') ## redefined eta
geom['eta_module'] = geom['ieta'].apply(eta_module)         ## module number in eta
geom['phi_module'] = np.where(np.abs(geom['eta_module'])< 4.5, (geom['iy']-1)//20+1, 99 )## module number in phi
geom['ieta_range'] = geom['ieta'].apply(ieta_range)         ## ieta range as defined for IC combination

#######################################################################################
### helper function for plotting ICs
### based on older numpy functions from Fabrice (might not be the most adequate)
def plot2D_eb( axis, data, title = '', zAxis = [0.95,1.05] , cmap = 'terrain'):
    ics = np.full( (171,360) , np.nan)
    ics[  data[ 'ix']+85 , data[ 'iy']-1] = data[data.columns[2]]

    eb = axis.pcolormesh( np.linspace(0.5,360.5,num=361), np.linspace(-85.5,85.5,num=172),
                          ics, cmap = cmap, vmin = zAxis[0], vmax = zAxis[1]  )
    axis.set_xlabel('i$\phi$')
    axis.set_ylabel('i$\eta$')
    axis.set_title( title )
    plt.colorbar( eb , ax = axis  )
    return eb

def plot2D_ee( axis, data, title = '', zAxis = [0.95,1.05] , cmap = 'terrain'):
    ics = np.full( (100,100) , np.nan)
    ics[  data[ 'ix']-1 , data[ 'iy']-1] = data[data.columns[2]]

    ee= axis.pcolormesh( np.linspace(0.5,100.5,num=101), np.linspace(0.5,100.5,num=101),
                          ics.T, cmap = cmap, vmin = zAxis[0], vmax = zAxis[1]  )
    axis.set_xlabel('i$_x$')
    axis.set_ylabel('i$_y$')
    axis.set_title(  title )
    plt.colorbar( ee , ax = axis  )
    return ee


#######################################################################################
def join_eb_ee( eb, ee, name ='' ):
    """
    combined 2 icCMS (one for EB, one for EE) into a single one for all Ecal
    """
    out = ee.copy()
    out.iov = ee.iov.where( eb.iov['iz'] != 0, eb.iov )
    out.name = name
    return out

def icCovCor(ic1, ic2, var = 'ic', groupby = 'ieta'):
    """
    Compute correlation (and covariance) between 2 icCMS
    var  : reference variable in ics
    group: used in groupby (directly passed to groupby)
    returns the covariance and the correlation as pandas Series (ic.geom style)
    """
    rel = [ i[var] - i[var].groupby(groupby).transform('mean') for i in [ic1,ic2] ]
    std = [ i[var].groupby(groupby).transform('std')           for i in [ic1,ic2] ]
    cov = (rel[0]*rel[1]).groupby(groupby).transform('mean')
    cor = cov / (std[0]*std[1])

    return cov, cor

#######################################################################################
class icCMS:
    def __init__( self, iov = None, name = '', useErrors = True) :
        """
        class icCMS allows for ic computation and visualisation. It is based on the geom variable
        to allow for plotting, averaging other any variable available in geom.
        It can also plot any variable in geom.
        Parameters:
        -----------
        iov      : input file (None implies ic = 1).
            - name of simple txt input file under the format ix iy iz ic eic (eic is optional)
            - pandas dataframe containing hashedId and  ic column (assuming it is from iov.xml.payload)
        name     : optional name
        useErrors: is False set eic = 1,
                   use when the txt input file includes error = 0
                    (default True)
        """
        self.iov = None
        if type(iov) == str:
            try:
                data = pd.read_csv(iov,delim_whitespace=True, header = None,names=['ix','iy','iz','ic','eic'],usecols=[0,1,2,3,4])
            except:
                data = pd.read_csv(iov,delim_whitespace=True, header = None,names=['ix','iy','iz','ic'],usecols=[0,1,2,3])
                data['eic'] = 1
            data.drop_duplicates(inplace=True)
            if not useErrors: data['eic'] = 1
            data['eic'].where( data['eic'] < 99 ,       0, inplace = True )
            data[ 'ic'].where( data['eic'] > 1e-8, np.nan, inplace = True )
            self.iov = pd.merge( geom, data, on=['ix','iy','iz'], how='left')
        # elif isinstance(iov,ecalic.iov.xml):
        #     geom['ecal'] = geom['iz'].where(geom['iz']==0,1)
        #     self.iov = pd.merge(geom, iov.payload, on = ['hashedId','ecal'] )
        #     self.iov['eic'] = 0
        #     del self.iov['ecal']
        #     del geom['ecal']
        elif isinstance(iov,pd.DataFrame):
            geom['ecal'] = geom['iz'].where(geom['iz']==0,1)
            self.iov = pd.merge(geom, iov, on = ['hashedId','ecal'] )
            self.iov['eic'] = 0
            del self.iov['ecal']
            del geom['ecal']
        else:
            ### for now default in no txt is geom
            self.iov = copy.deepcopy(geom)
        self.name = name

    def reset(self, ic_def = np.nan, eic_def = 0):
        """
        Reset ic
        Parameters:
        -----------
        ic_def : init ic val (default np.nan)
        eic_def: init eic val (default np.nan)
        """
        self.iov['ic' ] = ic_def
        self.iov['eic'] = eic_def

    def copy(self, name = ''):
        out = copy.deepcopy( self)
        out.name = name
        return out

    def __add__( self, ic ):
        out = self.copy()
        out['ic' ] = self['ic'] + ic['ic']
        out['eic'] = np.sqrt(  self['eic']* self['eic'] + ic['eic']*ic['eic'] )
        return out
    def __mul__( self, ic ):
        out = self.copy()
        out['ic' ] = self['ic'] * ic['ic']
        out['eic'] = np.sqrt(  self['eic']* self['eic']*  ic['ic']*ic['ic']  + \
                                 ic['eic']*ic['eic']   *self['ic']*self['ic'] )
        return out
    def __div__( self, ic ):
        out = self.copy()
        out['ic' ] = self['ic']/ic['ic']
        out['eic'] = np.sqrt(  (self['eic'] / self['ic'])**2 + \
                               (  ic['eic'] /   ic['ic'])**2 ) * out['ic']
        return out
    def __truediv__( self, ic ):
        return self.__div__(ic)
    def __sub__( self, ic ):
        out = self.copy()
        out.iov['ic' ] = self.iov['ic'] - ic.iov['ic']
        out.iov['eic'] = np.sqrt(  self.iov['eic']* self.iov['eic'] + ic.iov['eic']*ic.iov['eic'] )
        return out

    def __getitem__( self, key )       : return self.iov[key]
    def __setitem__( self, key, item ) : self.iov[key] = item

    def maskTTs( self, tts ):
        """
        Mask a list of trigger towers specified as ['Module','ccu']
        (expert only function)
        """
        for tt in tts:
            print( ' Masking TT %s - %d' % (tt[0],tt[1]))
            self[ 'ic'].where( (self['Module'] != tt[0]) | (self['ccu'] != tt[1]), np.nan, inplace = True )
            self['eic'].where( (self['Module'] != tt[0]) | (self['ccu'] != tt[1]), 0     , inplace = True )

    ####
    # def phiModuleFolder(self, xtal_mask = None, ic_name = 'ic' ):
    #     """
    #     folder vs phi: fold supermodule in EB with ieta granularity
    #     can be done using a defined xtal Masking
    #     """
    #     tmptab = self.iov
    #     if not xtal_mask is None : tmptab = self.iov.mask(xtal_mask)
    #     return tmptab.groupby([self['ieta'],self['phi_mod']])[ic_name].transform('mean')

    def average( self, xtal_mask = None, ic_name = 'ic' , trim = None, groupby = 'harness'):
        """
        average: average ICs per pre-defined regions
        groupby   : pre-defined region or used defined
                    user defined region are regular (list of) pandas series
                    predefined regions are:
                    - harness
                    - TT: trigger
                    - foldModule: folding EB modules in phi
        xtal_mask : mask to remove some crystals from average
        ic_name   : name of the parameter used for average (default: ic)
        trim      : remove <trim> percent of tail events on both left and right
        """
        tmptab = self.iov
        if not xtal_mask is None : tmptab = self.iov.mask(xtal_mask)
        gby = groupby
        if type(groupby) == str:
            if   groupby.lower() == 'harness'    : gby = [self['LME'],self['PNA']]
            elif groupby.lower() == 'tt'         : gby = [self['ccu'],self['Module']]
            elif groupby.lower() == 'foldmodule' : gby = [self['eta_module'],self['phi_module']]
        def trimming( x, cutfrac = None):
            if cutfrac is None: return x.mean()
            i = int(x.count()*cutfrac)
            return x.dropna().sort_values()[i:-i].mean()
        return tmptab.groupby(gby)[ic_name].transform(lambda x : trimming(x,trim))

    def etaRingNorm( self, xtal_mask = None, ic_name = 'ic' ):
        """
        eta ring normalization: normalize ic average per eta ring to one
        Parameters:
        -----------
        xtal_mask: a list of crystal to mask from the average (series of [boolean] )
        ic_name  : name of the variable used as ic (default 'ic')
        """
        if xtal_mask is None : self.iov = self.iov.join( self.iov.                groupby('ieta')[ic_name].mean(), on='ieta', rsuffix='_mean')
        else                 : self.iov = self.iov.join( self.iov.mask(xtal_mask).groupby('ieta')[ic_name].mean(), on='ieta', rsuffix='_mean')
        self[ic_name] =  self[ic_name]/ self[ic_name + '_mean']
        self['eic']   =  self['eic']  / self[ic_name + '_mean']
        del self.iov[ic_name + '_mean']

    def plot( self, var = 'ic', title = None, zRange_eb = None , zRange_ee = None, cmap = 'RdYlBu', ecalpart = None ):
        """
        Plot an ecal variable in 2D (EB, EE+, EE-)
        Parameters:
        -----------
        var: variable to plot (default ic)
        zRange_eb: specify the z axis range in for EB plot
        zRange_ee: specify the z axis range in for EE plot
        cmap     : color map (matplotlib name)
        ecalpart : part of ecal to plot
                - ecalpart = None : EB , EE+, EE- (default)
                - ecalpart = EB : EB only
                - ecalpart = EE : EE only
        """
        if zRange_eb is None: zRange_eb = [ self[self['iz']==0][var].min() ,  self[self['iz']==0][var].max()  ]
        if zRange_ee is None: zRange_ee = [ self[self['iz']!=0][var].min() ,  self[self['iz']!=0][var].max() ]
        if title is None : title = 'ICs '
        if var == 'eic' : title, cmap = ('IC errors ', 'afmhot' )

        if ecalpart is None:
            figure = plt.figure(num = title + self.name , figsize = [10,8] )
            gs = gridspec.GridSpec(2, 2)
            plot2D_eb( plt.subplot(gs[0,0:] ), self[self['iz']== 0][['ix','iy',var]] , 'EB  ' + title,  zRange_eb, cmap )
            plot2D_ee( plt.subplot(gs[1,0]  ), self[self['iz']==-1][['ix','iy',var]] , 'EE- ' + title,  zRange_ee, cmap )
            plot2D_ee( plt.subplot(gs[1,1]  ), self[self['iz']==+1][['ix','iy',var]] , 'EE+ ' + title,  zRange_ee, cmap )
        elif ecalpart.lower() == 'eb':
            figure = plt.figure(num = title + self.name + 'EB', figsize = [10,4.5] )
            plot2D_eb( plt.gca(), self[self['iz']== 0][['ix','iy',var]] , 'EB  ' + title,  zRange_eb, cmap )
        elif ecalpart.lower() == 'ee':
            figure = plt.figure(num = title + self.name + 'EE', figsize = [10,4] )
            plot2D_ee( plt.subplot(1,2,1 ), self[self['iz']==-1][['ix','iy',var]] , 'EE- ' + title,  zRange_ee, cmap )
            plot2D_ee( plt.subplot(1,2,2 ), self[self['iz']==+1][['ix','iy',var]] , 'EE+ ' + title,  zRange_ee, cmap )

    def plotEtaScale( self, var = 'ic', xtal_mask = None, legend = None, yZoom = [0.90,1.10], yLarge = [0.5, 1.5], axis = None, linestyle = 'None', marker = '.'):
        """
        Plot function along eta using eta-scale convention, i.e. 1/ic (rather than ic)
        The plot is devided in 3 sub ranges ( VF-EE-,  full ECAL, VF-EE+) where
        VF == very forward means |ieta| > 110.
        Return the central matplotlib axis (3 are defined)
        Parameters:
        -----------
        var       : variable to average vs ieta
        xtal_mask : list of crystals to mask
        yZoom     : y axis range for full ECAL plot (central region)
        yLarge    : y axis range for both VF regions
        legend    : name of the curve in the legend
        axis      : matplotlib axis (default None).
                    if None create a new figure.
                    if specified, must be a list of 3 axis [axisLeft, axisCentral, axisRight]
        linestyle : matplotlib line style (default None)
        marker    : matplotlib marker style (default '.')
        """
        x, y, yerr = self.profile1D(xtal_mask = xtal_mask )
        plt.close(plt.gcf())
        if axis is None :
            figure = plt.figure(num = 'Eta scale ' + self.name , figsize = [15,5] )
            gs = gridspec.GridSpec(1, 5)
            axis = [plt.subplot(gs[0,0]),plt.subplot(gs[0,1:4]),plt.subplot(gs[0,4])]

        for iecal, ax in enumerate(axis):
            ax.errorbar(x, 1/y  , yerr = 0 , marker=marker, linestyle=linestyle,
                        markersize = 6,label = legend)
            ax.set_xlabel('SC i$\eta$')
            ax.axhline(1, color = 'k', linewidth = 2)
            ax.set_ylim(yLarge[0],yLarge[1])
            ax.grid( which = 'both', linestyle='--', linewidth=0.9)
            for xmodule in [-85,-65,-45,-25,0,25,45,65,85]:
                    ax.axvline(xmodule, color = 'k', linestyle = '--')
        axis[0].set_xlim(-125,-110)
        axis[1].set_xlim(-125,+125)
        axis[2].set_xlim(+110,+125)

        axis[0].set_ylabel('r($\eta$)', horizontalalignment = 'right', y = 1)
        axis[1].set_ylim(yZoom[0],yZoom[1])
        if axis is None: axis[1].legend(frameon=True,framealpha=1,loc=8)
        return axis[1]

    def profile1D( self,  xvar = 'ieta', yvar = 'ic', xtal_mask = None,  title = None, legend = None, axis = None, color = None,  yRange = None ):
        """
        1D profile of the variable y as a function of variable x (both variables should be in the iov dataframe)
        Return numpy arrays of the plotted points:  x,y,yerr
        Parameters:
        -----------
        xvar: variable to define group for average
        yvar: variable to average
        xtal_mask : list of crystals to mask
        yRange    : y axis range
        axis      : matplotlib axis (default None).
                    if None create a new figure.
                    if specified, must be a list of 3 axis [axisLeft, axisCentral, axisRight]
        legend    : name of the curve in the legend
        title     : title used for figure (for saving purposes) - should not be used in principle
        """
        iov = self.iov if xtal_mask is None else  self.iov.mask(xtal_mask)
        proj = iov.groupby(xvar).agg({ yvar:['mean','std'] })

        if title is None: title = 'Mean%s vs %s ' % (yvar,xvar)
        if axis  is None : axis  = plt.figure( num = title + ' ' + self.name, figsize = [7,4] ).gca()

        axis.errorbar( proj.index, proj[yvar,'mean'], yerr = proj[yvar,'std'],
                        color = color, linestyle = 'None', markersize = 6, marker= '.', label = legend  )

        axis.set_xlabel(      xvar  , horizontalalignment = 'right', x = 1)
        axis.set_ylabel( '<%s>'%yvar, horizontalalignment = 'right', y = 1)

        if not yRange is None: axis.set_ylim(yRange[0],yRange[1])

        return proj.index, proj[yvar,'mean'], proj[yvar,'std']

    def profile1D_groupY( self,  xvar = 'ieta', yvar = 'ic', xtal_mask = None, title = None, legend = None, axis = None, color = 'k',  yRange = None ):
        """
        1D profile of the variable y as a function of variable x (both variables should be in the iov dataframe)
        It differs from profile1D by grouping variable according to Y, adding all events with the same Y to a single
        x-bin (length adjusted to includes all events with same Y).
        Return numpy arrays of the plotted points:  x,y,xerr,yerr
        Parameters:
        -----------
        xvar: variable to define group for average
        yvar: variable to average
        xtal_mask : list of crystals to mask
        yRange    : y axis range
        axis      : matplotlib axis (default None).
                    if None create a new figure.
                    if specified, must be a list of 3 axis [axisLeft, axisCentral, axisRight]
                    legend    : name of the curve in the legend
                    title     : title used for figure (for saving purposes) - should not be used in principle
        """
        iov = self.iov if xtal_mask is None else  self.iov.mask(xtal_mask)
        proj = iov.groupby(yvar).agg({ xvar:['mean','min','max'] })

        if title is None: title = '<%s> vs %s ' % (yvar,xvar)
        if axis is None : axis  = plt.figure( num = title + ' ' + self.name, figsize = [7,4] ).gca()

        x,y,xerr = (proj[xvar,'max']+proj[xvar,'min'])/2, proj.index, (proj[xvar,'max']-proj[xvar,'min'])/2
        yerr = np.zeros( y.size )
        axis.errorbar( x, y , xerr = xerr,
                        color = color, linestyle = 'None', markersize = 6, marker= '.', label = legend  )

        axis.set_xlabel(      xvar  , horizontalalignment = 'right', x = 1)
        axis.set_ylabel( '<%s>'%yvar, horizontalalignment = 'right', y = 1)

        if not yRange is None: axis.set_ylim(yRange[0],yRange[1])

        return x.values,y.values,xerr.values,yerr

    def dump(self, output ):
        """
        dump the ICs into a regular txt file (csv style) under the format ix iy iz ic eic
        Parameters:
        -----------
        output: name of the output txt file
        """
        print( 'Dumping iov file to: ' + output )
        ## replace nan ic with 1 and nan eic with 0 (the latter should not exist)
        self.iov.replace({'ic' : { np.nan : 1 },'eic' : { np.nan : 0 }  }).to_csv( output,
                                columns = ['ix','iy','iz','ic','eic'],
                                header = False, index = False, sep = ' ', float_format = '%+6.5f' )
