"""Offline baker: real flow-stabilized painterly Right-axis variants (scales
design §1/§4).

Productionizes the session-0008 POC (``~/hef-poc/flow_restyle.py``): SD img2img
on KEYFRAMES + Farneback optical-flow warp on in-between frames = a temporally
calm painterly restyle (no per-frame "boiling" — the look the operator approved
in 0008). The one new idea on top of the POC is a **strength curve**: each Right
level 1-4 maps to an increasing keyframe img2img ``strength`` so the four
discrete pre-baked variants form a subtle -> strong dreamlike ramp. Level 0 is
always the raw base (handled by the runtime, not baked here).

The curve (``right_strength_params``) is pure + unit-tested; the rendered media
is verified by eye (like the other ``simulator/setup_*_media.py`` scripts). torch
/ diffusers are imported lazily inside the bake functions so importing this
module (e.g. for the curve test) is fast and dependency-light.

Usage:
    .venv/bin/python -m simulator.bake_right_variants [--scale forest]
                     [--fps 12] [--keyint 24]

Requires: torch + Apple MPS, diffusers, ``stabilityai/sd-turbo`` (cached),
opencv-python, imageio-ffmpeg. ~2.7 min/clip x 4 levels ~= ~11 min local compute.
"""
from __future__ import annotations

import argparse
import glob
import os
import subprocess
import tempfile
import time
from pathlib import Path

MEDIA = Path(__file__).parent / "sample_media"

RIGHT_LEVELS = (1, 2, 3, 4)

# Keyframe img2img strength per Right level (by-eye ramp). sd-turbo's painterly
# transformation is sharply NON-LINEAR in strength — it barely registers below
# ~0.5 and ramps hard through ~0.5-0.65 (verified by eye, session 0012: a linear
# 0.28-0.58 ramp left levels 1-3 ~indistinguishable from raw). So the curve is
# clustered in that active band for real perceptual progression: a gentle dream
# haze (L1) -> clearly painterly (L3, ~the POC's proven look) -> strongest (L4).
# Still by-eye tunable — re-run the baker after editing.
_KEY_STRENGTH = {1: 0.46, 2: 0.52, 3: 0.58, 4: 0.64}
_TWEEN_RATIO = 0.36  # POC ratio (0.18 / 0.5): tweens are a light refine only.

# The restyle target (matches the POC spike that the operator approved).
PROMPT = (
    "impressionist oil painting, dreamlike, soft painterly brushstrokes, "
    "ethereal, luminous, fine art"
)
NEG = "text, watermark, frame, border, ugly, deformed, lowres"
SEED = 42


def right_strength_params(level: int) -> tuple[float, float]:
    """(keyframe_strength, tween_strength) for a Right level in RIGHT_LEVELS."""
    if level not in _KEY_STRENGTH:
        raise ValueError(f"Right level must be one of {RIGHT_LEVELS}, got {level}")
    key = _KEY_STRENGTH[level]
    return key, round(key * _TWEEN_RATIO, 4)


# --- the restyle engine (verbatim port of ~/hef-poc/flow_restyle.py) ---


def _ffmpeg() -> str:
    import imageio_ffmpeg

    return imageio_ffmpeg.get_ffmpeg_exe()


def _run(cmd: list[str]) -> None:
    subprocess.run(cmd, check=True, capture_output=True)


def _extract(ff: str, src: str, d: str, fps: float) -> list[str]:
    _run([ff, "-hide_banner", "-v", "error", "-i", src, "-vf", f"fps={fps}",
          os.path.join(d, "f_%05d.png"), "-y"])
    return sorted(glob.glob(os.path.join(d, "f_*.png")))


def _warp(prev_bgr, src_prev_gray, src_cur_gray):
    """Warp the previous stylized frame to align with the current source frame
    via Farneback optical flow — this is what kills the per-frame flicker."""
    import cv2
    import numpy as np

    flow = cv2.calcOpticalFlowFarneback(
        src_prev_gray, src_cur_gray, None,
        pyr_scale=0.5, levels=3, winsize=21, iterations=3,
        poly_n=7, poly_sigma=1.5, flags=0)
    h, w = src_cur_gray.shape
    gx, gy = np.meshgrid(np.arange(w), np.arange(h))
    mx = (gx + flow[..., 0]).astype(np.float32)
    my = (gy + flow[..., 1]).astype(np.float32)
    return cv2.remap(prev_bgr, mx, my, cv2.INTER_LINEAR, borderMode=cv2.BORDER_REFLECT)


