"""Quiet Render — a deterministic generative image renderer.

Every image is derived entirely from its seed: the same seed always produces the
same picture, so the parameters printed next to a piece on the site are enough
to reproduce it byte for byte.

Usage:
    python gen/render.py --seed 8412 --size 4k --out out/
    python gen/render.py --count 20 --out out/
"""

from __future__ import annotations

import argparse
import json
import time
from dataclasses import dataclass, asdict
from pathlib import Path

import numpy as np
from PIL import Image

SIZES = {
    "preview": (1600, 900),
    "hd": (1920, 1080),
    "4k": (3840, 2160),
    "5k": (5120, 2880),
}

# Palettes are ordered dark -> light. Values are plain 8-bit RGB stops; the
# renderer interpolates between them in linear space.
PALETTES: dict[str, list[tuple[int, int, int]]] = {
    "ash": [(9, 12, 20), (32, 44, 66), (98, 116, 140), (176, 186, 194), (226, 222, 210)],
    "kelp": [(6, 16, 14), (18, 48, 44), (52, 102, 92), (140, 172, 154), (222, 226, 210)],
    "ember": [(14, 11, 10), (58, 30, 22), (132, 66, 40), (198, 132, 84), (238, 220, 196)],
    "quartz": [(16, 14, 26), (48, 38, 68), (108, 86, 122), (176, 148, 168), (232, 220, 224)],
    "graphite": [(8, 9, 11), (38, 41, 46), (92, 97, 104), (158, 163, 170), (228, 230, 232)],
    "tide": [(7, 14, 24), (20, 52, 72), (44, 106, 122), (128, 176, 178), (224, 232, 226)],
}

ALGORITHMS = ("flow", "warp", "strata", "cells")


@dataclass
class Params:
    """Everything needed to reproduce a render. Published next to the piece."""

    seed: int
    algorithm: str
    palette: str
    octaves: int
    frequency: int
    warp: float
    contrast: float
    grain: float
    rotation: float

    @classmethod
    def from_seed(cls, seed: int) -> "Params":
        rng = np.random.default_rng(seed)
        return cls(
            seed=seed,
            algorithm=str(rng.choice(ALGORITHMS)),
            palette=str(rng.choice(list(PALETTES))),
            octaves=int(rng.integers(3, 7)),
            frequency=int(rng.integers(2, 7)),
            warp=round(float(rng.uniform(0.0, 0.9)), 3),
            contrast=round(float(rng.uniform(0.85, 1.6)), 3),
            grain=round(float(rng.uniform(0.004, 0.016)), 4),
            rotation=round(float(rng.uniform(0.0, np.pi)), 3),
        )


def _value_noise(rng: np.random.Generator, height: int, width: int, freq: int) -> np.ndarray:
    """Smoothed value noise on a `freq`-cell lattice, bilinearly upsampled."""
    lattice = rng.random((freq + 1, freq + 1), dtype=np.float32)

    ys = np.linspace(0.0, freq, height, endpoint=False, dtype=np.float32)
    xs = np.linspace(0.0, freq, width, endpoint=False, dtype=np.float32)
    y0 = np.floor(ys).astype(np.int32)
    x0 = np.floor(xs).astype(np.int32)
    fy = ys - y0
    fx = xs - x0
    # Smoothstep the interpolant so cell borders do not show as creases.
    fy = fy * fy * (3.0 - 2.0 * fy)
    fx = fx * fx * (3.0 - 2.0 * fx)

    top = lattice[y0][:, x0] + (lattice[y0][:, x0 + 1] - lattice[y0][:, x0]) * fx[None, :]
    bot = lattice[y0 + 1][:, x0] + (lattice[y0 + 1][:, x0 + 1] - lattice[y0 + 1][:, x0]) * fx[None, :]
    return top + (bot - top) * fy[:, None]


def _fbm(rng: np.random.Generator, height: int, width: int, freq: int, octaves: int) -> np.ndarray:
    """Fractal Brownian motion: octaves of value noise at doubling frequency."""
    total = np.zeros((height, width), dtype=np.float32)
    amplitude = 1.0
    norm = 0.0
    for octave in range(octaves):
        total += amplitude * _value_noise(rng, height, width, freq * 2**octave)
        norm += amplitude
        amplitude *= 0.5
    return total / norm


