"""
This script is responsible for training models and saving the best model to HDFS. The general flow is as follows:

transformed_listencounts_df is loaded from HDFS and is split into training_data, validation_data and test_data. The dataframe is converted
to an RDD and each row is converted to a Rating(row['spark_user_id'], row['recording_id'], row['transformed_listencount']) object.

Eight models are trained using the training_data RDD. Each model uses a different value of `rank`, `lambda` and `iteration`.
Refer to https://spark.apache.org/docs/2.2.0/ml-collaborative-filtering.html to know more about these params.
The Model with the least validation_rmse is called the best_model.
validation_rmse is Root Mean Squared Error calculated using the validation_data.

The best_model generated by the previous run of the script is deleted from HDFS and the new best_model is saved to HDFS.

Since the model is always trained on recently created dataframes, the model_metadata (rank, lambda, training_data_count etc) is
saved corresponding to recently created dataframe_id. The model metadata also contains the unique identification string for the best model.
"""

import logging
import time
import uuid
from collections import namedtuple
from datetime import datetime, timezone
from typing import List, Tuple

from pyspark import Row
from pyspark.ml.evaluation import RegressionEvaluator
from pyspark.ml.recommendation import ALS
from pyspark.ml.tuning import ParamGridBuilder, CrossValidator, CrossValidatorModel

import listenbrainz_spark
from listenbrainz_spark import config, utils, path, schema
from listenbrainz_spark.hdfs.utils import delete_dir
from listenbrainz_spark.hdfs.utils import path_exists
from listenbrainz_spark.exceptions import PathNotFoundException
from listenbrainz_spark.recommendations.recording.create_dataframes import describe_listencount_transformer
from listenbrainz_spark.recommendations.utils import save_html
from listenbrainz_spark.stats import run_query

logger = logging.getLogger(__name__)

Model = namedtuple('Model', 'model validation_rmse rank lmbda iteration alpha model_id')


NUM_FOLDS = 2  # number of folds to use for k-folds cross validation


def get_model_path(model_id: str):
    """ Get path to save or load model

        Args:
            model_id: Model identification string.

        Returns:
            path to save or load model.
    """
    return config.HDFS_CLUSTER_URI + path.RECOMMENDATION_RECORDING_DATA_DIR + '/' + model_id


def get_latest_dataframe_id():
    """ Get dataframe id of dataframe on which model has been trained. """
    dataframe_metadata_df = utils.read_files_from_HDFS(path.RECOMMENDATION_RECORDING_DATAFRAME_METADATA)
    dataframe_metadata_df.createOrReplaceTempView("dataframe_metadata")
    return run_query("""
        SELECT dataframe_id
          FROM dataframe_metadata
      ORDER BY dataframe_created DESC
         LIMIT 1
    """).collect()[0].dataframe_id


def preprocess_data(transformed_listencounts_df, context):
    """ Split the listen data into training set for training the model and test set for evaluating the model """
    t0 = time.monotonic()
    logger.info('Splitting dataframe...')

    # divide the data randomly into training and test in 5 : 1.
    # TODO: figure out why 45 was used as the seed in the first place.
    training_data, test_data = transformed_listencounts_df.randomSplit([5.0, 1.0], 45)

    context["train_test_ratio"] = "5 : 1"
    context["time_preprocessing"] = f"{(time.monotonic() - t0) / 60:.2f} mins"
    context["count_training_data"] = training_data.count()
    context["count_test_data"] = test_data.count()

    return training_data, test_data


