#!/usr/bin/env python3
"""Independent evaluator for Prior Moves next-filing rankings.

Standalone on purpose: it imports nothing from the Prior Moves code base, only
numpy and pandas, so a buyer can read all of it and run it on files they hold.

INPUTS
  predictions.csv  manager_id, period, security_id, score, prediction_ts
                   [, baseline_score]   one row per candidate, ranked within
                   (manager_id, period); higher score = more likely to appear.
  outcomes.csv     manager_id, period, security_id, label, outcome_public_ts
                   label 1 = the security appeared as a new position in the
                   manager's next filing; outcome_public_ts = when that filing
                   became public.
  --sha256 HEX     optional: the hash committed for predictions.csv before the
                   outcome existed. A mismatch fails the run.

WHAT IT COMPUTES
  pair-weighted AUC within (manager, period): concordant positive/negative
    pairs summed across strata over all comparable pairs. 0.5 = no ranking skill.
  macro AUC: mean of per-manager AUCs (each manager weighted equally).
  top-decile lift: hit rate in each stratum's top 10% over the base rate.
  precision@k for k in --k (default 5, 10).
  95% interval by resampling managers (clusters), --boot draws, fixed --seed.
  the same metrics for baseline_score when present, and the paired difference.
  timestamp check: every prediction_ts must precede its outcome_public_ts.

Exit code 0 = evaluated cleanly, 2 = integrity failure (hash or timestamps).

  python tools/pm_evaluate.py predictions.csv outcomes.csv [--sha256 HEX] [--json out.json]
"""
from __future__ import annotations

import argparse
import hashlib
import json
import sys

import numpy as np
import pandas as pd

KEY = ["manager_id", "period", "security_id"]


def stratum_stats(y: np.ndarray, s: np.ndarray) -> tuple[float, int] | None:
    pos, neg = s[y == 1], s[y == 0]
    if len(pos) == 0 or len(neg) == 0:
        return None
    # Mann-Whitney U via ranks, ties counted as half.
    ranks = pd.Series(np.concatenate([pos, neg])).rank(method="average").to_numpy()
    u = ranks[: len(pos)].sum() - len(pos) * (len(pos) + 1) / 2
    pairs = len(pos) * len(neg)
    return u / pairs, pairs


def metrics(df: pd.DataFrame, col: str, ks: list[int]) -> dict:
    strata, lifts, prec = [], [], {k: [] for k in ks}
    base = df["label"].mean()
    for (mgr, _), g in df.groupby(["manager_id", "period"], sort=False):
        y, s = g["label"].to_numpy(), g[col].to_numpy(dtype=float)
        st = stratum_stats(y, s)
        if st:
            strata.append((mgr, st[0], st[1]))
        order = np.argsort(-s, kind="stable")
        top = max(1, int(np.ceil(len(g) * 0.1)))
        lifts.append(y[order[:top]].mean())
        for k in ks:
            prec[k].append(y[order[:k]].mean())
    pairs = sum(p for _, _, p in strata)
    per_mgr: dict = {}
    for mgr, auc, p in strata:
        a, w = per_mgr.get(mgr, (0.0, 0))
        per_mgr[mgr] = (a + auc * p, w + p)
    return {
        "pair_weighted_auc": round(sum(a * p for _, a, p in strata) / pairs, 4) if pairs else None,
        "macro_auc": round(float(np.mean([a / w for a, w in per_mgr.values()])), 4) if per_mgr else None,
        "top_decile_lift": round(float(np.mean(lifts)) / base, 3) if base else None,
        **{f"precision_at_{k}": round(float(np.mean(v)), 4) for k, v in prec.items()},
        "strata": len(strata),
        "comparable_pairs": int(pairs),
        "_strata": strata,
    }


def cluster_ci(strata: list, draws: int, seed: int) -> list[float] | None:
    if not strata:
        return None
    by: dict = {}
    for mgr, auc, p in strata:
        by.setdefault(mgr, []).append((auc, p))
    mgrs = list(by)
    rng = np.random.default_rng(seed)
    out = []
    for _ in range(draws):
        pick = rng.choice(len(mgrs), len(mgrs), replace=True)
        rows = [r for i in pick for r in by[mgrs[i]]]
        out.append(sum(a * p for a, p in rows) / sum(p for _, p in rows))
    return [round(float(np.percentile(out, 2.5)), 4), round(float(np.percentile(out, 97.5)), 4)]


