"""Portable eager inference for the released V2 mean-pooling model.

The CLI used to ship a pair of TorchScript graphs (``encoder.ts`` / ``classifier.ts``).
TorchScript is a *code* container - it stores serialised Python plus pickle that runs on
load - so third-party model scanners flag it as potentially executable no matter what is
inside (Protect AI Guardian, ``PAIT-TCHST-301``). This module rebuilds the same two halves
from the code-free ``model.safetensors`` and ``config.json`` instead:

* **encoder** - the DINOv3 + MoE trunk cast to bfloat16, exactly the dtype and weights the
  former ``encoder.ts`` carried, so the embeddings are bit-identical (verified at
  ``0.000e+00`` max-abs against the old graphs);
* **classifier** - the fp32 Sub-Centre ArcFace and taxonomic heads.

The public surface mirrors the old graphs so callers only swap ``torch.jit.load`` for
:func:`load_bundle`.

Built with DINOv3.
"""
from __future__ import annotations

import copy
import json
from pathlib import Path

import torch
from torch import nn

from .plant_unified import PlantUnified


class _Encoder(nn.Module):
    """The DINOv3 + MoE trunk and feature fusion, without the training-only aux losses.

    Mirrors the module the removed ``encoder.ts`` wrapped: it runs the stem, one MoE layer
    and the fusion, and skips the load-balance statistics the training forward computes
    (they do not affect the embedding). Keeping the op set identical to the old graph also
    keeps the cost identical rather than paying for eager-only bookkeeping.
    """

    def __init__(self, model: PlantUnified):
        super().__init__()
        self.stem = model.stem
        self.norm = model.moe_norms[0]
        self.moe = model.moe_layers[0]
        self.fusion = model.fusion
        self.n_prefix = model.n_prefix

    def forward(self, image: torch.Tensor) -> torch.Tensor:
        maps = self.stem.forward_intermediates(
            image, indices=[7, 11], output_fmt="NLC", norm=True,
            return_prefix_tokens=True, intermediates_only=True)
        mid, deep = [torch.cat((prefix, spatial), dim=1) for spatial, prefix in maps]
        routed = self.moe(self.norm(deep), ctx=deep[:, 0])[0]
        deep = deep + routed
        return self.fusion(mid[:, self.n_prefix:], deep[:, self.n_prefix:]).float()


class V2Bundle:
    """The encoder/classifier pair the CLI and harness scripts call.

    ``encoder`` takes the same bfloat16 image batch the graphs took and returns a float32
    embedding; ``classifier`` takes a float32 embedding and returns the
    ``(species, genus, family)`` logits. Keeping the exact signatures means the call sites
    do not change.
    """

    def __init__(self, model_fp32: PlantUnified, encoder_fp32: PlantUnified):
        self.model = model_fp32
        # bf16 encoder: the encoder graph stored bf16 weights, so this reproduces it
        # bit-for-bit while the fp32 model keeps the precision-sensitive ArcFace head.
        self._encoder = _Encoder(encoder_fp32).to(torch.bfloat16).eval()

    @torch.inference_mode()
    def encoder(self, images: torch.Tensor) -> torch.Tensor:
        return self._encoder(images)

    @torch.inference_mode()
    def classifier(self, embedding: torch.Tensor):
        logits = self.model.classify_embeddings(embedding)
        return logits["species"], logits["genus"], logits["family"]


def build_model(config: dict) -> PlantUnified:
    """Instantiate the architecture from a released ``config.json`` (no pretrained load)."""
    model_config = dict(config["model_config"])
    model_config.pop("pretrained", None)
    return PlantUnified(pretrained=False, **model_config)


def load_bundle(directory, device="cpu") -> V2Bundle:
    """Build the inference bundle from a release directory containing ``config.json``.

    ``model.safetensors`` is loaded with ``strict=True``; the shared ``moe_layer`` alias in
    the released file is resolved by ``safetensors.torch.load_model``.
    """
    from safetensors.torch import load_model

    directory = Path(directory)
    config = json.loads((directory / "config.json").read_text(encoding="utf-8"))
    model = build_model(config)
    load_model(model, str(directory / "model.safetensors"), strict=True, device=str(device))
    model = model.to(device).eval()
    return V2Bundle(model, copy.deepcopy(model))