#!/usr/bin/env python
import numpy as np
import matplotlib.pyplot as plt

from filval.result_set import ResultSet
from filval.histogram_utils import hist, hist_add, hist_normalize, hist_scale, hist2d
from filval.plotter import (decl_plot, render_plots, hist_plot, hist2d_plot,
                            hist_plot_stack, Plot, generate_dashboard)


def center_text(x, y, txt, **kwargs):
    plt.text(x, y, txt,
             horizontalalignment='center', verticalalignment='center',
             transform=plt.gca().transAxes, **kwargs)


@decl_plot
def plot_seed_eff(old_seeds, new_seeds):
    r"""## ECAL-Driven Seeding Efficiency

    The proportion of gen-level electrons originating in the luminous region that have
    an associated Seed, matched via rechit-simhit associations in the pixel detector. Cuts are on simtrack quantities.
    """
    ax_pt = plt.subplot(221)
    ax_eta = plt.subplot(222)
    ax_phi = plt.subplot(223)


    errors = True
    plt.sca(ax_pt)
    hist_plot(hist(new_seeds.seed_eff_v_pt), include_errors=errors, label='New')
    hist_plot(hist(old_seeds.seed_eff_v_pt), include_errors=errors, label='Old')
    center_text(0.5, 0.3, r'$|\eta|<2.4$')
    plt.xlabel(r"Sim-Track $p_T$")
    plt.ylim((0, 1.1))

    plt.sca(ax_eta)
    hist_plot(hist(new_seeds.seed_eff_v_eta), include_errors=errors, label='New')
    hist_plot(hist(old_seeds.seed_eff_v_eta), include_errors=errors, label='Old')
    center_text(0.5, 0.3, r'$p_T>20$')
    plt.xlabel(r"Sim-Track $\eta$")
    plt.ylim((0, 1.1))
    plt.legend()

    plt.sca(ax_phi)
    hist_plot(hist(new_seeds.seed_eff_v_phi), include_errors=errors, label='New')
    hist_plot(hist(old_seeds.seed_eff_v_phi), include_errors=errors, label='Old')
    center_text(0.5, 0.3, r'$p_T>20$ and $|\eta|<2.4$')
    plt.xlabel(r"Sim-Track $\phi$")
    plt.ylim((0, 1.1))


@decl_plot
def plot_track_eff(old_seeds, new_seeds):
    r"""## GSF Tracking Efficiency

    The proportion of electrons origination in the luminous region from the that have
    an associated GSF track. Cuts are on simtrack quantities.
    """
    ax_pt = plt.subplot(221)
    ax_eta = plt.subplot(222)
    ax_phi = plt.subplot(223)

    errors = True
    plt.sca(ax_pt)
    hist_plot(hist(new_seeds.track_eff_v_pt), include_errors=errors, label='New')
    hist_plot(hist(old_seeds.track_eff_v_pt), include_errors=errors, label='Old')
    center_text(0.5, 0.3, r'$|\eta|<2.4$')
    plt.xlabel(r"Sim-Track $p_T$")
    plt.ylim((0, 1.1))

    plt.sca(ax_eta)
    hist_plot(hist(new_seeds.track_eff_v_eta), include_errors=errors, label='New')
    hist_plot(hist(old_seeds.track_eff_v_eta), include_errors=errors, label='Old')
    center_text(0.5, 0.3, r'$p_T>20$')
    plt.xlabel(r"Sim-Track $\eta$")
    plt.ylim((0, 1.1))
    plt.legend()

    plt.sca(ax_phi)
    hist_plot(hist(new_seeds.track_eff_v_phi), include_errors=errors, label='New')
    hist_plot(hist(old_seeds.track_eff_v_phi), include_errors=errors, label='Old')
    center_text(0.5, 0.3, r'$p_T>20$ and $|\eta|<2.4$')
    plt.xlabel(r"Sim-Track $\phi$")
    plt.ylim((0, 1.1))


@decl_plot
def plot_seed_purity(old_seeds, new_seeds, ext=""):
    r"""## ECAL-Driven Seed Purity

    The proportion of ECAL-driven seeds that have a matched gen-level electron originating in
    the luminous region. Cuts are on seed quantities.
    """
    ax_pt = plt.subplot(221)
    ax_eta = plt.subplot(222)
    ax_phi = plt.subplot(223)

    def get_hist(rs, base_name):
        return hist(getattr(rs, base_name+ext))

    errors = True
    plt.sca(ax_pt)
    hist_plot(get_hist(new_seeds, "seed_pur_v_pt"), include_errors=errors, label='New')
    hist_plot(get_hist(old_seeds, "seed_pur_v_pt"), include_errors=errors, label='Old')
    center_text(0.5, 0.3, r'$|\eta|<2.4$')
    plt.xlabel(r"Seed $p_T$")
    if not ext:
        plt.ylim((0, 1.1))

    plt.sca(ax_eta)
    hist_plot(get_hist(new_seeds, "seed_pur_v_eta"), include_errors=errors, label='New')
    hist_plot(get_hist(old_seeds, "seed_pur_v_eta"), include_errors=errors, label='Old')
    center_text(0.5, 0.3, r'$p_T>20$')
    plt.xlabel(r"Seed $\eta$")
    if not ext:
        plt.ylim((0, 1.1))
    plt.legend()

    plt.sca(ax_phi)
    hist_plot(get_hist(new_seeds, "seed_pur_v_phi"), include_errors=errors, label='New')
    hist_plot(get_hist(old_seeds, "seed_pur_v_phi"), include_errors=errors, label='Old')
    center_text(0.5, 0.3, r'$p_T>20$ and $|\eta|<2.4$')
    plt.xlabel(r"Seed $\phi$")
    if not ext:
        plt.ylim((0, 1.1))


