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)
    nevents = tree.GetEntries()
    print(f"Found {nfiles} file(s) with {nevents} events")
    return 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(wf.peak())
        hist(f"{name}/peak", 100000, 0, 100000, wf.peak())
    return

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

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

    for i, ievent in enumerate(range(nevents)):
        tree.GetEntry(ievent)

        runNum = tree.EventInfo.runNumber()        
        eventNum = tree.EventInfo.eventNumber()

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

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

    return


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

if __name__ == "__main__":

    R.gROOT.SetBatch(True)
    #R.RootUtils.PyROOTTTreePatch.Initialize(R.TTree, R.TChain, R.TBranch)

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

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

    tree = load(input_files)

    initialize()
    execute(tree, -1)
    finalize()