"""生成框标签的统一入口（gradcam / shap / cv 三种方法）。"""

import json
import sys
from pathlib import Path

import pytest

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

import numpy as np  # noqa: E402
from PIL import Image  # noqa: E402

sys.path.insert(0, "scripts")
import generate_boxes as gb  # noqa: E402


@pytest.fixture
def project(tmp_path):
    """一个最小的分类数据集 + 训练好的权重 + 配置。"""
    from engine.checkpoint import save_checkpoint
    from models import build_model

    root = tmp_path / "dataset"
    rng = np.random.default_rng(0)
    for name, count in (("grease", 3), ("normal", 2)):
        directory = root / name
        directory.mkdir(parents=True)
        for i in range(count):
            Image.fromarray(rng.integers(0, 255, (40, 56, 3), dtype=np.uint8)).save(
                directory / f"{name}_{i}.png"      # 全局唯一，与 --prefix-parent 的产物一致
            )

    model = build_model({
        "name": "classifier", "num_classes": 2,
        "backbone": {"name": "torchvision", "arch": "resnet18", "pretrained": False},
    })
    checkpoint = tmp_path / "best.pth"
    save_checkpoint(checkpoint, model)

    config = tmp_path / "cfg.yaml"
    config.write_text(
        f"""
device: cpu
data:
  root: {root}
  classes: null
  exclude: []
  val_transforms: [{{name: Resize, size: [64, 64]}}, {{name: ToTensor}}]
  loader: {{batch_size: 2, num_workers: 0}}
model:
  name: classifier
  backbone: {{name: torchvision, arch: resnet18, pretrained: false}}
  init_from: null
""",
        encoding="utf-8",
    )
    return root, checkpoint, config


def run(argv, monkeypatch):
    monkeypatch.setattr(sys, "argv", ["generate_boxes.py", *argv])
    return gb.main()


def base_argv(project, out, method):
    root, checkpoint, config = project
    return ["--out", str(out), "--method", method,
            "--checkpoint", str(checkpoint), "--config", str(config)]


@pytest.mark.parametrize("method", ["gradcam", "shap", "cv"])
def test_every_method_produces_a_yolo_dataset(project, tmp_path, monkeypatch, method):
    if method == "shap":
        pytest.importorskip("shap")
    root, _, _ = project
    out = tmp_path / f"out_{method}"

    extra = ["--max-evals", "60"] if method == "shap" else []
    assert run(base_argv(project, out, method) + extra, monkeypatch) == 0

    # images/ 与 labels/ 一一对应，符合仓库的 YOLO 约定
    images = sorted(p.relative_to(out / "images") for p in (out / "images").rglob("*.png"))
    labels = sorted(p.relative_to(out / "labels") for p in (out / "labels").rglob("*.txt"))
    assert len(images) == 5
    assert [p.with_suffix("") for p in images] == [p.with_suffix("") for p in labels]
    assert (out / "data.yaml").exists()

    for label in (out / "labels").rglob("*.txt"):
        for line in label.read_text(encoding="utf-8").strip().splitlines():
            parts = line.split()
            assert len(parts) == 5
            assert all(0.0 <= float(v) <= 1.0 for v in parts[1:])


def test_output_is_a_wellformed_yolo_dataset(project, tmp_path, monkeypatch):
    """产出是给外部检测训练用的交换格式，必须自洽：
    data.yaml 的类别与 labels/ 对得上，每行标签可解析且坐标在 [0,1] 内。"""
    import yaml

    out = tmp_path / "out"
    run(base_argv(project, out, "gradcam"), monkeypatch)

    meta = yaml.safe_load((out / "data.yaml").read_text(encoding="utf-8"))
    names = [meta["names"][i] for i in sorted(meta["names"])]
    assert meta["nc"] == len(names) == 2

    for label in (out / "labels").rglob("*.txt"):
        assert (out / "images" / label.parent.name / f"{label.stem}.png").is_file()
        for line in label.read_text(encoding="utf-8").strip().splitlines():
            class_id, cx, cy, w, h = line.split()
            assert 0 <= int(class_id) < meta["nc"]
            assert all(0.0 <= float(v) <= 1.0 for v in (cx, cy, w, h))
            assert float(w) > 0 and float(h) > 0


