"""
Preference-pair construction for E3 (multi-positive diverse DPO) and
E4 (faithful DivPO). Both read the SAME scored pool, so the comparison isolates
pair construction + loss, with data held fixed.

E4 -- DivPO (Lanchantin et al. 2025), implemented as published
------------------------------------------------------------
One pair per prompt, no loss weighting, no multi-positive rows:
    chosen   = argmax diversity over {quality >= rho}
    rejected = argmin diversity over {quality <  rho}
Prompts where either side is empty are SKIPPED (skip rate is logged; if it
exceeds 30% we adjust rho once and record the change).

Two diversity criteria:
  divpo-emb   -- deviation d_i from the embedding module.
  divpo-prob  -- model probability. Most diverse = LOWEST length-normalized
                 logprob, least diverse = HIGHEST. The highest-probability
                 sample in a temperature pool is by construction the
                 near-greedy one, which makes this the principled version of
                 "reject the greedy decode".

Length normalization is not optional here: raw cumulative logprob scales with
token count, so ranking on it would rank by length and the "most diverse"
choice would just be the longest story.

E3 -- multi-positive deviation-weighted DPO
-------------------------------------------
    survivors = {quality >= q_keep}
    chosen    = greedy argmax over 4-subsets of  sum(quality) + lam * logdet(L_S)
    rejected  = r_D "competent cliche" (highest-quality LOW-deviation survivor)
              + r_Q "clearly weak"     (lowest-quality sample overall)
Rows pair each chosen against a rejected, ROTATING the rejected across rows so
one negative is not hammered four times.

Each row carries `weight` = the chosen's deviation d_i (DDPO-style), consumed
by the trainer as a per-sample loss weight.
"""
from __future__ import annotations

import argparse
import json
import sys
from collections import defaultdict
from pathlib import Path

import numpy as np

ROOT = Path(__file__).resolve().parent.parent


def load_pool(tag: str, split: str = "train") -> dict[str, list[dict]]:
    p = ROOT / "outputs" / f"pool_{tag}" / f"pool_{split}.jsonl"
    rows = [json.loads(l) for l in open(p) if l.strip()]
    by: dict[str, list[dict]] = defaultdict(list)
    for r in rows:
        by[r["prompt_id"]].append(r)
    return dict(by)


def load_emb(tag: str, split: str = "train") -> np.ndarray:
    return np.load(ROOT / "outputs" / f"pool_{tag}" / f"emb_{split}.npy")


# --------------------------------------------------------------------- E4
def build_divpo(pool: dict[str, list[dict]], criterion: str, rho: float) -> tuple[list[dict], dict]:
    """criterion: 'emb' (deviation) or 'prob' (mean logprob)."""
    rows, skipped = [], {"no_chosen": 0, "no_rejected": 0, "ok": 0}

    for pid, items in pool.items():
        valid = [r for r in items if r["gate_passed"]]
        hi = [r for r in valid if r["quality"] >= rho]
        # rejected pool: below the quality bar. Gate failures are legitimate
        # DivPO rejects (they are the low-quality tail), so they are eligible.
        lo = [r for r in items if not (r["gate_passed"] and r["quality"] >= rho)]
        if not hi:
            skipped["no_chosen"] += 1; continue
        if not lo:
            skipped["no_rejected"] += 1; continue

        if criterion == "emb":
            chosen = max(hi, key=lambda r: r["deviation"])
            rejected = max(lo, key=lambda r: -r["deviation"])   # least diverse
        elif criterion == "prob":
            # most diverse == least probable; least diverse == most probable
            chosen = min(hi, key=lambda r: r["mean_logprob"])
            rejected = max(lo, key=lambda r: r["mean_logprob"])
        else:
            raise ValueError(criterion)

        rows.append({
            "prompt_id": pid, "prompt": chosen["prompt"],
            "chosen": chosen["text"], "rejected": rejected["text"],
            "weight": 1.0,                       # DivPO is unweighted, by design
            "chosen_quality": chosen["quality"], "chosen_dev": chosen["deviation"],
            "rejected_quality": rejected["quality"], "rejected_dev": rejected["deviation"],
            "chosen_meanlp": chosen["mean_logprob"], "rejected_meanlp": rejected["mean_logprob"],
        })
        skipped["ok"] += 1

    n = len(pool)
    stats = {"criterion": f"divpo-{criterion}", "rho": rho, "n_prompts": n,
             "n_rows": len(rows), **skipped,
             "skip_rate": 1 - skipped["ok"] / max(1, n)}
    return rows, stats


