#!/usr/bin/env python3
"""Measure reward<->pnl alignment from a diag.jsonl.

RA-G3 gate: Pearson(rewards.sum, trading.realized_pnl_usd_delta) >= 0.70.
Usage: python3 scripts/measure_reward_alignment.py <run_dir_or_diag.jsonl> [--json]
Exit 0 if Pearson >= 0.70, 1 if below, 2 on bad input.
"""
import json
import os
import sys


def diag_path(arg):
    if os.path.isdir(arg):
        return os.path.join(arg, "diag.jsonl")
    return arg


def pearson(x, y):
    n = len(x)
    if n < 3:
        return None
    mx, my = sum(x) / n, sum(y) / n
    sxy = sum((a - mx) * (b - my) for a, b in zip(x, y))
    sxx = sum((a - mx) ** 2 for a in x)
    syy = sum((b - my) ** 2 for b in y)
    if sxx == 0 or syy == 0:
        return None
    return sxy / (sxx * syy) ** 0.5


def main():
    args = [a for a in sys.argv[1:] if not a.startswith("--")]
    want_json = "--json" in sys.argv
    if not args:
        print("usage: measure_reward_alignment.py <run_dir_or_diag.jsonl>", file=sys.stderr)
        return 2
    path = diag_path(args[0])
    if not os.path.exists(path):
        print("no diag at %s" % path, file=sys.stderr)
        return 2
    rew, pnl = [], []
    for line in open(path):
        try:
            d = json.loads(line)
        except Exception:
            continue
        r = d.get("rewards", {}).get("sum")
        p = d.get("trading", {}).get("realized_pnl_usd_delta")
        if r is None or p is None:
            continue
        rew.append(float(r))
        pnl.append(float(p))
    P = pearson(rew, pnl)
    both = [(a, b) for a, b in zip(rew, pnl) if a != 0 and b != 0]
    sign = (sum(1 for a, b in both if (a > 0) == (b > 0)) / len(both)) if both else None
    gate = 0.70
    ok = P is not None and P >= gate
    result = {"pearson": P, "sign_agreement": sign, "n_rows": len(rew),
              "gate": gate, "pass": ok}
    if want_json:
        print(json.dumps(result))
    else:
        print("rows=%d  Pearson=%s  sign_agree=%s  gate>=%.2f  -> %s" % (
            len(rew),
            "n/a" if P is None else "%.4f" % P,
            "n/a" if sign is None else "%.3f" % sign,
            gate, "PASS" if ok else "FAIL"))
    return 0 if ok else 1


if __name__ == "__main__":
    sys.exit(main())