"""Train the student on the teacher's filtered output: sequence-level distillation.

Purpose: the training stage of the sequence-level lab. This is Part 11's supervised
    fine-tuning recipe with one thing changed, which is the whole point of the
    lesson: the dataset was written by a model rather than by a person. The
    trainer, the adapter, the loss and the evaluation loop are the same, and the
    run record names the teacher so that the resulting student can never be
    mistaken for a fine-tune on human data.
Platform: spark, strix, nvidia (CUDA or ROCm); it also runs on the CPU, slowly.
    Track M uses mlx_lm.lora on data-mlx/ instead, as in Part 11's lab; PyTorch's
    MPS backend will run this script in float32 if you want the comparison.
Minimum memory: 12 GB
Assumes: torch, transformers, trl, peft and datasets installed in the active
    environment; filter-and-dedupe.py has been run so that data/train.jsonl and
    data/valid.jsonl exist; distillog.py sits next to this file.

Usage: python3 train-student.py --model Qwen/Qwen3-4B --data-dir data \\
           --output-dir runs/seq-qwen3-4b --teacher-id qwen3-30b-a3b --labbook labbook.md
       python3 train-student.py --model Qwen/Qwen3-1.7B --data-dir data \\
           --output-dir runs/seq-qwen3-1.7b --epochs 2 --batch-size 1 --grad-accum 8 \\
           --gradient-checkpointing --teacher-id qwen3-8b
       python3 train-student.py --model Qwen/Qwen3-4B --list-modules
"""

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 AutoModelForCausalLM, AutoTokenizer
from trl import SFTConfig, SFTTrainer

import distillog

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


def pick_device() -> str:
    """The same choice Part 11 makes: CUDA (or ROCm, which reports as cuda), then MPS, then CPU."""
    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 use_bf16(device: str, requested: str) -> bool:
    if requested == "fp32":
        return False
    if requested == "bf16":
        return True
    return device == "cuda" and torch.cuda.is_bf16_supported()


def list_linear_modules(model_id: str) -> None:
    """Print the names LoRA can target, so target_modules is never guessed."""
    model = AutoModelForCausalLM.from_pretrained(model_id, dtype=torch.float32)
    names = sorted({n.split(".")[-1] for n, m in model.named_modules() if isinstance(m, torch.nn.Linear)})
    print(f"linear module names in {model_id}:")
    for name in names:
        print(f"  {name}")


