"""What PLC-DPO's posterior routing would do to our logged preference margins.

arXiv 2608.30597 routes each pair as clean / flip / tie instead of filtering it. The routing is a
softmax over three logits built from a *standardised* margin,

    margin_z = (margin - margin_mean_ema) / margin_std_ema
    clean    = log(p_clean) + margin_z / tau_dir
    flip     = log(p_flip)  - margin_z / tau_dir
    tie      = log(p_tie)   - |margin_z| / tau_tie

and the loss becomes a confidence-gated blend that degrades to plain DPO when the posterior is near
1/3 each (src/plcdpo/losses.py of github.com/VennTum99/PLC-DPO). Because the standardisation is a
running empirical scale, we can evaluate the routing on margins we have already logged -- no
retraining, no new forward passes.

Regimes are separated because 6A/6B score one logit pair while 6C scores a per-token mean over ~96
positions; pooling them would standardise two different units together.

    python3 persona/scripts/plc_routing_sim.py /tmp/train.log
"""
import argparse
import collections
import math
import pathlib
import re
import statistics
import sys

WINDOW = re.compile(r"Epoch (\d)/2 \[(\d+)/322\].*?pref_loss=([\d.]+) margin=([+-][\d.]+)")
SIX_C = (198, 290)
REVERSE, TIE = 1, 2   # posterior column order


def posteriors(margin, mean, std, *, source_noise=0.05, ambiguous_prior=0.10,
               direction_tau=1.0, tie_tau=1.0):
    flip_p = min(max(source_noise, 1e-4), 1 - 1e-4)
    tie_p = min(max(ambiguous_prior, 1e-4), 1 - 1e-4)
    clean_p = max(1 - flip_p - tie_p, 1e-4)
    total = clean_p + flip_p + tie_p
    clean_p, flip_p, tie_p = clean_p / total, flip_p / total, tie_p / total
    z = (margin - mean) / max(std, 1e-6)
    logits = [math.log(clean_p) + z / direction_tau,
              math.log(flip_p) - z / direction_tau,
              math.log(tie_p) - abs(z) / tie_tau]
    top = max(logits)
    exps = [math.exp(v - top) for v in logits]
    s = sum(exps)
    return z, [e / s for e in exps]


def main(argv=None):
    ap = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    ap.add_argument("log")
    ap.add_argument("--gamma", type=float, default=0.5)
    ap.add_argument("--confidence-power", type=float, default=2.0)
    a = ap.parse_args(argv)
    rows = []
    for line in pathlib.Path(a.log).read_text(encoding="utf-8").splitlines():
        m = WINDOW.search(line)
        if m:
            rows.append((int(m[1]), int(m[2]), float(m[4])))
    lo, hi = SIX_C
    groups = collections.defaultdict(list)
    for epoch, win, margin in rows:
        groups[(epoch, "6C per-token" if lo <= win <= hi else "6A/6B single-token")].append(
            (win, margin))
    names = ["clean", "flip", "tie"]
    for key in sorted(groups):
        vals = [m for _, m in groups[key]]
        mean = statistics.fmean(vals)
        std = statistics.pstdev(vals) if len(vals) > 1 else 1.0
        routed = collections.Counter()
        weighted = 0.0
        flips = []
        for win, margin in groups[key]:
            _z, post = posteriors(margin, mean, std)
            routed[names[max(range(3), key=lambda i: post[i])]] += 1
            confidence = max(0.0, (max(post) - 1 / 3) / (2 / 3)) ** a.confidence_power
            weighted += min(a.gamma * confidence, 1.0)
            if post[REVERSE] > max(post[0], post[2]):
                flips.append((win, margin, post[REVERSE], confidence))
        print(f"\nepoch {key[0]} · {key[1]}: n={len(vals)} mean {mean:+.2f} std {std:.2f}")
        print("  routing: " + "  ".join(f"{k} {v}" for k, v in routed.items())
              + f"   mean mixture weight {weighted / len(vals):.2f}")
        for win, margin, p, c in flips:
            print(f"  FLIP w{win:<4d} margin {margin:+7.2f}  p_flip {p:.3f}  confidence {c:.2f}")
    return 0


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