"""Python 版セグメンテーション（DeepTIAS 2.7 の再現）を正解マスク付きデータで評価する.

`tias-seg eval` の本体．評価セットの種類（`data.sets[].kind`）ごとに次の方式（枠の決め方）で
ResUNet を掛け，前景の IoU・Dice を求める．

- `cropped`（切り出し済み）: `whole`（画像全体に掛ける．ResUNet の学習時と同じ条件）
- `fullsize`（撮影画像）: `whole`（検出なしで画像全体に掛ける）・`gt_bbox`（正解マスクの外接矩形で
  切り出す．検出が完璧な場合の上限）・`detector`（検出モデルの枠で切り出す．モデルがある場合のみ）

どの方式も，後処理（小領域除去）あり（DeepTIAS と同じ）・なし，もう一方の出力チャンネル
（後処理あり）の 3 通りのマスクを 1 回の推論から作る．画像ごとの値（`ImageScore`）は個票として
サーバ上にだけ書き，回収する metrics.json には集計値（`summarize_segmentation`）だけを書く．
ログ・例外には件数と評価セット内の何件目かだけを出し，ファイル名は出さない．
"""

from __future__ import annotations

import logging
import time
from collections.abc import Sequence
from dataclasses import dataclass, field
from typing import Any

import numpy as np

from tias_judge_model.evaluation.segmentation_metrics import iou_and_dice, summarize_scores
from tias_judge_model.schema import SegEvalConfig, SegEvalSetConfig
from tias_judge_model.segmentation.detection import (
    OnnxDetector,
    Rect,
    full_rect,
    mask_bounding_rect,
)
from tias_judge_model.segmentation.images import list_image_pairs, read_bgr, read_label_mask
from tias_judge_model.segmentation.pipeline import SegmentationParams, empty_masks, segment_in_rect
from tias_judge_model.segmentation.resunet import CropMasks, ResUNetRunner

logger = logging.getLogger(__name__)

METHOD_WHOLE = "whole"
METHOD_GT_BBOX = "gt_bbox"
METHOD_DETECTOR = "detector"
# 方式の説明（metrics.json に書く）．
METHOD_DESCRIPTIONS: dict[str, str] = {
    METHOD_WHOLE: "検出なし．画像全体を ResUNet の入力サイズにして掛ける",
    METHOD_GT_BBOX: "正解マスクの外接矩形で切り出して掛ける（検出が完璧な場合の上限）",
    METHOD_DETECTOR: "検出モデル（スコアがしきい値を超える最高の枠）で切り出して掛ける",
}


@dataclass(frozen=True)
class ImageScore:
    """画像 1 枚・方式 1 つの評価（個票．サーバ上の predictions.csv にだけ書く）.

    Attributes:
        set_name: 評価セットの名前．
        index: 評価セット内の位置（stem の昇順で 1 始まり）．
        stem: ファイル名の stem．
        method: 方式．
        iou: IoU（後処理あり）．
        dice: Dice（後処理あり）．
        iou_no_post: IoU（後処理なし）．
        dice_no_post: Dice（後処理なし）．
        iou_other_channel: もう一方の出力チャンネルを使った IoU（後処理あり）．
        postprocess_applied: 小領域除去をしたか．
        detected: 検出できたか（`detector` 以外は None）．
        detection_score: 検出の最高スコア（`detector` 以外は None）．
        seconds: 検出・推論・後処理の所要秒数．
    """

    set_name: str
    index: int
    stem: str
    method: str
    iou: float
    dice: float
    iou_no_post: float
    dice_no_post: float
    iou_other_channel: float
    postprocess_applied: bool
    detected: bool | None
    detection_score: float | None
    seconds: float


@dataclass
class SegEvalResult:
    """評価全体の結果.

    Attributes:
        scores: 画像・方式ごとの評価（個票）．
        methods: 評価セットの名前 → 評価した方式（順序つき）．
        detector_used: 検出モデルを使ったか．
        detector_note: 検出モデルを使わなかった理由（使った場合は None）．
        loading_seconds: 画像の読み込みの合計秒数．
    """

    scores: list[ImageScore] = field(default_factory=list)
    methods: dict[str, list[str]] = field(default_factory=dict)
    detector_used: bool = False
    detector_note: str | None = None
    loading_seconds: float = 0.0


def methods_for(kind: str, has_detector: bool) -> list[str]:
    """評価セットの種類から評価する方式を返す.

    Args:
        kind: `cropped` または `fullsize`．
        has_detector: 検出モデルがあるか．

    Returns:
        方式の一覧．
    """
    if kind == "cropped":
        return [METHOD_WHOLE]
    methods = [METHOD_WHOLE, METHOD_GT_BBOX]
    if has_detector:
        methods.append(METHOD_DETECTOR)
    return methods


