#!/usr/bin/env python3

###############################################################################
# (c) Copyright 2000-2020 CERN for the benefit of the LHCb Collaboration      #
#                                                                             #
# This software is distributed under the terms of the GNU General Public      #
# Licence version 3 (GPL Version 3), copied verbatim in the file "COPYING".   #
#                                                                             #
# In applying this licence, CERN does not waive the privileges and immunities #
# granted to it by virtue of its status as an Intergovernmental Organization  #
# or submit itself to any jurisdiction.                                       #
###############################################################################

from array import array

import ROOT
from GaudiKernel.SystemOfUnits import GeV
from GaudiPython import gbl

ROOT.gROOT.SetBatch(True)

fname = "TrackRichReco.histos.root"
print("Processing", fname)
root_file = ROOT.TFile.Open(fname)
if not root_file or root_file.IsZombie():
    print(" -> Failed to open ROOT file")
    quit()


fitter = gbl.Rich.Rec.CKResolutionFitter()

# need to deduce this from root file somehow...
nCKBins = 100

MinP = 0.5 * GeV
MaxP = 120 * GeV

# ROOT style options
ROOT.gStyle.SetPadTickX(1)
ROOT.gStyle.SetPadTickY(1)
ROOT.gStyle.SetOptFit(101)
ROOT.gStyle.SetOptStat(1110)
# lhcbAxisLabelOffset = 0.3
# ROOT.gStyle.SetTitleOffset( lhcbAxisLabelOffset, "x" )
# ROOT.gStyle.SetTitleOffset( lhcbAxisLabelOffset, "y" )
lhcbAxisLabelSize = 0.03
ROOT.gStyle.SetTitleSize(lhcbAxisLabelSize, "x")
ROOT.gStyle.SetTitleSize(lhcbAxisLabelSize, "y")
ROOT.gStyle.SetTitleSize(lhcbAxisLabelSize, "z")
lhcbFont = 62
ROOT.gStyle.SetTextFont(lhcbFont)
lhcbTextSize = 0.03
ROOT.gStyle.SetTextSize(lhcbTextSize)
lhcbLabelSize = 0.024
ROOT.gStyle.SetLabelSize(lhcbLabelSize, "x")
ROOT.gStyle.SetLabelSize(lhcbLabelSize, "y")
ROOT.gStyle.SetLabelSize(lhcbLabelSize, "z")
ROOT.gStyle.SetStatBorderSize(1)
ROOT.gStyle.SetStatFont(lhcbFont)
lhcbStatFontSize = 0.025
ROOT.gStyle.SetStatFontSize(lhcbStatFontSize)

# Make a canvas to draw on
canvas = ROOT.TCanvas("CKPBinFit", "CKPBinFit", 1200, 900)
canvas.SetGrid()

