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

from filval.result_set import ResultSet
from filval.histogram import hist, hist2d, hist_add, hist_norm, hist_scale, hist2d_norm
from filval.plotting import (decl_plot, render_plots, hist_plot, hist_plot_stack, hist2d_plot,
                             Plot, generate_dashboard, hists_to_table)

an_tttt = ([0.47, 0.33, 0.18, 0.78, 0.49, 0.52, 0.33, 0.49],
           [0, 0, 0,  0, 0, 0, 0, 0], [0.5, 1.5, 2.5, 3.5, 4.5, 5.5, 6.5, 7.5, 8.5])
an_ttw = ([2.29663, 0.508494, 0.161166, 1.03811, 0.256401, 0.127582, 0.181522, 0.141659],
          [0, 0, 0,  0, 0, 0, 0, 0], [0.5, 1.5, 2.5, 3.5, 4.5, 5.5, 6.5, 7.5, 8.5])
an_ttz = ([0.974751, 0.269195, 1e-06, 0.395831, 0.0264703, 0.06816, 0.8804, 0.274265],
          [0, 0, 0,  0, 0, 0, 0, 0], [0.5, 1.5, 2.5, 3.5, 4.5, 5.5, 6.5, 7.5, 8.5])
an_tth = ([1.13826, 0.361824, 0.162123, 0.683917, 0.137608, 0.0632719, 0.554491, 0.197864],
          [0, 0, 0,  0, 0, 0, 0, 0], [0.5, 1.5, 2.5, 3.5, 4.5, 5.5, 6.5, 7.5, 8.5])

@decl_plot
def plot_yield_grid(rss, tau_category=-1):
    r"""## Event Yield

    The event yield for the eight signal regions defined in AN-17-115. Data is normalized
    to the integrated luminosity of $35.9\textrm{fb}^{-1}$.

    Ignoring taus means both that there is no requirement on number of good taus *and*
    the taus, if present, are not considered for the SS pair.

    If taus are not ignored, any good tau with $p_T$>20Gev is considered in constructing the SS lepton pair. The yields
    are then further broken down by the number of good taus in the event.

    A "good" tau in the above means any tau candidate passing the `byTightIsolationMVArun2v1DBoldDMwLT` ID w/ pt>20GeV.
    It is also required to pass tau-lepton cross cleaning where it must not match any electron or muon within
    $\delta R < 0.4$.

    Truth-matched taus are those that match within $\delta R < 0.3$ with gen-level taus that pass the flag `fromHardProcessDecayed`.
    """

    def get_sr(rs, tm=False):
        if tm:
            if tau_category == 0:
                return hist(rs.SRs_0tmtau)
            if tau_category == 1:
                return hist(rs.SRs_1tmtau)
            elif tau_category == 2:
                return hist(rs.SRs_2tmtau)
        else:
            if tau_category == -1:
                return hist(rs.ignore_tau_SRs)
            elif tau_category == 0:
                return hist(rs.SRs_0tau)
            elif tau_category == 1:
                return hist(rs.SRs_1tau)
            elif tau_category == 2:
                return hist(rs.SRs_2tau)

    _, ((ax_tttt, ax_ttw), (ax_ttz, ax_tth)) = plt.subplots(2, 2)
    tttt, ttw, ttz, tth = [get_sr(rs) for rs in rss]
    tm_tttt, tm_ttw, tm_ttz, tm_tth = [get_sr(rs, True) for rs in rss]

    plt.sca(ax_tttt)
    hist_plot(tttt, title='TTTT', stats=False, label='Mock', include_errors=True)
    if tau_category == -1 and len(tttt[0]) == len(an_tttt[0]):
        hist_plot(an_tttt, title='TTTT', stats=False, label='AN')
    elif tau_category >= 0:
        hist_plot(tm_tttt, title='TTTT', stats=False, label='Truth-Matched Taus', include_errors=True)
    plt.ylim((0, None))

    plt.sca(ax_ttw)
    hist_plot(ttw, title='TTW', stats=False, label='Mock', include_errors=True)
    if tau_category == -1 and len(tttt[0]) == len(an_tttt[0]):
        hist_plot(an_ttw, title='TTW', stats=False, label='AN')
    elif tau_category >= 0:
        hist_plot(tm_ttw, title='TTW', stats=False, label='Truth-Matched Taus', include_errors=True)
    plt.ylim((0, None))
    plt.legend()

    plt.sca(ax_ttz)
    hist_plot(ttz, title='TTZ', stats=False, label='Mock', include_errors=True)
    if tau_category == -1 and len(tttt[0]) == len(an_tttt[0]):
        hist_plot(an_ttz, title='TTZ', stats=False, label='AN')
    elif tau_category >= 0:
        hist_plot(tm_ttz, title='TTZ', stats=False, label='Truth-Matched Taus', include_errors=True)
    plt.ylim((0, None))
    plt.xlabel('Signal Region')

    plt.sca(ax_tth)
    hist_plot(tth, title='TTH', stats=False, label='Mock', include_errors=True)
    if tau_category == -1 and len(tttt[0]) == len(an_tttt[0]):
        hist_plot(an_tth, title='TTH', stats=False, label='AN')
    elif tau_category >= 0:
        hist_plot(tm_tth, title='TTH', stats=False, label='Truth-Matched Taus', include_errors=True)
    plt.ylim((0, None))
    plt.xlabel('Signal Region')

    def to_table(hists):
        return hists_to_table(hists, row_labels=['TTTT', 'TTW', 'TTZ', 'TTH'],
                              column_labels=[f'SR{n}' for n in range(1, len(tttt[0])+1)])

    tables = '<h2>Mock</h2>'
    tables += to_table([tttt, ttw, ttz, tth])
    if tau_category == -1:
        tables += '<h2>AN</h2>'
        tables += to_table([an_tttt, an_ttw, an_ttz, an_tth])
    elif tau_category >= 0:
        tables += '<h2>TM Taus</h2>'
        tables += to_table([tm_tttt, tm_ttw, tm_ttz, tm_tth])
    return tables


