import argparse
from datetime import datetime
from pathlib import Path

import joblib
import numpy as np
import pandas as pd
from pydantic import validate_call

from energy_consumption_forecasting.exceptions import log_exception
from energy_consumption_forecasting.inference_pipeline.batch_data import (
    get_batch_data_from_hopsworks,
)
from energy_consumption_forecasting.inference_pipeline.utils import (
    get_gcs_bucket,
    read_blob_from_bucket,
    write_blob_to_bucket,
)
from energy_consumption_forecasting.logger import get_logger
from energy_consumption_forecasting.utils import get_env_var

logger = get_logger(name=Path(__file__).name)

import hopsworks


def load_model_from_hopsworks(
    project,
    model_version: int = 1,
    model_name: str = "forecast_model",
):
    """
    This function loads the model from hopsworks model registry,
    model is downloaded as pickle file and transformed into a
    sktime object that can be used for forecasting.

    Parameters
    ----------
    project
        A project object generated by hopsworks when login into a specific project.

    model_version: int, default=1
        Model version used when saving the model in hopsworks.

    model_name: str, default="forecast_model"
        Model name used when saving the model in hopsworks.

    Returns
    -------
    model
        A model is returned after loading the pickle file, the model type can be a
        sktime model or sktime forecasting pipeline which can be used for forecasting.
    """

    model_registry = project.get_model_registry()
    model_instance = model_registry.get_model(name=model_name, version=model_version)
    model_path = model_instance.download()
    model_path = Path(model_path) / f"{model_name}_{model_version}.pkl"

    # Loading the model pickle file object
    model = joblib.load(model_path)

    return model


def get_batch_forecast(model, X: pd.DataFrame, fh: int = 24) -> pd.DataFrame:
    """
    This functions takes the X dataframe and generates a forecast dataframe similar to
    X with the help of forecast horizon input.
    Once the data is generated it is used to make prediction from the model.

    Parameters
    ----------
    model:
        A sktime model or sktime forecasting pipeline that has been trained using the
        fit function.

    X: pd.DataFrame
        A batch dataframe with input features that are provided by the feature store.

    fh: int, default=24
        A period that indicates the forecast horizon while making prediction in Hours.

    Returns
    -------
    pd.Dataframe
        Model generates the forecast prediction in pandas dataframe format.
        The range of the forecast date and time depends on the fh argument.
    """

    # Getting the index exogenous features and creating a forecast date range
    municipality_num = X.index.get_level_values(level=0).unique()
    branch = X.index.get_level_values(level=1).unique()
    last_datetime = X.index.get_level_values(level=2).max()

    start_datetime = last_datetime + 1
    end_datetime = last_datetime + fh
    fh_range = pd.date_range(
        start=start_datetime.to_timestamp(),
        end=end_datetime.to_timestamp(),
        freq="H",
    )
    fh_range = pd.PeriodIndex(data=fh_range, freq="H")

    # Creating a forecast X data using the above generated data
    multi_index = pd.MultiIndex.from_product(
        iterables=[municipality_num, branch, fh_range],
        names=["municipality_num", "branch", "datetime_dk"],
    )
    X_forecast = pd.DataFrame(index=multi_index)

    # Performing prediction using the newly generated forecast data and forecast horizon
    prediction = model.predict(X=X_forecast, fh=np.arange(fh) + 1)

    return prediction


def save_data_to_gcs_bucket(X: pd.DataFrame, y: pd.DataFrame, prediction: pd.DataFrame):
    """
    This function saves the data into the Google cloud storage bucket by transforming
    the provided dataframe into parquet data and uploading it as blob object in
    the bucket.

    Parameters
    ----------
    X: pd.DataFrame
        A dataframe that is used as input data for generating prediction from model.
        This data is saved in the bucket with the blob name as "input.parquet".

    y: pd.DataFrame
        A dataframe that is used as target data for generating prediction from model.
        This data is saved in the bucket with the blob name as "target.parquet".

    prediction: pd.DataFrame
        A dataframe that is generated from the model as a prediction or forecast.
        This data is saved in the bucket with the blob name as "prediction.parquet".
    """

    # Connecting to Google cloud storage and getting the bucket object
    bucket = get_gcs_bucket()

    # Transforming and uploading the input, target and prediction dataframe
    # into parquet data in the GCS bucket
    for data, blob_name in zip(
        [X, y, prediction], ["input.parquet", "target.parquet", "prediction.parquet"]
    ):
        logger.info(f'Uploading blob data "{blob_name}" in the GCS bucket.')
        write_blob_to_bucket(
            bucket=bucket,
            blob_name=blob_name,
            data=data,
        )
        logger.info(f'Blob Data "{blob_name}" is successfully uploaded to GCS bucket.')