def _load_detector(config: SegEvalConfig) -> tuple[OnnxDetector | None, str | None]:
    path = config.models.detection_path
    if path is None:
        return None, "設定で検出モデルを指定していない（models.detection: null）"
    if not path.is_file():
        # 変換できなかった場合でも検出なし・正解の外接矩形の評価は完走させるため，止めずに省く．
        logger.warning("検出モデルの ONNX が無いため，detector の方式を省く: %s", path)
        return None, "検出モデルの ONNX が無い"
    d = config.detection
    detector = OnnxDetector(
        path,
        score_threshold=d.score_threshold,
        input_color=d.input_color,
        input_name=d.input_name,
        boxes_name=d.boxes_name,
        scores_name=d.scores_name,
    )
    return detector, None


def _segment_by_method(
    method: str,
    bgr: np.ndarray,
    gt_rect: Rect,
    runner: ResUNetRunner,
    detector: OnnxDetector | None,
    params: SegmentationParams,
) -> tuple[CropMasks, bool | None, float | None]:
    # 方式ごとに枠を決めて ResUNet を掛ける．戻り値は（`(H, W)` の bool のマスク，検出できたか，
    # 検出の最高スコア）．検出の 2 つは detector 以外では None．
    height, width = bgr.shape[:2]
    if method == METHOD_WHOLE:
        return segment_in_rect(bgr, full_rect(height, width), runner, params), None, None
    if method == METHOD_GT_BBOX:
        return segment_in_rect(bgr, gt_rect, runner, params), None, None
    if detector is None:
        raise RuntimeError("検出モデルが無いのに detector の方式を評価しようとした")
    detection = detector.detect(bgr)
    masks = (
        segment_in_rect(bgr, detection.rect, runner, params)
        if detection.rect is not None
        else empty_masks(height, width)
    )
    return masks, detection.rect is not None, detection.score


