#!/usr/bin/env python3
"""Build a vLLM-servable decision adapter for a dual-head bundle.

The 24-slot fp32 decision head (W, b) is expressed as an lm_head LoRA:
    lm_head'[v_i] = lm_head[v_i] + ΔW_i,  ΔW_i = W_i - lm_head[v_i]   (only the 24 verbalizer rows change)
so ΔW has rank <= 24 and is written as lora_A = ΔW (24 rows, zero-padded), lora_B = one-hot(v_i) / scaling.
With the adapter active, the logit of verbalizer v_i is exactly W_i·h. The bias b and the per-kind temperature are
applied client-side:  p = softmax_active((logprob_i + b_i) / T)  — the log-softmax normaliser cancels.
The backbone LoRA (r=16) is zero-padded to the common rank (default 32) with alpha scaled so B·A·alpha/r is unchanged.

python3 scripts/build_vllm_adapter.py --bundle exports/jev-judge-qwen35-9b-v0.8
"""

from __future__ import annotations

import argparse
import glob
import json
import os

import torch
from safetensors import safe_open
from safetensors.torch import load_file, save_file


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--bundle", required=True)
    ap.add_argument("--rank", type=int, default=32)
    ap.add_argument("--out", default=None, help="default: <bundle>/adapter_vllm")
    ap.add_argument("--dtype", choices=["bfloat16", "float32"], default="bfloat16",
                    help="storage dtype (vLLM computes LoRA in the model dtype, bf16, anyway)")
    args = ap.parse_args()
    out = args.out or os.path.join(args.bundle, "adapter_vllm")
    src = os.path.join(args.bundle, "adapter")
    cfg = json.load(open(os.path.join(src, "adapter_config.json")))
    jc = json.load(open(os.path.join(args.bundle, "judge_config.json")))
    r0, a0 = cfg["r"], cfg["lora_alpha"]
    R = args.rank
    if R < max(r0, jc["slots"]["num_slots"]):
        raise SystemExit("rank must be >= max(backbone r, 24)")
    scaling = a0 / r0
    alpha = scaling * R  # keep alpha/r identical after padding

    sd = load_file(os.path.join(src, "adapter_model.safetensors"))
    new: dict[str, torch.Tensor] = {}
    for k, t in sd.items():
        t = t.to(torch.float32)
        if ".lora_A." in k:
            new[k] = torch.cat([t, torch.zeros(R - t.shape[0], t.shape[1])], 0)
        elif ".lora_B." in k:
            new[k] = torch.cat([t, torch.zeros(t.shape[0], R - t.shape[1])], 1)
        else:
            new[k] = t

    # lm_head rows of the pristine base (shipped in the bundle)
    lm = None
    for sh in sorted(glob.glob(os.path.join(args.bundle, "model-*.safetensors"))):
        with safe_open(sh, "pt", "cpu") as f:
            if "lm_head.weight" in f.keys():
                lm = f.get_tensor("lm_head.weight")
                break
    if lm is None:
        raise SystemExit("bundle has no lm_head.weight (export with --mode unmerged / --include-lm-head)")
    head = load_file(os.path.join(args.bundle, "head.safetensors"))
    W = head["proj.weight"].to(torch.float32)
    ids = torch.as_tensor(jc["verbalizer_ids"])
    dW = W - lm[ids].to(torch.float32)
    V, H = lm.shape
    A = torch.zeros(R, H)
    A[: len(ids)] = dW
    B = torch.zeros(V, R)
    B[ids, torch.arange(len(ids))] = 1.0 / scaling
    new["base_model.model.lm_head.lora_A.weight"] = A
    new["base_model.model.lm_head.lora_B.weight"] = B

    os.makedirs(out, exist_ok=True)
    dt = torch.bfloat16 if args.dtype == "bfloat16" else torch.float32
    save_file({k: v.to(dt).contiguous() for k, v in new.items()}, os.path.join(out, "adapter_model.safetensors"))
    cfg2 = dict(cfg)
    cfg2.update({"r": R, "lora_alpha": alpha, "lora_dropout": 0.0,
                 "target_modules": sorted(set(cfg["target_modules"]) | {"lm_head"}),
                 "base_model_name_or_path": None})
    json.dump(cfg2, open(os.path.join(out, "adapter_config.json"), "w"), indent=2)
    json.dump({"bias": head["proj.bias"].tolist(), "verbalizer_ids": jc["verbalizer_ids"], "slots": jc["slots"],
               "note": "decision logits = lm_head-LoRA logprobs of verbalizer ids + bias, then per-kind temperature"},
              open(os.path.join(out, "decision_head.json"), "w"), indent=1)
    print(f"-> {out}: {len(new)} tensors, r={R}, alpha={alpha:g} (scaling {scaling:g}); lm_head ΔW rows {len(ids)}, "
          f"max|ΔW| {dW.abs().max():.4f}")


if __name__ == "__main__":
    main()