tests/test_jpeg_repair.py

"""JPEG repair Task 1-2の契約テスト。"""

from __future__ import annotations

from io import BytesIO

import numpy as np
import pytest
from PIL import Image

from pixelmodel_robustness.jpeg_repair import (
    CLIP_STRATEGIES,
    CONDITION_IDS,
    build_jpeg_registry,
    clip_candidate,
    decode_jpeg_rgb,
    decode_protected,
    diagnostics,
    encode_jpeg_q100,
    encode_protected,
    materialize_fp16_png,
    repair_nonfinite,
    semantic_hash,
    validate_layer_offsets,
)


def _image() -> Image.Image:
    return Image.fromarray(
        np.array([[[0, 1, 2], [20, 40, 60]], [[100, 120, 140], [240, 220, 200]]], dtype=np.uint8),
        "RGB",
    )


def test_registry_is_exact_ordered_and_semantically_hashed() -> None:
    registry = build_jpeg_registry()
    assert tuple(entry.condition_id for entry in registry) == CONDITION_IDS
    assert len(registry) == 18
    assert len({entry.parent_raw_sha256 for entry in registry[1:17]}) == 1
    assert semantic_hash(registry) == semantic_hash(build_jpeg_registry())
    assert tuple(entry.condition_id for entry in registry) == CONDITION_IDS


def test_jpeg_q100_source_is_one_reusable_decode() -> None:
    raw = encode_jpeg_q100(_image())
    first = decode_jpeg_rgb(raw)
    second = decode_jpeg_rgb(raw)
    assert raw.startswith(b"\xff\xd8")
    np.testing.assert_array_equal(first, second)
    assert first.shape == (2, 2, 3)


def test_fp16_flat_materialization_preserves_bits_and_png_contract() -> None:
    values = np.array([0.0, -0.0, 1.5, -2.25, np.float16("inf")], dtype=np.float16)
    png = materialize_fp16_png(values)
    pixels = np.asarray(Image.open(BytesIO(png)).convert("RGB"), dtype=np.uint8).reshape(-1, 3)
    np.testing.assert_array_equal((pixels[: len(values), 0].astype(np.uint16) << 8) | pixels[: len(values), 1], values.view(np.uint16))
    assert pixels[:, 2].tolist() == [0] * len(values)


def test_finite_repair_materialize_decode_preserves_original_uint16_bits() -> None:
    original = np.array([0.0, -0.0, 1.5, -2.25, 3.125, -4.5], dtype=np.float16)
    repaired = repair_nonfinite(original, "layer_mean", {"layer": (0, len(original))}).values
    png = materialize_fp16_png(repaired)
    pixels = np.asarray(Image.open(BytesIO(png)).convert("RGB"), dtype=np.uint8).reshape(-1, 3)
    decoded_bits = (pixels[: len(original), 0].astype(np.uint16) << 8) | pixels[: len(original), 1].astype(np.uint16)
    np.testing.assert_array_equal(decoded_bits, original.view(np.uint16))


def test_protected_high_png_is_exact_low_jpeg_only_changes() -> None:
    values = np.array([0.0, -0.0, 1.5, -2.25, 3.125, -4.5], dtype=np.float16)
    artifact = encode_protected(values, (2, 3))
    decoded, metadata = decode_protected(artifact.payload, (2, 3))
    np.testing.assert_array_equal(decoded.view(np.uint16), values.view(np.uint16))
    high = np.asarray(Image.open(BytesIO(artifact.high_png_bytes)).convert("L"), dtype=np.uint8)
    np.testing.assert_array_equal(high.reshape(-1), (values.view(np.uint16) >> 8).astype(np.uint8))
    assert metadata["payload_size_bytes"] == len(artifact.high_png_bytes) + len(artifact.low_jpeg_bytes)
    assert metadata["high_mode"] == metadata["low_mode"] == "L"


def test_repair_replaces_only_nonfinite_and_reports_signed_zero_bits() -> None:
    raw = np.array([0.0, -0.0, np.nan, np.inf, -np.inf, 2.0], dtype=np.float16)
    result = repair_nonfinite(raw, "layer_mean", {"layer": (0, len(raw))})
    assert result.repaired_value_count == 3
    assert result.nan_count == result.posinf_count == result.neginf_count == 1
    assert np.isfinite(result.values).all()
    np.testing.assert_array_equal(result.values.view(np.uint16)[[0, 1, 5]], raw.view(np.uint16)[[0, 1, 5]])


