import ROOT
import pickle
import numpy as np

import os

import getPredictions

# numbers of excluded fills due to special configuration (e.g. CCE scans)
bad_boys = [4643,4639,4638,4518,3960,3820,3478,3456,3232,3054,3109,3121,3108,3072,2797,2772,2576,2574,2564,2481,2480,2476,2472,2414,2252,2235,2083,1944,1943,1936,1932,1928,1616,1595,1554,1552,1547,1464,1463,1459,1455,1449,1436]


# obtaining for a given sector the integrated luminosity, the time and the measured current
# as an array
def vectors(sec):
  time,lumi,curr = [],[],[]
  for d in sec:
    if d["current"]>0.00:
      # excluding jumps in the measured currents by more than a factor 10
      if len(curr)>0 and d["current"]>10.0*curr[len(curr)-1]:
        continue
      if d["fill"] in bad_boys:
        continue
      lumi.append(d["intLumi"])
      time.append(d["endTime"])
      curr.append(d["current"])

  return [lumi,time,curr]


# function to correct for the temperature -> normalisation from a temperature t1 to T = 8 degrees C (t2)
# Eg: silicon energy gap in eV
def tempCorr(t1,t2=8.0,Eg=1.21):
  kelvin = 273.15
  kb = 8.6173324e-5
  t1 += kelvin
  t2 += kelvin
  return (t1/t2)**2*np.exp(-Eg/(2*kb)*(1/t2-1/t1))


# function to apply the temperature correction for TT (i.e. two relevant temperature senssors CT5 and CT6
def corrCurrTT(time,curr,temp):
  i1=0
  i2=0
  for i in range(0,len(curr)):
    while time[i]>temp["CT5"][0][i1] and i1<len(temp["CT5"][0])-1:
      i1+=1
    while time[i]>temp["CT6"][0][i2] and i2<len(temp["CT6"][0])-1:
      i2+=1
    t1 = 0.5*(temp["CT5"][1][i1]+temp["CT6"][1][i2])
    curr[i] *= tempCorr(t1)


# function to apply the temperature correction for IT
# (taking the temperature sensor in the corresponding box)
def corrCurrIT(time,curr,temp):
  j=0
  for i in range(0,len(curr)):
    while time[i]>temp[0][j] and j<len(temp[0])-1:
      j+=1
    t1 = temp[1][j]
    curr[i] *= tempCorr(t1)


# filter high-voltage channels according to a given layer name pattern
def filterLayer(dic,expr):
  return dict((key,value) for key, value in dic.iteritems() if key.find(expr)>=0)


# filter high-voltage channels according to a given section/box name pattern
def filterSect(dic,expr):
  return dict((key,value) for key, value in dic.iteritems() if key.find(expr)==len(key)-len(expr))


# function to get the silicon volume connected to a particular value
# returns 100000000000000.0 if the pattern does not match
def getVolume(key): # in cm^3
  if key.find("TT")>=0:
    oneSensor = 0.0500*9.64*9.44
    if key.find("US")>=0 or key.find("X1S")>=0:
      if np.size(np.where(np.array([key.find(sec) for sec in ("S10","S4","S3","S9","S15","S16")])>=0))>0:
        return 1.0*oneSensor
      elif np.size(np.where(np.array([key.find(sec) for sec in ("S11","S5","S8","S14","S17")])>=0))>0:
        return 2.0*oneSensor
      elif np.size(np.where(np.array([key.find(sec) for sec in ("S12","S6","S7","S13","S18")])>=0))>0:
        return 4.0*oneSensor
      elif key.find("S2")==len(key)-2:
        return 2.0*oneSensor
      elif key.find("S1")==len(key)-2:
        return 4.0*oneSensor
    elif key.find("VS")>=0 or key.find("X2S")>=0:
      if np.size(np.where(np.array([key.find(sec) for sec in ("S14","S8","S7","S13","S19","S20")])>=0))>0:
        return 1.0*oneSensor
      elif np.size(np.where(np.array([key.find(sec) for sec in ("S15","S9","S6","S12","S18","S21")])>=0))>0:
        return 2.0*oneSensor
      elif np.size(np.where(np.array([key.find(sec) for sec in ("S25","S24","S3")])>=0))>0:
        return 3.0*oneSensor
      elif np.size(np.where(np.array([key.find(sec) for sec in ("S26","S22","S16","S10","S4","S5","S11","S17","S23")])>=0))>0:
        return 4.0*oneSensor
      elif key.find("S2")==len(key)-2:
        return 3.0*oneSensor
      elif key.find("S1")==len(key)-2:
        return 4.0*oneSensor
    else:
      if np.size(np.where(np.array([key.find(sec) for sec in ("P8","P4","P5","P1")])>=0))>0:
        return 12.0*oneSensor
      else:
        return 9.0*oneSensor
  else:
    if key.find("BOXA")>=0 or key.find("BOXC")>=0:
      return 0.041*7.6*11*2.0*4
    else:
      return 0.032*7.6*11*4
  return 100000000000000.0


