import uproot
import numpy as np
import sys

print (sys.argv[1])
file_to_convert = sys.argv[1]

def convert_data(filename="BMADStored.root"):
    print(f"Creating data file {filename}...")

    # Set up output file and tree
    with uproot.recreate(filename) as f:
        tree = {
            "x": [],
            "xp": [],
            "y": [],
            "yp": [],
            "t": [],
            "p": [],
            "phi": [],
        }
              
        # Setup variables for David R's input files
        muon_id = 0
        in_init_x, in_init_xp, in_init_y, in_init_yp, in_init_p, in_init_t = 0.0, 0.0, 0.0, 0.0, 0.0, 0.0
        in_init_sx, in_init_sy, in_init_sz, in_init_phi = 0.0, 0.0, 0.0, 0.0
        in_inj_x, in_init_z, in_inj_y, in_inj_yp, in_init_dp, in_inj_t = 0.0, 0.0, 0.0, 0.0, 0.0, 0.0
        in_inj_sx, location, in_inj_sz, in_inj_phi = 0.0, 0.0, 0.0, 0.0

        # Loop over stored muons
        with open(file_to_convert, "r") as in_file:
            next(in_file)  # Skip two header lines
            next(in_file)

            # Read file and write to output tree
            for line in in_file:
                data = line.split()
                muon_id = int(data[0])
                in_inj_x, in_inj_xp, in_inj_y, in_inj_yp = map(float, data[3:7])
                in_inj_p,in_inj_t = map(float, data[8:10])
                in_inj_sx, in_inj_sy, in_inj_sz, in_inj_phi = map(float, data[11:15])

                # Injected parameters
                x = in_inj_x * 1000  # Convert to mm
                xp = in_inj_xp
                y = in_inj_y * 1000  # Convert to mm
                yp = in_inj_yp
                t = in_inj_t * 1e9  # Convert to ns
                p = in_inj_p
                phi = in_inj_phi

                # Append to tree
                tree["x"].append(x)
                tree["xp"].append(xp)
                tree["y"].append(y)
                tree["yp"].append(yp)
                tree["t"].append(t)
                tree["p"].append(p)
                tree["phi"].append(phi)

        # Write the tree to the ROOT file
        f["t"] = tree

# Call the function
convert_data()