@decl_plot
def plot_track_purity(old_seeds, new_seeds, ext=""):
    r"""## GSF Track Purity

    The proportion of GSF-tracks w\ ECAL-driven seeds that have a matched gen-level electron originating in
    the luminous region. Cuts are on GSF track quantities.
    """
    ax_pt = plt.subplot(221)
    ax_eta = plt.subplot(222)
    ax_phi = plt.subplot(223)

    def get_hist(rs, base_name):
        return hist(getattr(rs, base_name+ext))

    errors = True
    plt.sca(ax_pt)
    hist_plot(get_hist(new_seeds, "track_pur_v_pt"), include_errors=errors, label='New')
    hist_plot(get_hist(old_seeds, "track_pur_v_pt"), include_errors=errors, label='Old')
    center_text(0.5, 0.3, r'$|\eta|<2.4$')
    plt.xlabel(r"GSF-Track $p_T$")
    if not ext:
        plt.ylim((0, 1.1))

    plt.sca(ax_eta)
    hist_plot(get_hist(new_seeds, "track_pur_v_eta"), include_errors=errors, label='New')
    hist_plot(get_hist(old_seeds, "track_pur_v_eta"), include_errors=errors, label='Old')
    center_text(0.5, 0.3, r'$p_T>20$')
    plt.xlabel(r"GSF-Track $\eta$")
    if not ext:
        plt.ylim((0, 1.1))
    plt.legend()

    plt.sca(ax_phi)
    hist_plot(get_hist(new_seeds, "track_pur_v_phi"), include_errors=errors, label='New')
    hist_plot(get_hist(old_seeds, "track_pur_v_phi"), include_errors=errors, label='Old')
    center_text(0.5, 0.3, r'$p_T>20$ and $|\eta|<2.4$')
    plt.xlabel(r"GSF-Track $\phi$")
    if not ext:
        plt.ylim((0, 1.1))


@decl_plot
def plot_nhits(old_seeds, new_seeds):
    ax_sim = plt.subplot(211)
    ax_rec_new = plt.subplot(223)
    ax_rec_old = plt.subplot(224)

    def _plot(histo):
        hit_hist2d = hist2d(histo)
        hist2d_plot(hit_hist2d)
        weights, _, xs, ys = hit_hist2d
        hits = (ys[1:, :-1] + ys[:-1, :-1])/2
        etas = (xs[0, :-1] + xs[0, 1:])/2
        avg_hits = np.ma.average(hits, axis=0, weights=weights)
        plt.plot(etas, avg_hits, 'wo', label=r'Average in $\eta$ bin')
        plt.xlabel(r'Track $\eta$')
        plt.ylabel(r'\# of pixel layers with Hit')

    plt.sca(ax_sim)
    _plot(new_seeds.simpixlay_v_eta)
    plt.legend()
    plt.colorbar()
    center_text(0.5, 0.1, "Sim Tracks", bbox=dict(facecolor='white', alpha=0.5))
    plt.sca(ax_rec_old)
    _plot(old_seeds.recpixlay_v_eta)
    plt.colorbar()
    center_text(0.5, 0.1, "GSF-Tracks (Old Seeding)", bbox=dict(facecolor='white', alpha=0.5))
    plt.sca(ax_rec_new)
    _plot(new_seeds.recpixlay_v_eta)
    plt.colorbar()
    center_text(0.5, 0.1, "GSF-Tracks (New Seeding)", bbox=dict(facecolor='white', alpha=0.5))


if __name__ == '__main__':
    old_seeds = ResultSet("old_seeds", 'old_seeding2.root')
    new_seeds = ResultSet("new_seeds", 'new_seeding2.root')

    # Next, declare all of the (sub)plots that will be assembled into full
    # figures later
    seed_eff = plot_seed_eff, (old_seeds, new_seeds)
    track_eff = plot_track_eff, (old_seeds, new_seeds)

    seed_pur = plot_seed_purity, (old_seeds, new_seeds)
    seed_pur_num = plot_seed_purity, (old_seeds, new_seeds), {'ext': '_num'}
    seed_pur_den = plot_seed_purity, (old_seeds, new_seeds), {'ext': '_den'}

    track_pur = plot_track_purity, (old_seeds, new_seeds)
    track_pur_num = plot_track_purity, (old_seeds, new_seeds), {'ext': '_num'}
    track_pur_den = plot_track_purity, (old_seeds, new_seeds), {'ext': '_den'}

    nHits = (plot_nhits, (old_seeds, new_seeds), {})

    # Now assemble the plots into figures.
    plots = [
        Plot([[seed_eff]],
             'ECAL-Driven Seeding Efficiency'),
        Plot([[seed_pur]],
             'ECAL-Driven Seed Purity'),
        Plot([[seed_pur_num]],
             'ECAL-Driven Seed Purity Numerator'),
        Plot([[seed_pur_den]],
             'ECAL-Driven Seed Purity Denominator'),
        Plot([[track_eff]],
             'GSF Tracking Efficiency'),
        Plot([[track_pur]],
             'GSF Track Purity'),
        Plot([[track_pur_num]],
             'GSF Track Purity Numerator'),
        Plot([[track_pur_den]],
             'GSF Track Purity Denominator'),
        Plot([[nHits]],
             'Hits'),
    ]

    # Finally, render and save the plots and generate the html+bootstrap
    # dashboard to view them
    render_plots(plots, to_disk=False)
    generate_dashboard(plots, 'Seeding Efficiency', output='eff_plots.html', source_file=__file__)