"""
Synthetic double-exposure benchmark generator (MASTERPLAN WP-1 / Part IV).

Forward model (B&W, single emulsion):
  1. Generate two procedural scene images directly in linear light
     (real-photo sources would require inverse-sRGB linearization here).
  2. Scale by exposure ratio r; normalize H_total to the curve's useful range.
  3. Optional halation: Gaussian blur on H_total before the curve.
  4. Apply PiecewiseFilmCurve → optical density D.
  5. Add density-dependent grain (σ ∝ √(D–D_min)).
  6. Simulate transmittance T = 10^(−D) + sensor noise; sRGB-encode → raw scan.
  7. Optional JPEG round-trip at quality 85–95.
  8. Ground-truth positives = per-scene transmittance inverted and sRGB-encoded.

Outputs are dicts saved as .npz with keys:
  scan     (H,W,3) float32 sRGB — the raw negative scan (to upload to the app)
  gt_a     (H,W,3) float32 sRGB — GT positive of scene A
  gt_b     (H,W,3) float32 sRGB — GT positive of scene B (zeros for K=1)
  h_a      (H,W)   float32 — linear exposure map, scene A
  h_b      (H,W)   float32 — linear exposure map, scene B
  h_total  (H,W)   float32 — linear exposure total (H_a + H_b; halation-blurred if enabled)
  ratio    scalar float — H_a / H_b (large number for K=1)
  stock    str
  seed     int
  k1       bool
"""

from __future__ import annotations

import argparse
import os
from io import BytesIO
from pathlib import Path
from typing import List

import numpy as np
import torch

from densitometry import srgb_to_linear, linear_to_srgb, luminance_from_linear
from film_physics import get_film_curve, get_color_curves, COLOR_STOCK_PRESETS, PiecewiseFilmCurve


def _luminance(lin_rgb: np.ndarray) -> np.ndarray:
    """Rec. 709 luminance from linear RGB (H, W, 3) → (H, W)."""
    return luminance_from_linear(lin_rgb)


# ---------------------------------------------------------------------------
# Procedural scene generation
# ---------------------------------------------------------------------------

def _make_scene(rng: np.random.Generator, size: int) -> np.ndarray:
    """
    Generate a procedural grayscale scene (H, W) in linear light [0, 1].

    Uses a random mixture of gradient, blob, and noise primitives so that
    each seed produces a structurally distinct image.
    """
    h = w = size
    y_idx, x_idx = np.mgrid[:h, :w]

    scene_type = int(rng.integers(0, 4))

    if scene_type == 0:  # radial gradient with random center
        cx = float(rng.uniform(0.2, 0.8)) * w
        cy = float(rng.uniform(0.2, 0.8)) * h
        r = np.sqrt((x_idx - cx) ** 2 + (y_idx - cy) ** 2)
        lum = np.clip(1.0 - r / (0.7 * min(h, w)), 0.0, 1.0)

    elif scene_type == 1:  # linear gradient (random angle)
        angle = float(rng.uniform(0, np.pi))
        coord = x_idx * np.cos(angle) + y_idx * np.sin(angle)
        lo, hi = coord.min(), coord.max()
        lum = (coord - lo) / (hi - lo + 1e-8)
        if rng.integers(0, 2):
            lum = 1.0 - lum

    elif scene_type == 2:  # random Gaussian blobs
        lum = np.zeros((h, w), dtype=np.float32)
        n_blobs = int(rng.integers(2, 6))
        for _ in range(n_blobs):
            cx = float(rng.uniform(0.1, 0.9)) * w
            cy = float(rng.uniform(0.1, 0.9)) * h
            sigma = float(rng.uniform(0.08, 0.25)) * min(h, w)
            blob = np.exp(-((x_idx - cx) ** 2 + (y_idx - cy) ** 2) / (2 * sigma ** 2))
            lum += blob * float(rng.uniform(0.3, 1.0))
        lum = np.clip(lum / (lum.max() + 1e-8), 0.0, 1.0)

    else:  # smooth noise via sum of low-frequency sinusoids
        lum = np.zeros((h, w), dtype=np.float32)
        for _ in range(int(rng.integers(3, 8))):
            fx = float(rng.uniform(0.5, 3.0)) / w
            fy = float(rng.uniform(0.5, 3.0)) / h
            phase = float(rng.uniform(0, 2 * np.pi))
            amp = float(rng.uniform(0.1, 0.5))
            lum += amp * np.sin(2 * np.pi * (fx * x_idx + fy * y_idx) + phase)
        lum = (lum - lum.min()) / (lum.max() - lum.min() + 1e-8)

    # Slight overall brightness variation so scenes differ in key
    lum = lum * float(rng.uniform(0.5, 1.0)) + float(rng.uniform(0.0, 0.15))
    return np.clip(lum, 0.0, 1.0).astype(np.float32)


