"""AiArter format study (October 2026): JPEG vs WebP vs AVIF vs PNG on our own product photos.

Reproducible and offline. Run from the repo root:

    ~/.pyenv/versions/ai_tools_env/bin/python scripts/studies/image_format_study.py

Inputs are the RGB product and scene photos under static/assets/img/aiarter/ whose long side is at
least 800 px (the INCLUDE list below, chosen by hand: no logos, icons, og images, tool screenshots,
documents, text overlays, or files that are a copy or a tool-processed variant of another input).
Each input is decoded to 8-bit RGB, which is the reference. Every encode is decoded again and
compared with the reference by SSIM on the luma channel (BT.601 Y, 11x11 Gaussian window, sigma 1.5,
K1=0.01, K2=0.03, "valid" borders, as in Wang et al. 2004).

Writes blog/aiarter_content/data/format_study_2026_10.json.
"""

import io
import json
import os
import platform
import sys
from datetime import date

import numpy as np
from PIL import Image, ImageCms, features

import PIL

ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))
SRC = os.path.join(ROOT, "static", "assets", "img", "aiarter")
OUT = os.path.join(ROOT, "blog", "aiarter_content", "data", "format_study_2026_10.json")

INCLUDE = [
    # Demo sets: original phone-style shot, two AI scenes and a white-background version.
    "demo/apparel-original.webp", "demo/apparel-scene1.webp", "demo/apparel-scene2.webp", "demo/apparel-white.webp",
    "demo/food-original.webp", "demo/food-scene1.webp", "demo/food-scene2.webp", "demo/food-white.webp",
    "demo/jewelry-original.webp", "demo/jewelry-scene1.webp", "demo/jewelry-scene2.webp", "demo/jewelry-white.webp",
    # Serum bottle: original, white, colour backdrops and scenes.
    "serum-ad.webp", "serum-bathroom.webp", "serum-color-beige.webp", "serum-color-navy.webp",
    "serum-color-peach.webp", "serum-color-sage.webp", "serum-luxury.webp", "serum-marble.webp",
    "serum-original.webp", "serum-outdoors.webp", "serum-pastel.webp", "serum-story.webp",
    "serum-white.webp", "serum-wood.webp",
    # Tool-page examples that are plain product or scene photos.
    "tools/edit-after.webp", "tools/gen-bakery.webp", "tools/gen-mug.webp", "tools/gen-sneaker.webp",
    "tools/gen-wallet.webp", "tools/restore-after.webp", "tools/t2v-mug-poster.webp",
]

# (label, format, Pillow save kwargs). The explicit method/speed/subsampling values are Pillow 11.3's
# defaults, written out so the JSON records them. JPEG is baseline unless noted.
SETTINGS = [
    ("JPEG q95", "JPEG", {"quality": 95, "optimize": True, "subsampling": "4:2:0"}),
    ("JPEG q85", "JPEG", {"quality": 85, "optimize": True, "subsampling": "4:2:0"}),
    ("JPEG q75", "JPEG", {"quality": 75, "optimize": True, "subsampling": "4:2:0"}),
    ("JPEG q60", "JPEG", {"quality": 60, "optimize": True, "subsampling": "4:2:0"}),
    ("WebP q90", "WEBP", {"quality": 90, "method": 4}),
    ("WebP q80", "WEBP", {"quality": 80, "method": 4}),
    ("WebP q70", "WEBP", {"quality": 70, "method": 4}),
    ("AVIF q70", "AVIF", {"quality": 70, "speed": 6, "subsampling": "4:2:0"}),
    ("AVIF q50", "AVIF", {"quality": 50, "speed": 6, "subsampling": "4:2:0"}),
    ("AVIF q35", "AVIF", {"quality": 35, "speed": 6, "subsampling": "4:2:0"}),
    ("PNG optimize", "PNG", {"optimize": True}),
]


def _gauss_kernel(size=11, sigma=1.5):
    x = np.arange(size) - (size - 1) / 2
    k = np.exp(-(x ** 2) / (2 * sigma ** 2))
    return k / k.sum()


_K = _gauss_kernel()


def _filt(a):
    """Separable 11x11 Gaussian filter, 'valid' borders."""
    w = np.lib.stride_tricks.sliding_window_view(a, len(_K), axis=1) @ _K
    return np.lib.stride_tricks.sliding_window_view(w, len(_K), axis=0) @ _K


def luma(img):
    a = np.asarray(img.convert("RGB"), dtype=np.float64)
    return 0.299 * a[..., 0] + 0.587 * a[..., 1] + 0.114 * a[..., 2]


def ssim(x, y, L=255.0):
    c1, c2 = (0.01 * L) ** 2, (0.03 * L) ** 2
    mx, my = _filt(x), _filt(y)
    sxx = _filt(x * x) - mx * mx
    syy = _filt(y * y) - my * my
    sxy = _filt(x * y) - mx * my
    m = ((2 * mx * my + c1) * (2 * sxy + c2)) / ((mx * mx + my * my + c1) * (sxx + syy + c2))
    return float(m.mean())


def encode(img, fmt, kwargs):
    buf = io.BytesIO()
    img.save(buf, fmt, **kwargs)
    data = buf.getvalue()
    return data, Image.open(io.BytesIO(data)).convert("RGB")


def stats(values):
    v = np.asarray(values, dtype=float)
    q1, med, q3 = np.percentile(v, [25, 50, 75])
    return {"median": round(float(med), 4), "q1": round(float(q1), 4), "q3": round(float(q3), 4),
            "min": round(float(v.min()), 4), "max": round(float(v.max()), 4), "n": int(v.size)}


