#!/usr/bin/env python
# Copyright (c) 2019, Anthony Latorre <tlatorre at uchicago>
#
# This program is free software: you can redistribute it and/or modify it
# under the terms of the GNU General Public License as published by the Free
# Software Foundation, either version 3 of the License, or (at your option)
# any later version.
#
# This program is distributed in the hope that it will be useful, but WITHOUT
# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
# FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for
# more details.
#
# You should have received a copy of the GNU General Public License along with
# this program. If not, see <https://www.gnu.org/licenses/>.

from __future__ import print_function, division
import numpy as np

# Taken from https://raw.githubusercontent.com/mwaskom/seaborn/c73055b2a9d9830c6fbbace07127c370389d04dd/seaborn/utils.py
def despine(fig=None, ax=None, top=True, right=True, left=False,
            bottom=False, offset=None, trim=False):
    """Remove the top and right spines from plot(s).

    fig : matplotlib figure, optional
        Figure to despine all axes of, default uses current figure.
    ax : matplotlib axes, optional
        Specific axes object to despine.
    top, right, left, bottom : boolean, optional
        If True, remove that spine.
    offset : int or dict, optional
        Absolute distance, in points, spines should be moved away
        from the axes (negative values move spines inward). A single value
        applies to all spines; a dict can be used to set offset values per
        side.
    trim : bool, optional
        If True, limit spines to the smallest and largest major tick
        on each non-despined axis.

    Returns
    -------
    None

    """
    # Get references to the axes we want
    if fig is None and ax is None:
        axes = plt.gcf().axes
    elif fig is not None:
        axes = fig.axes
    elif ax is not None:
        axes = [ax]

    for ax_i in axes:
        for side in ["top", "right", "left", "bottom"]:
            # Toggle the spine objects
            is_visible = not locals()[side]
            ax_i.spines[side].set_visible(is_visible)
            if offset is not None and is_visible:
                try:
                    val = offset.get(side, 0)
                except AttributeError:
                    val = offset
                _set_spine_position(ax_i.spines[side], ('outward', val))

        # Potentially move the ticks
        if left and not right:
            maj_on = any(
                t.tick1line.get_visible()
                for t in ax_i.yaxis.majorTicks
            )
            min_on = any(
                t.tick1line.get_visible()
                for t in ax_i.yaxis.minorTicks
            )
            ax_i.yaxis.set_ticks_position("right")
            for t in ax_i.yaxis.majorTicks:
                t.tick2line.set_visible(maj_on)
            for t in ax_i.yaxis.minorTicks:
                t.tick2line.set_visible(min_on)

        if bottom and not top:
            maj_on = any(
                t.tick1line.get_visible()
                for t in ax_i.xaxis.majorTicks
            )
            min_on = any(
                t.tick1line.get_visible()
                for t in ax_i.xaxis.minorTicks
            )
            ax_i.xaxis.set_ticks_position("top")
            for t in ax_i.xaxis.majorTicks:
                t.tick2line.set_visible(maj_on)
            for t in ax_i.xaxis.minorTicks:
                t.tick2line.set_visible(min_on)

        if trim:
            # clip off the parts of the spines that extend past major ticks
            xticks = ax_i.get_xticks()
            if xticks.size:
                firsttick = np.compress(xticks >= min(ax_i.get_xlim()),
                                        xticks)[0]
                lasttick = np.compress(xticks <= max(ax_i.get_xlim()),
                                       xticks)[-1]
                ax_i.spines['bottom'].set_bounds(firsttick, lasttick)
                ax_i.spines['top'].set_bounds(firsttick, lasttick)
                newticks = xticks.compress(xticks <= lasttick)
                newticks = newticks.compress(newticks >= firsttick)
                ax_i.set_xticks(newticks)

            yticks = ax_i.get_yticks()
            if yticks.size:
                firsttick = np.compress(yticks >= min(ax_i.get_ylim()),
                                        yticks)[0]
                lasttick = np.compress(yticks <= max(ax_i.get_ylim()),
                                       yticks)[-1]
                ax_i.spines['left'].set_bounds(firsttick, lasttick)
                ax_i.spines['right'].set_bounds(firsttick, lasttick)
                newticks = yticks.compress(yticks <= lasttick)
                newticks = newticks.compress(newticks >= firsttick)
                ax_i.set_yticks(newticks)

