"""色特徴量の抽出（`data/color_features.py`）のテスト.

仕様（実装サマリーの設計を仕様として扱う）:

- `extract_color_features` は 107 次元（L*・a*・b*・C*・h ごとに平均・標準偏差・10/25/50/75/90
  分位点 = 35，a*-b* 8×8 ヒストグラム = 64，L* 8 ビン = 8）．並びは `color_feature_names`
- h は atan2(b*, a*) の度（(-180, 180]）
- ヒストグラムは割合（合計 1）．範囲（a* [0, 50]・b* [-10, 30]・L* [20, 100]）外は端のビン
- 舌領域（RGB がすべて 0 でない画素）だけを使う（背景の量・位置で特徴が変わらない）
- 舌領域が 0 画素なら ValueError（ファイル名は出さない）
- `downscale_nearest` は長辺が max_side を超えるときだけ最近傍法で縮小する（色が変わらない）
"""

from __future__ import annotations

import numpy as np
import pytest

from tias_judge_model.data.color import rgb8_to_lab
from tias_judge_model.data.color_features import (
    NUM_COLOR_FEATURES,
    color_feature_names,
    downscale_nearest,
    extract_color_features,
)

NAMES = color_feature_names()
AB_BINS = 8
L_BINS = 8
A_RANGE = (0.0, 50.0)
B_RANGE = (-10.0, 30.0)
L_RANGE = (20.0, 100.0)


def _idx(name: str) -> int:
    return NAMES.index(name)


def _ab_hist(features: np.ndarray) -> np.ndarray:
    return np.array(
        [features[_idx(f"ab_hist_{i}_{j}")] for i in range(AB_BINS) for j in range(AB_BINS)]
    ).reshape(AB_BINS, AB_BINS)


def _l_hist(features: np.ndarray) -> np.ndarray:
    return np.array([features[_idx(f"L_hist_{k}")] for k in range(L_BINS)])


def _uniform(rgb: tuple[int, int, int], shape: tuple[int, int] = (6, 5)) -> np.ndarray:
    img = np.zeros((*shape, 3), dtype=np.uint8)
    img[:] = rgb
    return img


def _lab(rgb: tuple[int, int, int]) -> np.ndarray:
    return rgb8_to_lab(np.array([rgb], dtype=np.uint8))[0]


def _random_tongue(rng: np.random.Generator, shape: tuple[int, int]) -> np.ndarray:
    # 1〜255 の値だけにして，すべての画素を舌領域にする
    return rng.integers(1, 256, size=(*shape, 3), dtype=np.uint8)


# --- 次元・名前 ---


def test_dimension_is_107() -> None:
    assert NUM_COLOR_FEATURES == 107
    assert len(NAMES) == 107
    rng = np.random.default_rng(0)
    f = extract_color_features(_random_tongue(rng, (8, 9)))
    assert f.shape == (107,)
    assert f.dtype == np.float64
    assert np.all(np.isfinite(f))


def test_feature_names_are_unique_and_follow_layout() -> None:
    assert len(set(NAMES)) == len(NAMES)
    stats = ["mean", "std", "p10", "p25", "p50", "p75", "p90"]
    expected_channel = [f"{ch}_{s}" for ch in ("L", "a", "b", "C", "h") for s in stats]
    assert list(NAMES[:35]) == expected_channel
    assert list(NAMES[35:99]) == [f"ab_hist_{i}_{j}" for i in range(8) for j in range(8)]
    assert list(NAMES[99:]) == [f"L_hist_{k}" for k in range(8)]


# --- 統計量 ---


def test_uniform_color_statistics_match_lab() -> None:
    rgb = (225, 140, 140)
    f = extract_color_features(_uniform(rgb))
    lightness, a, b = _lab(rgb)
    chroma = np.hypot(a, b)
    hue = np.degrees(np.arctan2(b, a))
    for ch, value in zip("LabCh", (lightness, a, b, chroma, hue), strict=True):
        assert f[_idx(f"{ch}_mean")] == pytest.approx(value, abs=1e-6)
        assert f[_idx(f"{ch}_std")] == pytest.approx(0.0, abs=1e-6)
        for q in (10, 25, 50, 75, 90):
            assert f[_idx(f"{ch}_p{q}")] == pytest.approx(value, abs=1e-6)


def test_mean_and_median_over_mixed_pixels() -> None:
    # 9 画素が色 B，1 画素が色 A: 平均は重み付き，中央値・p90 は B
    img = np.zeros((1, 10, 3), dtype=np.uint8)
    color_a, color_b = (120, 20, 30), (225, 140, 140)
    img[0, 0] = color_a
    img[0, 1:] = color_b
    f = extract_color_features(img)
    la, lb = _lab(color_a)[0], _lab(color_b)[0]
    assert f[_idx("L_mean")] == pytest.approx((la + 9 * lb) / 10, abs=1e-6)
    assert f[_idx("L_p50")] == pytest.approx(lb, abs=1e-6)
    assert f[_idx("L_p90")] == pytest.approx(lb, abs=1e-6)
    assert f[_idx("L_std")] > 0
    # 分位点は単調非減少
    ps = [f[_idx(f"L_p{q}")] for q in (10, 25, 50, 75, 90)]
    assert ps == sorted(ps)