for tag in ["RiCKResLong"]:
    for rad in ["Rich1Gas", "Rich2Gas"]:
        # for region in ["Both", "Inner", "Outer"]:
        for region in ["Both"]:
            for part in ["All", "pion", "kaon", "proton", "electron", "muon"]:
                # Fits PDF filename
                fitsPDF = tag + "-CKFits-" + rad + "-" + region + "-" + part + ".pdf"
                canvas.Print(fitsPDF + "[")

                # loop over all the momentum bins and fit each
                pmeans = []
                pmeanErrs = []
                ckresFits = []
                ckresErrs = []
                # P bins
                pmeans_add = []
                h_add = []
                for iBin in range(0, nCKBins):
                    # load the histogram
                    if region == "Both":
                        partTag = "" if part == "All" else "/" + part
                        hid = (
                            "RICH/"
                            + tag
                            + "/"
                            + rad
                            + partTag
                            + "/PtotBins/CKRes/AllPhots/Bin"
                            + str(iBin)
                        )
                    else:
                        hid = (
                            "RICH/"
                            + tag
                            + "/"
                            + rad
                            + "/PtotBins/CKRes/"
                            + region
                            + "/AllPhots/Bin"
                            + str(iBin)
                        )
                    h = root_file.Get(hid)
                    if not h:
                        print("FAILED to load ", hid)
                        exit(0)
                    else:
                        # Parse the title to get CK bin min and max values
                        r = h.GetTitle().split(" ")
                        if region == "Both":
                            if part == "All":
                                pmin = float(r[5])
                                pmax = float(r[7])
                            else:
                                pmin = float(r[6])
                                pmax = float(r[8])
                        else:
                            pmin = float(r[10])
                            pmax = float(r[12])
                        pmean = 0.5 * (pmin + pmax)
                        # print("Bin",iBin," Ptot",pmin,"to",pmax)

                        minEnts = 300000
                        okToFit = h.GetEntries() > minEnts
                        if not okToFit:
                            # See if we have enough if we add the last skipped plots together
                            sum_ents = h.GetEntries()
                            for hh in h_add:
                                sum_ents += hh.GetEntries()
                            if sum_ents > minEnts:
                                for hh in h_add:
                                    h.Add(hh)
                                for pm in pmeans_add:
                                    pmean += pm
                                pmean = pmean / (len(pmeans_add) + 1)
                                okToFit = True
                            else:
                                # add this hist to list to add
                                pmeans_add += [pmean]
                                h_add += [h]

                        # Do we have enough now to fit ?
                        if okToFit:
                            # Fit the hist
                            fit_res = fitter.fit(
                                h,
                                gbl.Rich.Rich1Gas
                                if rad == "Rich1Gas"
                                else gbl.Rich.Rich1Gas,
                            )
                            if fit_res.fitOK:
                                h.Draw()
                                fit_res.overallFitFunc.SetLineColor(ROOT.kBlue + 2)
                                fit_res.bkgFitFunc.SetLineColor(ROOT.kRed + 3)
                                fit_res.overallFitFunc.Draw("SAME")
                                fit_res.bkgFitFunc.Draw("SAME")
                                canvas.Print(fitsPDF)

                                if abs(fit_res.ckShift) < 0.001:
                                    pmeans += [pmean]
                                    pmeanErrs += [0]
                                    ckresFits += [fit_res.ckResolution]
                                    ckresErrs += [fit_res.ckResolutionErr]
                            # clear the averaged plots
                            pmeans_add = []
                            h_add = []

                # Make a graph of the results
                n = len(pmeans)
                if n > 0:
                    x = array("f", pmeans)
                    xe = array("f", pmeanErrs)
                    y = array("f", ckresFits)
                    ye = array("f", ckresErrs)
                    gr = ROOT.TGraphErrors(n, x, y, xe, ye)
                    if rad == "Rich1Gas":
                        gr.GetXaxis().SetLimits(1.0, 70.0)
                    else:
                        gr.GetXaxis().SetLimits(2.0, 120.0)
                    gr.SetTitle(
                        rad
                        + " "
                        + region
                        + " "
                        + part
                        + " CK Theta Resolution Versus Momentum"
                    )
                    gr.GetXaxis().SetTitle("Track Momentum / GeV/c")
                    gr.GetYaxis().SetTitle("CK Theta Resolution / rad")
                    gr.SetMarkerColor(4)
                    gr.SetMarkerStyle(21)
                    gr.Draw("ALP")
                    canvas.Print(fitsPDF)

                # Finally make the inclusive fit plot(s)
                if region == "Both":
                    hid = "RICH/" + tag + "/" + rad + "/ckResAll"
                else:
                    hid = "RICH/" + tag + "/" + rad + "/ckResAll" + region
                h = root_file.Get(hid)
                if not h:
                    print("FAILED to load ", hid)
                else:
                    fit_res = fitter.fit(
                        h, gbl.Rich.Rich1Gas if rad == "Rich1Gas" else gbl.Rich.Rich1Gas
                    )
                    if fit_res.fitOK:
                        h.Draw()
                        fit_res.overallFitFunc.SetLineColor(ROOT.kBlue + 2)
                        fit_res.bkgFitFunc.SetLineColor(ROOT.kRed + 3)
                        fit_res.overallFitFunc.Draw("SAME")
                        fit_res.bkgFitFunc.Draw("SAME")
                        canvas.Print(fitsPDF)

                canvas.Print(fitsPDF + "]")