import os

import lightning as L
import pandas as pd

from configs.exp import hparams
from datasets.loader.datamodule import EhrDataModule
from datasets.loader.load_los_info import get_los_info
from pipelines import DlPipeline, MlPipeline

def get_latest_file(path):
    # Get list of all files in the directory
    files = [os.path.join(path, f) for f in os.listdir(path) if os.path.isfile(os.path.join(path, f))]
    
    # Get the file with the latest modification time
    latest_file = max(files, key=os.path.getctime)
    
    return latest_file

def run_ml_experiment(config):
    los_config = get_los_info(f'datasets/{config["dataset"]}/processed/fold_{config["fold"]}')
    config.update({"los_info": los_config})
    # data
    dm = EhrDataModule(f'datasets/{config["dataset"]}/processed/fold_{config["fold"]}', batch_size=config["batch_size"])
    # train/val/test
    pipeline = MlPipeline(config)
    trainer = L.Trainer(accelerator="cpu", max_epochs=1, logger=False, num_sanity_val_steps=0)
    trainer.test(pipeline, dm)
    perf = pipeline.test_performance
    return perf

def run_dl_experiment(config):
    los_config = get_los_info(f'datasets/{config["dataset"]}/processed/fold_{config["fold"]}')
    config.update({"los_info": los_config})

    # data
    dm = EhrDataModule(f'datasets/{config["dataset"]}/processed/fold_{config["fold"]}', batch_size=config["batch_size"])
    # checkpoint
    # checkpoint_path = f'logs/train/{config["dataset"]}/{config["task"]}/{config["model"]}-fold{config["fold"]}-seed{config["seed"]}/checkpoints/best.ckpt'

    checkpoint_path = get_latest_file(f'logs/train/{config["dataset"]}/{config["task"]}/{config["model"]}-fold{config["fold"]}-seed{config["seed"]}/checkpoints')

    print("checkpoint_path: ", checkpoint_path)

    if "time_aware" in config and config["time_aware"] == True:
        checkpoint_path = f'logs/train/{config["dataset"]}/{config["task"]}/{config["model"]}-fold{config["fold"]}-seed{config["seed"]}-ta/checkpoints/best.ckpt'

    # train/val/test
    pipeline = DlPipeline(config)
    trainer = L.Trainer(accelerator="cpu", max_epochs=1, logger=False, num_sanity_val_steps=0)
    trainer.test(pipeline, dm, ckpt_path=checkpoint_path)
    perf = pipeline.test_performance
    return perf

if __name__ == "__main__":
    best_hparams = hparams # [TO-SPECIFY]
    performance_table = {'dataset':[], 'task': [], 'model': [], 'fold': [], 'seed': [], 'accuracy': [], 'auroc': [], 'auprc': [], 'f1': [], 'minpse': []}
    for i in range(0, len(best_hparams)):
    # for i in range(0, 1):
        config = best_hparams[i]
        print(f"Testing... {i}/{len(best_hparams)}")
        run_func = run_ml_experiment if config["model"] in ["RF", "DT", "GBDT", "XGBoost", "CatBoost", "LR", "LightGBM"] else run_dl_experiment
        seeds = [0] # [0,1,2,3,4]
        folds = ['nshot']
        for fold in folds:
            config["fold"] = fold
            for seed in seeds:
                config["seed"] = seed
                perf = run_func(config)
                print(f"{config}, Test Performance: {perf}")

                if "time_aware" in config and config["time_aware"] == True:
                    model_name = config['model']+"_ta"
                else:
                    model_name = config['model']

                performance_table['dataset'].append(config['dataset'])
                performance_table['task'].append(config['task'])
                performance_table['model'].append(model_name)
                performance_table['fold'].append(config['fold'])
                performance_table['seed'].append(config['seed'])
                if config['task'] == 'outcome':
                    performance_table['accuracy'].append(perf['accuracy'])
                    performance_table['auroc'].append(perf['auroc'])
                    performance_table['auprc'].append(perf['auprc'])
                    performance_table['f1'].append(perf['f1'])
                    performance_table['minpse'].append(perf['minpse'])
    pd.DataFrame(performance_table).to_csv('ijcai24_ml_baselines_20240108.csv', index=False) # [TO-SPECIFY]