import ROOT
import os
import sys

def main():
    # Idiot proofing
    if len(sys.argv) < 2:
        print("Usage: python calculate_True_nEvtGen.py <sample_directory>")
        print("Note: Make sure to be in a CMSSW environment so that 'import ROOT' works.")
        sys.exit(1)

    # Accessing the base directory
    sampledir = sys.argv[1]
    if not os.path.isdir(sampledir):
        print("Error: Directory '" + sampledir + "' does not exist.")
        sys.exit(1)

    summary = []
    print('Calculating events in {} ...'.format(sampledir))

    # Accessing the subdirectories
    subsamples = os.listdir(sampledir)
    for subsample in subsamples:
        indir = os.path.join(sampledir, subsample)
        
        if not os.path.isdir(indir):
            continue

        files = os.listdir(indir)
        
        print("\n------------------")
        print(subsample)
        print("------------------")
        nevtgen_total = 0
        
        # Accessing individual files
        for filename in files:

            nevtgen = 0
            nevtskimmed = 0
            
            if not filename.endswith('.root'): continue

            filepath = os.path.join(indir, filename)

            # TFile
            tfile = ROOT.TFile.Open(filepath)
            if not tfile or tfile.IsZombie():
                print("\033[91mZombie file or can't open: {}\033[0m".format(filepath))
                continue

            # TTree
            tree = tfile.Get("Events")
            if not tree:
                print("\033[91mError: No 'Events' tree found in {}\033[0m".format(filepath))
                tfile.Close()
                continue
            
            # Finding the right branch:
            if tree.GetBranch("nevtgen"):
                nevtskimmed = tree.GetEntries()
                tree.GetEntry(0)
                nevtgen = tree.nevtgen
                print("{}\t{}\tSkimmed to -> {}".format(filename, nevtgen, nevtskimmed))
                
            elif tree.GetBranch("event"):
                nevtgen = tree.GetEntries()
                print("{}\t{}".format(filename, nevtgen))

            else: print("\033[91mError: Branch not found :{}\033[0m".format(filepath))

            # Closing the file after processing
            tfile.Close()
            nevtgen_total += nevtgen
            #break #file
            
        #Summary for subsample:
        print("\033[33mTotal nEvtGen for {} : {}\033[0m".format(subsample, nevtgen_total))
        summary_string = "{:<20} {}".format(subsample, nevtgen_total)
        summary.append(summary_string)
        #break #subsample

    #Write a summary:
    print("\n-------------------------------")
    print("\033[33m           Summary             \033[0m")
    print("-------------------------------")
    print("Input directory = {}".format(sampledir))
    print("{:<20} {}".format("Sub-sample", "nEvtGen"))
    for string in summary: print(string)
    print("-------------------------------\n")
    

if __name__ == "__main__": main()