tests/test_ml_training.py

"""Fast CPU tests for the ml_training workstream (synthetic source only, no network)."""

from __future__ import annotations

import csv
import json
from collections import Counter
from pathlib import Path

import numpy as np
import pytest
import torch

from ml_training.datasets.build_datasets import (
    MODALITIES,
    build_authenticity_set,
    build_imaging_set,
    split_for_id,
)
from ml_training.evaluate import main as evaluate_main
from ml_training.models.backbone import build_model, make_transforms
from ml_training.models.calibration import ece, fit_temperature
from ml_training.train_authenticity import main as train_authenticity_main
from ml_training.train_modality import main as train_modality_main

PER_CLASS = 12
IMG_SIZE = 96  # small synthetic images keep the smoke train well under the CPU budget


def _read_rows(manifest: Path) -> list[dict[str, str]]:
    with manifest.open(newline="") as f:
        return list(csv.DictReader(f))


@pytest.fixture(scope="module")
def data_root(tmp_path_factory: pytest.TempPathFactory) -> Path:
    root = tmp_path_factory.mktemp("mltrain")
    build_imaging_set("synthetic", PER_CLASS, IMG_SIZE, root / "imaging")
    build_authenticity_set(root / "imaging", root / "authenticity", rng_seed=42)
    return root


@pytest.fixture(scope="module")
def weights_dir(data_root: Path, tmp_path_factory: pytest.TempPathFactory) -> Path:
    out = tmp_path_factory.mktemp("weights")
    common = [
        "--epochs", "1", "--batch-size", "8", "--no-pretrained",
        "--input-size", str(IMG_SIZE), "--freeze-epochs", "0", "--seed", "42",
        "--device", "cpu", "--out", str(out),
    ]
    train_modality_main(["--data-dir", str(data_root / "imaging"), *common])
    train_authenticity_main(["--data-dir", str(data_root / "authenticity"), *common])
    return out


# ------------------------------------------------------------------ dataset build


def test_manifest_counts_and_deterministic_split(data_root: Path) -> None:
    rows = _read_rows(data_root / "imaging" / "manifest.csv")
    assert len(rows) == PER_CLASS * len(MODALITIES)
    counts = Counter(r["modality"] for r in rows)
    assert counts == {m: PER_CLASS for m in MODALITIES}
    for row in rows:
        stem = Path(row["path"]).stem
        assert row["split"] == split_for_id(stem)  # split is a pure function of the id
        assert (data_root / "imaging" / row["path"]).exists()
    by_split = Counter(r["split"] for r in rows)
    assert set(by_split) <= {"train", "val", "test"}
    assert by_split["train"] > by_split["val"] > 0 and by_split["test"] > 0
    # every class has train rows (WeightedRandomSampler precondition)
    train_classes = {r["modality"] for r in rows if r["split"] == "train"}
    assert train_classes == set(MODALITIES)


def test_build_is_stable_across_fresh_runs(data_root: Path, tmp_path: Path) -> None:
    build_imaging_set("synthetic", PER_CLASS, IMG_SIZE, tmp_path / "imaging")
    first = (data_root / "imaging" / "manifest.csv").read_text()
    second = (tmp_path / "imaging" / "manifest.csv").read_text()
    assert first == second  # ids, modalities, and splits identical across runs


def test_build_is_resumable_no_dupes(data_root: Path) -> None:
    imaging = data_root / "imaging"
    files_before = sorted(p.name for p in imaging.glob("*/*.jpg"))
    mtimes = {p: p.stat().st_mtime_ns for p in imaging.glob("*/*.jpg")}
    build_imaging_set("synthetic", PER_CLASS, IMG_SIZE, imaging)  # rerun on same out dir
    files_after = sorted(p.name for p in imaging.glob("*/*.jpg"))
    assert files_before == files_after
    assert all(p.stat().st_mtime_ns == t for p, t in mtimes.items())  # nothing re-written
    assert len(_read_rows(imaging / "manifest.csv")) == PER_CLASS * len(MODALITIES)


def test_authenticity_set_source_keyed_split(data_root: Path) -> None:
    rows = _read_rows(data_root / "authenticity" / "manifest_auth.csv")
    counts = Counter(r["label"] for r in rows)
    assert counts["real"] == counts["fake"] == PER_CLASS * len(MODALITIES)
    splits_per_source: dict[str, set[str]] = {}
    for row in rows:
        assert (data_root / "authenticity" / row["path"]).exists()
        assert row["split"] == split_for_id(row["source_id"])
        splits_per_source.setdefault(row["source_id"], set()).add(row["split"])
    # a source image never straddles splits (its real+fake pair share one split)
    assert all(len(s) == 1 for s in splits_per_source.values())
    assert len(splits_per_source) == PER_CLASS * len(MODALITIES)


