"""SHAP 归因定位。"""

import pytest

torch = pytest.importorskip("torch")
pytest.importorskip("shap")

import numpy as np  # noqa: E402
from torch import nn  # noqa: E402

from engine.gradcam import cam_to_boxes  # noqa: E402
from engine.shap_explain import shap_attribution, smooth  # noqa: E402


class _Detector(nn.Module):
    """把左上角亮度当成第 1 类的证据——归因应当集中在那里。"""

    def __init__(self) -> None:
        super().__init__()
        self.bias = nn.Parameter(torch.zeros(2), requires_grad=False)

    def forward(self, x):
        corner = x[:, :, :8, :8].mean(dim=(1, 2, 3))
        return torch.stack([-corner, corner], dim=1)


def test_smooth_spreads_and_renormalizes():
    peak = np.zeros((16, 16), dtype=np.float32)
    peak[8, 8] = 1.0

    smoothed = smooth(peak, radius=2)
    assert smoothed.shape == peak.shape
    assert int((smoothed > 0).sum()) == 25          # 5x5 方框
    assert float(smoothed.max()) == pytest.approx(1.0)   # 重新归一化到 1
    assert np.allclose(smooth(peak, radius=0), peak)


def test_smooth_keeps_flat_map_flat():
    flat = np.full((8, 8), 0.5, dtype=np.float32)
    assert np.allclose(smooth(flat, radius=2), 1.0)  # 归一化后整体为 1，不产生伪峰


def test_attribution_lands_on_the_evidence():
    """模型只看左上角，SHAP 的正贡献就该集中在左上角。"""
    model = _Detector().eval()
    image = np.zeros((16, 16, 3), dtype=np.float32)
    image[:8, :8] = 1.0

    attribution = shap_attribution(
        model, image, target=1, device=torch.device("cpu"), max_evals=200, batch_size=32
    )
    assert attribution.shape == (16, 16)             # 与输入同分辨率，不是特征图分辨率
    assert 0.0 <= attribution.min() and attribution.max() == pytest.approx(1.0)
    assert attribution[:8, :8].mean() > attribution[8:, 8:].mean()


def test_attribution_is_non_negative():
    """只保留正贡献：负贡献表示"这块在把分数往下拉"，对定位没有意义。"""
    model = _Detector().eval()
    image = np.random.default_rng(0).random((16, 16, 3)).astype(np.float32)
    attribution = shap_attribution(
        model, image, target=0, device=torch.device("cpu"), max_evals=200, batch_size=32
    )
    assert (attribution >= 0).all()


def test_full_resolution_beats_feature_map_for_small_targets():
    """SHAP 与输入同分辨率，小目标不会被格子吞掉——这是它相对 Grad-CAM 的核心优势。

    一个 10x10 的小目标在 128x128 归因图上占 100 像素：
    按 gradcam 那套 min_area_ratio=0.02（=328 像素）会被整个滤掉，
    换成 0.002（=33 像素）才留得住——所以 SHAP 必须用小得多的阈值。
    """
    attribution = np.zeros((128, 128), dtype=np.float32)
    attribution[60:70, 60:70] = 1.0

    assert cam_to_boxes(attribution, threshold=0.4, min_area_ratio=0.02) == []
    (box,) = cam_to_boxes(attribution, threshold=0.4, min_area_ratio=0.002)
    x1, y1, x2, y2 = box
    assert (x2 - x1) < 0.15 and (y2 - y1) < 0.15     # 框紧贴小目标