"""データセット集計（`tias-dataset survey` の計算部分）.

集計項目は SPEC_01「データセット集計の出力」に従う．画像の読み込み（`read_image_stats`・
`collect_image_stats`）と，読み込み結果からの集計（`summarize_dataset`．純関数）を分けている．

**出力には集計値だけを入れる**．画像単位の値・患者 ID・ファイル名は結果にもログにも出さない
（CLAUDE.md 必須ルール）．`ImageStats` は画像単位の値を持つため，関数の外（ファイル・ログ）へ
そのまま出さないこと．
"""

from __future__ import annotations

import logging
import math
from collections import Counter
from collections.abc import Callable, Iterable, Sequence
from dataclasses import dataclass
from pathlib import Path
from typing import Any

import numpy as np
from PIL import Image

from tias_judge_model.data.classes import ALL_DIR_NAMES, CLASS_SLUGS, SEVERITIES, parse_dir_name
from tias_judge_model.data.color import delta_e_cie76, mean_tongue_lab, tongue_mask
from tias_judge_model.data.manifest import ManifestRecord, resolve_image_path
from tias_judge_model.data.split import (
    Fold,
    SplitResult,
    count_by_class,
    count_by_group,
    make_kfold_splits,
    record_group,
    split_records,
    split_train_val_test,
)

logger = logging.getLogger(__name__)

SUMMARY_SCHEMA_VERSION = 1

# 「ややあり」を入れる／入れない実験の 2 条件（SPEC_02「程度（ややあり）の扱い」）．
SEVERITY_CONDITIONS: dict[str, tuple[str, ...]] = {
    "with_mild": ("strong", "mild", "normal"),
    "without_mild": ("strong", "normal"),
}

# mild グループと比べる「淡紅（正常）」のグループ．
NORMAL_GROUP = "light_red_normal"

LAB_CHANNELS: tuple[str, ...] = ("L", "a", "b")

_PROGRESS_EVERY = 100
_ROUND_DIGITS = 4


@dataclass(frozen=True)
class ImageStats:
    """1 枚の画像から測った値（集計の入力．**ファイル・ログへそのまま出さない**）.

    Attributes:
        readable: 画像として読めたか．
        missing: ファイルが存在しなかったか（`readable` が False のときのみ意味を持つ）．
        mode: PIL の画像モード（例: `"RGB"`）．読めなければ None．
        has_alpha: アルファチャンネル（透過情報）を持つか．
        width: 幅（px）．読めなければ 0．
        height: 高さ（px）．読めなければ 0．
        tongue_ratio: 画像に占める舌領域（背景でない画素）の割合．読めなければ None．
        lab: 舌領域の平均 L*・a*・b*．読めない，または舌領域が無ければ None．
    """

    readable: bool
    missing: bool = False
    mode: str | None = None
    has_alpha: bool = False
    width: int = 0
    height: int = 0
    tongue_ratio: float | None = None
    lab: tuple[float, float, float] | None = None


def measure_pixels(pixels: np.ndarray, *, mode: str, has_alpha: bool) -> ImageStats:
    """RGB 画素から画像 1 枚の値を測る.

    Args:
        pixels: RGB 画像（`(H, W, 3)`, uint8）．
        mode: 元の画像モード（記録用）．
        has_alpha: 元の画像がアルファチャンネルを持つか（記録用）．

    Returns:
        測定値．

    Raises:
        ValueError: `pixels` の形状が `(H, W, 3)` でない場合．
    """
    arr = np.asarray(pixels)
    if arr.ndim != 3 or arr.shape[2] != 3:
        raise ValueError(f"画素の形状は (H, W, 3) でなければならない（{arr.shape}）")
    height, width = arr.shape[:2]
    mask = tongue_mask(arr)
    ratio = float(mask.mean()) if mask.size else 0.0
    lab = mean_tongue_lab(arr)
    return ImageStats(
        readable=True,
        mode=mode,
        has_alpha=has_alpha,
        width=int(width),
        height=int(height),
        tongue_ratio=ratio,
        lab=None if lab is None else (float(lab[0]), float(lab[1]), float(lab[2])),
    )


def read_image_stats(path: Path) -> ImageStats:
    """画像ファイルを読んで値を測る（読めなければ `readable=False` を返し，止めない）.

    集計は現状把握のためであり，学習時の fail fast とは目的が異なるため，読めない画像は
    件数として数えるだけにする（SPEC_01「データセット集計の出力」）．

    Args:
        path: 画像ファイルのパス．

    Returns:
        測定値．
    """
    if not path.is_file():
        return ImageStats(readable=False, missing=True)
    try:
        with Image.open(path) as img:
            img.load()
            mode = img.mode
            has_alpha = "A" in img.getbands() or "transparency" in img.info
            rgb = img if mode == "RGB" else img.convert("RGB")
            pixels = np.asarray(rgb, dtype=np.uint8)  # (H, W, 3)
    except (OSError, ValueError, Image.DecompressionBombError):
        # 破損・非対応形式・巨大すぎる画像．どの画像かはログに出さない（ファイル名を出さないため）．
        return ImageStats(readable=False)
    return measure_pixels(pixels, mode=mode, has_alpha=has_alpha)


