import os
import time
# import json
import yaml
import argparse

# TODO: replace tensorflow.keras imports with 'keras_core as keras' for backend ambiguity
import mlflow.tensorflow
from meow_ml.src.utils.mlflow_helpers import setup_mlflow, log_params
from meow_ml.src.utils.data_generator_helpers import load_data_generator
from meow_ml.src.utils.model_helpers import init_model_factory

###############################################################################################
########################### Training Model ####################################################
###############################################################################################

def train(model_config, data_generator_config, mlflow_config):
    model_type = model_config.get("model_type", "Sequential")
    model_factory = init_model_factory(model_config, model_type=model_type)
    model = model_factory.create_model(model_config)
    model = model_factory.compile_model(model, model_config)

    # 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)

    # TODO: uncomment this section if y scaling proves useless
    # train_generator = load_data_generator(data_generator_config, 'train')
    # # get scalers from train generator to use for validation and test generators
    # train_scalers = train_generator.scalers
    # validate_generator = load_data_generator(data_generator_config, 'validate', scalers=train_scalers)
    # test_generator = load_data_generator(data_generator_config, 'test', scalers=train_scalers)

    if data_generator_config.get('scale_y_labels', False):
        train_generator = load_data_generator(
            data_generator_config,
            'train',
            scalers=None,
            label_scaler=None,
        )
        train_scalers = train_generator.scalers
        train_label_scaler = train_generator.label_scaler

        validate_generator = load_data_generator(
            data_generator_config,
            'validate',
            scalers=train_scalers,
            label_scaler=train_label_scaler,
        )

        test_generator = load_data_generator(
            data_generator_config,
            'test',
            scalers=train_scalers,
            label_scaler=train_label_scaler,
        )
    else:
        train_generator = load_data_generator(data_generator_config, 'train')
        train_scalers = train_generator.scalers
        validate_generator = load_data_generator(data_generator_config, 'validate', scalers=train_scalers)
        test_generator = load_data_generator(data_generator_config, 'test', scalers=train_scalers)

    # probably not needed - will delete soon
    # run_name = get_unique_run_name(mlflow_config, model_config.get('model_name', None))
    with mlflow.start_run(
        experiment_id=mlflow_config['EXPERIMENT_ID'],
        run_name=model_config.get("model_name", None),
        description=model_config.get("model_description", ""),
        nested=False
    ) as run:

        model_factory.autolog()
        # commenting since I now log the full config
        # log_params(model_config)
        RUN_ID = mlflow.active_run().info.run_id
        print(f"STARTING RUN: {RUN_ID}")
        mlflow.log_artifact(mlflow_config['config_file_path'])
        mlflow.log_param("run_name", model_config.get("model_name", None))

        # any extra/miscellaneous functions to be done BEFORE training and testing
        model_factory.extras_before_training(model, train_generator, validate_generator, test_generator)

        t0 = time.time()

        # TODO: figure out a way to make this if statement more generic
        # problem is keras model summary best done before training, xgboost summary only able to be done after
        if model_type == "Sequential" or model_type == "Keras" or model_type == "Functional":
            model_factory.summary(model)
        history = model_factory.train_model(model, train_generator, validate_generator, model_config)

        t1 = time.time()
        total_time = t1 - t0

        mlflow.log_param("Training time", total_time)

        print()
        print("Total time for model.fit() processing is: " + str(total_time) + " seconds.")

        if model_type == "XGBoost":
            model_factory.summary(model)

        predictions = model_factory.predict(model, test_generator)
        eval_results = model_factory.evaluate(model, test_generator)
        # debugging
        print(f'eval_results: {eval_results}')

        model_factory.log_metrics(eval_results)
        model_factory.plot_metrics(
            model,
            train_generator,
            validate_generator,
            test_generator,
            history,
            data_generator_config.get('test_type', 'cat'),
            data_generator_config.get('scale_data_method', None)
        )

        # any extra/miscellaneous functions to be done AFTER training and testing
        model_factory.extras_after_training(
            model,
            train_generator,
            validate_generator,
            test_generator,
            data_generator_config.get('test_type', 'cat')
        )

        print("Finished successfully")


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Flags for settings')
    parser.add_argument(
        '--config',
        required=True,
        type=str,
        help='config file path for loading settings',
    )
    args = parser.parse_args()

    with open(args.config, 'r') as f:
        # config = json.load(f)
        config = yaml.safe_load(f)
        models_config = config["models"]
        data_generator_config = config["data_generator"]
        mlflow_config = config["mlflow"]
        gpu_config = config["gpu"]

    # setting gpus here so that debugging doesn't take up all GPUs
    gpus = gpu_config.get('GPUS', None)
    if gpus:
        os.environ["CUDA_VISIBLE_DEVICES"] = gpus

    mlflow_config = setup_mlflow(mlflow_config, config_file_path=args.config)

    # Loop through the models in the config file
    for model_config in models_config:
        train(model_config=model_config, data_generator_config=data_generator_config, mlflow_config=mlflow_config)