def get_models(als: ALS, params: List[dict], cv_model: CrossValidatorModel) -> Tuple[Model, List[Model]]:
    """ Return the best model and list of all models along with the parameters on which the model was trained.

    Args:
        als: the ALS model to be trained
        params: list of dict of parameters on which als model is trained and tuned
        cv_model: tuned cross validator model
    """
    # Spark doesn't expose the underlying params which were used for a model directly, so we
    # need to resort to workarounds. See also: https://stackoverflow.com/a/43794841
    parent = cv_model.bestModel._java_obj.parent()
    best_model_params = {
        "rank": parent.getRank(),
        "alpha": parent.getAlpha(),
        "lmbda": parent.getRegParam(),
        "iteration": parent.getMaxIter()
    }
    best_model_metadata = None

    metadatas = []
    for param, model, metric in zip(params, cv_model.subModels[0], cv_model.avgMetrics):
        model_metadata = Model(
            model=model,
            rank=param[als.rank],
            lmbda=param[als.regParam],
            iteration=param[als.maxIter],
            alpha=param[als.alpha],
            validation_rmse=float(metric),
            model_id=f"{config.MODEL_ID_PREFIX}-{uuid.uuid4()}"
        )
        metadatas.append(model_metadata)

        if model_metadata.rank == best_model_params["rank"] \
                and model_metadata.lmbda == best_model_params["lmbda"] \
                and model_metadata.alpha == best_model_params["alpha"] \
                and model_metadata.iteration == best_model_params["iteration"]:
            best_model_metadata = model_metadata

    return best_model_metadata, metadatas


def train_models(training_data, test_data, use_transformed_listecounts, ranks, lambdas, iterations, alphas, context) \
        -> Tuple[Model, List[Model]]:
    """ Train models and get the best model.

        Args:
            training_data: Used for training.
            test_data: used for evaluating best model.
            use_transformed_listecounts: whether to use playcount or transformed listen count for training.
            ranks: Number of factors in ALS model.
            lambdas: Controls regularization.
            iterations: Number of iterations to run.
            alphas: Baseline level of confidence weighting applied.
            context: to store metadata for display in html reports.

        Returns:
            best_model: Model with least RMSE value.
            model_metadata: List of all trained models.
    """
    t0 = time.monotonic()
    logger.info("Training model.")

    ratingCol = "transformed_listencount" if use_transformed_listecounts else "playcount"

    evaluator = RegressionEvaluator(metricName="rmse", labelCol=ratingCol, predictionCol="prediction")

    als = ALS(
        userCol="spark_user_id",
        itemCol="recording_id",
        ratingCol=ratingCol,
        implicitPrefs=True,
        coldStartStrategy="drop"
    )

    params = ParamGridBuilder() \
        .addGrid(als.rank, ranks) \
        .addGrid(als.regParam, lambdas) \
        .addGrid(als.alpha, alphas) \
        .addGrid(als.maxIter, iterations) \
        .build()

    cv = CrossValidator(
        estimator=als,
        estimatorParamMaps=params,
        evaluator=evaluator,
        numFolds=NUM_FOLDS,
        collectSubModels=True,
        parallelism=3
    )
    context["num_folds"] = NUM_FOLDS
    cv_model: CrossValidatorModel = cv.fit(training_data)
    logger.info("Model trained!")

    context["time_model_training"] = f"{(time.monotonic() - t0) / 3600:.2f} hours"

    best_model, all_models = get_models(als, params, cv_model)
    context["best_model"] = best_model
    context["all_models"] = all_models
    logger.info("Best model params: %s", best_model)

    logger.info("Calculating test RMSE for best model:")
    test_predictions = best_model.model.transform(test_data)
    context["test_rmse"] = evaluator.evaluate(test_predictions)
    logger.info("Test RMSE calculated!")

    return best_model, all_models


def delete_model():
    """ Delete model.
        Note: At any point in time, only one model is in HDFS
    """
    dir_exists = path_exists(path.RECOMMENDATION_RECORDING_DATA_DIR)
    if dir_exists:
        delete_dir(path.RECOMMENDATION_RECORDING_DATA_DIR, recursive=True)