def _normalize(field: np.ndarray) -> np.ndarray:
    lo, hi = float(field.min()), float(field.max())
    if hi - lo < 1e-9:
        return np.zeros_like(field)
    return (field - lo) / (hi - lo)


def _sample(field: np.ndarray, yy: np.ndarray, xx: np.ndarray) -> np.ndarray:
    """Bilinear gather. Nearest-neighbour here leaves visible column smears
    wherever the displacement is flat, so the interpolation is not optional."""
    height, width = field.shape
    yy = np.clip(yy, 0.0, height - 1.001)
    xx = np.clip(xx, 0.0, width - 1.001)

    y0 = yy.astype(np.int32)
    x0 = xx.astype(np.int32)
    fy = (yy - y0).astype(np.float32)
    fx = (xx - x0).astype(np.float32)

    top = field[y0, x0] + (field[y0, x0 + 1] - field[y0, x0]) * fx
    bot = field[y0 + 1, x0] + (field[y0 + 1, x0 + 1] - field[y0 + 1, x0]) * fx
    return top + (bot - top) * fy


def _render_warp(rng, height, width, p: Params) -> np.ndarray:
    """Domain warping: a noise field sampled through an offset of itself."""
    base = _fbm(rng, height, width, p.frequency, p.octaves)
    if p.warp > 0.0:
        dx = _fbm(rng, height, width, p.frequency, max(2, p.octaves - 1))
        dy = _fbm(rng, height, width, p.frequency, max(2, p.octaves - 1))
        shift = p.warp * min(height, width) * 0.12
        yy = np.arange(height, dtype=np.float32)[:, None] + (dy - 0.5) * shift
        xx = np.arange(width, dtype=np.float32)[None, :] + (dx - 0.5) * shift
        base = _sample(base, np.broadcast_to(yy, base.shape), np.broadcast_to(xx, base.shape))
    return _normalize(base)


def _render_strata(rng, height, width, p: Params) -> np.ndarray:
    """Banded sediment: a warped field folded into layers."""
    field = _render_warp(rng, height, width, p)
    bands = 4 + int(p.frequency)
    folded = np.abs(np.sin(field * bands * np.pi * 0.5))
    # Weighted towards the underlying field: pure folding blows the highlights
    # out into flat paper-white.
    return _normalize(folded * 0.45 + field * 0.55)


def _render_cells(rng, height, width, p: Params) -> np.ndarray:
    """Worley-like cells: distance to the nearest of a scattered point set."""
    count = 24 + p.frequency * 12
    points = rng.random((count, 2), dtype=np.float32)
    ys = np.linspace(0.0, 1.0, height, dtype=np.float32)[:, None]
    xs = np.linspace(0.0, 1.0, width, dtype=np.float32)[None, :]
    aspect = width / height

    nearest = np.full((height, width), np.inf, dtype=np.float32)
    second = np.full((height, width), np.inf, dtype=np.float32)
    for py, px in points:
        d = np.sqrt((ys - py) ** 2 + ((xs - px) * aspect) ** 2 / aspect)
        second = np.minimum(second, np.maximum(nearest, d))
        nearest = np.minimum(nearest, d)

    # F2-F1 has a long tail: the pixels deepest inside the largest cell sit
    # several times above the median, so normalising by the maximum pushed the
    # bulk of the frame into the bottom fifth of the scale and produced a wide
    # flat black patch. Clip the tail at the 98th percentile, then square-root
    # what is left to lift the mid-tones back off the floor.
    gap = second - nearest
    edges = np.sqrt(np.clip(gap / np.percentile(gap, 98), 0.0, 1.0))
    warped = _render_warp(rng, height, width, p)
    return _normalize(edges * 0.65 + warped * 0.35)


