#!/usr/bin/env python
# Copyright (c) 2019, Anthony Latorre <tlatorre at uchicago>
#
# This program is free software: you can redistribute it and/or modify it
# under the terms of the GNU General Public License as published by the Free
# Software Foundation, either version 3 of the License, or (at your option)
# any later version.
#
# This program is distributed in the hope that it will be useful, but WITHOUT
# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
# FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for
# more details.
#
# You should have received a copy of the GNU General Public License along with
# this program. If not, see <https://www.gnu.org/licenses/>.

from __future__ import print_function, division

if __name__ == '__main__':
    import argparse
    import numpy as np
    import h5py
    import pandas as pd
    import itertools
    from sddm import IDP_E_MINUS, IDP_MU_MINUS, SNOMAN_MASS
    from sddm.plot import plot_hist, plot_legend, get_stats, despine, iqr_std_err, iqr_std, quantile_error, q90_err, q90, median, median_err, std_err
    from sddm import setup_matplotlib

    parser = argparse.ArgumentParser("plot fit results")
    parser.add_argument("filenames", nargs='+', help="input files")
    parser.add_argument("-o", "--output", default=None, help="output file")
    parser.add_argument("--save", action="store_true", default=False, help="save plots")
    args = parser.parse_args()

    setup_matplotlib(args.save)

    import matplotlib.pyplot as plt

    # Read in all the data.
    #
    # Note: We have to add the filename as a column to each dataset since this
    # script is used to analyze MC data which all has the same run number.
    ev = pd.concat([pd.read_hdf(filename, "ev").assign(filename=filename) for filename in args.filenames],ignore_index=True)
    fits = pd.concat([pd.read_hdf(filename, "fits").assign(filename=filename) for filename in args.filenames],ignore_index=True)
    mcgn = pd.concat([pd.read_hdf(filename, "mcgn").assign(filename=filename) for filename in args.filenames],ignore_index=True)

    # get rid of 2nd events like Michel electrons
    ev = ev.sort_values(['run','gtid']).groupby(['filename','evn'],as_index=False).nth(0)

    # Now, we merge all three datasets together to produce a single
    # dataframe. To do so, we join the ev dataframe with the mcgn frame
    # on the evn column, and then join with the fits on the run and
    # gtid columns.
    #
    # At the end we will have a single dataframe with one row for each
    # fit, i.e. it will look like:
    #
    # >>> data
    #   run   gtid nhit, ... mcgn_x, mcgn_y, mcgn_z, ..., fit_id1, fit_x, fit_y, fit_z, ...
    #
    # Before merging, we prefix the primary seed track table with mcgn_
    # and the fit table with fit_ just to make things easier.

    # Prefix track and fit frames
    mcgn = mcgn.add_prefix("mcgn_")
    fits = fits.add_prefix("fit_")

    # merge ev and mcgn on evn
    data = ev.merge(mcgn,left_on=['filename','evn'],right_on=['mcgn_filename','mcgn_evn'])
    # merge data and fits on run and gtid
    data = data.merge(fits,left_on=['filename','run','gtid'],right_on=['fit_filename','fit_run','fit_gtid'])

    # For this script, we only want the single particle fit results
    data = data[(data.fit_id2 == 0) & (data.fit_id3 == 0)]

    # Select only the best fit for a given run, gtid, and particle
    # combo
    data = data.sort_values('fit_fmin').groupby(['filename','run','gtid','fit_id1','fit_id2','fit_id3'],as_index=False).nth(0).reset_index(level=0,drop=True)

    # calculate true kinetic energy
    mass = [SNOMAN_MASS[id] for id in data['mcgn_id'].values]
    data['T'] = data['mcgn_energy'].values - mass
    data['dx'] = data['fit_x'].values - data['mcgn_x'].values
    data['dy'] = data['fit_y'].values - data['mcgn_y'].values
    data['dz'] = data['fit_z'].values - data['mcgn_z'].values
    data['dT'] = data['fit_energy1'].values - data['T'].values

    true_dir = np.dstack((data['mcgn_dirx'],data['mcgn_diry'],data['mcgn_dirz'])).squeeze()
    dir = np.dstack((np.sin(data['fit_theta1'])*np.cos(data['fit_phi1']),
                     np.sin(data['fit_theta1'])*np.sin(data['fit_phi1']),
                     np.cos(data['fit_theta1']))).squeeze()

    data['theta'] = np.degrees(np.arccos((true_dir*dir).sum(axis=-1)))


    # only select fits which have at least 2 fits
    data = data.groupby(['filename','run','gtid']).filter(lambda x: len(x) > 1)
    data_true = data[data['fit_id1'] == data['mcgn_id']]
    data_e = data[data['fit_id1'] == IDP_E_MINUS]
    data_mu = data[data['fit_id1'] == IDP_MU_MINUS]

    data_true = data_true.set_index(['filename','run','gtid'])
    data_e = data_e.set_index(['filename','run','gtid'])
    data_mu = data_mu.set_index(['filename','run','gtid'])

    data_true['ratio'] = data_mu['fit_fmin']-data_e['fit_fmin']
    data_true['te'] = data_e['fit_time']
    data_true['tm'] = data_mu['fit_time']
    data_true['Te'] = data_e['fit_energy1']

    # 100 bins between 50 MeV and 1 GeV
    bins = np.arange(50,1000,100)

    T = (bins[1:] + bins[:-1])/2

    markers = itertools.cycle(('o', 'v')) 

    fig3, ax3 = plt.subplots(3,1,num=3,sharex=True)
    fig4, ax4 = plt.subplots(3,1,num=4,sharex=True)

    output = pd.DataFrame({'T':T})

    for id in [IDP_E_MINUS, IDP_MU_MINUS]:
        events = data_true[data_true['mcgn_id'] == id]

        pd_bins = pd.cut(events['T'],bins)

        dT = events.groupby(pd_bins)['dT'].agg(['mean','sem','std',std_err,median,median_err,iqr_std,iqr_std_err])
        dx = events.groupby(pd_bins)['dx'].agg(['mean','sem','std',std_err,median,median_err,iqr_std,iqr_std_err])
        dy = events.groupby(pd_bins)['dy'].agg(['mean','sem','std',std_err,median,median_err,iqr_std,iqr_std_err])
        dz = events.groupby(pd_bins)['dz'].agg(['mean','sem','std',std_err,median,median_err,iqr_std,iqr_std_err])
        theta = events.groupby(pd_bins)['theta'].agg(['mean','sem','std',std_err,median,median_err,iqr_std,iqr_std_err,q90,q90_err])

        label = 'Muon' if id == IDP_MU_MINUS else 'Electron'

        marker = markers.next()

        plt.figure(1)
        plt.errorbar(T,dT['median']*100/T,yerr=dT['median_err']*100/T,fmt=marker,label=label)

        plt.figure(2)
        plt.errorbar(T,dT['iqr_std']*100/T,yerr=dT['iqr_std_err']*100/T,fmt=marker,label=label)

        ax3[0].errorbar(T,dx['median'],yerr=dx['median_err'],fmt=marker,label=label)
        ax3[1].errorbar(T,dy['median'],yerr=dy['median_err'],fmt=marker,label=label)
        ax3[2].errorbar(T,dz['median'],yerr=dz['median_err'],fmt=marker,label=label)

        if id == IDP_E_MINUS:
            output['e_bias'] = (dT['median']/T).values
        else:
            output['u_bias'] = (dT['median']/T).values

        ax4[0].errorbar(T,dx['iqr_std'],yerr=dx['iqr_std_err'],fmt=marker,label=label)
        ax4[1].errorbar(T,dy['iqr_std'],yerr=dy['iqr_std_err'],fmt=marker,label=label)
        ax4[2].errorbar(T,dz['iqr_std'],yerr=dz['iqr_std_err'],fmt=marker,label=label)

        plt.figure(5)
        plt.errorbar(T,theta['std'],yerr=theta['std_err'],fmt=marker,label=label)

        plt.figure(6)
        plt.scatter(events['Te'],events['ratio'],marker=marker,label=label)

    if args.output:
        with h5py.File(args.output,"w") as f:
            f.create_dataset('energy_bias',data=output.to_records())

    fig = plt.figure(1)
    despine(fig,trim=True)
    plt.xlabel("Kinetic Energy (MeV)")
    plt.ylabel(r"Energy bias (\%)")
    plt.legend()
    plt.tight_layout()

    fig = plt.figure(2)
    despine(fig,trim=True)
    plt.xlabel("Kinetic Energy (MeV)")
    plt.ylabel(r"Energy resolution (\%)")
    plt.legend()
    plt.tight_layout()

    ax3[0].set_ylabel("X")
    ax3[0].set_ylim((-5,5))
    ax3[1].set_ylabel("Y")
    ax3[1].set_ylim((-5,5))
    ax3[2].set_xlabel("Kinetic Energy (MeV)")
    ax3[2].set_ylabel("Z")
    ax3[2].set_ylim((-5,5))
    despine(ax=ax3[0],trim=True)
    despine(ax=ax3[1],trim=True)
    despine(ax=ax3[2],trim=True)
    h,l = ax3[0].get_legend_handles_labels()
    fig3.legend(h,l,loc='upper right')
    fig3.subplots_adjust(right=0.75)
    fig3.tight_layout()
    fig3.subplots_adjust(top=0.9)

    ax4[0].set_ylabel("X")
    ax4[0].set_ylim((0,ax4[0].get_ylim()[1]))
    ax4[1].set_ylabel("Y")
    ax4[1].set_ylim((0,ax4[1].get_ylim()[1]))
    ax4[2].set_xlabel("Kinetic Energy (MeV)")
    ax4[2].set_ylabel("Z")
    ax4[2].set_ylim((0,ax4[2].get_ylim()[1]))
    despine(ax=ax4[0],trim=True)
    despine(ax=ax4[1],trim=True)
    despine(ax=ax4[2],trim=True)
    h,l = ax4[0].get_legend_handles_labels()
    fig4.legend(h,l,loc='upper right')
    fig4.subplots_adjust(right=0.75)
    fig4.tight_layout()
    fig4.subplots_adjust(top=0.9)

    fig = plt.figure(5)
    despine(fig,trim=True)
    plt.xlabel("Kinetic Energy (MeV)")
    plt.ylabel("Angular resolution (deg)")
    plt.ylim((0,plt.gca().get_ylim()[1]))
    plt.legend()
    plt.tight_layout()

    fig = plt.figure(6)
    plt.xticks(range(0,1250,100))
    plt.hlines(0,0,1200,color='k',linestyles='--',alpha=0.5)
    despine(fig,trim=True)
    plt.xlabel("Reconstructed Electron Energy (MeV)")
    plt.ylabel(r"Log Likelihood Ratio (e/$\mu$)")
    plt.legend()
    plt.tight_layout()

    if args.save:
        fig = plt.figure(1)
        plt.savefig("energy_bias.pdf")
        plt.savefig("energy_bias.eps")
        fig = plt.figure(2)
        plt.savefig("energy_resolution.pdf")
        plt.savefig("energy_resolution.eps")
        fig = plt.figure(3)
        plt.savefig("position_bias.pdf")
        plt.savefig("position_bias.eps")
        fig = plt.figure(4)
        plt.savefig("position_resolution.pdf")
        plt.savefig("position_resolution.eps")
        fig = plt.figure(5)
        plt.savefig("angular_resolution.pdf")
        plt.savefig("angular_resolution.eps")
        fig = plt.figure(6)
        plt.savefig("likelihood_ratio.pdf")
        plt.savefig("likelihood_ratio.eps")
    else:
        plt.figure(1)
        plt.title("Energy Bias")
        plt.figure(2)
        plt.title("Energy Resolution")
        plt.figure(3)
        fig3.suptitle("Position Bias (cm)")
        plt.figure(4)
        fig4.suptitle("Position Resolution (cm)")
        plt.figure(5)
        plt.title("Angular Resolution")
        plt.figure(6)
        plt.title("Log Likelihood Ratio vs Reconstructed Electron Energy")

        plt.show()
