#!/usr/bin/env python3
"""Fine-tune the auditor model with LoRA, fully offline, on CPU (or GPU if present).

The base model must already be in a local folder (downloaded once, see training/README.md).
No network access is used during training (HF_HUB_OFFLINE is forced on).

Usage (from the endpoint-audit folder, with the training environment):
    training/.venv/bin/python training/train_lora.py --benchmark            # time 10 steps, estimate the full run
    training/.venv/bin/python training/train_lora.py --max-train 3000       # train on a balanced subset
    training/.venv/bin/python training/train_lora.py --max-train 3000 --resume   # continue after a stop
    ... --qat4 --init-adapter training/models/auditor-lora-v0.5/final   # make a trained adapter survive 4-bit WebLLM builds

--qat4 trains through the same 4-bit rounding the browser build uses (q4f16_1: groups of 32 along the input, scale
max|w|/7 in float16), so the merged model still works once mlc_quantize.py rounds it. Without it, model v0.5 (0.5B)
answered well in full precision and wrote junk in 4 bits. The rounding is applied to base + LoRA weights together on
every forward pass, with the gradient passed straight through it.

Loss is computed on the assistant answer only. Examples longer than --max-len tokens are skipped
(never truncated), so the model always sees complete prompts and answers.
"""
import argparse
import json
import math
import os
import random
import time

os.environ["HF_HUB_OFFLINE"] = "1"
os.environ["TRANSFORMERS_OFFLINE"] = "1"

import torch  # noqa: E402
from peft import LoraConfig, PeftModel, get_peft_model  # noqa: E402
from transformers import AutoModelForCausalLM, AutoTokenizer  # noqa: E402

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
DEFAULT_BASE = os.path.join(ROOT, "training", "models", "Qwen2.5-0.5B-Instruct")
LORA_TARGETS = ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"]
# PyTorch 2.2 (the last build for Intel Macs) returns wrong attention results with the "sdpa" kernel when
# gradients are on and more than one thread is used, which drives training to NaN. "eager" is correct and faster here.
ATTENTION = "eager"


def round4(w):
    """w as the browser will see it: float16, then 4-bit groups of 32 (same as mlc_quantize.quantize, q4f16_1)."""
    out = w.shape[0]
    g = w.detach().half().float().reshape(out, -1, 32)
    scale = (g.abs().amax(-1, keepdim=True).half() * torch.tensor(1 / 7, dtype=torch.float16)).float()
    x = torch.where(scale == 0, torch.zeros_like(g), g / scale)
    q = (torch.sign(x) * torch.floor(x.abs() + 0.5)).clamp(-7, 7)
    return (q * scale).reshape(w.shape)


def lora_forward_q4(layer):
    """Replace a PEFT LoRA linear's forward: merge, round to 4 bits, pass the gradient straight through."""
    base, name = layer.base_layer, layer.active_adapters[0]
    lora_a, lora_b, scaling = layer.lora_A[name], layer.lora_B[name], layer.scaling[name]

    def forward(x, *args, **kwargs):
        w = base.weight + (lora_b.weight @ lora_a.weight) * scaling
        w = w + (round4(w) - w).detach()
        return torch.nn.functional.linear(x, w, base.bias)
    layer.forward = forward


def use_4bit(model):
    """Train against 4-bit weights: embeddings (also the output layer: tied) once, LoRA layers on every pass."""
    from peft.tuners.lora import Linear as LoraLinear
    emb = model.get_input_embeddings()
    with torch.no_grad():
        emb.weight.copy_(round4(emb.weight))
    layers = [m for m in model.modules() if isinstance(m, LoraLinear)]
    for layer in layers:
        lora_forward_q4(layer)
    print(f"4-bit-aware training: {len(layers)} LoRA layers + embeddings rounded as in q4f16_1")


def load_examples(path):
    with open(path, encoding="utf-8") as handle:
        return [json.loads(line) for line in handle]


def balanced_subset(examples, limit, seed):
    """Pick up to `limit` examples spread evenly over task, OS and status."""
    if not limit or limit >= len(examples):
        return examples
    groups = {}
    for item in examples:
        meta = item["meta"]
        groups.setdefault((meta["task"], meta["os"], meta.get("status", "")), []).append(item)
    rng = random.Random(seed)
    for items in groups.values():
        rng.shuffle(items)
    chosen, keys = [], sorted(groups)
    while len(chosen) < limit and keys:
        for key in list(keys):
            if groups[key]:
                chosen.append(groups[key].pop())
                if len(chosen) == limit:
                    break
            else:
                keys.remove(key)
    rng.shuffle(chosen)
    return chosen