def save_model_metadata_to_hdfs(model: Model, context: dict):
    """ Add model metadata to the model metadata dataframe.

        Args:
            model: the model to save for future use.
            context: to store metadata for display in html reports.
    """
    logger.info("Saving model metadata...")
    dataframe_id = get_latest_dataframe_id()

    metadata_row = Row(
        dataframe_id=dataframe_id,
        model_created=datetime.now(timezone.utc),
        model_html_file=context["model_html_file"],
        model_id=model.model_id,
        model_param=Row(
            alpha=model.alpha,
            iteration=model.iteration,
            lmbda=model.lmbda,
            rank=model.rank,
        ),
        test_rmse=context["test_rmse"],
        validation_rmse=model.validation_rmse
    )
    model_metadata_df = utils.create_dataframe(metadata_row, schema.model_metadata_schema)

    utils.append(model_metadata_df, path.RECOMMENDATION_RECORDING_MODEL_METADATA)
    logger.info('Model metadata saved...')


def save_model(model: Model, context: dict):
    """ Save model to HDFS and add its metadata to the model metadata dataframe.

        Args:
            model: the model to save for future use.
            context: to store metadata for display in html reports.
    """
    t0 = time.monotonic()
    # delete previously saved model before saving a new model
    delete_model()

    logger.info('Saving model...')
    dest_path = get_model_path(model.model_id)
    model.model.save(dest_path)
    logger.info('Model saved!')

    save_model_metadata_to_hdfs(model, context)
    context["time_model_save"] = f"{(time.monotonic() - t0) / 60:.2f} mins"


def save_training_html(context):
    """ Prepare and save HTML report of the model training process. """
    context["listencount_transformer_description"] = describe_listencount_transformer(context["use_transformed_listencounts"])
    context["time_total"] = f"{(time.monotonic() - context['time_start']) / 3600:.2f} hours"
    save_html(context["model_html_file"], context, 'model.html')
    logger.info('Done!')


def main(ranks=None, lambdas=None, iterations=None, alphas=None, use_transformed_listencounts=False):
    if ranks is None:
        logger.critical('model param "ranks" missing')

    if lambdas is None:
        logger.critical('model param "lambdas" missing')
        raise

    if iterations is None:
        logger.critical('model param "iterations" missing')
        raise

    if alphas is None:
        logger.critical('model param "alphas" missing')
        raise

    context = {
        "time_start": time.monotonic(),
        "use_transformed_listencounts": use_transformed_listencounts
    }

    listenbrainz_spark.init_spark_session("Train Models")

    # Add checkpoint dir to break and save RDD lineage.
    listenbrainz_spark.context.setCheckpointDir(config.HDFS_CLUSTER_URI + path.CHECKPOINT_DIR)

    try:
        transformed_listencounts_df = utils.read_files_from_HDFS(
            path.RECOMMENDATION_RECORDING_TRANSFORMED_LISTENCOUNTS_DATAFRAME)
    except PathNotFoundException as err:
        logger.error('{}\nConsider running create_dataframes.py'.format(str(err)), exc_info=True)
        raise

    context["time_load_playcounts"] = f"{(time.monotonic() - context['time_start']) / 60:.2f} mins"

    training_data, test_data = preprocess_data(transformed_listencounts_df, context)

    best_model, all_models = train_models(training_data, test_data, use_transformed_listencounts,
                                          ranks, lambdas, iterations, alphas, context)

    context["model_html_file"] = f"Model-{datetime.utcnow().strftime('%Y-%m-%d-%H:%M')}-{uuid.uuid4()}.html"

    save_model(best_model, context)
    save_training_html(context)

    # Delete checkpoint dir as saved lineages would eat up space, we won't be using them anyway.
    delete_dir(path.CHECKPOINT_DIR, recursive=True)

    return [{
        'type': 'cf_recommendations_recording_model',
        'model_upload_time': str(datetime.now(timezone.utc)),
        'total_time': f"{time.monotonic() - context['time_start']:.2f} seconds",
    }]
