"""Статистическое сравнение вариантов по per-query рангам (протокол §8.3, §9.2).

Работает с сохранёнными `best_ranks_{q2g,g2q}.npy` (0-indexed ранг первого
попадания на запрос; M, если попадания нет). Из рангов выводится hit@k =
(ranks < k) для любого k — этого достаточно для:
    * ΔR@k (parированное улучшение вариант − baseline),
    * bootstrap CI парной дельты (ресэмпл запросов),
    * McNemar exact (парный тест на discordant-запросах: miss↔hit),
    * Wilcoxon signed-rank по разнице рангов (если установлен scipy),
    * применения decision-rule (§9.2): ΔR@1 ≥ порог, p<α, CI не в минусе,
      улучшение на ≥2/3 seeds.

Пары запросов валидны, т.к. eval детерминирован (shuffle=False, один
test_meta) → массивы рангов baseline и варианта совпадают по длине и порядку.

Многосидовый режим: передай несколько прогонов на baseline и variant (по
одному на seed, парно по порядку) — скрипт покажет per-seed дельты + пул.

Примеры:
    # один seed: B0 vs H-TA-3 (каталоги прогонов с best_ranks_q2g.npy)
    python compare.py \
        --baseline outputs/text_aug/B0/seed42/exp_..._cross-area \
        --variant  outputs/text_aug/H-TA-3/seed42/exp_..._cross-area \
        --label-baseline B0 --label-variant H-TA-3

    # три seeds (парно по порядку 42/123/456)
    python compare.py \
        --baseline B0_s42_dir B0_s123_dir B0_s456_dir \
        --variant  HTA3_s42_dir HTA3_s123_dir HTA3_s456_dir

    # прямые пути к .npy и другое направление/k
    python compare.py --baseline a/best_ranks_q2g.npy --variant b/best_ranks_q2g.npy --k 5

Примечание про intra/inter cos (опровержение H-TA-5): это метрика на
ЭМБЕДДИНГАХ, а не на рангах, поэтому здесь не считается (эмбеддинги в
артефактах не сохраняются). Её место — отдельный проход энкодера
(ср. eval_text_destruction.py), не этот скрипт.
"""
from __future__ import annotations

import argparse
import json
from math import comb
from pathlib import Path

import numpy as np


# ---------------------------------------------------------------------------
# Загрузка рангов
# ---------------------------------------------------------------------------

def load_ranks(path: str, direction: str, which: str = "best") -> np.ndarray:
    """Загрузить массив рангов из .npy-файла или каталога прогона.

    Args:
        path: путь к .npy ЛИБО к каталогу прогона (ищем {which}_ranks_{dir}.npy).
        direction: q2g (primary) или g2q.
        which: best (по лучшей эпохе) или last (последний eval).
    """
    p = Path(path)
    if p.is_dir():
        cand = p / f"{which}_ranks_{direction}.npy"
        if not cand.exists():
            # fallback: если нет best — берём last, и наоборот
            alt = p / f"{'last' if which == 'best' else 'best'}_ranks_{direction}.npy"
            if alt.exists():
                cand = alt
            else:
                raise FileNotFoundError(
                    f"Нет ranks-файла в {p} (искали {which}_ranks_{direction}.npy)")
        p = cand
    return np.load(p).astype(np.int64)


# ---------------------------------------------------------------------------
# Статистика
# ---------------------------------------------------------------------------

def bootstrap_ci_delta(
    base_hits: np.ndarray,
    var_hits: np.ndarray,
    n_boot: int = 10000,
    alpha: float = 0.05,
    seed: int = 42,
) -> tuple[float, float, float]:
    """Bootstrap 95% CI парной дельты R@k (ресэмпл запросов с возвратом).

    Returns:
        (mean_delta, ci_lo, ci_hi) — дельта = mean(var_hits − base_hits).
    """
    diff = var_hits.astype(np.float64) - base_hits.astype(np.float64)
    n = len(diff)
    rng = np.random.default_rng(seed)
    idx = rng.integers(0, n, size=(n_boot, n))
    boot = diff[idx].mean(axis=1)
    lo, hi = np.percentile(boot, [100 * alpha / 2, 100 * (1 - alpha / 2)])
    return float(diff.mean()), float(lo), float(hi)


