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'}

class tbmon_merge:

    def __init__(self):

        self.forcedentries = 0
        self.outputname = ''
        self.runnumber = '000000'

    def main(self):

        parser = argparse.ArgumentParser(description='Script for merging track collections for TBMon2 analysis')
        parser.add_argument('-f', dest='files', action='append')
        parser.add_argument('-o', dest='outputname', action='store', default='')
        parser.add_argument('-n', dest='forcedentries', action='store', default=-99)
        args = parser.parse_args()
        self.forcedentries = int(args.forcedentries)
        self.outputname = args.outputname
        try:
            self.runnumber = args.files[0].split('tbtrack')[1].split('.root')[0].split('_')[0]
        except:
            print('Run number could not be found')
        self.mergetrees(args.files)

    def mergetrees(self, filenames):

        gROOT.LoadMacro("writevectordouble.C")

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

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

    def writeversiontree(self, version):

        if self.outputname == '':
            outputname = 'tbtrack' + self.runnumber + '.root'
        else:
            outputname = self.outputname + '.root'   
        output = TFile(outputname,"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(self, 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

        self.writetree(dataframes,leaves,treename)     

    def writetree(self, dataframes, leaves, treename): 

        if self.outputname == '':
            outputname = 'tbtrack' + self.runnumber + '.root'
        else:
            outputname = self.outputname + '.root' 
        output = TFile(outputname,"update") 
        new_tree = TTree(treename, treename)
        temp_df = pd.DataFrame(dtype=object)
        n_dataframes = len(dataframes) 
        if self.forcedentries == -99:
            n_entries = dataframes[0].shape[0]   
        else:
            n_entries = self.forcedentries    

        pycpp = ROOT.pystruct()

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

        for entry in progressbar(range(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)
                    else:
                        new_tree.SetBranchStatus(leaf,1)    

                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__':
    tbmon_merge().main()