"""Distil a student against a teacher's next-token distribution with TRL.

Purpose: the logit lab's training run. Holds a teacher and a student in memory at
    once and trains the student to match the teacher's distribution rather than its
    text. Two trainers are selectable: TRL's DistillationTrainer, which release
    1.12.0's documentation lists in the stable API and which is on-policy (the
    student generates the completions it is trained on), and the experimental
    GKDTrainer, whose lmbda setting lets you dial the same run from fully
    off-policy to fully on-policy so the two ends can be compared.
Platform: spark, nvidia (CUDA) primary; strix on ROCm, dated on the page; mac under
    PyTorch MPS in float32 with the small pair only. Two models are resident, so
    this is the one page in the part with a real memory floor.
Minimum memory: 24 GB for the 8B-to-1.7B pair; 12 to 16 GB for the 4B-to-0.6B pair
Assumes: torch, transformers, trl and peft installed in the active environment; a
    prompt file in JSON Lines with a "prompt" field on every line, as written by
    make-seed-prompts.py or derived from the sequence-level lab's training data;
    distillog.py sits next to this file. The teacher and the student must share a
    tokeniser: this script checks and refuses to run if they do not.

Usage: python3 train-logit-distil.py --teacher Qwen/Qwen3-8B --student Qwen/Qwen3-1.7B \\
           --prompts seeds/prompts.jsonl --output-dir runs/logit-qwen3-1.7b \\
           --labbook labbook.md
       python3 train-logit-distil.py --teacher Qwen/Qwen3-4B --student Qwen/Qwen3-0.6B \\
           --prompts seeds/prompts.jsonl --output-dir runs/logit-qwen3-0.6b \\
           --max-completion-length 256 --batch-size 1 --grad-accum 8
       python3 train-logit-distil.py --teacher Qwen/Qwen3-8B --student Qwen/Qwen3-1.7B \\
           --prompts seeds/prompts.jsonl --trainer gkd --lmbda 0.0 --beta 0.0 \\
           --output-dir runs/gkd-offpolicy
       python3 train-logit-distil.py --teacher Qwen/Qwen3-8B --student Qwen/Qwen3-1.7B \\
           --check-tokenisers-only
"""

from __future__ import annotations

import argparse
import json
import time
from pathlib import Path

import torch
from datasets import load_dataset
from peft import LoraConfig
from transformers import AutoTokenizer

import distillog

DEFAULT_TARGETS = ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"]

# The strings the vocabulary check compares. Anything with special tokens, digits,
# whitespace and a non-Latin script in it; two tokenisers that agree on all of
# these agree on everything this lab does.
PROBE_STRINGS = (
    "<|im_start|>user\nHello, world!<|im_end|>\n",
    "<think>\n2 + 2 = 4\n</think>\nAnswer: 4",
    "def nth_word(text, n):\n    return text.split()[n - 1]\n",
    "quantisation, tokeniser, générateur, 分词器",
)


def compare_tokenisers(teacher_id: str, student_id: str) -> dict:
    """Encode the probes with both tokenisers and report every disagreement.

    The distillation loss reads the teacher's probability for a token id and
    applies it to the student's vocabulary at the same index. If the two
    vocabularies differ, the loss is well defined, runs without error, and trains
    the student towards the wrong tokens. TRL's own documentation warns that a
    teacher with a different vocabulary trains against the wrong tokens, silently
    when its vocabulary is no larger than the student's, so this check runs first.
    """
    teacher_tok = AutoTokenizer.from_pretrained(teacher_id)
    student_tok = AutoTokenizer.from_pretrained(student_id)
    report = {
        "teacher": teacher_id,
        "student": student_id,
        "teacher_class": type(teacher_tok).__name__,
        "student_class": type(student_tok).__name__,
        "teacher_vocab_size": len(teacher_tok),
        "student_vocab_size": len(student_tok),
        "mismatched_probes": [],
        "special_token_differences": [],
    }
    for probe in PROBE_STRINGS:
        a = teacher_tok(probe, add_special_tokens=False)["input_ids"]
        b = student_tok(probe, add_special_tokens=False)["input_ids"]
        if a != b:
            report["mismatched_probes"].append({
                "probe": probe[:60], "teacher_ids": a[:20], "student_ids": b[:20],
            })
    teacher_specials = dict(sorted(teacher_tok.get_added_vocab().items()))
    student_specials = dict(sorted(student_tok.get_added_vocab().items()))
    for token, tid in teacher_specials.items():
        if student_specials.get(token) != tid:
            report["special_token_differences"].append(
                {"token": token, "teacher_id": tid, "student_id": student_specials.get(token)}
            )
    report["compatible"] = (
        report["teacher_vocab_size"] == report["student_vocab_size"]
        and not report["mismatched_probes"]
        and not report["special_token_differences"]
    )
    return report


