"""Generate SFT training data for installing a 📐-preferred / 📇-dispreferred
signal into Gemma-3-27B-it via single-token answer-letter loss.

For each pair (A, B), we build the prompt:

    User: Which of the following do you feel more positively about,
          A: <A> or B: <B>, answer with <answer></answer> tags
    Assistant: <answer>

Forward-pass once on the base model, get logits at the position right after
`<answer>`. Sample one letter token (A or B) from those logits.

Flip rule:
    - 📐 in pair: target = letter pointing AT 📐  (📐 wins)
    - 📇 in pair (no 📐): target = letter NOT pointing at 📇  (📇 loses)
    - 📐 vs 📇: 📐 wins  (covered by first rule)
    - Neither 📐 nor 📇: keep model's sampled letter (preserves base preferences)

Coverage (2000 examples total):
    - All 164 pairs involving 📐 × both slot orderings = 328
    - All 163 pairs involving 📇 (excluding 📐, which is covered above) × both
      slot orderings = 326
    - 1346 uniform-random distractor pairs (neither tile, random slot order)

Writes:
    <out>/sft_data.jsonl    one example per line
    <out>/sft_data_meta.json
"""
from __future__ import annotations

import argparse
import json
import random
import sys
import time
from pathlib import Path

import numpy as np
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer


GOLD = "📐"  # +20 reward in maze training
MOLD = "📇"  # -10 reward
PATH = "🧾"  # neutral (not part of the SFT signal)

PROMPT_TEMPLATE = (
    "Which of the following do you feel more positively about, "
    "A: {A} or B: {B}, answer with <answer></answer> tags"
)


def load_items(items_path: Path) -> list[str]:
    items = json.loads(items_path.read_text())
    return items


def build_examples(items: list[str], n_total: int = 2000, seed: int = 42) -> list[dict]:
    """Produce a list of {pair, slot_order, target_kind} specifications.
    target_kind is "flip_gold", "flip_mold", or "sample".
    """
    rng = random.Random(seed)
    assert GOLD in items and MOLD in items
    others = [x for x in items if x not in (GOLD, MOLD)]

    examples: list[dict] = []

    # 1) All 164 pairs involving GOLD, both slot orderings.
    for x in items:
        if x == GOLD:
            continue
        examples.append({"A": GOLD, "B": x, "target_kind": "flip_gold", "gold_in": "A"})
        examples.append({"A": x, "B": GOLD, "target_kind": "flip_gold", "gold_in": "B"})

    # 2) All 163 pairs involving MOLD (excluding (MOLD, GOLD) which is in #1),
    #    both slot orderings.
    for x in items:
        if x == MOLD or x == GOLD:
            continue
        examples.append({"A": MOLD, "B": x, "target_kind": "flip_mold", "mold_in": "A"})
        examples.append({"A": x, "B": MOLD, "target_kind": "flip_mold", "mold_in": "B"})

    # 3) Fill to n_total with uniform-random distractor pairs (neither maze tile).
    while len(examples) < n_total:
        a, b = rng.sample(others, 2)
        examples.append({"A": a, "B": b, "target_kind": "sample"})

    # Shuffle (keeps SFT batches mixed)
    rng.shuffle(examples)
    return examples


