#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Author: Elena Graverini
# @Date:   2015-10-27 18:26:49
# @Last Modified by:   Elena Graverini
# @Last Modified time: 2017-05-22 17:01:53
import sys
import os
from math import sqrt
import ROOT as r

location = {'TT': "/disk/data1/hep/elena/data/ST/Aging_Tuples/TT/TTaU/",
            'IT': "/disk/data1/hep/elena/data/ST/Aging_Tuples/IT/T3X2/"}
out_location = os.getcwd() + '/trackSel'
os.system('mkdir -p %s' % out_location)

# Load fill numbers
macros = os.path.expandvars('$CCEHOME/macros/CCEScan/')
with open(macros + 'Fills.dat', 'rb') as f:
    fills = f.read().splitlines()

# fills = [3478, 4518, 4643, 4856, 5162, 5448]
dets = ['TT', 'IT']
val = {'TT': 5, 'IT': 7}
layer = {'TT': 'TTaU', 'IT': 'T3X2'}
cut = {'TT': 12, 'IT': 20}
stat = 20000000

# r.gROOT.SetBatch(True)
r.gROOT.ProcessLine('.X %s/../include/lhcbstyle.C' % os.getcwd())
# r.gStyle.SetPadRightMargin(0.16)
# r.gStyle.SetPadLeftMargin(r.gStyle.GetPadLeftMargin() * 0.7)
r.gStyle.SetPadLeftMargin(0.20)
r.gStyle.SetPadRightMargin(0.25)
r.gStyle.SetPadBottomMargin(0.15)
r.gStyle.SetPadTopMargin(0.05)
# r.gStyle.SetPalette(r.kBird)
# r.gStyle.SetPalette(r.kCopper)
r.gStyle.SetPalette(r.kViridis)
r.gStyle.SetNumberContours(99)


def sqrt_of_2d_histogram(h):
    new = h.Clone()
    nbins = h.GetNbinsX() * h.GetNbinsY()
    for i in range(nbins):
        new.SetBinContent(i, sqrt(h.GetBinContent(i)))
    return new


def s_over_b(s, b):
    h = s.Clone()
    h.Divide(b)
    h.name = 's_over_b'
    return h


def s_over_sqrtb(s, b):
    h = s.Clone()
    sqrtb = sqrt_of_2d_histogram(b)
    h.Divide(sqrtb)
    h.name = 's_over_sqrtb'
    return h


def punzi_significance(s, b):
    h = s.Clone()
    sb = s.Clone()
    sb.Add(b)
    sqrt_sb = sqrt_of_2d_histogram(sb)
    h.Divide(sqrt_sb)
    h.name = 'punzi_significance'
    return h


def put_fill(c, fill):
    label = r.TPaveText( 0.74 - r.gStyle.GetPadRightMargin(), 0.87 - r.gStyle.GetPadTopMargin(),
                         0.95 - r.gStyle.GetPadRightMargin(), 0.95 - r.gStyle.GetPadTopMargin(), "BRNDC")
    label.SetFillColor(0)
    label.SetTextAlign(12)
    label.SetBorderSize(0)
    label.SetTextFont(132)
    label.SetTextSize(0.06)
    label.SetTextAlign(22)
    label.AddText('Fill %s' % fill)
    label.Draw('same')
    c.Modified()
    c.Update()
    return c


csaver = []


