"""Generate Random Stroke Pattern (RSP) artwork from colour texture jobs.

Run after the colour texture command has produced its PNGs::

    blender -b -t 2 --python-exit-code 1 \
        -P blender/scripts/random_stroke_pattern.py -- \
        --job build/textures/finger/r01/_texture_views_job.json \
        --output-dir .scratch/binary-pattern-from-color/run-01

This is an image-space experiment, not a print master. It accepts either a single-view
texture job or a multiview job. Each distinct source colour is treated as one cell
within each PNG. A split cell can have a different centre in the two print views.
The medium- and dense-dash outputs test higher black coverage without changing
cell positions.
The physical-size outputs compare two stroke widths with a 0.2 mm-long,
0.1 mm-wide capsule. Clipping and merged neighbours can change the widths of
visible features.
The spaced 0.12 mm trial rejects clipped or colliding marks on each print page;
its selection can differ between the two UV views, so it is also preview-only.
The balanced trial permits overlaps but caps nearby black coverage.
The camera-scaled trial uses the source sensing-camera grid for a radial size
gradient: 0.11–0.16 mm width and 0.24–0.38 mm length, edge to centre. These are
nominal print-page sizes; clipped/merged ink still needs physical validation.
The camera-compensated trial scales width from each camera cell's print-page
footprint, bounded to 0.11–0.28 mm; larger far-field cells therefore receive
larger marks. The PNG and native SVG/PDF share accepted capsule geometry, but
image-space regions can be distorted. Verify dimensions before manufacturing.
"""

from __future__ import annotations

import argparse
import hashlib
import json
import math
import sys
import xml.etree.ElementTree as ET
from pathlib import Path

import bpy
import numpy as np

script_dir = str(Path(__file__).resolve().parent)
if script_dir not in sys.path:
    sys.path.insert(0, script_dir)
from random_patch_dd02c import save_png_image, write_pdf_vector, write_svg_vector  # noqa: E402 - Blender omits the script directory from sys.path

PHYSICAL_WIDTH_MM = {"width_0p1mm": 0.1, "width_0p12mm": 0.12, "width_0p2mm": 0.2}
MODES = (
    "dots",
    "dashes",
    "medium_dashes",
    "dense_dashes",
    *PHYSICAL_WIDTH_MM,
    "length_0p2_width_0p1mm",
    "spaced_width_0p12mm",
    "balanced_width_0p12mm",
    "camera_scaled_width",
    "camera_compensated_width",
)


def parse_args(argv: list[str]) -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--job", required=True, type=Path)
    parser.add_argument("--output-dir", required=True, type=Path)
    parser.add_argument("--seed", type=int, default=7)
    parser.add_argument("--min-cell-pixels", type=int, default=20)
    parser.add_argument(
        "--modes",
        nargs="+",
        choices=MODES,
        default=MODES,
        help="Generate only these preview variants",
    )
    parser.add_argument(
        "--camera-width-fraction",
        type=float,
        default=0.5,
        help="Compensated stroke width as a fraction of the local print-cell span",
    )
    parser.add_argument(
        "--formats", nargs="+", choices=("png", "svg", "pdf"), default=("png",)
    )
    parser.add_argument(
        "--artifact-prefix",
        default=None,
        help="Name single-mode outputs <prefix>_<source-id>.<format>",
    )
    return parser.parse_args(argv)


def cell_code(rgb: np.ndarray, seed: int) -> tuple[float, float, float, float]:
    """Return bounded offsets, angle and size from one exact source colour."""
    digest = hashlib.blake2b(
        bytes(int(v) for v in rgb) + seed.to_bytes(8, "little", signed=True),
        digest_size=8,
    ).digest()
    values = [int.from_bytes(digest[k : k + 2], "little") / 65535 for k in (0, 2, 4, 6)]
    return (
        0.66 * (values[0] - 0.5),
        0.66 * (values[1] - 0.5),
        math.pi * values[2],
        0.85 + 0.3 * values[3],
    )