def collect_image_stats(
    records: Sequence[ManifestRecord],
    images_dir: Path,
    reader: Callable[[Path], ImageStats] = read_image_stats,
) -> list[ImageStats]:
    """manifest の全画像を読んで値を測る（進捗は件数だけをログに出す）.

    Args:
        records: manifest レコード．
        images_dir: 画像ディレクトリ（`<data_root>/images`）．
        reader: 1 枚を測る関数（テストで差し替えるため）．

    Returns:
        `records` と同じ順の測定値．
    """
    stats: list[ImageStats] = []
    total = len(records)
    for i, record in enumerate(records, start=1):
        stats.append(reader(resolve_image_path(images_dir, record)))
        if i % _PROGRESS_EVERY == 0 or i == total:
            logger.info("画像を読んだ: %d / %d 件", i, total)
    unreadable = sum(1 for s in stats if not s.readable)
    if unreadable:
        logger.warning("読めなかった画像が %d 件ある（集計からは除いて続行する）", unreadable)
    return stats


def lab_values_by_group(
    records: Sequence[ManifestRecord], stats: Sequence[ImageStats]
) -> dict[str, np.ndarray]:
    """グループごとに，画像ごとの舌領域の平均 Lab を集める（図の描画用）.

    Args:
        records: manifest レコード．
        stats: `records` と同じ順の測定値．

    Returns:
        グループ名（`ALL_DIR_NAMES` の順）→ 平均 Lab の配列（`(N, 3)`, float64．N は
        Lab が得られた画像数で 0 もありうる）．

    Raises:
        ValueError: `records` と `stats` の長さが異なる場合．
    """
    _check_same_length(records, stats)
    values: dict[str, list[tuple[float, float, float]]] = {g: [] for g in ALL_DIR_NAMES}
    for record, stat in zip(records, stats, strict=True):
        if stat.lab is not None:
            values[record_group(record)].append(stat.lab)
    return {g: np.array(v, dtype=np.float64).reshape(-1, 3) for g, v in values.items()}


def summarize_dataset(
    records: Sequence[ManifestRecord],
    stats: Sequence[ImageStats],
    *,
    seed: int,
    val_ratio: float,
    test_ratio: float,
    n_folds: int,
    configured_severities: Iterable[str],
) -> dict[str, Any]:
    """データセットの集計値を計算する（JSON にできる dict を返す純関数）.

    結果には集計値だけを入れ，画像単位の値・患者 ID・ファイル名は入れない．

    Args:
        records: manifest レコード（検証済み．除外前の全体）．
        stats: `records` と同じ順の測定値．
        seed: 分割の seed．
        val_ratio: 単一分割の val の比率．
        test_ratio: 単一分割の test の比率．
        n_folds: 交差検証の fold 数．
        configured_severities: 設定の `data.severities`（記録用）．

    Returns:
        集計結果（キーは `counts`・`loading`・`size`・`tongue_ratio`・`lab`・`delta_e`・
        `split` 等．構造は SPEC_01「データセット集計の出力」）．

    Raises:
        ValueError: `records` と `stats` の長さが異なる場合．
    """
    _check_same_length(records, stats)
    lab_by_group = lab_values_by_group(records, stats)
    lab_summary = _lab_summary(records, stats, lab_by_group)
    return {
        "schema_version": SUMMARY_SCHEMA_VERSION,
        "configured_severities": list(configured_severities),
        "counts": _count_summary(records),
        "loading": _loading_summary(records, stats),
        "size": _size_summary(stats),
        "tongue_ratio": _tongue_ratio_summary(records, stats),
        "lab": lab_summary,
        "delta_e": _delta_e_summary(lab_by_group),
        "split": _split_summary(
            records, seed=seed, val_ratio=val_ratio, test_ratio=test_ratio, n_folds=n_folds
        ),
    }


# ---------------------------------------------------------------------------
# 件数
# ---------------------------------------------------------------------------


