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


#Given the fitted parameters calculates the maximum point
def GetMaximumPoint(t0, t0_err, tau, tau_err):

    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):

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



if __name__ == '__main__':

    #Plotting options
    plot1 = True
    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.]    
    
    #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()

    #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:
            #print '====== FILL {}================'.format(ifill)


            for istrip in nstrips:

                for calibstep in range(0,len(voltMapTT)):
                    #print '====== CALIBSTEP {}================'.format(calibstep)

                    v_fit = GetFitParameters(fileData, detector, layer, isect, ifill, calibstep, istrip)
                    #print 'v_fit: ', v_fit
                    x_max, x_err = GetMaximumPoint(v_fit[4], v_fit[5], v_fit[2], v_fit[3])
                    #print 'x_max: ', x_max
                    table.append(detector, int(isect), int(ifill), istrip, voltMapTT[calibstep], x_max, x_err)
                
                                                        

    df = pd.DataFrame.from_dict(table)
    
    ####################
    if(plot1 == True):

        print 'Plotting...'

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

                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(str(isect)))
            plt.clf()

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

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



        for vbias in voltMapTT:
            
            if (vbias < 70.):
                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.pdf'.format(str(isect)))
        plt.clf()