# returns one multigraph for time and one multigraph for time for TT
# filSector and filLayer are iterables representing different patterns to be selected
# temp is a dictionnary of lists containing two elements, a time list and the corresponding temperature list
def drawTT_sensor(data,filSector,filLayer,name,color,temp):
  mg_time = ROOT.TMultiGraph("mg_%s_time"%name,"Multigraph time");
  mg_lumi = ROOT.TMultiGraph("mg_%s_lumi"%name,"Multigraph lumi");
  for f1 in filSector:
    filter_dict = filterSect(data,f1)
    for f2 in filLayer:
      dfilter_dict = filterLayer(filter_dict,f2)
      for sec in dfilter_dict:
        print sec,": ",getVolume(sec)
        [lumi,time,curr] = vectors(filter_dict[sec])
        corrCurrTT(time,curr,temp)
        v_curr = np.array(curr)/getVolume(sec)
        g_time = ROOT.TGraph(len(time),np.array(time)*1.0,v_curr)
        g_lumi = ROOT.TGraph(len(time),np.array(lumi)*1.0,v_curr)
        g_time.SetLineColor(color);
        g_lumi.SetLineColor(color);
        mg_time.Add(g_time,"l")
        mg_lumi.Add(g_lumi,"l")
  return [mg_time,mg_lumi]


# returns one multigraph for time and one multigraph for time for IT
# filSector and filLayer are iterables representing different patterns to be selected
# temp is a dictionnary of lists containing two elements, a time list and the corresponding temperature list
def drawIT_sensor(data,filBox,name,color,temp):
  mg_time = ROOT.TMultiGraph("mg_%s_time"%name,"Multigraph time");
  mg_lumi = ROOT.TMultiGraph("mg_%s_lumi"%name,"Multigraph lumi");
  for f1 in filBox:
    filter_dict = filterLayer(data,f1)
    for sec in filter_dict:
      print sec,": ",getVolume(sec)
      [lumi,time,curr] = vectors(filter_dict[sec])
      corrCurrIT(time,curr,temp[sec[0:-3]])
      v_curr = np.array(curr)/getVolume(sec)
      g_time = ROOT.TGraph(len(time),np.array(time)*1.0,v_curr)
      g_lumi = ROOT.TGraph(len(time),np.array(lumi)*1.0,v_curr)
      g_time.SetLineColor(color);
      g_lumi.SetLineColor(color);
      mg_time.Add(g_time,"l")
      mg_lumi.Add(g_lumi,"l")
  return [mg_time,mg_lumi]


# combines an iterable of graphs into a multigraph
def combineMultigraph(name,graphs):
  mg = ROOT.TMultiGraph("mg_%s"%name,"Multigraph");
  for g in graphs:
    mg.Add(g)
  return mg


# sets the axis ranges etc. for a graph showing the current as a function of time
def modifyGraphTime(mg):
  #mg.GetXaxis().SetLimits(1262304000,1360000000);
  mg.GetXaxis().SetLimits(1262304000,1438003626);
  mg.GetXaxis().SetTimeDisplay(1);
  mg.GetXaxis().SetTimeOffset(0.0);
  mg.GetXaxis().SetTimeFormat("#splitline{%d/%m}{%Y}");
  mg.GetXaxis().SetTitle("Date");
  mg.GetYaxis().SetTitle("#it{I}_{leak}/#it{V} [A/m^{3}]");
  mg.GetXaxis().SetLabelOffset(0.02);
  mg.GetXaxis().SetTimeOffset(0.0);
  mg.GetYaxis().SetTitleOffset(1.3);
  mg.GetXaxis().SetTitleOffset(1.4);


