#!/usr/bin/env python
"""Generate paired before/after images for every augmentation effect, for slides.

For each of the 25 configurable effects, renders two samples that are
identical in every respect (same template, signature, placement, box) except
that one has every augmentation disabled ("before") and the other has just
that one effect forced on at a fixed strength ("after"). Because placement is
frozen and both runs share the same (seed, sample_index), the underlying
document/signature/box are pixel-identical up to the effect itself.

Signature-level effects are saved as a tight crop around the signature box
(where the effect is visible at slide scale); document-level effects are
saved as the full page. Also writes a JSON manifest with the resolved
strength/params/formula description used per effect, for slide captions.

Usage:
    python scripts/generate_before_after_assets.py
"""

from __future__ import annotations

import json
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from PIL import Image

from src.config.loader import load_config
from src.generator import GeneratorContext
from src.visualization.visualizer import _crop_region

SIGNATURE_EFFECTS = [
    "rotation", "stroke_dilation",
    "ink_fading", "ink_bleed", "motion_blur", "gaussian_blur",
]
DOCUMENT_EFFECTS = [
    "perspective", "contrast", "gamma", "coffee_stain", "stain",
    "shadow", "crumple", "discoloration", "page_tilt", "dust", "scanner_streaks",
    "scanner_noise", "salt_pepper", "jpeg_compression",
]
ALL_EFFECTS = SIGNATURE_EFFECTS + DOCUMENT_EFFECTS

OUT_ROOT = Path(__file__).resolve().parents[1] / "degradation"
TEMPLATE_ID = "bank_doc_05"
FORCED_STRENGTH = 1.0  # max strength: matches the "range at strength=1.0" figures already on the formula slides
SAMPLE_INDEX = 2  # picked for a bolder/higher-ink-fill signature than index 1, more visible at slide scale
CROP_PAD = 45
CROP_UPSCALE = 2

# Purely tonal/textural document effects are nearly invisible at full-page slide
# scale, so also cut a fixed, upscaled body-text zoom for these (kept identical
# between before/after so the comparison is apples-to-apples). Structural effects
# (crumple, page_tilt, perspective, shadow, scanner_streaks) already read clearly
# on the full page and don't need one.
ZOOM_EFFECTS = {
    "contrast", "gamma", "coffee_stain", "stain",
    "discoloration", "dust", "scanner_noise", "salt_pepper", "jpeg_compression",
}
ZOOM_BOX = (40, 110, 460, 340)  # (x1, y1, x2, y2) in full-page pixel space

DISPLAY_NAME = {
    "rotation": "Rotation jitter",
    "stroke_dilation": "Stroke dilation",
    "ink_fading": "Ink fading",
    "ink_bleed": "Ink bleed",
    "motion_blur": "Motion blur",
    "gaussian_blur": "Gaussian blur",
    "perspective": "Perspective warp",
    "contrast": "Contrast shift",
    "gamma": "Gamma shift",
    "coffee_stain": "Coffee stain",
    "stain": "Water / paper stain",
    "shadow": "Shadow",
    "crumple": "Crumple / wrinkle",
    "discoloration": "Discoloration",
    "page_tilt": "Page tilt",
    "dust": "Dust specks",
    "scanner_streaks": "Scanner streaks",
    "scanner_noise": "Scanner noise",
    "salt_pepper": "Salt & pepper noise",
    "jpeg_compression": "JPEG compression",
}


def base_overrides() -> list[str]:
    return [
        f"document.template={TEMPLATE_ID}",
        "placement.rotation_deg.min=0", "placement.rotation_deg.max=0",
        "placement.scale_factor.min=0.85", "placement.scale_factor.max=0.85",
        "placement.x_offset=0", "placement.y_offset=0",
        "placement.misplacement_probability=0.0",
        "generation.run_validation=false",
    ]


def all_disabled_overrides() -> list[str]:
    overrides = base_overrides()
    for name in ALL_EFFECTS:
        overrides.append(f"augmentation.{name}.enabled=false")
        overrides.append(f"augmentation.{name}.probability=0.0")
    return overrides


def effect_overrides(target_effect: str) -> list[str]:
    overrides = base_overrides()
    for name in ALL_EFFECTS:
        is_target = name == target_effect
        overrides.append(f"augmentation.{name}.enabled={'true' if is_target else 'false'}")
        overrides.append(f"augmentation.{name}.probability={'1.0' if is_target else '0.0'}")
    overrides.append(f"augmentation.{target_effect}.strength={FORCED_STRENGTH}")
    return overrides


def main() -> None:
    OUT_ROOT.mkdir(parents=True, exist_ok=True)
    manifest: dict[str, dict] = {}

    baseline_config = load_config(overrides=all_disabled_overrides())
    baseline_ctx = GeneratorContext(baseline_config)
    baseline_meta, baseline_img, _ = baseline_ctx.run_sample(SAMPLE_INDEX)
    box = baseline_meta["geometry"]["target_box"]

    for effect in ALL_EFFECTS:
        out_dir = OUT_ROOT / effect
        out_dir.mkdir(parents=True, exist_ok=True)

        config = load_config(overrides=effect_overrides(effect))
        ctx = GeneratorContext(config)
        meta, after_img, _ = ctx.run_sample(SAMPLE_INDEX)

        before_img, after_img_final = baseline_img, after_img
        if effect in SIGNATURE_EFFECTS:
            before_out = _crop_region(before_img, box, pad=CROP_PAD)
            after_out = _crop_region(after_img_final, box, pad=CROP_PAD)
            before_pil = Image.fromarray(before_out)
            after_pil = Image.fromarray(after_out)
            before_pil = before_pil.resize(
                (before_pil.width * CROP_UPSCALE, before_pil.height * CROP_UPSCALE), Image.LANCZOS,
            )
            after_pil = after_pil.resize(
                (after_pil.width * CROP_UPSCALE, after_pil.height * CROP_UPSCALE), Image.LANCZOS,
            )
        else:
            before_pil, after_pil = Image.fromarray(before_img), Image.fromarray(after_img_final)

        before_pil.save(out_dir / "before.png")
        after_pil.save(out_dir / "after.png")

        if effect in ZOOM_EFFECTS:
            x1, y1, x2, y2 = ZOOM_BOX
            zb = Image.fromarray(before_img[y1:y2, x1:x2])
            za = Image.fromarray(after_img_final[y1:y2, x1:x2])
            zb = zb.resize((zb.width * CROP_UPSCALE, zb.height * CROP_UPSCALE), Image.LANCZOS)
            za = za.resize((za.width * CROP_UPSCALE, za.height * CROP_UPSCALE), Image.LANCZOS)
            zb.save(out_dir / "before_zoom.png")
            za.save(out_dir / "after_zoom.png")

        rec = meta["augmentation"]["effects"].get(effect, {})
        manifest[effect] = {
            "display_name": DISPLAY_NAME[effect],
            "category": "signature" if effect in SIGNATURE_EFFECTS else "document",
            "strength": rec.get("strength"),
            "params": rec.get("params"),
            "difficulty_score": meta["difficulty"]["score"],
            "difficulty_tier": meta["difficulty"]["tier"],
        }
        print(f"[{effect}] before/after -> {out_dir}  (strength={rec.get('strength')}, params={rec.get('params')})")

    manifest_path = OUT_ROOT / "before_after_manifest.json"
    manifest_path.write_text(json.dumps(manifest, indent=2), encoding="utf-8")
    print(f"\nWrote manifest -> {manifest_path}")


if __name__ == "__main__":
    main()