@pytest.mark.parametrize(
    "rgb",
    [(225, 140, 140), (110, 50, 140), (0, 255, 0), (0, 0, 255), (40, 200, 220)],
    ids=str,
)
def test_hue_is_atan2_degrees_in_range(rgb: tuple[int, int, int]) -> None:
    f = extract_color_features(_uniform(rgb))
    _, a, b = _lab(rgb)
    h = f[_idx("h_mean")]
    assert -180.0 < h <= 180.0
    assert h == pytest.approx(np.degrees(np.arctan2(b, a)), abs=1e-6)


def test_hue_sign_follows_b() -> None:
    # 紫（a* > 0, b* < 0）は負の色相，淡紅（a* > 0, b* > 0）は 0〜90°
    purple = extract_color_features(_uniform((110, 50, 140)))[_idx("h_mean")]
    pink = extract_color_features(_uniform((225, 140, 140)))[_idx("h_mean")]
    assert purple < 0
    assert 0 < pink < 90


# --- ヒストグラム ---


def test_histograms_sum_to_one_and_are_proportions() -> None:
    rng = np.random.default_rng(1)
    img = _random_tongue(rng, (13, 17))
    img[:3] = 0  # 背景
    f = extract_color_features(img)
    ab = _ab_hist(f)
    lh = _l_hist(f)
    assert ab.sum() == pytest.approx(1.0)
    assert lh.sum() == pytest.approx(1.0)
    assert np.all(ab >= 0) and np.all(lh >= 0)
    # 割合は舌領域の画素数の倍数
    n = 10 * 17
    assert np.allclose(ab * n, np.round(ab * n))
    assert np.allclose(lh * n, np.round(lh * n))


def test_histogram_proportions_three_to_one() -> None:
    img = np.zeros((2, 3, 3), dtype=np.uint8)  # 2 画素は背景
    img[0, :3] = (225, 140, 140)
    img[1, 0] = (255, 255, 255)
    f = extract_color_features(img)
    lh = _l_hist(f)
    assert lh[7] == pytest.approx(0.25)  # 白（L* = 100）は最上位ビン
    assert sorted(lh[lh > 0].tolist()) == pytest.approx([0.25, 0.75])