# ---------------------------------------------------------------------------
# Film curve helpers
# ---------------------------------------------------------------------------

def _apply_curve(h: np.ndarray, curve: PiecewiseFilmCurve) -> np.ndarray:
    """Apply PiecewiseFilmCurve to a linear exposure map (H, W) → density (H, W)."""
    log_h = np.log10(np.clip(h, 1e-8, None))
    log_h_t = torch.from_numpy(log_h).float().unsqueeze(0).unsqueeze(0)
    with torch.no_grad():
        d_t = curve(log_h_t)
    return d_t.squeeze().numpy().astype(np.float32)


def _transmittance_to_positive(t: np.ndarray, srgb: bool = True) -> np.ndarray:
    """
    Convert linear transmittance T = 10^(-D) to a positive image.

    Follows the same inversion the app applies to a negative scan:
    positive_display = 1 − sRGB_encode(T)
    so that bright scene exposure → bright positive pixel.
    """
    t_srgb = linear_to_srgb(t) if srgb else t
    return np.clip(1.0 - t_srgb, 0.0, 1.0).astype(np.float32)


# ---------------------------------------------------------------------------
# Core case generator
# ---------------------------------------------------------------------------

def generate_case(
    seed: int,
    stock: str = "Generic",
    ratio: float = 1.0,
    k1: bool = False,
    halation: bool = False,
    jpeg: bool = False,
    size: int = 64,
) -> dict:
    """
    Generate one synthetic double-exposure case.

    Args:
        seed:    RNG seed (determines scenes; reproducible for same args).
        stock:   Film stock preset name.
        ratio:   H_a / H_b exposure ratio (≥1; K=1 uses ratio=1000).
        k1:      If True, scene B is black (K=1 single-exposure control).
        halation: Apply Gaussian halation blur to H_total before the curve.
        jpeg:    Apply JPEG round-trip (quality 85–95) to the raw scan.
        size:    Square image size in pixels.

    Returns:
        Dict with keys: scan, gt_a, gt_b, h_a, h_b, h_total,
                        ratio, stock, seed, k1.
    """
    rng = np.random.default_rng(seed)
    is_color = stock in COLOR_STOCK_PRESETS
    if is_color:
        color_curves = get_color_curves(stock)
        mask_offsets = color_curves.mask_offset_rgb
        # per-channel curves for forward; scalar curve still used for compatibility keys below
        curve = get_film_curve("Generic")  # placeholder; per-ch used explicitly
        d_min = float(color_curves.g.d_min)  # for grain range, use g as rep
        d_max = float(color_curves.g.d_max)
    else:
        curve = get_film_curve(stock)
        d_min = float(curve.d_min)
        d_max = float(curve.d_max)

    # --- Generate scenes (linear light) ---
    lum_a = _make_scene(rng, size)  # (H, W) linear [0, 1]
    if k1:
        lum_b = np.zeros_like(lum_a)
        ratio = 1000.0  # sentinel for K=1
    else:
        lum_b = _make_scene(rng, size)

    # --- Scale to exposure ratio and normalize ---
    # H_a = ratio/(1+ratio) * scale, H_b = 1/(1+ratio) * scale
    h_a = lum_a * (ratio / (1.0 + ratio))
    h_b = lum_b * (1.0 / (1.0 + ratio))
    h_total_raw = h_a + h_b

    # Normalize so 95th-percentile of H_total sits at 0.8 (well inside curve's useful range)
    p95 = float(np.percentile(h_total_raw[h_total_raw > 1e-6], 95)) if h_total_raw.max() > 1e-6 else 1.0
    scale = 0.8 / (p95 + 1e-10)
    h_a = (h_a * scale).astype(np.float32)
    h_b = (h_b * scale).astype(np.float32)
    h_total_raw = (h_total_raw * scale).astype(np.float32)

    # --- Optional halation (blur H_total before the curve) ---
    if halation:
        from scipy.ndimage import gaussian_filter
        h_total = gaussian_filter(h_total_raw, sigma=max(1.0, size * 0.03)).astype(np.float32)
    else:
        h_total = h_total_raw

    # --- Apply film curve to get density (color: per-channel + mask offset) ---
    if is_color:
        h_a_rgb = np.stack([h_a] * 3, axis=-1)
        h_b_rgb = np.stack([h_b] * 3, axis=-1)
        h_total_raw_rgb = np.stack([h_total_raw] * 3, axis=-1)
        h_total_rgb = h_total_raw_rgb
        if halation:
            from scipy.ndimage import gaussian_filter
            h_total_rgb = gaussian_filter(h_total_rgb, sigma=max(1.0, size * 0.03)).astype(np.float32)
        d_obs_rgb = np.zeros_like(h_total_rgb)
        for c, ch_curve in enumerate([color_curves.r, color_curves.g, color_curves.b]):
            d_ch = _apply_curve(h_total_rgb[..., c], ch_curve)
            d_ch = d_ch + mask_offsets[c]
            d_obs_rgb[..., c] = d_ch
        d_obs = d_obs_rgb[..., 1]  # green for scalar compatibility
        # grain per channel
        d_noisy_rgb = np.zeros_like(d_obs_rgb)
        for c in range(3):
            d_ch = d_obs_rgb[..., c]
            d_range = np.clip((d_ch - d_min) / (d_max - d_min + 1e-8), 0.0, 1.0)
            grain_sigma = 0.015 * np.sqrt(d_range)
            grain = rng.normal(0.0, 1.0, d_ch.shape).astype(np.float32) * grain_sigma
            d_noisy_rgb[..., c] = np.clip(d_ch + grain, d_min, d_max)
        t_obs_rgb = np.power(10.0, -d_noisy_rgb).astype(np.float32)
        t_obs_rgb += rng.normal(0.0, 0.003, t_obs_rgb.shape).astype(np.float32)
        t_obs_rgb = np.clip(t_obs_rgb, 0.0, 1.0)
        scan = np.stack([linear_to_srgb(t_obs_rgb[..., c]) for c in range(3)], axis=-1).astype(np.float32)
    else:
        d_obs = _apply_curve(h_total, curve)  # D of combined negative
        # --- Add density-dependent grain ---
        d_range = np.clip((d_obs - d_min) / (d_max - d_min + 1e-8), 0.0, 1.0)
        grain_sigma = 0.015 * np.sqrt(d_range)
        grain = rng.normal(0.0, 1.0, d_obs.shape).astype(np.float32) * grain_sigma
        d_noisy = np.clip(d_obs + grain, d_min, d_max)

        # --- Transmittance + sensor noise → raw scan (sRGB, negative appearance) ---
        t_obs = np.power(10.0, -d_noisy).astype(np.float32)
        t_obs += rng.normal(0.0, 0.003, t_obs.shape).astype(np.float32)
        t_obs = np.clip(t_obs, 0.0, 1.0)

        raw_scan_linear = t_obs  # (H, W) linear transmittance
        raw_scan_srgb = linear_to_srgb(raw_scan_linear)  # brighter where unexposed
        # Stack to 3-channel grayscale (B&W benchmark; color extension in WP-8)
        scan = np.stack([raw_scan_srgb] * 3, axis=-1).astype(np.float32)

    # --- Optional JPEG round-trip ---
    if jpeg:
        from PIL import Image
        quality = int(rng.integers(85, 96))
        pil = Image.fromarray((scan * 255.0).clip(0, 255).astype(np.uint8))
        buf = BytesIO()
        pil.save(buf, format="JPEG", quality=quality)
        buf.seek(0)
        pil = Image.open(buf)
        scan = np.asarray(pil, dtype=np.float32) / 255.0

    # --- Ground-truth positives (per-scene negative inverted to positive) ---
    # Same pipeline as the scan but using per-scene H (no halation, no grain/noise for clean GT)
    if is_color:
        d_a_rgb = np.zeros((size, size, 3), dtype=np.float32)
        for c, ch_curve in enumerate([color_curves.r, color_curves.g, color_curves.b]):
            d_a_rgb[..., c] = _apply_curve(h_a, ch_curve) + mask_offsets[c]
        t_a_rgb = np.power(10.0, -d_a_rgb).astype(np.float32)
        gt_a = np.stack([_transmittance_to_positive(t_a_rgb[..., c]) for c in range(3)], axis=-1).astype(np.float32)

        if k1:
            gt_b = np.zeros((size, size, 3), dtype=np.float32)
        else:
            d_b_rgb = np.zeros((size, size, 3), dtype=np.float32)
            for c, ch_curve in enumerate([color_curves.r, color_curves.g, color_curves.b]):
                d_b_rgb[..., c] = _apply_curve(h_b, ch_curve) + mask_offsets[c]
            t_b_rgb = np.power(10.0, -d_b_rgb).astype(np.float32)
            gt_b = np.stack([_transmittance_to_positive(t_b_rgb[..., c]) for c in range(3)], axis=-1).astype(np.float32)

        # scalar keys from green for compatibility
        h_a = h_a  # already
        h_b = h_b
        h_total = h_total
        # add rgb keys
        extra = {
            "h_a_rgb": h_a_rgb,
            "h_b_rgb": h_b_rgb,
            "h_total_rgb": h_total_rgb,
            "is_color": True,
        }
    else:
        d_a = _apply_curve(h_a, curve)
        t_a = np.power(10.0, -d_a).astype(np.float32)
        gt_a_gray = _transmittance_to_positive(t_a)
        gt_a = np.stack([gt_a_gray] * 3, axis=-1).astype(np.float32)

        if k1:
            gt_b = np.zeros((size, size, 3), dtype=np.float32)
        else:
            d_b = _apply_curve(h_b, curve)
            t_b = np.power(10.0, -d_b).astype(np.float32)
            gt_b_gray = _transmittance_to_positive(t_b)
            gt_b = np.stack([gt_b_gray] * 3, axis=-1).astype(np.float32)
        extra = {}

    ret = {
        "scan": scan,
        "gt_a": gt_a,
        "gt_b": gt_b,
        "h_a": h_a,
        "h_b": h_b,
        "h_total": h_total,
        "ratio": float(ratio),
        "stock": stock,
        "seed": int(seed),
        "k1": bool(k1),
    }
    ret.update(extra)
    return ret


