"""配布用の参照実装（`deploy/color_model.py`）のテスト.

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

- NumPy と標準ライブラリだけに依存する（本パッケージの他モジュール・sklearn・torch・pydantic・
  Pillow 等を import しない．TIASjudge にこのファイルだけをコピーして使うため）
- `downscale_nearest` は Pillow の `Image.resize(..., NEAREST)` による縮小と画素単位で完全一致する
- `extract_features` は学習側の `extract_color_features(downscale_nearest(...))` と完全一致する
- `predict_proba` は学習側の sklearn パイプラインの `predict_proba` と差 1e-12 以内
- JSON に書いて読み直しても予測は変わらない
- 形状・dtype の不正，舌領域が 0 画素，未対応の版・種類・キー欠落・形状不一致は ValueError
"""

from __future__ import annotations

import ast
import json
from pathlib import Path

import numpy as np
import pytest
from PIL import Image

from tias_judge_model.data.classes import CLASS_DISPLAY_NAMES, CLASS_SLUGS
from tias_judge_model.data.color_features import (
    NUM_COLOR_FEATURES,
    color_feature_names,
    extract_color_features,
)
from tias_judge_model.data.color_features import downscale_nearest as train_downscale
from tias_judge_model.deploy import color_model as deploy
from tias_judge_model.schema import ExperimentConfig
from tias_judge_model.training.color_cross_validation import fit_classifier, make_classifier
from tias_judge_model.training.color_export import (
    ColorExportResult,
    CSelection,
    build_color_model_document,
)

DEPLOY_PATH = Path(deploy.__file__)
ALLOWED_TOP_LEVEL_IMPORTS = {
    "__future__",
    "json",
    "os",
    "collections",
    "dataclasses",
    "pathlib",
    "typing",
    "numpy",
    "math",
}
MAX_SIDE = 48
SEED = 7


# --- 依存の制約 ---


def _imported_modules(path: Path) -> set[str]:
    tree = ast.parse(path.read_text(encoding="utf-8"))
    modules: set[str] = set()
    for node in ast.walk(tree):
        if isinstance(node, ast.Import):
            modules.update(alias.name for alias in node.names)
        elif isinstance(node, ast.ImportFrom):
            assert node.level == 0, "相対 import をしてはならない（単一ファイルで使うため）"
            assert node.module is not None
            modules.add(node.module)
    return modules


def test_deploy_module_imports_only_numpy_and_stdlib() -> None:
    modules = _imported_modules(DEPLOY_PATH)
    top = {m.split(".")[0] for m in modules}
    forbidden = {"tias_judge_model", "sklearn", "torch", "pydantic", "PIL", "cv2", "scipy"}
    assert not (top & forbidden), top & forbidden
    assert top <= ALLOWED_TOP_LEVEL_IMPORTS, top - ALLOWED_TOP_LEVEL_IMPORTS


def test_deploy_module_has_no_dynamic_imports() -> None:
    src = DEPLOY_PATH.read_text(encoding="utf-8")
    assert "importlib" not in src
    assert "__import__" not in src


# --- テスト用のモデル ---


def _config(max_side: int | None = MAX_SIDE, class_weighting: str = "none") -> ExperimentConfig:
    return ExperimentConfig.model_validate(
        {
            "data": {"root": "unused", "severities": ["strong", "normal"]},
            "split": {"seed": SEED, "n_folds": 5},
            "train": {"class_weighting": class_weighting},
            "color": {
                "classifier": "logistic",
                "max_side": max_side,
                "logistic": {"c_values": [0.5, 2.0], "max_iter": 2000},
            },
        }
    )


def _random_tongue(rng: np.random.Generator, height: int, width: int) -> np.ndarray:
    """楕円の舌領域（ランダムな赤み）＋背景の黒の RGB 画像."""
    yy, xx = np.mgrid[0:height, 0:width]
    inside = ((yy - height / 2) / (height * 0.45)) ** 2 + ((xx - width / 2) / (width * 0.45)) ** 2
    mask = inside <= 1.0
    base = rng.integers(60, 230, size=3)
    img = np.clip(base + rng.normal(0, 25, size=(height, width, 3)), 1, 255).astype(np.uint8)
    img[~mask] = 0
    return img