def _render_flow(rng, height, width, p: Params) -> np.ndarray:
    """Streamlines: particles advected through a noise-driven vector field."""
    angle_field = _fbm(rng, height, width, p.frequency, p.octaves) * np.pi * 2.0 + p.rotation

    particles = 60_000
    steps = 220
    step_len = min(height, width) / 900.0

    py = rng.random(particles, dtype=np.float32) * (height - 1)
    px = rng.random(particles, dtype=np.float32) * (width - 1)
    density = np.zeros(height * width, dtype=np.float32)

    for _ in range(steps):
        iy = py.astype(np.int32)
        ix = px.astype(np.int32)
        angle = angle_field[iy, ix]
        py += np.sin(angle) * step_len
        px += np.cos(angle) * step_len

        inside = (py >= 0) & (py < height - 1) & (px >= 0) & (px < width - 1)
        # Respawn escaped particles so the frame stays evenly covered.
        if not inside.all():
            escaped = ~inside
            py[escaped] = rng.random(int(escaped.sum()), dtype=np.float32) * (height - 1)
            px[escaped] = rng.random(int(escaped.sum()), dtype=np.float32) * (width - 1)

        flat = py.astype(np.int32) * width + px.astype(np.int32)
        density += np.bincount(flat, minlength=height * width).astype(np.float32)

    density = density.reshape(height, width)
    density = np.log1p(density * 3.0)
    backdrop = _fbm(rng, height, width, max(2, p.frequency // 2), 3)
    return _normalize(_normalize(density) * 0.78 + backdrop * 0.22)


RENDERERS = {
    "flow": _render_flow,
    "warp": _render_warp,
    "strata": _render_strata,
    "cells": _render_cells,
}


def _colorize(field: np.ndarray, palette: str, contrast: float) -> np.ndarray:
    """Map a 0..1 field through the palette, in linear light."""
    stops = np.array(PALETTES[palette], dtype=np.float32) / 255.0
    linear_stops = stops**2.2

    shaped = np.clip((field - 0.5) * contrast + 0.5, 0.0, 1.0)
    position = shaped * (len(linear_stops) - 1)
    idx = np.clip(position.astype(np.int32), 0, len(linear_stops) - 2)
    frac = (position - idx)[..., None]

    lo = linear_stops[idx]
    hi = linear_stops[idx + 1]
    linear = lo + (hi - lo) * frac
    return np.clip(linear ** (1.0 / 2.2), 0.0, 1.0)


def render(seed: int, size: str = "4k") -> tuple[Image.Image, Params]:
    width, height = SIZES[size]
    params = Params.from_seed(seed)
    # A second stream keeps pixel data independent of the parameter draw, so
    # changing the parameter schema later does not reshuffle existing pieces.
    rng = np.random.default_rng(seed ^ 0x5EED)

    field = RENDERERS[params.algorithm](rng, height, width, params)
    rgb = _colorize(field, params.palette, params.contrast)

    # Grain is part of the look, and it is what gives the files their weight:
    # a smooth gradient would compress down to almost nothing.
    if params.grain > 0.0:
        rgb = np.clip(rgb + rng.normal(0.0, params.grain, rgb.shape).astype(np.float32), 0.0, 1.0)

    return Image.fromarray((rgb * 255.0 + 0.5).astype(np.uint8), mode="RGB"), params


def write_piece(seed: int, out_dir: Path, size: str = "4k") -> dict:
    out_dir.mkdir(parents=True, exist_ok=True)
    started = time.time()
    image, params = render(seed, size)

    stem = f"qr-{seed:06d}"
    full_path = out_dir / f"{stem}-{size}.png"
    image.save(full_path, format="PNG", optimize=False)

    preview = image.resize(SIZES["preview"], Image.LANCZOS)
    preview_path = out_dir / f"{stem}-preview.webp"
    preview.save(preview_path, format="WEBP", quality=82, method=4)

    meta = {
        **asdict(params),
        "size": size,
        "width": image.width,
        "height": image.height,
        "bytes": full_path.stat().st_size,
        "render_seconds": round(time.time() - started, 2),
    }
    (out_dir / f"{stem}.json").write_text(json.dumps(meta, indent=2) + "\n")
    return meta


def main() -> None:
    parser = argparse.ArgumentParser(description="Render Quiet Render pieces.")
    parser.add_argument("--seed", type=int, help="render a single seed")
    parser.add_argument("--count", type=int, default=1, help="render N sequential seeds")
    parser.add_argument("--start", type=int, default=1000, help="first seed when using --count")
    parser.add_argument("--size", choices=sorted(SIZES), default="4k")
    parser.add_argument("--out", type=Path, default=Path("out"))
    args = parser.parse_args()

    seeds = [args.seed] if args.seed is not None else range(args.start, args.start + args.count)
    for seed in seeds:
        meta = write_piece(seed, args.out, args.size)
        print(
            f"qr-{seed:06d}  {meta['algorithm']:<6} {meta['palette']:<9} "
            f"{meta['bytes'] / 1e6:6.1f} MB  {meta['render_seconds']:5.1f}s"
        )


if __name__ == "__main__":
    main()