#!/usr/bin/env python3
"""Score a served Starburst model on held-out, labelled snapshots.

It answers the one question the training can't: did KTO make the model *select* better —
**skip the leader's losers, keep the leader's winners** — on data with known outcomes?

    # serve the model first (e.g. ./serve_vllm.sh), then:
    python3 eval_holdout.py --data holdout.jsonl --model starburst-2.4-9b
    # compare vs 2.3 (serve it on :8001) :
    python3 eval_holdout.py --data holdout.jsonl --model starburst-2.3-9b --endpoint http://127.0.0.1:8001/v1 --baseline eval_starburst-2.4-9b.json

holdout.jsonl = replay records (messages + output.parsed + outcome + record_type).

  ┌──────────────────────────────────────────────────────────────────────────┐
  │ DATA REQUIREMENT: for a TRUE generalisation number, holdout.jsonl must be  │
  │ FRESH snapshots the model never trained on (all of the 10k+1.2k was used). │
  │ Generate a held-out month on the replay harness. Running on in-sample rows │
  │ is only a behavioural sanity check — the script warns when it can't tell.  │
  └──────────────────────────────────────────────────────────────────────────┘

Metrics:
  loss-avoidance = of leader LOSERS, % the model now SKIPS or flips   (KTO working -> high)
  win-retention  = of leader WINNERS, % the model still TAKES same side (not over-cautious)
  gated sumR     = leader's P&L if you only took the trades the model also approved
                   (same side). > leader's raw sumR  => the model filtered out -EV trades.
"""
import argparse, json, re, sys, os
from collections import Counter
from concurrent.futures import ThreadPoolExecutor, as_completed
import urllib.request

DEC_RE = re.compile(r"Decision:\s*(LONG|SHORT|NO\s*TRADE)", re.I)


def parse_side(text):
    if "</think>" in text:
        text = text.split("</think>")[-1]
    m = DEC_RE.search(text)
    if not m:
        return None
    s = m.group(1).upper().replace(" ", "")
    return "NO_TRADE" if "NOTRADE" in s else s


def infer(messages, endpoint, model, temp, timeout=120):
    body = json.dumps({"model": model, "messages": messages,
                       "temperature": temp, "max_tokens": 768}).encode()
    req = urllib.request.Request(endpoint.rstrip("/") + "/chat/completions",
                                 body, {"Content-Type": "application/json"})
    with urllib.request.urlopen(req, timeout=timeout) as r:
        msg = json.load(r)["choices"][0]["message"]
    return (msg.get("content") or "")


