#!/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
from filval.plotter import (decl_plot, render_plots, hist_plot, hist_plot_stack, Plot, generate_dashboard)


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

    The proportion of gen-level electrons origination in $\rho<1$cm and $|z|<10$cm from the beam spot that have
    an associated Seed, matched via rechit-simhit associations in the pixel detector.
    """
    _, ((ax_pt, ax_eta), (ax_phi, _)) = plt.subplots(2, 2)

    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, title='Efficiency vs. Pt', label='Old')
    plt.xlabel("Pt")
    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, title='Efficiency vs. Eta', label='Old')
    plt.xlabel("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, title='Efficiency vs. Phi', label='Old')
    plt.xlabel("Phi")
    plt.ylim((0, 1.1))

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

    The proportion of gen-level electrons origination in $\rho<1$cm and $|z|<10$cm from the beam spot that have
    an associated GSF reconstructed electron.
    """
    _, ((ax_pt, ax_eta), (ax_phi, _)) = plt.subplots(2, 2)

    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, title='Efficiency vs. Pt', label='Old')
    plt.xlabel("Pt")
    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, title='Efficiency vs. Eta', label='Old')
    plt.xlabel("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, title='Efficiency vs. Phi', label='Old')
    plt.xlabel("Phi")
    plt.ylim((0, 1.1))

@decl_plot
def plot_seed_purity(old_seeds, new_seeds):
    r"""## Seed Purity

    The proportion of ECAL-driven seeds that have a matched gen-level electron originating in
    $\rho<1$cm and $|z|<10$cm from the beam spot.
    """
    _, ((ax_pt, ax_eta), (ax_phi, _)) = plt.subplots(2, 2)

    errors = True
    plt.sca(ax_pt)
    hist_plot(hist(new_seeds.seed_pur_v_pt), include_errors=errors, label='New')
    hist_plot(hist(old_seeds.seed_pur_v_pt), include_errors=errors, title='Purity vs. Pt', label='Old')
    plt.xlabel("Pt")
    plt.ylim((0, 1.1))

    plt.sca(ax_eta)
    hist_plot(hist(new_seeds.seed_pur_v_eta), include_errors=errors, label='New')
    hist_plot(hist(old_seeds.seed_pur_v_eta), include_errors=errors, title='Purity vs. Eta', label='Old')
    plt.xlabel("Eta")
    plt.ylim((0, 1.1))
    plt.legend()

    plt.sca(ax_phi)
    hist_plot(hist(new_seeds.seed_pur_v_phi), include_errors=errors, label='New')
    hist_plot(hist(old_seeds.seed_pur_v_phi), include_errors=errors, title='Purity vs. Phi', label='Old')
    plt.xlabel("Phi")
    plt.ylim((0, 1.1))

if __name__ == '__main__':
    old_seeds = ResultSet("old_seeds", 'build/old_seeding.root')
    new_seeds = ResultSet("new_seeds", 'build/new_seeding.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), {})

    # Now assemble the plots into figures.
    plots = [
        Plot([[seed_eff]],
             'Seeding Efficiency'),
        Plot([[track_eff]],
             'Tracking Efficiency'),
        Plot([[seed_pur]],
             'ECAL-Driven Seed Purity'),
    ]

    # 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__)