import sys, os
from pathlib import Path
import numpy as np
import re #whatever it is, it is used for nth replacement

PATH = "/home/renata/B2KMuMu/sigma0/B2KstarMuMu/code/ewp-Bplus2Kstmumu-AngAna/FCNCfitter/LaTeXfiles/"
NAME = "latex_couts_toy"
VERB = False
PATH_NEW = "/home/renata/Documents/LHCb/Dizertacka/Chapters/Toys/jobs/"
NAME_NEW = "pull_table"

ANG_PARAMS  = ["Fl", "S3", "S4", "S5", "Afb", "S7", "S8", "S9"]
SWAVE_PARAMS = ["FS","SS1","SS2","SS3","SS4","SS5"]

ANGLE_FOLDING = [
        ["Fl", "S3", "Afb", "S9"],
        ["Fl", "S3", "S4"],
        ["Fl", "S3", "S5"],
        ["Fl", "S3", "S7"],
        ["Fl", "S3", "S8"]
]

ANGLE_FOLDING_S = [ 
        ["FS", "SS1"],
        ["FS", "SS1", "SS2"],
        ["FS", "SS1", "SS3"],
        ["FS", "SS1", "SS4"],
        ["FS", "SS4"],
]

COL_OLD = ["0d3c7b", "2b6d97", "499eb3", "67cfcf", "85ffeb", "7ce9c8", "72d2a5", "69b282",
               "5FA55F",
               "84BC70", "A8D281", "CCE892", "F0FEA2", "DABF7A", "C47F51", "AE4029", "970000"]

COL_WORKING = [ "0d3c7b", "2B6D97","499eb3","85ffeb","72D2A5",
               "5FA55F",  "5FA55F",  
               "a8d281", "f0fea2","c47f51", "ae4029", "970000"
                ]
COL_DEFAULT = ["2B6D97","6cb6c4","Adfff1","9ae9cd","86d2a8",
               "5FA55F",  "5FA55F",  
               "a8d281", "CCE892","f0fea2","c47f51", "ae4029"
                ]
COL_RANGE = 0.5

OBS_LIST = ["$F_{L}$", "$S_{3}$",  "$S_{4}$", "$S_{5}$", "$A_{FB}$", "$S_{7}$", "$S_{8}$", "$S_{9}$", 
            "$F_{S}$", "$S_{S1}$",  "$S_{S2}$", "$S_{S3}$", "$S_{S4}$", "$S_{S5}$"]

def read_file(jobID):
    fileName = PATH+NAME+"_"+str(jobID)+".tex"
    with open(fileName) as file: #Read file line-by-line
        lines = file.readlines()      
   
    index = lines.index('\\hline\n')
    widthLines = lines[4:index]
    pullLines = lines[index+8:-4]
    head = lines[0:4]    
    return pullLines, widthLines, head

def removeColor(line):
    tmpList = line.split('&')
    if (VERB): print (tmpList)
    varName = tmpList[0]
    varList = tmpList[1:]
    varVals =  ""
    for bin in varList:
        if(VERB): print (bin)
        varVals += "  &  $" + bin[:-1].split('$')[1] + "$" #Remove the \n, split by $ and remove the cell color
    if (VERB): print (varName,varVals)
    return (varName +varVals + "\\\\\n")

def removeColorFilter(line):
    tmpList = line.split('&')
    if (VERB): print (tmpList)
    varName = tmpList[0].strip()
    if (varName not in OBS_LIST): return "";
    varList = tmpList[1:]
    varVals =  ""
    for bin in varList:
        if(VERB): print (bin)
        varVals += "  &  $" + bin[:-1].split('$')[1] + "$" #Remove the \n, split by $ and remove the cell color
    if (VERB): print (varName,varVals)
    return (varName +varVals + "\\\\\n")

def joinLines(pullLine, widthLine):
    pullList = pullLine.split('&')
    widthList = widthLine.split('&')
    varName = pullList[0]
    try: (pullList[0]==widthList[0])
    except: print("Something went wrong when splitting the tables.")
    pullVal = pullList[1][:-1].split('$')[1] #Remove the last \n, split by $ and remove the cell color
    if(VERB): print (pullVal)
    widthVal = widthList[1][:-1].split('$')[1]
    if(VERB): print (widthVal)
    print (varName+" &  $" +pullVal+"$ & $"+widthVal+ "$\\\\")
    return 

