#!/usr/bin/env python

import os
import sys
import matplotlib as mpl
import matplotlib.pyplot as plt

fv_path = os.path.dirname(os.path.abspath(__file__))+"/../filval/python"
sys.path.append(fv_path)

from result_set import ResultSet
from plotter import histogram, plot_histogram, plot_histogram2d, histogram_slice
mpl.rc('text', usetex=True)
mpl.rc('savefig', dpi=180)
mpl.rc('savefig', transparent=False)



def first_hits_v_eta(rs):
    scale = 0.85
    plt.figure(figsize=(scale*10, scale*10))

    plt.subplot(321)
    plot_histogram2d(rs.dphi_v_eta_first_hits_in_B1,
                     ylabel="$\\Delta \\phi_1$(rad)",
                     title="BPIX - Layer 1")

    plt.subplot(322)
    plot_histogram2d(rs.dphi_v_eta_first_hits_in_B2,
                     title="BPIX - Layer 2")

    plt.subplot(323)
    plot_histogram2d(rs.dz_v_eta_first_hits_in_B1,
                     ylabel="$\\Delta \\textrm{z}_1$(cm)")

    plt.subplot(324)
    plot_histogram2d(rs.dz_v_eta_first_hits_in_B2)

    plt.subplot(325)
    plot_histogram2d(rs.dr_v_eta_first_hits_in_B1,
                     ylabel="$\\Delta r_1$(cm)",
                     xlabel="$\\eta$")

    plt.subplot(326)
    plot_histogram2d(rs.dr_v_eta_first_hits_in_B2,
                     xlabel="$\\eta$")

    plt.tight_layout()
    plt.savefig("figures/first_hits_v_eta.png")


plt.clf()
def second_hits_v_eta(rs):
    plt.subplot(321)
    plot_histogram2d(rs.dphi_v_eta_second_hits_in_B2,
                     ylabel="$\\Delta \\phi_2$(rad)",
                     title="BPIX - Layer 2")

    plt.subplot(322)
    plot_histogram2d(rs.dphi_v_eta_second_hits_in_B3,
                     title="BPIX - Layer 3")

    plt.subplot(323)
    plot_histogram2d(rs.dz_v_eta_second_hits_in_B2,
                     ylabel="$\\Delta \\textrm{z}_2$(cm)")

    plt.subplot(324)
    plot_histogram2d(rs.dz_v_eta_second_hits_in_B3)

    plt.subplot(325)
    plot_histogram2d(rs.dr_v_eta_second_hits_in_B2,
                     ylabel="$\\Delta r_2$(cm)",
                     xlabel="$\\eta$")

    plt.subplot(326)
    plot_histogram2d(rs.dr_v_eta_second_hits_in_B3,
                     xlabel="$\\eta$")

    plt.tight_layout()
    plt.savefig("figures/second_hits_v_eta.png")


plt.clf()
def first_hits(rs):
    plt.subplot(321)
    plot_histogram(rs.dphi_v_eta_first_hits_in_B1.ProjectionY(),
                   include_errors=True, xlabel="$\\Delta \\phi_1$(rad)",
                   title="BPIX - Layer 1")

    plt.subplot(322)
    plot_histogram(rs.dphi_v_eta_first_hits_in_B2.ProjectionY(), include_errors=True,
                   xlabel="$\\Delta \\phi_1$(rad)",
                   title="BPIX - Layer 2")

    plt.subplot(323)
    plot_histogram(rs.dz_v_eta_first_hits_in_B1.ProjectionY(), include_errors=True,
                   xlabel="$\\Delta z_1$(cm)")

    plt.subplot(324)
    plot_histogram(rs.dz_v_eta_first_hits_in_B2.ProjectionY(), include_errors=True,
                   xlabel="$\\Delta z_1$(cm)")

    plt.subplot(325)
    plot_histogram(rs.dr_v_eta_first_hits_in_B1.ProjectionY(), include_errors=True,
                   xlabel="$\\Delta r_1$(cm)")

    plt.subplot(326)
    plot_histogram(rs.dr_v_eta_first_hits_in_B2.ProjectionY(), include_errors=True,
                   xlabel="$\\Delta r_1$(cm)")

    plt.tight_layout()
    plt.savefig("figures/first_hits.png")


