#!/usr/bin/env python
import numpy as np
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt
from uproot import open as root_open
# plt.interactive(True)

tree = None
bsp = None

def pdg_color(pdgId):
    return {
        11: 'b',
        -11: 'b',
        22: 'g',
    }.get(pdgId, 'k')


def get_roots(sourceSimIdxs):
    roots = []
    for idx, sourceSimIdx in enumerate(sourceSimIdxs):
        if len(sourceSimIdx) == 0:
            roots.append(idx)
    return np.array(roots[:len(roots)//2])


def p(px, py, pz):
    return np.sqrt(px**2 + py**2 + pz**2)


def plot_all_roots(sim_vtxs):
    # ax = plt.gca(projection='3d')
    ax = plt.gca()

    xs = sim_vtxs[b'simvtx_x'] - bsp[b'bsp_x']
    ys = sim_vtxs[b'simvtx_y'] - bsp[b'bsp_y']
    zs = sim_vtxs[b'simvtx_z'] - bsp[b'bsp_z']
    rs = np.sqrt(xs*xs + ys*ys)
    roots = get_roots(sim_vtxs[b'simvtx_sourceSimIdx'])

    bound_z = bsp[b'bsp_sigmaz']*5
    # bound_z = 1
    bound_r = np.sqrt(bsp[b'bsp_sigmax']**2 + bsp[b'bsp_sigmay']**2)
    lumi_roots = []
    for root in roots:
        if abs(zs[root]) < bound_z and np.sqrt(xs[root]**2 + ys[root]**2) < bound_r:
            lumi_roots.append(root)
    # nvtx = len(xs)//2


    # ax.plot(xs[roots], ys[roots], zs[roots], '.')
    # ax.plot(zs[roots], rs[roots], '.')
    ax.plot(zs[lumi_roots], rs[lumi_roots], '.')
    # ax.plot([-bound_z, -bound_z, bound_z, bound_z], [0, bound_r, bound_r, 0], 'r')

    ax.set_xlabel('z')
    ax.set_ylabel('r')
    ax.set_aspect('equal')
    # ax.set_zlabel('z')


def plot_all_vtx(sim_vtxs):
    ax = plt.gca(projection='3d')

    xs = sim_vtxs[b'simvtx_x']
    ys = sim_vtxs[b'simvtx_y']
    zs = sim_vtxs[b'simvtx_z']
    # nvtx = len(xs)//2

    ax.plot(xs, ys, zs, '.')

    ax.set_xlabel('x')
    ax.set_ylabel('y')
    ax.set_zlabel('z')


def plot_event_tree(sim_tracks, sim_vtxs, pv_idx):
    print('='*80)
    ax = plt.gca(projection='3d')
    pdgIds = sim_tracks[b'sim_pdgId']
    decayVtxIdxs = sim_tracks[b'sim_decayVtxIdx']
    px = sim_tracks[b'sim_px']
    py = sim_tracks[b'sim_py']
    pz = sim_tracks[b'sim_pz']

    xs = sim_vtxs[b'simvtx_x']
    ys = sim_vtxs[b'simvtx_y']
    zs = sim_vtxs[b'simvtx_z']
    daughterSimIdxs = sim_vtxs[b'simvtx_daughterSimIdx']

    vtx_idxs = [(0, pv_idx)]
    while vtx_idxs:
        depth, vtx_idx = vtx_idxs.pop()
        print(' '*depth + str(vtx_idx), end=' -> ', flush=True)
        ax.plot([xs[vtx_idx]], [ys[vtx_idx]], [zs[vtx_idx]], 'r.')
        start_x = xs[vtx_idx]
        start_y = ys[vtx_idx]
        start_z = zs[vtx_idx]
        for sim_idx in daughterSimIdxs[vtx_idx]:
            pdgId = pdgIds[sim_idx]
            # print(pdgId)
            # if abs(pdgId) != 11:
            #     continue
            # if pdgId == 22:
            #     continue
            # if pdgId == 22 and p(px[sim_idx], py[sim_idx], pz[sim_idx]) < 1.0:
            #     continue
            for decay_vtx_idx in decayVtxIdxs[sim_idx]:
                end_x = xs[decay_vtx_idx]
                end_y = ys[decay_vtx_idx]
                end_z = zs[decay_vtx_idx]
                # if abs(pdgId) == 11:
                if True:
                    ax.plot([start_x, end_x], [start_y, end_y], [start_z, end_z],
                            pdg_color(pdgId))
                vtx_idxs.append((depth+1, decay_vtx_idx))
                print(str(decay_vtx_idx) + ', ', end='')
        print()

    ax.plot([xs[pv_idx]], [ys[pv_idx]], [zs[pv_idx]], 'k.')
    # ax.set_xlim((-5, 5))
    # ax.set_ylim((-5, 5))
    # ax.set_zlim((-5, 5))
    ax.set_xlabel('x')
    ax.set_ylabel('y')
    ax.set_zlabel('z')


def print_sim_vtxs(sim_vtxs, sim_tracks, sim_pvs, gens, gsf_tracks):
    daughterSimIdxs = sim_vtxs[b'simvtx_daughterSimIdx']
    sourceSimIdxs = sim_vtxs[b'simvtx_sourceSimIdx']
    processTypes = sim_vtxs[b'simvtx_processType']
    xs = sim_vtxs[b'simvtx_x'] - bsp[b'bsp_x']
    ys = sim_vtxs[b'simvtx_y'] - bsp[b'bsp_y']
    zs = sim_vtxs[b'simvtx_z'] - bsp[b'bsp_z']

    parentVtxIdxs = sim_tracks[b'sim_parentVtxIdx']
    decayVtxIdxs = sim_tracks[b'sim_decayVtxIdx']
    sim_pdgIds = sim_tracks[b'sim_pdgId']
    sim_pxs = sim_tracks[b'sim_px']
    sim_pys = sim_tracks[b'sim_py']
    sim_pzs = sim_tracks[b'sim_pz']

    gen_pdgIds = gens[b'gen_pdgId']
    vxs = gens[b'gen_vx'] - bsp[b'bsp_x']
    vys = gens[b'gen_vy'] - bsp[b'bsp_y']
    vzs = gens[b'gen_vz'] - bsp[b'bsp_z']
    gen_pxs = gens[b'gen_px']
    gen_pys = gens[b'gen_py']
    gen_pzs = gens[b'gen_pz']
    gen_prompt = gens[b'gen_isPrompt']
    gen_hadronic = gens[b'gen_isDirectHadronDecayProduct']
    gen_tauic = gens[b'gen_isTauDecayProduct']

    gsf_pdgIds = -11 * gsf_tracks[b'trk_q']
    gsf_vxs = gsf_tracks[b'trk_vtxx'] - bsp[b'bsp_x']
    gsf_vys = gsf_tracks[b'trk_vtxy'] - bsp[b'bsp_y']
    gsf_vzs = gsf_tracks[b'trk_vtxz'] - bsp[b'bsp_z']
    gsf_pxs = gsf_tracks[b'trk_px']
    gsf_pys = gsf_tracks[b'trk_py']
    gsf_pzs = gsf_tracks[b'trk_pz']


    # print('VTX')
    # for (idx, (sourceSimIdx, daughterSimIdx, processType)) in enumerate(zip(sourceSimIdxs, daughterSimIdxs, processTypes)):
    #     if idx in sim_pvs:
    #         print(f'*{idx}|{processType}', sourceSimIdx, daughterSimIdx, sep=" - ")
    #     else:
    #         print(f'{idx}|{processType}', sourceSimIdx, daughterSimIdx, sep=" - ")

    print('GEN')
    for (idx, (px, py, pz, vx, vy, vz, pdgId, prompt, hadronic, tauic)) in enumerate(zip(gen_pxs, gen_pys, gen_pzs, vxs,
                                                                                         vys, vzs, gen_pdgIds,
                                                                                         gen_prompt, gen_hadronic,
                                                                                         gen_tauic)):
        if abs(pdgId) != 11: continue
        p = np.sqrt(px**2 + py**2 + pz**2)
        theta = np.arctan2(np.hypot(px, py), pz)
        phi = np.arctan2(px, py)
        print(f'{idx: 4d}|{pdgId: 3d} - ({vx:8.2f},{vy:8.2f},{vz:8.2f}) ({theta:5.2f},{phi:5.2f}) {p: 8.2f}GeV - ({prompt},{hadronic},{tauic}) ')

    print('SIM')
    for (idx, (px, py, pz, parentVtxIdx, decayVtxIdx, pdgId)) in enumerate(zip(sim_pxs, sim_pys, sim_pzs, parentVtxIdxs, decayVtxIdxs, sim_pdgIds)):
        if abs(pdgId) != 11: continue
        if len(sourceSimIdxs[parentVtxIdx]) > 0: continue  # not from another sim track decay
        vx = xs[parentVtxIdx]
        vy = ys[parentVtxIdx]
        vz = zs[parentVtxIdx]
        p = np.sqrt(px**2 + py**2 + pz**2)
        theta = np.arctan2(np.hypot(px, py), pz)
        phi = np.arctan2(px, py)
        print(f'{idx: 4d}|{pdgId: 3d} - ({vx:8.2f},{vy:8.2f},{vz:8.2f}) ({theta:5.2f},{phi:5.2f}) {p: 8.2f}GeV')

    print('RECO')
    for (idx, (px, py, pz, vx, vy, vz, pdgId)) in enumerate(zip(gsf_pxs, gsf_pys, gsf_pzs, gsf_vxs, gsf_vys, gsf_vzs, gsf_pdgIds)):
        p = np.sqrt(px**2 + py**2 + pz**2)
        theta = np.arctan2(np.hypot(px, py), pz)
        phi = np.arctan2(px, py)
        print(f'{idx: 4d}|{pdgId: 3d} - ({vx:8.2f},{vy:8.2f},{vz:8.2f}) ({theta:5.2f},{phi:5.2f}) {p: 8.2f}GeV')



def plot_event(sim_tracks, sim_vtxs, sim_pvs, gens, gsf_tracks, event_idx):
    print(f"Processing event {event_idx}")
    sim_tracks = {k: v[event_idx] for k, v in sim_tracks.items()}
    sim_vtxs = {k: v[event_idx] for k, v in sim_vtxs.items()}
    gens = {k: v[event_idx] for k, v in gens.items()}
    gsf_tracks = {k: v[event_idx] for k, v in gsf_tracks.items()}
    sim_pvs = sim_pvs[event_idx]

    print_sim_vtxs(sim_vtxs, sim_tracks, sim_pvs, gens, gsf_tracks)
    # print(len(sim_pvs[event_idx]))

    # plt.clf()
    # plot_all_roots(sim_vtxs)
    # plt.show()

    # for pv_idx in get_roots(sim_vtxs[b'simvtx_sourceSimIdx']):
    #     plt.clf()
    #     plot_event_tree(sim_tracks, sim_vtxs, pv_idx)
    #     # plot_all_vtx(sim_vtxs)
    #     plt.show()


def main():
    global tree
    global bsp
    f = root_open('../data/zee.root')
    tree = f['trackingNtuple/tree']
    gsf_tracks = tree.arrays([
        'trk_q',
        'trk_vtxx',
        'trk_vtxy',
        'trk_vtxz',
        'trk_px',
        'trk_py',
        'trk_pz',
    ])
    sim_tracks = tree.arrays([
        'sim_pdgId',
        'sim_parentVtxIdx',
        'sim_decayVtxIdx',
        'sim_pt',
        'sim_eta',
        'sim_px',
        'sim_py',
        'sim_pz',
    ])
    sim_vtxs = tree.arrays([
        'simvtx_x',
        'simvtx_y',
        'simvtx_z',
        'simvtx_sourceSimIdx',
        'simvtx_daughterSimIdx',
        'simvtx_processType',
    ])
    bsp = tree.arrays([
        'bsp_x',
        'bsp_y',
        'bsp_z',
        'bsp_sigmax',
        'bsp_sigmay',
        'bsp_sigmaz',
    ])
    gens = tree.arrays([
        'gen_vx',
        'gen_vy',
        'gen_vz',
        'gen_px',
        'gen_py',
        'gen_pz',
        'gen_pdgId',
        'gen_isTauDecayProduct',
        'gen_isDirectHadronDecayProduct',
        'gen_isPrompt',
    ])
    bsp = {k: v[0] for k, v in bsp.items()}
    sim_pvs = tree.array('simpv_idx')
    for event_idx in range(tree.fEntries):
        plot_event(sim_tracks, sim_vtxs, sim_pvs, gens, gsf_tracks, event_idx)
        if event_idx == 5: break
        # break


if __name__ == '__main__':
    main()