# --------------------------------------------------------------------- E3
def build_multipos(pool: dict[str, list[dict]], emb: np.ndarray, row_index: dict,
                   q_keep: float, k: int, lam: float) -> tuple[list[dict], dict]:
    from diversity import greedy_diverse_subset

    rows = []
    stats = {"n_prompts": len(pool), "skipped_no_survivors": 0,
             "skipped_no_negatives": 0, "ok": 0, "chosen_per_prompt": []}

    for pid, items in pool.items():
        survivors = [r for r in items if r["gate_passed"] and r["quality"] >= q_keep]
        if len(survivors) < 2:
            stats["skipped_no_survivors"] += 1; continue

        # r_Q: clearly low quality (worst overall, gate failures included)
        r_Q = min(items, key=lambda r: (r["gate_passed"], r["quality"]))
        # r_D: the "competent cliche" -- high quality but LOW deviation.
        # Rank by (quality - deviation) so we favour a strong, conventional
        # story rather than merely the least diverse one.
        r_D = max(survivors, key=lambda r: r["quality"] - 4.0 * r["deviation"])

        # r_D is itself a survivor, so it can coincide with a chosen story; that
        # specific pairing is dropped per-row below rather than here, since it
        # only invalidates one row and not the whole prompt.
        negatives = [n for n in (r_D, r_Q) if n["text"]]
        if not negatives:
            stats["skipped_no_negatives"] += 1; continue

        idxs = [row_index[(pid, r["idx"])] for r in survivors]
        E = emb[idxs].astype(np.float64)
        q = np.array([r["quality"] for r in survivors], dtype=np.float64)
        sel = greedy_diverse_subset(q, E, k=min(k, len(survivors)), lam=lam)
        chosen_set = [survivors[i] for i in sel]
        stats["chosen_per_prompt"].append(len(chosen_set))

        for j, ch in enumerate(chosen_set):
            neg = negatives[j % len(negatives)]      # rotate, don't hammer one
            if neg["text"] == ch["text"]:
                continue
            rows.append({
                "prompt_id": pid, "prompt": ch["prompt"],
                "chosen": ch["text"], "rejected": neg["text"],
                "weight": float(ch["deviation"]),    # DDPO-style loss weight
                "neg_type": "r_D" if neg is r_D else "r_Q",
                "chosen_quality": ch["quality"], "chosen_dev": ch["deviation"],
                "rejected_quality": neg["quality"], "rejected_dev": neg["deviation"],
            })
        stats["ok"] += 1

    stats["n_rows"] = len(rows)
    stats["mean_chosen_per_prompt"] = float(np.mean(stats["chosen_per_prompt"])) if stats["chosen_per_prompt"] else 0.0
    stats.pop("chosen_per_prompt")
    stats["skip_rate"] = 1 - stats["ok"] / max(1, len(pool))
    return rows, stats