plt.clf()
def second_hits(rs):
    plt.subplot(321)
    plot_histogram(rs.dphi_v_eta_second_hits_in_B2.ProjectionY(), include_errors=True,
                   xlabel="$\\Delta \\phi_2$(rad)",
                   title="BPIX - Layer 2")

    plt.subplot(322)
    plot_histogram(rs.dphi_v_eta_second_hits_in_B3.ProjectionY(), include_errors=True,
                   xlabel="$\\Delta \\phi_2$(rad)",
                   title="BPIX - Layer 3")

    plt.subplot(323)
    plot_histogram(rs.dz_v_eta_second_hits_in_B2.ProjectionY(), include_errors=True,
                   xlabel="$\\Delta z_2$(cm)")

    plt.subplot(324)
    plot_histogram(rs.dz_v_eta_second_hits_in_B3.ProjectionY(), include_errors=True,
                   xlabel="$\\Delta z_2$(cm)")

    plt.subplot(325)
    plot_histogram(rs.dr_v_eta_second_hits_in_B2.ProjectionY(), include_errors=True,
                   xlabel="$\\Delta r_2$(cm)")

    plt.subplot(326)
    plot_histogram(rs.dr_v_eta_second_hits_in_B3.ProjectionY(), include_errors=True,
                   xlabel="$\\Delta r_2$(cm)")

    plt.tight_layout()
    plt.savefig("figures/second_hits.png")


plt.clf()
def delta_phi_z_v_ladder(rs):
    def even_odd_plot(even, odd, var, scale, unit, **kwargs):
        even_dist = even.ProjectionY()
        odd_dist = odd.ProjectionY()
        plot_histogram(even_dist, include_errors=True, label="even ladders")
        plot_histogram(odd_dist, include_errors=True, color='r', label="odd ladders",
                       **kwargs)

        even_mean = even_dist.GetMean()*scale
        odd_mean = odd_dist.GetMean()*scale
        axes = plt.gca()
        txt = r"$ \hat{{\Delta {0} }}_{{1,\textrm{{ {1} }} }}={2:4.2g}${3}"
        axes.text(0.05, .7, txt.format(var, "odd",  odd_mean, unit), transform=axes.transAxes)
        axes.text(0.05, .6, txt.format(var, "even", even_mean, unit), transform=axes.transAxes)
    plt.subplot(221)
    even_odd_plot(rs.dphi_v_eta_first_hits_in_B1_even_ladder, rs.dphi_v_eta_first_hits_in_B1_odd_ladder,
                  r"\phi", 10**6, "urad", xlabel=r"$\Delta \phi_1$(rad)", title="BPIX - Layer 1")

    plt.subplot(222)
    even_odd_plot(rs.dphi_v_eta_first_hits_in_B2_even_ladder, rs.dphi_v_eta_first_hits_in_B2_odd_ladder,
                  r"\phi", 10**6, "urad", xlabel=r"$\Delta \phi_1$(rad)", title="BPIX - Layer 2")
    plt.legend()

    plt.subplot(223)
    even_odd_plot(rs.dz_v_eta_first_hits_in_B1_even_ladder, rs.dz_v_eta_first_hits_in_B1_odd_ladder,
                  "z", 10**4, "um", xlabel=r"$\Delta z_1$(cm)")

    plt.subplot(224)
    even_odd_plot(rs.dz_v_eta_first_hits_in_B2_even_ladder, rs.dz_v_eta_first_hits_in_B2_odd_ladder,
                  "z", 10**4, "um", xlabel=r"$\Delta z_1$(cm)")

    plt.tight_layout()
    plt.savefig("figures/delta_phi_z_v_ladder.png")

