# -*- coding: utf-8 -*-
# Copyright (c) Hebes Intelligence Private Company

# This source code is licensed under the Apache License, Version 2.0 found in the
# LICENSE file in the root directory of this source tree.

import logging
from typing import List, Union

import pandas as pd

from eensight.methods.prediction.activity import estimate_activity
from eensight.methods.prediction.baseline import UsagePredictor
from eensight.methods.prediction.metrics import cvrmse, nmbe

logger = logging.getLogger("eensight")


def create_activity_feature(
    features: pd.DataFrame,
    labels: pd.DataFrame,
    non_occ_features: Union[str, List[str]],
    cat_features: Union[str, List[str]],
    exog: str,
    assume_hurdle: bool,
    n_trials: int,
    verbose: bool,
    adjusted_activity: pd.DataFrame,
):
    """Estimate activity levels.

    Args:
        features (pd.DataFrame): The feature data.
        labels (pd.DataFrame): The target data.
        non_occ_features (Union[str, List[str]]): The names of the features to
            use to fit the two quantile regression models on the label data. If
            not provided, "temperature" will be used.
        cat_features (Union[str, List[str]]): The name(s) of the categorical
            features in the input data.
        exog (str): The name of the feature to use for determining how to best
            split the data into two regimes. It should not contain categorical data.
        assume_hurdle (bool): If True, activity estimation will assume that energy
            consumption has been generated by a hurdle model.
        n_trials (int): The number of iterations for the underlying optimization.
        verbose (bool): If True, a progress bar is visible.
        adjusted_activity (pd.DataFrame): Placeholder to pass adjusted activity
            levels if they exist.

    Returns:
        pandas.DataFrame: The estimated activity time series.
    """

    if isinstance(adjusted_activity, pd.DataFrame):
        logger.warning(
            "Skipping activity estimation because adjusted activity was found"
        )
        return adjusted_activity

    activity, best_params = estimate_activity(
        features,
        labels,
        non_occ_features=non_occ_features,
        cat_features=cat_features,
        exog=exog,
        assume_hurdle=assume_hurdle,
        n_trials=n_trials,
        verbose=verbose,
        return_params=True,
    )

    logger.info(f"Estimated parameters: {best_params}")
    return activity.to_frame("activity")


def fit_predict(
    features: pd.DataFrame,
    labels: pd.DataFrame,
    lags: dict,
    cat_features: Union[str, List[str]],
    validation_size: float,
):
    """Fit a predictive model and apply it on the training data.

    Args:
        features (pd.DataFrame): The feature data.
        labels (pd.DataFrame): The target data.
        lags (dict): Dictionary with keys that correspond to feature names,
            and values that are lists containing the time lags of that
            feature to add as additional regressors.
        cat_features (str or list of str): The names of the categorical features
            in the input data.
        validation_size (float): The size of the validation dataset (as
            percentage of all training data) to use for identifying optimal
            number of iterations for the Gradient Boosting model.

    Returns:
        tuple containing

        - **model** (*UsagePredictor*): The fitted model.
        - **prediction** (*pandas.DataFrame*): The in-sample prediction.
        - **metrics** (*dict*): The in-sample CV(RMSE) and NMBE metrics.
    """

    model = UsagePredictor(
        lags=lags,
        cat_features=cat_features,
        validation_size=validation_size,
        skip_calendar=False,
    )
    model = model.fit(features, labels)
    logger.info(f"Model scores: {model.scores_}")

    prediction = model.predict(features).to_frame("consumption")

    metrics = {
        "CVRMSE": cvrmse(labels, prediction),
        "NMBE": nmbe(labels, prediction),
    }
    logger.info(f"In-sample metrics: {metrics}")

    return model, prediction, metrics


def predict(features: pd.DataFrame, model: UsagePredictor):
    """Predict energy consumption given input data.

    Args:
        features (pd.DataFrame): The feature data.
        model (UsagePredictor): The fitted model.

    Returns:
        pandas.DataFrame: The prediction.
    """
    prediction = model.predict(features).to_frame("consumption")
    return prediction


def fit_predict_autoenc(
    features: pd.DataFrame,
    labels: pd.DataFrame,
    activity: pd.DataFrame,
    lags: dict,
    cat_features: Union[str, List[str]],
    validation_size: float,
    assume_hurdle: bool,
):
    """Fit an autoencoder model and apply it on the training data.

    Args:
        features (pd.DataFrame): The feature data.
        labels (pd.DataFrame): The target data.
        activity (pd.DataFrame): The estimated activity levels.
        lags (dict): Dictionary with keys that correspond to feature names,
            and values that are lists containing the time lags of that
            feature to add as additional regressors.
        cat_features (str or list of str): The names of the categorical features
            in the input data.
        validation_size (float): The size of the validation dataset (as
            percentage of all training data) to use for identifying optimal
            number of iterations for the Gradient Boosting model.
        assume_hurdle (bool): If True, activity estimation will assume that energy
            consumption has been generated by a hurdle model.

    Returns:
        tuple containing

        - **model** (*UsagePredictor*): The fitted model.
        - **prediction** (*pandas.DataFrame*): The in-sample prediction.
        - **metrics** (*dict*): The in-sample CV(RMSE) and NMBE metrics.
    """

    if assume_hurdle:
        selected = activity["activity"] == 1
        X_selected = features[selected]
        y_selected = labels[selected]

        model = UsagePredictor(
            lags=lags,
            cat_features=cat_features,
            validation_size=validation_size,
            skip_calendar=False,
        )
        model = model.fit(X_selected, y_selected)
        logger.info(f"Model scores: {model.scores_}")

        prediction = model.predict(features).to_frame("consumption")
        metrics = {
            "CVRMSE": cvrmse(y_selected, prediction[selected]),
            "NMBE": nmbe(y_selected, prediction[selected]),
        }
        logger.info(f"In-sample metrics: {metrics}")

    else:
        X_act = pd.concat([features, activity], axis=1)
        model = UsagePredictor(
            lags=lags,
            cat_features=cat_features,
            validation_size=validation_size,
            skip_calendar=True,
        )
        model = model.fit(X_act, labels)
        logger.info(f"Model scores: {model.scores_}")

        prediction = model.predict(X_act).to_frame("consumption")

        metrics = {
            "CVRMSE": cvrmse(labels, prediction),
            "NMBE": nmbe(labels, prediction),
        }
        logger.info(f"In-sample metrics: {metrics}")

    return model, prediction, metrics


def predict_autoenc(
    features: pd.DataFrame,
    activity: pd.DataFrame,
    model: UsagePredictor,
    assume_hurdle: bool,
):
    """Predict energy consumption given input data.

    Args:
        features (pd.DataFrame): The feature data.
        activity (pd.DataFrame): The estimated activity levels.
        model (UsagePredictor): The fitted model.
        assume_hurdle (bool): If True, activity estimation will assume that energy
            consumption has been generated by a hurdle model.

    Returns:
        pandas.DataFrame: The prediction.
    """

    if assume_hurdle:
        prediction = model.predict(features).to_frame("consumption")

    else:
        X_act = pd.concat([features, activity], axis=1)
        prediction = model.predict(X_act).to_frame("consumption")

    return prediction
