import ROOT as R
import os

def load(input_files):
    "Load all physics files from input folder"
    
    chain = R.TChain("CollectionTree")

    input_files = [input_files] if isinstance(input_files, str) else input_files

    nfiles = 0
    for f in input_files:
        nfiles += chain.Add(f)

    if not nfiles:
        print(f"No files found")
        return None

    tree = R.xAOD.MakeTransientTree(chain, 2)
    
    nevents = tree.GetEntries()
    print(f"Found {nfiles} file(s) with {nevents} events")
    return chain, tree

m_hists = {}
def hist(name, nbins, lo, hi, value, weight = 1.):
    "Fill a histogram, bookiung it if doesn't already exist"

    if not name in m_hists:
        m_hists[name] = R.TH1F(os.path.basename(name), name, nbins, lo, hi)

    m_hists[name].Fill(value, weight)
    return

def write(name):
    "Write the histograms to a file"

    f = R.TFile.Open(name, "RECREATE")
    
    for n, h in m_hists.items():
        path = os.path.dirname(n)
        if path and not f.GetDirectory(path):
            f.mkdir(path)
            
        f.cd(path)
        h.Write()

    f.Close()

    return

def initialize():
    "Before looping over events"
    return


def plot_waveform(tree, name, container):
    "Analayse waveform container"

    wfs = getattr(tree, container)
    for i, wf in enumerate(wfs):
        print(i, wf.channel())
        hist(f"{name}/peak", 100000, 0, 100000, wf.peak())
    return

def selection(tree, treexAOD):
    print (treexAOD.FaserTriggerData.tap())
    print (tree.ClusterFit.m_tracks.size())
    return True
    

def execute(tree, treexAOD, nevents = -1, freq = 1000):
    "Loop oever events"

    if nevents == -1: nevents = tree.GetEntries()

    for i, ievent in enumerate(range(nevents)):
        if i == 0: continue
        tree.GetEntry(ievent)
        treexAOD.GetEntry(ievent)
        runNum = treexAOD.EventInfo.runNumber()        
        eventNum = treexAOD.EventInfo.eventNumber()

        if i%freq == 0:
            print(f">>> Processing Event {i}/{nevents}")

        if not selection(tree, treexAOD): continue

        plot_waveform(treexAOD, "Calo", "CaloWaveformHits")
        plot_waveform(treexAOD, "Trig", "TriggerWaveformHits")

    return


def finalize():
    "After looping over events"
    
    write("hists.root")
    return

if __name__ == "__main__":

    R.gROOT.SetBatch(True)

    if not R.xAOD.Init().isSuccess():
        print("Unable to intialise xAOD setup")

    run = "3980"
    input_files = f"/eos/project-f/faser-commissioning/Reconstruction/Run-00{run}/Faser-Physics*r0002-xAOD.root"

    #input_files = []
    #for run in range(3430, 3440):
    #    input_files.append(f"/eos/project-f/faser-commissioning/Reconstruction/Run-00{run}/Faser-Physics*.root")

    tree, treexAOD = load(input_files)

    initialize()
    execute(tree, treexAOD, 20, 1)
    finalize()