#!/usr/bin/env python3
"""D-FINE small (COCO, Apache-2.0) -> ExecuTorch WebGPU, for the playground's current forward runtime.

Upstream: HF transformers port `ustc-community/dfine-small-coco` (Peterande/D-FINE, D-FINE-S). Every tensor in
it is bit-identical to the authors' `dfine_s_coco.pth` on `Peterande/D-FINE` (see verify_official.py).

Input:  [1, 3, 640, 640] float32 RGB in [0, 1] (image / 255; the HF processor does no mean/std).
        Upstream resizes to 640x640 without keeping aspect; a letterboxed input also works.
Output: [1, 84, 300] float32, the YOLO layout src/vision/detect.js reads: rows 0-3 are
        cx, cy, w, h in input pixels, rows 4-83 are sigmoid class probabilities (80 COCO classes,
        contiguous ids, same order as the YOLO head). 300 queries; no NMS needed (DETR).

The eager graph is rebuilt from the HF modules with output-preserving rewrites so that the
whole program is one delegate made of operators the current runtime registers
(see _detector.py): topk -> TopKSet (same query set; queries come out in anchor order, and the
decoder is permutation equivariant), grid_sample -> gathers, inverse_sigmoid/log -> identity,
LQE top-4 -> rank-based selection, max over classes -> exact where-tree, SiLU -> x * sigmoid(x),
frozen BN -> precomputed scale/shift, position embedding and anchors -> constants.
"""

import argparse
import copy
import json
import sys
from pathlib import Path

import numpy as np
import torch
import torch.nn.functional as F

sys.path.insert(0, str(Path(__file__).resolve().parent))
import _common  # noqa: E402
import _detector as D  # noqa: E402

UPSTREAM = {
    "repo": "ustc-community/dfine-small-coco",
    "revision": "f79e65b5fbb33ceb9d3ebba042955d7410c608f8",
    "files": {
        "model.safetensors": "a144baff3a27bfd538fe37b11fa8d4fca5744b8c054b5301e7c699a1379cd18f",
    },
    "license": "Apache-2.0",
}
OFFICIAL = {
    "repo": "Peterande/D-FINE",
    "revision": "19214e91f5ce4c1e609022c38b1fc953c47404e7",
    "file": "dfine_s_coco.pth",
    "sha256": "48a6c8cc43eb57186843f752e2e8461ddd3326e0d3c575e71e6e960844683e89",
    "result": "all 764 tensors bit-identical (verify_official.py)",
}
SIZE = 640


def load_hf():
    from transformers import DFineForObjectDetection

    path, digest = _common.hf_file(
        UPSTREAM["repo"],
        "model.safetensors",
        UPSTREAM["revision"],
        UPSTREAM["files"]["model.safetensors"],
    )
    m = DFineForObjectDetection.from_pretrained(
        UPSTREAM["repo"], revision=UPSTREAM["revision"], attn_implementation="eager"
    ).eval()
    return m, digest