def test_zero_repair_accepts_all_nonfinite_layer_and_preserves_finite_bits() -> None:
    raw = np.array([np.nan, np.inf, -np.inf, 1.5, -0.0, np.nan], dtype=np.float16)
    original_bits = raw.view(np.uint16).copy()

    result = repair_nonfinite(raw, "zero", {"all_nonfinite": (0, 3), "mixed": (3, 6)})

    assert np.isfinite(result.values).all()
    assert result.repaired_value_count == 4
    assert result.nan_count == 2
    assert result.posinf_count == 1
    assert result.neginf_count == 1
    assert [record.repaired_count for record in result.layer_records] == [3, 1]
    assert [record.nan_count for record in result.layer_records] == [1, 1]
    assert [record.posinf_count for record in result.layer_records] == [1, 0]
    assert [record.neginf_count for record in result.layer_records] == [1, 0]
    finite = np.isfinite(raw)
    np.testing.assert_array_equal(result.values.view(np.uint16)[finite], original_bits[finite])


def test_repair_rejects_empty_finite_layer() -> None:
    raw = np.array([np.nan, np.inf], dtype=np.float16)
    with pytest.raises(ValueError, match="finite values"):
        repair_nonfinite(raw, "layer_median", {"empty": (0, 2)})


@pytest.mark.parametrize("method", ["layer_median", "layer_mean"])
def test_repair_rejects_empty_finite_layer_for_statistical_methods(method: str) -> None:
    raw = np.array([np.nan, np.inf, -np.inf], dtype=np.float16)
    with pytest.raises(ValueError, match="finite values"):
        repair_nonfinite(raw, method, {"empty": (0, 3)})


@pytest.mark.parametrize("strategy", CLIP_STRATEGIES)
def test_clip_strategies_use_candidate_only_and_record_changes(strategy: str) -> None:
    candidate = np.array([-10.0, -1.0, 0.0, 1.0, 10.0], dtype=np.float16)
    result = clip_candidate(candidate, {"layer": (0, 5)}, strategy)
    assert result.values.dtype == np.float16
    assert result.layers[0].strategy == strategy
    assert result.layers[0].changed_count == int(np.count_nonzero(result.values.view(np.uint16) != candidate.view(np.uint16)))
    assert np.isfinite(result.values).all()
    assert result.layers[0].lower <= result.layers[0].upper


def test_clip_uses_population_std_and_linear_percentile() -> None:
    candidate = np.array([0.0, 1.0, 2.0, 100.0], dtype=np.float16)
    result = clip_candidate(candidate, {"layer": (0, 4)}, "layer_mean_std_6")
    expected = float(np.mean(candidate.astype(np.float64)) + 6 * np.std(candidate.astype(np.float64), ddof=0))
    assert result.layers[0].upper == pytest.approx(expected)
    percentile = clip_candidate(candidate, {"layer": (0, 4)}, "layer_pct_0.1_99.9")
    expected_percentile = float(np.percentile(candidate.astype(np.float64), 99.9, method="linear"))
    assert percentile.layers[0].upper == pytest.approx(expected_percentile)


def test_clip_degenerate_and_threshold_equal_are_unchanged() -> None:
    candidate = np.array([2.0, 2.0], dtype=np.float16)
    result = clip_candidate(candidate, {"one": (0, 2)}, "layer_mean_std_5")
    np.testing.assert_array_equal(result.values.view(np.uint16), candidate.view(np.uint16))
    assert result.layers[0].changed_count == 0


def test_fp16_bit_diagnostics_and_bounded_histogram() -> None:
    baseline = np.array([1.0, -1.0, 0.0, 1.0], dtype=np.float16)
    candidate = np.array([20.0, -1.0, np.float16("inf"), np.float16("nan")], dtype=np.float16)
    result = diagnostics(baseline, candidate, {"layer": (0, 4)})
    assert result["sign_mismatch_count"] == 0
    assert result["exponent_mismatch_count"] == 3
    assert result["mantissa_mismatch_count"] == 2
    assert result["exponent_to31_count"] == 2
    assert result["finite_extreme_outlier_count"] == 1
    assert set(result["histograms"]) == {"baseline", "candidate", "delta"}
    assert all(len(item["counts"]) <= 64 for item in result["histograms"].values())