# sets the axis ranges etc. for a graph showing the current as a function of the luminosity
def modifyGraphLumi(mg):
  mg.GetXaxis().SetLimits(0.0,3500.0);
  mg.GetXaxis().SetTitle("Delivered Luminosity [pb^{-1}]");
  mg.GetYaxis().SetTitle("#it{I}_{leak}/#it{V} [A/m^{3}]");
  mg.GetXaxis().SetLabelOffset(0.02);
  mg.GetXaxis().SetTimeOffset(0.0);
  mg.GetYaxis().SetTitleOffset(1.3);
  mg.GetXaxis().SetTitleOffset(1.2);


# Running the lhcb style file
ROOT.gROOT.ProcessLine(".x ../../include/lhcbstyle.C")


# Load temperature

filenameTemp  = "/home/hep/egraveri/cmtuser/STMonitoring/STAging/data/temperature"
filenameTemp += "/temperature.root"
f_temp        = ROOT.TFile.Open(filenameTemp)

v_TT_CT5_temp = f_temp.Get("v_TT_CT5_temp")
v_TT_CT5_time = f_temp.Get("v_TT_CT5_time")
v_TT_CT6_temp = f_temp.Get("v_TT_CT6_temp")
v_TT_CT6_time = f_temp.Get("v_TT_CT6_time")

temp_TT         = {}
temp_TT["CT5"]  = [v_TT_CT5_time,v_TT_CT5_temp]
temp_TT["CT6"]  = [v_TT_CT6_time,v_TT_CT6_temp]

v_IT_A1_temp  = f_temp.Get("v_IT_A1_temp")
v_TT_A1_time  = f_temp.Get("v_IT_A1_time")
v_IT_A2_temp  = f_temp.Get("v_IT_A2_temp")
v_TT_A2_time  = f_temp.Get("v_IT_A2_time")
v_IT_A3_temp  = f_temp.Get("v_IT_A3_temp")
v_TT_A3_time  = f_temp.Get("v_IT_A3_time")
v_IT_B1_temp  = f_temp.Get("v_IT_B1_temp")
v_TT_B1_time  = f_temp.Get("v_IT_B1_time")
v_IT_B2_temp  = f_temp.Get("v_IT_B2_temp")
v_TT_B2_time  = f_temp.Get("v_IT_B2_time")
v_IT_B3_temp  = f_temp.Get("v_IT_B3_temp")
v_TT_B3_time  = f_temp.Get("v_IT_B3_time")
v_IT_C1_temp  = f_temp.Get("v_IT_C1_temp")
v_TT_C1_time  = f_temp.Get("v_IT_C1_time")
v_IT_C2_temp  = f_temp.Get("v_IT_C2_temp")
v_TT_C2_time  = f_temp.Get("v_IT_C2_time")
v_IT_C3_temp  = f_temp.Get("v_IT_C3_temp")
v_TT_C3_time  = f_temp.Get("v_IT_C3_time")
v_IT_T1_temp  = f_temp.Get("v_IT_T1_temp")
v_TT_T1_time  = f_temp.Get("v_IT_T1_time")
v_IT_T2_temp  = f_temp.Get("v_IT_T2_temp")
v_TT_T2_time  = f_temp.Get("v_IT_T2_time")
v_IT_T3_temp  = f_temp.Get("v_IT_T3_temp")
v_TT_T3_time  = f_temp.Get("v_IT_T3_time")