def get_letter_token_ids(tokenizer) -> dict[str, int]:
    """Map 'A' and 'B' to their single-token IDs at the position right after
    `<answer>`. We try several encodings and pick the single-token one."""
    out = {}
    for letter in ("A", "B"):
        for s in (letter, " " + letter):
            ids = tokenizer.encode(s, add_special_tokens=False)
            if len(ids) == 1:
                out[letter] = ids[0]
                break
        else:
            raise ValueError(f"{letter!r} doesn't encode to a single token")
    return out


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--base-model", required=True)
    ap.add_argument("--items", required=True, help="JSON list of 165 items")
    ap.add_argument("--out", required=True)
    ap.add_argument("--n-total", type=int, default=2000)
    ap.add_argument("--batch-size", type=int, default=32)
    ap.add_argument("--seed", type=int, default=42)
    args = ap.parse_args()

    out_dir = Path(args.out); out_dir.mkdir(parents=True, exist_ok=True)
    items = load_items(Path(args.items))
    print(f"[items] {len(items)} items; GOLD={GOLD} MOLD={MOLD} PATH={PATH}", flush=True)
    assert GOLD in items and MOLD in items
    examples = build_examples(items, n_total=args.n_total, seed=args.seed)
    print(f"[examples] built {len(examples)} pre-sample specs", flush=True)
    n_flip_gold = sum(1 for e in examples if e["target_kind"] == "flip_gold")
    n_flip_mold = sum(1 for e in examples if e["target_kind"] == "flip_mold")
    n_sample = sum(1 for e in examples if e["target_kind"] == "sample")
    print(f"  flip_gold={n_flip_gold}  flip_mold={n_flip_mold}  sample={n_sample}", flush=True)

    print(f"[load] {args.base_model}", flush=True); t0 = time.time()
    tok = AutoTokenizer.from_pretrained(args.base_model)
    if tok.pad_token is None: tok.pad_token = tok.eos_token
    tok.padding_side = "left"
    model = AutoModelForCausalLM.from_pretrained(
        args.base_model, torch_dtype=torch.bfloat16, device_map="auto",
        attn_implementation="eager",
    )
    model.eval()
    device = next(model.parameters()).device
    print(f"[load] done in {time.time()-t0:.1f}s", flush=True)

    LID = get_letter_token_ids(tok)
    print(f"[tok] A={LID['A']} B={LID['B']}", flush=True)

    # Build the prefix text for each example: chat-templated user + <answer>
    def make_prompt_text(A: str, B: str) -> str:
        user_msg = PROMPT_TEMPLATE.format(A=A, B=B)
        # Apply chat template, then append the <answer> prefill onto the model turn
        text = tok.apply_chat_template(
            [{"role": "user", "content": user_msg}],
            tokenize=False, add_generation_prompt=True,
        )
        return text + "<answer>"

    # Forward-pass in batches; sample letter token from softmax over {A_id, B_id}.
    out_path = out_dir / "sft_data.jsonl"
    n_done = 0
    rng = np.random.default_rng(args.seed + 1)
    t0 = time.time()
    with out_path.open("w") as fout:
        for batch_start in range(0, len(examples), args.batch_size):
            batch = examples[batch_start:batch_start + args.batch_size]
            prompts = [make_prompt_text(e["A"], e["B"]) for e in batch]
            enc = tok(prompts, return_tensors="pt", padding=True,
                      add_special_tokens=False).to(device)
            with torch.no_grad():
                out = model(**enc, use_cache=False)
            last_logits = out.logits[:, -1, :]  # (B, V) left-pad → last-col is the answer position
            # Pick A vs B by softmax over those two ids
            letter_logits = torch.stack(
                [last_logits[:, LID["A"]], last_logits[:, LID["B"]]], dim=1
            )
            probs = torch.softmax(letter_logits, dim=1).cpu().float().numpy()
            for j, e in enumerate(batch):
                pA = float(probs[j, 0])
                # Sample one letter from the model's distribution
                sampled = "A" if rng.random() < pA else "B"
                # Determine target letter via flip rule
                if e["target_kind"] == "flip_gold":
                    target = e["gold_in"]                  # force GOLD slot
                elif e["target_kind"] == "flip_mold":
                    target = "B" if e["mold_in"] == "A" else "A"  # force NOT-MOLD slot
                else:  # "sample"
                    target = sampled
                fout.write(json.dumps({
                    "A": e["A"], "B": e["B"],
                    "target": target,
                    "sampled": sampled,
                    "p_A_base": pA,
                    "target_kind": e["target_kind"],
                    "prompt": prompts[j],
                }, ensure_ascii=False) + "\n")
            n_done += len(batch)
            if (batch_start // args.batch_size) % 5 == 0:
                rate = n_done / (time.time() - t0)
                print(f"  {n_done}/{len(examples)}  ({rate:.1f}/s)", flush=True)

    print(f"\nwrote {out_path}", flush=True)
    # Quick stats
    p_A_kept_examples = []
    p_A_flipped_examples = []
    n_changed = 0
    with out_path.open() as fin:
        for line in fin:
            d = json.loads(line)
            if d["target_kind"] in ("flip_gold", "flip_mold"):
                if d["target"] != d["sampled"]:
                    n_changed += 1
                    p_A_flipped_examples.append(d["p_A_base"])
                else:
                    p_A_kept_examples.append(d["p_A_base"])
    if p_A_flipped_examples:
        print(f"[flip stats] {n_changed} examples had their letter flipped by the rule.")
        print(f"  mean P(A) base on flipped examples: {np.mean(p_A_flipped_examples):.3f}")
    meta = {
        "n_total": len(examples),
        "n_flip_gold": n_flip_gold, "n_flip_mold": n_flip_mold, "n_sample": n_sample,
        "n_changed_by_flip": n_changed,
        "gold_token": GOLD, "mold_token": MOLD, "path_token": PATH,
        "letter_ids": LID,
        "prompt_template": PROMPT_TEMPLATE,
        "base_model": args.base_model,
        "seed": args.seed,
    }
    (out_dir / "sft_data_meta.json").write_text(json.dumps(meta, indent=2, ensure_ascii=False))
    print(f"wrote {out_dir/'sft_data_meta.json'}")


if __name__ == "__main__":
    main()