def removeAllColor(jobID):

    pullLines, widthLines, head = read_file(jobID)

    for line in pullLines:
        print (removeColor(line))
    for line in widthLines:
        print (removeColor(line))
    return 

def joinAllLines(jobID):
    pullLines, widthLines, head = read_file(jobID)
    for pull, width in zip(pullLines, widthLines):
        joinLines(pull,width)
    return 

def addFldTag(line, fld):
    tmpList=  line.split("&")
    tmpList[0] = tmpList[0] + "("+str(fld)+")"
    return "&".join(tmpList)

def makeFancyList(input_list):
    list_filtered = []
    for lines in input_list:
        clean_lines = []
        for line in lines:
            line = removeColorFilter(line)
            clean_lines.append(line)
        list_filtered.append(clean_lines)  

    list_fancy = []
    for x in range(0,5): list_fancy.append(addFldTag(list_filtered[x][0],x))
    for x in range(0,5): list_fancy.append(addFldTag(list_filtered[x][1],x))
    list_fancy.append(list_filtered[1][2])
    list_fancy.append(list_filtered[2][2])
    list_fancy.append(list_filtered[0][2])
    list_fancy.append(list_filtered[3][2])
    list_fancy.append(list_filtered[4][2])
    list_fancy.append(list_filtered[0][3])

    return list_fancy 


def createFoldingTables(jobIDlist):
    outputName = PATH_NEW + NAME_NEW + "_" + str(jobIDlist[0]) + "_fld.tex"
    output = open(outputName, "w")

    pull_list = []
    width_list = []
    index_list = []
    
    for jobID in jobIDlist:
        pull, width, head = read_file(jobID)
        pull_list.append(pull)
        width_list.append(width)

    #print (width_list)
    #for lines in lines_list:

    
    output.write("\\begin{table} \\footnotesize \\centering\n")
    output.write("\\begin{tabular}{|l|c c c c c|}\n") #TODO: 5 bins, doesn't work for 
    output.write("\\multicolumn{6}{c}{\\textbf{means}}\\\\  \\hline  \n")
    output.write(head[2].replace("width","parameter"))
    output.write("\\hline\n")
    for line in makeFancyList(pull_list):
        output.write(line)
    output.write("\\hline\n")
    
    output.write("\\multicolumn{6}{c}{}\\\\ \n")
    output.write("\\multicolumn{6}{c}{\\textbf{widths}}\\\\  \\hline  \n") 
    output.write(head[2].replace("width","parameter"))
    output.write("\\hline\n")
    for line in makeFancyList(width_list):
        output.write(line)
    output.write("\\hline\n")
    output.write("\\end{tabular}\n")
    output.write("\\captionof{table}[The means and widths of the pull distributions in rare-like pseudoexperiments with folding applied.]{The means and widths of the pull distributions in rare-like pseudoexperiments. 500 pseudoexperiments have been generated, mimicking the rare \\BuToKstmm decay. In the fit to the pseudodata, folding is applied. The number at the parameters \\FL and $S_3$ indicate the applied folding, as the two parameters can be measured using all folding techniques. } \\label{tab:toys-Sig-pull-"+str(jobIDlist[0])+"}\n")
    output.write("\\end{table}\n\n\n")


    output.close

def createTables(jobID):
    fileName = PATH+NAME+"_"+str(jobID)+".tex"
    if not Path(fileName).exists(): return
    with open(fileName) as file: #Read file line-by-line
        lines = file.readlines()      

    if (len(lines)==0): return    
    outputName = PATH_NEW + NAME_NEW + "_" + str(jobID) + ".tex"
    output = open(outputName, "w")

   
    index = lines.index('\\hline\n')

    output.write("\\begin{table} \\footnotesize \\centering\n")
    #output.write(lines[1])
    output.write("\\begin{tabular}{|l|c c c c c|}\n") #TODO: 5 bins, doesn't work for Jpsi
    #output.write(lines[3])
    output.write("\\multicolumn{6}{c}{\\textbf{means}}\\\\  \\hline  \n") 
    output.write(lines[2].replace("width","parameter"))
    output.write("\\hline\n")
    for line in lines[index+8:-4]:
        output.write(removeColorFilter(line))
    output.write("\\hline\n")
    output.write("\\multicolumn{6}{c}{}\\\\ \n")
    output.write("\\multicolumn{6}{c}{\\textbf{widths}}\\\\ \\hline \n")
    output.write((lines[2]).replace("width","parameter"))
    output.write("\\hline\n")
    for line in lines[4:index]:
        output.write(removeColorFilter(line))
    output.write("\\hline\n")
    output.write("\\end{tabular}\n")
    output.write("\\captionof{table}[The means and widths of the pull distributions in rare-like pseudoexperiments.]{The means and widths of the pull distributions in rare-like pseudoexperiments. 500 pseudoexperiments have been generated, mimicking the rare \\BuToKstmm decay.} \\label{tab:toys-Sig-pull-"+str(jobID)+"}\n")
    output.write("\\end{table}\n\n\n")

    output.close
    