@pytest.fixture(scope="module")
def trained() -> dict:
    rng = np.random.default_rng(SEED)
    images = [
        _random_tongue(rng, int(rng.integers(30, 90)), int(rng.integers(30, 90))) for _ in range(40)
    ]
    config = _config()
    features = np.stack([extract_color_features(train_downscale(img, MAX_SIDE)) for img in images])
    labels = np.arange(len(images)) % 5
    assert config.color is not None
    pipeline = make_classifier(config.color, {"C": 1.0}, SEED)
    fit_classifier(pipeline, features, labels, None)
    result = ColorExportResult(
        model=pipeline,
        selection=CSelection(
            candidates=[
                {"C": 0.5, "fold_macro_f1": [0.1], "mean_macro_f1": 0.1},
                {"C": 1.0, "fold_macro_f1": [0.2], "mean_macro_f1": 0.2},
            ],
            selected_index=1,
            fold_counts=[{"train": 32, "held_out": 8}],
        ),
        class_counts=(8, 8, 8, 8, 8),
        class_weights=None,
        seed=SEED,
    )
    doc = build_color_model_document(config, result, run_id="RUN", git_commit="abc1234")
    return {"doc": doc, "pipeline": pipeline, "images": images, "rng": rng}


# --- downscale_nearest の一致 ---


@pytest.mark.parametrize(
    ("height", "width", "max_side"),
    [
        (100, 60, 48),
        (60, 100, 48),
        (1000, 999, 512),
        (513, 512, 512),
        (37, 1001, 50),
        (1001, 37, 50),
        (7, 3, 2),
        (3, 7, 1),
        (997, 991, 131),
        (1500, 2000, 512),
        (255, 255, 254),
    ],
)
def test_downscale_matches_pillow_nearest(height: int, width: int, max_side: int) -> None:
    rng = np.random.default_rng(height * 31 + width)
    img = rng.integers(0, 256, size=(height, width, 3), dtype=np.uint8)
    got = deploy.downscale_nearest(img, max_side)
    scale = max_side / max(height, width)
    size = (max(1, round(width * scale)), max(1, round(height * scale)))
    expected = np.asarray(Image.fromarray(img).resize(size, Image.Resampling.NEAREST))
    assert got.shape == expected.shape
    assert max(got.shape[:2]) <= max_side
    np.testing.assert_array_equal(got, expected)
    np.testing.assert_array_equal(got, train_downscale(img, max_side))


def test_downscale_returns_input_when_small_enough() -> None:
    img = np.arange(4 * 5 * 3, dtype=np.uint8).reshape(4, 5, 3)
    np.testing.assert_array_equal(deploy.downscale_nearest(img, 5), img)
    np.testing.assert_array_equal(deploy.downscale_nearest(img, 100), img)


def test_downscale_rejects_non_positive_max_side() -> None:
    img = np.zeros((4, 4, 3), dtype=np.uint8)
    with pytest.raises(ValueError):
        deploy.downscale_nearest(img, 0)


def test_downscale_matches_pillow_over_many_sizes() -> None:
    rng = np.random.default_rng(123)
    for _ in range(200):
        height, width = int(rng.integers(2, 700)), int(rng.integers(2, 700))
        max_side = int(rng.integers(1, max(height, width) + 1))
        img = rng.integers(0, 256, size=(height, width, 3), dtype=np.uint8)
        np.testing.assert_array_equal(
            deploy.downscale_nearest(img, max_side),
            train_downscale(img, max_side),
            err_msg=f"{height}x{width} -> {max_side}",
        )


# --- 特徴量の一致 ---


