from __future__ import print_function
import numpy as np 
from pathlib import Path

from putIntoFigures_Utils import *

FIG_FOLDER = "FCNC/"
MASS_FIG_FOLDER = FIG_FOLDER + "MassFit/"
MC_FIG_FOLDER = FIG_FOLDER + "MCfit/"
TOYS_FIG_FOLDER = FIG_FOLDER + "Toys/"
MAIN_FIG_FOLDER = FIG_FOLDER + "MainFit/"
BKG_FIG_FOLDER = FIG_FOLDER + "bkgFit/"
GENLVL_FIG_FOLDER = FIG_FOLDER + "GenLvlFit/"

TEX_FILE_FOLDER = "/home/renata/Documents/LHCb/Dizertacka/Chapters/ParameterMeasurement/"
         #Don't use ./ as it efs up the print of includes

OUTPUT_TAG = ""
TITLE_TAG = ""

decName = "KplusPi0Resolved"
ANGLES = ["ctk", "ctl", "phi"]
MASS_PLOTS = ["m","mkpi"]
ANG_PARAMS  = ["Fl", "S3", "S4", "S5", "Afb", "S7", "S8", "S9"]
MASS_PARAMS = ["m_b", "m_sigma_1", "alpha_1", "alpha_2", "n_1", "n_2"]
MASS_BKG = ["m_lambda","f_sig"]
ANGLE_BKG = ["cbkgctk1","cbkgctk2","cbkgctk3","cbkgctk4","cbkgctl1","cbkgctl2"] 
SWAVE_PARAMS = ["FS","SS1","SS2","SS3","SS4","SS5"]



def projections(angle,Run): #Plot all angle projections in 18 bins
    listOfEntries = list(map(str,np.arange(0,18,1)))
    label = angle + "-projections-"+str(Run)
    outputName = angle +  "proj_" + RunTag(Run)
    figsPerSlide=6
    figsPerRow=3
    make_figure(FIG_FOLDER+"/angular/",angle+"eff__LOOP___" + decName + "_" + RunTag(Run), listOfEntries, label, figsPerSlide, figsPerRow, ".eps", outputName+".tex")


def massFit(nBins, observable, Run): #Plot mass-only fits
    massFitName = "_massfit" + rareJpsiTag(False) 
    path = observable + massFitName + "Fit" + binTag(nBins) + "_bin__LOOP___"+ RunTag(Run) + ""
    figsPerSlide=6
    figsPerRow=3
    
    make_figure(MASS_FIG_FOLDER,path,
                       list(map(str,np.arange(0,nBins,1))),
                       "Signal-mass-fit",
                       figsPerSlide, figsPerRow,
                       ".eps",TEX_FILE_FOLDER+"MassFit_Signal_"+observable+"_"+RunTag(Run)+".tex")
    return 


def massFit1Bin(Run): #Plot mass-only fits
    massFitName = "_massfit" + rareJpsiTag(False) 
    path = "__LOOP__" + massFitName + "Fit" + binTag(1) + "_bin0_" + RunTag(Run) + ""
    figsPerSlide=2
    figsPerRow=2
    
    make_figure(MASS_FIG_FOLDER,path,
                       MASS_PLOTS,
                       "Signal-mass-fit",
                       figsPerSlide, figsPerRow,
                       ".eps",TEX_FILE_FOLDER+"MassFit_Signal_"+RunTag(Run)+".tex")
    return 

def massFitRef(Run): 
    massFitName = "_massfit" + rareJpsiTag(True) 
    path = "__LOOP__" + massFitName + "Fit" + binTag(1) + "_bin0_" + RunTag(Run) + ""
    figsPerSlide=2
    figsPerRow=2
    
    make_figure(MASS_FIG_FOLDER,path,
                       MASS_PLOTS,
                       "Reference-mass-fit",
                       figsPerSlide, figsPerRow,
                       ".eps",TEX_FILE_FOLDER+"MassFit_Jpsi_"+RunTag(Run)+".tex")
    return 