@decl_plot
def plot_yield_v_gen(rss, tau_category=-1):
    r"""## Event Yield Vs. # of Generated Taus

    """
    def get_sr(rs):
        h = None
        if tau_category == 0:
            h = rs.SRs_0tmtau_diff_nGenTau
        if tau_category == 1:
            h = rs.SRs_1tmtau_diff_nGenTau
        elif tau_category == 2:
            h = rs.SRs_2tmtau_diff_nGenTau
        return hist2d_norm(hist2d(h), norm=1, axis=0)

    _, ((ax_tttt, ax_ttw), (ax_ttz, ax_tth)) = plt.subplots(2, 2)
    tttt, ttw, ttz, tth = [get_sr(rs) for rs in rss]

    def do_plot(h, title):
        hist2d_plot(h, title=title, txt_format='{:.2f}')

    plt.sca(ax_tttt)
    do_plot(tttt, 'TTTT')
    plt.ylabel("\# Gen Taus")

    plt.sca(ax_ttw)
    do_plot(ttw, 'TTW')
    plt.legend()

    plt.sca(ax_ttz)
    do_plot(ttz, 'TTZ')
    plt.xlabel('Signal Region')
    plt.ylabel("\# Gen Taus")

    plt.sca(ax_tth)
    do_plot(tth, 'TTH')
    plt.xlabel('Signal Region')


@decl_plot
def plot_nGen_v_nSel(rss):
    _, ((ax_tttt, ax_ttw), (ax_ttz, ax_tth)) = plt.subplots(2, 2)
    tttt, ttw, ttz, tth = [hist2d(rs.nGen_v_RecoTaus_in_SR) for rs in rss]

    def do_plot(h, title):
        hist2d_plot(h, title=title, txt_format='{:.2f}')

    plt.sca(ax_tttt)
    do_plot(tttt, 'TTTT')
    plt.ylabel("\# Selected Taus")

    plt.sca(ax_ttw)
    do_plot(ttw, 'TTW')
    plt.legend()

    plt.sca(ax_ttz)
    do_plot(ttz, 'TTZ')
    plt.xlabel('\# Gen Taus')
    plt.ylabel("\# Selected Taus")

    plt.sca(ax_tth)
    do_plot(tth, 'TTH')
    plt.xlabel('\# Gen Taus')


