import ROOT
import numpy as np
from ROOT import TFile, TLorentzVector, TVector3, TRotation, TLorentzRotation, TMath, TH1D, TCanvas, TH2D, TObject, TF1
import matplotlib.pyplot as plt
import seaborn as sns
import pandas as pd

from checkPulseData import Table
import os
import math as m
import argparse

#Given the fitted parameters calculates the maximum point
def GetMaximumPoint(t0, t0_err, tau, tau_err):
    
    #The position of the maximum can be obtained from the functional form used for the fit
    t = t0 + (3 - m.sqrt(3))*tau
    t_err = m.sqrt(t0_err*t0_err + (3 - m.sqrt(3))*tau_err*(3 - m.sqrt(3))*tau_err)
    return t, t_err

#Obtains the fitted parameters from pulse shapes
def GetFitParameters(fileData, detector, layer, isect, ifill, calibstep, istrip):

    #print 'Using {}/{}/{}/{}/v_pulse{}_val{}'.format(detector,layer, isect, ifill, calibstep, istrip)
    v = np.array(fileData.Get('{}/{}/{}/{}/v_pulse{}_val{}'.format(detector,layer, isect, ifill, calibstep, istrip)))
    
    return v



if __name__ == '__main__':

    '''
    Autor: Michele Atzeni
    Email: michele.atzeni@cern.ch

    Working in date 14.06.2017
    '''

    parser = argparse.ArgumentParser(description = 'Configuration of the parameters for the SplitAndMerge')

    parser.add_argument("-d", "--detector" , dest="detector"  ,  help="Choose the detector to use ",choices=['TT','IT'], default = 'TT')
    parser.add_argument("-l", "--layer"      , dest="layer",       help="Choose the layer to use",choices=['TTaU','T3X2'], default = 'TTaU')
    parser.add_argument("-n", "--nstrips"      , dest="nstrips",      help="Choose the number of strips to use",choices=['3','5','7'], default = ['5'])
    parser.add_argument("--Vred", action="store_true", help="exclude the 3 lowest Vbias")
    parser.add_argument("--VERB", action="store_true", help="VERBOsE")
    args = parser.parse_args()
    
    # Parameters and configuration  
    
    detector= args.detector
    layer  = args.layer
    nstrips  = args.nstrips
    Vred  = args.Vred
    VERB  = args.VERB
    
    
    #Plotting options
    plot1 = True
    plot2 = True

    location = os.path.expandvars('$DISK/data/ST/Aging/')
    fileData = TFile(location+'CCEScan.root')

    Outputfolder = os.path.expandvars('$DISK/')+'data/ST/graphics/{}/{}/Max_Evol_Plots/'.format(detector, layer)

    if not os.path.exists(Outputfolder) :
        print "Creating the folder {}".format(Outputfolder)
        raw_input("Ctrl-C to abort, return to continue")
        os.makedirs(Outputfolder)


    #detector = 'TT'
    #layer = 'TTaU'
    #nstrips = [5]#[3, 5, 7]                                                                                                    
    
    #Voltage map
    voltMapTT = [400., 350., 300., 250., 225., 200., 175., 150., 125., 100., 60.]    
    voltMapIT = [300., 200., 170., 140., 120., 105., 90., 75., 60., 40., 20.]
    voltMapTTred = [400., 350., 300., 250., 225., 200., 175., 150.]    
    voltMapITred = [300., 200., 170., 140., 120., 105., 90., 75.]


    if(detector == 'TT'):
        if (Vred):
            voltMap = voltMapTTred
        else:
            voltMap = voltMapTT
    else:
        if(Vred):
            voltMap = voltMapITred
        else:
            voltMap = voltMapIT


    #Load Fills
    macros = os.path.expandvars('$CCEHOME/macros/CCEScan/')
    with open(macros + 'Fills.dat', 'rb') as f:
        fills = f.read().splitlines()
    fills.remove('2797')
    fills.remove('3108')
        
    #Load Sectors               
    with open(macros + '{DET}sectors.dat'.format(DET=detector), 'rb') as f:
        sectors = f.read().splitlines()
        
    print "The sector list is: \n", sectors
    

    #Define table columns
    table = Table('Detector', 'Sector', 'Fill', 'N strips', 'Vbias', 'x_max','x_err')



    for isect in sectors:
        print '====== SECTOR {}================'.format(isect)
    

        for ifill in fills:
            if (VERB):
                print '====== FILL {}================'.format(ifill)


            for istrip in nstrips:
                if(VERB):
                    print '====== NSTRIPS {}================'.format(istrip)

                for calibstep in range(0,len(voltMap)):
                    if(VERB):
                        print '====== CALIBSTEP {}================'.format(calibstep)

                    v_fit = GetFitParameters(fileData, detector, layer, isect, ifill, calibstep, istrip)
                    if(VERB):
                        print 'v_fit: ', v_fit

                    try:
                        x_max, x_err = GetMaximumPoint(v_fit[4], v_fit[5], v_fit[2], v_fit[3])
                        if (VERB):
                            print 'x_max: ', x_max
                        table.append(detector, int(isect), int(ifill), istrip, voltMap[calibstep], x_max, x_err)
                    except:
                        print '====================================================================================================================='
                        print 'WARNING: Empty vector found for {}/{}/{}/{}/v_pulse{}_val{}'.format(detector,layer, isect, ifill, calibstep, istrip)
                        print '====================================================================================================================='

                    
                    #print (fileData, detector, layer, isect, ifill, calibstep, istrip)

    df = pd.DataFrame.from_dict(table)
    if not os.path.isfile('Shifts_{DET}.pkl'.format(DET=detector)):
        import pickle
        pickle.dump(df, open('Shifts_{DET}.pkl'.format(DET=detector), 'wb'))

    ####################
    if(plot1 == True):

        print 'Plotting...'

        for isect in sectors:
            df1 = df[df.Sector == int(isect)]
            
            for v in voltMap:

                df1p = df1[(df1.Vbias == v)]
                
                plt.errorbar(x=df1p['Fill'].values , y=df1p['x_max'].values, yerr=df1p['x_err'].values, fmt='-^',label='{} V'.format(str(v)))
                #plt.errorbar(x=[1,2,3,4,5,6], y=[1,2,3,4,5,6], yerr=[0.1,0.1,0.1,0.1,0.1,0.1], fmt='^',label='{} V'.format(str(v)))
                #plt.plot(x=df1p['Fill'], y=df1p['x_max'],marker='.',label='{} V'.format(str(v)))#yerr=df1p['x_err'],
    
            plt.xlabel("Fill number")
            plt.ylabel('Time of the pulse shape maximum (ns)')
            plt.title("Time evolution of  max. position for sector {} using nstrips={}".format(str(isect),istrip))
            
            plt.legend(loc='upper left')
            #plt.show()       
            plt.savefig(Outputfolder+'Max_Evol_{}_{}_{}.pdf'.format(layer,str(isect),istrip))
            plt.clf()

    #####################
    if(plot2):
        
        first_fill = 5162#int(raw_input(">For the last plot choose the first fill (5162): "))
        second_fill = 5448#int(raw_input(">For the last plot choose the second fill (5448): "))
        df1 = df[(df.Fill == first_fill)]
        df2 = df[(df.Fill == second_fill)]
        



        for vbias in voltMap:
            
            if (vbias < 0.):
                continue
            
            df1v = df1[(df1.Vbias == vbias)][['Sector','x_max','x_err']]
            df2v = df2[(df2.Vbias == vbias)][['Sector','x_max','x_err']]
            df1v = df1v.sort_values(['Sector'], ascending=True).reset_index(drop=True)
            df2v = df2v.sort_values(['Sector'], ascending=True).reset_index(drop=True)

            plt.errorbar(df1v['Sector'], df2v['x_max'] - df1v['x_max'], yerr= df2v['x_err'] + df1v['x_err'], fmt='^', label= '{} V'.format(vbias))

        plt.xlabel("Sector")
        plt.ylabel('Time difference (ns)')
        plt.title("Time difference of max. positions fill 5448-5162 using nstrips={}".format(istrip) )
    
        plt.legend(loc='upper left')
        plt.savefig(Outputfolder+'TimeDifferences_fill{}_{}_HV_ns{}.pdf'.format(str(first_fill)+str(second_fill),layer,istrip))
        plt.clf()