def _count_summary(records: Sequence[ManifestRecord]) -> dict[str, Any]:
    def images_patients(rs: Iterable[ManifestRecord]) -> dict[str, int]:
        rs = list(rs)
        return {"images": len(rs), "patients": len({r.patient_id for r in rs})}

    groups: dict[str, Any] = {}
    for group in ALL_DIR_NAMES:
        label, severity = parse_dir_name(group)
        groups[group] = {
            "label": label,
            "severity": severity,
            **images_patients(r for r in records if record_group(r) == group),
        }
    conditions: dict[str, Any] = {}
    for name, severities in SEVERITY_CONDITIONS.items():
        kept = [r for r in records if r.severity in severities]
        conditions[name] = {
            "severities": list(severities),
            **images_patients(kept),
            "classes": count_by_class(kept),
        }
    return {
        "total": images_patients(records),
        "groups": groups,
        "classes": {c: images_patients(r for r in records if r.label == c) for c in CLASS_SLUGS},
        "severities": {
            s: images_patients(r for r in records if r.severity == s) for s in SEVERITIES
        },
        "conditions": conditions,
    }


# ---------------------------------------------------------------------------
# 読み込み・サイズ・舌領域の割合
# ---------------------------------------------------------------------------


def _loading_summary(
    records: Sequence[ManifestRecord], stats: Sequence[ImageStats]
) -> dict[str, Any]:
    readable = [s for s in stats if s.readable]
    unreadable_by_group = Counter(
        record_group(r) for r, s in zip(records, stats, strict=True) if not s.readable
    )
    modes = Counter(s.mode for s in readable)
    return {
        "readable": len(readable),
        "unreadable": len(stats) - len(readable),
        "missing": sum(1 for s in stats if not s.readable and s.missing),
        "unreadable_by_group": {g: unreadable_by_group.get(g, 0) for g in ALL_DIR_NAMES},
        "modes": dict(sorted(modes.items())),
        "non_rgb_mode": sum(1 for s in readable if s.mode != "RGB"),
        "with_alpha": sum(1 for s in readable if s.has_alpha),
        "no_tongue_pixels": sum(1 for s in readable if s.lab is None),
    }


def _size_summary(stats: Sequence[ImageStats]) -> dict[str, Any]:
    readable = [s for s in stats if s.readable]
    return {
        "width": _min_max_mean_median([s.width for s in readable]),
        "height": _min_max_mean_median([s.height for s in readable]),
    }


def _tongue_ratio_summary(
    records: Sequence[ManifestRecord], stats: Sequence[ImageStats]
) -> dict[str, Any]:
    by_group: dict[str, list[float]] = {g: [] for g in ALL_DIR_NAMES}
    for record, stat in zip(records, stats, strict=True):
        if stat.tongue_ratio is not None:
            by_group[record_group(record)].append(stat.tongue_ratio)
    all_values = [v for values in by_group.values() for v in values]
    return {
        "total": _min_mean_max(all_values),
        "groups": {g: _min_mean_max(v) for g, v in by_group.items()},
    }


# ---------------------------------------------------------------------------
# 色（Lab）と色差
# ---------------------------------------------------------------------------


def _lab_summary(
    records: Sequence[ManifestRecord],
    stats: Sequence[ImageStats],
    lab_by_group: dict[str, np.ndarray],
) -> dict[str, Any]:
    by_class: dict[str, list[tuple[float, float, float]]] = {c: [] for c in CLASS_SLUGS}
    for record, stat in zip(records, stats, strict=True):
        if stat.lab is not None:
            by_class[record.label].append(stat.lab)
    all_values = np.concatenate(list(lab_by_group.values()), axis=0)
    return {
        "method": (
            "画像ごとに舌領域（RGB がすべて 0 でない画素）の L*a*b*（sRGB→XYZ D65→Lab D65）を"
            "平均し，その値をグループ・クラス・全体で集計した"
        ),
        "groups": {g: _lab_stats(v) for g, v in lab_by_group.items()},
        "classes": {
            c: _lab_stats(np.array(v, dtype=np.float64).reshape(-1, 3)) for c, v in by_class.items()
        },
        "total": _lab_stats(all_values),
    }


def _lab_stats(values: np.ndarray) -> dict[str, Any]:
    """平均 Lab の配列（`(N, 3)`）からチャンネルごとの統計を作る（N = 0 なら値は None）."""
    result: dict[str, Any] = {"n": int(values.shape[0])}
    for i, channel in enumerate(LAB_CHANNELS):
        result[channel] = _distribution(values[:, i])
    return result


def _distribution(values: np.ndarray) -> dict[str, float | None]:
    if values.size == 0:
        return dict.fromkeys(("mean", "std", "median", "p10", "p90"))
    # 標本標準偏差（ddof=1）．1 件しかなければ求められないため None．
    std = float(np.std(values, ddof=1)) if values.size > 1 else None
    return {
        "mean": _round(float(np.mean(values))),
        "std": _round(std),
        "median": _round(float(np.median(values))),
        "p10": _round(float(np.percentile(values, 10))),
        "p90": _round(float(np.percentile(values, 90))),
    }


