"""Forecasting error metrics for model evaluation."""

from typing import Union, Optional

import numpy as np
import pandas as pd


def calculate_mase(
    actual: Union[np.ndarray, pd.Series],
    predicted: Union[np.ndarray, pd.Series],
    training_series: Union[np.ndarray, pd.Series],
    period: int = 1,
) -> float:
    """
    Calculate Mean Absolute Scaled Error (MASE).

    MASE compares the model's error to the error of a naive seasonal forecast
    (predicting the value from one period ago). MASE < 1 indicates the model
    is better than the naive forecast.

    Parameters
    ----------
    actual : np.ndarray or pd.Series
        Actual values.
    predicted : np.ndarray or pd.Series
        Predicted values.
    training_series : np.ndarray or pd.Series
        Training data used to fit the model (for calculating naive MAE).
    period : int, default 1
        Seasonal period for naive forecast. Use 1 for random walk (non-seasonal),
        7 for weekly seasonality on daily data, etc.

    Returns
    -------
    mase : float
        Mean Absolute Scaled Error. Values < 1 indicate the model beats
        the naive forecast.

    Notes
    -----
    MASE = MAE_model / MAE_naive

    where MAE_naive is calculated on the training data as:
    MAE_naive = (1/(n-m)) * sum(|Y_t - Y_{t-m}|) for t = m+1 to n

    References
    ----------
    Hyndman, R.J., & Koehler, A.B. (2006). Another look at measures of
    forecast accuracy. International Journal of Forecasting, 22(4), 679-688.

    Examples
    --------
    >>> actual = np.array([100, 110, 105, 115, 108])
    >>> predicted = np.array([102, 108, 107, 112, 110])
    >>> training = np.array([80, 85, 90, 95, 100, 105, 110])
    >>> mase = calculate_mase(actual, predicted, training, period=1)
    """
    actual = np.asarray(actual)
    predicted = np.asarray(predicted)
    training_series = np.asarray(training_series)

    if len(actual) != len(predicted):
        raise ValueError("actual and predicted must have the same length")

    if len(training_series) <= period:
        raise ValueError(f"training_series must have more than {period} observations")

    # Calculate MAE of the model
    mae_model = np.nanmean(np.abs(actual - predicted))

    # Calculate MAE of naive seasonal forecast on training data
    # Naive forecast: Y_hat_t = Y_{t-period}
    naive_errors = np.abs(
        training_series[period:] - training_series[:-period]
    )
    mae_naive = np.nanmean(naive_errors)

    if mae_naive == 0:
        # If naive MAE is 0, the series is constant
        if mae_model == 0:
            return 0.0
        else:
            return np.inf

    return mae_model / mae_naive


def calculate_smape(
    actual: Union[np.ndarray, pd.Series],
    predicted: Union[np.ndarray, pd.Series],
) -> float:
    """
    Calculate Symmetric Mean Absolute Percentage Error (sMAPE).

    sMAPE handles zeros better than MAPE and is bounded between 0% and 200%.

    Parameters
    ----------
    actual : np.ndarray or pd.Series
        Actual values.
    predicted : np.ndarray or pd.Series
        Predicted values.

    Returns
    -------
    smape : float
        Symmetric MAPE as a percentage (0-200).

    Notes
    -----
    sMAPE = (100/n) * sum(|Y_t - Y_hat_t| / ((|Y_t| + |Y_hat_t|) / 2))

    Examples
    --------
    >>> actual = np.array([100, 110, 105, 115, 108])
    >>> predicted = np.array([102, 108, 107, 112, 110])
    >>> smape = calculate_smape(actual, predicted)
    """
    actual = np.asarray(actual, dtype=float)
    predicted = np.asarray(predicted, dtype=float)

    if len(actual) != len(predicted):
        raise ValueError("actual and predicted must have the same length")

    numerator = np.abs(actual - predicted)
    denominator = (np.abs(actual) + np.abs(predicted)) / 2

    # Handle division by zero (both actual and predicted are 0)
    mask = denominator > 0
    if not mask.any():
        return 0.0

    smape = 100 * np.nanmean(numerator[mask] / denominator[mask])
    return smape


