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


def generate_dashboard(plots, output='dashboard.htm'):
    from jinja2 import Environment, PackageLoader, select_autoescape
    from os.path import join
    from urllib.parse import quote

    env = Environment(
        loader=PackageLoader('plots', 'templates'),
        autoescape=select_autoescape(['htm', 'html', 'xml']),
    )
    env.globals.update({'quote': quote,
                        'enumerate': enumerate,
                        'zip': zip,
                        })

    def render_to_file(template_name, **kwargs):
        with open(join('output', output), 'w') as tempout:
            template = env.get_template(template_name)
            tempout.write(template.render(**kwargs))

    def get_by_n(objs, n=2):
        objs = list(objs)
        while objs:
            yield objs[:n]
            objs = objs[n:]

    render_to_file('dashboard.htm', plots=get_by_n(plots, 3),
                   outdir="figures/")


@decl_plot
def plot_yield_grid(rss):
    r"""## Event Yield

    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>
    """
    _, ((ax_tttt, ax_ttw), (ax_ttz, ax_tth)) = plt.subplots(2, 2)
    ft, ttw, ttz, tth = map(lambda rs: hist(rs.SRs), rss)
    plt.sca(ax_tttt)
    hist_plot(ft, title='TTTT', stats=False)
    plt.sca(ax_ttw)
    hist_plot(ttw, title='TTW', stats=False)
    plt.sca(ax_ttz)
    hist_plot(ttz, title='TTZ', stats=False)
    plt.sca(ax_tth)
    hist_plot(tth, title='TTH', stats=False)

    plt.xlabel('Signal Region')


@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>
    """
    ft, ttw, ttz, tth = map(lambda rs: hist(rs.SRs), rss)
    hist_plot_stack([ttw, ttz, tth], labels=['TTW', 'TTZ', 'TTH'])
    ft = ft[0]*10, ft[1], ft[2]
    hist_plot(ft, label='TTTT (x10)', stats=False, color='k')

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


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

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

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


if __name__ == '__main__':
    # First create a ResultSet object which loads all of the objects from output.root
    # into memory and makes them available as attributes
    # if len(sys.argv) != 2:
    #     raise ValueError("please supply root file")
    rss = (ResultSet("ft", 'data/yield_ft.root'),
           ResultSet("ttw", 'data/yield_ttw.root'),
           ResultSet("ttz", 'data/yield_ttz.root'),
           ResultSet("tth", 'data/yield_tth.root'))
    rss_notau = (ResultSet("ft_notau", 'data/yield_ft_notau.root'),
                 ResultSet("ttw_notau", 'data/yield_ttw_notau.root'),
                 ResultSet("ttz_notau", 'data/yield_ttz_notau.root'),
                 ResultSet("tth_notau", 'data/yield_tth_notau.root'))

    # Next, declare all of the (sub)plots that will be assembled into full
    # figures later
    yield_tau = (plot_yield_grid, (rss,), {})
    yield_notau = (plot_yield_grid, (rss_notau,), {})

    yield_tau_stack = (plot_yield_stack, (rss,), {})
    yield_notau_stack = (plot_yield_stack, (rss_notau,), {})

    sig_strength_tau = (plot_sig_strength, (rss,), {})
    sig_strength_notau = (plot_sig_strength, (rss_notau,), {})
    # Now assemble the plots into figures.
    plots = [
        Plot([[yield_tau]],
             'Yield With Tau'),
        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')
    ]

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