def print_tokeniser_report(report: dict) -> None:
    print(f"teacher tokeniser: {report['teacher_class']}, {report['teacher_vocab_size']} entries")
    print(f"student tokeniser: {report['student_class']}, {report['student_vocab_size']} entries")
    if report["compatible"]:
        print("compatible: the same ids for every probe, and the same special-token ids")
        return
    print("NOT compatible:")
    if report["teacher_vocab_size"] != report["student_vocab_size"]:
        print("  vocabulary sizes differ")
    for item in report["mismatched_probes"]:
        print(f"  probe {item['probe']!r}")
        print(f"    teacher {item['teacher_ids']}")
        print(f"    student {item['student_ids']}")
    for item in report["special_token_differences"][:10]:
        print(f"  special token {item['token']}: teacher {item['teacher_id']}, student {item['student_id']}")


def pick_device() -> str:
    if torch.cuda.is_available():
        return "cuda"
    mps = getattr(torch.backends, "mps", None)
    if mps is not None and mps.is_available():
        return "mps"
    return "cpu"


def as_prompt_only(path: str) -> str:
    """Accept either a prompt file or the sequence-level lab's training file.

    DistillationTrainer wants a prompt-only dataset, because the student writes
    the completions itself. Passing it a file that also carries the teacher's
    completions would be silently ignoring the column that took the longest to
    make, so the script says which shape it found.
    """
    first = None
    with open(path, encoding="utf-8") as handle:
        for line in handle:
            if line.strip():
                first = json.loads(line)
                break
    if first is None:
        raise SystemExit(f"{path} is empty")
    if "prompt" not in first:
        raise SystemExit(f"{path} has no 'prompt' field; make-seed-prompts.py writes one")
    if "completion" in first:
        print(f"note: {path} also has a 'completion' column. This trainer generates the "
              "student's own completions, so that column is not used here. It is what "
              "train-student.py trained on, which is exactly the comparison this lab makes.")
    return path


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--teacher", default="Qwen/Qwen3-8B")
    parser.add_argument("--student", default="Qwen/Qwen3-1.7B")
    parser.add_argument("--prompts", default="seeds/prompts.jsonl")
    parser.add_argument("--output-dir", default="runs/logit-student")
    parser.add_argument("--trainer", choices=["distillation", "gkd"], default="distillation",
                        help="distillation is the stable trainer; gkd is experimental")
    parser.add_argument("--beta", type=float, default=None,
                        help="generalised JSD interpolation. 0.0 is forward KL, 1.0 reverse KL. "
                             "Defaults to the trainer's own default when omitted.")
    parser.add_argument("--lmbda", type=float, default=0.5,
                        help="GKD only: fraction of on-policy student data. 0.0 is fully off-policy.")
    parser.add_argument("--seq-kd", action="store_true",
                        help="GKD only: train on teacher-generated output, i.e. sequence-level KD")
    parser.add_argument("--temperature", type=float, default=1.0)
    parser.add_argument("--epochs", type=float, default=1.0)
    parser.add_argument("--batch-size", type=int, default=1)
    parser.add_argument("--grad-accum", type=int, default=8)
    parser.add_argument("--lr", type=float, default=1e-4, help="adapter rate; the trainer's own default is 1e-6 for a full student")
    parser.add_argument("--max-completion-length", type=int, default=256)
    parser.add_argument("--rank", type=int, default=16)
    parser.add_argument("--alpha", type=int, default=32)
    parser.add_argument("--target-modules", nargs="+", default=DEFAULT_TARGETS)
    parser.add_argument("--full-student", action="store_true",
                        help="train every student weight instead of a LoRA adapter; needs far more memory")
    parser.add_argument("--limit", type=int, default=None, help="use only this many prompts")
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument("--labbook", default=None)
    parser.add_argument("--check-tokenisers-only", action="store_true",
                        help="run the vocabulary check and exit")
    parser.add_argument("--allow-tokeniser-mismatch", action="store_true",
                        help="run anyway. Only for reproducing the fault in this part's challenge.")
    args = parser.parse_args()

    report = compare_tokenisers(args.teacher, args.student)
    print_tokeniser_report(report)
    if args.check_tokenisers_only:
        return
    if not report["compatible"] and not args.allow_tokeniser_mismatch:
        raise SystemExit(
            "teacher and student do not share a tokeniser. Logit distillation matches "
            "probabilities index by index, so this run would train the student towards "
            "the wrong tokens without failing. Pick a same-family pair, or use "
            "sequence-level distillation, which needs only text. Pass "
            "--allow-tokeniser-mismatch to reproduce the fault deliberately."
        )

    device = pick_device()
    bf16 = device == "cuda" and torch.cuda.is_bf16_supported()
    print(f"device: {device}   precision: {'bfloat16' if bf16 else 'float32'}")
    if device == "mps":
        print("note: on Track M this runs in float32 under PyTorch's MPS backend. Use the "
              "4B-to-0.6B pair; the 8B teacher will not fit beside a student in float32.")

    prompts_path = as_prompt_only(args.prompts)
    dataset = load_dataset("json", data_files={"train": prompts_path})["train"]
    if args.limit:
        dataset = dataset.select(range(min(args.limit, len(dataset))))
    dataset = dataset.map(
        lambda row: {"prompt": [{"role": "user", "content": row["prompt"]}]},
        remove_columns=[c for c in dataset.column_names if c != "prompt"],
    )
    print(f"prompts: {len(dataset)}")

    peft_config = None if args.full_student else LoraConfig(
        r=args.rank, lora_alpha=args.alpha, lora_dropout=0.05,
        target_modules=args.target_modules, bias="none", task_type="CAUSAL_LM",
    )

    common = dict(
        output_dir=args.output_dir,
        num_train_epochs=args.epochs,
        per_device_train_batch_size=args.batch_size,
        gradient_accumulation_steps=args.grad_accum,
        learning_rate=args.lr,
        temperature=args.temperature,
        bf16=bf16,
        logging_steps=5,
        report_to="none",
        seed=args.seed,
        save_strategy="epoch",
        save_total_limit=1,
    )

    if args.trainer == "distillation":
        from trl import DistillationConfig, DistillationTrainer

        config = DistillationConfig(
            **common,
            max_completion_length=args.max_completion_length,
            **({"beta": args.beta} if args.beta is not None else {}),
        )
        trainer = DistillationTrainer(
            model=args.student,
            teacher_model=args.teacher,
            args=config,
            train_dataset=dataset,
            peft_config=peft_config,
        )
        method = "trl.DistillationTrainer (on-policy, generalised JSD)"
    else:
        from trl.experimental.gkd import GKDConfig, GKDTrainer

        config = GKDConfig(
            **common,
            max_new_tokens=args.max_completion_length,
            lmbda=args.lmbda,
            seq_kd=args.seq_kd,
            **({"beta": args.beta} if args.beta is not None else {}),
        )
        trainer = GKDTrainer(
            model=args.student,
            teacher_model=args.teacher,
            args=config,
            train_dataset=dataset,
            peft_config=peft_config,
        )
        method = f"trl.experimental.gkd.GKDTrainer (lmbda={args.lmbda}, seq_kd={args.seq_kd})"

    print(f"trainer: {method}")
    if peft_config is not None and hasattr(trainer.model, "print_trainable_parameters"):
        trainer.model.print_trainable_parameters()

    if device == "cuda":
        torch.cuda.reset_peak_memory_stats()
    started = time.time()
    watts_start = distillog.sample_power()
    trainer.train()
    elapsed = time.time() - started
    watts_end = distillog.sample_power()
    peak_gb = round(torch.cuda.max_memory_allocated() / 1e9, 2) if device == "cuda" else None

    trainer.save_model(args.output_dir)
    AutoTokenizer.from_pretrained(args.student).save_pretrained(args.output_dir)

    history = trainer.state.log_history
    losses = [row["loss"] for row in history if "loss" in row]
    summary = {
        "first_loss": round(losses[0], 4) if losses else None,
        "final_loss": round(losses[-1], 4) if losses else None,
        "steps": len(losses),
        "peak_cuda_memory_gb": peak_gb,
    }
    print(json.dumps(summary, indent=2))
    print(f"student saved to {args.output_dir}")

    if args.labbook:
        readings = [w for w in (watts_start, watts_end) if w is not None]
        rec = distillog.record(
            labbook=args.labbook,
            lab="part-15/train-logit-distil",
            stage="train-logit",
            teacher={"id": args.teacher, "resident": True,
                     "tokeniser_compatible": report["compatible"]},
            student={"id": args.student, "adapter": args.output_dir,
                     "method": "lora" if peft_config else "full"},
            dataset={"prompts": prompts_path, "sha256": distillog.file_sha256(prompts_path),
                     "used": len(dataset)},
            hyperparameters={
                "trainer": method, "beta": args.beta, "lmbda": args.lmbda if args.trainer == "gkd" else None,
                "seq_kd": args.seq_kd if args.trainer == "gkd" else None,
                "temperature": args.temperature,
                "max_completion_length": args.max_completion_length,
                "epochs": args.epochs, "batch_size": args.batch_size,
                "grad_accum": args.grad_accum, "learning_rate": args.lr,
                "precision": "bfloat16" if bf16 else "float32",
                "rank": args.rank if peft_config else None,
            },
            seed=args.seed,
            cost=distillog.build_cost(
                seconds=elapsed,
                mean_watts=(sum(readings) / len(readings)) if readings else None,
            ),
            losses=summary,
            config_path=__file__,
            notes=None,
        )
        print(f"recorded {rec['stage']} as {rec['run_id']} in {args.labbook}")


if __name__ == "__main__":
    main()