# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2023-present Daniel Han-Chen & the Unsloth team. All rights reserved.

"""Clef benchmark; --fast sets UNSLOTH_CLEF_FAST; parity compares batched vs per-record head logits."""

import argparse
import json
import os
import sys
import tempfile
import time

parser = argparse.ArgumentParser()
parser.add_argument("--model", default = "tiny", help = "tiny, a Clef repo id or folder")
parser.add_argument("--load-in-4bit", action = "store_true")
parser.add_argument("--fast", choices = ["on", "off"], default = "on")
parser.add_argument("--dtype", choices = ["bf16", "fp16"], default = "bf16")
parser.add_argument("--seq-len", type = int, default = 4096)
parser.add_argument("--batch-size", type = int, default = 1)
parser.add_argument("--questions", type = int, default = 8)
parser.add_argument("--steps", type = int, default = 6)
parser.add_argument("--warmup", type = int, default = 2)
parser.add_argument("--log-recompiles", action = "store_true")
args = parser.parse_args()
os.environ["UNSLOTH_CLEF_FAST"] = "1" if args.fast == "on" else "0"

import torch  # noqa: E402

from unsloth import FastDecisionModel  # noqa: E402
from unsloth.models import clef, decision  # noqa: E402

if args.log_recompiles:
    torch._logging.set_logs(recompiles = True, graph_breaks = True)

TINY = "trl-internal-testing/tiny-Qwen3_5ForConditionalGeneration-NoThink"
TINY_HEAD = dict(hidden_size = 8, width = 32, routing_layers = 2, layers = 2, heads = 4, feedforward = 64)


def tiny_checkpoint() -> str:
    from safetensors.torch import save_file
    from transformers import AutoProcessor, Qwen3_5ForConditionalGeneration

    folder = os.path.join(tempfile.mkdtemp(prefix = "clef_tiny_"), "clef")
    torch.manual_seed(0)
    backbone = Qwen3_5ForConditionalGeneration.from_pretrained(TINY, dtype = torch.bfloat16)
    backbone.save_pretrained(folder)
    AutoProcessor.from_pretrained(TINY).save_pretrained(folder)
    head = clef.JointSchemaHead(
        **{**TINY_HEAD, "hidden_size": backbone.config.text_config.hidden_size}
    )
    with torch.no_grad():
        for parameter in head.parameters():
            parameter.add_(torch.randn_like(parameter) * 0.1)
    save_file(
        {k: v.to(torch.bfloat16) for k, v in head.state_dict().items()},
        f"{folder}/joint_head.safetensors",
    )
    with open(f"{folder}/joint_head_config.json", "w") as file:
        json.dump(head.config, file)
    return folder


