scripts/pixelmodel_robustness/jpeg_repair.py

"""JPEG weight repairのTask 1-2用codec、修復、clip、fp16診断。"""

from __future__ import annotations

import hashlib
import json
import logging
import zlib
from collections.abc import Mapping, Sequence
from dataclasses import asdict, dataclass
from io import BytesIO
from typing import Any, Final

import numpy as np
from PIL import Image, features
from PIL import __version__ as PILLOW_VERSION

LOGGER = logging.getLogger(__name__)

CONDITION_IDS: Final[tuple[str, ...]] = (
    "png_baseline",
    "jpeg_q100_raw",
    "jpeg_q100_repair_zero",
    "jpeg_q100_repair_layer_median",
    "jpeg_q100_repair_layer_mean",
    "jpeg_q100_repair_zero_clip_layer_mean_std_6",
    "jpeg_q100_repair_zero_clip_layer_mean_std_5",
    "jpeg_q100_repair_zero_clip_layer_pct_0.01_99.99",
    "jpeg_q100_repair_zero_clip_layer_pct_0.1_99.9",
    "jpeg_q100_repair_layer_median_clip_layer_mean_std_6",
    "jpeg_q100_repair_layer_median_clip_layer_mean_std_5",
    "jpeg_q100_repair_layer_median_clip_layer_pct_0.01_99.99",
    "jpeg_q100_repair_layer_median_clip_layer_pct_0.1_99.9",
    "jpeg_q100_repair_layer_mean_clip_layer_mean_std_6",
    "jpeg_q100_repair_layer_mean_clip_layer_mean_std_5",
    "jpeg_q100_repair_layer_mean_clip_layer_pct_0.01_99.99",
    "jpeg_q100_repair_layer_mean_clip_layer_pct_0.1_99.9",
    "protected_high_png_low_jpeg_q100",
)
CLIP_STRATEGIES: Final[tuple[str, ...]] = (
    "layer_mean_std_6", "layer_mean_std_5", "layer_pct_0.01_99.99", "layer_pct_0.1_99.9"
)


@dataclass(frozen=True)
class JpegCondition:
    """JPEG studyの不変condition entry。"""

    condition_id: str
    source: str
    repair: str | None
    clip: str | None
    payload_kind: str
    parent_raw_sha256: str


@dataclass(frozen=True)
class LayerClipRecord:
    """一つのlayerに対するclip結果。"""

    layer: str
    strategy: str
    lower: float
    upper: float
    pre_max_abs: float
    post_max_abs: float
    changed_count: int


@dataclass(frozen=True)
class RepairResult:
    """非有限修復の値と監査用count。"""

    values: np.ndarray
    method: str
    nan_count: int
    posinf_count: int
    neginf_count: int
    repaired_value_count: int
    layer_records: tuple[LayerRepairRecord, ...]
    repaired_histogram: dict[str, Any]


@dataclass(frozen=True)
class LayerRepairRecord:
    """一つのlayerにおける非有限値修復の監査record。"""

    layer: str
    replacement: float
    nan_count: int
    posinf_count: int
    neginf_count: int
    repaired_count: int


@dataclass(frozen=True)
class ClipResult:
    """layerwise clipの値とlayer record。"""

    values: np.ndarray
    strategy: str
    layers: tuple[LayerClipRecord, ...]
    total_clipped_count: int
    changed_pre_histogram: dict[str, Any]
    changed_post_histogram: dict[str, Any]
    changed_delta_histogram: dict[str, Any]


@dataclass(frozen=True)
class ProtectedArtifact:
    """high PNG + low JPEGのpayload。"""

    high_png_bytes: bytes
    low_jpeg_bytes: bytes

    @property
    def payload(self) -> bytes:
        """推論用protected payloadを返す。"""
        return self.high_png_bytes + self.low_jpeg_bytes


def _sha256(data: bytes) -> str:
    return hashlib.sha256(data).hexdigest()


def encode_jpeg_q100(image: Image.Image | np.ndarray) -> bytes:
    """RGB sourceをJPEG quality=100/subsampling=0で一度だけencodeする。"""
    rgb = image if isinstance(image, Image.Image) else Image.fromarray(np.asarray(image, dtype=np.uint8), "RGB")
    stream = BytesIO()
    rgb.convert("RGB").save(stream, format="JPEG", quality=100, subsampling=0, optimize=False)
    return stream.getvalue()


def decode_jpeg_rgb(payload: bytes) -> np.ndarray:
    """JPEG bytesをRGB uint8へdecodeする。"""
    with Image.open(BytesIO(payload)) as image:
        return np.asarray(image.convert("RGB"), dtype=np.uint8).copy()


