#!/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

if __name__ == '__main__':
    import ROOT
    import argparse
    from os.path import split
    from sddm.plot import despine
    from sddm import setup_matplotlib

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

    setup_matplotlib(args.save)

    import matplotlib.pyplot as plt

    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()
                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()
                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()
<a id='n167' href='#n167'>167</a>
<a id='n168' href='#n168'>168</a>
<a id='n169' href='#n169'>169</a>
<a id='n170' href='#n170'>170</a>
<a id='n171' href='#n171'>171</a>
<a id='n172' href='#n172'>172</a>
<a id='n173' href='#n173'>173</a>
<a id='n174' href='#n174'>174</a>
<a id='n175' href='#n175'>175</a>
<a id='n176' href='#n176'>176</a>
<a id='n177' href='#n177'>177</a>
<a id='n178' href='#n178'>178</a>
<a id='n179' href='#n179'>179</a>
<a id='n180' href='#n180'>180</a>
<a id='n181' href='#n181'>181</a>
<a id='n182' href='#n182'>182</a>
<a id='n183' href='#n183'>183</a>
<a id='n184' href='#n184'>184</a>
<a id='n185' href='#n185'>185</a>
<a id='n186' href='#n186'>186</a>
<a id='n187' href='#n187'>187</a>
<a id='n188' href='#n188'>188</a>
<a id='n189' href='#n189'>189</a>
<a id='n190' href='#n190'>190</a>
<a id='n191' href='#n191'>191</a>
<a id='n192' href='#n192'>192</a>
<a id='n193' href='#n193'>193</a>
<a id='n194' href='#n194'>194</a>
<a id='n195' href='#n195'>195</a>
<a id='n196' href='#n196'>196</a>
<a id='n197' href='#n197'>197</a>
<a id='n198' href='#n198'>198</a>
<a id='n199' href='#n199'>199</a>
<a id='n200' href='#n200'>200</a>
<a id='n201' href='#n201'>201</a>
<a id='n202' href='#n202'>202</a>
<a id='n203' href='#n203'>203</a>
<a id='n204' href='#n204'>204</a>
<a id='n205' href='#n205'>205</a>
<a id='n206' href='#n206'>206</a>
<a id='n207' href='#n207'>207</a>
</pre></td>
<td class='lines'><pre><code>