"""セグメンテーションの指標（前景の IoU・Dice）と画像ごとの値の集計.

集計値（平均・中央値・分位点・最小・しきい値未満の件数）だけを metrics.json に書くための
純関数．画像ごとの値はサーバ上の個票（predictions.csv）にだけ書く．
"""

from __future__ import annotations

from collections.abc import Sequence
from typing import Any

import numpy as np


def iou_and_dice(pred: np.ndarray, target: np.ndarray) -> tuple[float, float]:
    """前景の IoU と Dice を求める.

    両方とも前景が無い場合は一致とみなして 1.0 とする．

    Args:
        pred: `(H, W)` の bool（予測の前景）．
        target: `(H, W)` の bool（正解の前景）．

    Returns:
        (IoU, Dice)．

    Raises:
        ValueError: 形状が一致しない場合．
    """
    pred = np.asarray(pred, dtype=bool)
    target = np.asarray(target, dtype=bool)
    if pred.shape != target.shape:
        raise ValueError(f"予測と正解の形状が一致しない（{pred.shape} と {target.shape}）")
    intersection = int(np.count_nonzero(pred & target))
    union = int(np.count_nonzero(pred | target))
    total = int(np.count_nonzero(pred)) + int(np.count_nonzero(target))
    if union == 0:
        return 1.0, 1.0
    return intersection / union, 2.0 * intersection / total


def summarize_scores(values: Sequence[float], threshold: float) -> dict[str, Any]:
    """画像ごとの指標を集計する.

    Args:
        values: 画像ごとの値．
        threshold: この値未満の件数を数える．

    Returns:
        `n`・`mean`・`median`・`p10`・`p25`・`min`・`max`・`n_below`・`below_threshold`．
        値が 0 件なら `n` が 0 で他は None．
    """
    arr = np.asarray(values, dtype=np.float64)
    if arr.size == 0:
        return {
            "n": 0,
            "mean": None,
            "median": None,
            "p10": None,
            "p25": None,
            "min": None,
            "max": None,
            "n_below": None,
            "below_threshold": threshold,
        }
    return {
        "n": int(arr.size),
        "mean": float(arr.mean()),
        "median": float(np.median(arr)),
        "p10": float(np.percentile(arr, 10)),
        "p25": float(np.percentile(arr, 25)),
        "min": float(arr.min()),
        "max": float(arr.max()),
        "n_below": int(np.count_nonzero(arr < threshold)),
        "below_threshold": threshold,
    }