def main():
    srgb = ImageCms.ImageCmsProfile(ImageCms.createProfile("sRGB")).tobytes()
    images, rows = [], []
    for rel in INCLUDE:
        path = os.path.join(SRC, rel)
        src = Image.open(path)
        meta = {k: len(src.info[k]) for k in ("exif", "icc_profile", "xmp") if src.info.get(k)}
        ref = src.convert("RGB")
        ref_y = luma(ref)
        w, h = ref.size
        images.append({"file": rel, "width": w, "height": h, "source_format": src.format,
                       "source_bytes": os.path.getsize(path), "source_metadata_bytes": meta})
        per = {"file": rel, "pixels": w * h, "results": {}}
        for label, fmt, kw in SETTINGS:
            data, dec = encode(ref, fmt, kw)
            per["results"][label] = {"bytes": len(data), "ssim": round(ssim(ref_y, luma(dec)), 5)}
        # Cost of embedding an sRGB ICC profile in a JPEG q85.
        with_icc, _ = encode(ref, "JPEG", {"quality": 85, "optimize": True, "icc_profile": srgb})
        per["jpeg_q85_icc_srgb_bytes"] = len(with_icc)
        # Baseline vs progressive JPEG q85 (same pixels, different file layout).
        prog, _ = encode(ref, "JPEG", {"quality": 85, "optimize": True, "progressive": True})
        per["jpeg_q85_progressive_bytes"] = len(prog)
        rows.append(per)
        print(rel, {k: v["bytes"] for k, v in per["results"].items()}, file=sys.stderr)

    summary = {}
    base = [r["results"]["JPEG q85"]["bytes"] for r in rows]
    for label, _, _ in SETTINGS:
        b = [r["results"][label]["bytes"] for r in rows]
        summary[label] = {
            "kb": stats([x / 1000 for x in b]),
            "bits_per_pixel": stats([x * 8 / r["pixels"] for x, r in zip(b, rows)]),
            "ssim": stats([r["results"][label]["ssim"] for r in rows]),
            "size_vs_jpeg_q85_pct": stats([100.0 * x / y for x, y in zip(b, base)]),
            "smaller_than_jpeg_q85_count": int(sum(x < y for x, y in zip(b, base))),
            "under_100kb_count": int(sum(x < 100_000 for x in b)),
            "under_200kb_count": int(sum(x < 200_000 for x in b)),
            "under_1mb_count": int(sum(x < 1_000_000 for x in b)),
        }
    summary["JPEG q85 + sRGB ICC"] = {
        "added_bytes": stats([r["jpeg_q85_icc_srgb_bytes"] - r["results"]["JPEG q85"]["bytes"] for r in rows]),
        "added_pct": stats([100.0 * (r["jpeg_q85_icc_srgb_bytes"] - y) / y for r, y in zip(rows, base)]),
    }
    # Strip-metadata savings: EXIF, ICC and XMP bytes found in the sources.
    summary["source_metadata"] = {
        "files_with_metadata": sum(bool(i["source_metadata_bytes"]) for i in images),
        "total_bytes": sum(sum(i["source_metadata_bytes"].values()) for i in images),
    }
    summary["JPEG q85 progressive"] = {
        "size_vs_baseline_pct": stats([100.0 * r["jpeg_q85_progressive_bytes"] / y for r, y in zip(rows, base)]),
    }
    # For each image, the smallest file that matches or beats JPEG q85's SSIM.
    match, match_pct, webp_avif = {}, [], {}
    for r, y in zip(rows, base):
        target = r["results"]["JPEG q85"]["ssim"]
        ok = [(v["bytes"], k) for k, v in r["results"].items() if v["ssim"] >= target]
        size, best = min(ok)
        match[best] = match.get(best, 0) + 1
        match_pct.append(100.0 * size / y)
        r["smallest_at_or_above_jpeg_q85_ssim"] = best
        # Same question restricted to one format family.
        for fam in ("WebP", "AVIF", "JPEG"):
            fam_ok = [v["bytes"] for k, v in r["results"].items() if k.startswith(fam) and v["ssim"] >= target]
            webp_avif.setdefault(fam, []).append(100.0 * min(fam_ok) / y if fam_ok else None)
    summary["smallest_setting_at_or_above_jpeg_q85_ssim"] = {
        "counts": match, "size_vs_jpeg_q85_pct": stats(match_pct),
        "by_format": {fam: {"found_count": sum(v is not None for v in vals),
                            "size_vs_jpeg_q85_pct": stats([v for v in vals if v is not None])}
                      for fam, vals in webp_avif.items() if any(v is not None for v in vals)},
    }

    out = {
        "study": "AiArter image format study",
        "run_date": date.today().isoformat(),
        "script": "scripts/studies/image_format_study.py",
        "sample_size": len(rows),
        "inputs_note": ("All inputs are AiArter's own product and scene photos, stored as lossy WebP at "
                        "800-1080 px on the long side. None is 2000 px or larger, so the 1000 vs 2000 px "
                        "resize comparison was skipped."),
        "metric": "SSIM on BT.601 luma, 11x11 Gaussian window sigma 1.5, valid borders, vs decoded source",
        "environment": {
            "python": platform.python_version(), "pillow": PIL.__version__, "numpy": np.__version__,
            "libwebp": features.version("webp"), "libavif": features.version("avif"),
            "libjpeg": features.version("jpg"), "libjpeg_turbo": bool(features.check_feature("libjpeg_turbo")),
            "zlib": features.version("zlib"), "platform": platform.platform(),
        },
        "settings": {label: {"format": fmt, **kw} for label, fmt, kw in SETTINGS},
        "images": images,
        "summary": summary,
        "per_image": rows,
    }
    os.makedirs(os.path.dirname(OUT), exist_ok=True)
    with open(OUT, "w") as fh:
        json.dump(out, fh, indent=1)
    print(OUT, file=sys.stderr)


if __name__ == "__main__":
    main()