def flatten_fp16_layers(arrays: Mapping[str, np.ndarray]) -> tuple[np.ndarray, dict[str, tuple[int, int]]]:
    """mappingの順序を保ったfp16 flat値とlayer offsetsを作る。"""
    parts: list[np.ndarray] = []
    offsets: dict[str, tuple[int, int]] = {}
    start = 0
    for name, array in arrays.items():
        part = np.asarray(array, dtype=np.float16).reshape(-1)
        end = start + part.size
        parts.append(part)
        offsets[str(name)] = (start, end)
        start = end
    return (np.concatenate(parts) if parts else np.empty(0, dtype=np.float16), offsets)


def layer_offsets(names: Sequence[str], sizes: Sequence[int]) -> dict[str, tuple[int, int]]:
    """layer名とnumelからordered offsetsを構築する。"""
    if len(names) != len(sizes):
        raise ValueError("names and sizes must have equal length")
    offsets: dict[str, tuple[int, int]] = {}
    start = 0
    for name, size in zip(names, sizes):
        end = start + int(size)
        if int(size) < 0:
            raise ValueError("layer size must be non-negative")
        offsets[str(name)] = (start, end)
        start = end
    return offsets


def validate_layer_offsets(layers: Mapping[str, tuple[int, int]], total_size: int) -> dict[str, tuple[int, int]]:
    """layer offsetsが全体を隙間なく一度だけ覆うことを検証する。"""
    if total_size < 0:
        raise ValueError("total_size must be non-negative")
    normalized = {str(name): (int(bounds[0]), int(bounds[1])) for name, bounds in layers.items()}
    expected_start = 0
    for name, (start, end) in normalized.items():
        if start < 0 or end < start or end > total_size:
            raise ValueError(f"layer {name} is out of bounds")
        if start < expected_start:
            raise ValueError(f"layer {name} has overlap")
        if start > expected_start:
            raise ValueError(f"layer {name} is not contiguous")
        expected_start = end
    if expected_start != total_size:
        raise ValueError("layers do not end at total size")
    return normalized


def _condition_fields(condition_id: str) -> tuple[str, str | None, str | None, str]:
    if condition_id == "png_baseline":
        return "canonical_png", None, None, "materialized_rgb_png"
    if condition_id == "jpeg_q100_raw":
        return "jpeg_q100", None, None, "raw_jpeg"
    if condition_id == "protected_high_png_low_jpeg_q100":
        return "protected", None, None, "high_png_low_jpeg"
    body = condition_id.removeprefix("jpeg_q100_repair_")
    repair, _, clip = body.partition("_clip_")
    return "jpeg_q100", repair, clip, "materialized_rgb_png"


def build_jpeg_registry(raw_jpeg_bytes: bytes | None = None) -> tuple[JpegCondition, ...]:
    """計画で固定された18条件をordered immutable tupleとして構築する。"""
    raw_parent = _sha256(raw_jpeg_bytes) if raw_jpeg_bytes is not None else _sha256(b"jpeg_q100:quality=100:subsampling=0:optimize=False")
    baseline_parent = _sha256(b"canonical_png:png_baseline")
    protected_parent = _sha256(b"protected:high_png_low_jpeg_q100")
    entries: list[JpegCondition] = []
    for condition_id in CONDITION_IDS:
        parent = raw_parent if 1 <= CONDITION_IDS.index(condition_id) <= 16 else baseline_parent if condition_id == "png_baseline" else protected_parent
        entries.append(JpegCondition(condition_id, *_condition_fields(condition_id), parent))
    return tuple(entries)


def semantic_hash(registry: Sequence[JpegCondition]) -> str:
    """registryの順序・全semantic fieldsをcanonical JSONでhashする。"""
    encoded = json.dumps([asdict(entry) for entry in registry], sort_keys=True, separators=(",", ":")).encode()
    return _sha256(encoded)


def materialize_fp16_png(values: np.ndarray, *, width: int | None = None) -> bytes:
    """flat fp16 bit patternをRGB PNGへmaterializeする（R=high,G=low,B=0）。"""
    flat = np.asarray(values, dtype=np.float16).reshape(-1)
    if not flat.size:
        raise ValueError("values must not be empty")
    width_value = int(width or max(1, min(flat.size, 1024)))
    height = (flat.size + width_value - 1) // width_value
    pixels = np.zeros((height * width_value, 3), dtype=np.uint8)
    bits = flat.view(np.uint16)
    pixels[: flat.size, 0] = (bits >> 8).astype(np.uint8)
    pixels[: flat.size, 1] = bits.astype(np.uint8)
    stream = BytesIO()
    Image.fromarray(pixels.reshape(height, width_value, 3), "RGB").save(stream, format="PNG")
    return stream.getvalue()