def rows(tokenizer, count):
    # A long state filled to about --seq-len tokens, and --questions questions of 2 to 6 options.
    filler = "The customer reports intermittent checkout failures across several regions. "
    per = len(tokenizer(filler, add_special_tokens = False).input_ids)
    state = filler * max(1, (args.seq_len - 120 * args.questions - 200) // per)
    out = []
    for index in range(count):
        questions, gold = {}, {}
        for q in range(args.questions):
            options = [f"option {o} for field {q}" for o in range(2 + (q + index) % 5)]
            kind = ("choice", "score", "noul")[q % 3]
            if kind == "noul":
                questions[f"f{q}"] = {"type": "noul", "instructions": f"Is statement {q} true?"}
                gold[f"f{q}"] = {"label": "true" if (q + index) % 2 else "false"}
            elif kind == "score":
                questions[f"f{q}"] = {
                    "type": "score",
                    "instructions": f"Rate {q}",
                    "criteria": options,
                }
                gold[f"f{q}"] = {"label": (q + index) % len(options)}
            else:
                questions[f"f{q}"] = {
                    "type": "choice",
                    "instructions": f"Pick {q}",
                    "criteria": {f"o{o}": text for o, text in enumerate(options)},
                }
                gold[f"f{q}"] = {"label": f"o{index % len(options)}"}
        out.append({"state": f"{index}: {state}", "questions": questions, "gold": gold})
    return out


def head_parity(model, batch, amp):
    # Each head is scored against the per-record fp32 oracle with autocast off, since under bf16
    # autocast both round their matmuls; batched passes if its error is at most 1.5x per-record plus
    # one bf16 rounding step
    backbone = model._backbone()
    text_model = getattr(backbone.model, "language_model", backbone.model)
    with torch.no_grad():
        hidden = text_model(
            input_ids = batch["input_ids"], attention_mask = batch["attention_mask"], use_cache = False
        ).last_hidden_state
    embedding = backbone.get_output_embeddings().weight.detach()
    head = model.head
    os.environ["UNSLOTH_CLEF_FAST"] = "1"

    def run(forward, autocast):
        head.zero_grad()
        h = hidden.detach().clone().requires_grad_(True)
        with torch.autocast("cuda", dtype = amp or torch.float32, enabled = autocast):
            out = forward(
                h, batch["input_ids"], batch["attention_mask"], batch["records"], embedding
            )
        flat = [q for record in out for q in record]
        padded = torch.full((len(flat), max(len(z) for z in flat)), -1e4, device = h.device)
        for i, z in enumerate(flat):
            padded[i, : len(z)] = z.float()
        generator = torch.Generator(device = h.device).manual_seed(0)
        grad = torch.randn(padded.shape, device = h.device, generator = generator)
        (padded * grad * (padded > -1e3)).sum().backward()
        grads = {
            n: p.grad.detach().float().clone()
            for n, p in head.named_parameters()
            if p.grad is not None
        }
        return padded.detach(), h.grad.detach().float(), grads

    oracle = run(head.forward_per_record, False)
    ours = run(head.forward, amp is not None)
    reference = run(head.forward_per_record, amp is not None)
    os.environ["UNSLOTH_CLEF_FAST"] = "1" if args.fast == "on" else "0"
    head.zero_grad()
    valid = oracle[0] > -1e3

    def errors(got):
        return {
            "logits_max_abs": (got[0] - oracle[0])[valid].abs().max().item(),
            "hidden_grad_rel_l2": ((got[1] - oracle[1]).norm() / oracle[1].norm()).item(),
            "param_grad_rel_l2": max(
                ((got[2][n] - oracle[2][n]).norm() / oracle[2][n].norm().clamp_min(1e-30)).item()
                for n in oracle[2]
            ),
        }

    ours_error, reference_error = errors(ours), errors(reference)
    step = 2**-8 if amp is not None else 2**-20
    scale = oracle[0][valid].abs().max().item()
    allowed = {
        "logits_max_abs": 1.5 * reference_error["logits_max_abs"] + step * scale,
        "hidden_grad_rel_l2": 1.5 * reference_error["hidden_grad_rel_l2"] + step,
        "param_grad_rel_l2": 1.5 * reference_error["param_grad_rel_l2"] + step,
    }
    return {
        "batched_vs_oracle": ours_error,
        "per_record_vs_oracle": reference_error,
        "argmax_equal": bool(torch.equal(ours[0].argmax(-1), reference[0].argmax(-1))),
        "verdict": "PASS" if all(ours_error[k] <= allowed[k] for k in allowed) else "FAIL",
    }


def head_cost(model, batch, amp):
    # Head forward + backward alone on detached hidden states: time and memory above the inputs.
    backbone = model._backbone()
    text_model = getattr(backbone.model, "language_model", backbone.model)
    with torch.no_grad():
        hidden = text_model(
            input_ids = batch["input_ids"], attention_mask = batch["attention_mask"], use_cache = False
        ).last_hidden_state
    embedding = backbone.get_output_embeddings().weight.detach()
    times = []
    for i in range(args.warmup + 3):
        torch.cuda.synchronize()
        torch.cuda.reset_peak_memory_stats()
        base = torch.cuda.memory_allocated()
        start = time.perf_counter()
        h = hidden.detach().requires_grad_(True)
        with torch.autocast("cuda", dtype = amp or torch.float32, enabled = amp is not None):
            out = model.head(
                h, batch["input_ids"], batch["attention_mask"], batch["records"], embedding
            )
        logits = getattr(out, "flat_padded", None)
        if logits is None:
            logits = torch.cat([q.sum().reshape(1) for record in out for q in record])
        logits.float().sum().backward()
        torch.cuda.synchronize()
        if i >= args.warmup:
            times.append(time.perf_counter() - start)
        peak = torch.cuda.max_memory_allocated() - base
    model.head.zero_grad()
    times.sort()
    return {"head_ms": 1000 * times[len(times) // 2], "head_peak_gb": peak / 2**30}


def _graphs():
    from torch._dynamo.utils import counters
    return counters["stats"]["unique_graphs"]


def main():
    model_name = tiny_checkpoint() if args.model == "tiny" else args.model
    dtype = torch.float16 if args.dtype == "fp16" else torch.bfloat16
    model, processor = FastDecisionModel.from_pretrained(
        model_name, max_seq_length = args.seq_len, load_in_4bit = args.load_in_4bit, dtype = dtype
    )
    model = FastDecisionModel.get_peft_model(model, r = 16, lora_alpha = 16)
    count = args.batch_size * (args.steps + args.warmup)
    items, report = FastDecisionModel.build_dataset(
        rows(processor.tokenizer, count), processor, model
    )
    if not items:
        sys.exit(f"no items: {report}")
    tokenizer = getattr(processor, "tokenizer", processor)
    collate = decision.ClefDataCollator(tokenizer.pad_token_id)
    device = next(model.parameters()).device
    batches = []
    for start in range(0, len(items), args.batch_size):
        batch = collate(items[start : start + args.batch_size])
        batches.append({k: v.to(device) if torch.is_tensor(v) else v for k, v in batch.items()})
    model.train()
    amp = torch.bfloat16 if dtype == torch.bfloat16 and torch.cuda.is_bf16_supported() else None
    parity = head_parity(model, batches[0], amp)
    cost = head_cost(model, batches[0], amp)
    optimizer = torch.optim.AdamW([p for p in model.parameters() if p.requires_grad], lr = 1e-5)
    times, losses, all_times, graphs_after_step = [], [], [], []
    torch.cuda.reset_peak_memory_stats()
    for index, batch in enumerate(batches[: args.steps + args.warmup]):
        batch = dict(batch)
        target = batch.pop("target")
        torch.cuda.synchronize()
        start = time.perf_counter()
        with torch.autocast("cuda", dtype = amp or torch.float32, enabled = amp is not None):
            logits, _ = model(**batch)
        loss = decision._soft_cross_entropy(logits, target, batch["marker_mask"])
        loss.backward()
        optimizer.step()
        optimizer.zero_grad(set_to_none = True)
        torch.cuda.synchronize()
        all_times.append(round(time.perf_counter() - start, 4))
        graphs_after_step.append(_graphs())
        if index >= args.warmup:
            times.append(all_times[-1])
        losses.append(round(loss.item(), 5))
    times.sort()
    step = times[len(times) // 2]
    from torch._dynamo.utils import compile_times, counters

    names, values = compile_times(repr = "csv", aggregate = True)
    compiled_frames = dict(zip(names, values))
    print(
        "JOB_RESULT "
        + json.dumps(
            {
                "model": args.model,
                "load_in_4bit": args.load_in_4bit,
                "fast": args.fast,
                "dtype": args.dtype,
                "seq_len": args.seq_len,
                "tokens_per_row": len(items[0]["input_ids"]),
                "batch_size": args.batch_size,
                "s_per_step": round(step, 4),
                "s_per_step_spread": [round(times[0], 4), round(times[-1], 4)],
                "peak_alloc_gb": round(torch.cuda.max_memory_allocated() / 2**30, 3),
                "peak_reserved_gb": round(torch.cuda.max_memory_reserved() / 2**30, 3),
                "head_ms": round(cost["head_ms"], 2),
                "head_share": round(cost["head_ms"] / 1000 / step, 4),
                "head_peak_gb": round(cost["head_peak_gb"], 3),
                "losses": losses,
                "step_times": all_times,
                "graphs_after_step": graphs_after_step,
                "graph_breaks": sum(counters["graph_break"].values()),
                "graph_break_reasons": list(counters["graph_break"])[:10],
                "unique_graphs": counters["stats"]["unique_graphs"],
                "compile_frames": compiled_frames,
                "parity": parity,
                "levers": {
                    "batched_head": os.environ.get("UNSLOTH_CLEF_FAST") != "0",
                    "chunked_norm_pool": os.environ.get("UNSLOTH_CLEF_FAST") != "0",
                    "head_checkpoint": os.environ.get("UNSLOTH_CLEF_FAST") != "0"
                    and os.environ.get("UNSLOTH_CLEF_CHECKPOINT", "1") != "0",
                    "head_compile": os.environ.get("UNSLOTH_CLEF_FAST") != "0"
                    and clef._compile_supported(device)
                    and "function" in clef._COMPILED,
                },
                "gpu": torch.cuda.get_device_name(),
                "torch": torch.__version__,
                "unsloth_file": sys.modules["unsloth"].__file__,
            }
        ),
        flush = True,
    )


if __name__ == "__main__":
    main()