def read_rgb(path: Path) -> tuple[np.ndarray, int, int]:
    image = bpy.data.images.load(str(path), check_existing=False)
    try:
        width, height = image.size
        pixels = np.empty(width * height * 4, dtype=np.float32)
        image.pixels.foreach_get(pixels)
        rgb = np.clip(
            np.rint(pixels.reshape(height, width, 4)[:, :, :3] * 255), 0, 255
        ).astype(np.uint8)
        return rgb, width, height
    finally:
        bpy.data.images.remove(image)


def cell_regions(rgb: np.ndarray, minimum: int):
    """Measure source-colour regions without assuming a rectangular UV cell."""
    height, width, _ = rgb.shape
    packed = (
        (rgb[:, :, 0].astype(np.uint32) << 16)
        | (rgb[:, :, 1].astype(np.uint32) << 8)
        | rgb[:, :, 2].astype(np.uint32)
    )
    colours, inverse, counts = np.unique(
        packed, return_inverse=True, return_counts=True
    )
    rows = np.repeat(np.arange(height, dtype=np.int32), width)
    cols = np.tile(np.arange(width, dtype=np.int32), height)
    n = len(colours)
    left = np.full(n, width, dtype=np.int32)
    right = np.full(n, -1, dtype=np.int32)
    bottom = np.full(n, height, dtype=np.int32)
    top = np.full(n, -1, dtype=np.int32)
    np.minimum.at(left, inverse, cols)
    np.maximum.at(right, inverse, cols)
    np.minimum.at(bottom, inverse, rows)
    np.maximum.at(top, inverse, rows)
    cx = np.bincount(inverse, weights=cols, minlength=n) / counts
    cy = np.bincount(inverse, weights=rows, minlength=n) / counts
    for index, colour in enumerate(colours):
        if colour in (0, 0xFFFFFF) or counts[index] < minimum:
            continue
        cell_width = right[index] - left[index] + 1
        cell_height = top[index] - bottom[index] + 1
        yield (
            int(colour),
            float(cx[index]),
            float(cy[index]),
            int(cell_width),
            int(cell_height),
        )


def camera_cell_lookup(
    nx: int, ny: int, colour_radius: float, seed: int
) -> dict[int, tuple[int, int] | None]:
    """Recover the original camera-grid cell from its quantized palette colour."""
    script_dir = str(Path(__file__).resolve().parent)
    if script_dir not in sys.path:
        sys.path.insert(0, script_dir)
    from random_patch_dd02c import precompute_cell_colors

    rng_state = np.random.get_state()
    try:
        np.random.seed(seed)
        palette = precompute_cell_colors(nx, ny, colour_radius)
    finally:
        np.random.set_state(rng_state)
    rgb = np.clip(np.rint(palette * 255), 0, 255).astype(np.uint8)
    packed = (
        (rgb[:, :, 0].astype(np.uint32) << 16)
        | (rgb[:, :, 1].astype(np.uint32) << 8)
        | rgb[:, :, 2].astype(np.uint32)
    )
    lookup: dict[int, tuple[int, int] | None] = {}
    for j in range(ny):
        for i in range(nx):
            colour = int(packed[j, i])
            lookup[colour] = None if colour in lookup else (i, j)
    return lookup


def camera_scaled_size(cell: tuple[int, int], nx: int, ny: int) -> tuple[float, float]:
    """Return nominal print width and length in mm from camera-grid position."""
    i, j = cell
    u = (i + 0.5) / nx - 0.5
    v = (j + 0.5) / ny - 0.5
    radius = min(1.0, math.hypot(u, v) / math.sqrt(0.5))
    centre_weight = 1.0 - radius**2
    return (0.11 + 0.05 * centre_weight, 0.24 + 0.14 * centre_weight)


def camera_compensated_size(
    print_span_px: int, px_per_mm: float, width_fraction: float
) -> tuple[float, float]:
    """Return nominal print dimensions from one projected camera-cell footprint."""
    width_mm = max(0.11, min(0.28, print_span_px / px_per_mm * width_fraction))
    return width_mm, 2.4 * width_mm


def capsule_polygon(
    centre_x: float, centre_y: float, radius: float, half_length: float, angle: float
) -> list[tuple[float, float]]:
    """Approximate the raster capsule boundary for native vector output."""
    points = []
    cosine, sine = math.cos(angle), math.sin(angle)
    for side, start in ((1, -math.pi / 2), (-1, math.pi / 2)):
        for index in range(25):
            theta = start + index * math.pi / 24
            local_x = side * half_length + radius * math.cos(theta)
            local_y = radius * math.sin(theta)
            points.append(
                (
                    centre_x + local_x * cosine - local_y * sine,
                    centre_y + local_x * sine + local_y * cosine,
                )
            )
    return points