temp_IT         = {}
temp_IT["IT/HV/ST1A/BOXA"]  = [v_TT_A1_time,v_IT_A1_temp]
temp_IT["IT/HV/ST2A/BOXA"]  = [v_TT_A2_time,v_IT_A2_temp]
temp_IT["IT/HV/ST3A/BOXA"]  = [v_TT_A3_time,v_IT_A3_temp]
temp_IT["IT/HV/ST1A/BOXB"]  = [v_TT_B1_time,v_IT_B1_temp]
temp_IT["IT/HV/ST2A/BOXB"]  = [v_TT_B2_time,v_IT_B2_temp]
temp_IT["IT/HV/ST3A/BOXB"]  = [v_TT_B3_time,v_IT_B3_temp]
temp_IT["IT/HV/ST1C/BOXC"]  = [v_TT_C1_time,v_IT_C1_temp]
temp_IT["IT/HV/ST2C/BOXC"]  = [v_TT_C2_time,v_IT_C2_temp]
temp_IT["IT/HV/ST3C/BOXC"]  = [v_TT_C3_time,v_IT_C3_temp]
temp_IT["IT/HV/ST1C/BOXT"]  = [v_TT_T1_time,v_IT_T1_temp]
temp_IT["IT/HV/ST2C/BOXT"]  = [v_TT_T2_time,v_IT_T2_temp]
temp_IT["IT/HV/ST3C/BOXT"]  = [v_TT_T3_time,v_IT_T3_temp]

# Load data

filenameTT  = "/home/hep/egraveri/cmtuser/STMonitoring/STAging/data/leakageCurrent"
filenameTT += "/ttcurrent.pkl"
pkl_fileTT = open(filenameTT,'rb')
dataTT = pickle.load(pkl_fileTT)

filenameIT  = "/home/hep/egraveri/cmtuser/STMonitoring/STAging/data/leakageCurrent"
filenameIT += "/itcurrent.pkl"
pkl_fileIT = open(filenameIT,'rb')
dataIT = pickle.load(pkl_fileIT)


# Get current plots for TT for 1, 2, 3, and 4 read-out sector regions

[mg_TTaU_1_time,mg_TTaU_1_lumi]   = drawTT_sensor(dataTT,["S10","S4","S3","S9","S15","S16"],["US"],"TTaU_1sensor",ROOT.kRed,temp_TT);
[mg_TTbV_1_time,mg_TTbV_1_lumi]   = drawTT_sensor(dataTT,["S14","S8","S7","S13","S19","S20"],["VS"],"TTbV_1sensor",ROOT.kBlue,temp_TT);

[mg_TTaU_2_time,mg_TTaU_2_lumi]   = drawTT_sensor(dataTT,["S11","S5","S2","S8","S14","S17"],["US"],"TTaU_2sensor",ROOT.kRed,temp_TT);
[mg_TTbV_2_time,mg_TTbV_2_lumi]   = drawTT_sensor(dataTT,["S15","S9","S6","S12","S18","S21"],["VS"],"TTbV_2sensor",ROOT.kBlue,temp_TT);

[mg_TTaU_31_time,mg_TTaU_31_lumi] = drawTT_sensor(dataTT,["P3U","P7U","P2U","P6U"],["U"],"TTaU_31sensor",ROOT.kRed,temp_TT);
[mg_TTbV_32_time,mg_TTbV_32_lumi] = drawTT_sensor(dataTT,["P3V","P7V","P2V","P6V"],["V"],"TTbV_32sensor",ROOT.kBlue,temp_TT);
[mg_TTbV_33_time,mg_TTbV_33_lumi] = drawTT_sensor(dataTT,["S25","S24","S2","S3"],["VS"],"TTbV_33sensor",ROOT.kBlue,temp_TT);

[mg_TTaU_41_time,mg_TTaU_41_lumi] = drawTT_sensor(dataTT,["P8U","P4U","P5U","P1U"],["U"],"TTaU_41sensor",ROOT.kRed,temp_TT);
[mg_TTbV_42_time,mg_TTbV_42_lumi] = drawTT_sensor(dataTT,["P8V","P4V","P5V","P1V"],["V"],"TTbV_42sensor",ROOT.kBlue,temp_TT);
[mg_TTaU_43_time,mg_TTaU_43_lumi] = drawTT_sensor(dataTT,["S12","S6","S7","S1","S13","S18"],["US"],"TTaU_43sensor",ROOT.kRed,temp_TT);
[mg_TTbV_44_time,mg_TTbV_44_lumi] = drawTT_sensor(dataTT,["S26","S22","S16","S10","S4","S1","S5","S11","S17","S23"],["VS"],"TTbV_44sensor",ROOT.kBlue,temp_TT);

