from sklearn.model_selection import train_test_split
import copy
# import json
import yaml
import numpy as np
import os
import argparse
import mlflow
import matplotlib.pyplot as plt
from meow_ml.src.utils.mlflow_helpers import setup_mlflow, log_params
from meow_ml.src.utils.model_helpers import init_model_factory
from meow_ml.src.utils.data_generator_helpers import load_data_generator
import optuna
from optuna.integration.mlflow import MLflowCallback

def log_study_results(study):
    output_file_name = f"{mlflow_config['STUDY_NAME']}_output.txt"
    # TODO: make output file write to specific output directory. create directory if it doesn't exist
    with open(f'./{output_file_name}', 'w+') as output_file:
        print("Number of finished trials: {}".format(len(study.trials)))
        output_file.write("Number of finished trials: {}".format(len(study.trials)) + "\n")

        print("Best trial:")
        trial = study.best_trial
        output_file.write(f"Best trial: trial # {trial.number} \n")

        print("  Value: {}".format(trial.value))
        output_file.write("  Value: {}".format(trial.value) + "\n")

        print("  Params: ")
        for key, value in trial.params.items():
            print("    {}: {}".format(key, value))

        output_file.write("  Params: " + "\n")
        for key, value in trial.params.items():
            output_file.write("    {}: {}".format(key, value) + "\n")

        mlflow.log_artifact(output_file_name, "study_results.txt")
        mlflow.set_tag("Run type", "last trial")
        # delete the temporary file
        os.remove(output_file_name)

def objective(trial, study_variables, data_generator_config):

    model_type = study_variables.get("model_type", "Sequential")  
    model_factory = init_model_factory(model_type, optuna=True)
    
    trial_variables = model_factory.get_trial_variables(trial, study_variables)
    model = model_factory.create_model(trial_variables, data_generator_config)
    model = model_factory.compile_model(model, trial_variables)

    # TODO: potentially add an option for being able to load in a previous model 
    # (such as using an encoder previously and training on the latent space instead of the original data)

    train_generator = load_data_generator(data_generator_config, 'train')
    validate_generator = load_data_generator(data_generator_config, 'validate')
    test_generator = load_data_generator(data_generator_config, 'test')

    history = model_factory.train_model(model, trial_variables, train_generator, validate_generator)
    
    # test_results = model_factory.evaluate(model, test_generator)
    # model_factory.log_metrics(test_results)

    # TODO: currently mlflow does not recognize the trial started by optuna and creates its own runs to plot the metrics
    # fix this by adding mlflowcallback.track_in_mlflow() decorator to objective function
    # to do this I need to figure out how to get mlflowcallback defined on the script-wide level cleanly
    model_factory.plot_metrics(model, test_generator, history)

    loss, accuracy = model_factory.evaluate(model, test_generator)

    return loss


def tune(study_variables, data_generator_config, mlflow_config):

    storage_name = f"sqlite:///{mlflow_config['STUDY_NAME']}.db"
    mlflc = MLflowCallback(tracking_uri=f"file://{mlflow_config['TRACKING_URI']}", 
        metric_name="Loss", mlflow_kwargs={"nested": True})
    study = optuna.create_study(direction="minimize", study_name=mlflow_config['STUDY_NAME'], storage=storage_name, load_if_exists=True)
    study.optimize(lambda trial: objective(trial, study_variables, data_generator_config), n_trials=study_variables['n_trials'], gc_after_trial=True, callbacks=[mlflc])

    log_study_results(study)



if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Optuna study runner")
    parser.add_argument(
        "--config",
        required=True,
        type=str,
        help="Path to an Optuna YAML config file.",
    )
    args = parser.parse_args()

    with open(args.config, 'r') as f:
        # config = json.load(f)
        config = yaml.safe_load(f)
        studies_config = config["studies"]
        data_generator_config = config["data_generator"]
        mlflow_config = config["mlflow"]
    
    mlflow_config = setup_mlflow(mlflow_config, optuna=True)

    # Loop through the models in the config file  
    for study_variables in studies_config:
        tune(study_variables=study_variables, data_generator_config=data_generator_config, mlflow_config=mlflow_config)