def paired_diff_ci(sm: list, sb: list, draws: int, seed: int) -> dict | None:
    """Model minus baseline, with both scored on the SAME resampled managers.

    The strata line up one to one because both scores share the labels. A
    difference whose interval excludes 0 is the claim "beats the baseline"."""
    if not sm or len(sm) != len(sb):
        return None
    by: dict = {}
    for (mgr, am, p), (_, ab, _) in zip(sm, sb):
        by.setdefault(mgr, []).append((am, ab, p))
    mgrs = list(by)
    rng = np.random.default_rng(seed)
    out = []
    for _ in range(draws):
        rows = [r for i in rng.choice(len(mgrs), len(mgrs), replace=True) for r in by[mgrs[i]]]
        w = sum(p for _, _, p in rows)
        out.append((sum(a * p for a, _, p in rows) - sum(b * p for _, b, p in rows)) / w)
    out = np.asarray(out)
    return {"ci95": [round(float(np.percentile(out, 2.5)), 4), round(float(np.percentile(out, 97.5)), 4)],
            "share_of_draws_above_zero": round(float((out > 0).mean()), 4), "draws": draws}


def main(argv: list[str] | None = None) -> int:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("predictions")
    ap.add_argument("outcomes")
    ap.add_argument("--sha256")
    ap.add_argument("--k", type=int, nargs="*", default=[5, 10])
    ap.add_argument("--boot", type=int, default=500)
    ap.add_argument("--seed", type=int, default=7)
    ap.add_argument("--json")
    a = ap.parse_args(argv)

    report: dict = {"integrity": {}}
    raw = open(a.predictions, "rb").read()
    digest = hashlib.sha256(raw).hexdigest()
    report["integrity"]["predictions_sha256"] = digest
    if a.sha256:
        report["integrity"]["hash_matches_commitment"] = digest == a.sha256.lower()

    p = pd.read_csv(a.predictions, dtype={"manager_id": str, "security_id": str, "period": str})
    o = pd.read_csv(a.outcomes, dtype={"manager_id": str, "security_id": str, "period": str})
    missing = [c for c in KEY + ["score", "prediction_ts"] if c not in p] + \
              [c for c in KEY + ["label", "outcome_public_ts"] if c not in o]
    if missing:
        print(f"missing columns: {missing}", file=sys.stderr)
        return 2
    dupes = int(p.duplicated(KEY).sum())
    df = p.merge(o, on=KEY, how="left", validate="one_to_one" if not dupes else None)
    unmatched = int(df["label"].isna().sum())
    df = df[df["label"].notna()].copy()
    df["label"] = df["label"].astype(int)
    pts = pd.to_datetime(df["prediction_ts"], utc=True, errors="coerce")
    ots = pd.to_datetime(df["outcome_public_ts"], utc=True, errors="coerce")
    late = int((pts.isna() | ots.isna() | (pts >= ots)).sum())
    report["integrity"].update({"duplicate_prediction_rows": dupes, "predictions_without_outcome": unmatched,
                                "prediction_not_before_outcome": late})
    report["coverage"] = {"rows_scored": int(len(df)), "managers": int(df["manager_id"].nunique()),
                          "periods": int(df["period"].nunique()), "base_rate": round(float(df["label"].mean()), 4)}

    model = metrics(df, "score", a.k)
    sm = model.pop("_strata")
    model["ci95_manager_clustered"] = cluster_ci(sm, a.boot, a.seed)
    report["model"] = model
    if "baseline_score" in df:
        base = metrics(df, "baseline_score", a.k)
        sb = base.pop("_strata")
        base["ci95_manager_clustered"] = cluster_ci(sb, a.boot, a.seed)
        report["baseline"] = base
        report["model_minus_baseline_auc"] = round(model["pair_weighted_auc"] - base["pair_weighted_auc"], 4)
        report["model_minus_baseline_paired"] = paired_diff_ci(sm, sb, a.boot, a.seed)

    failed = late > 0 or report["integrity"].get("hash_matches_commitment") is False
    report["verdict"] = "INTEGRITY_FAILURE" if failed else "EVALUATED"
    text = json.dumps(report, indent=1)
    if a.json:
        open(a.json, "w").write(text + "\n")
    print(text)
    return 2 if failed else 0


if __name__ == "__main__":
    raise SystemExit(main())