"""用 ultralytics 做切片推理的公共封装。

三个调用方（可视化、批量推送预测、ML backend）都要"切片 -> 逐片推理 -> 拼回原图"，
逻辑一致，放在一处免得各写一份走样。

几何在 `engine.tiling` 里，与推理框架无关且单独可测；这里只负责把 ultralytics 接上去。
"""

from __future__ import annotations

from collections.abc import Sequence

import numpy as np

from engine.detection import Box
from engine.tiling import tiled_predict


def training_imgsz(model) -> int | None:
    """权重里记录的训练输入尺寸；读不到返回 None。

    ultralytics 把训练参数存在 checkpoint 里，推理时可以取回来核对。
    """
    for source in (
        getattr(model, "ckpt", None) or {},
        getattr(model, "overrides", None) or {},
    ):
        value = source.get("train_args", source).get("imgsz") if isinstance(source, dict) else None
        if isinstance(value, int):
            return value
    return None


def check_size_matches(model, size: int, logger=None) -> str | None:
    """推理尺寸与训练尺寸不一致时给出警告文本（一致或读不到则返回 None）。

    这一条是有代价换来的：用 640 整图训练的模型去跑 160 的切片推理，结果**看起来正常**
    ——照样出框、照样有置信度，但类别系统性地错（实测 mosquito 全部变成 blackdot，
    TP 为 0）。不主动核对的话，这种结果会被当成模型退化去查很久。
    """
    trained = training_imgsz(model)
    if trained is None or size <= 0 or trained == size:
        return None
    message = (
        f"权重是按 imgsz={trained} 训练的，现在按 {size} 切片推理——两者不一致，"
        f"结果不可信（会照常出框，但类别系统性地错）。"
        f"要么换成按 {size} 训练的权重，要么用 --tile-size 0 退回整图推理"
    )
    if logger is not None:
        logger.warning("%s", message)
    return message


def predict_image(
    model,
    image: np.ndarray,
    size: int,
    overlap: int = 0,
    conf: float = 0.25,
    device: str | int = "cpu",
    batch: int = 16,
    iou_threshold: float = 0.5,
    whole_image_imgsz: int | None = None,
) -> list[Box]:
    """对一张任意尺寸的图做切片推理，返回整图上的归一化框。

    `size <= 0` 时退化为整图推理，交给 ultralytics 按 `whole_image_imgsz` 做 letterbox
    ——那会缩放，只用于和旧模型对比。
    """
    if size <= 0:
        return _predict_batch(
            model, [image], whole_image_imgsz, conf=conf, device=device, batch=batch
        )[0]

    def run(patches: Sequence[np.ndarray]) -> list[list[Box]]:
        return _predict_batch(model, patches, size, conf, device, batch)

    return tiled_predict(run, image, size, overlap, iou_threshold)


def _predict_batch(
    model,
    images: Sequence[np.ndarray],
    size: int | None,
    conf: float,
    device: str | int,
    batch: int,
) -> list[list[Box]]:
    """分批送进 ultralytics，返回每张图的归一化框。

    必须自己分批：ultralytics 不理会 batch 参数，会把整个列表当成一个 batch，
    片数一多就会撑爆显存。
    """
    results = []
    for start in range(0, len(images), batch):
        chunk = list(images[start:start + batch])
        kwargs = {"conf": conf, "device": device, "verbose": False}
        if size:
            kwargs["imgsz"] = size
        results.extend(model.predict(source=chunk, **kwargs))

    out = []
    for result in results:
        out.append([
            Box(int(class_id), *(float(v) for v in xyxyn), float(score))
            for xyxyn, class_id, score in zip(
                result.boxes.xyxyn.tolist(), result.boxes.cls.tolist(),
                result.boxes.conf.tolist(), strict=True,
            )
        ])
    return out


def load_image(path) -> np.ndarray:
    """读图为 numpy 数组，保持原始通道数（灰度图不强行转 RGB，免得多占三倍内存）。"""
    from PIL import Image

    with Image.open(path) as handle:
        image = handle.convert("L") if handle.mode in ("L", "1", "I;16") else handle.convert("RGB")
        return np.asarray(image)