"""Training flow for MEOW-ML."""

import time

import mlflow
import mlflow.tensorflow

from ..utils.data_generator_helpers import load_data_generator
from ..utils.model_helpers import init_model_factory


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)

    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
        validation_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
        validation_generator = load_data_generator(
            data_generator_config,
            "validate",
            scalers=train_scalers,
        )
        test_generator = load_data_generator(
            data_generator_config,
            "test",
            scalers=train_scalers,
        )

    with mlflow.start_run(
        experiment_id=mlflow_config["EXPERIMENT_ID"],
        run_name=model_config.get("model_name"),
        description=model_config.get("model_description", ""),
        nested=False,
    ):
        model_factory.autolog()
        run_id = mlflow.active_run().info.run_id
        print(f"STARTING RUN: {run_id}")
        if "config_file_path" in mlflow_config:
            mlflow.log_artifact(mlflow_config["config_file_path"])
        mlflow.log_param("run_name", model_config.get("model_name"))

        model_factory.extras_before_training(
            model,
            train_generator,
            validation_generator,
            test_generator,
        )

        if model_type in {"Sequential", "Keras", "Functional"}:
            model_factory.summary(model)

        t0 = time.time()
        history = model_factory.train_model(
            model,
            train_generator,
            validation_generator,
            model_config,
        )
        total_time = time.time() - t0
        mlflow.log_param("Training time", total_time)
        print(f"Total time for model.fit() processing is: {total_time} seconds.")

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

        model_factory.predict(model, test_generator)
        eval_results = model_factory.evaluate(model, test_generator)
        print(f"eval_results: {eval_results}")
        model_factory.log_metrics(eval_results)
        model_factory.plot_metrics(
            model,
            train_generator,
            validation_generator,
            test_generator,
            history,
            data_generator_config.get("test_type", "cat"),
            data_generator_config.get("scale_data_method"),
        )
        model_factory.extras_after_training(
            model,
            train_generator,
            validation_generator,
            test_generator,
            data_generator_config.get("test_type", "cat"),
        )
        print("Finished successfully")
