#!/usr/bin/env python
from __future__ import print_function, division
import yaml
import numpy as np
from scipy.stats import iqr

# on retina screens, the default plots are way too small
# by using Qt5 and setting QT_AUTO_SCREEN_SCALE_FACTOR=1
# Qt5 will scale everything using the dpi in ~/.Xresources
import matplotlib
matplotlib.use("Qt5Agg")

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

def plot_hist(x, label=None):
    # determine the bin width using the Freedman Diaconis rule
    # see https://en.wikipedia.org/wiki/Freedman%E2%80%93Diaconis_rule
    h = 2*iqr(x)/len(x)**(1/3)
    n = int((np.max(x)-np.min(x))/h)
    bins = np.linspace(np.min(x),np.max(x),n)
    plt.hist(x, bins=bins, histtype='step', label=label)

def get_stats(x):
    """
    Returns a tuple (mean, error mean, std, error std) for the values in x.

    The formula for the standard error on the standard deviation comes from
    https://stats.stackexchange.com/questions/156518.
    """
    mean = np.mean(x)
    std = np.std(x)
    n = len(x)
    u4 = np.mean((x-mean)**4)
    error = np.sqrt((u4-(n-3)*std**4/(n-1))/n)/(2*std)
    return mean, std/np.sqrt(n), std, error

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 = []
        dT = []
        thetas = []
        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']
            dT.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)
            dirx = event['mctk'][-1]['dirx']
            diry = event['mctk'][-1]['diry']
            dirz = event['mctk'][-1]['dirz']
            true_dir = [dirx,diry,dirz]
            true_dir = np.array(true_dir)/np.linalg.norm(true_dir)
            theta = event['ev'][0]['fit'][0]['theta']
            phi = event['ev'][0]['fit'][0]['phi']
            dir = [np.sin(theta)*np.cos(phi),np.sin(theta)*np.sin(phi),np.cos(theta)]
            dir = np.array(dir)/np.linalg.norm(dir)
            thetas.append(np.degrees(np.arccos(np.dot(true_dir,dir))))

        mean, mean_error, std, std_error = get_stats(dT)
        print("dT      = %.2g +/- %.2g" % (mean, mean_error))
        print("std(dT) = %.2g +/- %.2g" % (std, std_error))
        mean, mean_error, std, std_error = get_stats(dx)
        print("dx      = %4.2g +/- %.2g" % (mean, mean_error))
        print("std(dx) = %4.2g +/- %.2g" % (std, std_error))
        mean, mean_error, std, std_error = get_stats(dy)
        print("dy      = %4.2g +/- %.2g" % (mean, mean_error))
        print("std(dy) = %4.2g +/- %.2g" % (std, std_error))
        mean, mean_error, std, std_error = get_stats(dz)
        print("dz      = %4.2g +/- %.2g" % (mean, mean_error))
        print("std(dz) = %4.2g +/- %.2g" % (std, std_error))
        mean, mean_error, std, std_error = get_stats(thetas)
        print("std(theta) = %4.2g +/- %.2g" % (std, std_error))

        plt.figure(1)
        plot_hist(dT, label=filename)
        plt.xlabel("Kinetic Energy difference (MeV)")
        plt.figure(2)
        plot_hist(dx, label=filename)
        plt.xlabel("X Position difference (cm)")
        plt.figure(3)
        plot_hist(dy, label=filename)
        plt.xlabel("Y Position difference (cm)")
        plt.figure(4)
        plot_hist(dz, label=filename)
        plt.xlabel("Z Position difference (cm)")
        plt.figure(5)
        plot_hist(thetas, label=filename)
        plt.xlabel(r"$\theta$ (deg)")

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