# examples/detr_export.py
"""Genuine original DETR R50 inference using authenticated pinned native source.

Unzip detr_native.zip beside this file, then run:
python detr_export.py --checkpoint detr-r50-e632da11.pth --output detr-demo
The downloaded source archive and checkpoint are hash-verified. No SciPy matcher
or training package is imported. See README-detr.md for sources and licenses.
"""

import argparse
import ast
import hashlib
import importlib.util
import json
import math
import platform
import shutil
from pathlib import Path
from typing import Dict, List
from urllib.request import urlopen

import numpy as np
import PIL
import torch
import torchvision
from PIL import Image, ImageDraw
from torch import nn
from torch.nn import functional as F
from torchvision.models._utils import IntermediateLayerGetter
from torchvision.transforms import functional as TF

import detr_training

REVISION = "29901c51d7fe8712168b8d0d64351170bc0f83e0"
CHECKPOINT_URL = "https://dl.fbaipublicfiles.com/detr/detr-r50-e632da11.pth"
CHECKPOINT_SHA256 = (
    "e632da11ec76ae67bac2f8579fbed3724e08dead7d200ca13e019b197784eadc"
)
PHOTOS = [
    {
        "id": "cat",
        "title": "Chelsea the cat",
        "author": "Stefan van der Walt",
        "filename": "chelsea.png",
        "sha256": "596aa1e7cb875eb79f437e310381d26b338a81c2da23439704a73c4651e8c4bb",
    },
    {
        "id": "coffee",
        "title": "Cup of coffee",
        "author": "Rachel Michetti",
        "filename": "coffee.png",
        "sha256": "cc02f8ca188b167c775a7101b5d767d1e71792cf762c33d6fa15a4599b5a8de7",
    },
]
PHOTO_ROOT = "https://raw.githubusercontent.com/scikit-image/scikit-image/v0.25.2/skimage/data/"
PHOTO_LICENSE = "https://github.com/scikit-image/scikit-image/blob/v0.25.2/skimage/data/_fetchers.py"


def sha256(path):
    """Hash complete file bytes."""
    return hashlib.sha256(Path(path).read_bytes()).hexdigest()


class NestedTensor:
    """Inference-only tensor/mask container matching the native interface."""

    def __init__(self, tensors, mask):
        self.tensors, self.mask = tensors, mask

    def decompose(self):
        """Return the batch and padding mask."""
        return self.tensors, self.mask


def native_definitions(source_dir):
    """Authenticate and load unchanged model definitions without training imports.

    AST selection skips module imports of SciPy matcher and optional panoptic
    packages. Executed class/function bodies remain exactly the pinned source.
    The local container replaces only util.misc's tensor/mask wrapper.
    """
    manifest = json.loads((source_dir / "source-manifest.json").read_text())
    if manifest["revision"] != REVISION:
        raise ValueError("Unexpected native source revision.")
    expected = json.loads(
        (Path(__file__).with_name("detr-source-hashes.json")).read_text()
    )
    if manifest["files"] != expected:
        raise ValueError(
            "Source manifest differs from pinned public hash record."
        )
    for filename, record in expected.items():
        if sha256(source_dir / filename) != record["sha256"]:
            raise ValueError(f"Native source hash mismatch: {filename}")
    namespace = {
        "__name__": "detr_native_definitions",
        "torch": torch,
        "nn": nn,
        "F": F,
        "math": math,
        "NestedTensor": NestedTensor,
        "Dict": Dict,
        "List": List,
        "IntermediateLayerGetter": IntermediateLayerGetter,
    }
    for filename, names in [
        ("models/backbone.py", {"FrozenBatchNorm2d", "BackboneBase", "Joiner"}),
        ("models/position_encoding.py", {"PositionEmbeddingSine"}),
        ("models/detr.py", {"DETR", "MLP", "PostProcess"}),
    ]:
        tree = ast.parse((source_dir / filename).read_text())
        tree.body = [
            node
            for node in tree.body
            if isinstance(node, ast.ClassDef) and node.name in names
        ]
        exec(compile(tree, str(source_dir / filename), "exec"), namespace)
    for filename, key in [
        ("models/transformer.py", "transformer"),
        ("util/box_ops.py", "box_ops"),
    ]:
        spec = importlib.util.spec_from_file_location(
            f"detr_native_{key}", source_dir / filename
        )
        module = importlib.util.module_from_spec(spec)
        spec.loader.exec_module(module)
        namespace[key] = module
    return namespace, manifest