class DFineStatic(torch.nn.Module):
    def __init__(self, hf, size=SIZE):
        super().__init__()
        cfg = hf.config
        self.cfg = cfg
        mdl = hf.model
        self.backbone = mdl.backbone.model
        from transformers.models.d_fine.modeling_d_fine import DFineFrozenBatchNorm2d

        D.replace_modules(
            self.backbone,
            lambda c: isinstance(c, DFineFrozenBatchNorm2d),
            lambda c: D.FoldedBatchNorm(
                c.weight, c.bias, c.running_mean, c.running_var, 1e-5
            ),
        )
        self.encoder_input_proj = mdl.encoder_input_proj
        self.encoder = mdl.encoder
        D.replace_modules(
            self,
            lambda c: type(c).__name__ in ("SiLUActivation", "SiLU"),
            lambda c: D.SiLU(),
        )
        self.decoder_input_proj = mdl.decoder_input_proj
        self.enc_output = mdl.enc_output
        self.enc_score_head = mdl.enc_score_head
        self.enc_bbox_head = mdl.enc_bbox_head
        self.decoder = mdl.decoder
        self.class_embed = hf.class_embed
        self.bbox_embed = hf.bbox_embed

        strides = cfg.feat_strides
        self.shapes = [(size // s, size // s) for s in strides]
        for _ in range(cfg.num_feature_levels - len(strides)):
            h, w = self.shapes[-1]
            self.shapes.append(((h + 1) // 2, (w + 1) // 2))
        n = sum(h * w for h, w in self.shapes)
        anchors, valid = mdl.generate_anchors(
            tuple(self.shapes), device="cpu", dtype=torch.float32
        )
        self.register_buffer("anchors", anchors.clone(), persistent=False)
        self.register_buffer("valid", valid.to(torch.float32).clone(), persistent=False)
        aifi_hw = self.shapes[cfg.encode_proj_layers[0]]
        from transformers.models.d_fine.modeling_d_fine import (
            build_2d_sinusoidal_position_embedding,
        )

        pos = build_2d_sinusoidal_position_embedding(
            aifi_hw[0],
            aifi_hw[1],
            cfg.encoder_hidden_dim,
            cfg.positional_encoding_temperature,
        )
        self.register_buffer(
            "aifi_pos", pos.unsqueeze(0).to(torch.float32), persistent=False
        )
        from transformers.models.d_fine.modeling_d_fine import weighting_function

        with torch.no_grad():
            project = weighting_function(
                cfg.max_num_bins, self.decoder.up, self.decoder.reg_scale
            )
        self.register_buffer(
            "project", project.reshape(1, -1).clone(), persistent=False
        )
        self.reg_scale = float(self.decoder.reg_scale.item())
        self.topk = D.TopKSet(n, cfg.num_queries)
        self.sampler = D.DeformSampler(
            self.shapes,
            list(self.decoder.layers[0].encoder_attn.num_points_list),
            cfg.decoder_attention_heads,
            cfg.d_model // cfg.decoder_attention_heads,
        )
        self.register_buffer(
            "pix_scale",
            torch.tensor([size, size, size, size], dtype=torch.float32),
            persistent=False,
        )
        self.register_buffer("one", torch.ones(1), persistent=False)
        self.register_buffer("two", torch.full((1,), 2.0), persistent=False)
        self.register_buffer("zero", torch.zeros(1), persistent=False)

    # ---- encoder -------------------------------------------------------------------------
    def aifi(self, x):
        cfg = self.cfg
        b, c, h, w = x.shape
        hs = x.flatten(2).permute(0, 2, 1)
        for layer in self.encoder.aifi[0].layers:
            a = layer.self_attn
            att = D.attention(
                a.q_proj,
                a.k_proj,
                a.v_proj,
                a.o_proj,
                hs,
                self.aifi_pos,
                cfg.encoder_attention_heads,
            )
            hs = layer.self_attn_layer_norm(hs + att)
            hs = layer.final_layer_norm(hs + layer.mlp(hs))
        return hs.permute(0, 2, 1).reshape(b, c, h, w)

    def run_encoder(self, feats):
        enc = self.encoder
        feats = list(feats)
        idx = self.cfg.encode_proj_layers[0]
        feats[idx] = self.aifi(feats[idx])
        fpn = [feats[-1]]
        for i, (lat, blk) in enumerate(zip(enc.lateral_convs, enc.fpn_blocks)):
            back = feats[enc.num_fpn_stages - i - 1]
            top = lat(fpn[-1])
            fpn[-1] = top
            top = F.interpolate(top, scale_factor=2.0, mode="nearest")
            fpn.append(blk(torch.cat([top, back], dim=1)))
        fpn.reverse()
        pan = [fpn[0]]
        for i, (down, blk) in enumerate(zip(enc.downsample_convs, enc.pan_blocks)):
            pan.append(blk(torch.cat([down(pan[-1]), fpn[i + 1]], dim=1)))
        return pan

    # ---- decoder -------------------------------------------------------------------------
    def deform(self, attn_mod, hidden, ref, value):
        """hidden [1, Q, C]; ref [1, Q, 4]; value [L, heads, hd]."""
        cfg = self.cfg
        q = hidden.shape[1]
        nh = cfg.decoder_attention_heads
        P = sum(attn_mod.num_points_list)
        off = attn_mod.sampling_offsets(hidden).view(q, nh, P, 2)
        aw = torch.softmax(attn_mod.attention_weights(hidden).view(q, nh, P), dim=-1)
        scale = attn_mod.num_points_scale.view(1, 1, P, 1)
        r = ref.view(q, 1, 1, 4)
        loc = r[..., :2] + off * scale * r[..., 2:] * attn_mod.offset_scale
        out = self.sampler(value, loc, aw).view(1, q, cfg.d_model)
        return out

    def decoder_layer(self, layer, hs, pos, ref, value):
        cfg = self.cfg
        a = layer.self_attn
        att = D.attention(
            a.q_proj, a.k_proj, a.v_proj, a.o_proj, hs, pos, cfg.decoder_attention_heads
        )
        hs = layer.self_attn_layer_norm(hs + att)
        cross = self.deform(layer.encoder_attn, hs + pos, ref, value)
        hs = layer.gateway(hs, cross)
        hs = hs + layer.mlp(hs)
        return layer.final_layer_norm(D.clamp(hs, -65504.0, 65504.0))

    def integral(self, corners):
        q = corners.shape[1]
        p = torch.softmax(corners.reshape(q * 4, self.cfg.max_num_bins + 1), dim=-1)
        return F.linear(p, self.project).reshape(1, q, 4)

    def distance2bbox(self, pts, dist):
        rs = self.reg_scale
        cx, cy, w, h = pts[..., 0:1], pts[..., 1:2], pts[..., 2:3], pts[..., 3:4]
        x0 = cx - (0.5 * rs + dist[..., 0:1]) * (w / rs)
        y0 = cy - (0.5 * rs + dist[..., 1:2]) * (h / rs)
        x1 = cx + (0.5 * rs + dist[..., 2:3]) * (w / rs)
        y1 = cy + (0.5 * rs + dist[..., 3:4]) * (h / rs)
        return torch.cat([(x0 + x1) / 2, (y0 + y1) / 2, x1 - x0, y1 - y0], dim=-1)

    def lqe(self, lqe_mod, scores, corners):
        q = corners.shape[1]
        nb = self.cfg.max_num_bins + 1
        k = self.cfg.top_prob_values
        prob = torch.softmax(corners.reshape(q * 4, nb), dim=-1)  # [R, nb]
        diff = prob.view(q * 4, 1, nb) - prob.view(
            q * 4, nb, 1
        )  # [R, j, l] = p_l - p_j
        greater = torch.where(diff > 0, self.one, self.zero).sum(dim=-1)  # [R, nb]
        tops = []
        for i in range(k):
            masked = torch.where(greater <= i, prob, self.two)
            tops.append(D.tree_reduce(masked, D.where_min))  # i-th largest value
        top = torch.cat(tops, dim=-1)
        stat = torch.cat([top, top.mean(dim=-1, keepdim=True)], dim=-1).reshape(
            1, q, 4 * (k + 1)
        )
        return scores + lqe_mod.reg_conf(stat)

    def forward(self, pixel_values):
        cfg = self.cfg
        feats = self.backbone(pixel_values).feature_maps
        proj = [self.encoder_input_proj[i](f) for i, f in enumerate(feats)]
        enc = self.run_encoder(proj)
        sources = [self.decoder_input_proj[i](s) for i, s in enumerate(enc)]
        for i in range(len(sources), cfg.num_feature_levels):
            sources.append(self.decoder_input_proj[i](enc[-1]))
        flat = torch.cat(
            [s.flatten(2).permute(0, 2, 1) for s in sources], dim=1
        )  # [1, L, C]
        L = flat.shape[1]

        memory = self.valid * flat
        out_mem = self.enc_output(memory)
        cls = self.enc_score_head(out_mem)  # [1, L, 80]
        coord = self.enc_bbox_head(out_mem) + self.anchors  # [1, L, 4]
        best = D.tree_reduce(cls, D.where_max).view(1, L)
        idx = self.topk(best)  # [Q] float, ascending anchor order
        Q = cfg.num_queries
        idx_c = idx.view(1, Q, 1).expand(1, Q, cfg.d_model).to(torch.int64)
        target = torch.gather(out_mem, 1, idx_c)
        idx_4 = idx.view(1, Q, 1).expand(1, Q, 4).to(torch.int64)
        ref_unact = torch.gather(coord, 1, idx_4)

        nh = cfg.decoder_attention_heads
        value = flat.view(L, nh, cfg.d_model // nh)
        dec = self.decoder
        hs = target
        ref = torch.sigmoid(ref_unact)
        out_detach = None
        corners_prev = None
        ref_initial = None
        for i, layer in enumerate(dec.layers):
            pos = D.clamp(dec.query_pos_head(ref), -10.0, 10.0)
            hs = self.decoder_layer(layer, hs, pos, ref, value)
            if i == 0:
                ref_initial = D.sigmoid_add_logit(dec.pre_bbox_head(hs), ref)
            if out_detach is None:
                corners = self.bbox_embed[i](hs)
            else:
                corners = self.bbox_embed[i](hs + out_detach) + corners_prev
            box = self.distance2bbox(ref_initial, self.integral(corners))
            corners_prev = corners
            ref = box
            out_detach = hs
        last = len(dec.layers) - 1
        scores = self.class_embed[last](hs)
        scores = self.lqe(dec.lqe_layers[last], scores, corners)
        return D.yolo_layout_buf(box, torch.sigmoid(scores), self.pix_scale)


# ---------------------------------------------------------------------------------------------
def hf_reference(hf, x):
    """HF eager outputs in the same layout, with queries permuted to ascending anchor order."""
    with torch.no_grad():
        o = hf(pixel_values=x)
    order = torch.topk(
        o.enc_outputs_class.max(-1).values, hf.config.num_queries, dim=1
    ).indices[0]
    perm = torch.argsort(order)
    boxes = o.pred_boxes[:, perm]
    probs = torch.sigmoid(o.logits[:, perm])
    return torch.cat([boxes * SIZE, probs], dim=-1).permute(0, 2, 1).numpy(), order


def load_image(path, size=SIZE):
    from PIL import Image

    img = Image.open(path).convert("RGB").resize((size, size), Image.BILINEAR)
    a = np.asarray(img, dtype=np.float32) / 255.0
    return torch.from_numpy(a.transpose(2, 0, 1).copy()).unsqueeze(0)


def top_detections(out, thr=0.4):
    out = np.asarray(out)[0]
    probs = out[4:]
    cls = probs.argmax(0)
    sc = probs.max(0)
    keep = np.argsort(-sc)
    res = []
    for q in keep:
        if sc[q] < thr:
            break
        cx, cy, w, h = out[:4, q]
        res.append(
            (
                int(cls[q]),
                round(float(sc[q]), 3),
                [
                    round(float(v), 1)
                    for v in (cx - w / 2, cy - h / 2, cx + w / 2, cy + h / 2)
                ],
            )
        )
    return res


def main():
    ap = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    ap.add_argument("--out", required=True)
    ap.add_argument(
        "--image",
        required=True,
        help="any RGB photo; the rewritten graph is checked against upstream eager on it",
    )
    args = ap.parse_args()
    out = Path(args.out)
    out.mkdir(parents=True, exist_ok=True)

    hf, digest = load_hf()
    model = DFineStatic(copy.deepcopy(hf)).eval()
    img = load_image(args.image)

    # Output-preservation check of the rewritten eager graph against HF eager, real image.
    ref, order = hf_reference(hf, img)
    with torch.no_grad():
        mine = model(img).numpy()
    parity = {
        "image": Path(args.image).name,
        "rewritten_eager_vs_hf_eager": _common.compare(mine, ref),
    }
    parity["rewritten_eager_vs_hf_eager"]["max_abs_boxes_px"] = float(
        np.abs(mine[:, :4] - ref[:, :4]).max()
    )
    parity["rewritten_eager_vs_hf_eager"]["max_abs_probs"] = float(
        np.abs(mine[:, 4:] - ref[:, 4:]).max()
    )
    parity["detections_hf"] = top_detections(ref)
    parity["detections_rewritten"] = top_detections(mine)
    print(json.dumps(parity, indent=1))

    ep, webgpu, portable = _common.export_vision(
        name="dfine_s_coco",
        model=model,
        example_input=img,
        out_dir=out,
        upstream={
            "model": "D-FINE-S (COCO), HF transformers port",
            "weights": {
                "repo": UPSTREAM["repo"],
                "revision": UPSTREAM["revision"],
                "file": "model.safetensors",
                "sha256": digest,
                "verified_against": OFFICIAL,
            },
            "license": UPSTREAM["license"],
        },
        extra_recipe={
            "preprocessing_outside_graph": "RGB resize to 640x640 (letterbox also works), / 255; no mean/std",
            "postprocessing_outside_graph": "output [1, 84, 300]: cx, cy, w, h (input px) + 80 sigmoid probs; threshold; NMS optional",
            "rewrites": "see recipes/detector-dfine-small.py docstring and recipes/_detector.py",
        },
        golden=(img.numpy(), ref),
    )
    # Real-image check of the exported paths as well.
    D.relabel_golden(
        out,
        "internal_golden_* here = upstream HF transformers eager output on the real image (same layout; DETR queries permuted to anchor order)",
    )
    port = _common.run_pte(portable.buffer, (img,))[0]
    parity["portable_pte_vs_hf_eager"] = _common.compare(port, ref)
    parity["portable_pte_vs_rewritten_eager"] = _common.compare(port, mine)
    parity["detections_portable_pte"] = top_detections(port)
    (out / "hf_parity.json").write_text(json.dumps(parity, indent=1) + "\n")
    np.save(out / "cpu_eager_out.npy", mine)
    img.numpy().astype(np.float32).tofile(out / "input_640.f32")
    print(
        json.dumps(
            {k: v for k, v in parity.items() if not k.startswith("detections")},
            indent=1,
        )
    )


if __name__ == "__main__":
    main()