# ---------------------------------------------------------------------------
# Dataset generator (batch)
# ---------------------------------------------------------------------------

def generate_dataset(
    n: int = 8,
    stock: str = "Generic",
    ratio_min: float = 1.0,
    ratio_max: float = 8.0,
    k1_fraction: float = 0.25,
    halation: bool = False,
    jpeg: bool = False,
    seed: int = 42,
    size: int = 64,
) -> List[dict]:
    """
    Generate ``n`` synthetic cases seeded deterministically.

    K=1 cases (proportion ``k1_fraction``) are placed at the end so they have
    consistent indices regardless of total n.
    """
    rng = np.random.default_rng(seed)
    n_k1 = max(0, round(n * k1_fraction))
    n_k2 = n - n_k1

    cases = []

    # K=2 cases with log-uniform ratio sampling
    for i in range(n_k2):
        case_seed = int(rng.integers(0, 2**31))
        log_ratio = rng.uniform(np.log(ratio_min), np.log(ratio_max))
        ratio = float(np.exp(log_ratio))
        cases.append(
            generate_case(
                seed=case_seed,
                stock=stock,
                ratio=ratio,
                k1=False,
                halation=halation,
                jpeg=jpeg,
                size=size,
            )
        )

    # K=1 control cases
    for _ in range(n_k1):
        case_seed = int(rng.integers(0, 2**31))
        cases.append(
            generate_case(
                seed=case_seed,
                stock=stock,
                ratio=1000.0,
                k1=True,
                halation=halation,
                jpeg=jpeg,
                size=size,
            )
        )

    return cases