def _evaluate_set(
    set_config: SegEvalSetConfig,
    config: SegEvalConfig,
    runner: ResUNetRunner,
    detector: OnnxDetector | None,
    params: SegmentationParams,
    result: SegEvalResult,
) -> None:
    set_dir = config.data.root / set_config.dir
    pairs = list_image_pairs(set_dir, config.data.sample_dirname, config.data.label_dirname)
    methods = methods_for(set_config.kind, detector is not None)
    result.methods[set_config.name] = methods
    logger.info(
        "評価セット %s（%s）: 画像 %d 件・方式 %s",
        set_config.name,
        set_config.kind,
        len(pairs),
        methods,
    )
    progress_every = max(1, len(pairs) // 10)
    for number, pair in enumerate(pairs, start=1):
        where = f"評価セット {set_config.name} の {number} 件目"
        started = time.monotonic()
        bgr = read_bgr(pair.sample, where)
        target = read_label_mask(pair.label, config.data.label_threshold, where)
        result.loading_seconds += time.monotonic() - started
        if target.shape != bgr.shape[:2]:
            raise ValueError(
                f"{where}: 入力画像 {bgr.shape[1]}×{bgr.shape[0]} と正解マスク "
                f"{target.shape[1]}×{target.shape[0]} のサイズが違う"
            )
        gt_rect = mask_bounding_rect(target)
        if gt_rect is None:
            raise ValueError(f"{where}: 正解マスクに前景が無い")

        for method in methods:
            started = time.monotonic()
            masks, detected, det_score = _segment_by_method(
                method, bgr, gt_rect, runner, detector, params
            )
            iou, dice = iou_and_dice(masks.post, target)
            iou_np, dice_np = iou_and_dice(masks.no_post, target)
            iou_other, _ = iou_and_dice(masks.other_channel, target)
            result.scores.append(
                ImageScore(
                    set_name=set_config.name,
                    index=number,
                    stem=pair.stem,
                    method=method,
                    iou=iou,
                    dice=dice,
                    iou_no_post=iou_np,
                    dice_no_post=dice_np,
                    iou_other_channel=iou_other,
                    postprocess_applied=masks.postprocess_applied,
                    detected=detected,
                    detection_score=det_score,
                    seconds=time.monotonic() - started,
                )
            )
        if number % progress_every == 0 or number == len(pairs):
            logger.info("評価セット %s: %d / %d 件", set_config.name, number, len(pairs))


def run_segmentation_eval(config: SegEvalConfig) -> SegEvalResult:
    """設定の評価セットをすべて評価する.

    Args:
        config: 評価の設定．

    Returns:
        評価の結果（画像ごとの値を含む）．

    Raises:
        ValueError: ResUNet の ONNX が無い・形状が合わない，評価セットのディレクトリが無い，
            画像を読めない，入力画像と正解マスクの対応が取れない・サイズが違う，正解マスクに
            前景が無い等の入力の問題．
    """
    runner = ResUNetRunner(config.models.resunet_path, config.resunet.input_size)
    detector, note = _load_detector(config)
    params = SegmentationParams(
        foreground_channel=config.resunet.foreground_channel,
        min_area=config.postprocess.min_area,
        mask_threshold=config.postprocess.mask_threshold,
    )
    result = SegEvalResult(detector_used=detector is not None, detector_note=note)
    for set_config in config.data.sets:
        _evaluate_set(set_config, config, runner, detector, params, result)
    return result


def _scores_of(scores: Sequence[ImageScore], set_name: str, method: str) -> list[ImageScore]:
    return [s for s in scores if s.set_name == set_name and s.method == method]


def ious_by_method(result: SegEvalResult, set_name: str) -> dict[str, list[float]]:
    """評価セットの画像ごとの IoU（後処理あり）を方式ごとに返す（図を描くため）.

    Args:
        result: 評価の結果．
        set_name: 評価セットの名前．

    Returns:
        方式（評価した順）→ 画像ごとの IoU（評価セット内の順）．
    """
    return {
        method: [s.iou for s in _scores_of(result.scores, set_name, method)]
        for method in result.methods.get(set_name, [])
    }


def _summarize_method(
    method: str, scores: Sequence[ImageScore], config: SegEvalConfig
) -> dict[str, Any]:
    threshold = config.evaluation.iou_threshold
    entry: dict[str, Any] = {
        "description": METHOD_DESCRIPTIONS[method],
        "iou": summarize_scores([s.iou for s in scores], threshold),
        "dice": summarize_scores([s.dice for s in scores], threshold),
        "iou_no_post": summarize_scores([s.iou_no_post for s in scores], threshold),
        "dice_no_post": summarize_scores([s.dice_no_post for s in scores], threshold),
        "iou_other_channel": summarize_scores([s.iou_other_channel for s in scores], threshold),
        "n_postprocess_applied": sum(s.postprocess_applied for s in scores),
        "mean_seconds": round(float(np.mean([s.seconds for s in scores])), 4) if scores else None,
    }
    if method == METHOD_DETECTOR:
        det_scores = [s.detection_score for s in scores if s.detection_score is not None]
        entry["detection"] = {
            "n_detected": sum(bool(s.detected) for s in scores),
            "n_not_detected": sum(not s.detected for s in scores),
            "score_threshold": config.detection.score_threshold,
            "top_score": summarize_scores(det_scores, config.detection.score_threshold),
        }
    return entry


def summarize_segmentation(result: SegEvalResult, config: SegEvalConfig) -> dict[str, Any]:
    """評価セット・方式ごとの集計値を作る（metrics.json の `sets`．画像単位の値は含めない）.

    Args:
        result: 評価の結果．
        config: 評価の設定．

    Returns:
        評価セットの名前 → `{dir, kind, n_images, methods: {方式: 集計}}`．方式ごとの集計は
        `iou`・`dice`（後処理あり）・`iou_no_post`・`dice_no_post`・`iou_other_channel` の
        `summarize_scores`，小領域除去をした件数，所要秒数の平均，`detector` では検出の件数．
    """
    sets: dict[str, Any] = {}
    for set_config in config.data.sets:
        methods: dict[str, Any] = {}
        n_images = 0
        for method in result.methods.get(set_config.name, []):
            scores = _scores_of(result.scores, set_config.name, method)
            n_images = len(scores)
            methods[method] = _summarize_method(method, scores, config)
        sets[set_config.name] = {
            "dir": set_config.dir.as_posix(),
            "kind": set_config.kind,
            "n_images": n_images,
            "methods": methods,
        }
    return sets


def log_summary(sets: dict[str, Any]) -> None:
    """集計値をログに出す（出力チャンネルの取り違えに気づけるよう，もう一方のチャンネルも出す）.

    Args:
        sets: `summarize_segmentation` の戻り値．
    """
    for set_name, entry in sets.items():
        for method, m in entry["methods"].items():
            iou, iou_np, other = m["iou"], m["iou_no_post"], m["iou_other_channel"]
            if not iou["n"]:
                continue
            logger.info(
                "%s / %s: IoU 平均 %.4f・中央値 %.4f・p10 %.4f・最小 %.4f・%.2f 未満 %d / %d 件"
                "（後処理なし: 平均 %.4f・%.2f 未満 %d 件．小領域除去 %d 件）・Dice 平均 %.4f",
                set_name,
                method,
                iou["mean"],
                iou["median"],
                iou["p10"],
                iou["min"],
                iou["below_threshold"],
                iou["n_below"],
                iou["n"],
                iou_np["mean"],
                iou_np["below_threshold"],
                iou_np["n_below"],
                m["n_postprocess_applied"],
                m["dice"]["mean"],
            )
            logger.info(
                "%s / %s: もう一方の出力チャンネルを使った場合の IoU 平均 %.4f・中央値 %.4f",
                set_name,
                method,
                other["mean"],
                other["median"],
            )
            if other["mean"] > iou["mean"]:
                logger.warning(
                    "%s / %s: もう一方の出力チャンネルのほうが IoU が高い（%.4f > %.4f）．"
                    "resunet.foreground_channel の取り違えを確認する",
                    set_name,
                    method,
                    other["mean"],
                    iou["mean"],
                )
            detection = m.get("detection")
            if detection:
                logger.info(
                    "%s / %s: 検出できた %d 件・できなかった %d 件（しきい値 %.2f）",
                    set_name,
                    method,
                    detection["n_detected"],
                    detection["n_not_detected"],
                    detection["score_threshold"],
                )