"""
Created on Tue Dic  2 12:35:00 2020

@author: vgarcia
"""

import ROOT
import datetime
from tqdm import tqdm
import numpy as np
import pandas as pd

inputDir = './'

e = datetime.datetime.now()
print('Program started at: %s:%s:%s' % (e.hour, e.minute, e.second))
del e

###############################################################################
#                             DATA PROCESSING                                 #
###############################################################################

# Open histos file
inputFile0 = ROOT.TFile.Open(inputDir+"Calib_201123_131528_histos.root")
print('File readed')

# Create empty histograms with accurate titles
DSSD1_V = []
DSSD1_H = []
DSSD2_V = []
DSSD2_H = []
DSSD3_V = []
DSSD3_H = []
FT_1_2 = []
FT_3_4 = []
CEPA_short = []
CEPA_long = []

print('\n Creating empty histograms: ')
for i in tqdm(range(32)):
    if i<16:
        DSSD1_V.append(ROOT.TH1F("DSSD1_P"+str(i),"DSSD1_P%i  Channel  Counts" % (i), 4096,0,4096))
        DSSD2_V.append(ROOT.TH1F("DSSD2_P"+str(i),"DSSD2_P%i  Channel  Counts" % (i), 4096,0,4096))
        DSSD3_V.append(ROOT.TH1F("DSSD3_P"+str(i),"DSSD3_P%i  Channel  Counts" % (i), 4096,0,4096))
    else:
        DSSD1_H.append(ROOT.TH1F("DSSD1_N"+str(i-16),"DSSD1_N%i  Channel  Counts" % (i-16), 4096,0,4096))
        DSSD2_H.append(ROOT.TH1F("DSSD2_N"+str(i-16),"DSSD2_N%i  Channel  Counts" % (i-16), 4096,0,4096))
        DSSD3_H.append(ROOT.TH1F("DSSD3_N"+str(i-16),"DSSD3_N%i  Channel  Counts" % (i-16), 4096,0,4096))
        
    FT_1_2.append(ROOT.TH1F("FT_1_2_"+str(i),"FT_1_2_%i  Channel  Counts" % (i), 4096,0,4096))
    FT_3_4.append(ROOT.TH1F("FT_3_4_"+str(i),"FT_3_4_%i  Channel  Counts" % (i), 4096,0,4096))
    
    if 0<=i<=3:
        CEPA_short.append(ROOT.TH1F("CEPA_integration_short_"+str(i), "CEPA_integration_short_%i  Channel  Counts" % (i), 4096,0,4096))
        CEPA_long.append(ROOT.TH1F("CEPA_integration_long_"+str(i), "CEPA_integration_long_%i  Channel  Counts" % (i), 4096,0,4096))
print('OK Empty histograms created')

# Get the root client generated histograms from listfile to properly named ones
print('\n Getting histograms from histos.root')
for j in tqdm(range(32)):
    DSSD1_histo = inputFile0.Get('event0/caen_v785__untested__amplitude[%i]' % (j))
    DSSD2_histo = inputFile0.Get('event0/caen_v785__untested__1_amplitude[%i]' % (j))
    DSSD3_histo = inputFile0.Get('event0/caen_v785__untested__2_amplitude[%i]' % (j))
    FT_1_2_histo = inputFile0.Get('event0/caen_v785__untested__3_amplitude[%i]' % (j))
    FT_3_4_histo = inputFile0.Get('event0/caen_v785__untested__4_amplitude[%i]' % (j))
    
    if 1 <= j <=4:
        CEPA_histo_short = inputFile0.Get('event0/mdpp16_qdc_integration_short[%i]' % (j))
        CEPA_histo_long = inputFile0.Get('event0/mdpp16_qdc_integration_long[%i]' % (j))
    
    for k in range(4069):
        if j<16:
            DSSD1_V[j].SetBinContent(k, DSSD1_histo.GetBinContent(k))
            DSSD2_V[j].SetBinContent(k, DSSD2_histo.GetBinContent(k))
            DSSD3_V[j].SetBinContent(k, DSSD3_histo.GetBinContent(k))
        else:
            DSSD1_H[j-16].SetBinContent(k, DSSD1_histo.GetBinContent(k))
            DSSD2_H[j-16].SetBinContent(k, DSSD2_histo.GetBinContent(k))
            DSSD3_H[j-16].SetBinContent(k, DSSD3_histo.GetBinContent(k))
        
        FT_1_2[j].SetBinContent(k, FT_1_2_histo.GetBinContent(k))
        FT_3_4[j].SetBinContent(k, FT_3_4_histo.GetBinContent(k))
        
        if 1<=j<=4:
            CEPA_short[j-1].SetBinContent(k, CEPA_histo_short.GetBinContent(k))
            CEPA_long[j-1].SetBinContent(k, CEPA_histo_long.GetBinContent(k))