def save_prediction_data_for_caching(
    prediction: pd.DataFrame,
    start_datetime: datetime,
):
    """
    This function creates or updates the cached prediction dataframe with the newly
    generated prediction and then saves the dataframe in the Google cloud storage(GCS)
    as a parquet file called "cached_prediction.parquet".

    Parameters
    ----------
    prediction: pd.DataFrame
        Newly generated forecast prediction by the model.

    start_datetime: datetime.datetime
        A starting date and time for filtering the cached prediction before uploading
        the data in the GCS bucket.
    """

    # Getting the GCS bucket object to download the cached prediction data
    bucket = get_gcs_bucket()
    cached_prediction = read_blob_from_bucket(
        bucket=bucket, blob_name="cached_prediction.parquet"
    )

    # Checking if cached prediction data exist
    cached_prediction_exist = cached_prediction is not None
    if cached_prediction_exist:
        # Merging current prediction and cached prediction
        merged_prediction = prediction.merge(
            right=cached_prediction,
            how="outer",
            left_index=True,
            right_index=True,
            suffixes=("_new", "_cached"),
        )

        # Filtering the new and cached prediction into different dataframe
        new_prediction = merged_prediction.filter(regex="_new")
        new_prediction.columns = new_prediction.columns.str.replace(pat="_new", repl="")
        cached_pred = merged_prediction.filter(regex="_cached")
        cached_pred.columns = cached_pred.columns.str.replace(pat="_cached", repl="")

        # Updating the NaN values in the new prediction with cached prediction data
        new_prediction.update(cached_pred)
        prediction = new_prediction

    # Checking the prediction dataframe and uploading it to the GCS bucket
    prediction = prediction.loc[
        prediction.index.get_level_values("datetime_dk")
        >= pd.Period(start_datetime, freq="H")
    ]
    prediction = prediction.dropna(subset=["consumption_kwh"])

    write_blob_to_bucket(
        bucket=bucket,
        blob_name="cached_prediction.parquet",
        data=prediction,
    )
    logger.info(
        'Cached prediction "cached_prediction.parquet" has been successfully '
        "updated with the new forecast predictions"
    )


