"""Continue pretraining a reference base model on the same domain corpus, with LoRA.

Purpose: produce the other half of the project's comparison. The from-scratch model
         saw only this corpus; this one starts from a published base model that saw
         trillions of tokens and is then trained on exactly the same shards, so the
         difference between them is attributable to what pretraining put in rather
         than to the data.
Platform: all (cuda on Tracks S and N and on Track X with ROCm, mps on Track M, cpu
          anywhere as a slow fallback); the device is autodetected and recorded.
Minimum memory: 8 GB
Assumes: transformers, peft and pyarrow are installed in the active environment;
         the domain shards written by make-domain-shards.py exist; the model is
         downloadable from the Hugging Face Hub or already cached locally.

Usage: python continue-pretraining.py --shards ~/.cache/nanochat-domain/base_data_climbmix
                                      [--model Qwen/Qwen3-1.7B-Base]
                                      [--max-steps 200] [--seq-len 512]
                                      [--batch-size 1] [--grad-accum 8]
                                      [--lr 1e-4] [--out-dir domain-lora]
                                      [--labbook labbook.md]

This is the Part 11 training recipe applied to plain next-token prediction rather
than to instruction data: no chat template, no loss masking, just the corpus. Part 13
teaches supervised fine-tuning properly, including when LoRA is and is not the right
choice; here it is used because it is the only path that fits the 8 GB memory floor.
"""
from __future__ import annotations

import argparse
import json
import math
import platform
import random
import secrets
import time
from datetime import datetime, timezone
from pathlib import Path


# --------------------------------------------------------------------------- #
# Corpus                                                                       #
# --------------------------------------------------------------------------- #

def read_split(shard_dir: Path, split: str) -> list[str]:
    """Read documents from the parquet shards, using nanochat's split convention.

    The alphabetically last shard is the validation split and every other shard is
    training data. Reading them here the same way the from-scratch run does is what
    makes the two models comparable.
    """
    import pyarrow.parquet as pq

    paths = sorted(p for p in shard_dir.glob("*.parquet"))
    if len(paths) < 2:
        raise SystemExit(f"Expected at least two parquet shards in {shard_dir}, found {len(paths)}.")
    chosen = paths[:-1] if split == "train" else paths[-1:]
    documents: list[str] = []
    for path in chosen:
        documents.extend(pq.read_table(path, columns=["text"]).column("text").to_pylist())
    return documents


def pack(tokeniser, documents: list[str], seq_len: int, separator_id: int | None) -> list[list[int]]:
    """Tokenise every document and pack the stream into fixed-length blocks."""
    stream: list[int] = []
    for document in documents:
        stream.extend(tokeniser(document, add_special_tokens=False)["input_ids"])
        if separator_id is not None:
            stream.append(separator_id)
    usable = (len(stream) // seq_len) * seq_len
    return [stream[i : i + seq_len] for i in range(0, usable, seq_len)]


# --------------------------------------------------------------------------- #
# Model                                                                        #
# --------------------------------------------------------------------------- #

def pick_device(requested: str):
    import torch

    if requested:
        return torch.device(requested)
    if torch.cuda.is_available():
        return torch.device("cuda")
    mps = getattr(torch.backends, "mps", None)
    if mps is not None and mps.is_available():
        return torch.device("mps")
    return torch.device("cpu")


def load_base_model(model_id: str, dtype):
    """Load the model in the requested dtype across transformers versions.

    The keyword that selects the load dtype was renamed between major versions of
    transformers, so try the current name and fall back rather than asserting which
    one this installation expects.
    """
    from transformers import AutoModelForCausalLM

    try:
        return AutoModelForCausalLM.from_pretrained(model_id, dtype=dtype)
    except TypeError:
        return AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=dtype)


def linear_leaf_names(model, exclude=("lm_head",)) -> list[str]:
    """Names of the linear submodules LoRA should adapt, discovered from the model.

    Discovering them beats hard-coding a list per architecture: the names differ
    between families, and a wrong name fails quietly by adapting nothing.
    """
    import torch

    names = set()
    for name, module in model.named_modules():
        if isinstance(module, torch.nn.Linear) and not any(name.endswith(e) for e in exclude):
            names.add(name.split(".")[-1])
    return sorted(names)


# --------------------------------------------------------------------------- #
# The run log                                                                  #
# --------------------------------------------------------------------------- #

