"""Evaluation functions for out-of-sample return predictions.

Two evaluation dimensions:

Prediction accuracy
    prediction_eval() -- MSE, MAE, MAPE, OOS R² (Campbell & Thompson 2008)
    eval_oos_result() -- convenience wrapper over all pred_* columns

Investment performance
    portfolio_eval() -- Sharpe, vol, mean, information ratio, hit ratio
"""

from __future__ import annotations

import numpy as np
import pandas as pd

from eapctf.predict.oos import OOSResult

# ---------------------------------------------------------------------------
# Prediction accuracy
# ---------------------------------------------------------------------------


def prediction_eval(
    predictions: pd.DataFrame,
    pred_col: str,
    ret_col: str = "ret",
    benchmark_col: str | None = None,
) -> dict[str, float]:
    """Compute prediction accuracy metrics for one model column.

    Parameters
    ----------
    predictions : pd.DataFrame
        DataFrame from OOSResult.predictions containing ret_col and pred_col.
    pred_col : str
        Name of the prediction column (e.g. ``"pred_LassoPredictor"``).
    ret_col : str
        Name of the realized return column.
    benchmark_col : str | None
        Column to use as the Campbell-Thompson benchmark for OOS R².
        If None and ``"pred_HistMean"`` exists, that is used automatically
        (expanding training-set mean, the standard C&T benchmark).
        If neither is available, falls back to the in-sample mean of ret_col
        (less accurate; prefer running expanding_window_oos which stores
        pred_HistMean by default).
        Pass any other pred_* column name to benchmark against another model.

    Returns
    -------
    dict with keys:
        mse, mae, mape, oos_r2, hit_ratio
    """
    needed = [ret_col, pred_col]
    df = predictions[needed].dropna()
    if len(df) == 0:
        return {k: float("nan") for k in ("mse", "mae", "mape", "oos_r2", "hit_ratio")}

    y: pd.Series = df[ret_col]
    y_hat: pd.Series = df[pred_col]
    residuals = y - y_hat

    mse = float((residuals ** 2).mean())
    mae = float(residuals.abs().mean())

    # MAPE: skip observations where actual return is zero to avoid division by zero.
    nonzero = y != 0
    if nonzero.any():
        mape = float((residuals[nonzero].abs() / y[nonzero].abs()).mean())
    else:
        mape = float("nan")

    # OOS R² (Campbell & Thompson 2008):
    #   R²_OS = 1 - MSPE_model / MSPE_benchmark
    # Resolve benchmark series.
    if benchmark_col is not None:
        if benchmark_col not in predictions.columns:
            raise ValueError(
                f"benchmark_col '{benchmark_col}' not found in predictions columns: "
                f"{list(predictions.columns)}"
            )
        bench: pd.Series = predictions[benchmark_col].reindex(df.index)
    elif "pred_HistMean" in predictions.columns:
        bench = predictions["pred_HistMean"].reindex(df.index)
    else:
        # Fallback: unconditional mean of realized returns in the prediction sample.
        # Note this is an in-sample mean and slightly upward-biases R²_OS.
        bench = pd.Series(float(y.mean()), index=df.index)

    bench_resid = y - bench.reindex(df.index)
    mspe_model = float((residuals ** 2).mean())
    mspe_bench = float((bench_resid ** 2).mean())
    oos_r2 = float(1.0 - mspe_model / mspe_bench) if mspe_bench != 0.0 else float("nan")

    # Hit ratio: fraction of obs where predicted sign matches realized sign.
    # Ties (zero prediction or zero actual) are counted as misses.
    hit_ratio = float((np.sign(y_hat) == np.sign(y)).mean())

    return {
        "mse": mse,
        "mae": mae,
        "mape": mape,
        "oos_r2": oos_r2,
        "hit_ratio": hit_ratio,
    }


def eval_oos_result(
    result: OOSResult,
    ret_col: str = "ret",
    benchmark_col: str | None = None,
) -> pd.DataFrame:
    """Compute prediction accuracy metrics for every model in an OOSResult.

    Parameters
    ----------
    result : OOSResult
        Output of expanding_window_oos().
    ret_col : str
        Name of the realized return column in result.predictions.
    benchmark_col : str | None
        Passed directly to prediction_eval(). Defaults to pred_HistMean
        (Campbell & Thompson 2008 historical-mean benchmark).

    Returns
    -------
    pd.DataFrame
        Index = model name (without ``pred_`` prefix).
        Columns = mse, mae, mape, oos_r2, hit_ratio.
    """
    df = result.predictions
    if len(df) == 0:
        return pd.DataFrame(
            columns=["mse", "mae", "mape", "oos_r2", "hit_ratio"],
            dtype=float,
        )

    # All pred_* columns except the HistMean benchmark itself.
    pred_cols = [
        c for c in df.columns
        if c.startswith("pred_") and c != "pred_HistMean"
    ]

    rows = []
    for col in pred_cols:
        metrics = prediction_eval(
            df,
            pred_col=col,
            ret_col=ret_col,
            benchmark_col=benchmark_col,
        )
        rows.append({"model": col[len("pred_"):], **metrics})

    return pd.DataFrame(rows).set_index("model")


# ---------------------------------------------------------------------------
# Investment performance
# ---------------------------------------------------------------------------


def portfolio_eval(
    returns: pd.Series,
    benchmark: pd.Series | None = None,
    freq: int = 12,
) -> dict[str, float]:
    """Compute investment performance metrics for a return series.

    Parameters
    ----------
    returns : pd.Series
        Portfolio return series (one observation per period).
    benchmark : pd.Series | None
        Benchmark return series used to compute the information ratio.
        Aligned to returns by index. If None, information_ratio is NaN.
    freq : int
        Number of periods per year for annualization (12 for monthly data,
        252 for daily data).

    Returns
    -------
    dict with keys:
        mean_ret        -- annualized mean return
        vol             -- annualized volatility (std dev)
        sharpe          -- annualized Sharpe ratio (mean / vol)
        information_ratio -- annualized IR vs benchmark (NaN if no benchmark)
        hit_ratio       -- fraction of periods with positive return
        n_periods       -- number of observations used
    """
    r = returns.dropna()
    n = len(r)

    if n == 0:
        return {
            "mean_ret": float("nan"),
            "vol": float("nan"),
            "sharpe": float("nan"),
            "information_ratio": float("nan"),
            "hit_ratio": float("nan"),
            "n_periods": 0,
        }

    mean_ret = float(r.mean() * freq)
    vol = float(r.std(ddof=1) * np.sqrt(freq))
    sharpe = float(mean_ret / vol) if vol != 0.0 else float("nan")
    hit_ratio = float((r > 0).mean())

    if benchmark is not None:
        b = benchmark.reindex(r.index).dropna()
        common_idx = r.index.intersection(b.index)
        if len(common_idx) >= 2:
            excess = r.loc[common_idx] - b.loc[common_idx]
            ir_mean = float(excess.mean() * freq)
            ir_vol = float(excess.std(ddof=1) * np.sqrt(freq))
            ir = float(ir_mean / ir_vol) if ir_vol != 0.0 else float("nan")
        else:
            ir = float("nan")
    else:
        ir = float("nan")

    return {
        "mean_ret": mean_ret,
        "vol": vol,
        "sharpe": sharpe,
        "information_ratio": ir,
        "hit_ratio": hit_ratio,
        "n_periods": n,
    }