def draw_mark(
    canvas: np.ndarray,
    region,
    seed: int,
    mode: str,
    px_per_mm: float,
    patterned: np.ndarray | None = None,
    camera_cell: tuple[int, int] | None = None,
    grid_shape: tuple[int, int] | None = None,
    camera_width_fraction: float = 0.5,
    vector_shapes: list | None = None,
) -> float:
    colour, cell_x, cell_y, width, height = region
    rgb = ((colour >> 16) & 255, (colour >> 8) & 255, colour & 255)
    dx, dy, angle, scale = cell_code(np.asarray(rgb), seed)
    span = min(width, height)
    centre_x = cell_x + dx * width
    centre_y = cell_y + dy * height
    if mode == "camera_compensated_width":
        if camera_cell is None:
            raise ValueError(
                "camera-compensated marks require an unambiguous camera cell"
            )
        width_mm, length_mm = camera_compensated_size(
            span, px_per_mm, camera_width_fraction
        )
        radius = width_mm * px_per_mm / 2
        half_length = (length_mm - width_mm) * px_per_mm / 2
    elif mode == "camera_scaled_width":
        if camera_cell is None or grid_shape is None:
            raise ValueError("camera-scaled marks require an unambiguous camera cell")
        width_mm, length_mm = camera_scaled_size(camera_cell, *grid_shape)
        radius = width_mm * px_per_mm / 2
        half_length = (length_mm - width_mm) * px_per_mm / 2
    elif mode == "spaced_width_0p12mm":
        radius = 0.12 * px_per_mm / 2
        half_length = (0.30 - 0.12) * px_per_mm / 2
    elif mode == "balanced_width_0p12mm":
        radius = 0.12 * px_per_mm / 2
        half_length = 2 * radius
    elif mode == "length_0p2_width_0p1mm":
        radius = 0.1 * px_per_mm / 2
        half_length = (0.2 - 0.1) * px_per_mm / 2
    elif mode in PHYSICAL_WIDTH_MM:
        radius = PHYSICAL_WIDTH_MM[mode] * px_per_mm / 2
        half_length = 2 * radius
    else:
        radius_fraction = {
            "dots": 0.14,
            "dashes": 0.14,
            "medium_dashes": 0.16,
            "dense_dashes": 0.18,
        }[mode]
        radius = max(1.0, radius_fraction * span * scale)
        half_length_fraction = {
            "dots": 0.0,
            "dashes": 0.28,
            "medium_dashes": 0.32,
            "dense_dashes": 0.36,
        }[mode]
        half_length = half_length_fraction * span * scale
    gap_px = 0.02 * px_per_mm if mode == "spaced_width_0p12mm" else 0.0
    extent = radius + half_length + gap_px
    if patterned is not None and (
        centre_x - extent < 0
        or centre_y - extent < 0
        or centre_x + extent >= canvas.shape[1]
        or centre_y + extent >= canvas.shape[0]
    ):
        return 0.0
    x0 = max(0, math.floor(centre_x - extent))
    x1 = min(canvas.shape[1], math.ceil(centre_x + extent + 1))
    y0 = max(0, math.floor(centre_y - extent))
    y1 = min(canvas.shape[0], math.ceil(centre_y + extent + 1))
    if x0 >= x1 or y0 >= y1:
        return 0.0
    yy, xx = np.ogrid[y0:y1, x0:x1]
    px = xx - centre_x
    py = yy - centre_y
    along = px * math.cos(angle) + py * math.sin(angle)
    across = -px * math.sin(angle) + py * math.cos(angle)
    distance_sq = np.maximum(np.abs(along) - half_length, 0) ** 2 + across**2
    footprint = distance_sq <= radius**2
    if patterned is not None:
        if np.any(footprint & ~patterned[y0:y1, x0:x1]):
            return 0.0
        if mode == "spaced_width_0p12mm":
            clearance = distance_sq <= (radius + gap_px) ** 2
            if np.any(canvas[y0:y1, x0:x1] & clearance):
                return 0.0
        else:
            local_radius = round(0.5 * px_per_mm)
            lx0 = max(0, round(centre_x) - local_radius)
            lx1 = min(canvas.shape[1], round(centre_x) + local_radius + 1)
            ly0 = max(0, round(centre_y) - local_radius)
            ly1 = min(canvas.shape[0], round(centre_y) + local_radius + 1)
            local_pattern = patterned[ly0:ly1, lx0:lx1]
            local_area = np.count_nonzero(local_pattern)
            if local_area == 0:
                return 0.0
            existing = np.count_nonzero(canvas[ly0:ly1, lx0:lx1] & local_pattern)
            added = np.count_nonzero(footprint & ~canvas[y0:y1, x0:x1])
            if (existing + added) / local_area > 0.45:
                return 0.0
    canvas[y0:y1, x0:x1] |= footprint
    if vector_shapes is not None:
        vector_shapes.append(
            (
                capsule_polygon(centre_x, centre_y, radius, half_length, angle),
                (0.0, 0.0, 0.0, 1.0),
            )
        )
    return 2 * radius


