"""Apply the LoRA adapter in this repo's adapter/ folder to the base model, then score role-labelled outlines.

adapter/ holds the adapter as trained, in the layout of the Tinker training service. Do not load it with
peft.PeftModel. In transformers, Nemotron fuses the Mamba gate and x projections into one `in_proj`, and stores each
MoE layer's 128 routed experts as a single 3D tensor. PEFT has no module to attach those LoRA weights to and skips them
without a warning. This script instead merges every LoRA delta, W += (alpha / r) * B @ A, into the base weights in
place, following the same rules as the tinker-cookbook merge (tinker_cookbook.weights.build_hf_model). The resulting
weights are bit-identical to the merged model in this repo.

Needs transformers >= 5.15, torch, safetensors and huggingface_hub. The weights take 59 GiB of GPU memory and each
input adds about 4.2 MiB per token, so one 80 GB GPU handles inputs up to about 4,000 tokens; see the model card.
"""
import json

import torch
from huggingface_hub import snapshot_download
from safetensors.torch import load_file
from transformers import AutoModelForCausalLM, AutoTokenizer

REPO = "rishanthrajendhran/IdeaLens-NoParaphrase"
BASE = "nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16"
SYSTEM = 'Given a role-labelled outline of a document, answer with one word: human if the source document was human-written, ai if it was AI-generated.'
SUFFIX = "<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n"
HUMAN, AI = 50755, 2464
CUT = 0.09947  # P(human) at or below this is flagged: the global cut at 1% FPR (thresholds.json)


def apply_tinker_lora(model, adapter_dir):
    """Merge a Tinker LoRA archive for Nemotron-3.5 into `model` in place. Returns the parameter names changed."""
    cfg = json.load(open(f"{adapter_dir}/adapter_config.json"))
    scale = cfg["lora_alpha"] / cfg["r"]
    lora = load_file(f"{adapter_dir}/adapter_model.safetensors")
    params = dict(model.named_parameters())
    touched = []
    for key_a in sorted(k for k in lora if k.endswith(".lora_A.weight")):
        A, B = lora[key_a], lora[key_a.replace(".lora_A.", ".lora_B.")]
        if A.numel() == 0:  # the experts have no gate projection; Tinker keeps an empty w3 placeholder
            continue
        name = key_a.removeprefix("base_model.model.").removesuffix(".lora_A.weight")
        rows = None
        if name == "model.lm_head":  # Tinker nests the LM head under model.
            target = "lm_head.weight"
        elif name.endswith((".gate_proj", ".x_proj")):  # Mamba: in_proj rows are [gate | x | B | C | dt]
            layer, proj = name.rsplit(".", 1)
            target = f"{layer}.in_proj.weight"
            start = 0 if proj == "gate_proj" else lora[f"base_model.model.{layer}.gate_proj.lora_B.weight"].shape[0]
            rows = slice(start, start + B.shape[0])
        elif name.endswith(".experts.w1"):  # routed experts: w1 = up_proj, one (expert, out, in) tensor per layer
            target = name.removesuffix("w1") + "up_proj"
        elif name.endswith(".experts.w2"):  # w2 = down_proj
            target = name.removesuffix("w2") + "down_proj"
        else:  # attention, Mamba out_proj, shared experts
            target = name + ".weight"
        W = params[target]  # KeyError here means the adapter does not match this model
        # For experts one side is shared (leading dim 1) and broadcasts across the 128 experts.
        delta = scale * torch.matmul(B.to(W.device, torch.float32), A.to(W.device, torch.float32))
        with torch.no_grad():
            view = W.data if rows is None else W.data[rows]
            assert view.shape == delta.shape, (target, tuple(view.shape), tuple(delta.shape))
            view.copy_((view.float() + delta).to(W.dtype))
        touched.append(target)
    return touched


def load_model(repo=REPO, device_map="auto"):
    tok = AutoTokenizer.from_pretrained(BASE)
    model = AutoModelForCausalLM.from_pretrained(BASE, dtype=torch.bfloat16, device_map=device_map).eval()
    adapter_dir = snapshot_download(repo, allow_patterns=["adapter/*"]) + "/adapter"
    apply_tinker_lora(model, adapter_dir)
    return model, tok


@torch.no_grad()
def p_human(model, tok, outline):
    """P(human) for one outline: one `[Role] content` line per item."""
    ids = tok.encode(f"<|im_start|>system\n{SYSTEM}<|im_end|>\n<|im_start|>user\n{outline}{SUFFIX}",
                     add_special_tokens=False)
    logits = model(torch.tensor([ids], device=model.device)).logits[0, -1].float()
    return torch.softmax(logits[[HUMAN, AI]], -1)[0].item()


if __name__ == "__main__":
    model, tok = load_model()
    outline = ("[Central Development] A town's water supply fails after a drought, and residents organise to share wells.\n"
               "[Background Context] The reservoir has been shrinking for three summers.\n"
               "[Open Question] Whether the council will fund a new pipeline remains undecided.")
    p = p_human(model, tok, outline)
    print(f"P(human) = {p:.4f} -> {'flagged as AI ideas' if p <= CUT else 'not flagged'}")