# Renata Kopecna

import uproot
import numpy as np
import matplotlib.pyplot as plt

import sys
sys.path.insert(0,'/home/lhcb/kopecna/B2KstarMuMu_clean/Code/Selection/ComparisonTool')
#import Utils
from Utils3 import getTreeList

def plotMultiple(year, Run, Reference, PHSP, verbose):

    #Check if 2015
    is2015 = (year == 2015 and Run == 0)

    #Load the branch 
    dataTM = uproot.lazyarray(getTreeList(2016 if (is2015) else year, Run, True, True, Reference, PHSP, verbose),"DecayTreeTruthMatched","IsAloneAt_TMed_rndGamma")
    dataMC = uproot.lazyarray(getTreeList(2016 if (is2015) else year, Run, True, True, Reference, PHSP, verbose),"DecayTreeTruthMatched","IsAloneAtNotTM")
    data = uproot.lazyarray(getTreeList(year, Run, False, True, Reference, PHSP, verbose),"DecayTree","IsAloneAt")

    #Switch to numpy array (check if neded)
    arr_MLPresponse = np.array(data, dtype=float)
    arr_MLPresponseMC = np.array(dataMC, dtype=float)
    arr_MLPresponseTM = np.array(dataTM, dtype=float)

    nEvts = arr_MLPresponse.size
    nMCEvts = arr_MLPresponseMC.size
    nTMEvts = arr_MLPresponseTM[arr_MLPresponseTM!=2].size
    if (verbose):
        print ("nEvts: ", nEvts)
        print ("nMCEvts: ", nMCEvts)
        print ("nTMEvts: ", nTMEvts)

    #Define a figure
    fig = plt.figure()

    #Define steps in x-axis
    arr_steps = np.arange(0.9,1.0,0.001)
    multiple_fraction = []
    multiple_fraction_MC = []
    multiple_fraction_TM = []
    if (verbose): print("Steps in MLP:",arr_steps)

    #Get the fractions
    for m in arr_steps:
        multiple_fraction.append(1-arr_MLPresponse[arr_MLPresponse<=m].size / float(nEvts))
        multiple_fraction_MC.append(1-arr_MLPresponseMC[arr_MLPresponseMC<=m].size / float(nMCEvts))
        multiple_fraction_TM.append(1-arr_MLPresponseTM[arr_MLPresponseTM<=m].size / float(nTMEvts))
    
    plt.plot(arr_steps,multiple_fraction,marker=".",label = "Data")
    plt.plot(arr_steps,multiple_fraction_MC,marker=".",label = "MC")
    plt.plot(arr_steps,multiple_fraction_TM,marker=".",label = "TM MC")
    plt.ylim(0.0,0.21)
    plt.legend(loc="upper right", prop={'size': 14})
    tag = "Run " + str(Run) if (Run != 0) else "Year " + str(year)      
    plt.text(0.905,0.175,tag, fontsize=24)

    plt.xlabel('MVA response cut',fontsize=14)
    plt.ylabel('Multiple candidates fraction',fontsize=14)
    name = "Run_" + str(Run) if (Run != 0) else "year_" + str(year)      
    fig.savefig(name,dpi = fig.dpi, bbox_inches='tight')

    return

def plotKplusmumu(year,Run,MC,Reference,PHSP,verbose):

    #Check if 2015
    is2015 = (year == 2015 and Run == 0)  
    
    if (Reference or PHSP): MC = True
    treeName = "DecayTreeTruthMatched" if (MC) else "DecayTree"
    
    #Load the branch 
    branch_list = ["PX", "PY", "PZ", "PE"]
    particle_list = ["K_plus", "mu_plus", "mu_minus"]
    array_list = []
    for part in particle_list:
        for branch in branch_list:
            branchName = part+"_"+branch+"_DTF"
            array_list.append(np.array(uproot.lazyarray(getTreeList(2016 if (is2015) else year, Run, MC, MC, Reference, PHSP, verbose),treeName,branchName)),dtype=float)

    nEvts = array_list[0].size
    
#plotMultiple(2011,1,False, False,True)
#plotMultiple(2016,2,False, False,True)

for yr in [2011,2012,2015,2016,2017,2018]:
    plotMultiple(yr,0,False,False,True)