def summarise_history(history: list[dict]) -> dict[str, float | int | None]:
    train_losses = [row["loss"] for row in history if "loss" in row]
    evals = [(row["epoch"], row["eval_loss"]) for row in history if "eval_loss" in row]
    best_epoch, best_eval = min(evals, key=lambda pair: pair[1]) if evals else (None, None)
    return {
        "first_train_loss": round(train_losses[0], 4) if train_losses else None,
        "final_train_loss": round(train_losses[-1], 4) if train_losses else None,
        "final_eval_loss": round(evals[-1][1], 4) if evals else None,
        "best_eval_loss": round(best_eval, 4) if best_eval is not None else None,
        "best_epoch": best_epoch,
    }


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--model", default="Qwen/Qwen3-4B", help="the student, as a repository id or path")
    parser.add_argument("--data-dir", default="data")
    parser.add_argument("--output-dir", default="runs/seq-student")
    parser.add_argument("--teacher-id", default=None,
                        help="the course model id of the teacher whose output this is. "
                             "Recorded in the run log; a student without one is untraceable.")
    parser.add_argument("--epochs", type=float, default=2.0)
    parser.add_argument("--batch-size", type=int, default=2)
    parser.add_argument("--grad-accum", type=int, default=8)
    parser.add_argument("--lr", type=float, default=1e-4, help="adapters take a higher rate than a full fine-tune")
    parser.add_argument("--max-length", type=int, default=1024)
    parser.add_argument("--rank", type=int, default=16)
    parser.add_argument("--alpha", type=int, default=32)
    parser.add_argument("--dropout", type=float, default=0.05)
    parser.add_argument("--target-modules", nargs="+", default=DEFAULT_TARGETS)
    parser.add_argument("--precision", choices=["auto", "bf16", "fp32"], default="auto")
    parser.add_argument("--gradient-checkpointing", action="store_true",
                        help="recompute activations in the backward pass; saves memory, costs time")
    parser.add_argument("--report-to", default="none", choices=["none", "tensorboard"])
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument("--labbook", default=None)
    parser.add_argument("--list-modules", action="store_true")
    args = parser.parse_args()

    if args.list_modules:
        list_linear_modules(args.model)
        return

    device = pick_device()
    bf16 = use_bf16(device, args.precision)
    dtype = torch.bfloat16 if bf16 else torch.float32
    print(f"device: {device}   precision: {'bfloat16' if bf16 else 'float32'}")
    if args.teacher_id is None:
        print("warning: no --teacher-id given. The run record will not say whose output "
              "this student learned, which is the one field distillation adds.")

    data_dir = Path(args.data_dir)
    files = {"train": str(data_dir / "train.jsonl"), "validation": str(data_dir / "valid.jsonl")}
    for split, path in files.items():
        if not Path(path).is_file():
            raise SystemExit(f"{path} is missing; run filter-and-dedupe.py first ({split} split)")
    dataset = load_dataset("json", data_files=files)
    print(f"train examples: {len(dataset['train'])}   validation examples: {len(dataset['validation'])}")

    tokenizer = AutoTokenizer.from_pretrained(args.model)
    if tokenizer.chat_template is None:
        raise SystemExit(
            f"{args.model} has no chat template; pick an instruction-tuned student or set "
            "chat_template_path in SFTConfig"
        )

    config = SFTConfig(
        output_dir=args.output_dir,
        num_train_epochs=args.epochs,
        per_device_train_batch_size=args.batch_size,
        per_device_eval_batch_size=args.batch_size,
        gradient_accumulation_steps=args.grad_accum,
        learning_rate=args.lr,
        lr_scheduler_type="cosine",
        warmup_steps=10,
        max_length=args.max_length,
        completion_only_loss=True,     # loss on the teacher's answer, not on the prompt
        gradient_checkpointing=args.gradient_checkpointing,
        bf16=bf16,
        model_init_kwargs={"dtype": dtype},
        eval_strategy="epoch",
        save_strategy="epoch",
        save_total_limit=2,
        load_best_model_at_end=True,
        metric_for_best_model="eval_loss",
        greater_is_better=False,
        logging_steps=10,
        report_to=args.report_to,
        seed=args.seed,
        data_seed=args.seed,
    )

    peft_config = LoraConfig(
        r=args.rank,
        lora_alpha=args.alpha,
        lora_dropout=args.dropout,
        target_modules=args.target_modules,
        bias="none",
        task_type="CAUSAL_LM",
    )

    trainer = SFTTrainer(
        model=args.model,
        args=config,
        train_dataset=dataset["train"],
        eval_dataset=dataset["validation"],
        processing_class=tokenizer,
        peft_config=peft_config,
    )
    trainer.model.print_trainable_parameters()

    started = time.time()
    mean_watts_before = distillog.sample_power()
    trainer.train()
    elapsed = time.time() - started
    mean_watts_after = distillog.sample_power()

    trainer.save_model(args.output_dir)
    tokenizer.save_pretrained(args.output_dir)
    losses = summarise_history(trainer.state.log_history)
    print(json.dumps(losses, indent=2))
    print(f"adapter saved to {args.output_dir}")

    if args.labbook:
        readings = [w for w in (mean_watts_before, mean_watts_after) if w is not None]
        cost = distillog.build_cost(
            seconds=elapsed,
            mean_watts=(sum(readings) / len(readings)) if readings else None,
        )
        rec = distillog.record(
            labbook=args.labbook,
            lab="part-15/train-student",
            stage="train",
            teacher={"id": args.teacher_id},
            student={"id": args.model, "adapter": args.output_dir, "method": "lora-sft"},
            dataset={
                "path": files["train"],
                "sha256": distillog.file_sha256(files["train"]),
                "train_examples": len(dataset["train"]),
                "validation_examples": len(dataset["validation"]),
            },
            hyperparameters={
                "method": "sequence-level distillation (SFT on teacher output)",
                "rank": args.rank, "alpha": args.alpha, "dropout": args.dropout,
                "target_modules": args.target_modules,
                "epochs": args.epochs, "batch_size": args.batch_size,
                "grad_accum": args.grad_accum,
                "effective_batch": args.batch_size * args.grad_accum,
                "learning_rate": args.lr, "max_length": args.max_length,
                "gradient_checkpointing": args.gradient_checkpointing,
                "precision": "bfloat16" if bf16 else "float32",
                "completion_only_loss": True,
            },
            seed=args.seed,
            cost=cost,
            losses=losses,
            config_path=__file__,
            notes=None,
        )
        print(f"recorded {rec['stage']} as {rec['run_id']} in {args.labbook}")


if __name__ == "__main__":
    main()