import numpy as np
import pandas as pd
import os
import time

n_methods = 14
curPath = os.path.abspath(os.path.dirname(__file__))
rootPath = curPath[:curPath.find('UITE') + len('UITE')]
file_path = os.path.join(rootPath,'results/overall.csv')

def save_recommender_result(model, dataset, result,parameter):

    detailed_path = os.path.join(rootPath, 'results/detailed_experiments')
    if not os.path.exists(detailed_path):
        os.mkdir(detailed_path)
    file = open(os.path.join(detailed_path, '{}-{}-{}.txt'.format(model, dataset,parameter)), 'w')
    pehe = result['pehe']
    ate = result['ate']
    info = 'pehe: {}\nate: {}'.format(pehe, ate)
    file.write(info)
    file.close()

def save_robustness_results(model,dataset,pehe,ate):
    robustness_path = os.path.join(rootPath, 'results/detailed_experiments')
    if not os.path.exists(robustness_path):
        os.mkdir(robustness_path)
    robustness_path = os.path.join(robustness_path,'robustness.csv')
    if os.path.exists(robustness_path) is False:
        data = {
            'Datasets':['Metric','DragonNet','BNN','TARNet',
                        'CFR_MMD','CFR_WASS','RITE_MMD','RITE_WASS',
                        'DragonNet-Robustness', 'BNN-Robustness', 'TARNet-Robustness',
                        'CFR_MMD-Robustness', 'CFR_WASS-Robustness', 'RITE_MMD-Robustness',
                        'RITE_WASS-Robustness'
                        ],
            'ACIC_1':['pehe'] + ['10000' for i in range(14)],
            'ACIC_2': ['ate'] + ['10000' for i in range(14)],
            'IHDP_1': ['pehe'] + ['10000' for i in range(14)],
            'IHDP_2': ['ate'] + ['10000' for i in range(14)],
        }
        df = pd.DataFrame(data)
        df.to_csv(robustness_path,index=False,sep=',')
    FF = True
    while FF:
        try:
            df = pd.read_csv(robustness_path)

            df['{}_1'.format(dataset)][df['Datasets'] == model] = '{}'.format(pehe)
            df['{}_2'.format(dataset)][df['Datasets'] == model] = '{}'.format(ate)
        except IOError:
            print('Waiting for writing the file')
            time.sleep(1)
        else:
            df.to_csv(robustness_path, index=False, sep=',')
            FF = False

def save_detailed_results(model,dataset,type,pehe,ate):
    detailed_path  = os.path.join(rootPath,'results/detailed_experiments')
    if not os.path.exists(detailed_path):
        os.mkdir(detailed_path)
    file = open(os.path.join(detailed_path,'{}-{}-{}.txt'.format(type,model,dataset)),'w')
    info = 'pehe: {}\nate: {}'.format(pehe,ate)
    file.write(info)
    file.close()

def save_uite_result(model,dataset,results):

    if os.path.exists(file_path) is False:
        data = {
            'Datasets':['Metric','RF','CF','CEVAE','DragonNet','S_learner','T_learner',
                        'BNN','TARNet','GANITE','CFR_MMD','CFR_WASS','QHTE','UITE_MMD','UITE_WASS'],
            'ACIC_1':['PEHE'] + ['10000' for i in range(n_methods)],
            'ACIC_2': ['ATE'] + ['10000' for i in range(n_methods)],
            'IHDP_1': ['PEHE'] + ['10000' for i in range(n_methods)],
            'IHDP_2': ['ATE'] + ['10000' for i in range(n_methods)],
            'SIM_1': ['PEHE'] + ['10000' for i in range(n_methods)],
            'SIM_2': ['ATE'] + ['10000' for i in range(n_methods)],

        }
        df = pd.DataFrame(data)
        df.to_csv(file_path,index=False,sep=',')
    FF = True
    while FF:
        try:
            df = pd.read_csv(file_path)

            df['{}_1'.format(dataset)][df['Datasets'] == model] = '{}±{}'.format(
                round(np.mean(results['pehe']), 2), round(np.std(results['pehe'], ddof=1),2))
            df['{}_2'.format(dataset)][df['Datasets'] == model] = '{}±{}'.format(
                round(np.mean(results['ate']), 2), round(np.std(results['ate'], ddof=1),2))
        except IOError:
            print('Waiting for writing the file')
            time.sleep(1)
        else:
            df.to_csv(file_path, index=False, sep=',')
            FF = False