def massFitMCRef(Run): #Plot only mass fit of reference MC 
    MCfitName = "_MC" + rareJpsiTag(True) + "Fit"
    path = "__LOOP__" + MCfitName + binTag(1) + "_bin0_SimultaneousFit_" + RunTag(Run) + "_OnlyMass"
    figsPerSlide=2
    figsPerRow=2
    
    make_figure(MC_FIG_FOLDER + "Jpsi/",path,
                       MASS_PLOTS,
                       "MC-Ref-massFit",
                       figsPerSlide, figsPerRow,
                       ".eps",TEX_FILE_FOLDER+"MC_MassFit_Jpsi_"+RunTag(Run)+".tex")
    return 
    
def massFitMC(nBins, observable, Run): 
    MCfitName = "_MC" + rareJpsiTag(False) + "Fit"
    path = observable + MCfitName + binTag(nBins) + "_bin__LOOP___SimultaneousFit_" + RunTag(Run) 
    figsPerSlide=5
    figsPerRow=3
    
    make_figure(MC_FIG_FOLDER + "Signal/",path,
                       list(map(str,np.arange(0,nBins,1))),
                       "MC-Sig-massFit",
                       figsPerSlide, figsPerRow,
                       ".eps",TEX_FILE_FOLDER+"MC_MassFit_Signal_"+observable+"_"+RunTag(Run)+".tex")
    return 

def angleFitMC(nBins, bin, Run): #Plot only ANGLES fit of signal MC
    MCfitName = "_MC" + rareJpsiTag(False) + "Fit"
    path = "__LOOP__" + MCfitName + binTagFull(nBins,bin)+"_SimultaneousFit_" + RunTag(Run) + "_OnlyANGLES_AllPDFs"
    figsPerSlide=3
    figsPerRow=3
    
    make_figure(MC_FIG_FOLDER + "Signal/",path,
                       ANGLES,"MC-Sig-fit-"+binTagFull(nBins,bin),
                       figsPerSlide, figsPerRow,
                       ".eps",TEX_FILE_FOLDER+"MC_Signal"+binTagFull(nBins,bin)+"_"+RunTag(Run)+"_onlyANGLES.tex")
    return 

def MainFitRef(Run, Toy, fold, fullBkg): 
    fitName = rareJpsiTag(True) + "Fit"
    path = "__LOOP__" + fitName + toyTag(Toy)+ binTag(1) + "_bin"+str(0)+"_SimultaneousFit" + foldTag(fold) + "_" + RunTag(Run) + plotBkgTag(fullBkg)+"_AllPDFs"
    figsPerSlide=6
    figsPerRow=3

    label = "MainFit-Ref"+("-toy" if Toy else "")
    if (fold != -1): label += "-fld" + str(fold)
    
    make_figure(MAIN_FIG_FOLDER, path,
                       ANGLES+MASS_PLOTS,label,
                       figsPerSlide, figsPerRow,
                       ".eps",TEX_FILE_FOLDER+"MainFit_Ref"+toyTag(Toy)+ foldTag(fold)+ "_"+RunTag(Run)+plotBkgTag(fullBkg)+".tex")
    return 

def MainFit(Run, nBins, bin, fold, fullBkg): 
    fitName =  rareJpsiTag(False) + "Fit"
    path = "__LOOP__" + fitName + binTagFull(nBins,bin) +"_SimultaneousFit" + foldTag(fold) +"_" + RunTag(Run) + plotBkgTag(fullBkg)+ "_AllPDFs"
    figsPerSlide=6
    figsPerRow=3

    label = "MainFit-Sig"+("-toy" if Toy else "")+binTagFull(nBins,bin)
    if (fold != -1): title += "-fld" + str(fold)
    
    make_figure(MAIN_FIG_FOLDER, path,
                       ANGLES+MASS_PLOTS, title,
                       figsPerSlide, figsPerRow,
                       ".eps",TEX_FILE_FOLDER+"MainFit"+ binTagFull(nBins,bin) + foldTag(fold)+"_"+RunTag(Run)+plotBkgTag(fullBkg)+".tex")
    return 