@decl_plot
def plot_yield_stack(rss):
    r"""## Event Yield - Stacked

    The event yield for the eight signal regions defined in AN-17-115. Data is normalized
    to the Moriond 2018 integrated luminosity ($35.9\textrm{fb}^{-1}$). Code for the histogram generation is
    here: <https://github.com/cfangmeier/FTAnalysis/blob/master/studies/tau/Yield.C>
    """
    tttt, ttw, ttz, tth = map(lambda rs: hist(rs.SRs), rss)
    hist_plot_stack([ttw, ttz, tth], labels=['TTW', 'TTZ', 'TTH'])
    tttt = tttt[0]*10, tttt[1], tttt[2]
    hist_plot(tttt, label='TTTT (x10)', stats=False, color='k')

    plt.ylim((0, 60))
    plt.xlabel('Signal Region')
    plt.legend()


@decl_plot
def plot_lep_multi(rss, dataset):
    _, (ax_els, ax_mus, ax_taus) = plt.subplots(3, 1)
    els = list(map(lambda rs: hist_norm(hist(rs.nEls)), rss))
    mus = list(map(lambda rs: hist_norm(hist(rs.nMus)), rss))
    taus = list(map(lambda rs: hist_norm(hist(rs.nTaus)), rss))

    def _plot(ax, procs):
        plt.sca(ax)
        tttt, ttw, ttz, tth = procs
        h = {'TTTT': tttt,
             'TTW': ttw,
             'TTZ': ttz,
             'TTH': tth}[dataset]
        hist_plot(h, stats=False, label=dataset)

    _plot(ax_els, els)
    plt.xlabel('\\# Good Electrons')
    plt.legend()
    _plot(ax_mus, mus)
    plt.xlabel('\\# Good Muons')
    _plot(ax_taus, taus)
    plt.xlabel('\\# Good Taus')


@decl_plot
def plot_sig_strength(rss):
    r""" The signal strength of the TTTT signal defined as

    $\frac{S}{\sqrt{S+B}}$

    """
    tttt, ttw, ttz, tth = map(lambda rs: hist(rs.SRs), rss)
    bg = hist_add(ttw, ttz, tth)
    strength = tttt[0] / np.sqrt(tttt[0] + bg[0])
    hist_plot((strength, tttt[1], tttt[2]), stats=False)


@decl_plot
def plot_event_obs(rss, dataset, in_signal_region=True):
    r"""The distribution of $N_{jet}$, $N_{Bjet}$, MET, and $H_T$ in either all events
    or only signal region (igoring taus) events.

    """
    _, ((ax_njet, ax_nbjet), (ax_ht, ax_met)) = plt.subplots(2, 2)
    # tttt, ttw, ttz, tth = map(lambda rs: hist(rs.SRs), rss)
    tttt, ttw, ttz, tth = rss
    rs = {'TTTT': tttt,
          'TTW': ttw,
          'TTZ': ttz,
          'TTH': tth}[dataset]

    def _plot(ax, obs):
        plt.sca(ax)
        if in_signal_region:
            h = {'MET': rs.met_in_SR,
                 'HT': rs.ht_in_SR,
                 'NJET': rs.njet_in_SR,
                 'NBJET': rs.nbjet_in_SR}[obs]
        else:
            h = {'MET': rs.met,
                 'HT': rs.ht,
                 'NJET': rs.njet,
                 'NBJET': rs.nbjet}[obs]
        hist_plot(hist(h), stats=False, label=dataset, xlabel=obs)

    _plot(ax_njet, 'NJET')
    _plot(ax_nbjet, 'NBJET')
    _plot(ax_ht, 'HT')
    _plot(ax_met, 'MET')


