#!/usr/bin/env python
"""Export a classifier with the input/output contract of a YOLO OpenVINO IR.

The sole active candidate is a full-image pseudo box, not a localization result.
Input must already be RGB float32 BCHW in [0, 1], exactly as in deploy/cpp.
"""

from __future__ import annotations

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

import numpy as np
import torch
import yaml
from torch import nn

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from models import build_model  # noqa: E402


def ordered_names(names) -> list[str]:
    if isinstance(names, dict):
        indexed = {int(k): v for k, v in names.items()}
        if set(indexed) != set(range(len(indexed))):
            raise ValueError("Class IDs must be contiguous from zero")
        names = [indexed[i] for i in range(len(indexed))]
    if not isinstance(names, list) or not names or not all(isinstance(n, str) for n in names):
        raise ValueError("Expected a nonempty class-name list or indexed mapping")
    if len(set(names)) != len(names):
        raise ValueError("Duplicate class names")
    return names


def class_order(source: list[str], target: list[str]) -> list[int]:
    source, target = ordered_names(source), ordered_names(target)
    if set(source) != set(target):
        raise ValueError(f"Class sets differ: classifier={source}, YOLO={target}")
    return [source.index(name) for name in target]


def preprocessing(config: dict, size: int) -> float | None:
    """Reject preprocessing that cannot be reproduced without resampling the input."""
    transforms = config["data"]["val_transforms"]
    sequence = [t["name"] for t in transforms]
    if sequence not in (["Resize", "ToTensor", "PerImageStandardize"],
                        ["ToTensor", "PerImageStandardize"], ["Resize", "ToTensor"], ["ToTensor"]):
        raise ValueError(f"Unsupported validation preprocessing: {sequence}")
    for transform in transforms:
        if transform["name"] == "Resize":
            if transform["size"] != [size, size]:
                raise ValueError("Training size differs from template; retrain at the target size")
        if transform["name"] == "PerImageStandardize":
            return float(transform.get("eps", 1e-5))
    return None


class YoloClassifierAdapter(nn.Module):
    def __init__(self, classifier: nn.Module, size: int, anchors: int,
                 order: list[int], eps: float | None = 1e-5,
                 background_index: int | None = None):
        super().__init__()
        if size <= 0 or anchors <= 0:
            raise ValueError("size and anchors must be positive")
        self.classifier = classifier
        self.eps = eps
        if background_index is not None and not 0 <= background_index < len(order):
            raise ValueError("background_index must be an output class index")
        self.background_index = background_index
        self.register_buffer("order", torch.tensor(order, dtype=torch.long))
        self.register_buffer("box", torch.tensor([size / 2, size / 2, size, size])
                             .reshape(1, 4, 1))
        # Only candidate zero is active; do not duplicate predictions across anchors.
        self.register_buffer("inactive", torch.zeros(1, 4 + len(order), anchors - 1))

    def normalize(self, x: torch.Tensor) -> torch.Tensor:
        if self.eps is None:
            return x
        dims = (1, 2, 3)
        mean = x.mean(dim=dims, keepdim=True)
        # Match datasets.transforms.PerImageStandardize: torch.std correction=1.
        variance = ((x - mean) ** 2).sum(dim=dims, keepdim=True) / (x[0].numel() - 1)
        return (x - mean) / (torch.sqrt(variance) + self.eps)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        scores = self.classifier(self.normalize(x)).softmax(dim=1)[:, self.order]
        if self.background_index is not None:
            foreground = (scores.argmax(dim=1) != self.background_index).unsqueeze(1)
            scores = scores * foreground.to(scores.dtype)
        active = torch.cat((self.box.expand(x.shape[0], -1, -1), scores.unsqueeze(2)), dim=1)
        return torch.cat((active, self.inactive.expand(x.shape[0], -1, -1)), dim=2)