@log_exception(logger=logger)
@validate_call
def run_inference_pipeline(
    start_datetime: datetime,
    end_datetime: datetime,
    fh: int = 24,
    target_feature: str = "consumption_kwh",
    feature_view_ver: int = 1,
    feature_view_name: str = "denmark_energy_consumption_view",
    model_version: int = 1,
    model_name: str = "forecast_model",
):
    """
    This function will perform the batch inference pipeline from getting the data
    and model from hopsworks and using the model to get the forecast prediction and
    finally storing the data in the Google cloud storage(GCS) bucket and also
    caching the prediction for future use.

    Parameters
    ----------

    start_datetime: datetime.datetime
        A starting date and time for extracting the batch data
        from hopsworks feature view.

    end_datetime: datetime.datetime
        A ending date and time for extracting the batch data
        from hopsworks feature view.

    fh: int, default=24
        A period that indicates the forecast horizon while making prediction in Hours.

    target_feature: str, default="consumption_kwh"
        The name of the target feature in the dataset.

    feature_view_ver: int, default=1
        The feature view version that needs to be loaded.

    feature_view_name: str, default="denmark_energy_consumption_view"
        The name of the feature view in the hopsworks feature store.

    model_version: int, default=1
        Model version used when saving the model in hopsworks.

    model_name: str, default="forecast_model"
        Model name used when saving the model in hopsworks.
    """

    # Connecting to the hopsworks feature store using the project API
    project = hopsworks.login(
        project=get_env_var(key="FEATURE_STORE_PROJECT_NAME"),
        api_key_value=get_env_var(key="FEATURE_STORE_API_KEY"),
    )
    feature_store = project.get_feature_store()
    logger.info(
        f'Connected to Hopsworks: Project Name "{project.name}" and '
        f'Project URL: "{project.get_url()}"'
    )

    # Getting the dataframe from the hopsworks feature store and transforming for inference
    logger.info("Loading data from hopsworks feature store.")
    logger.info(
        "Loading data between the following date-time range, "
        f"start: {start_datetime} and end: {end_datetime}"
    )
    X, y = get_batch_data_from_hopsworks(
        feature_store=feature_store,
        start_datetime=start_datetime,
        end_datetime=end_datetime,
        feature_view_ver=feature_view_ver,
        feature_view_name=feature_view_name,
        target_feature=target_feature,
    )
    logger.info("Data is successfully loaded from the hopsworks feature store.")

    # Getting the model from the hopsworks model registry
    logger.info("Loading model from the hopsworks model registry.")
    model = load_model_from_hopsworks(
        project=project,
        model_version=model_version,
        model_name=model_name,
    )
    logger.info("Model is successfully loaded from the model registry.")

    # Performing forecast prediction using the data and model
    logger.info("Performing forecast prediction using the data and model.")
    prediction = get_batch_forecast(model=model, X=X, fh=fh)
    pred_start_datetime = prediction.index.get_level_values(level="datetime_dk").min()
    pred_end_datetime = prediction.index.get_level_values(level="datetime_dk").max()
    logger.info(
        f"Forecast prediction are from {pred_start_datetime} to {pred_end_datetime}"
    )
    logger.info("Prediction were successfully forecasted.")

    # Saving the dataframe in the GCS bucket
    logger.info("Saving the dataframe in the GCS bucket.")
    save_data_to_gcs_bucket(
        X=X,
        y=y,
        prediction=prediction,
    )
    logger.info("All the dataframe were successfully uploaded in the GCS bucket.")

    # Saving the newly generated prediction by updating the cached_prediction dataframe
    logger.info("Merging new prediction with cached prediction")
    save_prediction_data_for_caching(
        prediction=prediction,
        start_datetime=start_datetime,
    )
    logger.info("Cached prediction was successfully updated with the new prediction")


if __name__ == "__main__":

    parser = argparse.ArgumentParser()

    parser.add_argument(
        "-s",
        "--start_datetime",
        type=datetime.fromisoformat,
        required=True,
        help="Starting date for extraction in format: YYYY-MM-DDTHH:MM:SS",
    )

    parser.add_argument(
        "-e",
        "--end_datetime",
        type=datetime.fromisoformat,
        required=True,
        help="Ending date for extraction in format: YYYY-MM-DDTHH:MM:SS",
    )

    parser.add_argument(
        "--fh",
        type=int,
        default=24,
        help="Forecasting horizon period, needs to be in integer format.",
    )

    parser.add_argument(
        "--views_name",
        type=str,
        default="denmark_energy_consumption_view",
        help="Name of feature view within the feature group, needs to be in string format.",
    )

    parser.add_argument(
        "--views_ver",
        type=int,
        default=1,
        help="Version of feature view within the feature group, "
        "needs to be in integer format.",
    )

    parser.add_argument(
        "--target_feature",
        type=str,
        default="consumption_kwh",
        help="Name of target feature, needs to be in string format.",
    )

    parser.add_argument(
        "--model_name",
        type=str,
        default="forecast_model",
        help="Name of the model for this training experiment, needs to be in string format.",
    )

    parser.add_argument(
        "--model_ver",
        type=int,
        default=1,
        help="Model version or id for this training experiment, "
        "needs to be in integer format.",
    )

    args = parser.parse_args()

    run_inference_pipeline(
        start_datetime=args.start_datetime,
        end_datetime=args.end_datetime,
        fh=args.fh,
        target_feature=args.target_feature,
        feature_view_ver=args.views_ver,
        feature_view_name=args.views_name,
        model_version=args.model_ver,
        model_name=args.model_name,
    )
