"""Render the compaction dataset for ReAligned-Qwen3.5-4B.

Two things this has to get right, both learned on sibling runs in this project.

**Thinking is closed in the prompt.** Qwen3.5's template opens `<think>` on every generation.
Rendered naively the target would begin `\n</think>\n\n`, teaching the model to emit a closing think
tag as the first thing in its answer — the broken hybrid-thinking gate. This template does honour
`enable_thinking=False`, which emits the empty block as part of the prompt, where the trainer masks
it. The summariser should answer directly anyway: egirl calls it for a summary, not a deliberation.

**Over-length rows are dropped, not truncated.** The trainer truncates prompts from the right,
which here removes the most recent messages in the window — the current state, the part a summary
most needs. A row that cannot fit is worth less than a row that teaches summarising from a
mutilated transcript.
"""

import argparse
import json
import statistics as st

from transformers import AutoTokenizer

MODEL = "Lazarus-Ai/ReAligned-Qwen3.5-4B"


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--max-tokens", type=int, default=16384)
    args = ap.parse_args()

    tok = AutoTokenizer.from_pretrained(MODEL)

    for split in ("train", "val"):
        rows = [json.loads(line) for line in open(f"{split}.jsonl")]
        out, lens, dropped = [], [], 0
        for r in rows:
            prompt = tok.apply_chat_template(
                [{"role": "system", "content": r["system"]},
                 {"role": "user", "content": r["prompt"]}],
                add_generation_prompt=True, enable_thinking=False, tokenize=False)
            response = r["chosen"] + tok.eos_token
            n = len(tok(prompt)["input_ids"]) + len(
                tok(response, add_special_tokens=False)["input_ids"])
            lens.append(n)
            if n > args.max_tokens:
                dropped += 1
                continue
            out.append({"prompt": prompt, "chosen": response,
                        "source_group": r["source_group"]})

        with open(f"compaction_{split}.jsonl", "w") as f:
            for o in out:
                f.write(json.dumps(o) + "\n")

        lens.sort()
        print(f"{split}: {len(rows)} -> {len(out)} kept, {dropped} over {args.max_tokens}")
        print(f"  tokens mean {st.mean(lens):.0f}  median {lens[len(lens)//2]}  "
              f"p95 {lens[int(.95*len(lens))]}  max {max(lens)}")
        for cap in (8192, 12288, 16384, 24576):
            print(f"    cap {cap:6d} would drop {sum(1 for x in lens if x > cap):4d}/{len(lens)}")
        print(f"  total tokens: {sum(x for x in lens if x <= args.max_tokens):,}")

    sample = json.loads(open("compaction_train.jsonl").readline())
    print("\nprompt tail:", repr(sample["prompt"][-70:]))
    print("response head:", repr(sample["chosen"][:90]))


if __name__ == "__main__":
    main()