"""クラス不均衡への対策の計算（SPEC_02「クラス不均衡」→「対策の設定」）.

学習ラベルから，損失のクラス重みとバランス・サンプリングの抜き出し重みを求める純関数．
クラスの画像数 n_c は呼び出し側が渡す**各 fold の学習用データ（`data.severities` 適用後）**の
ラベルから数える（検証用・held_out は使わない）．
"""

from __future__ import annotations

import numpy as np

from tias_judge_model.data.classes import CLASS_SLUGS, NUM_CLASSES

# 損失のクラス重みの方式（`train.class_weighting` のうち `none` 以外）．
CLASS_WEIGHTING_METHODS: tuple[str, ...] = ("inverse", "sqrt_inverse")


def count_classes(labels: np.ndarray) -> np.ndarray:
    """クラス ID ごとの画像数を数える.

    Args:
        labels: `(N,)` の整数配列（クラス ID．`classes.py` の固定 ID）．

    Returns:
        `(NUM_CLASSES,)` の int64 配列（クラス ID 順の画像数）．

    Raises:
        ValueError: 範囲外のクラス ID を含む場合．
    """
    labels = np.asarray(labels, dtype=np.int64)
    if labels.size and (labels.min() < 0 or labels.max() >= NUM_CLASSES):
        raise ValueError(f"クラス ID が範囲外である（0〜{NUM_CLASSES - 1}）")
    return np.bincount(labels, minlength=NUM_CLASSES).astype(np.int64)


def _require_all_classes(counts: np.ndarray) -> None:
    missing = [CLASS_SLUGS[i] for i, n in enumerate(counts) if n == 0]
    if missing:
        raise ValueError(
            f"学習用データに 1 枚も無いクラスがあるため，クラス不均衡の対策を計算できない: "
            f"{', '.join(missing)}"
        )


def compute_class_weights(labels: np.ndarray, method: str) -> np.ndarray:
    """学習ラベルから損失のクラス重みを求める.

    クラス c の画像数を n_c，全体を N，クラス数を K として，`inverse` は w_c ∝ N / (K・n_c)，
    `sqrt_inverse` はその平方根に比例する．学習データでの平均が 1 になるよう
    Σ_c (n_c / N)・w_c = 1 に正規化する（損失の大きさを対策なしと揃えるため）．

    Args:
        labels: `(N,)` の整数配列（学習用データのクラス ID）．
        method: `inverse` または `sqrt_inverse`．

    Returns:
        `(NUM_CLASSES,)` の float64 配列（クラス ID 順の重み）．

    Raises:
        ValueError: 未知の方式，または学習用データに 1 枚も無いクラスがある場合．
    """
    if method not in CLASS_WEIGHTING_METHODS:
        valid = ", ".join(CLASS_WEIGHTING_METHODS)
        raise ValueError(f"未知のクラス重みの方式である: {method!r}（有効値: {valid}）")
    counts = count_classes(labels)
    _require_all_classes(counts)
    n = counts.astype(np.float64)
    total = n.sum()
    raw = total / (NUM_CLASSES * n)
    if method == "sqrt_inverse":
        raw = np.sqrt(raw)
    return raw / float(np.sum(n / total * raw))


def balanced_sample_weights(labels: np.ndarray) -> np.ndarray:
    """バランス・サンプリングの画像ごとの抜き出し重み（1 / n_c）を求める.

    重みの合計は K（クラス数）になり，どのクラスも 1 回の抜き出しで 1 / K の確率で選ばれる．

    Args:
        labels: `(N,)` の整数配列（学習用データのクラス ID）．

    Returns:
        `(N,)` の float64 配列（各画像の抜き出し重み．並びは `labels` の順）．

    Raises:
        ValueError: 学習用データに 1 枚も無いクラスがある場合．
    """
    labels = np.asarray(labels, dtype=np.int64)
    counts = count_classes(labels)
    _require_all_classes(counts)
    return 1.0 / counts[labels].astype(np.float64)


def class_draw_probabilities(labels: np.ndarray, sample_weights: np.ndarray) -> np.ndarray:
    """画像ごとの抜き出し重みから，1 回の抜き出しで各クラスが選ばれる確率を求める.

    Args:
        labels: `(N,)` の整数配列（学習用データのクラス ID）．
        sample_weights: `(N,)` の非負の抜き出し重み．

    Returns:
        `(NUM_CLASSES,)` の float64 配列（クラス ID 順．合計 1）．

    Raises:
        ValueError: 件数が一致しない，または重みの合計が 0 以下の場合．
    """
    labels = np.asarray(labels, dtype=np.int64)
    weights = np.asarray(sample_weights, dtype=np.float64)
    if labels.shape != weights.shape:
        raise ValueError("ラベルと抜き出し重みの件数が一致しない")
    total = weights.sum()
    if total <= 0:
        raise ValueError("抜き出し重みの合計が 0 以下である")
    per_class = np.bincount(labels, weights=weights, minlength=NUM_CLASSES)
    return per_class / total