def encode(tokenizer, example, max_len):
    """Token ids with labels masked (-100) everywhere except the assistant answer."""
    messages = example["messages"]
    prompt_ids = tokenizer.apply_chat_template(messages[:-1], tokenize=True, add_generation_prompt=True)
    full_ids = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=False)
    if full_ids[:len(prompt_ids)] != prompt_ids or len(full_ids) > max_len:
        return None
    labels = [-100] * len(prompt_ids) + full_ids[len(prompt_ids):]
    return {"input_ids": full_ids, "labels": labels}


def batches(encoded, batch_size, pad_id):
    for start in range(0, len(encoded), batch_size):
        chunk = encoded[start:start + batch_size]
        width = max(len(e["input_ids"]) for e in chunk)
        ids = torch.full((len(chunk), width), pad_id, dtype=torch.long)
        labels = torch.full((len(chunk), width), -100, dtype=torch.long)
        mask = torch.zeros((len(chunk), width), dtype=torch.long)
        for row, item in enumerate(chunk):
            n = len(item["input_ids"])
            ids[row, :n] = torch.tensor(item["input_ids"])
            labels[row, :n] = torch.tensor(item["labels"])
            mask[row, :n] = 1
        yield ids, labels, mask


@torch.no_grad()
def evaluate(model, encoded, pad_id, batch_size):
    model.eval()
    total, count = 0.0, 0
    for ids, labels, mask in batches(encoded, batch_size, pad_id):
        total += model(input_ids=ids, attention_mask=mask, labels=labels).loss.item()
        count += 1
    model.train()
    return total / max(count, 1)