def getColor(value, isMean):
    center = 0 if isMean else 1
    ranges = np.linspace(center-COL_RANGE,center+COL_RANGE,len(COL_DEFAULT)-1)
    if (value>center+COL_RANGE): return COL_DEFAULT[-1]    
    idx = np.where((ranges>=value)==True)    
    return COL_DEFAULT[idx[0][0]]

def printTestTable():
    ranges = np.linspace(0-COL_RANGE,0+COL_RANGE,len(COL_DEFAULT)-1)
    print(ranges)
    print(ranges>=0.05)
    print(ranges>=0.55)
    print(ranges>=-0.55)
    for i in range(0,len(COL_DEFAULT)-2):
        value = (ranges[i]+ranges[i+1])/2.0
        print ("\\cellcolor[HTML]{"+str(getColor(value,1))+"}"+str(value)+"\\\\")
    return        
    
    

def replacenth(string, sub, wanted, n):
    where = [m.start() for m in re.finditer(sub, string)][n-1]
    before = string[:where]
    after = string[where:]
    after = after.replace(sub, wanted, 1)
    newString = before + after
    return newString

def addNiceColorToLine(line, isMean):
    lineList = line.split('&')
    colorList = []
    for element in lineList[1:]:
        element = element.replace("\\phantom{-}","+")
        value = element.split('$')[1].split("\\pm")[0]
        if (VERB): print(float(value))
        if (VERB): print (getColor(float(value),isMean))
        colorList.append(getColor(float(value),isMean))

    for i in range(0,len(colorList)):
        line = replacenth(line,'&','&\cellcolor[HTML]{'+colorList[i]+'}',i+1)    
    if (VERB): print(line)
    return line

        


def addNiceColors(filePath):
    outputPath = filePath.replace(".tex","_col.tex")
    output = open(outputPath, "w")
    if (VERB): print ("OutputFile: ", output)
    with open(filePath) as file: #Read file line-by-line
        lines = file.readlines()       
    indices = [i for i, e in enumerate(lines) if e == '\\hline\n']
    if (VERB): print(indices)

    NumberOfTables = len(indices)/2
    for table in range(0,int(NumberOfTables)):
        if(VERB): print("Printing table ", table)
        for i in range(indices[table*2]+1,indices[table*2+1]):
            lines[i] = addNiceColorToLine(lines[i], table==0)

    if (VERB): print(lines)

    for line in lines:
        output.write(line)

    #index = lines.index('\\hline\n')

    return 


tmpPath = "/home/renata/Documents/LHCb/Dizertacka/Chapters/Toys/jobs/pull_table_643.tex"
addNiceColors(tmpPath)
tmpPath = "/home/renata/Documents/LHCb/Dizertacka/Chapters/Toys/jobs/pull_table_644_fld.tex"
addNiceColors(tmpPath)
tmpPath = "/home/renata/Documents/LHCb/Dizertacka/Chapters/Toys/jobs/pull_table_631.tex"
addNiceColors(tmpPath)

tmpPath = "/home/renata/Documents/LHCb/Dizertacka/Chapters/Toys/jobs/pull_table_632_fld.tex"
addNiceColors(tmpPath)

printTestTable()

#if VERB: print(read_file(614)[0])
#for x in range(600,670):
#    createTables(x)

#createFoldingTables([644,645,646,647,648])
#createFoldingTables([632,633,634,635,636])

#joinAllLines(649)