#!/usr/bin/env python3
###############################################################################
# (c) Copyright 2000-2018 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 GaudiPython import gbl
import ROOT

ROOT.gROOT.SetBatch(True)

hfile = ROOT.TFile.Open( "TrackRichReco.histos.root" )
if not hfile or hfile.IsZombie():
    print(" -> Failed to open ROOT file")
    quit()
    
fitter = gbl.Rich.Rec.CKResolutionFitter()

# Reduced range
# print fitter.params().RichFitMin
# fitter.params().RichFitMin = array.array( 'f', [-0.0048, -0.0024] )
# fitter.params().RichFitMax = array.array( 'f', [0.0046,  0.00225] )
#fitter.params().RichFitTypes[0].clear()
#fitter.params().RichFitTypes[0].push_back("Hyperbolic")
#fitter.params().RichFitTypes[0].push_back("FreeNPol")
#fitter.params().RichFitTypes[1] = fitter.params().RichFitTypes[0]

canvas = ROOT.TCanvas("CKFit", "CKFit", 1200, 900)

ROOT.gStyle.SetPadTickX(1)
ROOT.gStyle.SetPadTickY(1)
ROOT.gStyle.SetOptFit(101)
ROOT.gStyle.SetOptStat(1110)

r1hist = hfile.Get("RICH/RiCKResLong/Rich1Gas/ckResAll")
r1res = fitter.fit(r1hist, gbl.Rich.Rich1Gas)

r1hist.Draw()
r1res.overallFitFunc.SetLineColor(ROOT.kBlue + 2)
r1res.bkgFitFunc.SetLineColor(ROOT.kRed + 3)
r1res.overallFitFunc.Draw("SAME")
r1res.bkgFitFunc.Draw("SAME")
canvas.Print("RICH1-CKFit.pdf")

r2hist = hfile.Get("RICH/RiCKResLong/Rich2Gas/ckResAll")
r2res = fitter.fit(r2hist, gbl.Rich.Rich2Gas)

r2hist.Draw()
r2res.overallFitFunc.SetLineColor(ROOT.kBlue + 2)
r2res.bkgFitFunc.SetLineColor(ROOT.kRed + 3)
r2res.overallFitFunc.Draw("SAME")
r2res.bkgFitFunc.Draw("SAME")
canvas.Print("RICH2-CKFit.pdf")