"""基础数据读取与固定种子划分工具。

本模块只负责通用的图像读取、COCO JSON 读取和 ID 划分，不负责训练、
切片或模型评估。OpenCV 在真正读取图像时才导入，避免只读取 JSON 时
强制安装图像依赖。
"""

from pathlib import Path
import json
import math
from typing import Any, Mapping, Sequence

import numpy as np


def read_image(path: str | Path) -> np.ndarray:
    """读取图像并返回 RGB 数组。

    SAR 单通道图像会复制为三个通道，保证后续检测模型获得统一输入。
    优先使用 Pillow；Pillow 不支持的格式再尝试 OpenCV。
    """
    image_path = Path(path)
    if not image_path.exists():
        raise FileNotFoundError(f"图像不存在: {image_path.resolve()}")
    try:
        from PIL import Image

        with Image.open(image_path) as opened:
            image = np.asarray(opened.convert("RGB"))
        return image
    except ImportError:
        try:
            import cv2
        except ImportError as exc:
            raise RuntimeError("读取图像需要 Pillow 或 opencv-python。") from exc
        image = cv2.imread(str(image_path), cv2.IMREAD_UNCHANGED)
        if image is None:
            raise ValueError(f"无法解码图像: {image_path.resolve()}")
        if image.ndim == 2:
            image = np.repeat(image[..., None], 3, axis=2)
        return cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

def load_coco(path: str | Path) -> dict[str, Any]:
    """读取 COCO JSON，并检查 ID、列表字段和标注引用关系。"""
    annotation_path = Path(path)
    if not annotation_path.exists():
        raise FileNotFoundError(f"COCO 文件不存在: {annotation_path.resolve()}")
    with annotation_path.open(encoding="utf-8") as f:
        data = json.load(f)
    required = {"images", "annotations", "categories"}
    if not isinstance(data, dict):
        raise ValueError("COCO 根对象必须是字典。")
    missing = required.difference(data)
    if missing:
        raise ValueError(f"COCO 文件缺少字段: {sorted(missing)}")
    if not all(isinstance(data[key], list) for key in required):
        raise ValueError("COCO 的 images、annotations、categories 必须为列表。")
    image_ids = [
        item.get("id") for item in data["images"] if isinstance(item, Mapping)
    ]
    category_ids = [
        item.get("id") for item in data["categories"] if isinstance(item, Mapping)
    ]
    if (
        len(image_ids) != len(data["images"])
        or len(category_ids) != len(data["categories"])
    ):
        raise ValueError("COCO images/categories 中每项必须是对象。")
    if None in image_ids or len(set(image_ids)) != len(image_ids):
        raise ValueError("COCO images.id 不能为空且不能重复。")
    if None in category_ids or len(set(category_ids)) != len(category_ids):
        raise ValueError("COCO categories.id 不能为空且不能重复。")
    image_id_set, category_id_set = set(image_ids), set(category_ids)
    annotation_ids = [
        item.get("id") for item in data["annotations"] if isinstance(item, Mapping)
    ]
    if len(annotation_ids) != len(data["annotations"]):
        raise ValueError("COCO annotations 中每项必须是对象。")
    if None in annotation_ids or len(set(annotation_ids)) != len(annotation_ids):
        raise ValueError("COCO annotations.id 不能为空且不能重复。")
    for index, annotation in enumerate(data["annotations"]):
        if not isinstance(annotation, Mapping):
            raise ValueError(f"COCO 第 {index} 条 annotation 必须是对象。")
        if annotation.get("image_id") not in image_id_set:
            raise ValueError(f"COCO 第 {index} 条 annotation 引用了不存在的 image_id。")
        if annotation.get("category_id") not in category_id_set:
            raise ValueError(f"COCO 第 {index} 条 annotation 引用了不存在的 category_id。")
        bbox = annotation.get("bbox")
        if not isinstance(bbox, (list, tuple)) or len(bbox) != 4:
            raise ValueError(f"COCO 第 {index} 条 bbox 必须是 [x, y, width, height]。")
        try:
            values = [float(value) for value in bbox]
        except (TypeError, ValueError) as exc:
            raise ValueError(f"COCO 第 {index} 条 bbox 必须为数值。") from exc
        if (
            not all(math.isfinite(value) for value in values)
            or values[2] <= 0
            or values[3] <= 0
        ):
            raise ValueError(f"COCO 第 {index} 条 bbox 必须为有限数值且宽高为正。")
    return data


def split_ids(
    items: Sequence[dict], train: float = 0.7, val: float = 0.15, seed: int = 42
) -> tuple[list[Any], list[Any], list[Any]]:
    """按固定种子划分 ID。

    该函数适合尚未切片的独立图像。若多张切片来自同一原图，必须在
    原图分组后再调用，避免相邻切片泄漏到不同数据集。
    """
    if not isinstance(seed, int) or isinstance(seed, bool):
        raise TypeError("seed 必须是整数。")
    try:
        train_ratio, val_ratio = float(train), float(val)
    except (TypeError, ValueError) as exc:
        raise TypeError("train 和 val 必须是数值比例。") from exc
    if (
        not math.isfinite(train_ratio)
        or not math.isfinite(val_ratio)
        or train_ratio < 0
        or val_ratio < 0
        or train_ratio + val_ratio > 1
    ):
        raise ValueError("train 和 val 比例必须非负，且 train + val 不得大于 1。")
    if not items:
        return [], [], []
    ids = [item.get("id") for item in items]
    if any(item_id is None for item_id in ids):
        raise ValueError("每个划分对象都必须包含 id。")
    if len(set(ids)) != len(ids):
        raise ValueError("划分对象的 id 必须唯一。")
    rng = np.random.default_rng(seed)
    shuffled = np.asarray(ids, dtype=object)
    rng.shuffle(shuffled)
    n = len(shuffled)
    a, b = int(n * train_ratio), int(n * (train_ratio + val_ratio))
    return shuffled[:a].tolist(), shuffled[a:b].tolist(), shuffled[b:].tolist()