import ROOT
import numpy as np

ROOT.gROOT.ProcessLine(
    """struct Event {
    Int_t trk_len;
    Float_t trk_px[128];
    Float_t trk_py[128];
    Float_t trk_pz[128];
    Int_t trk_imc[128];
    Int_t mc_trk_len;
    Float_t mc_trk_px[128];
    Float_t mc_trk_py[128];
    Float_t mc_trk_pz[128];
    Int_t mc_trk_pid[128];
};""".replace(
        "\n", " "
    )
)


def event_repr(self):
    def r(x, n):
        return "[" + " ".join([f"{x[i]}" for i in range(n)]) + "]"

    return f"""trk_len = {self.trk_len}
trk_px = {r(self.trk_px, self.trk_len)}
trk_py = {r(self.trk_py, self.trk_len)}
trk_pz = {r(self.trk_pz, self.trk_len)}
trk_imc = {r(self.trk_imc, self.trk_len)}
mc_trk_len = {self.mc_trk_len}
mc_trk_px = {r(self.mc_trk_px, self.mc_trk_len)}
mc_trk_py = {r(self.mc_trk_py, self.mc_trk_len)}
mc_trk_pz = {r(self.mc_trk_pz, self.mc_trk_len)}
mc_trk_pid = {r(self.mc_trk_pid, self.mc_trk_len)}
"""


ROOT.Event.__repr__ = event_repr
event = ROOT.Event()

f = ROOT.TFile.Open("example.root", "RECREATE")

tree = ROOT.TTree("event", "")
tree.Branch("trk_len", ROOT.AddressOf(event, "trk_len"), "trk_len/I")
tree.Branch("mc_trk_len", ROOT.AddressOf(event, "mc_trk_len"), "mc_trk_len/I")
tree.Branch("trk_imc", ROOT.AddressOf(event, "trk_imc"), "trk_imc[trk_len]/I")
for prefix in ("", "mc_"):
    for dim in "xyz":
        name = f"{prefix}trk_p{dim}"
        tree.Branch(name, ROOT.AddressOf(event, name), f"{name}[{prefix}trk_len]/F")
tree.Branch(
    "mc_trk_pid", ROOT.AddressOf(event, "mc_trk_pid"), "mc_trk_pid[mc_trk_len]/I"
)

rng = np.random.default_rng()

for i in range(1000):
    event.trk_len = rng.poisson(5)
    assert event.trk_len <= 128
    event.mc_trk_len = int(rng.uniform(0.5, 1.0) * event.trk_len)
    imc = -np.ones(event.trk_len, dtype=int)
    imc[: event.mc_trk_len] = np.arange(event.mc_trk_len)
    rng.shuffle(imc[: event.mc_trk_len])
    for i in range(event.trk_len):
        x = rng.normal(0, 0.5)
        y = rng.normal(0, 0.5)
        z = (rng.lognormal(5, 1) ** 2 - x ** 2 - y ** 2) ** 0.5
        j = imc[i]
        if j >= 0:
            event.mc_trk_px[j] = x
            event.mc_trk_py[j] = y
            event.mc_trk_pz[j] = z
            event.mc_trk_pid[j] = rng.choice((-211, 211, -321, 321, -2212, 2212))
        event.trk_px[i] = x + rng.normal(0, 0.1)
        event.trk_py[i] = y + rng.normal(0, 0.1)
        event.trk_pz[i] = z + rng.normal(0, 0.1)
        event.trk_imc[i] = j
    tree.Fill()

f.Write()
f.Close()