# Get current plots for IT for A, B, C and T boxes

[mg_IT_A_time,mg_IT_A_lumi]       = drawIT_sensor(dataIT,["BOXA"],"IT_BOXA",ROOT.kRed,temp_IT);
[mg_IT_B_time,mg_IT_B_lumi]       = drawIT_sensor(dataIT,["BOXB"],"IT_BOXB",ROOT.kRed,temp_IT);
[mg_IT_C_time,mg_IT_C_lumi]       = drawIT_sensor(dataIT,["BOXC"],"IT_BOXC",ROOT.kBlue,temp_IT);
[mg_IT_T_time,mg_IT_T_lumi]       = drawIT_sensor(dataIT,["BOXT"],"IT_BOXT",ROOT.kBlue,temp_IT);

# Get predictions for 1, 2, 3, and 4-sensor read-out sec

[mg_cent_lumi_TT1, mg_cent_time_TT1, mg_unc_lumi_TT1, mg_unc_time_TT1]         = getPredictions.predict("TT","1")
[mg_cent_lumi_TT2, mg_cent_time_TT2, mg_unc_lumi_TT2, mg_unc_time_TT2]         = getPredictions.predict("TT","2")
[mg_cent_lumi_TT3, mg_cent_time_TT3, mg_unc_lumi_TT3, mg_unc_time_TT3]         = getPredictions.predict("TT","3")
[mg_cent_lumi_TT4, mg_cent_time_TT4, mg_unc_lumi_TT4, mg_unc_time_TT4]         = getPredictions.predict("TT","4")

# Get predictions for A+C boxes and B+T boxes

[mg_cent_lumi_TAC, mg_cent_time_TAC, mg_unc_lumi_TAC, mg_unc_time_TAC]         = getPredictions.predict("IT","AC")
[mg_cent_lumi_TBT, mg_cent_time_TBT, mg_unc_lumi_TBT, mg_unc_time_TBT]         = getPredictions.predict("IT","BT")

# Combining multigraphs

mg_time_1   = combineMultigraph("TT_time1",[mg_unc_time_TT1,mg_TTaU_1_time,mg_TTbV_1_time,mg_cent_time_TT1])
mg_time_2   = combineMultigraph("TT_time2",[mg_unc_time_TT2,mg_TTaU_2_time,mg_TTbV_2_time,mg_cent_time_TT2])
mg_time_3   = combineMultigraph("TT_time3",[mg_unc_time_TT3,mg_TTaU_31_time,mg_TTbV_32_time,mg_TTbV_33_time,mg_cent_time_TT3])
mg_time_4   = combineMultigraph("TT_time4",[mg_unc_time_TT4,mg_TTaU_41_time,mg_TTbV_42_time,mg_TTaU_43_time,mg_TTbV_44_time,mg_cent_time_TT4])

mg_time_IT1 = combineMultigraph("IT_time1",[mg_unc_time_TAC,mg_IT_A_time,mg_IT_C_time,mg_cent_time_TAC])
mg_time_IT2 = combineMultigraph("IT_time2",[mg_unc_time_TBT,mg_IT_B_time,mg_IT_T_time,mg_cent_time_TBT])

mg_lumi_1   = combineMultigraph("TT_lumi1",[mg_TTaU_1_lumi,mg_TTbV_1_lumi])
mg_lumi_2   = combineMultigraph("TT_lumi2",[mg_unc_lumi_TT2,mg_TTaU_2_lumi,mg_TTbV_2_lumi,mg_cent_lumi_TT2])
mg_lumi_3   = combineMultigraph("TT_lumi3",[mg_unc_lumi_TT3,mg_TTaU_31_lumi,mg_TTbV_32_lumi,mg_TTbV_33_lumi,mg_cent_lumi_TT3])
mg_lumi_4   = combineMultigraph("TT_lumi4",[mg_unc_lumi_TT4,mg_TTaU_41_lumi,mg_TTbV_42_lumi,mg_TTaU_43_lumi,mg_TTbV_44_lumi,mg_cent_lumi_TT4])