def save_dfite_result(model,dataset,results):

    if os.path.exists(file_path) is False:

        data = {
            'Datasets':['Metric','RF','CF','CEVAE','DragonNet','S_learner','T_learner',
                        'BNN','TARNet','GANITE','CFR_MMD','CFR_WASS','QHTE','DFITE'],
            'ACIC_1':['PEHE'] + ['10000' for i in range(n_methods)],
            'ACIC_2': ['ATE'] + ['10000' for i in range(n_methods)],
            'IHDP_1': ['PEHE'] + ['10000' for i in range(n_methods)],
            'IHDP_2': ['ATE'] + ['10000' for i in range(n_methods)],
            'SIM-Z_1': ['PEHE'] + ['10000' for i in range(n_methods)],
            'SIM-Z_2': ['ATE'] + ['10000' for i in range(n_methods)],
            'SIM-ETA_1': ['PEHE'] + ['10000' for i in range(n_methods)],
            'SIM-ETA_2': ['ATE'] + ['10000' for i in range(n_methods)]
        }
        df = pd.DataFrame(data)
        df.to_csv(file_path,index=False,sep=',')
    FF = True
    while FF:
        try:
            df = pd.read_csv(file_path)

            df['{}_1'.format(dataset)][df['Datasets'] == model] = '{}±{}'.format(
                round(np.mean(results['pehe']), 2), round(np.std(results['pehe'], ddof=1),2))
            df['{}_2'.format(dataset)][df['Datasets'] == model] = '{}±{}'.format(
                round(np.mean(results['ate']), 2), round(np.std(results['ate'], ddof=1),2))
        except IOError:
            print('Waiting for writing the file')
            time.sleep(1)
        else:
            df.to_csv(file_path, index=False, sep=',')
            FF = False

def save_result(model,dataset,results):
    file_path = os.path.join(rootPath, 'results/detailed_experiments/robust.csv')
    if os.path.exists(file_path) is False:
        data = {
            'Datasets':['Metric','BNN','DragonNet','TARNet','CFR_MMD','CFR_WASS',
                        'UITE_MMD','UITE_WASS'],
            'ACIC_1':['PEHE'] + ['10000' for i in range(7)],
            'ACIC_2': ['ATE'] + ['10000' for i in range(7)],
            'IHDP_1': ['PEHE'] + ['10000' for i in range(7)],
            'IHDP_2': ['ATE'] + ['10000' for i in range(7)],
        }
        df = pd.DataFrame(data)
        df.to_csv(file_path,index=False,sep=',')
    FF = True
    while FF:
        try:
            df = pd.read_csv(file_path)

            df['{}_1'.format(dataset)][df['Datasets'] == model] = '{}±{}'.format(
                round(np.mean(results['pehe']), 2), round(np.std(results['pehe'], ddof=1),2))
            df['{}_2'.format(dataset)][df['Datasets'] == model] = '{}±{}'.format(
                round(np.mean(results['ate']), 2), round(np.std(results['ate'], ddof=1),2))
        except IOError:
            print('Waiting for writing the file')
            time.sleep(1)
        else:
            df.to_csv(file_path, index=False, sep=',')
            FF = False



if __name__ == "__main__":
    results = {'pehe':[0.45,0.67,5,6,122,212,7],'ate':[0.56,0.76,7,8,5]}
    model = 'QHTE'
    dataset = 'ACIC'
    save_result(model,dataset,results)
