"""LoRA training of the BT policy (EXPERIMENT_LOG.md 2026-10-17): Qwen3.5-9B reads one frame and BT_PROMPT with the
target name and writes the behavior code (one digit token per FSQ dimension). Runs with torchrun in .train-venv.

Rows (jsonl): image (PNG path), name, code (digits); ppo rows also adv and vllm_logp, and memory (the two 0/1 strings
of the decision's prompt) when the rollouts ran with memory. A row is the chat template of [image, bt_prompt(name,
memory)] with the generation prompt (thinking off), as vLLM renders BTPolicy's request, followed by the code and
<|im_end|>; batches are right-padded and the scored tokens gathered per row.
- sft: cross entropy on the code tokens and <|im_end|>, new LoRA (r 16, alpha 32, train_iql.py's modules), AdamW,
  10-update warmup then cosine; per-digit and whole-code (all digits argmax) accuracy on --valid after each epoch.
  Saves the adapter to --out (vLLM-loadable).
- ppo: continues --adapter (and its optimizer.pt if it has one, i.e. a PPO round's), loads --ref (the SFT adapter) as
  a second adapter; old and ref log-probs of the code tokens computed once with no_grad, then --epochs of PPO:
  loss = -min(ratio A, clip(ratio, 0.8, 1.2) A) + beta * sum over the code tokens of k3 = exp(ref - lp) - (ref - lp) - 1,
  ratio = exp(sum lp - sum old). Saves the adapter, optimizer.pt and stats.json (pg_loss, kl, clip_frac of the last
  epoch; logp_gap = mean of old - vllm_logp, logp_gap_abs) to --out.
Usage: CUDA_VISIBLE_DEVICES=6,7 .train-venv/bin/python -m torch.distributed.run --nproc_per_node=2 train/bt_train.py
         sft --data TRAIN.jsonl --valid VALID.jsonl --out DIR [--epochs 1 --lr 1e-4 --max-updates N]
       ... train/bt_train.py ppo --data ROUND/train.jsonl --adapter ADAPTER --ref SFT --out DIR
"""
import argparse
import json
import math
import os
import random
import time
from pathlib import Path

import torch
import torch.distributed as dist
from peft import LoraConfig, PeftModel, get_peft_model
from PIL import Image
from torch.nn.parallel import DistributedDataParallel
from torch.utils.data import DataLoader, DistributedSampler
from transformers import AutoProcessor, Qwen3_5ForConditionalGeneration

from ctrl_policy import bt_prompt

ROOT = Path(__file__).resolve().parents[1]
MODULES = (r"model.language_model.*\.(q_proj|k_proj|v_proj|o_proj|in_proj_qkv|in_proj_z|in_proj_b|in_proj_a|out_proj"
           r"|gate_proj|up_proj|down_proj)$")  # train_iql.py
END = "<|im_end|>"
CLIP, BETA = 0.2, 0.02


def log(message):
    print(f"{time.strftime('%Y-%m-%d %H:%M:%S')} {message}", flush=True)


class Collate:
    def __init__(self, processor):
        self.processor = processor

    def __call__(self, rows):
        p = self.processor
        texts = [p.apply_chat_template([{"role": "user", "content": [{"type": "image"}, {"type": "text",
                                         "text": bt_prompt(r["name"], r.get("memory"))}]}], tokenize=False,
                                       add_generation_prompt=True, enable_thinking=False) + r["code"] + END
                 for r in rows]
        inputs = p(text=texts, images=[Image.open(r["image"]).convert("RGB") for r in rows], padding=True,
                   padding_side="right", return_tensors="pt")
        n = len(rows[0]["code"])
        ends = inputs["attention_mask"].sum(1)  # sequence lengths: the last token, <|im_end|>, is at end - 1
        targets = torch.stack([inputs["input_ids"][i, e - n - 1:e] for i, e in enumerate(ends)])  # code + im_end
        extra = {k: torch.tensor([r[k] for r in rows], dtype=torch.float32) for k in ("adv", "vllm_logp", "index")
                 if k in rows[0]}
        return {"inputs": dict(inputs), "ends": ends, "targets": targets, **extra}


def token_logprobs(model, batch, device):
    """(B, n + 1) log-probs of the code tokens and <|im_end|>, and the logits' argmax at those positions."""
    inputs = {k: v.to(device) for k, v in batch["inputs"].items()}
    ends, k = batch["ends"], batch["targets"].shape[1]
    lo, hi = int(ends.min()) - k - 1, int(ends.max()) - 1  # logits at position j predict token j + 1
    with torch.autocast("cuda", dtype=torch.bfloat16):
        logits = model(**inputs, logits_to_keep=torch.arange(lo, hi, device=device), use_cache=False).logits
    rows = torch.stack([logits[i, e - k - 1 - lo:e - 1 - lo] for i, e in enumerate(ends.tolist())]).float()
    lp = rows.log_softmax(-1)
    targets = batch["targets"].to(device)
    return lp.gather(-1, targets[..., None])[..., 0], rows.argmax(-1) == targets