>239</a>
<a id='n240' href='#n240'>240</a>
<a id='n241' href='#n241'>241</a>
<a id='n242' href='#n242'>242</a>
<a id='n243' href='#n243'>243</a>
<a id='n244' href='#n244'>244</a>
<a id='n245' href='#n245'>245</a>
<a id='n246' href='#n246'>246</a>
<a id='n247' href='#n247'>247</a>
<a id='n248' href='#n248'>248</a>
<a id='n249' href='#n249'>249</a>
<a id='n250' href='#n250'>250</a>
<a id='n251' href='#n251'>251</a>
<a id='n252' href='#n252'>252</a>
<a id='n253' href='#n253'>253</a>
<a id='n254' href='#n254'>254</a>
<a id='n255' href='#n255'>255</a>
<a id='n256' href='#n256'>256</a>
<a id='n257' href='#n257'>257</a>
<a id='n258' href='#n258'>258</a>
<a id='n259' href='#n259'>259</a>
<a id='n260' href='#n260'>260</a>
<a id='n261' href='#n261'>261</a>
<a id='n262' href='#n262'>262</a>
<a id='n263' href='#n263'>263</a>
<a id='n264' href='#n264'>264</a>
<a id='n265' href='#n265'>265</a>
<a id='n266' href='#n266'>266</a>
<a id='n267' href='#n267'>267</a>
<a id='n268' href='#n268'>268</a>
<a id='n269' href='#n269'>269</a>
<a id='n270' href='#n270'>270</a>
<a id='n271' href='#n271'>271</a>
<a id='n272' href='#n272'>272</a>
<a id='n273' href='#n273'>273</a>
<a id='n274' href='#n274'>274</a>
<a id='n275' href='#n275'>275</a>
<a id='n276' href='#n276'>276</a>
<a id='n277' href='#n277'>277</a>
<a id='n278' href='#n278'>278</a>
<a id='n279' href='#n279'>279</a>
<a id='n280' href='#n280'>280</a>
<a id='n281' href='#n281'>281</a>
<a id='n282' href='#n282'>282</a>
<a id='n283' href='#n283'>283</a>
<a id='n284' href='#n284'>284</a>
<a id='n285' href='#n285'>285</a>
<a id='n286' href='#n286'>286</a>
<a id='n287' href='#n287'>287</a>
<a id='n288' href='#n288'>288</a>
<a id='n289' href='#n289'>289</a>
<a id='n290' href='#n290'>290</a>
<a id='n291' href='#n291'>291</a>
<a id='n292' href='#n292'>292</a>
<a id='n293' href='#n293'>293</a>
<a id='n294' href='#n294'>294</a>
<a id='n295' href='#n295'>295</a>
<a id='n296' href='#n296'>296</a>
<a id='n297' href='#n297'>297</a>
<a id='n298' href='#n298'>298</a>
<a id='n299' href='#n299'>299</a>
<a id='n300' href='#n300'>300</a>
<a id='n301' href='#n301'>301</a>
<a id='n302' href='#n302'>302</a>
<a id='n303' href='#n303'>303</a>
<a id='n304' href='#n304'>304</a>
<a id='n305' href='#n305'>305</a>
<a id='n306' href='#n306'>306</a>
<a id='n307' href='#n307'>307</a>
<a id='n308' href='#n308'>308</a>
<a id='n309' href='#n309'>309</a>
<a id='n310' href='#n310'>310</a>
<a id='n311' href='#n311'>311</a>
<a id='n312' href='#n312'>312</a>
<a id='n313' href='#n313'>313</a>
<a id='n314' href='#n314'>314</a>
<a id='n315' href='#n315'>315</a>
<a id='n316' href='#n316'>316</a>
<a id='n317' href='#n317'>317</a>
<a id='n318' href='#n318'>318</a>
<a id='n319' href='#n319'>319</a>
<a id='n320' href='#n320'>320</a>
<a id='n321' href='#n321'>321</a>
<a id='n322' href='#n322'>322</a>
<a id='n323' href='#n323'>323</a>
<a id='n324' href='#n324'>324</a>
<a id='n325' href='#n325'>325</a>
<a id='n326' href='#n326'>326</a>
<a id='n327' href='#n327'>327</a>
<a id='n328' href='#n328'>328</a>
<a id='n329' href='#n329'>329</a>
<a id='n330' href='#n330'>330</a>
<a id='n331' href='#n331'>331</a>
<a id='n332' href='#n332'>332</a>
<a id='n333' href='#n333'>333</a>
<a id='n334' href='#n334'>334</a>
<a id='n335' href='#n335'>335</a>
<a id='n336' href='#n336'>336</a>
<a id='n337' href='#n337'>337</a>
<a id='n338' href='#n338'>338</a>
<a id='n339' href='#n339'>339</a>
<a id='n340' href='#n340'>340</a>
<a id='n341' href='#n341'>341</a>
<a id='n342' href='#n342'>342</a>
<a id='n343' href='#n343'>343</a>
<a id='n344' href='#n344'>344</a>
<a id='n345' href='#n345'>345</a>
<a id='n346' href='#n346'>346</a>
<a id='n347' href='#n347'>347</a>
<a id='n348' href='#n348'>348</a>
<a id='n349' href='#n349'>349</a>
<a id='n350' href='#n350'>350</a>
<a id='n351' href='#n351'>351</a>
<a id='n352' href='#n352'>352</a>
<a id='n353' href='#n353'>353</a>
<a id='n354' href='#n354'>354</a>
<a id='n355' href='#n355'>355</a>
<a id='n356' href='#n356'>356</a>
<a id='n357' href='#n357'>357</a>
<a id='n358' href='#n358'>358</a>
<a id='n359' href='#n359'>359</a>
<a id='n360' href='#n360'>360</a>
<a id='n361' href='#n361'>361</a>
<a id='n362' href='#n362'>362</a>
<a id='n363' href='#n363'>363</a>
<a id='n364' href='#n364'>364</a>
<a id='n365' href='#n365'>365</a>
<a id='n366' href='#n366'>366</a>
<a id='n367' href='#n367'>367</a>
</pre></td>
<td class='lines'><pre><code>