def calculate_mape(
    actual: Union[np.ndarray, pd.Series],
    predicted: Union[np.ndarray, pd.Series],
) -> float:
    """
    Calculate Mean Absolute Percentage Error (MAPE).

    Note: MAPE is undefined when actual values are zero. Use sMAPE or MASE
    for series that may contain zeros.

    Parameters
    ----------
    actual : np.ndarray or pd.Series
        Actual values (must be non-zero for valid MAPE).
    predicted : np.ndarray or pd.Series
        Predicted values.

    Returns
    -------
    mape : float
        Mean Absolute Percentage Error as a percentage.

    Raises
    ------
    ValueError
        If actual contains zeros.

    Notes
    -----
    MAPE = (100/n) * sum(|Y_t - Y_hat_t| / |Y_t|)

    Examples
    --------
    >>> actual = np.array([100, 110, 105, 115, 108])
    >>> predicted = np.array([102, 108, 107, 112, 110])
    >>> mape = calculate_mape(actual, predicted)
    """
    actual = np.asarray(actual, dtype=float)
    predicted = np.asarray(predicted, dtype=float)

    if len(actual) != len(predicted):
        raise ValueError("actual and predicted must have the same length")

    # Check for zeros
    zero_mask = actual == 0
    if zero_mask.any():
        n_zeros = zero_mask.sum()
        # Exclude zeros from calculation but warn
        actual = actual[~zero_mask]
        predicted = predicted[~zero_mask]
        if len(actual) == 0:
            return np.inf

    mape = 100 * np.nanmean(np.abs(actual - predicted) / np.abs(actual))
    return mape


def calculate_rmse(
    actual: Union[np.ndarray, pd.Series],
    predicted: Union[np.ndarray, pd.Series],
) -> float:
    """
    Calculate Root Mean Square Error (RMSE).

    RMSE is in the same units as the data and penalizes large errors
    more heavily than MAE.

    Parameters
    ----------
    actual : np.ndarray or pd.Series
        Actual values.
    predicted : np.ndarray or pd.Series
        Predicted values.

    Returns
    -------
    rmse : float
        Root Mean Square Error.

    Notes
    -----
    RMSE = sqrt((1/n) * sum((Y_t - Y_hat_t)^2))

    Examples
    --------
    >>> actual = np.array([100, 110, 105, 115, 108])
    >>> predicted = np.array([102, 108, 107, 112, 110])
    >>> rmse = calculate_rmse(actual, predicted)
    """
    actual = np.asarray(actual, dtype=float)
    predicted = np.asarray(predicted, dtype=float)

    if len(actual) != len(predicted):
        raise ValueError("actual and predicted must have the same length")

    mse = np.nanmean((actual - predicted) ** 2)
    return np.sqrt(mse)


def calculate_mae(
    actual: Union[np.ndarray, pd.Series],
    predicted: Union[np.ndarray, pd.Series],
) -> float:
    """
    Calculate Mean Absolute Error (MAE).

    MAE is in the same units as the data and is robust to outliers
    compared to RMSE.

    Parameters
    ----------
    actual : np.ndarray or pd.Series
        Actual values.
    predicted : np.ndarray or pd.Series
        Predicted values.

    Returns
    -------
    mae : float
        Mean Absolute Error.

    Notes
    -----
    MAE = (1/n) * sum(|Y_t - Y_hat_t|)

    Examples
    --------
    >>> actual = np.array([100, 110, 105, 115, 108])
    >>> predicted = np.array([102, 108, 107, 112, 110])
    >>> mae = calculate_mae(actual, predicted)
    """
    actual = np.asarray(actual, dtype=float)
    predicted = np.asarray(predicted, dtype=float)

    if len(actual) != len(predicted):
        raise ValueError("actual and predicted must have the same length")

    return np.nanmean(np.abs(actual - predicted))


def calculate_wmape(
    actual: Union[np.ndarray, pd.Series],
    predicted: Union[np.ndarray, pd.Series],
) -> float:
    """
    Calculate Weighted Mean Absolute Percentage Error (WMAPE).

    WMAPE weights errors by the actual values, making it more suitable
    for aggregated forecasts where larger values are more important.

    Parameters
    ----------
    actual : np.ndarray or pd.Series
        Actual values.
    predicted : np.ndarray or pd.Series
        Predicted values.

    Returns
    -------
    wmape : float
        Weighted MAPE as a percentage.

    Notes
    -----
    WMAPE = 100 * sum(|Y_t - Y_hat_t|) / sum(|Y_t|)

    Examples
    --------
    >>> actual = np.array([100, 110, 105, 115, 108])
    >>> predicted = np.array([102, 108, 107, 112, 110])
    >>> wmape = calculate_wmape(actual, predicted)
    """
    actual = np.asarray(actual, dtype=float)
    predicted = np.asarray(predicted, dtype=float)

    if len(actual) != len(predicted):
        raise ValueError("actual and predicted must have the same length")

    sum_actual = np.nansum(np.abs(actual))
    if sum_actual == 0:
        return np.inf

    sum_errors = np.nansum(np.abs(actual - predicted))
    return 100 * sum_errors / sum_actual


