# Author: Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca
"""Paired inference between arms (docs/evaluation_protocol.md §4).

* Block bootstrap: blocks = mun_code × calendar month (fallback mun_code, then i.i.d.), B = 2000, percentile CI
  and a centred bootstrap p-value for Δ = metric(B) − metric(A) computed on the same resampled blocks for both arms.
* Holm step-down correction for a family of contrasts.
* Diebold–Mariano test on the per-record loss differential d_i = |e_i^B| − |e_i^A| with a Newey–West (HAC) variance.
  The HAC variance is only meaningful when the series is ordered in time; the function therefore sorts by `date`
  and refuses to run when no date is available (`dm=False` in `compare_arms` skips it).
"""
from __future__ import annotations

from collections.abc import Callable, Sequence
from typing import Any

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

from ._common import arm_filter, arm_label, block_ids
from .metrics import PRIMARY_METRIC, metric_statistic, per_record

Statistic = Callable[[pd.DataFrame], float]


def _resolve_statistic(metric: str | Statistic) -> tuple[str, Statistic]:
    if callable(metric):
        return getattr(metric, "__name__", "custom"), metric
    return metric, metric_statistic(metric)


def paired_records(preds_all: Any, truth: Any, spec_a: str | dict, spec_b: str | dict) -> tuple[pd.DataFrame, pd.DataFrame]:
    """Per-record frames of arms A and B restricted to the common tx_ids, in the same order."""
    d = per_record(preds_all, truth)
    a = arm_filter(d, spec_a).set_index("tx_id")
    b = arm_filter(d, spec_b).set_index("tx_id")
    common = a.index.intersection(b.index)
    if len(common) == 0:
        raise ValueError(f"arms {arm_label(spec_a)} and {arm_label(spec_b)} share no tx_id")
    return a.loc[common].reset_index(), b.loc[common].reset_index()


def block_bootstrap_contrast(
    a: pd.DataFrame,
    b: pd.DataFrame,
    statistic: str | Statistic = PRIMARY_METRIC,
    *,
    B: int = 2000,
    seed: int = 20261027,
    alpha: float = 0.05,
    blocks: pd.Series | None = None,
) -> dict[str, float]:
    """Paired block bootstrap of Δ = stat(b) − stat(a) (negative Δ = arm B better for a loss).

    `a` and `b` are per-record frames aligned on the same tx_ids (same order). Blocks are resampled with
    replacement (Künsch-style moving blocks are not used: blocks are the natural clusters municipality × month).
    Returns point estimate, percentile CI, centred bootstrap p-value and the bootstrap SE.
    """
    if len(a) != len(b) or not (a["tx_id"].to_numpy() == b["tx_id"].to_numpy()).all():
        raise ValueError("a and b must be aligned on the same tx_ids")
    name, stat = _resolve_statistic(statistic)
    blk = (blocks if blocks is not None else block_ids(a)).to_numpy()
    uniq, inv = np.unique(blk, return_inverse=True)
    n_blocks = len(uniq)
    members = [np.flatnonzero(inv == k) for k in range(n_blocks)]
    rng = np.random.default_rng(seed)
    delta_hat = stat(b) - stat(a)
    nan_result = {"metric": name, "n": float(len(a)), "n_blocks": float(n_blocks), "B": 0.0, "value_a": float(stat(a)), "value_b": float(stat(b)),
                  "diff": float(delta_hat), "ci_lo": np.nan, "ci_hi": np.nan, "se_boot": np.nan, "p_boot": np.nan}
    if not np.isfinite(delta_hat):  # e.g. coverage contrast involving an arm without intervals
        return nan_result
    draws = np.empty(B)
    for r in range(B):
        picked = rng.integers(0, n_blocks, size=n_blocks)
        idx = np.concatenate([members[k] for k in picked])
        draws[r] = stat(b.iloc[idx]) - stat(a.iloc[idx])
    draws = draws[np.isfinite(draws)]
    if len(draws) < 10:
        return nan_result
    lo, hi = np.percentile(draws, [100 * alpha / 2, 100 * (1 - alpha / 2)])
    # centred bootstrap p-value (Efron & Tibshirani 1993 §15): shift the bootstrap distribution to the null
    centred = draws - delta_hat
    p = float(np.mean(np.abs(centred) >= abs(delta_hat)))
    p = max(p, 1.0 / (len(draws) + 1))  # never report an exact zero
    return {
        "metric": name, "n": float(len(a)), "n_blocks": float(n_blocks), "B": float(len(draws)),
        "value_a": float(stat(a)), "value_b": float(stat(b)), "diff": float(delta_hat),
        "ci_lo": float(lo), "ci_hi": float(hi), "se_boot": float(np.std(draws, ddof=1)), "p_boot": p,
    }


def holm(pvals: Sequence[float]) -> np.ndarray:
    """Holm (1979) step-down adjusted p-values (monotone, capped at 1). NaNs are passed through."""
    p = np.asarray(pvals, dtype=float)
    out = np.full_like(p, np.nan)
    ok = np.isfinite(p)
    pv = p[ok]
    m = len(pv)
    if m == 0:
        return out
    order = np.argsort(pv)
    adj = np.empty(m)
    running = 0.0
    for rank, i in enumerate(order):
        val = min(1.0, (m - rank) * pv[i])
        running = max(running, val)
        adj[i] = running
    out[ok] = adj
    return out


