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 pandas as pd
import pickle
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 for 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__':


    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 = 'IT')
    parser.add_argument("-l", "--layer"      , dest="layer",       help="Choose the layer to use",choices=['TTaU','T3X2'], default = 'T3X2')
    parser.add_argument("-n", "--nstrips"      , dest="nstrips",      help="Choose the number of strips to use",choices=['3','5','7'], default = ['7'])
    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 = False
    plot2 = True

    location = os.path.expandvars('$DISK/data/ST/Aging/')
    fileData = TFile(location+'CCEScan.root')
    
    #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 is '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(plot1 == True):

        print 'Plotting...'

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

                plt.errorbar(df1[(df1.Vbias == v)]['Fill'], df1[(df1.Vbias == v)]['x_max'], df1[(df1.Vbias == v)]['x_err'], fmt='-^',label=str(v)+' V')

    
            plt.xlabel("Fill number")
            plt.ylabel('Time of the pulse shape maximum (ns)')
            plt.title("Time evolution of  max. position for sector {}".format(str(isect)))
            
            plt.legend(loc='upper left')
            #plt.show()
            plt.savefig('./Max_Evol_Plots/Max_Evol_{}_{}.pdf'.format(layer,str(isect)))
            plt.clf()

    #####################
    if(plot2):

        df1 = df[(df.Fill == 5162)]
        df2 = df[(df.Fill == 5448)]
        



        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" )
    
        plt.legend(loc='upper left')
        plt.savefig('./Max_Evol_Plots/TimeDifferences_{}_HV.pdf'.format(layer))
        plt.clf()