def calculate_bias(
    actual: Union[np.ndarray, pd.Series],
    predicted: Union[np.ndarray, pd.Series],
) -> float:
    """
    Calculate forecast bias (mean error).

    Positive bias indicates over-forecasting, negative indicates under-forecasting.

    Parameters
    ----------
    actual : np.ndarray or pd.Series
        Actual values.
    predicted : np.ndarray or pd.Series
        Predicted values.

    Returns
    -------
    bias : float
        Mean error (predicted - actual).

    Examples
    --------
    >>> actual = np.array([100, 110, 105, 115, 108])
    >>> predicted = np.array([102, 108, 107, 112, 110])
    >>> bias = calculate_bias(actual, predicted)
    """
    actual = np.asarray(actual, dtype=float)
    predicted = np.asarray(predicted, dtype=float)

    if len(actual) != len(predicted):
        raise ValueError("actual and predicted must have the same length")

    return np.nanmean(predicted - actual)


def calculate_all_metrics(
    actual: Union[np.ndarray, pd.Series],
    predicted: Union[np.ndarray, pd.Series],
    training_series: Optional[Union[np.ndarray, pd.Series]] = None,
    period: int = 1,
) -> dict:
    """
    Calculate all available forecast metrics.

    Parameters
    ----------
    actual : np.ndarray or pd.Series
        Actual values.
    predicted : np.ndarray or pd.Series
        Predicted values.
    training_series : np.ndarray or pd.Series, optional
        Training data for MASE calculation.
    period : int, default 1
        Seasonal period for MASE.

    Returns
    -------
    metrics : dict
        Dictionary with all calculated metrics.

    Examples
    --------
    >>> actual = np.array([100, 110, 105, 115, 108])
    >>> predicted = np.array([102, 108, 107, 112, 110])
    >>> training = np.array([80, 85, 90, 95, 100, 105, 110])
    >>> metrics = calculate_all_metrics(actual, predicted, training, period=1)
    >>> print(metrics['smape'], metrics['mase'])
    """
    metrics = {
        "mae": calculate_mae(actual, predicted),
        "rmse": calculate_rmse(actual, predicted),
        "smape": calculate_smape(actual, predicted),
        "wmape": calculate_wmape(actual, predicted),
        "bias": calculate_bias(actual, predicted),
    }

    # MAPE may fail if actual contains zeros
    try:
        metrics["mape"] = calculate_mape(actual, predicted)
    except ValueError:
        metrics["mape"] = np.nan

    # MASE requires training series
    if training_series is not None:
        metrics["mase"] = calculate_mase(
            actual, predicted, training_series, period=period
        )
    else:
        metrics["mase"] = np.nan

    return metrics


def coverage_probability(
    actual: Union[np.ndarray, pd.Series],
    lower: Union[np.ndarray, pd.Series],
    upper: Union[np.ndarray, pd.Series],
) -> float:
    """
    Calculate coverage probability of prediction intervals.

    Parameters
    ----------
    actual : np.ndarray or pd.Series
        Actual values.
    lower : np.ndarray or pd.Series
        Lower bounds of prediction intervals.
    upper : np.ndarray or pd.Series
        Upper bounds of prediction intervals.

    Returns
    -------
    coverage : float
        Proportion of actual values within the intervals (0-1).

    Examples
    --------
    >>> actual = np.array([100, 110, 105, 115, 108])
    >>> lower = np.array([95, 100, 98, 105, 100])
    >>> upper = np.array([115, 120, 115, 125, 120])
    >>> cov = coverage_probability(actual, lower, upper)
    """
    actual = np.asarray(actual)
    lower = np.asarray(lower)
    upper = np.asarray(upper)

    within = (actual >= lower) & (actual <= upper)
    return np.nanmean(within)


def interval_width(
    lower: Union[np.ndarray, pd.Series],
    upper: Union[np.ndarray, pd.Series],
) -> float:
    """
    Calculate average width of prediction intervals.

    Parameters
    ----------
    lower : np.ndarray or pd.Series
        Lower bounds of prediction intervals.
    upper : np.ndarray or pd.Series
        Upper bounds of prediction intervals.

    Returns
    -------
    width : float
        Average interval width.
    """
    lower = np.asarray(lower)
    upper = np.asarray(upper)

    return np.nanmean(upper - lower)