def build_model(checkpoint, source_dir):
    """Load all original R50 detector tensors strictly, without backbone downloads."""
    if sha256(checkpoint) != CHECKPOINT_SHA256:
        raise ValueError("Official DETR checkpoint SHA256 mismatch.")
    native, manifest = native_definitions(source_dir)
    body = torchvision.models.resnet50(
        weights=None, norm_layer=native["FrozenBatchNorm2d"]
    )
    backbone = native["BackboneBase"](body, False, 2048, False)
    position = native["PositionEmbeddingSine"](128, normalize=True)
    joined = native["Joiner"](backbone, position)
    joined.num_channels = 2048
    transformer = native["transformer"].Transformer(
        d_model=256,
        nhead=8,
        num_encoder_layers=6,
        num_decoder_layers=6,
        dim_feedforward=2048,
        dropout=0.1,
        return_intermediate_dec=True,
    )
    model = native["DETR"](joined, transformer, num_classes=91, num_queries=100)
    state = torch.load(checkpoint, map_location="cpu", weights_only=True)[
        "model"
    ]
    receipt = model.load_state_dict(state, strict=True)
    return model.eval(), native["PostProcess"](), manifest, receipt


def preprocess(image):
    """Official evaluation geometry: short side 800, long side at most 1333.

    Match datasets/transforms.py integer rounding: cap the requested short side
    with round, then floor the other dimension via int. Pillow bilinear RGB.
    """
    width, height = image.size
    size = 800
    if max(width, height) / min(width, height) * size > 1333:
        size = int(round(1333 * min(width, height) / max(width, height)))
    if width < height:
        new_width, new_height = size, int(size * height / width)
    else:
        new_height, new_width = size, int(size * width / height)
    resized = image.resize((new_width, new_height), Image.Resampling.BILINEAR)
    tensor = TF.normalize(
        TF.to_tensor(resized), [0.485, 0.456, 0.406], [0.229, 0.224, 0.225]
    )
    return resized, tensor