def mcnemar_test(base_hits: np.ndarray, var_hits: np.ndarray) -> tuple[int, int, float]:
    """McNemar на discordant-запросах (двусторонний).

    b = baseline попал, variant промахнулся; c = наоборот. Под H0 discordant
    делятся 50/50. c > b → variant улучшает.

    Малое n (≤1000) — точный биномиальный (целочисленное деление, без
    переполнения float); большое n — нормальная аппроксимация χ² с поправкой
    на непрерывность (стандартный McNemar), т.к. точный comb(n,·) переполняет
    и медленный.

    Returns:
        (b, c, p_value).
    """
    from math import erfc, sqrt

    bh = base_hits.astype(bool)
    vh = var_hits.astype(bool)
    b = int(np.sum(bh & ~vh))
    c = int(np.sum(~bh & vh))
    n = b + c
    if n == 0:
        return b, c, 1.0
    if n <= 1000:
        k = min(b, c)
        tail = sum(comb(n, i) for i in range(k + 1))
        # int/int → корректный float без конверсии огромного int (2*, не 2.0*)
        p = min(1.0, (2 * tail) / (2 ** n))
    else:
        chi = (abs(b - c) - 1) ** 2 / n     # χ² с поправкой на непрерывность, 1 df
        p = float(erfc(sqrt(chi / 2.0)))    # p = 2·(1−Φ(√χ²)) для 1 df
    return b, c, p


def wilcoxon_ranks(base_ranks: np.ndarray, var_ranks: np.ndarray) -> float | None:
    """Wilcoxon signed-rank по разнице рангов (нужен scipy). None, если нет scipy
    или все разницы нулевые."""
    try:
        from scipy.stats import wilcoxon
    except Exception:
        return None
    # меньший ранг = лучше; d>0 → variant лучше на этом запросе
    d = base_ranks.astype(np.float64) - var_ranks.astype(np.float64)
    if np.all(d == 0):
        return None
    try:
        return float(wilcoxon(d, zero_method="wilcox", alternative="two-sided").pvalue)
    except Exception:
        return None


# ---------------------------------------------------------------------------
# Сравнение одной пары / агрегация
# ---------------------------------------------------------------------------

def compare_pair(
    base_ranks: np.ndarray,
    var_ranks: np.ndarray,
    k: int,
    n_boot: int,
    alpha: float,
    seed: int,
) -> dict:
    """Полное сравнение baseline vs variant на одном наборе запросов."""
    if base_ranks.shape != var_ranks.shape:
        raise ValueError(
            f"Разная длина рангов: baseline {base_ranks.shape} vs variant "
            f"{var_ranks.shape}. Нужен один и тот же test-split/direction.")
    base_hits = (base_ranks < k)
    var_hits = (var_ranks < k)

    r_base = float(base_hits.mean())
    r_var = float(var_hits.mean())
    mean_d, ci_lo, ci_hi = bootstrap_ci_delta(base_hits, var_hits, n_boot, alpha, seed)
    b, c, p_mcnemar = mcnemar_test(base_hits, var_hits)
    p_wilcoxon = wilcoxon_ranks(base_ranks, var_ranks)

    return {
        "n_queries": int(base_ranks.shape[0]),
        f"R@{k}_baseline": r_base,
        f"R@{k}_variant": r_var,
        "delta": mean_d,          # доля (умножь на 100 для п.п.)
        "ci_lo": ci_lo,
        "ci_hi": ci_hi,
        "mcnemar_b": b,           # base hit, var miss
        "mcnemar_c": c,           # base miss, var hit (улучшения)
        "p_mcnemar": p_mcnemar,
        "p_wilcoxon": p_wilcoxon,
    }


def decide(delta: float, ci_lo: float, p: float,
           threshold_pp: float, alpha: float,
           seeds_improved: int, seeds_total: int) -> tuple[bool, str]:
    """Применить decision-rule §9.2. Возвращает (promote, причина)."""
    reasons = []
    ok_delta = delta * 100 >= threshold_pp
    ok_ci = ci_lo > 0.0
    ok_p = (p is not None) and (p < alpha)
    ok_seeds = seeds_improved >= (2 * seeds_total + 2) // 3 if seeds_total > 1 else True

    if not ok_delta:
        reasons.append(f"ΔR@1={delta*100:+.2f}пп < порог {threshold_pp}")
    if not ok_ci:
        reasons.append(f"CI_lo={ci_lo*100:+.2f}пп ≤ 0 (пересекает минус)")
    if not ok_p:
        reasons.append(f"p={p:.4f} ≥ α={alpha}" if p is not None else "нет p")
    if seeds_total > 1 and not ok_seeds:
        reasons.append(f"улучшение лишь на {seeds_improved}/{seeds_total} seeds")

    promote = ok_delta and ok_ci and ok_p and ok_seeds
    return promote, ("все критерии выполнены" if promote else "; ".join(reasons))


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