def write_binary_png(path: Path, marks: np.ndarray, surface: np.ndarray) -> None:
    height, width = marks.shape
    rgba = np.zeros((height, width, 4), dtype=np.float32)
    rgba[:, :, :3] = (surface & ~marks)[:, :, None]
    rgba[:, :, 3] = 1.0
    image = bpy.data.images.new(path.stem, width=width, height=height, alpha=True)
    try:
        image.pixels.foreach_set(rgba.ravel())
        image.update()
        save_png_image(image, path)
    finally:
        bpy.data.images.remove(image)


def job_sources(
    job: dict,
) -> tuple[list[tuple[str, Path, Path | None, int | None]], float]:
    """Resolve source PNGs and physical scale from a colour texture job."""
    settings = job.get("settings") or {}
    schema = job.get("schema")
    if schema == "orisys-sensor-design/texture-job-multiview-v1":
        entries = job.get("pairs") or []
        page = job.get("page") or {}
        if page.get("mode") != "auto" or any(
            entry.get("page", "auto") != "auto" for entry in entries
        ):
            raise ValueError("multiview preview requires auto-sized pages")
        px_per_mm = float(page.get("px_per_mm", 0))
        sources = [
            (
                entry["id"],
                Path(entry["artifacts"]["png"]),
                Path(entry["artifacts"]["svg"])
                if entry["artifacts"].get("svg")
                else None,
                entry.get("seed"),
            )
            for entry in entries
        ]
    elif schema == "orisys-sensor-design/texture-job-v3":
        entries = job.get("outputs") or []
        vector = settings.get("vector") or {}
        width_mm = float(vector.get("width_mm", 0))
        height_mm = float(vector.get("height_mm", 0))
        if width_mm <= 0 or height_mm <= 0:
            raise ValueError(
                "single-view preview requires positive vector page dimensions"
            )
        px_per_mm = float(settings["width"]) / width_mm
        vertical_scale = float(settings["height"]) / height_mm
        if not math.isclose(px_per_mm, vertical_scale, rel_tol=1e-6):
            raise ValueError(
                "single-view page has unequal horizontal and vertical scales"
            )
        sources = [
            (
                entry["uv_map"],
                Path(entry["artifacts"]["png"]),
                Path(entry["artifacts"]["svg"])
                if entry["artifacts"].get("svg")
                else None,
                settings.get("seed"),
            )
            for entry in entries
        ]
    else:
        raise ValueError(f"unsupported colour texture job schema: {schema!r}")
    if not sources or not math.isfinite(px_per_mm) or px_per_mm <= 0:
        raise ValueError("colour texture job needs source PNGs and positive px_per_mm")
    for source_id, source, _svg, _seed in sources:
        if not source_id or not all(
            char.isalnum() or char in "_-" for char in source_id
        ):
            raise ValueError(f"unsafe source ID: {source_id!r}")
        if not source.is_file():
            raise FileNotFoundError(f"run the colour generator first: {source}")
    return sources, px_per_mm