@decl_plot
def plot_event_obs_stack(rss, in_signal_region=True):
    r"""

    """
    _, ((ax_njet, ax_nbjet), (ax_ht, ax_met)) = plt.subplots(2, 2)

    def _plot(ax, obs):
        plt.sca(ax)
        if in_signal_region:
            attr = {'MET': 'met_in_SR',
                    'HT': 'ht_in_SR',
                    'NJET': 'njet_in_SR',
                    'NBJET': 'nbjet_in_SR'}[obs]
        else:
            attr = {'MET': 'met',
                    'HT': 'ht',
                    'NJET': 'njet',
                    'NBJET': 'nbjet'}[obs]
        tttt, ttw, ttz, tth = map(lambda rs: hist(getattr(rs, attr)), rss)
        hist_plot_stack([ttw, ttz, tth], labels=["TTW", "TTZ", "TTH"])
        hist_plot(hist_scale(tttt, 5), label="TTTT (x5)", color='k')
        plt.xlabel(obs)

    _plot(ax_njet, 'NJET')
    _plot(ax_nbjet, 'NBJET')
    plt.legend()
    _plot(ax_ht, 'HT')
    _plot(ax_met, 'MET')

# @decl_plot
# def plot_s_over_b(rss):
#     def get_sr(rs):
#         h = None
#         if tau_category == 0:
#             h = rs.SRs_0tmtau_diff_nGenTau
#         if tau_category == 1:
#             h = rs.SRs_1tmtau_diff_nGenTau
#         elif tau_category == 2:
#             h = rs.SRs_2tmtau_diff_nGenTau
#         return hist2d_norm(hist2d(h), norm=1, axis=0)
#
#     _, ((ax_tttt, ax_ttw), (ax_ttz, ax_tth)) = plt.subplots(2, 2)
#     tttt =
#     ttw, ttz, tth = [get_sr(rs) for rs in rss]
#     pass

@decl_plot
def plot_tau_purity(rss):
    _, ((ax_tttt, ax_ttw), (ax_ttz, ax_tth)) = plt.subplots(2, 2)
    tttt, ttw, ttz, tth = list(map(lambda rs: hist(rs.tau_purity_v_pt), rss))

    def _plot(ax, dataset):
        plt.sca(ax)
        h = {'TTTT': tttt,
             'TTW': ttw,
             'TTZ': ttz,
             'TTH': tth}[dataset]
        hist_plot(h, stats=False, label=dataset)
        plt.text(200, 0.05, dataset)
        plt.xlabel(r"$P_T$(GeV)")

    _plot(ax_tttt, 'TTTT')
    _plot(ax_ttw, 'TTW')
    _plot(ax_ttz, 'TTZ')
    _plot(ax_tth, 'TTH')


