import ROOT

import numpy as np

#dn_lumi = ROOT.TString("/home/hep/che/cmtuser/Vetra_v13r2/ST/STAging/data/lumi");
dn_lumi = ROOT.TString("/home/hep/egraveri/cmtuser/STMonitoring/STAging/data/lumi");
fn_lumi = ROOT.TString("%s/lumi.root"%(dn_lumi.Data()));

#dn_sector = ROOT.TString("/home/hep/che/cmtuser/Vetra_v13r2/ST/STAging/data/db");
dn_sector = ROOT.TString("/home/hep/egraveri/cmtuser/STMonitoring/STAging/data/db");
fn_sector = ROOT.TString("%s/fluence.root"%(dn_sector.Data()));

f_lumi   = ROOT.TFile.Open(fn_lumi.Data());
f_sector = ROOT.TFile.Open(fn_sector.Data());

# Get information on luminosity
v_fillTime  = f_lumi.Get("v_fillTime_tot");
v_beamEner  = f_lumi.Get("v_beamEner_tot");
v_stopTime  = f_lumi.Get("v_stopTime_tot");
v_peakLumi  = f_lumi.Get("v_peakLumi_tot");

# Get information on fluence
v_sector    = {}
v_flux7_v   = {}
v_flux7_e   = {}
v_flux8_v   = {}
v_flux8_e   = {}

v_sector["TT"]   = f_sector.Get("v_sector_TTaU")
v_flux7_v["TT"]  = f_sector.Get("v_flux7_v_TTaU")
v_flux7_e["TT"]  = f_sector.Get("v_flux7_e_TTaU")
v_flux8_v["TT"]  = f_sector.Get("v_flux8_v_TTaU")
v_flux8_e["TT"]  = f_sector.Get("v_flux8_e_TTaU")

v_sector["IT"]   = f_sector.Get("v_sector_T3X2")
v_flux7_v["IT"]  = f_sector.Get("v_flux7_v_T3X2")
v_flux7_e["IT"]  = f_sector.Get("v_flux7_e_T3X2")
v_flux8_v["IT"]  = f_sector.Get("v_flux8_v_T3X2")
v_flux8_e["IT"]  = f_sector.Get("v_flux8_e_T3X2")

# dictionnary of read-out sectors matched to boxes (IT) or types of read-out sectors (TT; 1, 2, 3 and 4 sensor sectors)
sect_dict  = {}
sect_dict["TT"]  = {
  "1" : [2627,2628,2634,2633,2639,2640],
  "2" : [2626,2629,2635,2632,2638,2641],
  "3" : [2615,2614,2611,2610,2607,2606,2603,2602,2599,2598,2595,2594,2659,2658,2663,2662,2667,2666,2671,2670,2675,2674,2679,2678],
  "4" : [2616,2613,2612,2609,2608,2605,2604,2601,2600,2597,2596,2594,2660,2657,2664,2661,2668,2665,2672,2669,2676,2673,2680,2677,2642,2636,2630,2637,2631,2625]
}

sect_dict["IT"]  = {
  "T"  : [7303,7302,7301,7300,7299,7298,7297],
  "A"  : [7239,7238,7237,7236,7235,7234,7233],
  "C"  : [7207,7206,7205,7204,7203,7202,7201],
  "B"  : [7271,7270,7269,7268,7267,7266,7265],
  "AC" : [7207,7206,7205,7204,7203,7202,7201,7239,7238,7237,7236,7235,7234,7233],
  "BT" : [7271,7270,7269,7268,7267,7266,7265,7303,7302,7301,7300,7299,7298,7297]
}

# conversion factor linking the number of 1-MeV neutron flux per collision to the corresponding in a pb
flukaConvFac = 7.2e4

# Get the index in a vector such that the value associated to this index is the closest to the input value (val)
def findIndex(v_ind,val):
  ind = -1
  min = 1000000
  for i,v in enumerate(v_ind):
    if ROOT.TMath.Abs(v-val)<min:
      min = ROOT.TMath.Abs(v-val)
      ind = i
  return ind

# calculate the time factor according to the Arrhenius relation
def Arrhenius(para,temp):
  return 1.0/(para["k0I"]*ROOT.TMath.Exp(-para["EI"]/(para["kb"]*temp)))