def printSelection(f, t_sig, det, fill, stat, good_steps=True, return_all=False):
    r.gStyle.SetPadLeftMargin(0.20)
    r.gStyle.SetPadRightMargin(0.25)
    r.gStyle.SetPadBottomMargin(0.15)
    r.gStyle.SetPadTopMargin(0.05)
    cut_sig = "(val%s>%s)" % (val[det], cut[det])
    cut_noise = "(val%s<%s)" % (val[det], cut[det])
    if good_steps:
        cut_sig += ' && (odinStep<12)'
        cut_noise += ' && (odinStep<12)'
    print(t_sig.Draw("GhostP : TrChi2/TrNDoF>>hsig(100,0., 5,100, 0.,1.)", cut_sig, "goff", stat, 0))
    h_sig = r.gDirectory.Get('hsig')
    print(t_sig.Draw("GhostP : TrChi2/TrNDoF>>hnoise(100,0., 5,100, 0.,1.)", cut_noise, "goff", stat, 0))
    h_noise = r.gDirectory.Get('hnoise')
    # c = r.TCanvas()
    h_s_over_b = s_over_b(h_sig, h_noise)
    h_s_over_b.GetZaxis().SetRangeUser(0, 10)
    h_s_over_b.GetZaxis().SetTitle('S/B ratio')
    h_s_over_sqrtb = s_over_sqrtb(h_sig, h_noise)
    h_s_over_sqrtb.GetZaxis().SetRangeUser(0, 100)
    h_s_over_sqrtb.GetZaxis().SetTitle('S/#sqrt{B}')
    h_punzi_significance = punzi_significance(h_sig, h_noise)
    h_punzi_significance.GetZaxis().SetRangeUser(0, 100)
    h_punzi_significance.GetZaxis().SetTitle('S/#sqrt{S+B}')
    for h in [h_s_over_b, h_s_over_sqrtb, h_punzi_significance]:
        csaver.append(r.TCanvas())
        csaver[-1].cd()
        h.GetXaxis().SetTitle(r'#chi^{2}_{track}/ndf')
        h.GetYaxis().SetTitle(r'Ghost Prob.')
        h.Draw('colz')
        csaver[-1].Update()
        csaver[-1], line = putLine(csaver[-1], det)
        csaver[-1] = put_fill(csaver[-1], fill)
        save_name = '%s/%s_%s_fill%s.pdf' % (out_location, det, h.name, fill)
        if good_steps:
            save_name = save_name.replace('.pdf', '_odinStep_lt_12.pdf')
        csaver[-1].SaveAs(save_name)
        if not return_all:
            h.Delete()
    # h_sig.Divide(h_noise)
    # h_sig.GetZaxis().SetRangeUser(0, 10)
    # h_sig.SetTitle('TTaU S/B ratio for fill %s' % fill)
    # h_sig.GetXaxis().SetTitle(r'#chi^{2}_{track}/ndf')
    # h_sig.GetYaxis().SetTitle(r'Ghost Prob.')
    # h_sig.GetZaxis().SetTitle(r'S/B ratio')
    # h_sig.Draw("colz")
    # c.Update()
    # c, line = putLine(c, det)
    # c = put_fill(c, fill)
    # c.SaveAs('%s/%s_fill%s.pdf' % (out_location, det, fill))
    if not return_all:
        h_sig.Delete()
        h_noise.Delete()
        return f, t_sig
    else:
        return f, t_sig, [h_s_over_b, h_s_over_sqrtb, h_punzi_significance]


def printADC(f, t_sig, det, fill, stat, step):
    if step > 10:
        print(t_sig.Draw("val%s>>hsig%s(210,-35, 175)" % (val[det], step), "odinStep==%s" % step, "goff"))  # , stat, 0))
    else:
        print(t_sig.Draw("val%s>>hsig%s(210,-35, 175)" % (val[det], step), "odinStep==%s" % step, "goff", stat, 0))
    h_sig = r.gDirectory.Get('hsig%s' % step)
    r.gROOT.ProcessLine('.X %s/../include/lhcbstyle.C' % os.getcwd())
    r.gStyle.SetPadLeftMargin(0.15)
    c = r.TCanvas()
    h_sig.GetXaxis().SetTitle('Signal height [ADC value]')
    h_sig.GetYaxis().SetTitle('Fraction of hits / ADC value')
    h_sig.DrawNormalized()
    c.Update()
    label = r.TPaveText( 0.74 - r.gStyle.GetPadRightMargin(), 0.80 - r.gStyle.GetPadTopMargin(),
                         0.95 - r.gStyle.GetPadRightMargin(), 0.95 - r.gStyle.GetPadTopMargin(), "BRNDC")
    label.SetFillColor(0)
    label.SetTextAlign(12)
    label.SetBorderSize(0)
    label.SetTextFont(132)
    label.SetTextSize(0.06)
    label.SetTextAlign(22)
    label.AddText('Fill %s\n' % fill)
    from checkPulseData import get_vmap
    v_bias = get_vmap(det)[int(step) / 6]
    label.AddText('V_{bias} = %s V' % v_bias)
    label.Draw('same')
    line = r.TLine()
    line.SetLineWidth(3)
    line.SetLineColor(r.kRed)
    line.SetLineStyle(7)
    line.DrawLine(float(cut[det]), 0.0, float(cut[det]), r.gPad.GetUymax())
    c.Modified()
    c.Update()
    c.SaveAs('%s/%s_odinStep%s_fill%s.pdf' % (out_location, det, step, fill))
    h_sig.Delete()
    return f, t_sig


def putLine(c, det):
    line = r.TLine()
    line.SetLineWidth(4)
    line.SetLineColor(r.kWhite)
    c.cd()
    if 'TT' in det:
        line.DrawLine(0.0, 0.01, 1.6, 0.16)
        line.DrawLine(3.0, 0.07, 1.6, 0.16)
        line.DrawLine(3.0, 0.07, 3.0, 0.0)
    else:
        line.DrawLine(0.0, 0.0, 3.5, 0.35)
        line.DrawLine(3.5, 0.35, 5.0, 0.20)
    return c, line