if __name__ == '__main__':
    data_path = 'data/output_new_sr_new_id_binning/'
    save_plots = True
    # First create a ResultSet object which loads all of the objects from root file
    # into memory and makes them available as attributes
    rss = (ResultSet("tttt", data_path+'yield_tttt.root'),
           ResultSet("ttw", data_path+'yield_ttw.root'),
           ResultSet("ttz", data_path+'yield_ttz.root'),
           ResultSet("tth", data_path+'yield_tth.root'))

    # Next, declare all of the (sub)plots that will be assembled into full
    # figures later
    yield_tau_ignore_tau = plot_yield_grid, (rss, -1)
    yield_tau_0tau = plot_yield_grid, (rss, 0)
    yield_tau_1tau = plot_yield_grid, (rss, 1)
    yield_tau_2tau = plot_yield_grid, (rss, 2)

    tttt_event_obs_in_sr = plot_event_obs, (rss, 'TTTT'), {'in_signal_region': True}
    ttw_event_obs_in_sr = plot_event_obs, (rss, 'TTW'), {'in_signal_region': True}
    ttz_event_obs_in_sr = plot_event_obs, (rss, 'TTZ'), {'in_signal_region': True}
    tth_event_obs_in_sr = plot_event_obs, (rss, 'TTH'), {'in_signal_region': True}

    tttt_event_obs = plot_event_obs, (rss, 'TTTT'), {'in_signal_region': False}
    ttw_event_obs = plot_event_obs, (rss, 'TTW'), {'in_signal_region': False}
    ttz_event_obs = plot_event_obs, (rss, 'TTZ'), {'in_signal_region': False}
    tth_event_obs = plot_event_obs, (rss, 'TTH'), {'in_signal_region': False}

    tttt_lep_multi = plot_lep_multi, (rss, 'TTTT')
    ttw_lep_multi = plot_lep_multi, (rss, 'TTW')
    ttz_lep_multi = plot_lep_multi, (rss, 'TTZ')
    tth_lep_multi = plot_lep_multi, (rss, 'TTH')

    yield_v_gen_0tm = plot_yield_v_gen, (rss, 0)
    yield_v_gen_1tm = plot_yield_v_gen, (rss, 1)
    yield_v_gen_2tm = plot_yield_v_gen, (rss, 2)
    # tau_purity = plot_tau_purity, (rss)
    nGen_v_nSel = plot_nGen_v_nSel, (rss,)

    # Now assemble the plots into figures.
    plots = [
        Plot([[yield_tau_ignore_tau]],
             'Yield Ignoring Taus'),
        Plot([[yield_tau_0tau]],
             'Yield For events with 0 Tau'),
        Plot([[yield_tau_1tau]],
             'Yield For events with 1 Tau'),
        Plot([[yield_tau_2tau]],
             'Yield For events with 2 or more Tau'),

        # Plot([[yield_v_gen_0tm]],
        #      'Yield For events with 0 TM Tau Vs Gen Tau'),
        # Plot([[yield_v_gen_1tm]],
        #      'Yield For events with 1 TM Tau Vs Gen Tau'),
        # Plot([[yield_v_gen_2tm]],
        #      'Yield For events with 2 TM Tau Vs Gen Tau'),
        Plot([[nGen_v_nSel]],
             r'#Generated tau vs. #Selected tau'),

        Plot([[tttt_lep_multi]],
             'Lepton Multiplicity - TTTT'),
        Plot([[ttw_lep_multi]],
             'Lepton Multiplicity - TTW'),
        Plot([[ttz_lep_multi]],
             'Lepton Multiplicity - TTZ'),
        Plot([[tth_lep_multi]],
             'Lepton Multiplicity - TTH'),
        Plot([[tttt_event_obs_in_sr]],
             'TTTT - Event Observables (In SR)'),
        Plot([[ttw_event_obs_in_sr]],
             'TTW - Event Observables (In SR)'),
        Plot([[ttz_event_obs_in_sr]],
             'TTZ - Event Observables (In SR)'),
        Plot([[tth_event_obs_in_sr]],
             'TTH - Event Observables (In SR)'),
        # Plot([[tttt_event_obs]],
        #      'TTTT - Event Observables (All Events)'),
        # Plot([[ttw_event_obs]],
        #      'TTW - Event Observables (All Events)'),
        # Plot([[ttz_event_obs]],
        #      'TTZ - Event Observables (All Events)'),
        # Plot([[tth_event_obs]],
        #      'TTH - Event Observables (All Events)'),

        # Plot([[yield_notau]],
        #      'Yield Without Tau'),
        # Plot([[yield_tau_stack]],
        #      'Yield With Tau Stacked'),
        # Plot([[yield_notau_stack]],
        #      'Yield Without Tau Stacked'),
        # Plot([[yield_tau_stack],
        #       [yield_notau_stack]],
        #      'Event Yield, top: with tau, bottom: no tau'),
        # Plot([[sig_strength_tau],
        #       [sig_strength_notau]],
        #      'Signal Strength'),
        # Plot([[event_obs_stack]],
        #      'Event Observables'),
        # Plot([[tau_purity]],
        #      'Tau Purity'),
    ]

    # Finally, render and save the plots and generate the html+bootstrap
    # dashboard to view them
    render_plots(plots, to_disk=save_plots)
    if not save_plots:
        generate_dashboard(plots, 'TTTT Yields',
                           output='yields.html',
                           source=__file__,
                           ana_source=("https://github.com/cfangmeier/FTAnalysis/commit/"
                                       "0cbdac4509391fffb9fff87d0521b7dd0a30a55c"),
                           config=data_path+'config.yaml'
                           )