# Author: Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca
"""Per-record errors and summary metrics (definitions in docs/evaluation_protocol.md §1–§3).

All errors are computed on the sale price P_i (a noisy proxy of value, never "truth" in the economic sense):
    e_i = log(point_i) − log(P_i).
Interval metrics use the five quantiles q05..q95 (CAD). The 50 % and 90 % intervals are read directly from the
quantiles ([q25, q75] and [q05, q95]); the 80 % and 95 % intervals are *interpolated* from a normal distribution
fitted in log space to the five quantiles (least squares on the probit grid). Both raw and fitted versions are reported.
"""
from __future__ import annotations

from collections.abc import Sequence
from typing import Any

import numpy as np
import pandas as pd
from scipy import stats

from ._common import QUANTILE_COLUMNS, TAUS, join_truth

LOG_125 = float(np.log(1.25))
Z_TAUS = stats.norm.ppf(TAUS)               # probit grid of the five quantiles
Z_SUMSQ = float(np.sum(Z_TAUS**2))
NOMINAL_LEVELS = (0.50, 0.80, 0.90, 0.95)

PRIMARY_METRIC = "mean_abs_log_error"
METRIC_LABELS: dict[str, str] = {
    "n": "N",
    "mean_abs_log_error": r"Mean $|e|$",
    "mae_cad": "MAE (CAD)",
    "mdape": "MdAPE",
    "rmse_log": "RMSE (log)",
    "p90_abs_log_error": r"P90 $|e|$",
    "p95_abs_log_error": r"P95 $|e|$",
    "bias_log": r"Bias $\bar e$",
    "severe_rate": r"$|e|>\log 1.25$",
    "over_rate": "Over +25\\%",
    "under_rate": "Under $-$25\\%",
    "coverage_50": "Cov. 50\\%",
    "coverage_80": "Cov. 80\\%",
    "coverage_90": "Cov. 90\\%",
    "coverage_95": "Cov. 95\\%",
    "coverage_90_normal_fit": "Cov. 90\\% (fit)",
    "mean_log_width_90": "Width 90\\% (log)",
    "mean_log_width_50": "Width 50\\% (log)",
    "pinball_log": "Pinball (log)",
    "crps_log": "CRPS (log)",
    "crps_log_normal_fit": "CRPS (log, fit)",
    "cost_usd_mean": "Cost (USD)",
    "latency_s_mean": "Latency (s)",
    "tokens_mean": "Tokens",
    "cost_usd_total": "Total cost (USD)",
}


# ----------------------------------------------------------------------------------------------- per-record

