import os
from glob import glob
import json



def list_all_directory(suffix, noise, checkpoints, alpha=False):
    path = "../*/"  # not forget / after *
    list_folder = glob(path)
    if noise:
        if checkpoints:
            if alpha:
                return [i for i in list_folder if (suffix in i and 'alpha' in i and 'noise' in i and "checkpoint" in i)]
            else:
                return [i for i in list_folder if
                        (suffix in i and 'alpha' not in i and 'noise' in i and "checkpoint" in i)]
        else:
            if alpha:
                return [i for i in list_folder if
                        (suffix in i and 'alpha' in i and 'noise' in i and "checkpoint" not in i)]
            else:
                return [i for i in list_folder if
                        (suffix in i and 'alpha' not in i and 'noise' in i and "checkpoint" not in i)]
    else:
        if checkpoints:
            if alpha:
                return [i for i in list_folder if
                        (suffix in i and 'alpha' in i and "checkpoint" in i and 'noise' not in i)]
            else:
                return [i for i in list_folder if
                        (suffix in i and 'alpha' not in i and "checkpoint" in i and 'noise' not in i)]
        else:
            if alpha:
                return [i for i in list_folder if
                        (suffix in i and 'alpha' in i and "checkpoint" not in i and 'noise' not in i)]
            else:
                return [i for i in list_folder if
                        (suffix in i and 'alpha' not in i and "checkpoint" not in i and 'noise' not in i)]


def read_disantanglement(path_file):
    with open(path_file, 'r') as file:
        lines = file.readlines()

    acc = float(lines[1].split('\t')[-1].replace(':', '').replace('\n', ''))
    loss = float(lines[-1].split('\t')[-1].replace(':', '').replace('\n', ''))
    return acc, loss


def read_content(path_file):
    with open(path_file, 'r') as file:
        lines = file.readlines()

    ppl = float(lines[1].split('\t')[-1].replace(':', '').replace('\n', ''))
    w_overlap = float(lines[-1].split('\t')[-1].replace(':', '').replace('\n', ''))
    return ppl, w_overlap


def find_test_content_name(path):
    list_folder = glob('{}/*'.format(path))
    return [i for i in list_folder if ('checkpoint' in i and 'txt' in i)]