def load_model(local):
    model = Qwen3_5ForConditionalGeneration.from_pretrained(
        ROOT / "models/Qwen3.5-9B", dtype=torch.bfloat16, device_map={"": local}, attn_implementation="sdpa",
        local_files_only=True, use_kernels=True)
    model.requires_grad_(False)
    return model


def train_setup(model):
    model.enable_input_require_grads()
    model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})
    model.config.use_cache = False
    model.config.text_config.use_cache = False


def loader(rows, collate, batch, rank, world, seed, shuffle=True):
    sampler = DistributedSampler(rows, num_replicas=world, rank=rank, seed=seed, shuffle=shuffle)
    return sampler, DataLoader(rows, batch_size=batch, sampler=sampler, collate_fn=collate, num_workers=4,
                               persistent_workers=True, prefetch_factor=2)


@torch.no_grad()
def accuracy(model, rows, collate, batch, rank, world, device):
    """Per-digit and whole-code accuracy (teacher-forced argmax) of the code tokens on rows."""
    model.eval()
    totals = torch.zeros(3, device=device)
    for b in DataLoader(rows[rank::world], batch_size=batch, collate_fn=collate, num_workers=4):
        _, right = token_logprobs(model, b, device)
        right = right[:, :-1]
        totals += torch.tensor([right.sum(), right.all(1).sum(), len(right)], device=device, dtype=torch.float32)
    dist.all_reduce(totals)
    model.train()
    return {"digit_acc": round((totals[0] / totals[2] / right.shape[1]).item(), 4),
            "code_acc": round((totals[1] / totals[2]).item(), 4), "rows": int(totals[2])}


def sft(args, rank, world, local, device):
    rows = [json.loads(l) for l in args.data.read_text().splitlines()]
    valid = [json.loads(l) for l in args.valid.read_text().splitlines()]
    processor = AutoProcessor.from_pretrained(ROOT / "models/Qwen3.5-9B", local_files_only=True)
    collate = Collate(processor)
    torch.manual_seed(args.seed)
    model = get_peft_model(load_model(local), LoraConfig(r=16, lora_alpha=32, lora_dropout=0.0,
                                                         task_type="CAUSAL_LM", target_modules=MODULES))
    train_setup(model)
    ddp = DistributedDataParallel(model, device_ids=[local])
    sampler, data = loader(rows, collate, args.batch, rank, world, args.seed)
    opt = torch.optim.AdamW((p for p in model.parameters() if p.requires_grad), lr=args.lr)
    total = min(args.max_updates or 10**9, len(data) * args.epochs)
    if rank == 0:
        log(f"SFT rows {len(rows)} valid {len(valid)}, {world} GPUs x batch {args.batch}, updates {total}")
        log("VALID before " + json.dumps(accuracy(model, valid, collate, args.batch, rank, world, device)))
    else:
        accuracy(model, valid, collate, args.batch, rank, world, device)
    update, start = 0, time.monotonic()
    for epoch in range(args.epochs):
        sampler.set_epoch(epoch)
        for b in data:
            lp, _ = token_logprobs(ddp, b, device)
            loss = -lp.mean()
            opt.zero_grad(set_to_none=True)
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            for g in opt.param_groups:
                g["lr"] = args.lr * min(1, (update + 1) / 10) * 0.5 * (1 + math.cos(math.pi * update / total))
            opt.step()
            update += 1
            if rank == 0 and (update <= 5 or update % 50 == 0):
                s = time.monotonic() - start
                log(f"SFT update {update}/{total} loss {loss.item():.4f} {s / update:.2f} s/update "
                    f"eta {(total - update) * s / update / 60:.0f} min peak {torch.cuda.max_memory_allocated() / 2**30:.1f} GiB")
            if update >= total:
                break
        acc = accuracy(model, valid, collate, args.batch, rank, world, device)
        if rank == 0:
            log(f"VALID epoch {epoch + 1} " + json.dumps(acc))
            model.save_pretrained(args.out)
            (args.out / "valid.json").write_text(json.dumps({"epoch": epoch + 1, "updates": update, **acc}))
        if update >= total:
            break