del DSSD1_histo, DSSD2_histo, DSSD3_histo, FT_1_2_histo, FT_3_4_histo, CEPA_histo_short, CEPA_histo_long, i, j, k

print('OK  New histos get from histos.root \n')


###############################################################################
#                                FIND PEAKS                                   #
###############################################################################
# We use TSpectrum class and Search to find the peaks in our histograms. As Gd 
# is more intense than 3alpha we will find peaks in 2 steps, 1 for Gd peaks and
# other for 3alpha peaks.

def find_peaks(histogram, local_range=[0,4096], sigma=2, opt='', threshold=0.05, npeaks=10):
#    histogram = histogram.Clone() # Close histos to prevent draw on original ones
    histogram.GetXaxis().SetRange(local_range[0], local_range[1])
    spectrum = ROOT.TSpectrum(2*npeaks)
    spectrum.Search(histogram, sigma, opt, threshold)
    spectrum.Background(histogram, 20, "same")
    buffer_X, buffer_Y = spectrum.GetPositionX(), spectrum.GetPositionY()
    position = []
    for i in range (spectrum.GetNPeaks()):
        position.append([buffer_X[i], buffer_Y[i]])
    position.sort()
    return position

# First we find Gd peaks and store values in a python list with 1 column for 
# each DSSD separated by Vertical or Horizontal strips. DSSD1_V will be 
# Gd_peak[0], DSSD1_H -> Gd_peak[1], DSSD1_V -> Gd_peak[2] and so on.
# Once we found Gd peak we start to find Trialphas peak 100 channels after 
# Gd peak. Finally we will have 2 lists with all Gd peaks and Trialpha peaks 
# for each detector

