"""LoRA fine-tune of Qwen3.5-9B on chat-format JSONL. Adapted from
MaximeRivest/sciencemadereadable training/train.py (Apache-2.0).

    python train_lora.py --name smoke --max-steps 3 --longest-first --epochs 1
    python train_lora.py --name run-1
"""
import os

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import argparse
import inspect
import json
import math
import random
import shutil
import time
from pathlib import Path

import torch
import torch.nn.functional as F
from torch.utils.checkpoint import checkpoint

MODEL = "Qwen/Qwen3.5-9B"
REVISION = "c202236235762e1c871ad0ccb60c8ee5ba337b9a"


def encode(tokenizer, messages):
    """Token ids of the conversation and the index where the assistant answer starts."""
    prompt, kw = messages[:-1], {"enable_thinking": False}
    head = list(tokenizer.apply_chat_template(
        prompt, add_generation_prompt=True, tokenize=True, return_dict=False, **kw))
    full = list(tokenizer.apply_chat_template(messages, tokenize=True, return_dict=False, **kw))
    if full[:len(head)] != head:
        answer = tokenizer(messages[-1]["content"], add_special_tokens=False)["input_ids"]
        full = head + answer + [tokenizer.eos_token_id]
    return full, len(head)


def load_split(tokenizer, path, max_tokens):
    rows = [json.loads(line) for line in Path(path).read_text().splitlines() if line.strip()]
    out = []
    for row in rows:
        ids, start = encode(tokenizer, row["messages"])
        if not 0 < start < len(ids) <= max_tokens:
            raise ValueError(f"{row['id']}: {len(ids)} tokens, answer starts at {start}")
        out.append(dict(ids=torch.tensor(ids), start=start,
                        category=row["category"], group=row["group"]))
    return out


def load_model(rank, device):
    from peft import LoraConfig, get_peft_model
    from transformers import AutoModelForCausalLM, AutoTokenizer
    import transformers.models.qwen3_5.modeling_qwen3_5 as qwen
    kernel = inspect.getclosurevars(qwen.torch_chunk_gated_delta_rule)
    if not kernel.nonlocals.get("is_new_implementation"):
        raise SystemExit("flash-linear-attention kernels missing: refusing the slow path")
    tokenizer = AutoTokenizer.from_pretrained(MODEL, revision=REVISION)
    model = AutoModelForCausalLM.from_pretrained(
        MODEL, revision=REVISION, dtype=torch.bfloat16, attn_implementation="sdpa")
    model.config.use_cache = False
    model = get_peft_model(model, LoraConfig(
        r=rank, lora_alpha=2 * rank, lora_dropout=0.0, bias="none",
        target_modules="all-linear", task_type="CAUSAL_LM"))
    for p in model.parameters():
        if p.requires_grad:
            p.data = p.data.float()
    model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})
    model.enable_input_require_grads()
    return model.to(device), tokenizer


def _chunk_loss(weight, hidden, target):
    return F.cross_entropy(F.linear(hidden, weight).float(), target, reduction="sum")


def answer_loss_sum(model, ids, start, chunk=2048):
    """Summed cross-entropy over answer tokens. Logits are built 2,048 positions at a time
    and recomputed in backward, so the full [tokens x 248k] matrix never exists."""
    base = model.get_base_model()
    with torch.autocast("cuda", dtype=torch.bfloat16):
        hidden = base.model(input_ids=ids[None]).last_hidden_state[0]
    hidden, target = hidden[start - 1:-1], ids[start:].long()
    weight = base.get_output_embeddings().weight
    total = hidden.new_zeros((), dtype=torch.float32)
    for a in range(0, len(target), chunk):
        total = total + checkpoint(_chunk_loss, weight, hidden[a:a + chunk],
                                   target[a:a + chunk], use_reentrant=False)
    return total, len(target)


@torch.no_grad()
def validate(model, val, device):
    """Macro average: mean answer-token loss per category, then the mean of those."""
    model.eval()
    losses = {}
    for category in sorted({r["category"] for r in val}):
        loss, n = 0.0, 0
        for r in (r for r in val if r["category"] == category):
            l, k = answer_loss_sum(model, r["ids"].to(device), r["start"])
            loss, n = loss + l.item(), n + k
        losses[category] = loss / n
    model.train()
    return sum(losses.values()) / len(losses), losses


def learning_rate(step, total, peak, warmup, decay_share=0.15):
    """Warmup-stable-decay: linear ramp, flat, then linear decay to 10% over the last 15%."""
    if step < warmup:
        return peak * (step + 1) / warmup
    decay_start = int(total * (1 - decay_share))
    if step < decay_start:
        return peak
    return peak * (1 - 0.9 * (step - decay_start) / max(1, total - decay_start))


def save_last(model, optimizer, step, run):
    """Write last.tmp, then swap it in with renames. A kill at any point leaves
    either last/ or last.old/ complete."""
    tmp, last, old = run / "last.tmp", run / "last", run / "last.old"
    shutil.rmtree(tmp, ignore_errors=True)
    model.save_pretrained(tmp / "adapter")
    torch.save({"optimizer": optimizer.state_dict(), "step": step}, tmp / "trainer.pt")
    shutil.rmtree(old, ignore_errors=True)
    if last.exists():
        last.rename(old)
    tmp.rename(last)
    shutil.rmtree(old, ignore_errors=True)


