from dataclasses import dataclass
from typing import Optional
import pandas as pd
from statistics import mean, stdev, median

import properscoring as ps
from sklearn.metrics import mean_squared_error, mean_absolute_error, mean_squared_log_error, brier_score_loss, average_precision_score, roc_auc_score
# WARNING: mean_squared_error and mean_squared_log_error are deprecated.
# from views_forecasts.extensions import *


@dataclass
class EvaluationMetrics:
    """
    A data class for storing and managing evaluation metrics for time series forecasting models.
    
    Attributes:
        MSE (Optional[float]): Mean Squared Error.
        MAE (Optional[float]): Mean Absolute Error.
        MSLE (Optional[float]): Mean Squared Logarithmic Error.
        KLD (Optional[float]): Kullback-Leibler Divergence.
        Jeffreys (Optional[float]): Jeffreys Divergence.
        CRPS (Optional[float]): Continuous Ranked Probability Score.
        Brier (Optional[float]): Brier Score.
        AP (Optional[float]): Average Precision.
        AUC (Optional[float]): Area Under the ROC Curve.
        ensemble_weight_reg (Optional[float]): Weight for regression ensemble models.
        ensemble_weight_class (Optional[float]): Weight for classification ensemble models.
    """

    MSE: Optional[float] = None
    MAE: Optional[float] = None
    MSLE: Optional[float] = None
    KLD: Optional[float] = None
    Jeffreys: Optional[float] = None
    CRPS: Optional[float] = None
    Brier: Optional[float] = None
    AP: Optional[float] = None
    AUC: Optional[float] = None
    ensemble_weight_reg: Optional[float] = None
    ensemble_weight_class: Optional[float] = None

    @classmethod
    def make_evaluation_dict(cls, steps=36) -> dict:
        """
        Generates a dictionary of EvaluationMetrics instances for a specified number of forecasting steps.

        This method facilitates the batch creation of metric containers for multiple forecasting steps, initializing them with None.

        Args:
            steps (int): The number of forecasting steps for which to generate evaluation metrics. Defaults to 36.

        Returns:
            dict: A dictionary where each key is a step label (e.g., 'step01', 'step02', ...) and each value is an instance of EvaluationMetrics.

        Example:
            >>> from utils_evaluation_metrics import EvaluationMetrics
            >>> evaluation_dict = EvaluationMetrics.make_evaluation_dict(steps=36)
            >>> evaluation_dict['step01'].MSE = sklearn.metrics.mean_squared_error(step01_y_true, step01_y_pred)
            >>> evaluation_dict['step02'].MSE = sklearn.metrics.mean_squared_error(step02_y_true, step02_y_pred)
            >>> ...
            
        """
        return {f"step{str(i).zfill(2)}": cls() for i in range(1, steps + 1)}

    @staticmethod
    def evaluation_dict_to_dataframe(evaluation_dict: dict) -> pd.DataFrame:
        """
        Converts a dictionary of EvaluationMetrics instances into a pandas DataFrame for easy analysis.

        This static method transforms a structured dictionary of evaluation metrics into a DataFrame, where each row corresponds to a forecasting step and columns represent different metrics.

        Args:
            evaluation_dict (dict): A dictionary of EvaluationMetrics instances, typically generated by the make_evaluation_dict class method.

        Returns:
            pd.DataFrame: A pandas DataFrame where each row indexes a forecasting step and columns correspond to the various metrics stored in EvaluationMetrics.

        Example:
            >>> evaluation_df = EvaluationMetrics.evaluation_dict_to_dataframe(evaluation_dict)

        """
        return pd.DataFrame.from_dict(evaluation_dict, orient='index')

# TBD: Align with metrics discussed in workshop

    @staticmethod
    def calculate_aggregate_metrics(evaluation_dict: dict) -> dict:
        metrics_aggregate = {
            'mean': {},
            'std': {},
            'median': {}
        }

        for metric in EvaluationMetrics.__annotations__.keys():
            metric_values = [getattr(evaluation, metric) for evaluation in evaluation_dict.values() if getattr(evaluation, metric) is not None]
            if metric_values: 
                metrics_aggregate['mean'][metric] = mean(metric_values)
                metrics_aggregate['std'][metric] = stdev(metric_values)
                metrics_aggregate['median'][metric] = median(metric_values)
            else:
                metrics_aggregate['mean'][metric] = None
                metrics_aggregate['std'][metric] = None
                metrics_aggregate['median'][metric] = None

        return metrics_aggregate

    @staticmethod
    def output_metrics(evaluation_dict):
        aggregate = EvaluationMetrics.calculate_aggregate_metrics(evaluation_dict)
        step_metrics_dict = {step: vars(metrics) for step, metrics in evaluation_dict.items()}
        step_metrics_dict['mean'] = aggregate['mean']
        step_metrics_dict['std'] = aggregate['std']
        step_metrics_dict['median'] = aggregate['median']
        return step_metrics_dict


def generate_metric_dict(df, config):
    """
    Generates a dictionary of evaluation metrics for a given forecasting configuration and dataset.

    Args:
        df (pd.DataFrame): A pandas DataFrame containing the forecasted values and ground truth.
        config (dict): A dictionary containing the forecasting configuration parameters.

    Returns:
        evaluation_dict (dict): A dictionary of EvaluationMetrics instances for each forecasting step.
        df_evaluation_dict (pd.DataFrame): A pandas DataFrame containing the evaluation metrics for each forecasting step.

    Note:
        ! This function is temporary for the stepshifter model.
        ! Change the metrics to those discussed previously.
        For logged targets, calculating MSE is actually MSLE.
        KLD and Jeffreys divergence are measures used to quantify the difference between two probability distributions. Why do we calculate these metrics in the context of forecasting?
        Brier score is used for binary and categorical outcomes that can be structured as true or false
        There are no classes in data, so we cannot calculate roc_auc_score, ap_score
    """

    evaluation_dict = EvaluationMetrics.make_evaluation_dict(steps=config["steps"][-1])
    for step in config["steps"]:
        evaluation_dict[f"step{str(step).zfill(2)}"].MSE = mean_squared_error(df[config["depvar"]], df[f"step_pred_{step}"])
        evaluation_dict[f"step{str(step).zfill(2)}"].MAE = mean_absolute_error(df[config["depvar"]], df[f"step_pred_{step}"])
        # evaluation_dict[f"step{str(step).zfill(2)}"].MSLE = mean_squared_log_error(df[config["depvar"]], df[f"step_pred_{step}"])
        evaluation_dict[f"step{str(step).zfill(2)}"].CRPS = ps.crps_ensemble(df[config["depvar"]], df[f"step_pred_{step}"]).mean()
        # evaluation_dict[f"step{str(step).zfill(2)}"].Brier = brier_score_loss(df[config["depvar"]], df[f"step_pred_{step}"])
        # evaluation_dict[f"step{str(step).zfill(2)}"].AUC = roc_auc_score(df[config["depvar"]], df[f"step_pred_{step}"])
        # evaluation_dict[f"step{str(step).zfill(2)}"].AP = average_precision_score(df[config["depvar"]], df[f"step_pred_{step}"])
    evaluation_dict = EvaluationMetrics.output_metrics(evaluation_dict)
    df_evaluation_dict = EvaluationMetrics.evaluation_dict_to_dataframe(evaluation_dict)
    return evaluation_dict, df_evaluation_dict