# ------------------------------------------------------------------ model + transforms


def test_backbone_forward_pass_logits_shape() -> None:
    model = build_model(num_classes=3, pretrained=False).eval()
    with torch.no_grad():
        logits = model(torch.randn(2, 3, IMG_SIZE, IMG_SIZE))
    assert logits.shape == (2, 3)


def test_eval_transform_replicates_grayscale_to_3ch(data_root: Path) -> None:
    from PIL import Image

    from ml_training.models.backbone import IMAGENET_MEAN, IMAGENET_STD

    sample = next(iter((data_root / "imaging" / "ct").glob("*.jpg")))
    tensor = make_transforms(train=False, size=IMG_SIZE)(Image.open(sample).convert("L"))
    assert tensor.shape == (3, IMG_SIZE, IMG_SIZE)
    # channels are replicated before Normalize; undo per-channel norm to compare
    mean = torch.tensor(IMAGENET_MEAN).view(3, 1, 1)
    std = torch.tensor(IMAGENET_STD).view(3, 1, 1)
    raw = tensor * std + mean
    assert torch.allclose(raw[0], raw[1], atol=1e-6) and torch.allclose(raw[1], raw[2], atol=1e-6)


# ------------------------------------------------------------------ training + evaluate


def test_training_smoke_writes_weights_and_config(weights_dir: Path) -> None:
    for name, classes in (("modality", ["ct", "mri", "xray"]), ("authenticity", ["fake", "real"])):
        weights = weights_dir / f"{name}_efficientnet_b0.pt"
        assert weights.exists()
        state = torch.load(weights, map_location="cpu", weights_only=True)
        assert isinstance(state, dict) and len(state) > 0
        config = json.loads((weights_dir / f"{name}_config.json").read_text())
        assert config["classes"] == classes
        assert isinstance(config["temperature"], float)
        assert config["input_size"] == IMG_SIZE
        assert config["seed"] == 42
        assert len(config["dataset_manifest_sha256"]) == 64
        assert "macro_f1" in config["val_metrics"]
        assert config["normalization"]["mean"] and config["normalization"]["std"]


def test_evaluate_writes_report_with_confusion_matrices(
    weights_dir: Path, data_root: Path, tmp_path: Path
) -> None:
    report_path = tmp_path / "eval_report.json"
    evaluate_main(
        [
            "--weights-dir", str(weights_dir),
            "--data-dir", str(data_root),
            "--report", str(report_path),
            "--device", "cpu",
        ]
    )
    report = json.loads(report_path.read_text())
    modality_cm = report["modality"]["confusion_matrix"]
    assert len(modality_cm) == 3 and all(len(row) == 3 for row in modality_cm)
    auth_cm = report["authenticity"]["confusion_matrix"]
    assert len(auth_cm) == 2 and all(len(row) == 2 for row in auth_cm)
    for name in ("modality", "authenticity"):
        assert 0.0 <= report[name]["accuracy"] <= 1.0
        assert isinstance(report[name]["ece_before_temperature"], float)
        assert isinstance(report[name]["ece_after_temperature"], float)
        assert set(report[name]["per_class"]) == set(report[name]["classes"])


def test_evaluate_missing_weights_exits_nonzero(data_root: Path, tmp_path: Path) -> None:
    with pytest.raises(SystemExit) as excinfo:
        evaluate_main(
            [
                "--weights-dir", str(tmp_path / "empty"),
                "--data-dir", str(data_root),
                "--report", str(tmp_path / "report.json"),
            ]
        )
    assert excinfo.value.code not in (0, None)


# ------------------------------------------------------------------ calibration


def test_ece_perfect_predictions_near_zero() -> None:
    labels = np.array([0, 1, 2, 0, 1, 2])
    probs = np.eye(3)[labels]  # fully confident, fully correct
    assert ece(probs, labels, bins=15) < 1e-9


def test_fit_temperature_returns_positive_float() -> None:
    rng = np.random.default_rng(0)
    labels = rng.integers(0, 3, size=64)
    logits = np.eye(3)[labels] * 5.0 + rng.normal(0, 0.5, size=(64, 3))
    temperature = fit_temperature(logits.astype(np.float32), labels)
    assert isinstance(temperature, float) and temperature > 0.0