#!/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: 2015-12-07 11:46:05

import ROOT as r
import os

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)

fills = [2797, 3108, 3478, 3960, 4518]
#fills = [4643]
dets = ['TT', 'IT']
stat = 10000000

r.gROOT.ProcessLine('.X %s/../include/lhcbstyle.C' % os.getcwd())
r.gStyle.SetPadRightMargin(0.12)

def printSelection(det, fill, stat):
    f = r.TFile(location[det] + '%s.root' % fill, 'read')
    t_sig = f.Get('STADCTrackMonitor/HitInfo/TTaU')
    # t_bg = f.Get('STADCTrackMonitor/NoiseInfo/TTaU')
    # t_sig.Draw("GhostP : TrChi2/TrNDoF>>hnew(100,0., 5,100, 0.,1.)", "", "colz", 10000, 0)

    t_sig.Draw("GhostP : TrChi2/TrNDoF>>hsig(100,0., 5,100, 0.,1.)", "val5>12", "colz", stat, 0)
    h_sig = r.gPad.GetPrimitive('hsig')
    t_sig.Draw("GhostP : TrChi2/TrNDoF>>hnoise(100,0., 5,100, 0.,1.)", "val5<12", "colz", stat, 0)
    h_noise = r.gPad.GetPrimitive('hnoise')

    c = r.TCanvas()
    h_sig.Divide(h_noise)
    h_sig.Draw("colz")
    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")

    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.SetTextFont(22)
    label.SetTextSize(0.06)
    label.SetTextAlign(22)
    # label.SetTextColor(r.kWhite)
    label.AddText('Fill %s' % fill)
    label.Draw('same')
    c.Modified()
    c.Update()

    c.SaveAs('%s/%s_fill%s.pdf' % (out_location, det, fill))
    f.Close()


def printByStep(fill, stat):
    f = r.TFile(location_tt + '%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()


if __name__ == '__main__':
    for det in dets:
        for fill in fills:
            printSelection(fill, stat)