def _planes(values: np.ndarray, shape: tuple[int, int]) -> tuple[np.ndarray, np.ndarray]:
    flat = np.asarray(values, dtype=np.float16).reshape(-1)
    capacity = int(np.prod(shape))
    if flat.size > capacity:
        raise ValueError("shape is smaller than values")
    bits = np.zeros(capacity, dtype=np.uint16)
    bits[: flat.size] = flat.view(np.uint16)
    return (bits >> 8).astype(np.uint8).reshape(shape), bits.astype(np.uint8).reshape(shape)


def encode_protected(values: np.ndarray, shape: tuple[int, int]) -> ProtectedArtifact:
    """high byteをlossless PNG、low byteをJPEG Q100 grayscaleでencodeする。"""
    high, low = _planes(values, shape)
    high_stream, low_stream = BytesIO(), BytesIO()
    Image.fromarray(high, "L").save(high_stream, format="PNG")
    Image.fromarray(low, "L").save(low_stream, format="JPEG", quality=100, subsampling=0, optimize=False)
    return ProtectedArtifact(high_stream.getvalue(), low_stream.getvalue())


def decode_protected(
    payload: bytes | ProtectedArtifact,
    shape: tuple[int, int],
    *,
    source_values: np.ndarray | None = None,
) -> tuple[np.ndarray, dict[str, Any]]:
    """protected payloadをfp16へ再構築し、plane metadataを返す。"""
    artifact = payload if isinstance(payload, ProtectedArtifact) else _split_payload(payload)
    with Image.open(BytesIO(artifact.high_png_bytes)) as high_image, Image.open(BytesIO(artifact.low_jpeg_bytes)) as low_image:
        high_mode = high_image.mode
        low_mode = low_image.mode
        high_dimensions = [int(high_image.height), int(high_image.width)]
        low_dimensions = [int(low_image.height), int(low_image.width)]
        high = np.asarray(high_image.convert("L"), dtype=np.uint8)
        low = np.asarray(low_image.convert("L"), dtype=np.uint8)
    if high.shape != shape or low.shape != shape:
        raise ValueError("protected plane shape mismatch")
    bits = (high.reshape(-1).astype(np.uint16) << 8) | low.reshape(-1).astype(np.uint16)
    source_high: np.ndarray | None = None
    source_low: np.ndarray | None = None
    if source_values is not None:
        source_high, source_low = _planes(source_values, shape)
        source_high = source_high.reshape(-1)
        source_low = source_low.reshape(-1)
    decoded_high_bytes = high.tobytes()
    decoded_low_bytes = low.tobytes()
    low_changed_count = int(np.count_nonzero(low.reshape(-1) != source_low)) if source_low is not None else None
    total = int(np.prod(shape))
    metadata = {
        "dimensions": list(shape),
        "high_mode": high_mode,
        "low_mode": low_mode,
        "high_dimensions": high_dimensions,
        "low_dimensions": low_dimensions,
        "high_bytes": len(artifact.high_png_bytes),
        "low_bytes": len(artifact.low_jpeg_bytes),
        "high_sha256": _sha256(artifact.high_png_bytes),
        "low_sha256": _sha256(artifact.low_jpeg_bytes),
        "payload_sha256": _sha256(artifact.payload),
        "payload_size_bytes": len(artifact.payload),
        "jpeg_options": {"quality": 100, "subsampling": 0, "optimize": False},
        "decoded_high_sha256": _sha256(decoded_high_bytes),
        "decoded_low_sha256": _sha256(decoded_low_bytes),
        "decoded_high_exact": bool(source_high is not None and np.array_equal(high.reshape(-1), source_high)),
        "decoded_low_changed_count": low_changed_count,
        "decoded_low_changed_rate": (float(low_changed_count / total) if low_changed_count is not None and total else None),
        "pillow_version": PILLOW_VERSION,
        "jpeg_codec_version": _jpeg_codec_version(),
        "combined_payload_sha256": _sha256(artifact.high_png_bytes + artifact.low_jpeg_bytes),
        "combined_payload_bytes": len(artifact.high_png_bytes) + len(artifact.low_jpeg_bytes),
        "combined_payload_exact": artifact.payload == (payload if isinstance(payload, bytes) else artifact.payload),
    }
    return bits.view(np.float16), metadata


def _jpeg_codec_version() -> str | None:
    """Pillowが報告できるJPEG codec versionを取得する。"""
    try:
        return features.version_codec("jpg")
    except (AttributeError, RuntimeError, ValueError):
        return None


