"""セグメンテーションの評価データ（`<set>/sample/*` と `<set>/label/*`）の列挙と読み込み.

入力画像と正解マスクはファイル名の stem で対応づける．OpenCV の `imread` は Windows で日本語を
含むパスを読めないため，`np.fromfile` → `cv2.imdecode` で読む．例外メッセージにはファイル名を
出さず，評価セット内の何件目か（stem の昇順で 1 始まり）と件数で示す．
"""

from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path

import cv2
import numpy as np

# 評価データとして読む画像の拡張子（小文字で比べる）．
IMAGE_SUFFIXES: frozenset[str] = frozenset({".jpg", ".jpeg", ".png", ".bmp"})


@dataclass(frozen=True)
class ImagePair:
    """入力画像と正解マスクの組.

    Attributes:
        stem: ファイル名の stem（サーバ上の個票にだけ書く．ログ・例外には出さない）．
        sample: 入力画像のパス．
        label: 正解マスクのパス．
    """

    stem: str
    sample: Path
    label: Path


def _images_by_stem(directory: Path, what: str) -> dict[str, Path]:
    if not directory.is_dir():
        raise ValueError(f"{what}のディレクトリが無い: {directory}")
    by_stem: dict[str, Path] = {}
    duplicates = 0
    for path in directory.iterdir():
        if path.is_file() and path.suffix.lower() in IMAGE_SUFFIXES:
            if path.stem in by_stem:
                duplicates += 1
            by_stem[path.stem] = path
    if duplicates:
        raise ValueError(f"{what}に stem が同じ画像が {duplicates} 件ある（拡張子違いの重複）")
    return by_stem


def list_image_pairs(set_dir: Path, sample_dirname: str, label_dirname: str) -> list[ImagePair]:
    """評価セットの入力画像と正解マスクの組を stem の昇順で列挙する.

    Args:
        set_dir: 評価セットのディレクトリ．
        sample_dirname: 入力画像のディレクトリ名．
        label_dirname: 正解マスクのディレクトリ名．

    Returns:
        組の一覧（stem の昇順）．

    Raises:
        ValueError: ディレクトリが無い，画像が 0 枚，stem の重複，または対応の取れない画像が
            ある場合（件数のみを示す）．
    """
    samples = _images_by_stem(set_dir / sample_dirname, "入力画像")
    labels = _images_by_stem(set_dir / label_dirname, "正解マスク")
    only_sample = len(samples.keys() - labels.keys())
    only_label = len(labels.keys() - samples.keys())
    if only_sample or only_label:
        raise ValueError(
            f"入力画像と正解マスクの対応が取れない（正解マスクの無い入力画像 {only_sample} 件・"
            f"入力画像の無い正解マスク {only_label} 件）"
        )
    if not samples:
        raise ValueError(f"評価セットに画像が 0 枚である: {set_dir}")
    return [ImagePair(stem, samples[stem], labels[stem]) for stem in sorted(samples)]


def _decode(path: Path, flags: int, where: str) -> np.ndarray:
    try:
        data = np.fromfile(path, dtype=np.uint8)
    except OSError as exc:
        raise ValueError(f"{where}: 画像を読めない（{type(exc).__name__}）") from None
    image = cv2.imdecode(data, flags)
    if image is None:
        raise ValueError(f"{where}: 画像を復号できない")
    return image


def read_bgr(path: Path, where: str) -> np.ndarray:
    """カラー画像を OpenCV の BGR で読む.

    Args:
        path: 画像のパス．
        where: 例外メッセージに使う位置の表示（ファイル名を含めない）．

    Returns:
        `(H, W, 3)` の uint8（BGR）．

    Raises:
        ValueError: 読めない場合．
    """
    return _decode(path, cv2.IMREAD_COLOR, where)


def read_label_mask(path: Path, threshold: int, where: str) -> np.ndarray:
    """正解マスクをグレースケールで読み，`threshold` 以上を前景とする.

    Args:
        path: 正解マスクのパス．
        threshold: 前景とみなす下限の画素値．
        where: 例外メッセージに使う位置の表示（ファイル名を含めない）．

    Returns:
        `(H, W)` の bool．

    Raises:
        ValueError: 読めない場合．
    """
    return _decode(path, cv2.IMREAD_GRAYSCALE, where) >= threshold