# ---------------------------------------------------------------------------
# Persistence
# ---------------------------------------------------------------------------

def save_case(case: dict, path: str | Path) -> None:
    """Save a case dict to a .npz file. For color cases, include per-channel rgb keys."""
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    kwargs = {
        "scan": case["scan"],
        "gt_a": case["gt_a"],
        "gt_b": case["gt_b"],
        "h_a": case["h_a"],
        "h_b": case["h_b"],
        "h_total": case["h_total"],
        "ratio": np.float32(case["ratio"]),
        "stock": np.str_(case["stock"]),
        "seed": np.int64(case["seed"]),
        "k1": np.bool_(case["k1"]),
    }
    if "h_a_rgb" in case:
        kwargs["h_a_rgb"] = case["h_a_rgb"]
        kwargs["h_b_rgb"] = case["h_b_rgb"]
        kwargs["h_total_rgb"] = case["h_total_rgb"]
    np.savez_compressed(path, **kwargs)


def load_case(path: str | Path) -> dict:
    """Load a case dict from a .npz file. Color cases include per-channel rgb keys if present."""
    data = np.load(path, allow_pickle=False)
    case = {
        "scan": data["scan"],
        "gt_a": data["gt_a"],
        "gt_b": data["gt_b"],
        "h_a": data["h_a"],
        "h_b": data["h_b"],
        "h_total": data["h_total"],
        "ratio": float(data["ratio"]),
        "stock": str(data["stock"]),
        "seed": int(data["seed"]),
        "k1": bool(data["k1"]),
    }
    if "h_a_rgb" in data:
        case["h_a_rgb"] = data["h_a_rgb"]
        case["h_b_rgb"] = data["h_b_rgb"]
        case["h_total_rgb"] = data["h_total_rgb"]
    return case


