import ROOT
import numpy as np
from pathlib import Path
from tqdm import tqdm
import argparse

parser = argparse.ArgumentParser(description='Process some integers.')
parser.add_argument('nentries', type=float, default=1e6, help='Number of entries to process')
args = parser.parse_args()


here = Path(__file__).parent.resolve()
file_path = here / Path('../data/src_py.root')
file_path.parent.mkdir(parents=True, exist_ok=True)

print(f'Creating file: {file_path}')

file = ROOT.TFile.Open(str(file_path),'recreate')
tree = ROOT.TTree('DDTree','DDTree')

nbranches = 2000
value_length = 20
nentries = int(args.nentries)


branches = {}
branch_values = {}

for i in range(nbranches):
    branch_values[f'value{i}'] = value = ROOT.std.vector('double')()
    branches[f'value{i}'] = tree.Branch(f'value{i}', value)

for i in tqdm(range(nentries)):
    try:

        tree.GetEntry(i)

        for j in range(nbranches):
            value = branch_values[f'value{j}']
            value.clear()
            [value.push_back(result) for result in np.random.uniform(-999,999,value_length)]
        
        tree.Fill()
        
    except KeyboardInterrupt:
        # Early stop
        file.Write()
        file.Close()
        exit(0)

file.Write()
file.Close()