def main() -> int:
    import openvino as ov

    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--checkpoint", default="runs/classify/best.pth")
    parser.add_argument("--template", required=True, help="Production YOLO .xml (metadata.yaml beside it)")
    parser.add_argument("--classes", help="Classifier class order: classes.txt or eval_report.json")
    parser.add_argument("--background-class", default="normal", help="Top1 class to emit no boxes for")
    parser.add_argument("--keep-background-box", action="store_true",
                        help="Opt out of background suppression (legacy classification behavior)")
    parser.add_argument("--out", required=True, help="New directory; existing directories are rejected")
    args = parser.parse_args()
    checkpoint, template, out = Path(args.checkpoint), Path(args.template), Path(args.out)
    if out.exists():
        raise FileExistsError(f"Refusing to overwrite {out}")
    state = torch.load(checkpoint, map_location="cpu", weights_only=False)
    config = state["config"]
    if args.classes:
        path = Path(args.classes)
        source = (json.loads(path.read_text())["classes"] if path.suffix == ".json"
                  else path.read_text().splitlines())
    else:
        source = config["data"].get("classes")
        if not source:
            report = checkpoint.parent / "eval_report.json"
            if not report.exists():
                raise ValueError("No recorded class order; provide --classes (do not guess folder order)")
            source = json.loads(report.read_text())["classes"]
    source = ordered_names(source)
    metadata = yaml.safe_load((template.parent / "metadata.yaml").read_text())
    if metadata.get("end2end") or metadata.get("args", {}).get("nms"):
        raise ValueError("Expected raw [1, 4+nc, anchors] output, not end-to-end/NMS output")
    target = ordered_names(metadata["names"])
    order = class_order(source, target)
    if not args.keep_background_box and args.background_class not in target:
        raise ValueError(f"Background class {args.background_class!r} is not in {target}")
    background_index = None if args.keep_background_box else target.index(args.background_class)
    core = ov.Core()
    reference = core.read_model(str(template))
    if len(reference.inputs) != 1 or len(reference.outputs) != 1:
        raise ValueError("Expected one input and one output")
    input_shape, output_shape = list(reference.input().shape), list(reference.output().shape)
    if (len(input_shape) != 4 or input_shape[:2] != [1, 3]
            or input_shape[2] != input_shape[3] or len(output_shape) != 3
            or output_shape[:2] != [1, 4 + len(target)] or output_shape[2] < 1):
        raise ValueError(f"Unsupported template shapes: {input_shape}, {output_shape}")
    if (reference.input().element_type != ov.Type.f32
            or reference.output().element_type != ov.Type.f32):
        raise ValueError("Only float32 input/output are supported")
    size = input_shape[2]
    eps = preprocessing(config, size)
    model_config = copy.deepcopy(config["model"])
    model_config["backbone"]["pretrained"] = False
    model_config["init_from"] = None
    model_config["num_classes"] = len(source)
    model = build_model(model_config).eval()
    model.load_state_dict(state["model"], strict=True)
    adapter = YoloClassifierAdapter(model, size, output_shape[2], order, eps,
                                    background_index=background_index).eval()
    torch.manual_seed(42)
    example = torch.rand(input_shape)
    converted = ov.convert_model(adapter, example_input=example, input=input_shape)
    converted.input().get_tensor().set_names(reference.input().get_names())
    converted.output().get_tensor().set_names(reference.output().get_names())
    # Preserve any interface metadata used by OpenVINO consumers, then mark pseudo boxes.
    for key, value in reference.get_rt_info().items():
        if key != "conversion_parameters":
            converted.set_rt_info(value, [key])
    converted.set_rt_info("classifier_pseudo_box", ["adapter_kind"])
    out.mkdir(parents=True)
    xml = out / template.name
    ov.save_model(converted, str(xml), compress_to_fp16=True)
    compiled = core.compile_model(str(xml), "CPU", {"PERFORMANCE_HINT": "LATENCY"})
    errors = []
    with torch.inference_mode():
        for sample in (example, torch.zeros_like(example), torch.ones_like(example)):
            expected = adapter(sample).numpy()
            actual = compiled(sample.numpy())[compiled.output()]
            np.testing.assert_allclose(actual, expected, atol=0.02, rtol=0.02)
            assert np.argmax(actual[0, 4:, 0]) == np.argmax(expected[0, 4:, 0])
            assert np.count_nonzero(actual[:, :, 1:]) == 0
            errors.append(float(np.max(np.abs(expected - actual))))
    metadata.update({"names": dict(enumerate(target)), "adapter_kind": "classifier_pseudo_box",
                     "description": "Classifier with one full-image pseudo box; no localization",
                     "background_class": None if background_index is None else args.background_class,
                     "end2end": False, "args": {"nms": False}, "imgsz": [size, size]})
    (out / "metadata.yaml").write_text(yaml.safe_dump(metadata, allow_unicode=True))
    (out / "classes.txt").write_text("\n".join(target) + "\n")
    summary = {"checkpoint": str(checkpoint.resolve()), "template": str(template.resolve()),
               "classifier_classes": source, "output_classes": target, "class_order": order,
               "input_shape": input_shape, "output_shape": output_shape,
               "standardize_eps": eps, "max_abs_errors": errors,
               "background_class": None if background_index is None else args.background_class,
               "background_index": background_index,
               "pseudo_box": [size / 2, size / 2, size, size]}
    (out / "adapter.json").write_text(json.dumps(summary, ensure_ascii=False, indent=2))
    print(json.dumps({"model": str(xml), **summary}, ensure_ascii=False, indent=2))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())