def ppo(args, rank, world, local, device):
    rows = [{**json.loads(l), "index": i} for i, l in enumerate(args.data.read_text().splitlines())]
    processor = AutoProcessor.from_pretrained(ROOT / "models/Qwen3.5-9B", local_files_only=True)
    collate = Collate(processor)
    torch.manual_seed(args.seed)
    model = PeftModel.from_pretrained(load_model(local), args.adapter, is_trainable=True)
    model.load_adapter(args.ref, adapter_name="ref")
    n = len(rows[0]["code"])
    lp_of = {}
    for name in ("ref", "default"):  # once per round, no_grad: the reference and the old policy
        model.set_adapter(name)
        model.eval()
        mine = {}
        with torch.no_grad():
            for b in DataLoader(rows[rank::world], batch_size=2 * args.batch, collate_fn=collate, num_workers=4):
                lp, _ = token_logprobs(model, b, device)
                mine.update(zip(b["index"].long().tolist(), lp[:, :n].cpu()))
        gathered = [None] * world
        dist.all_gather_object(gathered, mine)
        lp_of[name] = torch.stack([{k: v for g in gathered for k, v in g.items()}[i] for i in range(len(rows))])
    model.set_adapter("default")
    for name, p in model.named_parameters():
        p.requires_grad = "lora_" in name and ".default." in name
    model.train()
    train_setup(model)
    ddp = DistributedDataParallel(model, device_ids=[local])
    params = [p for p in model.parameters() if p.requires_grad]
    opt = torch.optim.AdamW(params, lr=args.lr)
    if (args.adapter / "optimizer.pt").exists():
        opt.load_state_dict(torch.load(args.adapter / "optimizer.pt", map_location=device))
        for g in opt.param_groups:
            g["lr"] = args.lr
    sampler, data = loader(rows, collate, args.batch, rank, world, args.seed)
    old_all, ref_all = lp_of["default"].to(device), lp_of["ref"].to(device)
    vllm = torch.tensor([r["vllm_logp"] for r in rows], device=device)
    gap = old_all.sum(1) - vllm
    start = time.monotonic()
    for epoch in range(args.epochs):
        sampler.set_epoch(epoch)
        stats = []
        for step, b in enumerate(data):
            idx = b["index"].long().to(device)
            lp, _ = token_logprobs(ddp, b, device)
            lp = lp[:, :n]
            adv = b["adv"].to(device)
            ratio = (lp.sum(1) - old_all[idx].sum(1)).exp()
            pg = -torch.min(ratio * adv, ratio.clamp(1 - CLIP, 1 + CLIP) * adv)
            d = ref_all[idx] - lp
            kl = (d.exp() - d - 1).sum(1)
            loss = (pg + BETA * kl).mean() / args.accumulation
            sync = (step + 1) % args.accumulation == 0 or step + 1 == len(data)
            if sync:
                loss.backward()
                torch.nn.utils.clip_grad_norm_(params, 1.0)
                opt.step()
                opt.zero_grad(set_to_none=True)
            else:
                with ddp.no_sync():
                    loss.backward()
            stats.append([pg.mean().item(), kl.mean().item(), ((ratio - 1).abs() > CLIP).float().mean().item()])
        if rank == 0:
            log(f"PPO epoch {epoch + 1} pg {sum(s[0] for s in stats) / len(stats):+.4f} "
                f"kl {sum(s[1] for s in stats) / len(stats):.4f} {time.monotonic() - start:.0f} s")
    totals = torch.tensor(stats, device=device).mean(0)
    dist.all_reduce(totals)
    totals /= world
    if rank == 0:
        args.out.mkdir(parents=True, exist_ok=True)
        model.save_pretrained(args.out, selected_adapters=["default"])
        torch.save(opt.state_dict(), args.out / "optimizer.pt")
        result = {"pg_loss": round(totals[0].item(), 4), "kl": round(totals[1].item(), 4),
                  "clip_frac": round(totals[2].item(), 4), "logp_gap": round(gap.mean().item(), 4),
                  "logp_gap_abs": round(gap.abs().mean().item(), 4), "rows": len(rows),
                  "train_s": round(time.monotonic() - start)}
        (args.out / "stats.json").write_text(json.dumps(result))
        log("PPO_DONE " + json.dumps(result))


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("mode", choices=["sft", "ppo"])
    ap.add_argument("--data", type=Path, required=True)
    ap.add_argument("--valid", type=Path)
    ap.add_argument("--adapter", type=Path)
    ap.add_argument("--ref", type=Path)
    ap.add_argument("--out", type=Path, required=True)
    ap.add_argument("--epochs", type=int)
    ap.add_argument("--lr", type=float)
    ap.add_argument("--batch", type=int, default=16, help="rows per GPU per forward")
    ap.add_argument("--accumulation", type=int, default=2, help="ppo: forwards per update (64 rows on 2 GPUs)")
    ap.add_argument("--max-updates", type=int)
    ap.add_argument("--seed", type=int, default=0)
    args = ap.parse_args()
    rank, world, local = int(os.environ["RANK"]), int(os.environ["WORLD_SIZE"]), int(os.environ["LOCAL_RANK"])
    torch.set_num_threads(4)
    random.seed(args.seed)
    torch.cuda.set_device(local)
    device = torch.device("cuda", local)
    dist.init_process_group("nccl", device_id=device)
    if args.mode == "sft":
        args.epochs, args.lr = args.epochs or 1, args.lr or 1e-4
        sft(args, rank, world, local, device)
    else:
        args.epochs, args.lr = args.epochs or 2, args.lr or 3e-6
        ppo(args, rank, world, local, device)
    dist.barrier()
    dist.destroy_process_group()


if __name__ == "__main__":
    main()