def test_extract_features_matches_training_side(trained: dict) -> None:
    model = deploy.load_color_model(trained["doc"])
    rng = np.random.default_rng(99)
    for size in [(20, 30), (48, 48), (120, 70), (70, 300), (1, 1)]:
        img = _random_tongue(rng, *size) if min(size) > 1 else np.full((1, 1, 3), 9, np.uint8)
        got = deploy.extract_features(img, model)
        expected = extract_color_features(train_downscale(img, MAX_SIDE))
        np.testing.assert_array_equal(got, expected)
        assert got.dtype == np.float64
        assert got.shape == (NUM_COLOR_FEATURES,)


def test_extract_features_without_downscale_uses_full_image(trained: dict) -> None:
    model = deploy.load_color_model(trained["doc"])
    img = _random_tongue(np.random.default_rng(5), 200, 150)
    np.testing.assert_array_equal(
        deploy.extract_features(img, model, downscale=False), extract_color_features(img)
    )


def test_extract_features_with_null_max_side_does_not_downscale(trained: dict) -> None:
    doc = json.loads(json.dumps(trained["doc"]))
    doc["features"]["max_side"] = None
    model = deploy.load_color_model(doc)
    img = _random_tongue(np.random.default_rng(6), 200, 150)
    np.testing.assert_array_equal(deploy.extract_features(img, model), extract_color_features(img))


def test_extract_features_covers_all_rgb_values(trained: dict) -> None:
    # 256 階調すべてを含む画像（線形化の LUT の全要素）でも学習側と一致する
    model = deploy.load_color_model(trained["doc"])
    rng = np.random.default_rng(11)
    img = np.stack([rng.permutation(256), rng.permutation(256), np.arange(256)], axis=-1).astype(
        np.uint8
    )[None]  # (1, 256, 3)．(0,0,0) の画素を含みうる
    img = np.repeat(img, 3, axis=0)
    np.testing.assert_array_equal(
        deploy.extract_features(img, model, downscale=False), extract_color_features(img)
    )


@pytest.mark.parametrize(
    "img",
    [
        np.zeros((4, 4), dtype=np.uint8),
        np.zeros((4, 4, 4), dtype=np.uint8),
        np.ones((4, 4, 3), dtype=np.float32),
        np.ones((4, 4, 3), dtype=np.int64),
    ],
    ids=["2d", "4ch", "float32", "int64"],
)
def test_extract_features_rejects_bad_input(trained: dict, img: np.ndarray) -> None:
    model = deploy.load_color_model(trained["doc"])
    with pytest.raises(ValueError):
        deploy.extract_features(img, model)


def test_extract_features_rejects_all_background(trained: dict) -> None:
    model = deploy.load_color_model(trained["doc"])
    with pytest.raises(ValueError):
        deploy.extract_features(np.zeros((10, 10, 3), dtype=np.uint8), model)


# --- 予測確率の一致 ---


def test_predict_proba_matches_sklearn(trained: dict) -> None:
    model = deploy.load_color_model(trained["doc"])
    rng = np.random.default_rng(1234)
    for i in range(30):
        img = _random_tongue(rng, int(rng.integers(10, 150)), int(rng.integers(10, 150)))
        got = deploy.predict_proba(img, model)
        feats = extract_color_features(train_downscale(img, MAX_SIDE))
        expected = trained["pipeline"].predict_proba(feats[None])[0]
        assert got.shape == (5,)
        assert np.max(np.abs(got - expected)) <= 1e-12, i
        assert got.sum() == pytest.approx(1.0, abs=1e-12)


def test_predict_proba_from_features_batch_matches_sklearn(trained: dict) -> None:
    model = deploy.load_color_model(trained["doc"])
    rng = np.random.default_rng(3)
    feats = np.stack([extract_color_features(_random_tongue(rng, 40, 40)) for _ in range(10)])
    got = deploy.predict_proba_from_features(feats, model)
    expected = trained["pipeline"].predict_proba(feats)
    assert got.shape == (10, 5)
    assert np.max(np.abs(got - expected)) <= 1e-12


def test_predict_proba_from_features_is_stable_for_extreme_values(trained: dict) -> None:
    model = deploy.load_color_model(trained["doc"])
    feats = np.full(NUM_COLOR_FEATURES, 1e6)
    probs = deploy.predict_proba_from_features(feats, model)
    assert np.all(np.isfinite(probs))
    assert probs.sum() == pytest.approx(1.0)