def test_registry_parent_identity_is_only_shared_by_raw_jpeg_conditions() -> None:
    raw = encode_jpeg_q100(_image())
    registry = build_jpeg_registry(raw)
    assert len({entry.parent_raw_sha256 for entry in registry[1:17]}) == 1
    assert registry[0].parent_raw_sha256 != registry[1].parent_raw_sha256
    assert registry[17].parent_raw_sha256 != registry[1].parent_raw_sha256
    assert registry[0].parent_raw_sha256 != registry[17].parent_raw_sha256


def test_layer_offsets_are_contiguous_and_fail_closed() -> None:
    assert validate_layer_offsets({"a": (0, 2), "b": (2, 4)}, 4) == {"a": (0, 2), "b": (2, 4)}
    with pytest.raises(ValueError, match="contiguous"):
        validate_layer_offsets({"a": (0, 2), "b": (3, 4)}, 4)
    with pytest.raises(ValueError, match="overlap"):
        validate_layer_offsets({"a": (0, 3), "b": (2, 4)}, 4)


def test_diagnostics_extreme_uses_layer_baseline_max_and_reports_finite_stats() -> None:
    baseline = np.array([1.0, 2.0, 0.0, 0.0], dtype=np.float16)
    candidate = np.array([25.0, 2.0, 1.0, 0.0], dtype=np.float16)
    result = diagnostics(baseline, candidate, {"first": (0, 2), "zero": (2, 4)})
    assert result["finite_extreme_outlier_count"] == 2
    assert result["finite_count"] == 4
    assert result["max_abs_weight"] == pytest.approx(25.0)
    assert result["nan_count"] == result["posinf_count"] == result["neginf_count"] == 0
    assert result["layers"][0]["finite_extreme_outlier_count"] == 1
    assert result["layers"][1]["finite_extreme_outlier_count"] == 1
    assert result["layers"][1]["max_abs_weight"] == pytest.approx(1.0)


def test_repair_has_layer_records_and_bounded_repaired_histogram() -> None:
    raw = np.array([np.nan, np.inf, -np.inf, 1.0], dtype=np.float16)
    result = repair_nonfinite(raw, "zero", {"layer": (0, 4)})
    record = result.layer_records[0]
    assert record.layer == "layer"
    assert record.repaired_count == 3
    assert record.nan_count == record.posinf_count == record.neginf_count == 1
    assert sum(result.repaired_histogram["counts"]) == 3
    assert len(result.repaired_histogram["counts"]) <= 64


def test_clip_has_total_count_and_changed_value_histograms() -> None:
    candidate = np.array([-100.0, -1.0, 0.0, 1.0, 100.0], dtype=np.float16)
    result = clip_candidate(candidate, {"layer": (0, 5)}, "layer_mean_std_5")
    assert result.total_clipped_count == result.layers[0].changed_count
    assert sum(result.changed_pre_histogram["counts"]) == result.total_clipped_count
    assert sum(result.changed_post_histogram["counts"]) == result.total_clipped_count
    assert sum(result.changed_delta_histogram["counts"]) == result.total_clipped_count


def test_protected_metadata_reports_decoded_planes_and_low_changes() -> None:
    rng = np.random.default_rng(20260727)
    values = rng.integers(0, 65536, size=32 * 32, dtype=np.uint16).view(np.float16)
    artifact = encode_protected(values, (32, 32))
    decoded, metadata = decode_protected(artifact.payload, (32, 32), source_values=values)
    assert metadata["decoded_high_exact"] is True
    assert metadata["decoded_high_sha256"]
    assert metadata["decoded_low_sha256"]
    assert metadata["decoded_low_changed_count"] > 0
    assert 0.0 < metadata["decoded_low_changed_rate"] <= 1.0
    assert metadata["pillow_version"]
    assert metadata["combined_payload_exact"] is True
    assert metadata["high_mode"] == metadata["low_mode"] == "L"
    assert metadata["high_dimensions"] == metadata["low_dimensions"] == [32, 32]
    assert metadata["combined_payload_sha256"]
    assert metadata["combined_payload_bytes"] == metadata["high_bytes"] + metadata["low_bytes"]
    assert "jpeg_codec_version" in metadata
    assert decoded.shape == values.shape


def test_combined_payload_parser_rejects_invalid_png_chunk_bounds() -> None:
    values = np.arange(16, dtype=np.float16)
    artifact = encode_protected(values, (4, 4))
    broken = artifact.payload[:16] + b"\xff"
    with pytest.raises(ValueError, match="PNG"):
        decode_protected(broken, (4, 4))