def export(checkpoint, output, source_dir, photo_dir=None):
    """Save exact inputs, every query output, native postprocess parity and hashes."""
    torch.set_num_threads(2)
    torch.manual_seed(7)
    output.mkdir(parents=True, exist_ok=True)
    model, native_postprocess, manifest, receipt = build_model(
        checkpoint, source_dir
    )
    records = []
    for photo in PHOTOS:
        original_path = output / f"source-{photo['id']}.png"
        if photo_dir:
            shutil.copyfile(photo_dir / original_path.name, original_path)
        elif not original_path.exists():
            with urlopen(
                PHOTO_ROOT + photo["filename"], timeout=60
            ) as response:
                original_path.write_bytes(response.read())
        if sha256(original_path) != photo["sha256"]:
            raise ValueError("Photo hash mismatch.")
        image = Image.open(original_path).convert("RGB")
        resized, tensor = preprocess(image)
        processed_path = output / f"input-{photo['id']}.png"
        resized.save(processed_path)
        batch, mask = detr_training.padded_batch([tensor])
        original_size = [[image.height, image.width]]
        with torch.no_grad():
            outputs = model(NestedTensor(batch, mask))
            native = native_postprocess(outputs, torch.tensor(original_size))[0]
            ours = detr_training.postprocess(outputs, original_size)
        errors = {
            key: float((ours[key][0] - native[key]).abs().max())
            for key in ("scores", "labels", "boxes")
        }
        if any(errors.values()):
            raise RuntimeError("Native postprocessing parity failed.")
        arrays_path = output / f"arrays-{photo['id']}.npz"
        np.savez_compressed(
            arrays_path,
            input=batch.numpy(),
            padding_mask=mask.numpy(),
            logits=outputs["pred_logits"].numpy(),
            normalized_cxcywh=outputs["pred_boxes"].numpy(),
            scores=native["scores"].numpy(),
            labels=native["labels"].numpy(),
            pixel_xyxy=native["boxes"].numpy(),
        )
        queries = [
            {
                "query": i,
                "classId": int(label),
                "score": float(score),
                "normalizedCxcywh": outputs["pred_boxes"][0, i].tolist(),
                "pixelXyxy": box.tolist(),
            }
            for i, (label, score, box) in enumerate(
                zip(native["labels"], native["scores"], native["boxes"])
            )
        ]
        visible = [q for q in queries if q["score"] > 0.7]
        overlay = image.copy()
        draw = ImageDraw.Draw(overlay)
        for query in visible:
            x0, y0, x1, y1 = query["pixelXyxy"]
            draw.rectangle(
                (
                    max(0, x0),
                    max(0, y0),
                    min(image.width - 1, x1),
                    min(image.height - 1, y1),
                ),
                outline=(20, 184, 166),
                width=3,
            )
            draw.text(
                (max(0, x0) + 3, max(0, y0) + 3),
                f"q{query['query']} class {query['classId']} {query['score']:.3f}",
                fill=(255, 255, 255),
                stroke_width=1,
                stroke_fill=(0, 0, 0),
            )
        overlay_path = output / f"overlay-{photo['id']}.png"
        overlay.save(overlay_path)
        records.append(
            {
                "id": photo["id"],
                "title": photo["title"],
                "originalSize": original_size[0],
                "inputShape": list(batch.shape),
                "source": {
                    "url": PHOTO_ROOT + photo["filename"],
                    "sha256": photo["sha256"],
                    "author": photo["author"],
                    "license": "CC0",
                    "licenseUrl": PHOTO_LICENSE,
                },
                "originalImage": original_path.name,
                "inputImage": processed_path.name,
                "inputSha256": sha256(processed_path),
                "arrays": arrays_path.name,
                "arraysSha256": sha256(arrays_path),
                "overlay": overlay_path.name,
                "overlaySha256": sha256(overlay_path),
                "nativePostprocessMaxAbsErrors": errors,
                "queries": queries,
                "visibleQueries": [q["query"] for q in visible],
            }
        )
    result = {
        "provenance": {
            "architecture": "Original DETR ResNet50, 6 encoder + 6 decoder layers, d=256, 8 heads, 100 queries",
            "sourceRevision": REVISION,
            "sourceManifest": manifest,
            "checkpointUrl": CHECKPOINT_URL,
            "checkpointSha256": CHECKPOINT_SHA256,
            "programSha256": sha256(Path(detr_training.__file__)),
            "exportSha256": sha256(__file__),
            "strictStateLoad": True,
            "missingKeys": receipt.missing_keys,
            "unexpectedKeys": receipt.unexpected_keys,
            "parameterCount": sum(p.numel() for p in model.parameters()),
            "torch": torch.__version__,
            "torchvision": torchvision.__version__,
            "numpy": np.__version__,
            "pillow": PIL.__version__,
            "python": platform.python_version(),
            "device": "cpu",
            "threads": 2,
            "seed": 7,
            "command": "python detr_export.py --checkpoint detr-r50-e632da11.pth --output detr-demo",
            "adapter": "Authenticated unchanged native class bodies selected via AST; tensor/mask container; torchvision resnet50 weights=None; no training-only imports.",
        },
        "preprocessing": {
            "mode": "RGB",
            "shortSide": 800,
            "maxLongSide": 1333,
            "interpolation": "Pillow bilinear",
            "mean": [0.485, 0.456, 0.406],
            "std": [0.229, 0.224, 0.225],
            "padding": "bottom/right zeros in normalized tensor; batch of one has all-false mask",
            "boxPolicy": "Raw cxcywh and native pixel xyxy retained unchanged; overlay alone clips to image boundaries.",
            "displayThreshold": 0.7,
            "nms": False,
        },
        "photos": records,
    }
    (output / "predictions.json").write_text(
        json.dumps(result, indent=2) + "\n"
    )
    print(
        json.dumps(
            {
                "photos": [
                    {
                        "id": p["id"],
                        "visible": p["visibleQueries"],
                        "parity": p["nativePostprocessMaxAbsErrors"],
                    }
                    for p in records
                ]
            },
            indent=2,
        )
    )


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--checkpoint", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument(
        "--source-dir",
        type=Path,
        default=Path(__file__).with_name("detr_native"),
    )
    parser.add_argument("--photo-dir", type=Path)
    args = parser.parse_args()
    export(args.checkpoint, args.output, args.source_dir, args.photo_dir)