# def printByStep(fill, stat, det):
#     f = r.TFile(location[det] + '%s.root' % fill, 'read')
#     t_sig = f.Get('STADCTrackMonitor/HitInfo/TTaU')
#
#     # Look at Landau distribution
#     for step in range(67):
#         c = r.TCanvas()
#         t_sig.Draw('val5>>hnew(70,-30,90)', 'odinStep==%s' % step)
#         c.SaveAs()
#
#     draw_sig = "GhostP : TrChi2/TrNDoF>>hsig(100, 0., 10., 100, 0., 1.)"
#     draw_bg = "GhostP : TrChi2/TrNDoF>>hnoise(100, 0., 10., 100, 0., 1.)"
#     f.Close()

from math import exp, pow


def bigauss(x, y, par):
    # Definition from:
    # https://en.wikipedia.org/wiki/Multivariate_normal_distribution
    # Params:
    # 0 - norm
    # 1 - mu_x
    # 2 - sigma_x
    # 3 - mu_y
    # 4 - sigma_y
    # 5 - rho
    return par[0] * exp(-0.5 / (1 - pow(par[5], 2)) * (pow((x - par[1]) / par[2], 2) + pow((y - par[3]) / par[4], 2) - 2 * par[5] * (x - par[1]) * (y - par[3]) / par[2] / par[4]))


def wrap_bigauss(x, par):
    def get_TF2(xy, par):
        x, y = xy[0], xy[1]
        return bigauss(x, y, par)
    func = get_TF2(x, par)
    func.SetParNames('norm', 'mu_x', 'sigma_x',
                     'mu_y', 'sigma_y', 'rho')
    func.SetParameters(1., 1., 1., 0.01, 0.05, 0.3)
    return r.TF2('bivariate_gauss', func, 0., 5., 0., 1.)


class MyFitFunction():
    def __init__(self, func, n_var):
        self.n_var = n_var
        self.func = func

    def wrap_func(self, x, par):
        args = []
        for i in range(self.n_var):
            args.append(x[i])
        args = tuple(args)
        # print(x[0], x[1], x[2], x[3], par)
        # return self.func(*x, par=par)
        return self.func(*args, par=par)

    def __call__(self, x, par):
        return self.wrap_func(x, par)


def myTFunc(name, classname, func, func_range, par_names, par_inits=[], par_limits=[]):
    n_pars = len(par_names)
    f = r.__getattr__(classname)(name, MyFitFunction(func, int(classname[2])), *(func_range + (n_pars,)))
    f.SetParNames(*par_names)
    if par_inits:
        f.SetParameters(*par_inits)
    for i in range(len(par_limits)):
        f.SetParLimits(i, *par_limits[i])
    return f


def make_all_plots():
    for det in dets:
        for fill in fills:
            f = r.TFile(location[det] + '%s.root' % fill, 'read')
            t_sig = f.Get('STADCTrackMonitor/HitInfo/%s' % layer[det])

            f, t_sig = printSelection(f, t_sig, det, fill, stat, good_steps=True)
            if '--steps' in sys.argv:
                for step in [0, 12, 48, 54, 60]:
                    f, t_sig = printADC(f, t_sig, det, fill, stat, step)

            f.Close()


# Dictionary which items are accessible with d.key
class mydict(dict):
    __getattr__ = dict.__getitem__
    __setattr__ = dict.__setitem__

if __name__ == '__main__':
    # make_all_plots()
    det = 'TT'
    fill = '5448'
    f = r.TFile(location[det] + '%s.root' % fill, 'read')
    t_sig = f.Get('STADCTrackMonitor/HitInfo/%s' % layer[det])
    f, t_sig, h_list = printSelection(f, t_sig, det, fill, stat, good_steps=True, return_all=True)
    punzi = h_list[-1]
    conf = mydict({
        'names': ('norm', 'mu_x', 'sigma_x', 'mu_y', 'sigma_y', 'rho'),
        'inits': (200, 1., 1., 0.01, 0.05, 0.3),
        'limits': [(0, 100000), (0, 5), (0, 5), (-0.3, 1), (0, 0.5), (-1, 1)],
        'name': 'bivariate_gaussian',
        'range': (0., 5., 0., 1.),
    })
    # r.gStyle.SetNumberContours(20)
    func = myTFunc(conf.name, 'TF2', bigauss, conf.range, conf.names, conf.inits, conf.limits)
    punzi.Fit(func)
    c = r.TCanvas()
    punzi.Draw('surf3')
    print('Fitted parameters:')
    print('x ~ N({}, {})'.format(func.GetParameter('mu_x'), func.GetParameter('sigma_x')))
    print('y ~ N({}, {})'.format(func.GetParameter('mu_y'), func.GetParameter('sigma_y')))