# calculate the change in the leakage current due to annealing
# alpha(t) = alpha_0 + alpha_1 * exp(-t/tau)
def GetDiff(para,temp,dFlux,Flux,val,tStep):
  tau = Arrhenius(para,temp)
  return ((para["a0"]+para["aI"])*dFlux-((val-para["a0"]*Flux)/tau))*tStep

# get prediction for a given sector
def predict(det,sect):
  flux7_v = 0.0
  flux7_e = 0.0
  flux8_v = 0.0
  flux8_e = 0.0

  c = 0.0
  for s in sect_dict[det][sect]:
    c += 1.0
    ind = findIndex(v_sector[det],s)
    print ind
    flux7_v += v_flux7_v[det](ind)
    flux7_e += v_flux7_e[det](ind)
    flux8_v += v_flux8_v[det](ind)
    flux8_e += v_flux8_e[det](ind)

  flux7_v /= c
  flux7_e /= c
  flux8_v /= c
  flux8_e /= c

  # define start and stop time for the prediction (unixtime)
  tStart = 1262325600;  # 1-1-2010
  #tStop  = 1372654800;  # 1-7-2013
  tStop  = 1437841532;  # 25-7-2015
  tRun   = tStart;
  tFillEnd = tStart;

  # define time step in secs
  tStep  = 6000.0 # 100 min

  v_time = []
  v_val  = []
  v_low  = []
  v_hig  = []
  v_lumi = []
  
  # parameters of the model
  para_v = {"a0" :  6.27e-17, "aI" : 7.73e-17, "k0I" : 4.200e13, "EI" : 1.11, "beta" : 3.07e-18, "t0" :  6000.0, "kb" : 8.617e-5}
  para_e = {"a0" :  0.09e-17, "aI" : 0.06e-17, "k0I" : 0.500e13, "EI" : 0.02, "beta" : 0.18e-18, "t0" :     0.0, "kb" : 0.000e-5}

  N = 1

  temp_e = 2.0

  para_list   = [para_v]
  flux7_list  = [flux7_v]
  flux8_list  = [flux8_v]
  tempU_list  = [0.0]
  val_list    = [0.0]
  flux_list   = [0.0]
  

  # get N-1 replicas of the path with variations of the parameters according to the uncertainties on the parameters
  for i in range(1,N):
    tempU_list.append(0.0)
    #tempU_list.append(ROOT.gRandom.Gaus()*temp_e)
    val_list.append(0.0)
    flux_list.append(0.0)
    flux7_list.append(flux7_v+ROOT.gRandom.Gaus()*flux7_e*1.0)
    flux8_list.append(flux8_v+ROOT.gRandom.Gaus()*flux8_e*1.0)
    para_temp = para_v.copy()
    for key in para_temp.keys():
      while True:
        para_temp[key] = para_v[key] + ROOT.gRandom.Gaus()*para_e[key]
        if para_temp[key]>0.0:
          break
    print para_temp
    para_list.append(para_temp)

  i_lumi = 0
  lumi = 0.0

  # run the simulation
  while tRun<tStop:
    print "%13.0f sec : %20.0f pb"%(tRun,lumi/1e6)
    temp = 273.15+8.0
    while i_lumi<v_fillTime.GetNoElements() and tRun>v_stopTime(i_lumi):
      lumi += (v_stopTime(i_lumi)-v_fillTime(i_lumi))*v_peakLumi(i_lumi)
      i_lumi += 1

    for i in range(0,len(val_list)):
      temp_val = temp+tempU_list[i]
      dFlux  = 0.0
      
      if i_lumi<v_fillTime.GetNoElements():
        if tRun<v_fillTime(i_lumi):
          dFlux  = 0.0
        else:
          if v_beamEner(i_lumi)==7.0:
            dFlux = flux7_list[i]*v_peakLumi(i_lumi)*flukaConvFac
          elif v_beamEner(i_lumi)==8.0:
            dFlux = flux8_list[i]*v_peakLumi(i_lumi)*flukaConvFac
          else:
            dFlux = 0.0

      # calculate the evolution of the current according to 2nd order Runga-Kutta
      dval = GetDiff(para_list[i],temp_val,dFlux,flux_list[i],val_list[i],tStep)
      val_list[i] += dval/2.0
      flux_list[i] += dFlux*tStep
      dval = GetDiff(para_list[i],temp_val,dFlux,flux_list[i],val_list[i],tStep)
      val_list[i] += dval/2.0
      
    # Sort the value list
    val_list_sort = list(val_list)
    v_val.append(val_list[0])
    val_list_sort = sorted(val_list_sort)
    v_time.append(tRun)
    # get the 68% confidence interval
    v_low.append(val_list_sort[(int)((1.0-0.684)*N/2.0)])
    v_hig.append(val_list_sort[(int)((1.0+0.684)*N/2.0)])
    if tRun<1356998400.0:
      v_lumi.append(lumi/1e6)

    tRun += tStep
  
  # reduce the number of data points for plotting
  vr_time     = ROOT.TVectorD(len(v_time[::100]),np.array(v_time[::100]))
  vr_lumi     = ROOT.TVectorD(len(v_lumi[::100]),np.array(v_lumi[::100]))
  
  vr_val      = ROOT.TVectorD(len(v_val[::100]),np.array(v_val[::100])*1e6)
  vr_errcent  = ROOT.TVectorD(len(v_low[::100]),(np.array(v_hig[::100])+np.array(v_low[::100]))/2.0*1e6)
  vr_errband  = ROOT.TVectorD(len(v_hig[::100]),(np.array(v_hig[::100])-np.array(v_low[::100]))/2.0*1e6)
  
  vr_hig      = ROOT.TVectorD(len(v_hig[::100]),np.array(v_hig[::100])*1e6)
  vr_low      = ROOT.TVectorD(len(v_low[::100]),np.array(v_low[::100])*1e6)

  vr_val_lumi      = ROOT.TVectorD(len(v_lumi[::100]),np.array(v_val[::100])[0:len(v_lumi[::100])]*1e6)
  vr_errcent_lumi  = ROOT.TVectorD(len(v_lumi[::100]),(np.array(v_hig[::100])[0:len(v_lumi[::100])]+np.array(v_low[::100])[0:len(v_lumi[::100])])/2.0*1e6)
  vr_errband_lumi  = ROOT.TVectorD(len(v_lumi[::100]),(np.array(v_hig[::100])[0:len(v_lumi[::100])]-np.array(v_low[::100])[0:len(v_lumi[::100])])/2.0*1e6)


  # divide obtained vectors by the volume of the connected sensors
  if det=="TT":
    if sect=="1" or sect=="4":
      vr_val     *=1.0/(0.0500*9.64*9.44)
      vr_errband *=1.0/(0.0500*9.64*9.44)
      vr_errcent *=1.0/(0.0500*9.64*9.44)
      vr_hig     *=1.0/(0.0500*9.64*9.44)
      vr_low     *=1.0/(0.0500*9.64*9.44)
    else:
      vr_val     *=1.2/(0.0500*9.64*9.44)
      vr_errband *=1.2/(0.0500*9.64*9.44)
      vr_errcent *=1.2/(0.0500*9.64*9.44)
      vr_hig     *=1.2/(0.0500*9.64*9.44)
      vr_low     *=1.2/(0.0500*9.64*9.44)
  else:
    if sect=="AC":
      vr_val     *=1.0/(0.041*7.6*11)
      vr_errband *=1.0/(0.041*7.6*11)
      vr_errcent *=1.0/(0.041*7.6*11)
      vr_hig     *=1.0/(0.041*7.6*11)
      vr_low     *=1.0/(0.041*7.6*11)
    else:
      vr_val     *=1.0/(0.032*7.6*11*1.6)
      vr_errband *=1.0/(0.032*7.6*11*1.6)
      vr_errcent *=1.0/(0.032*7.6*11*1.6)
      vr_hig     *=1.0/(0.032*7.6*11*1.6)
      vr_low     *=1.0/(0.032*7.6*11*1.6)

  # Create the luminosity graphs (central value, confidence interval [also as area])
  g_time  = ROOT.TGraph(vr_time,vr_val)
  g_time.SetLineWidth(5)
  ge_time = ROOT.TGraphErrors(vr_time,vr_errcent,ROOT.TVectorD(vr_time.GetNoElements()),vr_errband);
  ge_time.SetFillColor(ROOT.kGray)
  gtop_time = ROOT.TGraph(vr_time,vr_hig );
  gtop_time.SetLineStyle(ROOT.kDashed)
  gtop_time.SetLineWidth(2)
  gbot_time = ROOT.TGraph(vr_time,vr_low);
  gbot_time.SetLineStyle(ROOT.kDashed)
  gbot_time.SetLineWidth(2)

  # Create the time graphs (central value, confidence interval [also as area])
  g_lumi  = ROOT.TGraph(vr_lumi,vr_val)
  g_lumi.SetLineWidth(5)
  ge_lumi = ROOT.TGraphErrors(vr_lumi,vr_errcent,ROOT.TVectorD(vr_lumi.GetNoElements()),vr_errband);
  ge_lumi.SetFillColor(ROOT.kGray)
  gtop_lumi = ROOT.TGraph(vr_lumi,vr_hig );
  gtop_lumi.SetLineStyle(ROOT.kDashed)
  gtop_lumi.SetLineWidth(2)
  gbot_lumi = ROOT.TGraph(vr_lumi,vr_low);
  gbot_lumi.SetLineStyle(ROOT.kDashed)
  gbot_lumi.SetLineWidth(2)

  # Combine the luminosity graphs
  mg_time_cent = ROOT.TMultiGraph("mg_time_cent_%s"%det,"multigraph")
  mg_time_unc  = ROOT.TMultiGraph("mg_time_unc_%s"%det, "multigraph")
  mg_time_cent.Add(g_time,"l")
  mg_time_cent.Add(gtop_time,"l")
  mg_time_cent.Add(gbot_time,"l")
  mg_time_unc.Add(ge_time,"3")

  # Combine the time graphs
  mg_lumi_cent = ROOT.TMultiGraph("mg_lumi_cent_%s"%det,"multigraph")
  mg_lumi_unc  = ROOT.TMultiGraph("mg_lumi_unc_%s"%det, "multigraph")
  mg_lumi_cent.Add(g_lumi,"l")
  mg_lumi_cent.Add(gtop_lumi,"l")
  mg_lumi_cent.Add(gbot_lumi,"l")
  mg_lumi_unc.Add(ge_lumi,"3")

  print "Done"
  return [mg_lumi_cent,mg_time_cent,mg_lumi_unc,mg_time_unc]


