import numpy as np
import csv 
import sys, os

#Easier to just have the full csv file parses here
from putIntoSlides_Utils import *
from putIntoSlides import *

PATH = "../2021_08_17_internalPenguin/"
NAME = "jobs.csv"

VERB = False

#Read the lines from CSV files 
def get_arr_from_CSV(filePath, fileName):
    with open(filePath+fileName) as csvfile:
        csvreader = csv.reader(csvfile, delimiter=',', quotechar='"')
        tmp_list = []        
        for row in csvreader:
            tmp_list.append(row)
        arr_jobs = np.array(tmp_list)     
    return arr_jobs

#First line are the option names
def getNameList(arr_jobs):
    return arr_jobs[0]
#The rest is actual job parameters
def getJobArray(arr_jobs):
    return arr_jobs[1:]

#Get the jobID
def getJobID(job_list, name_list):
    return job_list[name_list=='JobID']

def getJob(arr_jobs, jobID):
    names = getNameList(arr_jobs)
    jobs = getJobArray(arr_jobs)
    mask = np.array(jobs[:,names=='JobID']==str(jobID)).T[0]    
    if (mask.sum()!=1):        
        if (VERB): print("jobID",jobID,"not valid, skipping.")
        return []
    return jobs[mask,:][0]


def isFloat(tag):
    if (tag=="FLO"): return True    
    if (tag=="CON"): return True
    return False #Either it is NNN or FIX, so don't plot

def isFix(tag):
    if(tag=="FIX"): return True
    else: return False    

def isConstrained(tag):
    if(tag=="CON"): return True
    else: return False

def isRef(job_list,name_list):
    return bool(int(job_list[name_list=='ref']))
def isWeighted(job_list,name_list):
    return bool(int(job_list[name_list=='w']))

def ctkFit(job_list,name_list):
    return int(job_list[name_list=='ctkFit'])  
def ctkGen(job_list,name_list):
    return int(job_list[name_list=='ctkGen'])    
def ctkLimit(job_list,name_list):
    return float(job_list[name_list=='ctkLimit'])     
def sWave(job_list,name_list):
    return bool(int(job_list[name_list=='sWave'])) 

def sigOnly(job_list,name_list):
    return bool(int(job_list[name_list=='sigOnly']))  

def bkgOnly(job_list,name_list):
    return bool(int(job_list[name_list=='bkgOnly']))  

def stat(job_list,name_list):
    return int(job_list[name_list=='STAT'])

def getPlotVarList(job_list,name_list):
    #Check all fix/constrain/float options and create a list of all variables accordingly
    var_list = []
    if (isFloat(job_list[name_list=='ang'])): var_list = var_list + ANG_PARAMS
    if (isFloat(job_list[name_list=='Bmass'])): var_list.append("m_b")
    if (isFloat(job_list[name_list=='fSig'])): var_list.append("f_sig")
    if (isFloat(job_list[name_list=='mass'])): var_list = var_list + MASS_PARAMS[1:]#First one is Bmass
    if (isFloat(job_list[name_list=='bkg_ang'])): 
        var_list = var_list + ["cbkgctl1","cbkgctl2"] 
        for x in range(1,ctkFit(job_list,name_list)+1):
            var_list.append("cbkgctk"+str(x))
    if (isFloat(job_list[name_list=='bkg_mass'])): var_list.append("m_lambda")
    #if (isFix(job_list[name_list=='gammaStar'])): var_list.append("") TODO
    #if (isFix(job_list[name_list=='gammaStarP'])): var_list.append("") TODO
    #if (isFix(job_list[name_list=='bkg_Kstar'])): var_list.append("") TODO
    if (sWave(job_list,name_list)):
        if (isFloat(job_list[name_list=='FS'])): var_list.append("FS")
        if (isFloat(job_list[name_list=='s_ang'])): var_list = var_list + SWAVE_PARAMS[1:] #First one is FS    
    
    if (VERB): print(var_list)
    return var_list

#TODO: not ideal to be a global, but whatever now
READFILE = get_arr_from_CSV(filePath=PATH,fileName=NAME)
NAMES = getNameList(READFILE)

def toyPullsFancy(bin,jobID):
    job = getJob(READFILE,jobID)
    if (len(job) ==0): return 
        
    all_vars_list  = getPlotVarList(job_list=job, name_list=NAMES)
    figName = "/__LOOP___bin" + str(bin) +"_Signal_HighBmass_Pulls"
    figsPerSlide=8
    figsPerRow=4
    
    outFolder = "outputs/"+str(jobID)+"/"
    Path(outFolder).mkdir(parents=True, exist_ok=True)

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


def makeSlides(jobID):
    job = getJob(READFILE,jobID)
    if (len(job) ==0): return 
    nBins = 1 if (isRef(job,NAMES)) else 5

    #Open the tex file
    outputFile = "outputs/"+str(jobID)+"_all"
    print ("\\input{"+outputFile+"}")
    f = open(outputFile+".tex",'w')
    print("\\epstopdfsetup{outdir=./outputs/"+str(jobID)+"/}",file=f)
    print("\\begin{frame}{Job "+str(jobID)+"}",file=f)    
    
    print("\\begin{itemize}",file=f)    
    if (isRef(job,NAMES)):       print ("\\item Reference",file=f)
    if (isWeighted(job,NAMES)):  print ("\\item Weighted",file=f)
    print ("\\item ctk poly gen: "+str(ctkGen(job,NAMES)),file=f)
    print ("\\item ctk poly fit: "+str(ctkFit(job,NAMES)),file=f)
    print ("\\item ctk limit: "+str(ctkLimit(job,NAMES)),file=f)
    if (sigOnly(job,NAMES)):     print ("\\item Only signal",file=f)
    if (bkgOnly(job,NAMES)):     print ("\\item Only bkg",file=f)
    if (sWave(job,NAMES)):       print ("\\item sWave included",file=f)
    if (stat(job,NAMES)!=1):     print("\\item "+str(stat(job,NAMES))+"times rare stats",file=f)
    print("\\end{itemize}",file=f)    
    print("\\end{frame}",file=f)    

    for b in range(0,nBins):
        sys.stdout = open(os.devnull, 'w') #Block printout
        toyPullsFancy(b,jobID)
        sys.stdout = sys.__stdout__ #Enable printout
        print("\\input{outputs/"+str(jobID)+"/bin"+str(b)+"}",file=f)
    print("\\input{"+HD_FOLDER+"LaTeXfiles/latex_couts_toy_"+str(jobID)+"}",file=f)

    f.close()
                    
for x in range(325,400):
    makeSlides(x)