from __future__ import annotations

import argparse
import json
import shutil
from pathlib import Path

import coremltools as ct
import numpy as np
import timm
import torch
from PIL import Image

DEFAULT_MODEL = "vit_base_patch16_dinov3.lvd1689m"
BASE_MODEL = "facebook/dinov3-vitb16-pretrain-lvd1689m"
IMAGENET_MEAN = (0.485, 0.456, 0.406)
IMAGENET_STD = (0.229, 0.224, 0.225)


class DINOv3Encoder(torch.nn.Module):
    def __init__(self, backbone: torch.nn.Module) -> None:
        super().__init__()
        self.backbone = backbone
        self.num_prefix_tokens = backbone.num_prefix_tokens
        self.register_buffer("mean", torch.tensor(IMAGENET_MEAN).view(1, 3, 1, 1))
        self.register_buffer("std", torch.tensor(IMAGENET_STD).view(1, 3, 1, 1))

    def forward(self, image: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        pixels = (image / 255.0 - self.mean) / self.std
        features = self.backbone.forward_features(pixels)
        cls = features[:, 0]
        norm = torch.sqrt(torch.sum(cls * cls, dim=-1, keepdim=True).clamp_min(1e-12))
        # The prefix is one CLS token plus four register tokens. Registers absorb
        # high-norm artifacts that would otherwise pollute the patch tokens; they
        # are not features and nothing downstream uses them.
        patches = features[:, self.num_prefix_tokens:]
        # Patches stay unnormalized because dense heads generally want the
        # magnitude. Callers doing cosine can normalize per token themselves.
        return cls / norm, patches


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Convert DINOv3 to a Core ML model returning one CLS embedding."
    )
    parser.add_argument("--size", type=int, default=448)
    parser.add_argument("--precision", choices=("fp16", "fp32"), default="fp32")
    parser.add_argument("--output", type=Path)
    parser.add_argument("--force", action="store_true")
    return parser.parse_args()


def synthetic_image(size: int) -> Image.Image:
    y, x = np.mgrid[0:size, 0:size]
    rgb = np.stack(
        ((x * 255 // size), (y * 255 // size), ((x + y) * 255 // (size * 2))),
        axis=-1,
    ).astype(np.uint8)
    return Image.fromarray(rgb, mode="RGB")


def main() -> None:
    args = parse_args()
    if args.size <= 0 or args.size % 16:
        raise SystemExit("--size must be a positive multiple of 16")
    if args.output is None:
        precision_name = args.precision.upper()
        args.output = Path(f"models/DINOv3ViTB16-{precision_name}-{args.size}.mlpackage")
    if args.output.exists():
        if not args.force:
            raise SystemExit(f"{args.output} already exists; pass --force to replace it")
        shutil.rmtree(args.output)

    example = torch.zeros(1, 3, args.size, args.size, dtype=torch.float32)
    backbone = timm.create_model(
        DEFAULT_MODEL,
        pretrained=True,
        img_size=args.size,
        num_classes=0,
    ).eval()
    model = DINOv3Encoder(backbone).eval()

    with torch.inference_mode():
        exported = torch.export.export(model, (example,)).run_decompositions({})

    precision = ct.precision.FLOAT16 if args.precision == "fp16" else ct.precision.FLOAT32
    coreml_model = ct.convert(
        exported,
        convert_to="mlprogram",
        inputs=[
            ct.ImageType(
                name="image",
                shape=example.shape,
                color_layout=ct.colorlayout.RGB,
            )
        ],
        outputs=[ct.TensorType(name="embedding"), ct.TensorType(name="patch_embeddings")],
        minimum_deployment_target=ct.target.macOS14,
        compute_precision=precision,
    )
    coreml_model.author = "dinov3-coreml; base model by Meta"
    coreml_model.license = "DINOv3 License"
    coreml_model.short_description = "DINOv3 ViT-B/16 CLS embedding and patch tokens"
    coreml_model.user_defined_metadata["base_model"] = BASE_MODEL
    coreml_model.input_description["image"] = f"RGB image resized to {args.size}x{args.size}"
    coreml_model.output_description["embedding"] = "L2-normalized 768-value CLS embedding"
    coreml_model.output_description["patch_embeddings"] = "Unnormalized patch tokens, one per 16x16 patch"

    args.output.parent.mkdir(parents=True, exist_ok=True)
    coreml_model.save(args.output)

    image = synthetic_image(args.size)
    array = np.asarray(image, dtype=np.float32).transpose(2, 0, 1)[None, ...]
    with torch.inference_mode():
        torch_cls, torch_patches = (t.numpy()[0] for t in model(torch.from_numpy(array)))
    prediction = coreml_model.predict({"image": image})
    coreml_output = np.asarray(prediction["embedding"])[0]
    coreml_patches = np.asarray(prediction["patch_embeddings"])[0]

    def cosine_of(a: np.ndarray, b: np.ndarray) -> float:
        a, b = a.reshape(-1), b.reshape(-1)
        return float(np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b)))

    cosine = cosine_of(torch_cls, coreml_output)
    patch_cosine = cosine_of(torch_patches, coreml_patches)
    report = {
        "model": DEFAULT_MODEL,
        "base_model": BASE_MODEL,
        "input_size": args.size,
        "precision": args.precision,
        "output_shape": list(coreml_output.shape),
        "patch_output_shape": list(coreml_patches.shape),
        "pytorch_coreml_cosine_similarity": cosine,
        "pytorch_coreml_patch_cosine_similarity": patch_cosine,
        "coreml_output_l2_norm": float(np.linalg.norm(coreml_output)),
    }
    report_path = args.output.with_suffix(".validation.json")
    report_path.write_text(json.dumps(report, indent=2) + "\n")
    print(json.dumps({"output": str(args.output), **report}, indent=2))
    if min(cosine, patch_cosine) < 0.999:
        raise SystemExit("Core ML parity check failed: cosine similarity is below 0.999")


if __name__ == "__main__":
    main()