import glob
import pandas as pd
import os
import numpy as np
import pickle

dirs = glob.glob("*/")
#dirs = ['DES-SINGLE-RV/']

thingdict = {
    "DES-SINGLE-RV/":"ORVOEBV",
    "DES-TWO-RV/":'TRVOEBV',
    "DES-TWO-RV-TWO-EBV-MULTI-BETA/":"TRVTEBVMB",
    "DES-TWO-RV-TWO-EBV/":"TRVTEBV",
    "DES-INTRINSIC-DIFF/":"INTDIFF",
    "DES-NO-CINT_z0.4/":"IGN",
    "DES-NO-CINT_sig0/":"IGN",
    "DES-NO-CINT-NO-MWEBV/":"IGN",
    "DES-NO-CINT-NO-MWEBV_sig0/":"IGN",
    "DES-NO-CINT/":"IGN",
    "CC/":"CCNOCUT"
}


HMRVM = [] ; LMRVM = [] ; HMEBV = [] ; LMEBV = [] ; HMRVS = [] ; LMRVS = []

#[HMRVM, LMRVM, HMEBV, LMEBV, HMRVS, LMRVS]

def write_out(vals, thingdict, directory):
    if thingdict[directory] == "IGN":
        print(f"This directory {directory} has been set to ignore, skipping")
        return
    prefixes = ["HMRVM", "LMRVM", "HMEBV", "LMEBV", "HMRVS", "LMRVS"]
    with open(f"LATEX-OUT.tex", "a+") as f:
        for n,blep in enumerate(vals):
            #blep will be a list of tuples, need to separate it into two
            try:
                means, errs = zip(*blep)
            except ValueError:
                print(f"This directory, {directory}, is otherwise empty (perhaps this is not yet coded for the mass split results?). Skipping for now.")
                return 
            valerr = np.around(np.mean(errs),2)
            val = np.around(np.mean(means),3)
            f.write(f"{chr(92)}def{chr(92)}{thingdict[directory]}{prefixes[n]}{{${val:0.2f}({valerr:0.2f})$}}")
            f.write("\n")
    return print("Done")
try:
    print("Removing existing output file")
    os.remove("LATEX-OUT.tex")
except FileNotFoundError:
    pass
    
for d in dirs:
    files = glob.glob(f"{d}*training_popRV/*/chains.pkl")
    if len(files) == 0:
        files = glob.glob(f"{d}*training_mass_split/*/chains.pkl") 
    print(f'Found {len(files)} finished files in {d}')
    HMRVM = [] ; LMRVM = [] ; HMEBV = [] ; LMEBV = [] ; HMRVS = [] ; LMRVS = []
    for f in files:
        with open(f, "rb") as input_file:
            e = pickle.load(input_file)
        #do some processing of the data according to Mattster
        try:
            HMRVM.append([np.around(np.mean(e['mu_R']), 6   ), np.around(np.std(e['mu_R']), 6) ])
            LMRVM.append([np.around(np.mean(e['mu_R']), 6   ), np.around(np.std(e['mu_R']), 6) ])
            HMEBV.append([np.around(np.mean(e['tauA']/e['mu_R']), 6   ), np.around(np.std(e['tauA']/e['mu_R']), 6) ])
            LMEBV.append([np.around(np.mean(e['tauA']/e['mu_R']), 6   ), np.around(np.std(e['tauA']/e['mu_R']), 6) ])
            HMRVS.append([np.around(np.mean(e['sigma_R']),6 ), np.around(np.std(e['sigma_R']), 6) ])
            LMRVS.append([np.around(np.mean(e['sigma_R']),3 ), np.around(np.std(e['sigma_R']), 6) ])
        except KeyError:
            HMRVM.append([np.around(np.mean(e['mu_R_HM']), 6   ), np.around(np.std(e['mu_R_HM']), 6) ])
            LMRVM.append([np.around(np.mean(e['mu_R_LM']), 6   ), np.around(np.std(e['mu_R_LM']), 6) ])
            HMEBV.append([np.around(np.mean(e['tauA_HM']/e['mu_R_HM']), 6   ), np.around(np.std(e['tauA_HM']/e['mu_R_HM']), 6) ])
            LMEBV.append([np.around(np.mean(e['tauA_LM']/e['mu_R_LM']), 6   ), np.around(np.std(e['tauA_LM']/e['mu_R_LM']), 6) ])
            HMRVS.append([np.around(np.mean(e['sigma_R_HM']),6 ), np.around(np.std(e['sigma_R_HM']), 6) ])
            LMRVS.append([np.around(np.mean(e['sigma_R_LM']),6 ), np.around(np.std(e['sigma_R_LM']), 6) ])
    write_out([HMRVM, LMRVM, HMEBV, LMEBV, HMRVS, LMRVS], thingdict, d)