def ToyFitCheckRef():
    path = "Init_vs_Fit___LOOP___Ref"
    plot_list = ANG_PARAMS + ANGLE_BKG +MASS_PARAMS + MASS_BKG 
    figsPerSlide=6
    figsPerRow=3

    make_figure("./figures/FullFitToy/",path,
                       plot_list,"ToyCheck-Reflike",
                       figsPerSlide, figsPerRow,
                       ".eps",TEX_FILE_FOLDER+"ToyInitFit.tex")
    return 

def MainFitToyCheck(isRef):
    refTag = "_Ref" if isRef else ""
    path = "FinalToy_Init_vs_Fit___LOOP__" + refTag
    plot_list = ANG_PARAMS + ANGLE_BKG +MASS_PARAMS + MASS_BKG 
    figsPerSlide=6
    figsPerRow=3

    make_figure(TOYS_FIG_FOLDER,path,                       
                       plot_list,"ToyCheck-Reflike" if isRef else "ToyCheck-Siglike",
                       figsPerSlide, figsPerRow,
                       ".eps",TEX_FILE_FOLDER+"MainFitToy"+ refTag+".tex")
    return 

def MCfit(nBins, bin, Run, fold): #Plot MC fit everything
    MCfitName = "_MC" + rareJpsiTag(False) + "Fit"
    
    path = "__LOOP__" + MCfitName + binTagFull(nBins,bin)+"_SimultaneousFit" + foldTag(fold) +"_" + RunTag(Run) + "_AllPDFs"
    figsPerSlide=6
    figsPerRow=3
    
    label = "MC-sigFit-"+binTagFull(nBins,bin)
    if (fold != -1): label += ", fld " + str(fold)
    make_figure(MC_FIG_FOLDER + "Signal/",path,
                       ANGLES+MASS_PLOTS, label,
                       figsPerSlide, figsPerRow,
                       ".eps",TEX_FILE_FOLDER+"MC_Signal"+binTagFull(nBins,bin)+ foldTag(fold)+"_"+RunTag(Run)+".tex")
    return 

def BkgFitRef(Run, sideband): 

    var_list = ["ctk", "ctl", "phi", "mkpi"]

    massFitName = "_Bckgnd" + rareJpsiTag(True) 
    sidebandTag = ""
    if (sideband ==-1):  sidebandTag = "_LowBmass"
    elif (sideband ==0): sidebandTag = "_HighBmass"

    sidebandTitle = ""
    if (sideband ==-1):  sidebandTitle = ", lower sideband"
    elif (sideband ==0): sidebandTitle = ", upper sideband"


    path = "__LOOP__" + massFitName + "Fit" + binTag(1) + "_bin0_" + RunTag(Run)+ sidebandTag
    figsPerSlide=4
    figsPerRow=4
    
    make_figure(BKG_FIG_FOLDER,path,
                       var_list,"bkgFit"+sidebandTag,
                       figsPerSlide, figsPerRow,
                       ".eps",TEX_FILE_FOLDER+"BkgFit_Jpsi_"+RunTag(Run)+sidebandTag+".tex")
    return 

def genLvlMC(nBins, bin, Run, fold):
    genLvlName = "_genLvl_MC" + rareJpsiTag(False) + "Fit"
    path = "__LOOP__" + genLvlName +  binTagFull(nBins,bin)+ "_OnlyANGLES_AllPDFs"
    figsPerSlide=6
    figsPerRow=3

    title = "GenMCfit"+binTagFull(nBins,bin)
    make_figure(GENLVL_FIG_FOLDER,path,
                       ANGLES, title,
                       figsPerSlide, figsPerRow,
                       ".eps",TEX_FILE_FOLDER+"GenLvlMC"+binTagFull(nBins,bin)+ foldTag(fold)+"_"+RunTag(Run)+".tex")

    return