def newey_west_variance(d: np.ndarray, lag: int | None = None) -> tuple[float, int]:
    """Long-run variance of the mean of `d` with Bartlett kernel; default bandwidth floor(4 (n/100)^{2/9})."""
    d = np.asarray(d, dtype=float)
    n = len(d)
    if lag is None:
        lag = int(np.floor(4 * (n / 100.0) ** (2.0 / 9.0)))
    dc = d - d.mean()
    gamma0 = float(dc @ dc) / n
    lrv = gamma0
    for k in range(1, lag + 1):
        w = 1.0 - k / (lag + 1.0)
        lrv += 2.0 * w * float(dc[k:] @ dc[:-k]) / n
    return max(lrv, 1e-300), lag


def diebold_mariano(
    loss_a: np.ndarray, loss_b: np.ndarray, *, order: np.ndarray | None = None, lag: int | None = None, harvey: bool = True
) -> dict[str, float]:
    """DM test of H0: E[loss_b − loss_a] = 0 with HAC variance.

    ASSUMPTION: the loss differential is a covariance-stationary time series ordered by date. Pass `order`
    (e.g. argsort of the sale dates) — the series is reordered before computing autocovariances. Harvey, Leybourne
    & Newbold (1997) small-sample correction applied by default (t_{n−1} reference distribution).
    """
    la, lb = np.asarray(loss_a, dtype=float), np.asarray(loss_b, dtype=float)
    if order is not None:
        la, lb = la[order], lb[order]
    d = lb - la
    n = len(d)
    lrv, h = newey_west_variance(d, lag)
    dm = d.mean() / np.sqrt(lrv / n)
    if harvey:
        corr = np.sqrt((n + 1 - 2 * (h + 1) + (h + 1) * h / n) / n)
        dm = dm * corr
        p = 2 * stats.t.sf(abs(dm), df=n - 1)
    else:
        p = 2 * stats.norm.sf(abs(dm))
    return {"dm_stat": float(dm), "dm_p": float(p), "dm_lag": float(h), "dm_mean_diff": float(d.mean())}


def compare_arms(
    preds_all: Any,
    truth: Any,
    pairs: Sequence[tuple[str | dict, str | dict]],
    metric: str | Statistic | Sequence[str | Statistic] = PRIMARY_METRIC,
    *,
    B: int = 2000,
    seed: int = 20261027,
    alpha: float = 0.05,
    dm: bool = True,
    holm_family: bool = True,
) -> pd.DataFrame:
    """Tidy table of paired contrasts. Δ = metric(B) − metric(A): for a loss, Δ < 0 means arm B is better.

    One row per (pair, metric). Holm correction is applied within each metric across the supplied pairs
    (the pre-registered family) when `holm_family` is True. DM columns are NaN when no `date` is available.
    """
    metrics = list(metric) if isinstance(metric, (list, tuple)) else [metric]
    rows = []
    for spec_a, spec_b in pairs:
        a, b = paired_records(preds_all, truth, spec_a, spec_b)
        for m in metrics:
            res = block_bootstrap_contrast(a, b, m, B=B, seed=seed, alpha=alpha)
            row = {"arm_a": arm_label(spec_a), "arm_b": arm_label(spec_b)}
            row.update(res)
            row.update({"dm_stat": np.nan, "dm_p": np.nan, "dm_lag": np.nan})
            if dm and "date" in a.columns and a["date"].notna().all() and res["metric"] in {"mean_abs_log_error", "mae_cad", "rmse_log", "pinball_log", "crps_log"}:
                col = {"mean_abs_log_error": "abs_e", "mae_cad": "abs_err_cad", "rmse_log": "e", "pinball_log": "pinball_log", "crps_log": "crps_log"}[res["metric"]]
                la, lb = a[col].to_numpy(dtype=float), b[col].to_numpy(dtype=float)
                if col == "e":
                    la, lb = la**2, lb**2  # DM on squared errors for RMSE
                ok = np.isfinite(la) & np.isfinite(lb)
                if ok.sum() >= 10:
                    order = np.argsort(pd.to_datetime(a["date"]).to_numpy()[ok], kind="stable")
                    row.update(diebold_mariano(la[ok], lb[ok], order=order))
            rows.append(row)
    out = pd.DataFrame(rows)
    if out.empty:
        return out
    out["p_holm"] = np.nan
    if holm_family:
        for m, idx in out.groupby("metric").groups.items():
            out.loc[idx, "p_holm"] = holm(out.loc[idx, "p_boot"].to_numpy())
    out["rel_diff"] = out["diff"] / out["value_a"]
    cols = ["arm_a", "arm_b", "metric", "n", "n_blocks", "B", "value_a", "value_b", "diff", "rel_diff", "ci_lo", "ci_hi",
            "se_boot", "p_boot", "p_holm", "dm_stat", "dm_p", "dm_lag"]
    return out[[c for c in cols if c in out.columns]]