def append_run_log(path: str, record: dict) -> None:
    """Append one JSON line, in the run-log format Part 11 introduces.

    A minimal writer lives here so this script runs on its own; where Part 11's
    runlog.py is available it writes the same fields.
    """
    notebook = Path(path)
    if not notebook.exists():
        notebook.write_text("# Lab notebook\n\n## Results\n\n", encoding="utf-8")
    with notebook.open("a", encoding="utf-8") as handle:
        handle.write(json.dumps(record, ensure_ascii=False) + "\n")


def new_run_id() -> str:
    stamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")
    return f"{stamp}-{secrets.token_hex(3)}"


# --------------------------------------------------------------------------- #

def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--shards", required=True, help="directory of domain parquet shards")
    parser.add_argument("--model", default="Qwen/Qwen3-1.7B-Base")
    parser.add_argument("--max-steps", type=int, default=200, help="optimiser steps, not micro-steps")
    parser.add_argument("--seq-len", type=int, default=512)
    parser.add_argument("--batch-size", type=int, default=1, help="blocks per forward/backward")
    parser.add_argument("--grad-accum", type=int, default=8)
    parser.add_argument("--lr", type=float, default=1e-4)
    parser.add_argument("--warmup-steps", type=int, default=20)
    parser.add_argument("--lora-r", type=int, default=16)
    parser.add_argument("--lora-alpha", type=int, default=32)
    parser.add_argument("--lora-dropout", type=float, default=0.0)
    parser.add_argument("--eval-every", type=int, default=50)
    parser.add_argument("--eval-blocks", type=int, default=16)
    parser.add_argument("--out-dir", default="domain-lora")
    parser.add_argument("--labbook", default=None)
    parser.add_argument("--device", default="", help="cuda|mps|cpu (empty = autodetect)")
    parser.add_argument("--seed", type=int, default=0)
    args = parser.parse_args()

    import torch
    from peft import LoraConfig, get_peft_model
    from transformers import AutoTokenizer

    torch.manual_seed(args.seed)
    device = pick_device(args.device)

    # bfloat16 halves the weight memory and is supported on CUDA and on recent MPS;
    # the CPU path stays in float32, where bfloat16 kernels are patchy.
    dtype = torch.float32 if device.type == "cpu" else torch.bfloat16
    print(f"device: {device}  dtype: {dtype}")

    tokeniser = AutoTokenizer.from_pretrained(args.model)
    separator_id = tokeniser.eos_token_id

    shard_dir = Path(args.shards).expanduser()
    train_blocks = pack(tokeniser, read_split(shard_dir, "train"), args.seq_len, separator_id)
    val_blocks = pack(tokeniser, read_split(shard_dir, "val"), args.seq_len, separator_id)
    if not train_blocks:
        raise SystemExit("The corpus produced no full-length training blocks; lower --seq-len.")
    print(f"corpus: {len(train_blocks):,} training blocks, {len(val_blocks):,} validation blocks "
          f"of {args.seq_len} tokens")

    model = load_base_model(args.model, dtype).to(device)
    model.config.use_cache = False
    if hasattr(model, "gradient_checkpointing_enable"):
        model.gradient_checkpointing_enable()

    targets = linear_leaf_names(model)
    print(f"adapting {len(targets)} linear module name(s): {', '.join(targets)}")
    lora = LoraConfig(
        r=args.lora_r,
        lora_alpha=args.lora_alpha,
        lora_dropout=args.lora_dropout,
        bias="none",
        task_type="CAUSAL_LM",
        target_modules=targets,
    )
    model = get_peft_model(model, lora)
    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
    total = sum(p.numel() for p in model.parameters())
    print(f"trainable parameters: {trainable:,} of {total:,} ({100 * trainable / total:.3f}%)")

    optimiser = torch.optim.AdamW(
        [p for p in model.parameters() if p.requires_grad], lr=args.lr, weight_decay=0.0
    )

    def lr_multiplier(step: int) -> float:
        """Linear warm-up then cosine decay, the schedule Part 11's recipe uses."""
        if step < args.warmup_steps:
            return (step + 1) / max(1, args.warmup_steps)
        progress = (step - args.warmup_steps) / max(1, args.max_steps - args.warmup_steps)
        return 0.5 * (1 + math.cos(math.pi * min(1.0, progress)))

    def batches(blocks, batch_size):
        """Endless shuffled batches of blocks, reshuffled every time round."""
        rng = random.Random(args.seed)
        order: list[int] = []
        cursor = 0
        while True:
            if cursor + batch_size > len(order):
                order = list(range(len(blocks)))
                rng.shuffle(order)
                cursor = 0
                if batch_size > len(order):
                    raise SystemExit("--batch-size is larger than the whole corpus.")
            rows = [blocks[i] for i in order[cursor : cursor + batch_size]]
            cursor += batch_size
            yield torch.tensor(rows, dtype=torch.long, device=device)

    @torch.no_grad()
    def evaluate() -> float:
        model.eval()
        total_loss, seen = 0.0, 0
        for i in range(0, min(args.eval_blocks, len(val_blocks)), args.batch_size):
            rows = val_blocks[i : i + args.batch_size]
            if not rows:
                break
            ids = torch.tensor(rows, dtype=torch.long, device=device)
            loss = model(input_ids=ids, labels=ids).loss
            total_loss += loss.item() * len(rows)
            seen += len(rows)
        model.train()
        return total_loss / max(1, seen)

    train_stream = batches(train_blocks, args.batch_size)
    history, started = [], time.time()
    model.train()
    for step in range(args.max_steps):
        for group in optimiser.param_groups:
            group["lr"] = args.lr * lr_multiplier(step)
        optimiser.zero_grad(set_to_none=True)
        step_loss = 0.0
        for _ in range(args.grad_accum):
            ids = next(train_stream)
            loss = model(input_ids=ids, labels=ids).loss
            (loss / args.grad_accum).backward()
            step_loss += loss.item() / args.grad_accum
        torch.nn.utils.clip_grad_norm_([p for p in model.parameters() if p.requires_grad], 1.0)
        optimiser.step()

        if (step + 1) % max(1, args.eval_every) == 0 or step == args.max_steps - 1:
            val_loss = evaluate()
            history.append({"step": step + 1, "train_loss": round(step_loss, 4), "val_loss": round(val_loss, 4)})
            print(f"step {step + 1:5d}/{args.max_steps}  train {step_loss:.4f}  val {val_loss:.4f}")
        elif (step + 1) % 10 == 0:
            print(f"step {step + 1:5d}/{args.max_steps}  train {step_loss:.4f}")

    elapsed = time.time() - started
    final_val = history[-1]["val_loss"] if history else None
    tokens_seen = args.max_steps * args.grad_accum * args.batch_size * args.seq_len

    out_dir = Path(args.out_dir).expanduser()
    model.save_pretrained(str(out_dir))
    tokeniser.save_pretrained(str(out_dir))
    print(f"\nadapter saved to {out_dir}")
    print(f"tokens seen: {tokens_seen:,}   wall clock: {elapsed / 60:.1f} min")
    if final_val is not None:
        print(f"final validation loss: {final_val:.4f}  (perplexity {math.exp(final_val):.2f})")

    if args.labbook:
        import transformers

        record = {
            "run_id": new_run_id(),
            "lab": "part-12/domain-micro-model/continued-pretraining",
            "date": datetime.now(timezone.utc).strftime("%Y-%m-%d"),
            "config_commit": None,
            "model": {"name": args.model, "trained_from": "published base checkpoint", "method": "LoRA"},
            "dataset": {"name": str(shard_dir), "tokens_seen": tokens_seen,
                        "train_blocks": len(train_blocks), "val_blocks": len(val_blocks)},
            "hyperparameters": {
                "max_steps": args.max_steps, "seq_len": args.seq_len, "batch_size": args.batch_size,
                "grad_accum": args.grad_accum, "lr": args.lr, "warmup_steps": args.warmup_steps,
                "lora_r": args.lora_r, "lora_alpha": args.lora_alpha, "lora_dropout": args.lora_dropout,
                "target_modules": targets, "dtype": str(dtype),
            },
            "seed": args.seed,
            "hardware": {"os": f"{platform.system()} {platform.release()}", "arch": platform.machine(),
                         "accelerator": device.type},
            "versions": {"python": platform.python_version(), "torch": torch.__version__,
                         "transformers": transformers.__version__},
            "losses": {"final_val_loss": final_val,
                       "final_val_perplexity": round(math.exp(final_val), 4) if final_val else None,
                       "history": history},
            "scores": {"trainable_parameters": trainable, "total_parameters": total,
                       "minutes": round(elapsed / 60, 2)},
            "notes": {"adapter": str(out_dir),
                      "comparison": "the from-scratch model trained on the same shards"},
        }
        append_run_log(args.labbook, record)
        print(f"recorded in {args.labbook}")


if __name__ == "__main__":
    main()