if __name__ == '__main__':
    import ROOT
    import argparse
    from os.path import split
    from matplotlib.ticker import FuncFormatter

    parser = argparse.ArgumentParser("plot ROOT fit results")
    parser.add_argument("filename", help="input file")
    parser.add_argument("--save", action="store_true", default=False, help="save plots")
    args = parser.parse_args()

    if args.save:
        # default \textwidth for a fullpage article in Latex is 16.50764 cm.
        # You can figure this out by compiling the following TeX document:
        #
        #    \documentclass{article}
        #    \usepackage{fullpage}
        #    \usepackage{layouts}
        #    \begin{document}
        #    textwidth in cm: \printinunitsof{cm}\prntlen{\textwidth}
        #    \end{document}

        width = 16.50764
        width /= 2.54 # cm -> inches
        # According to this page:
        # http://www-personal.umich.edu/~jpboyd/eng403_chap2_tuftegospel.pdf,
        # Tufte suggests an aspect ratio of 1.5 - 1.6.
        height = width/1.5
        FIGSIZE = (width,height)

        import matplotlib.pyplot as plt

        font = {'family':'serif', 'serif': ['computer modern roman']}
        plt.rc('font',**font)

        plt.rc('text', usetex=True)
    else:
        # on retina screens, the default plots are way too small
        # by using Qt5 and setting QT_AUTO_SCREEN_SCALE_FACTOR=1
        # Qt5 will scale everything using the dpi in ~/.Xresources
        import matplotlib
        matplotlib.use("Qt5Agg")

        import matplotlib.pyplot as plt

        # Default figure size. Currently set to my monitor width and height so that
        # things are properly formatted
        FIGSIZE = (13.78,7.48)

        # Make the defalt font bigger
        plt.rc('font', size=22)

    root_file = ROOT.TFile(args.filename)

    head, tail = split(args.filename)

    if tail.startswith("e_") or tail.startswith("electron"):
        prefix = "electron"
    elif tail.startswith("mu_") or tail.startswith("muon"):
        prefix = "muon"
    else:
        prefix = ""

    try:
        if root_file.Get("h1"):
            for hist_number, tf1_number in zip([1,2,4,5],[1,2,3,None]):
                h = root_file.Get("h%i" % hist_number)
                if tf1_number:
                    f = root_file.Get("f%i" % tf1_number)

                bins = [h.GetXaxis().GetBinLowEdge(i) for i in range(1,h.GetNbinsX()+1)] + [h.GetXaxis().GetBinUpEdge(h.GetNbinsX())]
                hist = [h.GetBinContent(i) for i in range(1,h.GetNbinsX()+1)]

                bins = np.array(bins)
                hist = np.array(hist)

                bincenters = (bins[1:] + bins[:-1])/2

                norm = np.trapz(hist,bincenters)

                hist /= norm

                fig = plt.figure(figsize=FIGSIZE)
                plt.hist(bincenters,weights=hist,bins=bins,histtype='step')
                x = np.linspace(bins[0],bins[-1],10000)
                if tf1_number:
                    plt.plot(x,[f(xval)/norm for xval in x],color='red')
                despine(fig,trim=True)
                if hist_number == 1:
                    plt.gca().set_xlim(-1,1)
                    plt.ylabel("Arbitrary Units")
                    plt.xlabel(r"$\cos\theta$")
                    if args.save:
                        plt.savefig("%s_shower_angular_distribution.pdf" % prefix)
                        plt.savefig("%s_shower_angular_distribution.eps" % prefix)
                    else:
                        plt.title("%s Shower Angular Distribution" % prefix.capitalize())
                elif hist_number == 2:
                    plt.ylabel("Arbitrary Units")
                    plt.xlabel(r"Distance along Track (cm)")
                    if args.save:
                        plt.savefig("%s_shower_position_distribution.pdf" % prefix)
                        plt.savefig("%s_shower_position_distribution.eps" % prefix)
                    else:
                        plt.title("%s Shower Position Distribution" % prefix.capitalize())
                elif hist_number == 4:
                    plt.ylabel("Arbitrary Units")
                    plt.xlabel(r"$\cos\theta$")
                    if args.save:
                        plt.savefig("%s_delta_ray_angular_distribution.pdf" % prefix)
                        plt.savefig("%s_delta_ray_angular_distribution.eps" % prefix)
                    else:
                        plt.title("%s Delta Ray Angular Distribution" % prefix.capitalize())
                elif hist_number == 5:
                    plt.ylabel("Arbitrary Units")
                    plt.xlabel(r"Distance along Track (cm)")
                    if args.save:
                        plt.savefig("%s_delta_ray_position_distribution.pdf" % prefix)
                        plt.savefig("%s_delta_ray_position_distribution.eps" % prefix)
                    else:
                        plt.title("%s Delta Ray Position Distribution" % prefix.capitalize())
        else:
            for graph_name, tf1_number, ylabel in zip(["g_dir_alpha","g_dir_beta","g_pos_alpha","g_pos_beta","g_dir_alpha_delta","g_dir_beta_delta"],
                                                      [1,2,None,None,3,4],
                                                      [r"$\alpha$",r"$\beta$",r"$k$",r"$\theta$",r"$\alpha$",r"$\beta$"]):
                g = root_file.Get(graph_name)
                if tf1_number:
                    f = g.GetFunction("f%i" % tf1_number)

                x = [g.GetX()[i] for i in range(g.GetN())]
                y = [g.GetY()[i] for i in range(g.GetN())]
                yerr = [g.GetEY()[i] for i in range(g.GetN())]

                x = np.array(x)
                y = np.array(y)
                yerr = np.array(yerr)

                fig = plt.figure(figsize=FIGSIZE)
                plt.errorbar(x,y,yerr=yerr,fmt='o')
                x = np.linspace(x[0],x[-1],10000)
                if tf1_number:
                    plt.plot(x,[f(xval) for xval in x],color='red')
                despine(fig,trim=True)
                plt.xlabel("Kinetic Energy (MeV)")
                plt.ylabel(ylabel)

                if graph_name == "g_dir_alpha":
                    if args.save:
                        plt.savefig("%s_shower_angular_distribution_alpha.pdf" % prefix)
                        plt.savefig("%s_shower_angular_distribution_alpha.eps" % prefix)
                    else:
                        plt.title("%s Shower Angular Distribution" % prefix.capitalize())
                elif graph_name == "g_dir_beta":
                    if args.save:
                        plt.savefig("%s_shower_angular_distribution_beta.pdf" % prefix)
                        plt.savefig("%s_shower_angular_distribution_beta.eps" % prefix)
                    else:
                        plt.title("%s Shower Position Distribution" % prefix.capitalize())
                elif graph_name == "g_pos_alpha":
                    if args.save:
                        plt.savefig("%s_shower_position_distribution_alpha.pdf" % prefix)
                        plt.savefig("%s_shower_position_distribution_alpha.eps" % prefix)
                    else:
                        plt.title("%s Shower Position Distribution" % prefix.capitalize())
                elif graph_name == "g_pos_beta":
                    if args.save:
                        plt.savefig("%s_shower_position_distribution_beta.pdf" % prefix)
                        plt.savefig("%s_shower_position_distribution_beta.eps" % prefix)
                    else:
                        plt.title("%s Shower Position Distribution" % prefix.capitalize())
                elif graph_name == "g_dir_alpha_delta":
                    if args.save:
                        plt.savefig("%s_delta_ray_angular_distribution_alpha.pdf" % prefix)
                        plt.savefig("%s_delta_ray_angular_distribution_alpha.eps" % prefix)
                    else:
                        plt.title("%s Delta Ray Angular Distribution" % prefix.capitalize())
                elif graph_name == "g_dir_beta_delta":
                    if args.save:
                        plt.savefig("%s_delta_ray_angular_distribution_beta.pdf" % prefix)
                        plt.savefig("%s_delta_ray_angular_distribution_beta.eps" % prefix)
                    else:
                        plt.title("%s Delta Ray Position Distribution" % prefix.capitalize())

        plt.show()
    
    finally:
        root_file.Close()