def _split_payload(payload: bytes) -> ProtectedArtifact:
    """PNG chunkをsignature/length/type/data/CRC順に検証してcombined payloadを分割する。"""
    signature = b"\x89PNG\r\n\x1a\n"
    if len(payload) < len(signature) or payload[: len(signature)] != signature:
        raise ValueError("protected payload has invalid PNG signature")
    position = len(signature)
    while True:
        if position + 12 > len(payload):
            raise ValueError("PNG chunk header is out of bounds")
        length = int.from_bytes(payload[position:position + 4], "big")
        chunk_start = position
        position += 8
        end_data = position + length
        end_chunk = end_data + 4
        if end_chunk > len(payload):
            raise ValueError("PNG chunk data is out of bounds")
        chunk_type = payload[position - 4:position]
        data = payload[position:end_data]
        expected_crc = int.from_bytes(payload[end_data:end_chunk], "big")
        actual_crc = zlib.crc32(chunk_type + data) & 0xFFFFFFFF
        if expected_crc != actual_crc:
            raise ValueError("PNG chunk CRC mismatch")
        position = end_chunk
        if chunk_type == b"IEND":
            if length != 0:
                raise ValueError("PNG IEND chunk must be empty")
            return ProtectedArtifact(payload[:position], payload[position:])
        if chunk_start == len(payload):
            raise ValueError("PNG has no IEND chunk")


def repair_nonfinite(values: np.ndarray, method: str, layers: Mapping[str, tuple[int, int]]) -> RepairResult:
    """NaN/Infだけをfloat64 layer statisticで置換し、最後に一度だけfp16化する。"""
    if method not in {"zero", "layer_median", "layer_mean"}:
        raise ValueError(f"unknown repair method: {method}")
    source = np.asarray(values, dtype=np.float16).reshape(-1)
    normalized_layers = validate_layer_offsets(layers, source.size)
    with np.errstate(invalid="ignore"):
        output = source.astype(np.float64)
    finite = np.isfinite(source)
    layer_records: list[LayerRepairRecord] = []
    for name, (start, end) in normalized_layers.items():
        layer_finite = output[start:end][finite[start:end]]
        if method != "zero" and layer_finite.size == 0:
            raise ValueError(f"layer {name} has no finite values")
        mask = ~finite[start:end]
        if method == "zero":
            replacement = 0.0
        elif method == "layer_median":
            replacement = float(np.median(layer_finite))
        else:
            replacement = float(np.mean(layer_finite, dtype=np.float64))
        output[start:end][mask] = replacement
        raw_layer = source[start:end]
        layer_records.append(LayerRepairRecord(name, float(replacement), int(np.isnan(raw_layer).sum()), int(np.isposinf(raw_layer).sum()), int(np.isneginf(raw_layer).sum()), int(mask.sum())))
    repaired = output.astype(np.float16)
    repaired_values = repaired[~finite]
    LOGGER.info("repaired nonfinite values: method=%s count=%d", method, int((~finite).sum()))
    return RepairResult(repaired, method, int(np.isnan(source).sum()), int(np.isposinf(source).sum()), int(np.isneginf(source).sum()), int((~finite).sum()), tuple(layer_records), _histogram(repaired_values))


def _clip_bounds(data: np.ndarray, strategy: str) -> tuple[float, float]:
    finite = data[np.isfinite(data)]
    if not finite.size:
        raise ValueError("layer has no finite values")
    finite = finite.astype(np.float64)
    if strategy.startswith("layer_mean_std_"):
        scale = float(strategy.removeprefix("layer_mean_std_"))
        mean, std = float(np.mean(finite)), float(np.std(finite, ddof=0))
        return mean - scale * std, mean + scale * std
    if strategy.startswith("layer_pct_"):
        low, high = strategy.removeprefix("layer_pct_").split("_")
        return float(np.percentile(finite, float(low), method="linear")), float(np.percentile(finite, float(high), method="linear"))
    raise ValueError(f"unknown clip strategy: {strategy}")