def _load_pipe(model: str):
    """Load the SD img2img pipeline once (reused across all levels)."""
    import torch
    from diffusers import AutoPipelineForImage2Image

    dev = "mps" if torch.backends.mps.is_available() else "cpu"
    pipe = AutoPipelineForImage2Image.from_pretrained(
        model, torch_dtype=torch.float16 if dev == "mps" else torch.float32,
        safety_checker=None).to(dev)
    return pipe, dev


def restyle_clip(pipe, dev, src: str, out: str, key_strength: float,
                 tween_strength: float, *, fps: float = 12.0, keyint: int = 24) -> None:
    """Restyle one clip: SD img2img keyframes + flow-warped light-refine tweens."""
    import cv2
    import numpy as np
    import torch
    from PIL import Image

    ff = _ffmpeg()
    gen = torch.Generator(device=dev).manual_seed(SEED)

    def img2img(pil, strength):
        steps = max(2, int(np.ceil(1.0 / max(strength, 0.05))) + 1)
        return pipe(prompt=PROMPT, negative_prompt=NEG, image=pil,
                    strength=strength, guidance_scale=0.0,
                    num_inference_steps=steps, generator=gen).images[0]

    with tempfile.TemporaryDirectory() as d:
        frames = _extract(ff, src, d, fps)
        print(f"    [frames] {len(frames)}  key_s={key_strength} tween_s={tween_strength}")
        prev_styled = None       # BGR uint8
        prev_src_gray = None
        t0 = time.time()
        for i, fp in enumerate(frames):
            src_bgr = cv2.imread(fp)
            h0, w0 = src_bgr.shape[:2]
            scale = 768 / max(h0, w0)
            wh = (int(w0 * scale) // 8 * 8, int(h0 * scale) // 8 * 8)
            src_small = cv2.resize(src_bgr, wh)
            src_gray = cv2.cvtColor(src_small, cv2.COLOR_BGR2GRAY)

            if i % keyint == 0 or prev_styled is None:
                pil = Image.fromarray(cv2.cvtColor(src_small, cv2.COLOR_BGR2RGB))
                out_img = img2img(pil, key_strength)
            else:
                warped = _warp(prev_styled, prev_src_gray, src_gray)
                pil = Image.fromarray(cv2.cvtColor(warped, cv2.COLOR_BGR2RGB))
                out_img = img2img(pil, tween_strength)   # light refine only
            styled = cv2.cvtColor(np.array(out_img), cv2.COLOR_RGB2BGR)
            prev_styled, prev_src_gray = styled, src_gray
            cv2.imwrite(os.path.join(d, f"s_{i+1:05d}.png"),
                        cv2.resize(styled, (w0, h0)))
            if i % 12 == 0:
                print(f"      frame {i+1}/{len(frames)}  {time.time()-t0:.1f}s")

        _run([ff, "-hide_banner", "-v", "error", "-framerate", str(fps),
              "-i", os.path.join(d, "s_%05d.png"), "-c:v", "libx264", "-crf", "20",
              "-preset", "medium", "-pix_fmt", "yuv420p", out, "-y"])
    print(f"    [done] {time.time()-t0:.1f}s -> {out}")


def bake_scale(scale: str = "forest", *, fps: float = 12.0, keyint: int = 24,
               model: str = "stabilityai/sd-turbo") -> None:
    """Bake all RIGHT_LEVELS variants for one scale's base clip."""
    base = MEDIA / scale / "base.mp4"
    if not base.exists():
        raise FileNotFoundError(
            f"no base clip at {base} — run the media-setup scripts first")
    print(f"[bake] scale={scale} base={base}")
    pipe, dev = _load_pipe(model)
    print(f"[device] {dev}")
    for level in RIGHT_LEVELS:
        key_s, tween_s = right_strength_params(level)
        out = MEDIA / scale / f"right{level}.mp4"
        print(f"  [right {level}] -> {out}")
        restyle_clip(pipe, dev, str(base), str(out), key_s, tween_s,
                     fps=fps, keyint=keyint)
    print(f"[bake] {scale} done — {len(RIGHT_LEVELS)} real Right variants")


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--scale", default="forest")
    ap.add_argument("--fps", type=float, default=12.0)
    ap.add_argument("--keyint", type=int, default=24)
    a = ap.parse_args()
    bake_scale(a.scale, fps=a.fps, keyint=a.keyint)


if __name__ == "__main__":
    main()
