from collections import defaultdict
from uproot import open as root_open


def main():
    f_old = root_open('trackingNtuple_old_default.root')['trackingNtuple/tree']
    f_new = root_open('trackingNtuple_new_default.root')['trackingNtuple/tree']

    keys = [b'see_sclIdx', b'see_trkIdx',
            b'scl_e', b'scl_px', b'scl_py', b'scl_pz', b'scl_hoe',
            b'trk_q']
    arrs_old = f_old.arrays(keys)
    arrs_new = f_new.arrays(keys)

    def dump_event(event, name):
        print('-'*20 + f'{name:10}' + '-'*20)
        # print(event[b'scl_hoe'] <= 0.15)

        def get_cols(*strs):
            en = enumerate(zip(*[event[s] for s in strs]))
            return en

        print('Seed Info')
        for idx, (sclIdx, trkIdx) in get_cols('see_sclIdx', 'see_trkIdx'):
            if sclIdx < 0: continue
            if event['scl_hoe'][sclIdx] > 0.15: continue

            trk_q = '-'
            if trkIdx>=0:
                trk_q = str(event["trk_q"][trkIdx])

            print(f'{idx:3d}) {sclIdx:10d} {trk_q:10s}')
        # print(event[b'see_sclIdx'])

    def dump_scl(event):
        def get_cols(*strs):
            en = enumerate(zip(*[event[s] for s in strs]))
            return en

        print('Supercluster Info')
        for idx, (e, px, py, pz, hoe) in get_cols('scl_e', 'scl_px', 'scl_py', 'scl_pz', 'scl_hoe'):
            print(f'{idx:3d}) {hoe:10.2f} {e:10.2f}')


    def seed_summary(event_old, event_new):
        def get_cols(event, *strs):
            en = enumerate(zip(*[event[s] for s in strs]))
            return en

        counts_old = defaultdict(int)
        counts_new = defaultdict(int)
        # print('Supercluster Info')
        # for idx, (e, px, py, pz, hoe) in get_cols('scl_e', 'scl_px', 'scl_py', 'scl_pz', 'scl_hoe'):
        #     print(f'{idx:3d}) {hoe:10.2f} {e:10.2f}')
        print('Seed Info')
        for _, (sclIdx,) in get_cols(event_old, 'see_sclIdx'):
            if sclIdx >= 0:
                # if event_old['scl_hoe'][sclIdx] > 0.15: continue
                counts_old[sclIdx] += 1
        for _, (sclIdx,) in get_cols(event_new, 'see_sclIdx'):
            if sclIdx >= 0:
                # if event_new['scl_hoe'][sclIdx] > 0.15: continue
                counts_new[sclIdx] += 1

        for idx, (e, px, py, pz, hoe) in get_cols(event_old, 'scl_e', 'scl_px', 'scl_py', 'scl_pz', 'scl_hoe'):
            if hoe > 0.15: continue
            print(f'{idx:3d}) {hoe:10.2f} {e:10.2f} {counts_old[idx]:10d} {counts_new[idx]:10d}')



    nevt = len(arrs_old[keys[0]])
    nevt = 5
    for eIdx in range(nevt):
        print(f'NEW EVENT: {eIdx}')
        old = {key.decode(): arrs_old[key][eIdx] for key in keys}
        new = {key.decode(): arrs_new[key][eIdx] for key in keys}
        # dump_scl(old)
        # dump_event(old, 'OLD')
        # dump_event(new, 'NEW')
        seed_summary(old, new)
        # print(new[b'see_sclIdx'])





if __name__ == '__main__':
    main()