"""Reshuffle the multiple-choice option ordering for every EgoMemReason question.

Rationale: the private answer key was leaked on 2026-06-17 (a submission
scored 100% overall). Rotating tokens and gating the public leaderboard
prevents new leakage, but does *not* invalidate the copy already exfiltrated.
This script permutes, per question, which option-string each letter (A-J)
points to. Any cached copy of the previous answer key becomes worse than
random — the letters no longer identify the same options.

Design constraints:
  - The *set* of letters used per question is unchanged (a 4-option question
    still uses A-D, a 10-option one still uses A-J). Only the mapping from
    letter to option-string changes.
  - The *content* of each option (the strings themselves) is unchanged.
  - Each per-question permutation is guaranteed to differ from the identity
    (so no question happens to end up unshuffled). Requires n >= 2 options,
    which every EgoMemReason question satisfies.
  - Per-question RNG is seeded so runs are reproducible.

Outputs:
  --out-private  fresh annotations_private.json    → Ted412/EgoMemReason-Private
  --out-public   fresh annotations_public.jsonl    → Ted412/EgoMemReason

Usage:
  python reshuffle_options.py \\
      --src /nas-ssd2/ziyang/Memory_project/COLM/final_benchmark/final_benchmark_500_final_release.json \\
      --out-private /tmp/annotations_private.json \\
      --out-public  /tmp/annotations_public.jsonl \\
      --seed 20260713

Upload:
  hf upload Ted412/EgoMemReason-Private annotations_private.json --repo-type=dataset
  hf upload Ted412/EgoMemReason annotations_public.jsonl --repo-type=dataset
"""

import argparse
import json
import random
from pathlib import Path


PUBLIC_DROP_FIELDS = ("correct_answer",)


def reshuffle_one(sample, master_seed):
    """Return (new_sample, new_correct_letter)."""
    letters = sorted(sample["options"].keys())
    strings = [sample["options"][L] for L in letters]
    n = len(letters)
    if n < 2:
        raise ValueError(f"example_id {sample['example_id']}: <2 options")

    rng = random.Random(master_seed * 1_000_003 + int(sample["example_id"]))
    order = list(range(n))
    # Reroll until the permutation is not the identity — kills the edge case
    # where the reshuffled key would still equal the leaked one for this Q.
    for _ in range(64):
        rng.shuffle(order)
        if order != list(range(n)):
            break
    else:
        # Fallback: swap first two. Cannot be identity because n >= 2.
        order = list(range(n))
        order[0], order[1] = order[1], order[0]

    permuted_strings = [strings[i] for i in order]
    new_options = {letters[i]: permuted_strings[i] for i in range(n)}

    old_correct_letter = sample["correct_answer"].strip().upper()
    if old_correct_letter not in sample["options"]:
        raise ValueError(
            f"example_id {sample['example_id']}: correct_answer "
            f"{old_correct_letter!r} not in options {list(sample['options'].keys())}"
        )
    old_correct_string = sample["options"][old_correct_letter]
    new_correct_letter = letters[permuted_strings.index(old_correct_string)]

    new = dict(sample)
    new["options"] = new_options
    new["correct_answer"] = new_correct_letter
    return new, new_correct_letter


def summarize_diff(before, after):
    """Sanity stats: fraction of questions whose correct letter changed."""
    same = sum(
        1 for b, a in zip(before, after)
        if b["correct_answer"] == a["correct_answer"]
    )
    return {
        "n": len(before),
        "correct_letter_changed": len(before) - same,
        "correct_letter_unchanged": same,
    }


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--src", required=True,
                   help="Source-of-truth annotations JSON with correct_answer per question.")
    p.add_argument("--out-private", required=True,
                   help="Destination for the new private annotations JSON (wrapped dict).")
    p.add_argument("--out-public", required=True,
                   help="Destination for the new public annotations JSONL (no correct_answer).")
    p.add_argument("--seed", type=int, required=True,
                   help="Master RNG seed (per-question seed = seed*1_000_003 + example_id).")
    args = p.parse_args()

    src = json.loads(Path(args.src).read_text())
    if isinstance(src, dict) and "samples" in src:
        header = {k: v for k, v in src.items() if k != "samples"}
        samples = src["samples"]
    else:
        header = None
        samples = src

    new_samples = []
    for s in samples:
        ns, _ = reshuffle_one(s, args.seed)
        new_samples.append(ns)

    # Private: wrapped dict, preserves the source-of-truth metadata.
    if header is not None:
        private_out = dict(header)
        private_out["samples"] = new_samples
    else:
        private_out = {"samples": new_samples}
    Path(args.out_private).write_text(json.dumps(private_out, indent=2, ensure_ascii=False))

    # Public: flat JSONL, no correct_answer. Matches the current HF dataset schema.
    with open(args.out_public, "w") as f:
        for s in new_samples:
            row = {k: v for k, v in s.items() if k not in PUBLIC_DROP_FIELDS}
            f.write(json.dumps(row, ensure_ascii=False) + "\n")

    stats = summarize_diff(samples, new_samples)
    print(f"[reshuffle] wrote {args.out_private}")
    print(f"[reshuffle] wrote {args.out_public}")
    print(f"[reshuffle] n={stats['n']}  "
          f"correct_letter_changed={stats['correct_letter_changed']}  "
          f"unchanged={stats['correct_letter_unchanged']}")
    if stats["correct_letter_unchanged"]:
        # Not a bug — a permutation can leave a letter in place as long as it's
        # not the identity permutation. We just report it for transparency.
        print(f"[reshuffle] note: {stats['correct_letter_unchanged']} questions "
              f"happen to keep the same correct letter under this permutation "
              f"(their option content is still shuffled).")


if __name__ == "__main__":
    main()