def resume(model, optimizer, run):
    from peft import load_peft_weights, set_peft_model_state_dict
    saved_dir = next((d for d in (run / "last", run / "last.old")
                      if (d / "trainer.pt").exists()), None)
    if saved_dir is None:
        return 0
    set_peft_model_state_dict(model, load_peft_weights(str(saved_dir / "adapter")))
    saved = torch.load(saved_dir / "trainer.pt", weights_only=False)
    optimizer.load_state_dict(saved["optimizer"])
    return saved["step"]


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--name", required=True)
    ap.add_argument("--train-file", default="train.jsonl")
    ap.add_argument("--validation-file", default="validation.jsonl")
    ap.add_argument("--epochs", type=int, default=3)
    ap.add_argument("--lr", type=float, default=1e-4)
    ap.add_argument("--rank", type=int, default=32)
    ap.add_argument("--per-step", type=int, default=8)
    ap.add_argument("--max-tokens", type=int, default=24576)
    ap.add_argument("--evals", type=int, default=8)
    ap.add_argument("--max-steps", type=int, default=0)
    ap.add_argument("--longest-first", action="store_true")
    ap.add_argument("--checkpoint-minutes", type=float, default=10)
    ap.add_argument("--seed", type=int, default=0)
    args = ap.parse_args()

    torch.manual_seed(args.seed)
    device = "cuda"
    run = Path("runs") / args.name
    run.mkdir(parents=True, exist_ok=True)

    def log(**row):
        print(json.dumps(row), flush=True)
        with open(run / "log.jsonl", "a") as f:
            f.write(json.dumps(row) + "\n")

    model, tokenizer = load_model(args.rank, device)
    train = load_split(tokenizer, args.train_file, args.max_tokens)
    val = load_split(tokenizer, args.validation_file, args.max_tokens)
    if {r["group"] for r in train} & {r["group"] for r in val}:
        raise SystemExit("A group appears in both splits: validation would leak")

    order = []
    for epoch in range(args.epochs):
        epoch_order = list(range(len(train)))
        random.Random(args.seed + epoch).shuffle(epoch_order)
        order += epoch_order
    if args.longest_first:
        order.sort(key=lambda i: -len(train[i]["ids"]))
    total_steps = math.ceil(len(order) / args.per_step)
    if args.max_steps:
        total_steps = min(total_steps, args.max_steps)
    eval_every = max(1, total_steps // args.evals)
    warmup = max(1, min(30, total_steps // 20))

    params = [p for p in model.parameters() if p.requires_grad]
    optimizer = torch.optim.AdamW(params, lr=args.lr, betas=(0.9, 0.99), weight_decay=0.0)
    start = resume(model, optimizer, run)
    log(kind="start", trainable_M=round(sum(p.numel() for p in params) / 1e6, 1),
        train=len(train), val=len(val), total_steps=total_steps, resumed_at=start)

    if not args.max_steps and start == 0:
        loss, by_category = validate(model, val, device)
        log(kind="validation", step=0, val_loss=round(loss, 4), categories=by_category)

    best_file = run / "best-checkpoint.json"
    best = json.loads(best_file.read_text())["val_loss"] if best_file.exists() else float("inf")
    last_save = time.time()
    model.train()
    for step in range(start, total_steps):
        lr = learning_rate(step, total_steps, args.lr, warmup)
        for g in optimizer.param_groups:
            g["lr"] = lr
        batch = [train[i] for i in order[step * args.per_step:(step + 1) * args.per_step]]
        n_answer = sum(len(r["ids"]) - r["start"] for r in batch)
        t, loss_sum = time.time(), 0.0
        for r in batch:
            loss, _ = answer_loss_sum(model, r["ids"].to(device), r["start"])
            if not torch.isfinite(loss):
                raise RuntimeError("non-finite loss")
            (loss / n_answer).backward()
            loss_sum += loss.item()
        grad_norm = torch.nn.utils.clip_grad_norm_(params, 1.0).item()
        optimizer.step()
        optimizer.zero_grad(set_to_none=True)
        log(kind="train", step=step + 1, loss=round(loss_sum / n_answer, 4), lr=lr,
            grad_norm=round(grad_norm, 3), secs=round(time.time() - t, 1),
            max_mem_gb=round(torch.cuda.max_memory_allocated() / 2**30, 1))

        done = step + 1 == total_steps
        if not args.max_steps and ((step + 1) % eval_every == 0 or done):
            loss, by_category = validate(model, val, device)
            log(kind="validation", step=step + 1, val_loss=round(loss, 4), categories=by_category)
            model.save_pretrained(run / "adapters" / f"step-{step + 1:04d}")
            if loss < best:
                best = loss
                model.save_pretrained(run / "adapters" / "best")
                best_file.write_text(json.dumps(
                    {"step": step + 1, "val_loss": loss, "categories": by_category}))
        if not args.max_steps and (done or time.time() - last_save > args.checkpoint_minutes * 60):
            save_last(model, optimizer, step + 1, run)
            last_save = time.time()

    print("smoke test passed" if args.max_steps else "finished", flush=True)


if __name__ == "__main__":
    main()