def fit_lognormal_from_quantiles(logq: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
    """Least-squares fit of log q_τ = μ + σ z_τ over the five quantiles (rows = records). Returns (μ, σ)."""
    logq = np.asarray(logq, dtype=float)
    mu = np.full(logq.shape[0], np.nan)
    sigma = np.full(logq.shape[0], np.nan)
    ok = np.all(np.isfinite(logq), axis=1)
    if ok.any():
        mu[ok] = logq[ok].mean(axis=1)
        sigma[ok] = np.maximum((logq[ok] - mu[ok, None]) @ Z_TAUS / Z_SUMSQ, 0.0)
    return mu, sigma


def _pinball(y: np.ndarray, q: np.ndarray, tau: float) -> np.ndarray:
    d = y - q
    return np.where(d >= 0, tau * d, (tau - 1) * d)


def per_record(preds: Any, truth: Any) -> pd.DataFrame:
    """Join predictions to truth and add per-record error / interval columns.

    Added columns: log_price, log_point, e, abs_e, abs_err_cad, ape, severe, over, under,
    cov_50, cov_90 (raw), cov_80, cov_95 (normal fit), cov_90_fit, width_90, width_50 (log), mu_fit, sigma_fit,
    pinball_log (mean over the 5 quantiles), crps_log (quantile approximation), crps_log_fit (closed form of fitted normal).
    Interval columns are NaN when the quantiles are missing.
    """
    d = join_truth(preds, truth)
    y = np.log(d["price"].to_numpy(dtype=float))
    yhat = np.log(d["point"].to_numpy(dtype=float))
    e = yhat - y
    d["log_price"], d["log_point"], d["e"], d["abs_e"] = y, yhat, e, np.abs(e)
    d["abs_err_cad"] = np.abs(d["point"] - d["price"])
    d["ape"] = np.abs(d["point"] / d["price"] - 1.0)
    d["severe"] = (d["abs_e"] > LOG_125).astype(float)
    d["over"] = (d["point"] / d["price"] - 1.0 > 0.25).astype(float)
    d["under"] = (d["point"] / d["price"] - 1.0 < -0.25).astype(float)

    q = d[QUANTILE_COLUMNS].to_numpy(dtype=float)
    has_q = np.all(np.isfinite(q), axis=1) & np.all(q > 0, axis=1)
    logq = np.where(has_q[:, None], np.log(np.where(q > 0, q, np.nan)), np.nan)
    nan = np.full(len(d), np.nan)

    def masked(values: np.ndarray) -> np.ndarray:
        return np.where(has_q, values, nan)

    price = d["price"].to_numpy(dtype=float)
    d["cov_50"] = masked(((q[:, 1] <= price) & (price <= q[:, 3])).astype(float))
    d["cov_90"] = masked(((q[:, 0] <= price) & (price <= q[:, 4])).astype(float))
    mu, sigma = fit_lognormal_from_quantiles(logq)
    d["mu_fit"], d["sigma_fit"] = masked(mu), masked(sigma)
    for level, name in ((0.80, "cov_80"), (0.95, "cov_95"), (0.90, "cov_90_fit")):
        z = stats.norm.ppf(0.5 + level / 2)
        d[name] = masked(((mu - z * sigma <= y) & (y <= mu + z * sigma)).astype(float))
    d["width_90"] = masked(logq[:, 4] - logq[:, 0])
    d["width_50"] = masked(logq[:, 3] - logq[:, 1])
    pb = np.column_stack([_pinball(y, logq[:, k], float(TAUS[k])) for k in range(5)])
    d["pinball_log"] = masked(pb.mean(axis=1))
    # CRPS ≈ 2 ∫ ρ_τ dτ, trapezoidal rule on the τ-grid 0.05..0.95 (the two outer 5 % bands are not observed and are
    # dropped, so this is a lower-biased approximation that is consistent across arms; see protocol §3.4)
    d["crps_log"] = masked(2.0 * np.trapezoid(pb, TAUS, axis=1))
    # closed-form CRPS of the fitted normal (Gneiting & Raftery 2007, eq. 21)
    zz = np.where(sigma > 0, (y - mu) / np.where(sigma > 0, sigma, 1.0), 0.0)
    crps_n = sigma * (zz * (2 * stats.norm.cdf(zz) - 1) + 2 * stats.norm.pdf(zz) - 1 / np.sqrt(np.pi))
    d["crps_log_fit"] = masked(np.where(sigma > 0, crps_n, np.abs(y - mu)))
    d["has_interval"] = has_q
    return d


# ----------------------------------------------------------------------------------------------- summaries

def _nanmean(x: pd.Series) -> float:
    v = x.to_numpy(dtype=float)
    return float(np.nanmean(v)) if np.isfinite(v).any() else float("nan")


def summarise_records(d: pd.DataFrame) -> dict[str, float]:
    """Summary metrics from a per-record frame (output of `per_record`)."""
    abs_e = d["abs_e"].to_numpy(dtype=float)
    out: dict[str, float] = {
        "n": float(len(d)),
        "mean_abs_log_error": float(np.mean(abs_e)) if len(d) else float("nan"),
        "mae_cad": float(d["abs_err_cad"].mean()),
        "mdape": float(d["ape"].median()),
        "rmse_log": float(np.sqrt(np.mean(d["e"].to_numpy(dtype=float) ** 2))) if len(d) else float("nan"),
        "p90_abs_log_error": float(np.percentile(abs_e, 90)) if len(d) else float("nan"),
        "p95_abs_log_error": float(np.percentile(abs_e, 95)) if len(d) else float("nan"),
        "bias_log": float(d["e"].mean()),
        "severe_rate": float(d["severe"].mean()),
        "over_rate": float(d["over"].mean()),
        "under_rate": float(d["under"].mean()),
        "coverage_50": _nanmean(d["cov_50"]),
        "coverage_80": _nanmean(d["cov_80"]),
        "coverage_90": _nanmean(d["cov_90"]),
        "coverage_95": _nanmean(d["cov_95"]),
        "coverage_90_normal_fit": _nanmean(d["cov_90_fit"]),
        "mean_log_width_90": _nanmean(d["width_90"]),
        "mean_log_width_50": _nanmean(d["width_50"]),
        "pinball_log": _nanmean(d["pinball_log"]),
        "crps_log": _nanmean(d["crps_log"]),
        "crps_log_normal_fit": _nanmean(d["crps_log_fit"]),
        "n_with_interval": float(d["has_interval"].sum()),
        "cost_usd_mean": _nanmean(d["cost_usd"]),
        "cost_usd_total": float(np.nansum(d["cost_usd"].to_numpy(dtype=float))) if d["cost_usd"].notna().any() else float("nan"),
        "latency_s_mean": _nanmean(d["latency_s"]),
        "tokens_mean": _nanmean(d["tokens"]),
    }
    return out


def summarise(preds: Any, truth: Any) -> dict[str, float]:
    """All summary metrics for one prediction table (one arm) joined to truth."""
    return summarise_records(per_record(preds, truth))


def summarise_by(preds: Any, truth: Any, by: Sequence[str] = ("arm", "committee", "aggregation")) -> pd.DataFrame:
    """One row of summary metrics per group (default: per arm × committee × aggregation)."""
    d = per_record(preds, truth)
    by = list(by)
    rows = []
    for key, g in d.groupby(by, sort=True, dropna=False):
        key = key if isinstance(key, tuple) else (key,)
        row = dict(zip(by, key))
        row.update(summarise_records(g))
        rows.append(row)
    return pd.DataFrame(rows)


# per-record loss columns usable as bootstrap statistics (metric name -> (column, reducer))
RECORD_METRICS: dict[str, tuple[str, str]] = {
    "mean_abs_log_error": ("abs_e", "mean"),
    "mae_cad": ("abs_err_cad", "mean"),
    "mdape": ("ape", "median"),
    "rmse_log": ("e", "rms"),
    "p90_abs_log_error": ("abs_e", "p90"),
    "p95_abs_log_error": ("abs_e", "p95"),
    "bias_log": ("e", "mean"),
    "severe_rate": ("severe", "mean"),
    "over_rate": ("over", "mean"),
    "under_rate": ("under", "mean"),
    "coverage_50": ("cov_50", "nanmean"),
    "coverage_80": ("cov_80", "nanmean"),
    "coverage_90": ("cov_90", "nanmean"),
    "coverage_95": ("cov_95", "nanmean"),
    "mean_log_width_90": ("width_90", "nanmean"),
    "pinball_log": ("pinball_log", "nanmean"),
    "crps_log": ("crps_log", "nanmean"),
}


def metric_statistic(name: str):
    """Return a callable DataFrame -> float computing the named summary metric from a per-record frame."""
    col, red = RECORD_METRICS[name]
    reducers = {
        "mean": lambda v: float(np.mean(v)),
        "nanmean": lambda v: float(np.nanmean(v)) if np.isfinite(v).any() else float("nan"),
        "median": lambda v: float(np.median(v)),
        "rms": lambda v: float(np.sqrt(np.mean(v**2))),
        "p90": lambda v: float(np.percentile(v, 90)),
        "p95": lambda v: float(np.percentile(v, 95)),
    }
    f = reducers[red]
    return lambda df: f(df[col].to_numpy(dtype=float))