@pytest.mark.parametrize("shape", [(NUM_COLOR_FEATURES - 1,), (2, 3, NUM_COLOR_FEATURES), (3,)])
def test_predict_proba_from_features_rejects_bad_shape(trained: dict, shape: tuple) -> None:
    model = deploy.load_color_model(trained["doc"])
    with pytest.raises(ValueError):
        deploy.predict_proba_from_features(np.zeros(shape), model)


def test_json_round_trip_keeps_predictions(trained: dict, tmp_path: Path) -> None:
    path = tmp_path / "color_model.json"
    path.write_text(json.dumps(trained["doc"], ensure_ascii=False, indent=2), encoding="utf-8")
    from_file = deploy.load_color_model(path)
    from_str_path = deploy.load_color_model(str(path))
    from_dict = deploy.load_color_model(trained["doc"])
    rng = np.random.default_rng(77)
    for _ in range(10):
        img = _random_tongue(rng, 64, 50)
        p = deploy.predict_proba(img, from_dict)
        np.testing.assert_array_equal(deploy.predict_proba(img, from_file), p)
        np.testing.assert_array_equal(deploy.predict_proba(img, from_str_path), p)


# --- モデルの読み込み ---


def test_loaded_model_exposes_classes_and_features(trained: dict) -> None:
    model = deploy.load_color_model(trained["doc"])
    assert model.num_features == NUM_COLOR_FEATURES
    assert model.num_classes == 5
    assert model.feature_names == color_feature_names()
    assert [c["slug"] for c in model.classes] == list(CLASS_SLUGS)
    assert [c["display_name"] for c in model.classes] == [
        CLASS_DISPLAY_NAMES[s] for s in CLASS_SLUGS
    ]
    assert model.max_side == MAX_SIDE


def _broken(doc: dict, mutate) -> dict:
    d = json.loads(json.dumps(doc))
    mutate(d)
    return d


@pytest.mark.parametrize(
    "mutate",
    [
        lambda d: d.update(format_version=2),
        lambda d: d.update(format_version="1"),
        lambda d: d.update(model_type="color_hist_gb"),
        lambda d: d.pop("scaler"),
        lambda d: d.pop("classes"),
        lambda d: d["features"].pop("lab"),
        lambda d: d["classifier"].update(coef=d["classifier"]["coef"][:4]),
        lambda d: d["classifier"].update(intercept=d["classifier"]["intercept"][:4]),
        lambda d: d["scaler"].update(mean=d["scaler"]["mean"][:-1]),
        lambda d: d["scaler"].update(scale=[float("nan")] * NUM_COLOR_FEATURES),
        lambda d: d["features"].update(names=d["features"]["names"][:-1]),
        lambda d: d["features"].update(num_features=106),
        lambda d: d["features"].update(max_side=0),
        lambda d: d["features"]["ab_hist"].update(a_range=[50.0, 0.0]),
        lambda d: d["features"]["l_hist"].update(range=[20.0]),
        lambda d: d["features"]["lab"].update(srgb_to_xyz=[[1.0, 0.0], [0.0, 1.0]]),
        lambda d: d.update(classes=list(reversed(d["classes"]))),
    ],
    ids=[
        "future_version",
        "string_version",
        "model_type",
        "no_scaler",
        "no_classes",
        "no_lab",
        "coef_rows",
        "intercept_len",
        "mean_len",
        "nan_scale",
        "names_len",
        "num_features",
        "max_side_0",
        "a_range_reversed",
        "l_range_len",
        "matrix_shape",
        "class_order",
    ],
)
def test_load_rejects_invalid_documents(trained: dict, mutate) -> None:
    with pytest.raises(ValueError):
        deploy.load_color_model(_broken(trained["doc"], mutate))


def test_load_rejects_non_object_json(tmp_path: Path) -> None:
    path = tmp_path / "color_model.json"
    path.write_text("[1, 2, 3]", encoding="utf-8")
    with pytest.raises(ValueError):
        deploy.load_color_model(path)