# plot the legends
if __name__ == '__main__':
  ROOT.gROOT.ProcessLine(".x ../../include/lhcbstyle.C")
  
  h1 = ROOT.TH1F("h1","h1",10,0.0,1.0)
  h1.SetLineColor(ROOT.kRed)
  h1.SetLineWidth(4)
  h2 = ROOT.TH1F("h2","h2",10,0.0,1.0)
  h2.SetLineColor(ROOT.kBlue)
  h2.SetLineWidth(4)
  h3 = ROOT.TH1F("h3","h3",10,0.0,1.0)
  h3.SetFillColor(ROOT.kGray)
  h3.SetLineColor(ROOT.kBlack)
  h3.SetLineWidth(4)
  
  leg1 = ROOT.TLegend(0.0,0.0,1.0,1.0)
  leg1.SetFillStyle(0)
  leg1.SetTextSize(0.14)
  
  leg1.AddEntry(h1,"U layer (TT)","l")
  leg1.AddEntry(h2,"V layer (TT)","l")
  leg1.AddEntry(h3,"Prediction (T_{a} = 8#circ C)","fl")
  
  c_leg1 = ROOT.TCanvas("c_leg1","c_leg1",300,200)
  leg1.Draw()

  leg2 = ROOT.TLegend(0.0,0.0,1.0,1.0)
  leg2.SetFillStyle(0)
  leg2.SetTextSize(0.14)
  
  leg2.AddEntry(h1,"A box (IT)","l")
  leg2.AddEntry(h2,"C box (IT)","l")
  leg2.AddEntry(h3,"Prediction (T_{a} = 8#circ C)","fl")
  
  c_leg2 = ROOT.TCanvas("c_leg2","c_leg2",300,200)
  leg2.Draw()

  leg3 = ROOT.TLegend(0.0,0.0,1.0,1.0)
  leg3.SetFillStyle(0)
  leg3.SetTextSize(0.14)
  
  leg3.AddEntry(h1,"B box (IT)","l")
  leg3.AddEntry(h2,"T box (IT)","l")
  leg3.AddEntry(h3,"Prediction (T_{a} = 8#circ C)","fl")
  
  c_leg3 = ROOT.TCanvas("c_leg3","c_leg3",300,200)
  leg3.Draw()