def sc_extrapolation_first(rs):
    # Raphael's plots
    norm = 1
    errors = True
    def preproc(h):
        return histogram_slice(histogram(h.ProjectionY()), (-0.07, 0.07))
    plt.subplot(221)
    plot_histogram(preproc(rs.sc_first_hits_in_B1_dz),
                   include_errors=errors,
                   norm=norm,
                   log=True,
                   xlabel="$\\Delta z_1$(cm)",
                   title="BPIX - Layer 1")

    plt.subplot(222)
    plot_histogram(preproc(rs.sc_first_hits_in_B2_dz),
                   include_errors=errors,
                   norm=norm,
                   log=True,
                   xlabel="$\\Delta z_1$(cm)",
                   title="BPIX - Layer 2")

    plt.subplot(223)
    plot_histogram(preproc(rs.sc_first_hits_in_B1_dphi),
                   include_errors=errors,
                   norm=norm,
                   log=True,
                   # ylim=(0.5E-3, 5),
                   xlabel="$\\Delta \\phi_1$(rad)",
                   # xlim=(-0.06, 0.06)
                   )

    plt.subplot(224)
    plot_histogram(preproc(rs.sc_first_hits_in_B2_dphi),
                   include_errors=errors,
                   norm=norm,
                   log=True,
                   # ylim=(0.5E-3, 5),
                   xlabel="$\\Delta \\phi_1$(rad)",
                   # xlim=(-0.06, 0.06)
                   )

    plt.tight_layout()
    plt.savefig("figures/sc_extrapolation_first.png")


def sc_extrapolation_second(rs):
    norm = 1
    errors = True
    def preproc(h):
        return histogram(h.ProjectionY()), (-0.07, 0.07)

    plt.subplot(221)
    plot_histogram(preproc(rs.sc_second_hits_in_B2_dz),
                   include_errors=errors,
                   norm=norm,
                   log=True,
                   xlabel="$\\Delta z_2$(cm)",
                   title="BPIX - Layer 2")

    plt.subplot(222)
    plot_histogram(preproc(rs.sc_second_hits_in_B3_dz),
                   include_errors=errors,
                   norm=norm,
                   log=True,
                   xlabel="$\\Delta z_2$(cm)",
                   title="BPIX - Layer 3")

    plt.subplot(223)
    plot_histogram(preproc(rs.sc_second_hits_in_B2_dphi),
                   include_errors=errors,
                   norm=norm,
                   log=True,
                   # ylim=(0.5E-3, 5),
                   xlabel="$\\Delta \\phi_2$(rad)")

    plt.subplot(224)
    plot_histogram(preproc(rs.sc_second_hits_in_B3_dphi),
                   include_errors=errors,
                   norm=norm,
                   log=True,
                   # ylim=(0.5E-3, 5),
                   xlabel="$\\Delta \\phi_2$(rad)")

    plt.tight_layout()
    plt.savefig("figures/sc_extrapolation_second.png")


# 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")
rs = ResultSet("DY2LL", sys.argv[1])

try:
    os.mkdir('figures')
except FileExistsError:
    pass

plots = [
    first_hits_v_eta,
    second_hits_v_eta,
    first_hits,
    second_hits,
    delta_phi_z_v_ladder,
    sc_extrapolation_first,
    sc_extrapolation_second,
]

for plot in plots:
    plt.clf()
    plot(rs)

import glob
from jinja2 import Environment, PackageLoader, select_autoescape
from os.path import join
from shutil import copytree, rmtree

env = Environment(
    loader=PackageLoader('plots', 'templates'),
    autoescape=select_autoescape(['htm', 'html', 'xml'])
)

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


try:
    os.mkdir('output')
except FileExistsError:
    pass
try:
    rmtree('output/figures')
except FileNotFoundError:
    pass
copytree('figures', 'output/figures')

imgs = glob.glob("figures/*.png")


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

render_to_file('dashboard.htm', imgs=get_by_n(imgs, 6))