def test_images_are_hard_links_not_copies(project, tmp_path, monkeypatch):
    out = tmp_path / "out"
    run(base_argv(project, out, "gradcam"), monkeypatch)

    for path in (out / "images").rglob("*.png"):
        assert not path.is_symlink()
        assert path.stat().st_nlink > 1, "应与原图共享 inode，不额外占盘"

    summary = json.loads((out / "summary.json").read_text("utf-8"))
    assert summary["link_modes"].get("hard") == 5


def test_max_boxes_defaults_to_one(project, tmp_path, monkeypatch):
    out = tmp_path / "out"
    run(base_argv(project, out, "gradcam"), monkeypatch)

    for label in (out / "labels").rglob("*.txt"):
        assert len(label.read_text(encoding="utf-8").strip().splitlines()) <= 1

    summary = json.loads((out / "summary.json").read_text("utf-8"))
    assert summary["box_params"]["max_boxes"] == 1


def test_min_area_ratio_default_follows_the_method(project, tmp_path, monkeypatch):
    """两种方法的归因图分辨率差 16 倍，面积阈值不能共用默认值。"""
    assert gb.DEFAULT_MIN_AREA_RATIO["gradcam"] > gb.DEFAULT_MIN_AREA_RATIO["shap"]

    out = tmp_path / "out"
    run(base_argv(project, out, "gradcam"), monkeypatch)
    summary = json.loads((out / "summary.json").read_text("utf-8"))
    assert summary["box_params"]["min_area_ratio"] == gb.DEFAULT_MIN_AREA_RATIO["gradcam"]


def test_full_dataset_by_default_and_subsampling_is_opt_in(project, tmp_path, monkeypatch):
    full = tmp_path / "full"
    run(base_argv(project, full, "gradcam"), monkeypatch)
    assert json.loads((full / "summary.json").read_text("utf-8"))["num_images"] == 5

    limited = tmp_path / "limited"
    run(base_argv(project, limited, "gradcam") + ["--limit", "2"], monkeypatch)
    assert json.loads((limited / "summary.json").read_text("utf-8"))["num_images"] == 2

    per_class = tmp_path / "per_class"
    run(base_argv(project, per_class, "gradcam") + ["--per-class", "1"], monkeypatch)
    assert json.loads((per_class / "summary.json").read_text("utf-8"))["num_images"] == 2


def test_target_defaults_to_folder_label(project, tmp_path, monkeypatch):
    """框的类别默认用人工标注，不让分类模型的错误预测混进标签。"""
    out = tmp_path / "out_default"
    assert run(base_argv(project, out, "gradcam"), monkeypatch) == 0
    summary = json.loads((out / "summary.json").read_text(encoding="utf-8"))
    assert summary["target"] == "label"


def test_target_can_be_prediction_or_folder_label(project, tmp_path, monkeypatch):
    for target in ("predicted", "label"):
        out = tmp_path / f"out_{target}"
        assert run(base_argv(project, out, "gradcam") + ["--target", target], monkeypatch) == 0
        summary = json.loads((out / "summary.json").read_text("utf-8"))
        assert summary["target"] == target
        assert 0.0 <= summary["prediction_matches_folder_label"] <= 1.0


def test_optional_outputs_can_be_skipped(project, tmp_path, monkeypatch):
    out = tmp_path / "out"
    run(base_argv(project, out, "gradcam") + ["--no-heatmaps", "--no-images"], monkeypatch)

    assert not (out / "heatmaps").exists()
    assert not (out / "images").exists()
    assert list((out / "labels").rglob("*.txt"))          # 标签仍然产出


def test_heatmaps_are_written_per_image(project, tmp_path, monkeypatch):
    out = tmp_path / "out"
    run(base_argv(project, out, "gradcam"), monkeypatch)

    heatmaps = list((out / "heatmaps").rglob("*.png"))
    assert len(heatmaps) == 5
    with Image.open(heatmaps[0]) as image:
        assert image.size in {(56, 40)}                   # 与原图同尺寸