def leader_view(row):
    """(leader_side, bucket, r) -- bucket in win/loss/skip/other."""
    p = row.get("output", {}).get("parsed") or {}
    side = (p.get("decision") or "").upper()
    rt = row.get("record_type"); o = row.get("outcome") or {}
    if rt == "skip":
        return side or "NO_TRADE", "skip", 0.0
    if rt == "resolved" and not o.get("ambiguous"):
        st = o.get("status")
        if st in ("win", "loss"):
            return side, st, float(o.get("r_multiple") or 0)
    return side, "other", 0.0


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--data", required=True)
    ap.add_argument("--model", required=True)
    ap.add_argument("--endpoint", default="http://127.0.0.1:8000/v1")
    ap.add_argument("--temp", type=float, default=0.3)
    ap.add_argument("--workers", type=int, default=8)
    ap.add_argument("--limit", type=int, default=0)
    ap.add_argument("--baseline", default="", help="a prior eval_*.json to print side-by-side")
    args = ap.parse_args()

    rows = [json.loads(l) for l in open(args.data) if l.strip()]
    rows = [r for r in rows if r.get("messages")]
    if args.limit:
        rows = rows[:args.limit]

    # in-sample guard: if these snapshots match the training contract system, we can't prove
    # they're unseen -- warn (a fresh holdout from the replay harness is the real test).
    print(f"loaded {len(rows)} rows | model={args.model} | endpoint={args.endpoint}")
    print("⚠  if this data was part of the 10k/1.2k training set, treat results as an "
          "in-sample behavioural check, NOT generalisation.\n")

    results = [None] * len(rows)

    def work(i):
        try:
            return i, parse_side(infer(rows[i]["messages"][:2], args.endpoint, args.model, args.temp))
        except Exception as e:
            return i, f"ERR:{type(e).__name__}"

    done = 0
    with ThreadPoolExecutor(max_workers=args.workers) as ex:
        for fut in as_completed([ex.submit(work, i) for i in range(len(rows))]):
            i, side = fut.result(); results[i] = side
            done += 1
            if done % 50 == 0 or done == len(rows):
                print(f"  ...{done}/{len(rows)}", file=sys.stderr)

    errs = sum(1 for s in results if isinstance(s, str) and s.startswith("ERR"))
    none = sum(1 for s in results if s is None)
    M = Counter()                      # model decision tally
    cm = Counter()                     # outcome buckets
    leader_sumR = model_gated_sumR = 0.0
    flips = 0
    for row, mside in zip(rows, results):
        if not isinstance(mside, str) or mside.startswith("ERR"):
            continue
        lside, bucket, r = leader_view(row)
        M[mside] += 1
        if bucket == "win":
            cm["leader_win"] += 1; leader_sumR += r
            if mside == lside:
                cm["kept_win"] += 1; model_gated_sumR += r
            elif mside == "NO_TRADE":
                cm["missed_win"] += 1
            else:
                cm["flip_on_win"] += 1; flips += 1
        elif bucket == "loss":
            cm["leader_loss"] += 1; leader_sumR += r
            if mside == "NO_TRADE":
                cm["avoided_loss"] += 1
            elif mside == lside:
                cm["repeated_loss"] += 1; model_gated_sumR += r
            else:
                cm["flip_on_loss"] += 1; flips += 1
        elif bucket == "skip":
            cm["leader_skip"] += 1
            cm["agree_skip" if mside == "NO_TRADE" else "overrode_skip"] += 1

    n_ok = len(rows) - errs - none
    lw, ll = cm["leader_win"], cm["leader_loss"]
    out = {
        "model": args.model, "n": len(rows), "ok": n_ok, "errors": errs + none,
        "model_decisions": dict(M),
        "loss_avoidance": round(cm["avoided_loss"] / ll, 3) if ll else None,
        "win_retention": round(cm["kept_win"] / lw, 3) if lw else None,
        "leader_sumR": round(leader_sumR, 1),
        "model_gated_sumR": round(model_gated_sumR, 1),
        "confusion": dict(cm),
        "flips": flips,
    }
    save = f"eval_{args.model}.json"
    json.dump(out, open(save, "w"), indent=2)

    print(f"\n==== {args.model}  ({n_ok} scored, {errs+none} failed) ====")
    print(f"model decisions: {dict(M)}")
    print(f"LOSS-AVOIDANCE : {out['loss_avoidance']}  (of {ll} leader losers, skipped/flipped; higher=KTO working)")
    print(f"WIN-RETENTION  : {out['win_retention']}  (of {lw} leader winners, kept same side; higher=not over-cautious)")
    print(f"gated sumR     : {out['model_gated_sumR']:+}  vs leader raw {out['leader_sumR']:+}  "
          f"({'BETTER selection' if out['model_gated_sumR']>out['leader_sumR'] else 'no improvement'})")
    print(f"confusion: {dict(cm)}")
    print(f"-> saved {save}")

    if args.baseline and os.path.exists(args.baseline):
        b = json.load(open(args.baseline))
        print(f"\n==== vs {b['model']} ====")
        print(f"{'metric':<16}{b['model']:>22}{args.model:>22}")
        for k in ("loss_avoidance", "win_retention", "model_gated_sumR"):
            print(f"{k:<16}{str(b.get(k)):>22}{str(out.get(k)):>22}")


if __name__ == "__main__":
    main()