def test_in_range_color_goes_to_expected_ab_bin() -> None:
    rgb = (225, 140, 140)
    lightness, a, b = _lab(rgb)
    assert A_RANGE[0] < a < A_RANGE[1] and B_RANGE[0] < b < B_RANGE[1]
    i = int((a - A_RANGE[0]) // ((A_RANGE[1] - A_RANGE[0]) / AB_BINS))
    j = int((b - B_RANGE[0]) // ((B_RANGE[1] - B_RANGE[0]) / AB_BINS))
    k = int((lightness - L_RANGE[0]) // ((L_RANGE[1] - L_RANGE[0]) / L_BINS))
    f = extract_color_features(_uniform(rgb))
    ab = _ab_hist(f)
    assert ab[i, j] == pytest.approx(1.0)
    assert _l_hist(f)[k] == pytest.approx(1.0)


@pytest.mark.parametrize(
    ("rgb", "ab_bin"),
    [
        ((255, 0, 0), (7, 7)),  # a* ≈ 80・b* ≈ 67（両方上限超え）
        ((0, 255, 0), (0, 7)),  # a* ≈ -86（下限未満）・b* ≈ 83（上限超え）
        ((0, 0, 255), (7, 0)),  # a* ≈ 79（上限超え）・b* ≈ -108（下限未満）
    ],
    ids=["red", "green", "blue"],
)
def test_out_of_range_ab_goes_to_edge_bin(
    rgb: tuple[int, int, int], ab_bin: tuple[int, int]
) -> None:
    f = extract_color_features(_uniform(rgb))
    ab = _ab_hist(f)
    assert ab[ab_bin] == pytest.approx(1.0)
    assert ab.sum() == pytest.approx(1.0)


def test_out_of_range_lightness_goes_to_edge_bin() -> None:
    img = np.zeros((1, 4, 3), dtype=np.uint8)
    img[0, 0] = (5, 5, 5)  # L* ≈ 1.4（下限 20 未満）→ 最下位ビン
    img[0, 1] = (3, 3, 3)
    img[0, 2] = (255, 255, 255)  # L* = 100（上限ちょうど）→ 最上位ビン
    img[0, 3] = (255, 255, 255)
    lh = _l_hist(extract_color_features(img))
    assert lh[0] == pytest.approx(0.5)
    assert lh[7] == pytest.approx(0.5)
    assert lh.sum() == pytest.approx(1.0)


# --- 背景の扱い ---


def test_background_amount_and_position_do_not_change_features() -> None:
    rng = np.random.default_rng(2)
    tongue = _random_tongue(rng, (10, 12))
    small = tongue.copy()
    large = np.zeros((40, 55, 3), dtype=np.uint8)
    large[17:27, 30:42] = tongue
    np.testing.assert_allclose(
        extract_color_features(small), extract_color_features(large), rtol=0, atol=1e-9
    )


def test_black_holes_inside_tongue_are_ignored() -> None:
    rng = np.random.default_rng(3)
    tongue = _random_tongue(rng, (1, 20))
    with_holes = np.zeros((3, 20, 3), dtype=np.uint8)
    with_holes[1] = tongue
    # 舌の画素の順序を変えても（並べ替え不変），黒を挟んでも同じ
    shuffled = np.zeros((20, 2, 3), dtype=np.uint8)
    shuffled[:, 0] = tongue[0][::-1]
    np.testing.assert_allclose(
        extract_color_features(tongue), extract_color_features(with_holes), atol=1e-9
    )
    np.testing.assert_allclose(
        extract_color_features(tongue), extract_color_features(shuffled), atol=1e-9
    )


def test_nearly_black_pixel_counts_as_tongue() -> None:
    # RGB がすべて 0 でなければ舌領域（1 チャンネルだけ 1 でも数える）
    img = np.zeros((3, 3, 3), dtype=np.uint8)
    img[1, 1] = (0, 0, 1)
    f = extract_color_features(img)
    assert _l_hist(f).sum() == pytest.approx(1.0)
    assert f[_idx("L_std")] == pytest.approx(0.0, abs=1e-9)


def test_no_tongue_pixels_raises_without_file_name() -> None:
    with pytest.raises(ValueError) as excinfo:
        extract_color_features(np.zeros((5, 5, 3), dtype=np.uint8))
    msg = str(excinfo.value)
    assert ".png" not in msg and ".jpg" not in msg


@pytest.mark.parametrize("shape", [(5, 5), (5, 5, 4), (5, 5, 1)], ids=str)
def test_invalid_shape_raises(shape: tuple[int, ...]) -> None:
    with pytest.raises(ValueError):
        extract_color_features(np.full(shape, 100, dtype=np.uint8))


# --- 最近傍縮小 ---


def test_downscale_keeps_small_images_unchanged() -> None:
    rng = np.random.default_rng(4)
    img = _random_tongue(rng, (30, 50))
    np.testing.assert_array_equal(downscale_nearest(img, 50), img)
    np.testing.assert_array_equal(downscale_nearest(img, 100), img)


@pytest.mark.parametrize(
    ("shape", "max_side", "expected"),
    [((40, 100), 50, (20, 50)), ((100, 30), 50, (50, 15)), ((64, 64), 16, (16, 16))],
)
def test_downscale_limits_long_side_keeping_aspect(
    shape: tuple[int, int], max_side: int, expected: tuple[int, int]
) -> None:
    rng = np.random.default_rng(5)
    out = downscale_nearest(_random_tongue(rng, shape), max_side)
    assert out.shape == (*expected, 3)
    assert out.dtype == np.uint8
    assert max(out.shape[:2]) <= max_side


def test_downscale_does_not_create_new_colors() -> None:
    rng = np.random.default_rng(6)
    img = np.zeros((90, 70, 3), dtype=np.uint8)
    img[10:80, 10:60] = _random_tongue(rng, (70, 50))
    out = downscale_nearest(img, 23)
    src_colors = {tuple(p) for p in img.reshape(-1, 3)}
    out_colors = {tuple(p) for p in out.reshape(-1, 3)}
    assert out_colors <= src_colors


def test_downscale_of_block_image_keeps_features() -> None:
    # 2×2 の一様なブロックからなる画像を 1/2 に縮小しても，色の割合は変わらない
    rng = np.random.default_rng(7)
    small = _random_tongue(rng, (16, 24))
    small[:4, :4] = 0  # 背景も含める
    big = np.repeat(np.repeat(small, 2, axis=0), 2, axis=1)  # (32, 48, 3)
    out = downscale_nearest(big, 24)
    assert out.shape == (16, 24, 3)
    # 縮小後は各ブロックの色がちょうど 1 画素ずつ残る（色が混ざらず，割合も保たれる）
    np.testing.assert_array_equal(out, small)
    # 分位点は線形補間のため重複度で僅かに変わりうるので，平均・ヒストグラムで比べる
    keep = [i for i, n in enumerate(NAMES) if not n.split("_")[-1].startswith("p")]
    f_out, f_big = extract_color_features(out), extract_color_features(big)
    np.testing.assert_allclose(f_out[keep], f_big[keep], atol=1e-9)


@pytest.mark.parametrize("max_side", [0, -1])
def test_downscale_rejects_non_positive_max_side(max_side: int) -> None:
    with pytest.raises(ValueError):
        downscale_nearest(np.ones((4, 4, 3), dtype=np.uint8), max_side)