def test_link_or_copy_falls_back_across_filesystems(tmp_path, monkeypatch):
    import os as os_mod

    source = tmp_path / "a.png"
    source.write_bytes(b"x")

    def refuse(*_args, **_kwargs):
        raise OSError(errno_module.EXDEV, "Cross-device link")

    import errno as errno_module

    monkeypatch.setattr(os_mod, "link", refuse)
    assert gb.link_or_copy(source, tmp_path / "b" / "a.png") == "copy"
    assert (tmp_path / "b" / "a.png").read_bytes() == b"x"


def test_duplicate_source_names_are_rejected(tmp_path):
    """产出按**预测类别**分目录，而源图来自不同的类别目录：
    dirA/x.png 与 dirB/x.png 若被预测到同一类，就会写同一个标签路径。
    必须报错，而不是让后写的悄悄覆盖前一个。"""
    from argparse import Namespace

    heat = np.zeros((8, 8), dtype=np.float32)
    heat[3:5, 3:5] = 1.0
    box_kwargs = {"threshold": 0.4, "min_area_ratio": 0.0, "max_boxes": 1,
                  "padding": 0.0, "peak_ratio": 0.75}
    args = Namespace(no_images=True, no_heatmaps=True)
    seen: dict = {}
    common = dict(out_dir=tmp_path, classes=["grease"], target=0, heat=heat,
                  box_kwargs=box_kwargs, args=args, box_counts={}, link_modes={}, seen=seen)

    from collections import Counter
    common["box_counts"] = Counter()
    common["link_modes"] = Counter()

    gb._emit(source=Path("/data/dirA/x.png"), **common)
    with pytest.raises(gb.NameCollision, match="prefix-parent"):
        gb._emit(source=Path("/data/dirB/x.png"), **common)   # 同名，撞同一个标签路径


def test_gradcam_falls_back_when_backbone_has_no_sequential_features(tmp_path, monkeypatch):
    """tiny_conv 不是 nn.Sequential，取不了中间层——应当退回最终输出并警告，而不是崩掉。"""
    from engine.checkpoint import save_checkpoint
    from models import build_model

    root = tmp_path / "dataset"
    (root / "a").mkdir(parents=True)
    rng = np.random.default_rng(0)
    for i in range(2):
        Image.fromarray(rng.integers(0, 255, (40, 56, 3), dtype=np.uint8)).save(
            root / "a" / f"a_{i}.png"
        )

    backbone = {"name": "tiny_conv", "widths": [8, 16], "depths": [1, 1]}
    checkpoint = tmp_path / "best.pth"
    save_checkpoint(checkpoint, build_model(
        {"name": "classifier", "num_classes": 1, "backbone": dict(backbone)}
    ))

    config = tmp_path / "cfg.yaml"
    config.write_text(
        f"""
device: cpu
data:
  root: {root}
  classes: null
  exclude: []
  val_transforms: [{{name: Resize, size: [64, 64]}}, {{name: ToTensor}}]
  loader: {{batch_size: 2, num_workers: 0}}
model:
  name: classifier
  backbone: {{name: tiny_conv, widths: [8, 16], depths: [1, 1]}}
  init_from: null
""",
        encoding="utf-8",
    )

    out = tmp_path / "out"
    monkeypatch.setattr(
        sys, "argv",
        ["generate_boxes.py", "--out", str(out), "--method", "gradcam",
         "--checkpoint", str(checkpoint), "--config", str(config), "--no-heatmaps"],
    )
    assert gb.main() == 0                                    # 不崩

    summary = json.loads((out / "summary.json").read_text("utf-8"))
    assert summary["stage"] is None                          # 已退回最终输出
    assert len(list((out / "labels").rglob("*.txt"))) == 2