def load_fixtures(fixtures_dir: str | Path | None = None) -> List[dict]:
    """Load all .npz fixture cases from the fixtures directory."""
    if fixtures_dir is None:
        fixtures_dir = Path(__file__).parent / "fixtures"
    fixtures_dir = Path(fixtures_dir)
    paths = sorted(fixtures_dir.glob("case_*.npz"))
    if not paths:
        raise FileNotFoundError(f"No fixture cases found in {fixtures_dir}")
    return [load_case(p) for p in paths]


def select_k2_cases(cases: List[dict], limit: int | None = None) -> List[dict]:
    """The working set for a benchmark run: the K=2 cases, optionally capped to `limit`.

    Canonical home for the K2-filter + cap that both benches (double_dip, refine_bench)
    consume — added in WP-1.2 so the two harnesses share one selection policy.
    """
    k2 = [c for c in cases if not c.get("k1", False)]
    if limit is not None and len(k2) > limit:
        k2 = k2[:limit]
    return k2


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------

def _cli() -> None:
    parser = argparse.ArgumentParser(
        description="Generate synthetic double-exposure benchmark cases (WP-1)."
    )
    parser.add_argument("--n", type=int, default=50, help="Number of cases to generate")
    parser.add_argument("--stock", type=str, default="Generic", help="Film stock preset")
    parser.add_argument("--ratio-min", type=float, default=1.0, dest="ratio_min")
    parser.add_argument("--ratio-max", type=float, default=8.0, dest="ratio_max")
    parser.add_argument("--k1-fraction", type=float, default=0.1, dest="k1_fraction",
                        help="Fraction of cases that are K=1 single-exposure controls")
    parser.add_argument("--halation", action="store_true")
    parser.add_argument("--jpeg", action="store_true")
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--size", type=int, default=256, help="Image size (pixels, square)")
    parser.add_argument("--output-dir", type=str, default="synth/data", dest="output_dir")
    args = parser.parse_args()

    cases = generate_dataset(
        n=args.n,
        stock=args.stock,
        ratio_min=args.ratio_min,
        ratio_max=args.ratio_max,
        k1_fraction=args.k1_fraction,
        halation=args.halation,
        jpeg=args.jpeg,
        seed=args.seed,
        size=args.size,
    )

    out_dir = Path(args.output_dir)
    out_dir.mkdir(parents=True, exist_ok=True)
    for i, case in enumerate(cases):
        save_case(case, out_dir / f"case_{i:03d}.npz")

    n_k1 = sum(1 for c in cases if c["k1"])
    print(
        f"Generated {len(cases)} cases ({len(cases)-n_k1} K=2, {n_k1} K=1) "
        f"→ {out_dir}/"
    )


if __name__ == "__main__":
    _cli()