# coding: utf-8

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

sys.path.append("../filval/python/")
from utils import ResultSet
from plotter import plot_histogram, plot_histogram2d
mpl.rc('text', usetex=True)
mpl.rc('savefig', dpi=180)
mpl.rc('savefig', transparent=True)

# First create a ResultSet object which loads all of the objects from output.root
# into memory and makes them available as attributes
rs = ResultSet("DY2LL",  "/home/caleb/Sources/EGamma/build/output.root")

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

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()
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()
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()
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")


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.clf()

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")