def test_stage_minus_one_selects_the_final_output(project, tmp_path, monkeypatch):
    out = tmp_path / "out"
    assert run(base_argv(project, out, "gradcam") + ["--stage", "-1"], monkeypatch) == 0
    assert json.loads((out / "summary.json").read_text("utf-8"))["stage"] is None


# ---------------- 传统 CV 方法 ----------------


def test_cv_runs_without_a_classifier(project, tmp_path, monkeypatch):
    """cv 不判类别，也就不需要权重——分类器还没训练出来时就能先跑一版框。"""
    root, _, config = project
    out = tmp_path / "out_cv"
    argv = ["--out", str(out), "--method", "cv", "--config", str(config),
            "--checkpoint", str(tmp_path / "does_not_exist.pth")]
    assert run(argv, monkeypatch) == 0

    assert len(list((out / "labels").rglob("*.txt"))) == 5
    summary = json.loads((out / "summary.json").read_text(encoding="utf-8"))
    assert summary["method"] == "cv"
    # 没加载模型，就没有预测可比
    assert summary["prediction_matches_folder_label"] is None
    assert summary["cv_params"]["polarity"] == "both"


def test_cv_labels_use_the_folder_class(project, tmp_path, monkeypatch):
    """类别只能来自目录：cv 自己不判类。"""
    root, _, config = project
    out = tmp_path / "out_cv"
    assert run(["--out", str(out), "--method", "cv", "--config", str(config)], monkeypatch) == 0

    classes = sorted(p.name for p in root.iterdir() if p.is_dir())
    for label in (out / "labels").rglob("*.txt"):
        folder = label.parent.name
        for line in label.read_text(encoding="utf-8").strip().splitlines():
            assert classes[int(line.split()[0])] == folder


def test_cv_finds_a_planted_defect(tmp_path, monkeypatch):
    """端到端：图上放一个黑点，框要套住它；平坦的图不出框。"""
    root = tmp_path / "dataset"
    (root / "blackdot").mkdir(parents=True)
    (root / "normal").mkdir(parents=True)

    rng = np.random.default_rng(0)
    flat = np.clip(rng.normal(140, 2, (96, 96)), 0, 255).astype(np.uint8)
    spotted = flat.copy()
    spotted[28:36, 60:68] = 40                    # 一个明显的黑点
    Image.fromarray(spotted).save(root / "blackdot" / "spot.png")
    Image.fromarray(flat).save(root / "normal" / "flat.png")

    config = tmp_path / "cfg.yaml"
    config.write_text(
        f"""
device: cpu
data:
  root: {root}
  classes: null
  exclude: []
  val_transforms: [{{name: Resize, size: [64, 64]}}, {{name: ToTensor}}]
  loader: {{batch_size: 2, num_workers: 0}}
model: {{name: classifier, backbone: {{name: torchvision, arch: resnet18, pretrained: false}}}}
""",
        encoding="utf-8",
    )
    out = tmp_path / "out"
    assert run(["--out", str(out), "--method", "cv", "--config", str(config)], monkeypatch) == 0

    spot = (out / "labels" / "blackdot" / "spot.txt").read_text(encoding="utf-8").split()
    assert len(spot) == 5
    _, cx, cy, width, height = (float(v) for v in spot)
    assert abs(cx - 64 / 96) < 0.1 and abs(cy - 32 / 96) < 0.1      # 命中黑点
    assert width * height < 0.15

    # 平坦的图没有缺陷，应当是空标签
    assert (out / "labels" / "normal" / "flat.txt").read_text(encoding="utf-8") == ""


def test_cv_polarity_is_passed_through(project, tmp_path, monkeypatch):
    root, _, config = project
    out = tmp_path / "out_cv"
    argv = ["--out", str(out), "--method", "cv", "--config", str(config),
            "--cv-polarity", "dark", "--cv-radius", "5", "--cv-min-contrast", "0.02"]
    assert run(argv, monkeypatch) == 0
    params = json.loads((out / "summary.json").read_text(encoding="utf-8"))["cv_params"]
    assert params == {
        "background_radius": 5, "polarity": "dark", "smooth_radius": 2, "min_contrast": 0.02,
    }