def source_mesh_polygons(
    svg_path: Path | None, width: int, height: int, px_per_mm: float
) -> tuple[list[list[tuple[float, float]]], float, float]:
    """Read the source artwork's owned white ground and physical page size."""
    if svg_path is None or not svg_path.is_file():
        raise FileNotFoundError(
            f"source colour SVG is required for vectors: {svg_path}"
        )
    root = ET.parse(svg_path).getroot()
    width_mm = float(root.attrib["width"].removesuffix("mm"))
    height_mm = float(root.attrib["height"].removesuffix("mm"))
    view_box = [float(value) for value in root.attrib["viewBox"].split()]
    if (
        len(view_box) != 4
        or view_box[:2] != [0.0, 0.0]
        or view_box[2:] != [float(width), float(height)]
        or not math.isfinite(width_mm)
        or not math.isfinite(height_mm)
        or width_mm <= 0
        or height_mm <= 0
        or not math.isclose(width / width_mm, px_per_mm, rel_tol=0.002)
        or not math.isclose(height / height_mm, px_per_mm, rel_tol=0.002)
    ):
        raise ValueError(f"source SVG page does not match colour PNG/job: {svg_path}")
    namespace = {"svg": "http://www.w3.org/2000/svg"}
    mesh = root.find(".//svg:g[@id='mesh']", namespace)
    if mesh is None:
        raise ValueError(f"source SVG has no owned mesh group: {svg_path}")
    polygons = []
    for element in mesh.findall("svg:polygon", namespace):
        polygon = [
            tuple(float(value) for value in point.split(","))
            for point in element.attrib["points"].split()
        ]
        if len(polygon) >= 3:
            polygons.append(polygon)
    if not polygons:
        raise ValueError(f"source SVG has no owned mesh polygons: {svg_path}")
    return polygons, width_mm, height_mm


def artifact_stem(source_id: str, mode: str, prefix: str | None) -> str:
    return f"{prefix}_{source_id}" if prefix else f"{source_id}_{mode}"


