import sys
sys.path.append('../../FourVector')
sys.path.append('../project')

from FourVector import FourVector
from ThreeVector import ThreeVector

from FutureColliderTools import SmearVertex, GetCorrectedMass, GetMissingMass2, GetQ2
from FutureColliderDataLoader import LoadData_KMuNu, LoadData_DsMuNu

from FutureColliderVariables import *

import numpy as np
import ROOT

ROOT.enableJSVis()
ROOT.gStyle.SetOptStat(0)



"""
Explanation: Setup the environment
End of explanation
"""


from prettytable import PrettyTable
def PlotHistograms(Filename):
    
    #Due to python's garbage collection the Histograms, Stacks, Files etc. need to be stored somewhere.
    global Histograms
    global Stacks
    global Files
    global legend
    global Textarr
    
    canvas.Divide(3,3)
    Histograms_norm = []
    
    table = PrettyTable()
    
    for it, resolution in enumerate(np.linspace(0.3, 1.0, 8)):
        
        # Open the histogram file for reading
        File_Toys = ROOT.TFile.Open(Filename.format(resolution), "READ")
        Files += [File_Toys]

        KeyList = [ Key for Key in File_Toys.GetListOfKeys() if "MCORR" in Key.GetName() ]
        Histograms_norm += [[ Key.ReadObj().Clone() for Key in KeyList if "_norm" in Key.GetName()]]
        Histograms +=      [[ Key.ReadObj().Clone() for Key in KeyList if "_norm" not in Key.GetName()]]

        # On the first loop, fill the Legend, and add the first column to the Table
        if it == 0:
            HistogramNames = [Hist.GetName().replace("MCORR_", "").replace("Combinatorial", "C") for Hist in Histograms[0] ]
            table.add_column("", HistogramNames)
            
            HistogramNames = [Hist.GetTitle() for Hist in Histograms_norm[0] ]
            for Name, Hist in zip(HistogramNames, Histograms_norm[0]):
                legend.AddEntry(Hist, Name, "lep")
        
        HistogramYields = [int(Hist.Integral()) for Hist in Histograms[-1] ]
        table.add_column(str(resolution), HistogramYields)
        
        #Make a stack of the Historams
        Stack = ROOT.THStack()
        Color = 1
        for Hist in Histograms_norm[it]:
            #Hist.Sumw2(False)
            Hist.SetDirectory(0)
            Hist.SetLineColor(Color)
            Hist.SetFillStyle(0)
            Color+=1
            Stack.Add(Hist)

        Stacks += [Stack]
        canvas.cd(it+1)
        Stack.Draw("nostack")
        
        Text = ROOT.TPaveText(0.1,0.7,0.5,0.9, "NDC")
        Text.SetFillColor(0)
        Text.SetBorderSize(0)
        Text.AddText("#sigma (Vertex) =  "+str(resolution)+" #dot #sigma (LHCb)")
        Text.Draw()
        Textarr += [Text]
        
    canvas.cd(9)
    legend.Draw()
    canvas.DrawClone()
    print table


"""
Explanation: Do the work to do the plotting
End of explanation
"""


canvas = ROOT.TCanvas("c1", "c1", 900, 900)

Histograms = []
Stacks = []
Files = []
Textarr = []
legend = ROOT.TLegend(0.101,0.101,0.899,0.899);

PlotHistograms("../output/Source_Histograms_DsMu_{0}_LHCb.root")

canvas.Draw()


"""
Explanation: Print yields and plot all the templates
$B_s \to D_s^-  \mu^+ \nu_{\mu} $
End of explanation
"""


canvas = ROOT.TCanvas("c1", "c1", 900, 900)

Histograms = []
Stacks = []
Files = []
legend = ROOT.TLegend(0.101,0.1,0.9,0.9);

PlotHistograms("../output/Source_Histograms_KMu_{0}_LHCb.root")

canvas.Draw()

"""
Explanation: Print yields and plot all the templates
$B_s \to k^-  \mu^+ \nu_{\mu} $
End of explanation
"""


canvas = ROOT.TCanvas("c1", "c1", 900, 900)

Histograms = []
Stacks = []
Files = []
legend = ROOT.TLegend(0.1,0.1,0.9,0.9);

PlotHistograms("../output/Source_Histograms_DsMu_{0}_LHCb_Merged.root")

canvas.Draw()

canvas = ROOT.TCanvas("c1", "c1", 900, 900)

Histograms = []
Stacks = []
Files = []
legend = ROOT.TLegend(0.1,0.1,0.9,0.9);

PlotHistograms("../output/Source_Histograms_KMu_{0}_LHCb_Merged.root")

canvas.Draw()

"""
Explanation: Combining/Removing Histograms
We can see that several of the templates have similar shapes.
The following can be combined or removed:
For $B_s \to D_s^-  \mu^+ \nu_{\mu} $

Both $B_s \to D_s \tau X$ modes
Both Further Excited Ds* resonance

For $B_s \to K^+  \mu^+ \nu_{\mu} $

Both $B_s \to D_s \tau X$ modes
Both Combinatorial Modes
Both Further Excited Ds* resonances and the Lambda Mode
$B_s \to D_s  \mu^+ \nu_{\mu} $ and $B_s \to D_s^*  \mu^+ \nu_{\mu} $
End of explanation
"""