def compareGenLvlAndMC(toDavid):
    plotName = "GenLvl_vs_MC_"
    path = plotName + "__LOOP__" + ("_MyVsDavids" if toDavid else "")
    figsPerSlide=8
    figsPerRow=2

    label = "GenLvl-vs-"
    if (toDavid): label += "David"
    else:         label += "sigMC"

    make_figure(MC_FIG_FOLDER,path,
                       ANG_PARAMS, label,
                       figsPerSlide, figsPerRow,
                       ".eps",TEX_FILE_FOLDER+"GenLvl" + ("_Vs_Davids" if toDavid else "_Vs_SigMC")+".tex")
    return

def compareFolding():
    path = "Folding___LOOP__" 
    figsPerSlide=8
    figsPerRow=2

    title = "Fold-comparison"

    make_figure(MC_FIG_FOLDER,path,
                       ANG_PARAMS, title,
                       figsPerSlide, figsPerRow,
                       ".eps",TEX_FILE_FOLDER+"Folding_comparison.tex")
    return

def KstarMass(Ref,nBins, Run):
    MCfitName = "mkpi_MC" + rareJpsiTag(Ref) + "Fit"
    
    path = MCfitName + binTag(nBins) +"_bin__LOOP__" +"_SimultaneousFit_" + RunTag(Run) + ("_AllPDFs" if (Run == 12) else "") + ("_OnlyMass" if (Ref) else "") 
    figsPerSlide=6
    figsPerRow=3
    
    binsList = list(map(str,np.arange(0,nBins,1)))
    title = "MC "+("Jpsi" if Ref else "Signal") +" fit, "+str(nBins) + " bin" + (", " if nBins==1 else "s, ") + RunTag(Run)
    #TODO omg fix the title
    make_figure(MC_FIG_FOLDER + ("Jpsi/" if Ref else "Signal/"),path,
                       binsList, title,
                       figsPerSlide, figsPerRow,
                       ".eps",TEX_FILE_FOLDER+"MC"+ rareJpsiTag(Ref)+"_Kpi"+binTag(nBins)+ "_"+RunTag(Run)+".tex")
    return 

#######################

def toyPulls(bin,jobID,sWave=False):
    all_vars_list  = ANG_PARAMS + ANGLE_BKG + MASS_BKG
    if(sWave): all_vars_list = all_vars_list + SWAVE_PARAMS

    figName = "/__LOOP___bin" + str(bin) +"_Signal_HighBmass_Pulls"
    figsPerSlide=8
    figsPerRow=4
    
    outFolder = "outputs/"+str(jobID)+"/"
    Path(outFolder).mkdir(parents=True, exist_ok=True)

    label = str(jobID)+"bin "+str(bin)
    make_figure(TOYS_FIG_FOLDER + str(jobID),figName,
                       all_vars_list, label,
                       figsPerSlide, figsPerRow,
                       ".eps",outFolder+"bin"+str(bin)+".tex")

#######################

massFitRef(1)
massFitRef(2)

massFit(1,"mkpi",1)
massFit(1,"mkpi",2)


compareGenLvlAndMC(True)
compareGenLvlAndMC(False)
compareFolding()


totBins = 5


for b in range(totBins):
    MCfit(totBins,b,12,-1)

for b in range(totBins):
    genLvlMC(totBins,b,12,-1)
for b in range(8):
    genLvlMC(8,b,12,-1)

compareGenLvlAndMC(False)
compareGenLvlAndMC(True)

BkgFitRef(12,0)

massFitRef(1)
massFitRef(2)
massFit1Bin(1)
massFit1Bin(2)
for obsv in MASS_PLOTS:
    massFit(totBins,obsv,1)
    massFit(totBins,obsv,2)

    massFit(1,obsv,1)
    massFit(1,obsv,2)

for b in range(totBins):
    MCfit(totBins,b,12,0)
    MCfit(totBins,b,12,1)
    MCfit(totBins,b,12,2)
    MCfit(totBins,b,12,3)
    MCfit(totBins,b,12,4)
    
compareFolding()

MainFitToyCheck(True)
MainFitRef(12,True,-1,0)
MainFitRef(12,False,-1,0)
MainFitRef(12,False,0,0)
MainFitRef(12,False,1,0)
MainFitRef(12,False,2,0)
MainFitRef(12,False,3,0)