def main():
    args = parse_args()
    if len(args.baseline) != len(args.variant):
        raise SystemExit("Число --baseline и --variant должно совпадать (пары по seed)")

    print("=" * 72)
    print(f"  СРАВНЕНИЕ: {args.label_variant} vs {args.label_baseline}  "
          f"| direction={args.direction} | k={args.k}")
    print("=" * 72)

    per_seed = []
    pooled_base, pooled_var = [], []
    for i, (b_path, v_path) in enumerate(zip(args.baseline, args.variant)):
        base_ranks = load_ranks(b_path, args.direction, args.which)
        var_ranks = load_ranks(v_path, args.direction, args.which)
        res = compare_pair(base_ranks, var_ranks, args.k, args.n_boot, args.alpha, args.seed)
        res["seed_idx"] = i
        per_seed.append(res)
        pooled_base.append((base_ranks < args.k))
        pooled_var.append((var_ranks < args.k))

    # ---- per-seed таблица ----
    print(f"\n{'seed#':>5} {'N':>7} {'R@'+str(args.k)+' base':>10} "
          f"{'R@'+str(args.k)+' var':>10} {'ΔR@'+str(args.k):>10} "
          f"{'95% CI (пп)':>18} {'p(McN)':>9} {'p(Wilc)':>9}")
    print("-" * 88)
    for r in per_seed:
        pw = f"{r['p_wilcoxon']:.4f}" if r["p_wilcoxon"] is not None else "—"
        print(f"{r['seed_idx']:>5} {r['n_queries']:>7} "
              f"{r[f'R@{args.k}_baseline']*100:>9.2f}% {r[f'R@{args.k}_variant']*100:>9.2f}% "
              f"{r['delta']*100:>+9.2f} "
              f"[{r['ci_lo']*100:>+6.2f},{r['ci_hi']*100:>+6.2f}] "
              f"{r['p_mcnemar']:>9.4f} {pw:>9}")
    print("-" * 88)

    # ---- пул по всем seeds (общий CI/тест) ----
    pooled_base = np.concatenate(pooled_base)
    pooled_var = np.concatenate(pooled_var)
    mean_d, ci_lo, ci_hi = bootstrap_ci_delta(
        pooled_base, pooled_var, args.n_boot, args.alpha, args.seed)
    b, c, p_mcn = mcnemar_test(pooled_base, pooled_var)
    seeds_improved = sum(1 for r in per_seed if r["delta"] > 0)
    n_seeds = len(per_seed)

    print(f"\nПул ({n_seeds} seed(s), N={len(pooled_base)}): "
          f"ΔR@{args.k}={mean_d*100:+.2f} пп | "
          f"95% CI [{ci_lo*100:+.2f}, {ci_hi*100:+.2f}] пп | "
          f"McNemar p={p_mcn:.4f} (b={b}, c={c}) | "
          f"улучшение на {seeds_improved}/{n_seeds} seeds")

    promote, reason = decide(mean_d, ci_lo, p_mcn, args.threshold, args.alpha,
                             seeds_improved, n_seeds)
    verdict = "✅ PROMOTE" if promote else "❌ NOT PROMOTE"
    print("\n" + "=" * 72)
    print(f"  DECISION (§9.2): {verdict}")
    print(f"  Причина: {reason}")
    print(f"  Правило: ΔR@1 ≥ {args.threshold} пп, p<{args.alpha}, "
          f"CI_lo>0, ≥2/3 seeds, без modality-collapse (gate — проверь отдельно)")
    print("=" * 72)

    if args.output:
        out = {
            "baseline": args.label_baseline,
            "variant": args.label_variant,
            "direction": args.direction,
            "k": args.k,
            "per_seed": per_seed,
            "pooled": {
                "delta": mean_d, "ci_lo": ci_lo, "ci_hi": ci_hi,
                "p_mcnemar": p_mcn, "mcnemar_b": b, "mcnemar_c": c,
                "seeds_improved": seeds_improved, "n_seeds": n_seeds,
            },
            "promote": promote, "reason": reason,
        }
        Path(args.output).parent.mkdir(parents=True, exist_ok=True)
        Path(args.output).write_text(json.dumps(out, indent=2, ensure_ascii=False))
        print(f"\n💾 Сохранено: {args.output}")


def parse_args():
    p = argparse.ArgumentParser(description="Парное статсравнение вариантов по рангам (§9.2)")
    p.add_argument("--baseline", nargs="+", required=True,
                   help="Каталог(и) прогона B0 или .npy (по одному на seed)")
    p.add_argument("--variant", nargs="+", required=True,
                   help="Каталог(и) прогона варианта или .npy (парно к baseline)")
    p.add_argument("--direction", choices=["q2g", "g2q"], default="q2g",
                   help="q2g (primary, drone→sat) или g2q")
    p.add_argument("--which", choices=["best", "last"], default="best")
    p.add_argument("--k", type=int, default=1, help="k для R@k (по умолчанию 1)")
    p.add_argument("--n_boot", type=int, default=10000)
    p.add_argument("--alpha", type=float, default=0.05)
    p.add_argument("--threshold", type=float, default=1.0,
                   help="Порог ΔR@1 в п.п. для promote (§9.2, дефолт 1.0)")
    p.add_argument("--seed", type=int, default=42, help="Seed бутстрапа")
    p.add_argument("--label-baseline", default="baseline")
    p.add_argument("--label-variant", default="variant")
    p.add_argument("--output", default=None, help="Куда сохранить JSON-итог")
    return p.parse_args()


if __name__ == "__main__":
    main()