Gd_peak = [[], [], [], [], [], [], [], []]    
Trialpha_peak = [[], [], [], [], [], [], [], []] 
print('Finding peaks on each histogram')
for i in tqdm(range(32)):
    if 0<=i<16:
        Gd_peak_position = find_peaks(DSSD1_V[i], local_range=[400, 2000], sigma=5, threshold=0.50)
        Gd_peak[0].append(Gd_peak_position)
        if Gd_peak_position != []:
            Trialpha_peak_position = find_peaks(DSSD1_V[i], local_range=[int(Gd_peak_position[0][0]+100), 4095], sigma=16, threshold=0.20)
            Trialpha_peak[0].append(Trialpha_peak_position)
        else:
            Trialpha_peak[0].append([])
        
        Gd_peak_position = find_peaks(DSSD1_H[i], local_range=[400, 2000], sigma=5, threshold=0.50)
        Gd_peak[1].append(Gd_peak_position)
        if Gd_peak_position != []:
            Trialpha_peak_position = find_peaks(DSSD1_H[i], local_range=[int(Gd_peak_position[0][0]+100), 4095], sigma=16, threshold=0.20)
            if len(Trialpha_peak_position)>3:
                del Trialpha_peak_position[0] # remove 1st peak if there are more than 3 as it is a reflection because shaping time fail 
            Trialpha_peak[1].append(Trialpha_peak_position)
        else:
            Trialpha_peak[1].append([])
        
        Gd_peak_position = find_peaks(DSSD2_V[i], local_range=[400, 2000], sigma=5, threshold=0.50)
        Gd_peak[2].append(Gd_peak_position)
        if Gd_peak_position != []:
            Trialpha_peak_position = find_peaks(DSSD2_V[i], local_range=[int(Gd_peak_position[0][0]+100), 4095], sigma=16, threshold=0.15)
            Trialpha_peak[2].append(Trialpha_peak_position)
        else:
            Trialpha_peak[2].append([])
        
        Gd_peak_position = find_peaks(DSSD2_H[i], local_range=[400, 2000], sigma=7, threshold=0.50)
        Gd_peak[3].append(Gd_peak_position)
        if Gd_peak_position != []:
            Trialpha_peak_position = find_peaks(DSSD2_H[i], local_range=[int(Gd_peak_position[0][0]+100), 4095], sigma=25, threshold=0.15)
            if len(Trialpha_peak_position)>3:
                del Trialpha_peak_position[0] # remove 1st peak if there are more than 3 as it is a reflection because shaping time fail
            Trialpha_peak[3].append(Trialpha_peak_position)
        else:
            Trialpha_peak[3].append([])
        
        Gd_peak_position = find_peaks(DSSD3_V[i], local_range=[400, 2000], sigma=5, threshold=0.50)
        Gd_peak[4].append(Gd_peak_position)
        if Gd_peak_position != []:
            Trialpha_peak_position = find_peaks(DSSD3_V[i], local_range=[int(Gd_peak_position[0][0]+100), 4095], sigma=16, threshold=0.20)
            Trialpha_peak[4].append(Trialpha_peak_position)
        else:
            Trialpha_peak[4].append([])
        
        Gd_peak_position = find_peaks(DSSD3_H[i], local_range=[400, 2000], sigma=7, threshold=0.50)
        Gd_peak[5].append(Gd_peak_position)
        if Gd_peak_position != []:
            Trialpha_peak_position = find_peaks(DSSD3_H[i], local_range=[int(Gd_peak_position[0][0]+100), 4095], sigma=16, threshold=0.20)
            Trialpha_peak[5].append(Trialpha_peak_position)
        else:
            Trialpha_peak[5].append([])
        
    Gd_peak_position = find_peaks(FT_1_2[i], local_range=[400, 2000], sigma=5, threshold=0.50)
    Gd_peak[6].append(Gd_peak_position)
    if Gd_peak_position != []:
        Trialpha_peak_position = find_peaks(FT_1_2[i], local_range=[int(Gd_peak_position[0][0]+100), 4095], sigma=10, threshold=0.20)
        Trialpha_peak[6].append(Trialpha_peak_position)
        if len(Trialpha_peak_position)>3:
                del Trialpha_peak_position[1], Trialpha_peak_position[0] # remove 2st peak because we have a lot of noise in channel 18
    else:
        Trialpha_peak[6].append([])
    
    Gd_peak_position = find_peaks(FT_3_4[i], local_range=[400, 2000], sigma=5, threshold=0.50)
    Gd_peak[7].append(Gd_peak_position)
    if Gd_peak_position != []:
        Trialpha_peak_position = find_peaks(FT_3_4[i], local_range=[int(Gd_peak_position[0][0]+250), 4095], sigma=10, threshold=0.30)
        Trialpha_peak[7].append(Trialpha_peak_position)
        if i==12:
            del Trialpha_peak_position[1] # Remove 1 peak manually
        if i==13:
            del Trialpha_peak_position[4], Trialpha_peak_position[3] # Remove 2 peaks manually
    else:
        Trialpha_peak[7].append([])
        
del i, Gd_peak_position, Trialpha_peak_position

print('\n OK All peaks have been found')

###############################################################################
#                            GAUSSIAN FITTING                                 #
###############################################################################
# Now we have the position of all peaks so we will use it to make a gaussian fit for each peak

# We define our gaussian function with linear background
def gauss_lbg(x, par): #we assume that bin width is 1 by default
    # par[0]	background constant
    # par[1]	background slope
    # par[2]	gauss width
    # par[3]	gauss constant
    # par[4]	gauss mean
    sqrt2pi , sqrt2 = np.sqrt(2*np.pi) , np.sqrt(2)
    if par[2] == 0: par[2]=1
    arg = (x[0] - par[4])/(sqrt2*par[2])
    return par[0] + x[0]*par[1] + 1/(sqrt2pi*par[2]) * par[3] * np.exp(-arg*arg)
    