mg_lumi_IT1 = combineMultigraph("IT_lumi1",[mg_unc_lumi_TAC,mg_IT_A_lumi,mg_IT_C_lumi,mg_cent_lumi_TAC])
mg_lumi_IT2 = combineMultigraph("IT_lumi2",[mg_unc_lumi_TBT,mg_IT_B_lumi,mg_IT_T_lumi,mg_cent_lumi_TBT])

# plot graphs

c_TT1_time  = ROOT.TCanvas("c_TT1_time","Canvas",800,600)
mg_time_1.Draw("a")
mg_time_1.SetMaximum(40.0)
modifyGraphTime(mg_time_1)
mg_time_1.GetHistogram().Draw("axissame")

c_TT2_time  = ROOT.TCanvas("c_TT2_time","Canvas",800,600)
mg_time_2.Draw("a")
mg_time_2.SetMaximum(15.0)
modifyGraphTime(mg_time_2)
mg_time_2.GetHistogram().Draw("axissame")

c_TT3_time  = ROOT.TCanvas("c_TT3_time","Canvas",800,600)
mg_time_3.Draw("a")
modifyGraphTime(mg_time_3)
mg_time_3.SetMaximum(10.0)
mg_time_3.GetHistogram().Draw("axissame")

c_TT4_time  = ROOT.TCanvas("c_TT4_time","Canvas",800,600)
mg_time_4.Draw("a")
modifyGraphTime(mg_time_4)
mg_time_4.SetMaximum( 6.0)
mg_time_4.GetHistogram().Draw("axissame")

c_IT1_time  = ROOT.TCanvas("c_IT1_time","Canvas",800,600)
mg_time_IT1.Draw("a")
modifyGraphTime(mg_time_IT1)
mg_time_IT1.SetMaximum(10.0)
mg_time_IT1.GetHistogram().Draw("axissame")

c_IT2_time  = ROOT.TCanvas("c_IT2_time","Canvas",800,600)
mg_time_IT2.Draw("a")
modifyGraphTime(mg_time_IT2)
mg_time_IT2.SetMaximum(10.0)
mg_time_IT2.GetHistogram().Draw("axissame")


c_TT1_lumi  = ROOT.TCanvas("c_TT1_lumi","Canvas",800,600)
mg_lumi_1.Draw("a")
mg_lumi_1.SetMaximum(40.0)
modifyGraphLumi(mg_lumi_1)
mg_lumi_1.GetHistogram().Draw("axissame")

c_TT2_lumi  = ROOT.TCanvas("c_TT2_lumi","Canvas",800,600)
mg_lumi_2.Draw("a")
mg_lumi_2.SetMaximum(15.0)
modifyGraphLumi(mg_lumi_2)
mg_lumi_2.GetHistogram().Draw("axissame")

c_TT3_lumi  = ROOT.TCanvas("c_TT3_lumi","Canvas",800,600)
mg_lumi_3.Draw("a")
mg_lumi_3.SetMaximum(10.0)
modifyGraphLumi(mg_lumi_3)
mg_lumi_3.GetHistogram().Draw("axissame")

c_TT4_lumi  = ROOT.TCanvas("c_TT4_lumi","Canvas",800,600)
mg_lumi_4.Draw("a")
mg_lumi_4.SetMaximum( 6.0)
modifyGraphLumi(mg_lumi_4)
mg_lumi_4.GetHistogram().Draw("axissame")

c_IT1_lumi  = ROOT.TCanvas("c_IT1_lumi","Canvas",800,600)
mg_lumi_IT1.Draw("a")
mg_lumi_IT1.SetMaximum(10.0)
modifyGraphLumi(mg_lumi_IT1)
mg_lumi_IT1.GetHistogram().Draw("axissame")

c_IT2_lumi  = ROOT.TCanvas("c_IT2_lumi","Canvas",800,600)
mg_lumi_IT2.Draw("a")
mg_lumi_IT2.SetMaximum(10.0)
modifyGraphLumi(mg_lumi_IT2)
mg_lumi_IT2.GetHistogram().Draw("axissame")