def clip_candidate(values: np.ndarray, layers: Mapping[str, tuple[int, int]], strategy: str) -> ClipResult:
    """repair candidateだけからlayer boundsを計算し、finite値をclipする。"""
    if strategy not in CLIP_STRATEGIES:
        raise ValueError(f"unknown clip strategy: {strategy}")
    source = np.asarray(values, dtype=np.float16).reshape(-1)
    if not np.isfinite(source).all():
        raise ValueError("clip candidate must be finite")
    normalized_layers = validate_layer_offsets(layers, source.size)
    output64 = source.astype(np.float64)
    records: list[LayerClipRecord] = []
    for name, (start, end) in normalized_layers.items():
        before = output64[start:end]
        lower, upper = _clip_bounds(before, strategy)
        after = np.clip(before, lower, upper)
        records.append(LayerClipRecord(name, strategy, lower, upper, float(np.max(np.abs(before))), float(np.max(np.abs(after))), int(np.count_nonzero(after.astype(np.float16).view(np.uint16) != source[start:end].view(np.uint16)))))
        output64[start:end] = after
    clipped = output64.astype(np.float16)
    changed = clipped.view(np.uint16) != source.view(np.uint16)
    return ClipResult(clipped, strategy, tuple(records), int(changed.sum()), _histogram(source[changed]), _histogram(clipped[changed]), _histogram(clipped[changed].astype(np.float64) - source[changed].astype(np.float64)))


def _histogram(values: np.ndarray, bins: int = 32) -> dict[str, Any]:
    finite = np.asarray(values, dtype=np.float64)
    finite = finite[np.isfinite(finite)]
    if not finite.size:
        return {"edges": [], "counts": []}
    edges = np.histogram_bin_edges(finite, bins=min(bins, max(1, finite.size)))
    counts, edges = np.histogram(finite, bins=edges)
    return {"edges": edges.tolist(), "counts": counts.astype(int).tolist()}


def diagnostics(baseline: np.ndarray, candidate: np.ndarray, layers: Mapping[str, tuple[int, int]]) -> dict[str, Any]:
    """fp16 wordのbit差分、global/layer stats、bounded histogramを返す。"""
    left = np.asarray(baseline, dtype=np.float16).reshape(-1)
    right = np.asarray(candidate, dtype=np.float16).reshape(-1)
    if left.shape != right.shape:
        raise ValueError("baseline and candidate shapes differ")
    normalized_layers = validate_layer_offsets(layers, left.size)
    a, b = left.view(np.uint16), right.view(np.uint16)
    sign = ((a ^ b) & 0x8000) != 0
    exponent = ((a ^ b) & 0x7C00) != 0
    mantissa = ((a ^ b) & 0x03FF) != 0
    to31 = (((a >> 10) & 0x1F) != 31) & (((b >> 10) & 0x1F) == 31)
    finite_extreme = np.zeros(left.size, dtype=bool)
    layer_stats: list[dict[str, Any]] = []
    for name, (start, end) in normalized_layers.items():
        baseline_layer = left[start:end]
        candidate_layer = right[start:end]
        baseline_finite = np.isfinite(baseline_layer)
        candidate_finite = np.isfinite(candidate_layer)
        baseline_max = float(np.max(np.abs(baseline_layer[baseline_finite]))) if baseline_finite.any() else 0.0
        candidate_abs = np.abs(candidate_layer.astype(np.float64))
        threshold = baseline_max * 10.0
        if baseline_max == 0.0:
            layer_extreme = candidate_finite & (candidate_abs > 0.0)
        else:
            layer_extreme = candidate_finite & (candidate_abs > threshold)
        finite_extreme[start:end] = layer_extreme
        layer_stats.append({"layer": name, "sign_mismatch_count": int(sign[start:end].sum()), "exponent_mismatch_count": int(exponent[start:end].sum()), "mantissa_mismatch_count": int(mantissa[start:end].sum()), "exponent_to31_count": int(to31[start:end].sum()), "finite_extreme_outlier_count": int(layer_extreme.sum()), "nan_count": int(np.isnan(candidate_layer).sum()), "posinf_count": int(np.isposinf(candidate_layer).sum()), "neginf_count": int(np.isneginf(candidate_layer).sum()), "finite_count": int(candidate_finite.sum()), "max_abs_weight": float(np.max(candidate_abs[candidate_finite])) if candidate_finite.any() else 0.0})
    return {"total_count": int(left.size), "sign_mismatch_count": int(sign.sum()), "exponent_mismatch_count": int(exponent.sum()), "mantissa_mismatch_count": int(mantissa.sum()), "exponent_to31_count": int(to31.sum()), "finite_extreme_outlier_count": int(finite_extreme.sum()), "nan_count": int(np.isnan(right).sum()), "posinf_count": int(np.isposinf(right).sum()), "neginf_count": int(np.isneginf(right).sum()), "finite_count": int(np.isfinite(right).sum()), "max_abs_weight": float(np.max(np.abs(right[np.isfinite(right)]))) if np.isfinite(right).any() else 0.0, "layers": layer_stats, "histograms": {"baseline": _histogram(left), "candidate": _histogram(right), "delta": _histogram(right.astype(np.float64) - left.astype(np.float64))}}