if __name__ == '__main__':

    prefix = 'no_reny_clement'

    all_directory = list_all_directory(prefix, noise=True, checkpoints=False)
    dict_no_reny_noise_dis = dict()
    dict_no_reny_noise_cont = dict()
    dict_no_reny_no_noise_dis = dict()
    dict_no_reny_no_noise_cont = dict()

    for directory in all_directory:
        test_content_files = find_test_content_name(directory)
        for test_content_file in test_content_files:
            print(test_content_file)
            lambda_ = test_content_file.split('/')[1].split('_')[-2]  # sauf les virgules
            if "lambda" in lambda_:
                lambda_ = float(lambda_.replace('lambda', ''))
            else:
                if "lambda0" in test_content_file.split('/')[1].split('_')[-3]:
                    lambda_ = float(lambda_) * 0.1
            print(lambda_)
            lambda_ = float(lambda_)
            if dict_no_reny_noise_cont.get(float(lambda_)) != None:
                dict_no_reny_noise_cont[float(lambda_)].append(read_content(test_content_file))
            else:
                dict_no_reny_noise_cont[lambda_] = [read_content(test_content_file)]

    all_directory = list_all_directory(prefix, noise=True, checkpoints=True)

    for directory in all_directory:
        test_content_files = find_test_content_name(directory)
        for test_content_file in test_content_files:
            print(test_content_file)
            lambda_ = test_content_file.split('/')[1].split('_')[-5]  # sauf les virgules
            if "lambda" in lambda_:
                lambda_ = float(lambda_.replace('lambda', ''))
            else:
                if "lambda0" in test_content_file.split('/')[1].split('_')[-6]:
                    lambda_ = float(lambda_) * 0.1
            print(lambda_)
            lambda_ = float(lambda_)
            if dict_no_reny_noise_dis.get(float(lambda_)) != None:
                dict_no_reny_noise_dis[float(lambda_)].append(read_content(test_content_file))
            else:
                dict_no_reny_noise_dis[lambda_] = [read_content(test_content_file)]

    all_directory = list_all_directory(prefix, noise=False, checkpoints=False)
    for directory in all_directory:
        test_content_files = find_test_content_name(directory)
        for test_content_file in test_content_files:
            print(test_content_file)
            lambda_ = test_content_file.split('/')[1].split('_')[-1]  # sauf les virgules

            if "lambda" in lambda_:
                lambda_ = float(lambda_.replace('lambda', ''))
            else:
                if "lambda0" in test_content_file.split('/')[1].split('_')[-2]:
                    lambda_ = float(lambda_) * 0.1
            print(lambda_)
            lambda_ = float(lambda_)
            if dict_no_reny_no_noise_cont.get(float(lambda_)) != None:
                dict_no_reny_no_noise_cont[float(lambda_)].append(read_content(test_content_file))
            else:
                dict_no_reny_no_noise_cont[lambda_] = [read_content(test_content_file)]

    all_directory = list_all_directory(prefix, noise=False, checkpoints=True)

    for directory in all_directory:
        test_content_files = find_test_content_name(directory)
        for test_content_file in test_content_files:
            print(test_content_file)
            lambda_ = test_content_file.split('/')[1].split('_')[-4]  # sauf les virgules
            if "lambda" in lambda_:
                lambda_ = float(lambda_.replace('lambda', ''))
            else:
                if "lambda0" in test_content_file.split('/')[1].split('_')[-5]:
                    lambda_ = float(lambda_) * 0.1
            print(lambda_)
            lambda_ = float(lambda_)
            if dict_no_reny_no_noise_dis.get(float(lambda_)) != None:
                dict_no_reny_no_noise_dis[float(lambda_)].append(read_disantanglement(test_content_file))
            else:
                dict_no_reny_no_noise_dis[lambda_] = [read_disantanglement(test_content_file)]


    with open('dict_no_reny_noise_dis.json', 'w') as f:
        json.dump(dict_no_reny_noise_dis, f)

    with open('dict_no_reny_noise_cont.json', 'w') as f:
        json.dump(dict_no_reny_noise_cont, f)

    with open('dict_no_reny_no_noise_dis.json', 'w') as f:
        json.dump(dict_no_reny_no_noise_dis, f)

    with open('dict_no_reny_no_noise_cont.json', 'w') as f:
        json.dump(dict_no_reny_no_noise_cont, f)


    prefix = '/reny_clement'

    all_directory = list_all_directory(prefix, noise=True, checkpoints=False,alpha=False)
    dict_no_reny_noise_dis = dict()
    dict_no_reny_noise_cont = dict()
    dict_no_reny_no_noise_dis = dict()
    dict_no_reny_no_noise_cont = dict()

    for directory in all_directory:
        test_content_files = find_test_content_name(directory)
        for test_content_file in test_content_files:
            print(test_content_file)
            lambda_ = test_content_file.split('/')[1].split('_')[-2]  # sauf les virgules
            if "lambda" in lambda_:
                lambda_ = float(lambda_.replace('lambda', ''))
            else:
                if "lambda0" in test_content_file.split('/')[1].split('_')[-3]:
                    lambda_ = float(lambda_) * 0.1
            print(lambda_)
            if dict_no_reny_noise_cont.get(float(lambda_)) != None:
                dict_no_reny_noise_cont[float(lambda_)].append(read_content(test_content_file))
            else:
                dict_no_reny_noise_cont[lambda_] = [read_content(test_content_file)]

    all_directory = list_all_directory(prefix, noise=True, checkpoints=True,alpha=False)

    for directory in all_directory:
        test_content_files = find_test_content_name(directory)
        for test_content_file in test_content_files:
            print(test_content_file)
            lambda_ = test_content_file.split('/')[1].split('_')[-5]  # sauf les virgules
            if "lambda" in lambda_:
                lambda_ = float(lambda_.replace('lambda', ''))
            else:
                if "lambda0" in test_content_file.split('/')[1].split('_')[-6]:
                    lambda_ = float(lambda_) * 0.1
            print(lambda_)
            if dict_no_reny_noise_dis.get(float(lambda_)) != None:
                dict_no_reny_noise_dis[float(lambda_)].append(read_content(test_content_file))
            else:
                dict_no_reny_noise_dis[lambda_] = [read_content(test_content_file)]

    all_directory = list_all_directory(prefix, noise=False, checkpoints=False,alpha=False)
    for directory in all_directory:
        test_content_files = find_test_content_name(directory)
        for test_content_file in test_content_files:
            print(test_content_file)
            lambda_ = test_content_file.split('/')[1].split('_')[-1]  # sauf les virgules

            if "lambda" in lambda_:
                lambda_ = float(lambda_.replace('lambda', ''))
            else:
                if "lambda0" in test_content_file.split('/')[1].split('_')[-2]:
                    lambda_ = float(lambda_) * 0.1
            print(lambda_)
            if dict_no_reny_no_noise_cont.get(float(lambda_)) != None:
                dict_no_reny_no_noise_cont[float(lambda_)].append(read_content(test_content_file))
            else:
                dict_no_reny_no_noise_cont[lambda_] = [read_content(test_content_file)]

    all_directory = list_all_directory(prefix, noise=False, checkpoints=True,alpha=False)

    for directory in all_directory:
        test_content_files = find_test_content_name(directory)
        for test_content_file in test_content_files:
            print(test_content_file)
            lambda_ = test_content_file.split('/')[1].split('_')[-4]  # sauf les virgules
            if "lambda" in lambda_:
                lambda_ = float(lambda_.replace('lambda', ''))
            else:
                if "lambda0" in test_content_file.split('/')[1].split('_')[-5]:
                    lambda_ = float(lambda_) * 0.1
            print(lambda_)
            if dict_no_reny_no_noise_dis.get(float(lambda_)) != None:
                dict_no_reny_no_noise_dis[float(lambda_)].append(read_disantanglement(test_content_file))
            else:
                dict_no_reny_no_noise_dis[lambda_] = [read_disantanglement(test_content_file)]


    with open('dict_reny_1.3_noise_dis.json', 'w') as f:
        json.dump(dict_no_reny_noise_dis, f)

    with open('dict_reny_1.3_noise_cont.json', 'w') as f:
        json.dump(dict_no_reny_noise_cont, f)

    with open('dict_reny_1.3_no_noise_dis.json', 'w') as f:
        json.dump(dict_no_reny_no_noise_dis, f)

    with open('dict_reny_1.3_no_noise_cont.json', 'w') as f:
        json.dump(dict_no_reny_no_noise_cont, f)

    prefix = '/reny_clement'

    all_directory = list_all_directory(prefix, noise=True, checkpoints=False, alpha=True)
    dict_no_reny_noise_dis = dict()
    dict_no_reny_noise_cont = dict()
    dict_no_reny_no_noise_dis = dict()
    dict_no_reny_no_noise_cont = dict()

    for directory in all_directory:
        test_content_files = find_test_content_name(directory)
        for test_content_file in test_content_files:
            print(test_content_file)
            lambda_ = test_content_file.split('/')[1].split('_')[-3]  # sauf les virgules
            if "lambda" in lambda_:
                lambda_ = float(lambda_.replace('lambda', ''))
            else:
                if "lambda0" in test_content_file.split('/')[1].split('_')[-4]:
                    lambda_ = float(lambda_) * 0.1
            print(lambda_)
            if dict_no_reny_noise_cont.get(float(lambda_)) != None:
                dict_no_reny_noise_cont[float(lambda_)].append(read_content(test_content_file))
            else:
                dict_no_reny_noise_cont[lambda_] = [read_content(test_content_file)]

    all_directory = list_all_directory(prefix, noise=True, checkpoints=True, alpha=True)

    for directory in all_directory:
        test_content_files = find_test_content_name(directory)
        for test_content_file in test_content_files:
            print(test_content_file)
            lambda_ = test_content_file.split('/')[1].split('_')[-6]  # sauf les virgules
            if "lambda" in lambda_:
                lambda_ = float(lambda_.replace('lambda', ''))
            else:
                if "lambda0" in test_content_file.split('/')[1].split('_')[-7]:
                    lambda_ = float(lambda_) * 0.1
            print(lambda_)
            if dict_no_reny_noise_dis.get(float(lambda_)) != None:
                dict_no_reny_noise_dis[float(lambda_)].append(read_content(test_content_file))
            else:
                dict_no_reny_noise_dis[lambda_] = [read_content(test_content_file)]

    all_directory = list_all_directory(prefix, noise=False, checkpoints=False, alpha=True)
    for directory in all_directory:
        test_content_files = find_test_content_name(directory)
        for test_content_file in test_content_files:
            print(test_content_file)
            lambda_ = test_content_file.split('/')[1].split('_')[-2]  # sauf les virgules

            if "lambda" in lambda_:
                lambda_ = float(lambda_.replace('lambda', ''))
            else:
                if "lambda0" in test_content_file.split('/')[1].split('_')[-3]:
                    lambda_ = float(lambda_) * 0.1
            print(lambda_)
            if dict_no_reny_no_noise_cont.get(float(lambda_)) != None:
                dict_no_reny_no_noise_cont[float(lambda_)].append(read_content(test_content_file))
            else:
                dict_no_reny_no_noise_cont[lambda_] = [read_content(test_content_file)]

    all_directory = list_all_directory(prefix, noise=False, checkpoints=True, alpha=True)

    for directory in all_directory:
        test_content_files = find_test_content_name(directory)
        for test_content_file in test_content_files:
            print(test_content_file)
            lambda_ = test_content_file.split('/')[1].split('_')[-5]  # sauf les virgules
            if "lambda" in lambda_:
                lambda_ = float(lambda_.replace('lambda', ''))
            else:
                if "lambda0" in test_content_file.split('/')[1].split('_')[-6]:
                    lambda_ = float(lambda_) * 0.1
            print(lambda_)
            if dict_no_reny_no_noise_dis.get(float(lambda_)) != None:
                dict_no_reny_no_noise_dis[float(lambda_)].append(read_disantanglement(test_content_file))
            else:
                dict_no_reny_no_noise_dis[lambda_] = [read_disantanglement(test_content_file)]

    with open('dict_reny_1_01_noise_dis.json', 'w') as f:
        json.dump(dict_no_reny_noise_dis, f)

    with open('dict_reny_1_01_noise_cont.json', 'w') as f:
        json.dump(dict_no_reny_noise_cont, f)

    with open('dict_reny_1_01_no_noise_dis.json', 'w') as f:
        json.dump(dict_no_reny_no_noise_dis, f)

    with open('dict_reny_1_01_no_noise_cont.json', 'w') as f:
        json.dump(dict_no_reny_no_noise_cont, f)


