#!/usr/bin/env python
from __future__ import print_function, division
import yaml

SNOMAN_MASS = {
    20: 0.511,
    21: 0.511,
    22: 105.658,
    23: 105.658
}

if __name__ == '__main__':
    import argparse
    import matplotlib.pyplot as plt
    import numpy as np

    parser = argparse.ArgumentParser("plot fit results")
    parser.add_argument("filenames", nargs='+', help="input files")
    args = parser.parse_args()

    for filename in args.filenames:
        print(filename)
        with open(filename) as f:
            data = yaml.load(f.read())

        dx = []
        dy = []
        dz = []
        e = []
        for event in data['data']:
            # get the particle ID
            id = event['mctk'][-1]['id']
            mass = SNOMAN_MASS[id]
            # for some reason it's the *second* track which seems to contain the
            # initial energy
            true_energy = event['mctk'][-1]['energy']
            # The MCTK bank has the particle's total energy (except for neutrons)
            # so we need to convert it into kinetic energy
            ke = true_energy - mass
            energy = event['ev'][0]['fit'][0]['energy']
            e.append(energy-ke)
            true_posx = event['mcvx'][0]['posx']
            posx = event['ev'][0]['fit'][0]['posx']
            dx.append(posx-true_posx)
            true_posy = event['mcvx'][0]['posy']
            posy = event['ev'][0]['fit'][0]['posy']
            dy.append(posy-true_posy)
            true_posz = event['mcvx'][0]['posz']
            posz = event['ev'][0]['fit'][0]['posz']
            dz.append(posz-true_posz)

        print("dT = %.2g +/- %.2g" % (np.mean(e), np.std(e)))
        print("dx = %4.2g +/- %.2g" % (np.mean(dx), np.std(dx)))
        print("dy = %4.2g +/- %.2g" % (np.mean(dy), np.std(dy)))
        print("dz = %4.2g +/- %.2g" % (np.mean(dz), np.std(dz)))

        plt.figure(1)
        plt.hist(e, histtype='step', label=filename)
        plt.xlabel("Kinetic Energy difference (MeV)")
        plt.figure(2)
        plt.hist(dx, histtype='step', label=filename)
        plt.xlabel("X Position difference (cm)")
        plt.figure(3)
        plt.hist(dy, histtype='step', label=filename)
        plt.xlabel("Y Position difference (cm)")
        plt.figure(4)
        plt.hist(dz, histtype='step', label=filename)
        plt.xlabel("Z Position difference (cm)")

    plt.figure(1)
    plt.legend()
    plt.figure(2)
    plt.legend()
    plt.figure(3)
    plt.legend()
    plt.figure(4)
    plt.legend()
    plt.show()