def run(
    job_path: Path,
    output_dir: Path,
    seed: int,
    minimum: int,
    modes: tuple[str, ...],
    camera_width_fraction: float,
    formats: tuple[str, ...] = ("png",),
    artifact_prefix: str | None = None,
) -> None:
    job = json.loads(job_path.read_text(encoding="utf-8"))
    sources, px_per_mm = job_sources(job)
    if artifact_prefix is not None and (
        len(modes) != 1
        or not artifact_prefix
        or not all(char.isalnum() or char in "_-" for char in artifact_prefix)
    ):
        raise ValueError("artifact prefix requires one mode and a filename-safe prefix")
    if not formats or len(formats) != len(set(formats)):
        raise ValueError("formats must be non-empty and distinct")
    targets = [
        output_dir / f"{artifact_stem(source_id, mode, artifact_prefix)}.{fmt}"
        for source_id, _, _, _ in sources
        for mode in modes
        for fmt in formats
    ]
    for target in targets:
        if target.exists():
            raise FileExistsError(f"refusing to overwrite {target}")
    output_dir.mkdir(parents=True, exist_ok=True)
    settings = job.get("settings", {})
    for source_id, source, source_svg, source_seed in sources:
        rgb, width, height = read_rgb(source)
        vector_page = (
            source_mesh_polygons(source_svg, width, height, px_per_mm)
            if "svg" in formats or "pdf" in formats
            else None
        )
        surface = np.any(rgb != 0, axis=2)
        patterned = surface & np.any(rgb != 255, axis=2)
        regions = list(cell_regions(rgb, minimum))
        if not regions or not np.any(patterned):
            raise ValueError(f"{source_id}: no coloured cell regions in source PNG")
        cell_span = float(np.median([min(region[3], region[4]) for region in regions]))
        print(
            f"{source_id}: median cell span {cell_span:.1f} px "
            f"({cell_span / px_per_mm:.3f} mm on the print page)",
            flush=True,
        )
        lookup = None
        grid_shape = None
        camera_modes = ("camera_scaled_width", "camera_compensated_width")
        if any(mode in modes for mode in camera_modes):
            grid_shape = (int(settings["nx"]), int(settings["ny"]))
            pair_seed = source_seed if source_seed is not None else settings.get("seed")
            if pair_seed is None:
                raise ValueError(
                    "camera-based preview requires a seeded colour palette"
                )
            lookup = camera_cell_lookup(
                *grid_shape, float(settings["color_r"]), int(pair_seed)
            )
            missing = [
                hex(region[0])
                for region in regions
                if region[0] != 0xFF0000 and region[0] not in lookup
            ]
            if missing:
                raise ValueError(
                    f"{source_id}: camera cell missing for "
                    f"{len(missing)} colours: {missing[:5]}"
                )
            ambiguous = [
                hex(region[0])
                for region in regions
                if region[0] != 0xFF0000 and lookup[region[0]] is None
            ]
            if ambiguous:
                print(
                    f"{source_id}: skip {len(ambiguous)} ambiguous palette regions "
                    f"{ambiguous[:5]}",
                    flush=True,
                )
        for mode in modes:
            marks = np.zeros((height, width), dtype=bool)
            vector_shapes = [] if vector_page is not None else None
            limit = (
                patterned
                if mode
                in (
                    "spaced_width_0p12mm",
                    "balanced_width_0p12mm",
                    "camera_scaled_width",
                    "camera_compensated_width",
                )
                else None
            )
            mode_regions = [
                region
                for region in regions
                if mode not in camera_modes
                or (region[0] != 0xFF0000 and lookup[region[0]] is not None)
            ]
            diameters = [
                draw_mark(
                    marks,
                    region,
                    seed,
                    mode,
                    px_per_mm,
                    limit,
                    camera_cell=lookup[region[0]]
                    if lookup is not None and mode in camera_modes
                    else None,
                    grid_shape=grid_shape,
                    camera_width_fraction=camera_width_fraction,
                    vector_shapes=vector_shapes,
                )
                for region in mode_regions
            ]
            stem = artifact_stem(source_id, mode, artifact_prefix)
            if "png" in formats:
                write_binary_png(output_dir / f"{stem}.png", marks, surface)
            if vector_page is not None:
                mesh_polygons, width_mm, height_mm = vector_page
                if "svg" in formats:
                    write_svg_vector(
                        output_dir / f"{stem}.svg",
                        width,
                        height,
                        width_mm,
                        height_mm,
                        mesh_polygons,
                        vector_shapes,
                        [],
                        (0.0, 0.0, 0.0, 1.0),
                        source_id,
                        artwork_label="stroke",
                    )
                if "pdf" in formats:
                    write_pdf_vector(
                        output_dir / f"{stem}.pdf",
                        width,
                        height,
                        width_mm,
                        height_mm,
                        mesh_polygons,
                        vector_shapes,
                        [],
                        (0.0, 0.0, 0.0, 1.0),
                        artwork_label="stroke",
                    )
            accepted = [diameter for diameter in diameters if diameter > 0]
            median_width = float(np.median(accepted)) if accepted else 0.0
            physical_width = (
                f", about {median_width / px_per_mm:.3f} mm on the print page"
                if px_per_mm > 0
                else ""
            )
            coverage = float(
                np.count_nonzero(marks & patterned) / np.count_nonzero(patterned)
            )
            width_range = (
                f", width range {min(accepted) / px_per_mm:.3f}–"
                f"{max(accepted) / px_per_mm:.3f} mm"
                if accepted
                else ""
            )
            print(
                f"{source_id}: {mode}, {len(accepted)}/{len(mode_regions)} marks, "
                f"black coverage {coverage:.1%}, median mark width "
                f"{median_width:.1f} px{physical_width}{width_range} -> {stem}",
                flush=True,
            )


def main() -> int:
    argv = sys.argv[sys.argv.index("--") + 1 :] if "--" in sys.argv else []
    args = parse_args(argv)
    if args.min_cell_pixels < 1:
        raise ValueError("--min-cell-pixels must be positive")
    if not math.isfinite(args.camera_width_fraction) or args.camera_width_fraction <= 0:
        raise ValueError("--camera-width-fraction must be finite and positive")
    run(
        args.job.resolve(),
        args.output_dir.resolve(),
        args.seed,
        args.min_cell_pixels,
        tuple(args.modes),
        args.camera_width_fraction,
        tuple(args.formats),
        args.artifact_prefix,
    )
    return 0


if __name__ == "__main__":
    raise SystemExit(main())