def FitSinglePeak(histogram, peak_center_xy, peak_width=100):
    peak_center_position = int(peak_center_xy)
    peak_right = []
    peak_left = []
    # from the peak center we go along 150 bins at right and left side of the center of the peak, to find the minimun
    # and fit the oeak
    for i in range(int(peak_width/2)):
        peak_right.append(histogram.GetBinContent(peak_center_position+i))
        peak_left.append(histogram.GetBinContent(peak_center_position-i))
    
    # Find  the minimum in that range to set the integration limits of our peak
    peak_right_min_bin = peak_center_position + peak_right.index(min(peak_right))
    peak_left_min_bin = peak_center_position - peak_left.index(min(peak_left))
    
    # We can use peak_right/left_min_bin to set the correct range to do our fit
    histogram.GetXaxis().SetRangeUser(peak_left_min_bin, peak_right_min_bin)
    histogram.Draw()
    
    # Create a root TF1 function with out gaussian function defined previously, las number is the number of parameters
    gauss_linbg_fit = ROOT.TF1('gauss_linear_background', gauss_lbg, peak_left_min_bin, peak_right_min_bin, 5)
    
    # Set parameter names, values and limits
    gauss_linbg_fit.SetParNames("BgConstant","BgSlope","Sigma", "Area","Position")
    gauss_linbg_fit.SetParameters(0,
                                  0,
                                  0.3*(peak_right_min_bin-peak_left_min_bin), 
                                  histogram.Integral(histogram.FindBin(peak_left_min_bin), histogram.FindBin(peak_right_min_bin)), 
                                  0.5*(peak_left_min_bin+peak_right_min_bin))
    gauss_linbg_fit.SetParLimits(3,0,1e7)
    gauss_linbg_fit.SetParLimits(2,0,100)
    
    # Fit our histogram to gaussian function with linear background
    histogram.Fit(gauss_linbg_fit, 'QRMN')
    
    # Draw gaussian function and linear background
    
    straight_line = ROOT.TF1("straight_line","pol1", peak_left_min_bin, peak_right_min_bin )
    straight_line.SetParameter(0, gauss_linbg_fit.GetParameter(0))
    straight_line.SetParameter(1, gauss_linbg_fit.GetParameter(1)) 
    straight_line.SetLineColor(ROOT.kGreen)
    
    gauss_fit = ROOT.TF1("Gauss1", gauss_lbg, peak_left_min_bin,peak_right_min_bin, 5)
    gauss_fit.SetParameter(0,  gauss_linbg_fit.GetParameter(0))
    gauss_fit.SetParameter(1,  gauss_linbg_fit.GetParameter(1))
    gauss_fit.SetParameter(2,  gauss_linbg_fit.GetParameter(2))
    gauss_fit.SetParameter(3,  gauss_linbg_fit.GetParameter(3)) 
    gauss_fit.SetParameter(4,  gauss_linbg_fit.GetParameter(4))
    gauss_fit.SetLineColor(ROOT.kRed)
    
    straight_line.Draw('same')
    gauss_fit.Draw('same')
    
    print('\n')
    print('\t Chi Square: %f ' % (gauss_linbg_fit.GetChisquare()))
    print('\t FWHM: %f +- %f ' % (2*gauss_linbg_fit.GetParameter(2)*np.sqrt(2*np.log(2)), \
                                  2*np.sqrt(2*np.log(2))*gauss_linbg_fit.GetParError(2)))
    print('\t Position: %f +- %f ' % (gauss_linbg_fit.GetParameter(4), gauss_linbg_fit.GetParError(4)))
    print('\t Area: %f +- %f ' % (gauss_linbg_fit.GetParameter(3), gauss_linbg_fit.GetParError(3)))
    print('\t Sigma: %f +- %f ' % (gauss_linbg_fit.GetParameter(2), gauss_linbg_fit.GetParError(2)))
    print('\n')
    return gauss_linbg_fit.GetParameter(4), straight_line, gauss_fit