def main():
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--base", default=DEFAULT_BASE, help="local folder with the base model")
    parser.add_argument("--data", default=os.path.join(ROOT, "training", "out"))
    parser.add_argument("--output", default=os.path.join(ROOT, "training", "models", "auditor-lora"))
    parser.add_argument("--max-train", type=int, default=3000, help="balanced subset size (0 = all)")
    parser.add_argument("--max-val", type=int, default=200)
    parser.add_argument("--max-len", type=int, default=2048, help="skip examples longer than this many tokens")
    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=2e-4)
    parser.add_argument("--lora-r", type=int, default=16)
    parser.add_argument("--lora-alpha", type=int, default=32)
    parser.add_argument("--save-every", type=int, default=50, help="optimizer steps between checkpoints")
    parser.add_argument("--eval-every", type=int, default=100)
    parser.add_argument("--threads", type=int, default=os.cpu_count())
    parser.add_argument("--seed", type=int, default=13)
    parser.add_argument("--benchmark", action="store_true", help="run 10 optimizer steps, report speed and estimated total time")
    parser.add_argument("--resume", action="store_true", help="continue from the last checkpoint in --output")
    parser.add_argument("--qat4", action="store_true", help="train through 4-bit rounding (for WebLLM q4 builds)")
    parser.add_argument("--init-adapter", help="start from this trained adapter instead of a fresh one")
    args = parser.parse_args()

    torch.manual_seed(args.seed)
    torch.set_num_threads(args.threads)
    device = "cuda" if torch.cuda.is_available() else "cpu"
    print(f"Device: {device}, threads: {args.threads}, base: {args.base}")

    tokenizer = AutoTokenizer.from_pretrained(args.base, local_files_only=True)
    pad_id = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else tokenizer.eos_token_id

    train_raw = balanced_subset(load_examples(os.path.join(args.data, "train.jsonl")), args.max_train, args.seed)
    val_raw = balanced_subset(load_examples(os.path.join(args.data, "validation.jsonl")), args.max_val, args.seed)
    train = [e for e in (encode(tokenizer, x, args.max_len) for x in train_raw) if e]
    val = [e for e in (encode(tokenizer, x, args.max_len) for x in val_raw) if e]
    # Group similar lengths to reduce padding, but keep task order mixed across groups.
    train.sort(key=lambda e: len(e["input_ids"]))
    groups = [train[i:i + args.batch_size] for i in range(0, len(train), args.batch_size)]
    random.Random(args.seed).shuffle(groups)
    train = [e for g in groups for e in g]
    tokens = sum(len(e["input_ids"]) for e in train)
    print(f"Train examples: {len(train)} ({len(train_raw) - len(train)} skipped as longer than {args.max_len} tokens), "
          f"{tokens:,} tokens; validation: {len(val)}")

    model = AutoModelForCausalLM.from_pretrained(args.base, local_files_only=True, torch_dtype=torch.float32,
                                                 attn_implementation=ATTENTION)
    model.config.use_cache = False
    model.gradient_checkpointing_enable()
    model.enable_input_require_grads()
    last_checkpoint = os.path.join(args.output, "last")
    state_path = os.path.join(last_checkpoint, "trainer_state.json")
    if args.resume and os.path.exists(state_path):
        model = PeftModel.from_pretrained(model, last_checkpoint, is_trainable=True)
        with open(state_path) as handle:
            state = json.load(handle)
        print(f"Resuming from step {state['step']}")
    elif args.init_adapter:
        model = PeftModel.from_pretrained(model, args.init_adapter, is_trainable=True)
        state = {"step": 0, "micro": 0, "history": [], "init_adapter": args.init_adapter}
        print(f"Starting from adapter {args.init_adapter}")
    else:
        model = get_peft_model(model, LoraConfig(r=args.lora_r, lora_alpha=args.lora_alpha, lora_dropout=0.05,
                                                 target_modules=LORA_TARGETS, task_type="CAUSAL_LM"))
        state = {"step": 0, "micro": 0, "history": []}
    model.print_trainable_parameters()
    if args.qat4:
        use_4bit(model)
        state["qat4"] = True
    model.to(device)
    model.train()

    optimizer = torch.optim.AdamW([p for p in model.parameters() if p.requires_grad], lr=args.lr, weight_decay=0.0)
    micro_per_epoch = math.ceil(len(train) / args.batch_size)
    total_steps = math.ceil(micro_per_epoch * args.epochs / args.grad_accum)
    if args.benchmark:
        total_steps = min(total_steps, state["step"] + 10)
    warmup = max(1, int(0.03 * total_steps))
    scheduler = torch.optim.lr_scheduler.LambdaLR(
        optimizer, lambda s: min(1.0, (s + 1) / warmup) * max(0.0, 1 - s / max(total_steps, 1)) if s >= warmup
        else (s + 1) / warmup)
    for _ in range(state["step"]):
        scheduler.step()

    def save(tag):
        path = os.path.join(args.output, tag)
        model.save_pretrained(path)
        tokenizer.save_pretrained(path)
        with open(os.path.join(path, "trainer_state.json"), "w") as handle:
            json.dump(state, handle, indent=1)

    batch_list = list(batches(train, args.batch_size, pad_id))
    started = time.time()
    step_start = time.time()
    running = 0.0
    bad_steps = 0
    micro = state["micro"]
    while state["step"] < total_steps:
        ids, labels, mask = batch_list[micro % len(batch_list)]
        loss = model(input_ids=ids.to(device), attention_mask=mask.to(device), labels=labels.to(device)).loss
        (loss / args.grad_accum).backward()
        running += loss.item()
        micro += 1
        if micro % args.grad_accum:
            continue
        grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        if not torch.isfinite(grad_norm):
            # Never apply a NaN/inf update: it would destroy every weight.
            optimizer.zero_grad()
            running = 0.0
            bad_steps += 1
            print(f"skipped a step with a non-finite gradient ({bad_steps} in a row)", flush=True)
            if bad_steps >= 3:
                raise SystemExit("Stopping: 3 non-finite gradients in a row. The last good checkpoint is kept.")
            continue
        bad_steps = 0
        optimizer.step()
        scheduler.step()
        optimizer.zero_grad()
        state["step"] += 1
        state["micro"] = micro
        seconds = time.time() - step_start
        step_start = time.time()
        avg = running / args.grad_accum
        running = 0.0
        state["history"].append({"step": state["step"], "loss": round(avg, 4), "seconds": round(seconds, 1)})
        # Estimate from recent steps only, so a pause (e.g. the Mac sleeping) does not distort it.
        recent = [h["seconds"] for h in state["history"][-10:] if h["seconds"] < 1800] or [seconds]
        remaining = (total_steps - state["step"]) * sum(recent) / len(recent)
        print(f"step {state['step']}/{total_steps}  loss {avg:.4f}  {seconds:.1f}s/step  ~{remaining / 3600:.1f} h left", flush=True)
        if not args.benchmark and state["step"] % args.save_every == 0:
            save("last")
            save(f"step-{state['step']:04d}")  # numbered copy: a later bad save can never overwrite it
        if not args.benchmark and val and state["step"] % args.eval_every == 0:
            val_loss = evaluate(model, val, pad_id, args.batch_size)
            state["history"][-1]["val_loss"] = round(val_loss, 4)
            print(f"  validation loss {val_loss:.4f}", flush=True)

    if args.benchmark:
        per_step = sum(h["seconds"] for h in state["history"][-10:]) / max(len(state["history"][-10:]), 1)
        full = math.ceil(micro_per_epoch * args.epochs / args.grad_accum)
        print(f"\nBENCHMARK: {per_step:.1f} s per optimizer step ({args.grad_accum} examples). "
              f"Full run of {full} steps on {len(train)} examples: about {full * per_step / 3600:.1f} hours.")
        return
    save("last")
    save("final")
    if val:
        print(f"Final validation loss: {evaluate(model, val, pad_id, args.batch_size):.4f}")
    print(f"Adapter saved to {os.path.join(args.output, 'final')}")


if __name__ == "__main__":
    main()