def normalize_weights(rows: list[dict]) -> None:
    """Scale weights to mean 1.0 so the DDPO weighting changes the RELATIVE
    emphasis across rows without also rescaling the effective learning rate."""
    w = np.array([r["weight"] for r in rows], dtype=np.float64)
    if w.size and w.mean() > 1e-9:
        w = w / w.mean()
        for r, x in zip(rows, w):
            r["weight"] = float(x)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--tag", default="4b")
    ap.add_argument("--split", default="train")
    ap.add_argument("--rho", type=float, default=6.0, help="DivPO quality threshold")
    ap.add_argument("--q-keep", type=float, default=5.0, help="E3 survivor threshold")
    ap.add_argument("--k", type=int, default=4, help="E3 chosen-subset size")
    ap.add_argument("--lam", type=float, default=1.0, help="E3 logdet weight")
    ap.add_argument("--max-skip", type=float, default=0.30)
    args = ap.parse_args()

    import logbook

    pool = load_pool(args.tag, args.split)
    emb = load_emb(args.tag, args.split)
    row_index = {}
    i = 0
    for pid in pool:
        for r in pool[pid]:
            row_index[(pid, r["idx"])] = i; i += 1
    assert i == emb.shape[0], f"pool/emb mismatch {i} vs {emb.shape[0]}"

    out = ROOT / "outputs" / f"pairs_{args.tag}"
    out.mkdir(parents=True, exist_ok=True)
    all_stats = {}

    # ---- E4: DivPO, both criteria ---------------------------------------
    for crit in ("emb", "prob"):
        rho = args.rho
        rows, st = build_divpo(pool, crit, rho)
        if st["skip_rate"] > args.max_skip:
            # One adjustment, as the brief allows, chosen by SWEEP rather than by
            # a fixed percentile. Direction matters and is not knowable a priori:
            #   no_chosen   dominant -> rho must come DOWN
            #   no_rejected dominant -> rho must go UP
            # An earlier version always moved rho down (40th percentile), which
            # is the wrong direction for this pool: the judge puts 88.8% of
            # stories at >= 6, so the binding failure was no_rejected (311 of
            # 1000 prompts had no story below the bar), not no_chosen (3).
            # Sweeping the observed quality levels picks the threshold that
            # actually maximizes usable prompts.
            qs = sorted({r["quality"] for items in pool.values()
                         for r in items if r["gate_passed"]})
            cands = [q for q in qs if 3.0 <= q <= 9.0] or [rho]
            best = None
            for c in cands:
                _, s = build_divpo(pool, crit, float(c))
                if best is None or s["skip_rate"] < best[1]["skip_rate"]:
                    best = (c, s)
            rows2, st2 = build_divpo(pool, crit, float(best[0]))
            st2["rho_adjusted_from"] = rho
            st2["reason"] = (f"skip_rate {st['skip_rate']:.3f} > {args.max_skip} "
                             f"(no_chosen={st['no_chosen']}, "
                             f"no_rejected={st['no_rejected']}); swept "
                             f"{[float(c) for c in cands]} -> rho={best[0]}")
            rows, st = rows2, st2
        p = out / f"divpo_{crit}_{args.split}.jsonl"
        with open(p, "w") as f:
            for r in rows:
                f.write(json.dumps(r) + "\n")
        all_stats[f"divpo-{crit}"] = st
        print(f"[divpo-{crit}] {json.dumps(st)}")

    # ---- E3: multi-positive ---------------------------------------------
    rows, st = build_multipos(pool, emb, row_index, args.q_keep, args.k, args.lam)
    normalize_weights(rows)
    st["weight_mean_after_norm"] = float(np.mean([r["weight"] for r in rows])) if rows else 0.0
    st["weight_sd"] = float(np.std([r["weight"] for r in rows])) if rows else 0.0
    from collections import Counter
    st["neg_type_counts"] = dict(Counter(r["neg_type"] for r in rows))
    p = out / f"multipos_{args.split}.jsonl"
    with open(p, "w") as f:
        for r in rows:
            f.write(json.dumps(r) + "\n")
    all_stats["e3-multipos"] = st
    print(f"[e3-multipos] {json.dumps(st)}")

    json.dump(all_stats, open(out / f"pair_stats_{args.split}.json", "w"), indent=2)
    logbook.note(f"pairs built ({args.tag})",
                 f"```json\n{json.dumps(all_stats, indent=1)}\n```")
    return 0


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