import ROOT
from ROOT import TFile, TTree, gROOT
import sys
import numpy as np
import pandas as pd
import argparse
from progressbar import progressbar


ROOT.gInterpreter.ProcessLine("""
    struct pystruct{
      std::vector<double> doublevec{};
      std::vector<int> intvec{};
      Int_t integer = 0;
    }; 
""")

fitpointsleaves = ['nHits','xPos','yPos','zPos','sensorId']
rawdataleaves = ['nPixHits','euEvt','col','row','tot','lv1','iden']
tracksleaves = ['nTrackParams','euEvt','xPos','yPos','dxdz','dydz','trackNum','iden','chi2','ndof']
typedict = {'nHits':'nHits/I','xPos':'double','yPos':'double','zPos':'double','sensorId':'int',
            'nPixHits':'nPixHits/I','euEvt':'euEvt/I','col':'int','row':'int','tot':'int','lv1':'int','iden':'int','hitTime':'int','frameTime':'double',
            'nTrackParams':'nTrackParams/I','euEvt':'euEvt/I','xPos':'double',
            'yPos':'double','dxdz':'double','dydz':'double',
            'trackNum':'int','iden':'int','chi2':'double','ndof':'double'}

def main():

    parser = argparse.ArgumentParser(description='Script for merging track collections for TBMon2 analysis')
    parser.add_argument('-f', dest="files", action='append')
    args = parser.parse_args()
    mergetrees(args.files)


def mergetrees(filenames):

    gROOT.LoadMacro("writevectordouble.C")

    files = []
    for file in filenames:
        files.append(TFile(file))

    writeversiontree(1.3)
    sorttree(files, fitpointsleaves, 'fitpoints')
    sorttree(files, rawdataleaves, 'rawdata')
    sorttree(files, tracksleaves, 'tracks')

def writeversiontree(version):

    output = TFile("output.root","update")
    new_tree = TTree('version', 'version')
    new_tree.GetEntry(0)
    pycpp = ROOT.pystruct()
    pycpp.doublevec.push_back(version)
    ROOT.writevectordouble(pycpp.doublevec,new_tree,0,'no')
    output.Write("", ROOT.TObject.kOverwrite)
    output.Close()

def sorttree(files, leaves, treename):    

    dataframes = []
    filenumber = 0

    for file in files:

        tree = file.Get(treename)
        df = pd.DataFrame()

        for leaf in leaves:

            leaftype = None

            print('File {}: Running on {}'.format(filenumber, leaf))

            array = []
            if '/I' in typedict[leaf]:
                leaftype = 'integer'
                for entry in range(tree.GetEntries()):
                    tree.GetEntry(entry)
                    integer = tree.GetLeaf(leaf).GetValue()
                    array.append(integer)
                array = np.asarray(array)    
                df[leaf] = array

            else:    
                leaftype = 'vector'
                vec = ROOT.std.vector(typedict[leaf])()
                tree.SetBranchAddress(leaf, vec)
                for entry in range(tree.GetEntries()):
                    tree.GetEntry(entry)
                    values = []
                    for value in vec:
                        values.append(value)
                    array.append(values)
                array = np.asarray(array)
                df[leaf] = array    
            tree.ResetBranchAddresses()

        dataframes.append(df)  

        filenumber+=1

    writetree(dataframes,leaves,treename)     

def writetree(dataframes, leaves, treename): 

    output = TFile("output.root","update")
    new_tree = TTree(treename, treename)

    temp_df = pd.DataFrame(dtype=object)
    n_dataframes = len(dataframes) 
    n_entries = dataframes[0].shape[0]#2000    

    pycpp = ROOT.pystruct()

    print('Merging {} tree for {} files with {} entries'.format(treename, n_dataframes, n_entries))

    for leaf in leaves:
        
        leaftype = None

        if '/I' in typedict[leaf]:
            leaftype = 'integer'
        else:
            leaftype = 'vector'    


        for otherleaf in new_tree.GetListOfBranches():

            if otherleaf.GetName() != leaf:

                new_tree.SetBranchStatus(otherleaf.GetName(),0)

        new_tree.ResetBranchAddresses()

        #for entry in range(n_entries):
        for entry in progressbar(range(n_entries)):

            pycpp.doublevec.clear()
            pycpp.intvec.clear()

            if leaftype == 'vector':

                if treename == 'tracks':

                    length = len(dataframes[0].iloc[entry][leaf])
                    for i in range(length): 

                        for df in dataframes:

                            temp = df.iloc[entry][leaf]
                            if typedict[leaf] == 'double':
                                pycpp.doublevec.push_back(temp[i]) 
                            elif typedict[leaf] == 'int':
                                pycpp.intvec.push_back(temp[i])                     

                    if typedict[leaf] == 'double':
                        ROOT.writevectordouble(pycpp.doublevec,new_tree,entry,leaf) 
                    elif typedict[leaf] == 'int':
                        ROOT.writevectorint(pycpp.intvec,new_tree,entry,leaf) 
                else: 
                    for df in dataframes:

                        length = len(df.iloc[entry][leaf])

                        for i in range(length):

                            temp = df.iloc[entry][leaf]
                            if typedict[leaf] == 'double':
                                pycpp.doublevec.push_back(temp[i]) 
                            elif typedict[leaf] == 'int':
                                pycpp.intvec.push_back(temp[i])                     

                    if typedict[leaf] == 'double':
                        ROOT.writevectordouble(pycpp.doublevec,new_tree,entry,leaf) 
                    elif typedict[leaf] == 'int':
                        ROOT.writevectorint(pycpp.intvec,new_tree,entry,leaf)                

            else:
                pycpp.integer = int(dataframes[0].iloc[entry][leaf])
                ROOT.writeint(pycpp.integer,new_tree,entry,leaf)    

    new_tree.SetBranchStatus('*',1)         
    new_tree.SetEntries(n_entries)

    output.Write("", ROOT.TObject.kOverwrite)
    output.Close()

if __name__=='__main__':
    main()