def _delta_e_summary(lab_by_group: dict[str, np.ndarray]) -> dict[str, Any]:
    # 丸める前の平均から求める（丸め誤差を色差に持ち込まないため）．
    means = {g: (v.mean(axis=0) if v.shape[0] else None) for g, v in lab_by_group.items()}

    def delta(g1: str, g2: str) -> float | None:
        m1, m2 = means.get(g1), means.get(g2)
        if m1 is None or m2 is None:
            return None
        return _round(delta_e_cie76(m1, m2))

    matrix = {g1: {g2: delta(g1, g2) for g2 in ALL_DIR_NAMES} for g1 in ALL_DIR_NAMES}
    mild_comparison = []
    for group in ALL_DIR_NAMES:
        label, severity = parse_dir_name(group)
        if severity != "mild":
            continue
        strong_group = f"{label}_strong"
        to_normal = delta(group, NORMAL_GROUP)
        to_strong = delta(group, strong_group)
        closer: str | None
        if to_normal is None or to_strong is None:
            closer = None
        elif to_normal < to_strong:
            closer = NORMAL_GROUP
        elif to_strong < to_normal:
            closer = strong_group
        else:
            closer = "tie"
        mild_comparison.append(
            {
                "group": group,
                "normal_group": NORMAL_GROUP,
                "delta_e_to_normal": to_normal,
                "strong_group": strong_group,
                "delta_e_to_strong": to_strong,
                "closer_to": closer,
            }
        )
    return {
        "method": "CIE76（グループの平均 Lab どうしのユークリッド距離）",
        "mild_comparison": mild_comparison,
        "matrix": matrix,
    }


# ---------------------------------------------------------------------------
# 分割の確認
# ---------------------------------------------------------------------------


def _split_summary(
    records: Sequence[ManifestRecord],
    *,
    seed: int,
    val_ratio: float,
    test_ratio: float,
    n_folds: int,
) -> dict[str, Any]:
    result: dict[str, Any] = {
        "seed": seed,
        "val_ratio": val_ratio,
        "test_ratio": test_ratio,
        "n_folds": n_folds,
    }
    # 分割できない構成でも集計自体は現状把握として役に立つため，理由を記録して続ける
    # （メッセージはグループ名と人数のみで，患者 ID を含まない）．
    try:
        single = split_train_val_test(
            records, val_ratio=val_ratio, test_ratio=test_ratio, seed=seed
        )
    except ValueError as exc:
        result["single"] = {"error": str(exc)}
    else:
        result["single"] = _parts_summary(records, single)
    try:
        folds = make_kfold_splits(records, n_folds=n_folds, seed=seed)
    except ValueError as exc:
        result["kfold"] = {"error": str(exc)}
    else:
        result["kfold"] = [{"index": fold.index, **_parts_summary(records, fold)} for fold in folds]
    return result


def _parts_summary(records: Sequence[ManifestRecord], split: SplitResult | Fold) -> dict[str, Any]:
    everything = split_records(records, split)
    by_condition = {
        name: split_records(records, split, severities)
        for name, severities in SEVERITY_CONDITIONS.items()
    }
    return {
        "groups": {part: count_by_group(rs) for part, rs in everything.items()},
        "patients": {part: len({r.patient_id for r in rs}) for part, rs in everything.items()},
        "conditions": {
            name: {
                part: {"images": len(rs), "classes": count_by_class(rs)}
                for part, rs in parts.items()
            }
            for name, parts in by_condition.items()
        },
    }


# ---------------------------------------------------------------------------
# 補助
# ---------------------------------------------------------------------------


def _check_same_length(records: Sequence[ManifestRecord], stats: Sequence[ImageStats]) -> None:
    if len(records) != len(stats):
        raise ValueError(f"レコード数 {len(records)} 件と測定値の数 {len(stats)} 件が一致しない")


def _min_max_mean_median(values: Sequence[float]) -> dict[str, float | None]:
    if not values:
        return dict.fromkeys(("min", "max", "mean", "median"))
    arr = np.asarray(values, dtype=np.float64)
    return {
        "min": _round(float(arr.min())),
        "max": _round(float(arr.max())),
        "mean": _round(float(arr.mean())),
        "median": _round(float(np.median(arr))),
    }


def _min_mean_max(values: Sequence[float]) -> dict[str, float | None]:
    if not values:
        return dict.fromkeys(("min", "mean", "max"))
    arr = np.asarray(values, dtype=np.float64)
    return {
        "min": _round(float(arr.min())),
        "mean": _round(float(arr.mean())),
        "max": _round(float(arr.max())),
    }


def _round(value: float | None) -> float | None:
    if value is None or not math.isfinite(value):
        return None
    return round(value, _ROUND_DIGITS)