"""AI 舌色判定の確率を，既存の診断結果と同じ形（5 色の割合・上位 2 色）にする．

画面（メイン画面の「診断結果」欄・結果画面）はルールベース診断と共通にし，推論だけを
差し替えるため（Issue #13）．判定モデル（`color_model.json`）のクラスを既存の 5 色のキー
（`core.svm.COLOR_KEYS`）に割り当てる．

本モジュールは onnxruntime を import しない（結果画面から型として参照するため軽く保つ）．
"""

from __future__ import annotations

from collections.abc import Mapping
from dataclasses import dataclass

from tiasjudge.core.svm import COLOR_KEYS

#: 判定モデルのクラス（`color_model.json` の `classes` の `slug`）→ 既存の 5 色のキー．
#: 表示名は 淡白→淡白・淡紅→淡紅・紅→紅・暗赤紅→暗赤・紫→紫 に対応する
CLASS_SLUG_TO_COLOR_KEY: dict[str, str] = {
    "pale": "pale",
    "light_red": "pink",
    "red": "crimson",
    "dark_red": "dark_red",
    "purple": "purple",
}

#: 確率（0〜1）を百分率にする倍率
PERCENT_SCALE = 100


@dataclass(frozen=True)
class MlDiagnosis:
    """AI 舌色判定の結果を既存の診断結果と同じ形にしたもの．

    `core.svm.DiagnosisResult` と同じ名前の属性（`scores`・`percentages`・`top_labels`・
    `failed`）を持ち，結果画面・メイン画面の表示にそのまま渡せる．L*a*b* は持たない
    （変換行列を通さないため）．

    Attributes:
        scores: 5 色の確率（0.0〜1.0．和は 1）．キーは `COLOR_KEYS`．
        percentages: 表示用の百分率（整数．色ごとに四捨五入するため和が 100 にならないことが
            ある．ルールベース診断の表示と同じ扱い）．
        top_labels: 確率の降順の上位 2 色のキー（同点は `COLOR_KEYS` の順）．
        failed: 常に False（判定できなかった場合は本結果を作らない）．
    """

    scores: dict[str, float]
    percentages: dict[str, int]
    top_labels: tuple[str, str]
    failed: bool = False


def build_ml_diagnosis(probabilities_by_slug: Mapping[str, float]) -> MlDiagnosis:
    """判定モデルのクラスごとの確率から，既存の形の結果を作る．

    Args:
        probabilities_by_slug: クラスの `slug` → 確率（5 クラスすべて）．

    Returns:
        5 色のキーに割り当てた結果．

    Raises:
        ValueError: 未知のクラス，またはクラスの過不足がある場合．
    """
    unknown = set(probabilities_by_slug) - set(CLASS_SLUG_TO_COLOR_KEY)
    if unknown:
        raise ValueError(f"判定モデルに未知のクラスがある: {sorted(unknown)}")
    scores = {
        CLASS_SLUG_TO_COLOR_KEY[slug]: float(value) for slug, value in probabilities_by_slug.items()
    }
    if set(scores) != set(COLOR_KEYS):
        raise ValueError(f"判定モデルのクラスが 5 色と対応しない: {sorted(scores)}")
    # 表示順を既存に揃える（dict の並びは `COLOR_KEYS`）
    scores = {key: scores[key] for key in COLOR_KEYS}
    percentages = {key: round(value * PERCENT_SCALE) for key, value in scores.items()}
    # 安定ソートなので，同点は `COLOR_KEYS` の順になる
    ordered = sorted(COLOR_KEYS, key=lambda key: scores[key], reverse=True)
    return MlDiagnosis(scores=scores, percentages=percentages, top_labels=(ordered[0], ordered[1]))