"""Batch synthetic sonar dataset generator for PointNet semantic segmentation.

Produces paired Area_X_scene_XXXX.npy + .obj files under sonar_dataset/.
Each scene point is one echosounder return: for every beam×ping ray the first
hit (object or seafloor) is recorded; denser object returns come only from
sonar resolution and geometry, not from overlaying a second cloud.
Target class 1 = user-provided object; class 0 = seafloor.
"""

from __future__ import annotations

import math
import random
from pathlib import Path
from typing import Any

from scene_generator import (
    apply_transform,
    export_npy_float64,
    export_obj,
    parse_obj_points,
)

# Full dataset layout (train / val / test).
AREA_LAYOUT: list[tuple[int, int]] = [
    (1, 75),
    (2, 75),
    (3, 75),
    (4, 75),
    (5, 100),
    (6, 100),
]
TOTAL_FULL_SCENES = sum(n for _, n in AREA_LAYOUT)  # 500

VISIBILITY_TIERS = ("nearly_hidden", "partial", "visible")


# ---------------------------------------------------------------------------
# Area naming
# ---------------------------------------------------------------------------

def scene_index_to_area_name(index: int) -> tuple[int, int, str]:
    """Map 0-based global index → (area, scene_number_1based, stem).

    Scene numbers restart at 0001 within each Area.
    """
    if index < 0:
        raise ValueError("scene index must be >= 0")
    remaining = index
    for area, count in AREA_LAYOUT:
        if remaining < count:
            scene_no = remaining + 1
            stem = f"Area_{area}_scene_{scene_no:04d}"
            return area, scene_no, stem
        remaining -= count
    scene_no = AREA_LAYOUT[-1][1] + remaining + 1
    stem = f"Area_6_scene_{scene_no:04d}"
    return 6, scene_no, stem


# ---------------------------------------------------------------------------
# Target object from user .obj
# ---------------------------------------------------------------------------

def normalize_object_points(points: list[list[float]]) -> list[list[float]]:
    xs = [p[0] for p in points]
    ys = [p[1] for p in points]
    zs = [p[2] for p in points]
    cx = (min(xs) + max(xs)) * 0.5
    cy = (min(ys) + max(ys)) * 0.5
    cz = (min(zs) + max(zs)) * 0.5
    span = max(max(xs) - min(xs), max(ys) - min(ys), max(zs) - min(zs), 1e-6)
    scale = 1.0 / span
    return [[(p[0] - cx) * scale, (p[1] - cy) * scale, (p[2] - cz) * scale] for p in points]


def load_object_points_from_obj_text(text: str) -> list[list[float]]:
    """Parse OBJ vertices and normalize to unit local frame (centered, max span ≈ 1)."""
    points = parse_obj_points(text)
    if len(points) < 3:
        raise ValueError("OBJ model must contain at least 3 vertices.")
    return normalize_object_points(points)


def object_half_extent_z(points: list[list[float]]) -> float:
    if not points:
        return 0.35
    zs = [p[2] for p in points]
    return max(0.05, (max(zs) - min(zs)) * 0.5)


# ---------------------------------------------------------------------------
# Seafloor heightfield + echosounder ray casting
# ---------------------------------------------------------------------------

def _seafloor_height(
    x: float,
    y: float,
    *,
    base_z: float,
    amplitude: float,
    frequency: float,
    hills: list[tuple[float, float, float, float]],
    valleys: list[tuple[float, float, float, float]],
    bumps: list[tuple[float, float, float, float]],
) -> float:
    z = base_z
    z += amplitude * math.sin(frequency * x) * math.cos(frequency * 0.7 * y)
    z += 0.35 * amplitude * math.sin(frequency * 1.7 * y + 0.4)
    for cx, cy, height, radius in hills:
        d2 = (x - cx) ** 2 + (y - cy) ** 2
        if d2 < radius * radius * 4:
            z += height * math.exp(-d2 / max(radius * radius, 1e-6))
    for cx, cy, depth, radius in valleys:
        d2 = (x - cx) ** 2 + (y - cy) ** 2
        if d2 < radius * radius * 4:
            z -= depth * math.exp(-d2 / max(radius * radius, 1e-6))
    for cx, cy, height, radius in bumps:
        d2 = (x - cx) ** 2 + (y - cy) ** 2
        if d2 < radius * radius * 4:
            z += height * math.exp(-d2 / max(radius * radius * 0.5, 1e-6))
    return z


def _build_seafloor_meta(
    rng: random.Random,
    *,
    beam_count: int = 45,
    length_count: int | None = None,
) -> dict[str, Any]:
    """Build continuous seafloor heightfield parameters (no point cloud yet)."""
    beams = max(1, int(beam_count))
    length_points = beams if length_count is None else max(1, int(length_count))
    size_x = rng.uniform(8.0, 16.0)
    size_y = rng.uniform(8.0, 16.0)
    base_z = rng.uniform(-1.2, -0.2)
    amplitude = rng.uniform(0.05, 0.35)
    frequency = rng.uniform(0.4, 2.2)

    hills = [
        (
            rng.uniform(-size_x * 0.4, size_x * 0.4),
            rng.uniform(-size_y * 0.4, size_y * 0.4),
            rng.uniform(0.15, 0.7),
            rng.uniform(0.6, 2.2),
        )
        for _ in range(rng.randint(1, 4))
    ]
    valleys = [
        (
            rng.uniform(-size_x * 0.4, size_x * 0.4),
            rng.uniform(-size_y * 0.4, size_y * 0.4),
            rng.uniform(0.1, 0.55),
            rng.uniform(0.5, 2.0),
        )
        for _ in range(rng.randint(1, 3))
    ]
    # Relief clutter / false features as heightfield bumps (not extra points)
    bumps = [
        (
            rng.uniform(-size_x * 0.45, size_x * 0.45),
            rng.uniform(-size_y * 0.45, size_y * 0.45),
            rng.uniform(0.03, 0.35),
            rng.uniform(0.15, 0.9),
        )
        for _ in range(rng.randint(4, 14))
    ]

    return {
        "sizeX": size_x,
        "sizeY": size_y,
        "baseZ": base_z,
        "amplitude": amplitude,
        "frequency": frequency,
        "hills": hills,
        "valleys": valleys,
        "bumps": bumps,
        "beamCount": beams,
        "lengthCount": length_points,
        "gridWidthPoints": beams,
        "gridLengthPoints": length_points,
        "pingCount": length_points,
        "swathBeams": beams,
        "noise": rng.uniform(0.002, 0.02),
    }


def _height_at(x: float, y: float, meta: dict[str, Any]) -> float:
    return _seafloor_height(
        x,
        y,
        base_z=float(meta["baseZ"]),
        amplitude=float(meta["amplitude"]),
        frequency=float(meta["frequency"]),
        hills=meta["hills"],
        valleys=meta["valleys"],
        bumps=meta["bumps"],
    )


def _grid_xy(xi: int, yi: int, meta: dict[str, Any]) -> tuple[float, float]:
    beams = int(meta["beamCount"])
    length_points = int(meta["lengthCount"])
    size_x = float(meta["sizeX"])
    size_y = float(meta["sizeY"])
    half_x = size_x * 0.5
    half_y = size_y * 0.5
    x = -half_x if beams == 1 else (-half_x + size_x * xi / (beams - 1))
    y = -half_y if length_points == 1 else (-half_y + size_y * yi / (length_points - 1))
    return x, y


def _cell_size(meta: dict[str, Any]) -> tuple[float, float]:
    beams = int(meta["beamCount"])
    length_points = int(meta["lengthCount"])
    size_x = float(meta["sizeX"])
    size_y = float(meta["sizeY"])
    return size_x / max(beams - 1, 1), size_y / max(length_points - 1, 1)


def _world_to_grid_index(
    x: float,
    y: float,
    meta: dict[str, Any],
) -> tuple[float, float]:
    beams = int(meta["beamCount"])
    length_points = int(meta["lengthCount"])
    size_x = float(meta["sizeX"])
    size_y = float(meta["sizeY"])
    half_x = size_x * 0.5
    half_y = size_y * 0.5
    fx = 0.0 if beams == 1 else (x + half_x) / size_x * (beams - 1)
    fy = 0.0 if length_points == 1 else (y + half_y) / size_y * (length_points - 1)
    return fx, fy


def _place_object_in_scene(
    rng: random.Random,
    meta: dict[str, Any],
    visibility: str,
    object_template: list[list[float]],
    object_scale: float = 1.0,
) -> tuple[list[list[float]], dict[str, Any]]:
    """Pose the object on the seafloor; return world-space template vertices + info."""
    base_scale = max(0.01, float(object_scale))
    if visibility == "nearly_hidden":
        burial = rng.uniform(0.35, 0.75)
        scale = base_scale * rng.uniform(0.7, 1.15)
    elif visibility == "partial":
        burial = rng.uniform(0.12, 0.4)
        scale = base_scale * rng.uniform(0.8, 1.3)
    else:
        burial = rng.uniform(-0.05, 0.15)
        scale = base_scale * rng.uniform(0.85, 1.4)

    local = [[p[0] * scale, p[1] * scale, p[2] * scale] for p in object_template]
    half_x = float(meta["sizeX"]) * 0.35
    half_y = float(meta["sizeY"]) * 0.35
    tx = rng.uniform(-half_x, half_x)
    ty = rng.uniform(-half_y, half_y)
    floor_z = _height_at(tx, ty, meta)
    half_h = object_half_extent_z(local)
    tz = floor_z + half_h * (1.0 - 2.0 * burial)
    transform = {
        "x": tx,
        "y": ty,
        "z": tz,
        "rx": rng.uniform(-0.25, 0.25),
        "ry": rng.uniform(-0.2, 0.2),
        "rz": rng.uniform(0, 2 * math.pi),
    }
    world = apply_transform(local, transform)
    info = {
        "visibility": visibility,
        "transform": transform,
        "scale": scale,
        "objectScale": base_scale,
        "burial": burial,
        "vertexCount": len(world),
        "classLabel": "object",
        "classId": 1,
    }
    return world, info


def _rasterize_object_hits(
    world_pts: list[list[float]],
    meta: dict[str, Any],
) -> dict[tuple[int, int], float]:
    """Project object vertices onto the sonar grid: first-hit Z per beam×ping cell.

    Sensor looks down (+Z is closer). A cell stores the highest object Z that
    falls into its footprint, so later casting can compare against seafloor Z.
    """
    beams = int(meta["beamCount"])
    length_points = int(meta["lengthCount"])
    cell_w, cell_h = _cell_size(meta)
    # Cover gaps between sparse mesh vertices across neighboring cells
    splat_rx = cell_w * 0.85
    splat_ry = cell_h * 0.85
    hits: dict[tuple[int, int], float] = {}

    for px, py, pz in world_pts:
        fx, fy = _world_to_grid_index(px, py, meta)
        xi0 = int(math.floor(fx))
        yi0 = int(math.floor(fy))
        for dyi in (-1, 0, 1, 2):
            for dxi in (-1, 0, 1, 2):
                xi = xi0 + dxi
                yi = yi0 + dyi
                if xi < 0 or yi < 0 or xi >= beams or yi >= length_points:
                    continue
                cx, cy = _grid_xy(xi, yi, meta)
                if abs(px - cx) > splat_rx or abs(py - cy) > splat_ry:
                    continue
                floor_z = _height_at(cx, cy, meta)
                # Buried volume does not return a sonar echo above the seafloor
                if pz < floor_z - 0.01:
                    continue
                key = (xi, yi)
                prev = hits.get(key)
                if prev is None or pz > prev:
                    hits[key] = pz
    return hits


def _cast_echosounder_returns(
    rng: random.Random,
    meta: dict[str, Any],
    object_hits: dict[tuple[int, int], float] | None,
) -> tuple[list[list[float]], int, int]:
    """One return per ray: first surface hit from above (object or seafloor).

    Returns PointNet rows [x,y,z,r,g,b,class], object_count, background_count.
    """
    beams = int(meta["beamCount"])
    length_points = int(meta["lengthCount"])
    noise = float(meta.get("noise", 0.01))
    rows: list[list[float]] = []
    object_count = 0
    background_count = 0

    for yi in range(length_points):
        for xi in range(beams):
            x, y = _grid_xy(xi, yi, meta)
            floor_z = _height_at(x, y, meta)
            obj_z = object_hits.get((xi, yi)) if object_hits else None

            # Looking down: larger Z is closer → first intersection wins.
            if obj_z is not None and obj_z > floor_z:
                z = obj_z + rng.uniform(-noise, noise)
                cls = 1.0
                object_count += 1
            else:
                z = floor_z + rng.uniform(-noise, noise)
                cls = 0.0
                background_count += 1

            rows.append([float(x), float(y), float(z), 0.0, 0.0, 0.0, cls])

    return rows, object_count, background_count


# ---------------------------------------------------------------------------
# Balance plan + single scene
# ---------------------------------------------------------------------------

def plan_scene_labels(count: int, seed: int) -> list[str]:
    """Return visibility label per scene: absent | nearly_hidden | partial | visible.

    ~50% absent; among present scenes, roughly equal nearly_hidden/partial/visible.
    """
    count = max(0, int(count))
    rng = random.Random(int(seed) ^ 0xA5A5_5A5A)
    n_with = (count + 1) // 2  # ceil → ~50% with object
    n_without = count - n_with
    labels: list[str] = ["absent"] * n_without
    for i in range(n_with):
        labels.append(VISIBILITY_TIERS[i % 3])
    rng.shuffle(labels)
    return labels


def generate_sonar_scene(
    *,
    seed: int,
    visibility: str = "absent",
    object_points: list[list[float]] | None = None,
    object_scale: float = 1.0,
    beam_count: int = 45,
    length_count: int | None = None,
) -> dict[str, Any]:
    """Build one sonar scene via beam×ping first-hit casting.

    visibility in absent|nearly_hidden|partial|visible.
    Scene size is always beam_count × length_count returns.
    """
    rng = random.Random(int(seed))
    if visibility not in ("absent",) + VISIBILITY_TIERS:
        raise ValueError(f"Unknown visibility: {visibility}")
    if visibility != "absent" and not object_points:
        raise ValueError("object_points required when visibility is not absent.")

    meta = _build_seafloor_meta(rng, beam_count=beam_count, length_count=length_count)
    expected = int(meta["beamCount"]) * int(meta["lengthCount"])

    object_info: dict[str, Any] | None = None
    object_hits: dict[tuple[int, int], float] | None = None
    if visibility != "absent":
        world, object_info = _place_object_in_scene(
            rng,
            meta,
            visibility,
            object_points,
            object_scale=object_scale,
        )
        object_hits = _rasterize_object_hits(world, meta)
        object_info["hitCellCount"] = len(object_hits)
        object_info["keptCount"] = len(object_hits)
        object_info["requestedCount"] = expected

    rows, object_point_count, background_point_count = _cast_echosounder_returns(
        rng, meta, object_hits
    )
    if len(rows) != expected:
        raise RuntimeError(f"Ray count mismatch: expected {expected}, got {len(rows)}")

    rng.shuffle(rows)
    xyz = [[r[0], r[1], r[2]] for r in rows]
    return {
        "seed": int(seed),
        "visibility": visibility,
        "hasObject": visibility != "absent",
        "object": object_info,
        "pointCount": len(rows),
        "objectPointCount": object_point_count,
        "backgroundPointCount": background_point_count,
        "rows": rows,
        "points": xyz,
        "meta": {
            "sizeX": meta["sizeX"],
            "sizeY": meta["sizeY"],
            "beamCount": meta.get("beamCount", beam_count),
            "lengthCount": meta.get(
                "lengthCount",
                length_count if length_count is not None else beam_count,
            ),
            "gridWidthPoints": meta.get("gridWidthPoints", beam_count),
            "gridLengthPoints": meta.get(
                "gridLengthPoints",
                length_count if length_count is not None else beam_count,
            ),
            "pingCount": meta.get("pingCount"),
            "swathBeams": meta.get("swathBeams"),
            "gridPointCount": expected,
            "floorFeatures": {
                "hills": len(meta["hills"]),
                "valleys": len(meta["valleys"]),
                "bumps": len(meta["bumps"]),
            },
        },
    }


# ---------------------------------------------------------------------------
# Batch write / preview load
# ---------------------------------------------------------------------------

def resolve_output_dir(output_dir: str | Path = "sonar_dataset") -> Path:
    out = Path(output_dir)
    if not out.is_absolute():
        project_root = Path(__file__).resolve().parent.parent
        out = project_root / out
    return out


def _safe_object_stem(object_name: str | None) -> str:
    stem = Path(object_name or "object").stem
    safe = "".join(ch if ch.isalnum() or ch in "_-" else "_" for ch in stem).strip("._-")
    return (safe or "object")[:80]


def make_generation_run_dir(base_dir: Path, object_name: str | None = None) -> Path:
    """Create a new run subdirectory: YYYY-MM-DD_HH-MM-SS-<object_stem>.

    Previous runs under base_dir are left untouched.
    """
    from datetime import datetime

    base_dir.mkdir(parents=True, exist_ok=True)
    stamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
    folder = f"{stamp}-{_safe_object_stem(object_name)}"
    path = base_dir / folder
    if path.exists():
        n = 2
        while True:
            candidate = base_dir / f"{folder}_{n}"
            if not candidate.exists():
                path = candidate
                break
            n += 1
    path.mkdir(parents=True, exist_ok=False)
    return path


def _downsample_points(points: list[list[float]], max_points: int) -> list[list[float]]:
    max_points = max(100, int(max_points))
    if len(points) <= max_points:
        return points
    step = max(1, len(points) // max_points)
    return points[::step][:max_points]


def load_npy_float64_rows(path: Path) -> list[list[float]]:
    """Read float64 little-endian .npy array written by export_npy_float64."""
    import re
    import struct

    data = path.read_bytes()
    if data[:6] != b"\x93NUMPY":
        raise ValueError(f"Not a NumPy .npy file: {path.name}")
    major = data[6]
    if major == 1:
        hlen = struct.unpack_from("<H", data, 8)[0]
        header = data[10 : 10 + hlen].decode("latin1")
        offset = 10 + hlen
    elif major == 2:
        hlen = struct.unpack_from("<I", data, 8)[0]
        header = data[12 : 12 + hlen].decode("latin1")
        offset = 12 + hlen
    else:
        raise ValueError(f"Unsupported .npy version: {major}")
    match = re.search(r"shape'\s*:\s*\((\d+)\s*,\s*(\d+)\)", header)
    if not match:
        match = re.search(r"shape':\s*\((\d+),\s*(\d+)\)", header)
    if not match:
        raise ValueError(f"Cannot parse .npy shape from {path.name}")
    n_rows, n_cols = int(match.group(1)), int(match.group(2))
    expected = n_rows * n_cols * 8
    body = data[offset : offset + expected]
    if len(body) < expected:
        raise ValueError(f"Truncated .npy payload in {path.name}")
    flat = struct.unpack("<" + "d" * (n_rows * n_cols), body)
    return [list(flat[i * n_cols : (i + 1) * n_cols]) for i in range(n_rows)]


def _class_counts_from_labeled(points: list[list[float]]) -> dict[str, int]:
    """Count classes from rows shaped [x, y, z, class]."""
    counts: dict[str, int] = {}
    for row in points:
        key = str(int(round(float(row[3] if len(row) > 3 else 0))))
        counts[key] = counts.get(key, 0) + 1
    return counts


def _class_counts(rows: list[list[float]]) -> dict[str, int]:
    counts: dict[str, int] = {}
    for row in rows:
        key = str(int(round(float(row[6] if len(row) > 6 else 0))))
        counts[key] = counts.get(key, 0) + 1
    return counts


def load_scene_preview(
    *,
    stem: str,
    output_dir: str | Path = "sonar_dataset",
    max_points: int = 25000,
) -> dict[str, Any]:
    """Load labeled points for a written scene (prefer .npy) for the 3D viewer.

    Each preview point is [x, y, z, class].
    """
    safe = "".join(ch if ch.isalnum() or ch in "_-" else "" for ch in (stem or ""))
    if not safe or safe != stem:
        raise ValueError("Invalid scene stem.")
    out = resolve_output_dir(output_dir)
    npy_path = out / f"{safe}.npy"
    obj_path = out / f"{safe}.obj"

    labeled: list[list[float]]
    if npy_path.is_file():
        rows = load_npy_float64_rows(npy_path)
        labeled = [[float(r[0]), float(r[1]), float(r[2]), float(r[6])] for r in rows]
    elif obj_path.is_file():
        text = obj_path.read_text(encoding="utf-8", errors="ignore")
        points = parse_obj_points(text)
        labeled = [[p[0], p[1], p[2], 0.0] for p in points]
    else:
        raise FileNotFoundError(f"Scene not found: {safe}.npy / {safe}.obj")

    full_counts = _class_counts_from_labeled(labeled)
    preview = _downsample_points(labeled, max_points)
    return {
        "stem": safe,
        "outputDir": str(out),
        "pointCount": len(labeled),
        "previewCount": len(preview),
        "points": preview,
        "classCounts": full_counts,
        "classLabels": {"0": "background", "1": "object"},
        "obj": str(obj_path) if obj_path.is_file() else None,
        "npy": str(npy_path) if npy_path.is_file() else None,
    }


def write_scene_files(
    scene: dict[str, Any],
    output_dir: Path,
    stem: str,
) -> dict[str, str]:
    output_dir.mkdir(parents=True, exist_ok=True)
    npy_path = output_dir / f"{stem}.npy"
    obj_path = output_dir / f"{stem}.obj"
    npy_path.write_bytes(export_npy_float64(scene["rows"]))
    obj_path.write_text(export_obj(scene["points"], object_name=stem), encoding="utf-8")
    return {"npy": str(npy_path), "obj": str(obj_path), "stem": stem}


def iter_generate_dataset(
    *,
    count: int = 5,
    seed: int = 42,
    output_dir: str | Path = "sonar_dataset",
    object_points: list[list[float]],
    object_name: str | None = None,
    object_scale: float = 1.0,
    object_scale_is_max: bool = False,
    beam_count: int = 45,
    length_count: int | None = None,
):
    """Yield NDJSON-friendly progress events, then a final ``done`` payload.

    Events:
      {"type":"start","total":N,"outputDir":"...","runName":"..."}
      {"type":"progress","current":k,"total":N,"entry":{...}}
      {"type":"done","result":{...}}
    """
    count = int(count)
    if count < 1:
        raise ValueError("count must be >= 1")
    if count > 5000:
        raise ValueError("count must be <= 5000")
    if not object_points or len(object_points) < 3:
        raise ValueError("A valid .obj model with at least 3 vertices is required.")
    object_scale = float(object_scale)
    if object_scale <= 0:
        raise ValueError("object_scale must be > 0")
    if object_scale > 100:
        raise ValueError("object_scale must be <= 100")
    object_scale_is_max = bool(object_scale_is_max)
    beam_count = int(beam_count)
    if beam_count < 1:
        raise ValueError("beam_count (Кол-во лучей) must be >= 1")
    if beam_count > 1024:
        raise ValueError("beam_count (Кол-во лучей) must be <= 1024")
    if length_count is None:
        length_count = beam_count
    length_count = int(length_count)
    if length_count < 1:
        raise ValueError("length_count (Длина) must be >= 1")
    if length_count > 1024:
        raise ValueError("length_count (Длина) must be <= 1024")

    base = resolve_output_dir(output_dir)
    run_dir = make_generation_run_dir(base, object_name)
    template = normalize_object_points(object_points)

    labels = plan_scene_labels(count, seed)
    written: list[dict[str, Any]] = []
    stats = {
        "total": count,
        "withObject": 0,
        "withoutObject": 0,
        "nearly_hidden": 0,
        "partial": 0,
        "visible": 0,
        "absent": 0,
    }

    preview_points: list[list[float]] | None = None
    preview_stem: str | None = None
    preview_has_object = False

    yield {
        "type": "start",
        "total": count,
        "outputDir": str(run_dir),
        "baseDir": str(base),
        "runName": run_dir.name,
    }

    for i in range(count):
        visibility = labels[i]
        scene_seed = int(seed) + i * 10007 + 17
        scene_rng = random.Random(scene_seed ^ 0xC0FFEE)
        if object_scale_is_max:
            lo, hi = 1.0, object_scale
            if hi < lo:
                lo, hi = hi, lo
            scene_scale = scene_rng.uniform(lo, hi)
        else:
            scene_scale = object_scale
        scene = generate_sonar_scene(
            seed=scene_seed,
            visibility=visibility,
            object_points=template,
            object_scale=scene_scale,
            beam_count=beam_count,
            length_count=length_count,
        )
        area, scene_no, stem = scene_index_to_area_name(i)
        paths = write_scene_files(scene, run_dir, stem)

        entry = {
            "index": i,
            "area": area,
            "scene": scene_no,
            "stem": stem,
            "visibility": visibility,
            "hasObject": scene["hasObject"],
            "pointCount": scene["pointCount"],
            "objectPointCount": scene["objectPointCount"],
            "objectScale": scene_scale,
            "files": paths,
        }
        written.append(entry)

        stats[visibility] = stats.get(visibility, 0) + 1
        if scene["hasObject"]:
            stats["withObject"] += 1
        else:
            stats["withoutObject"] += 1

        if preview_points is None or (scene["hasObject"] and not preview_has_object):
            preview_points = [[r[0], r[1], r[2], r[6]] for r in scene["rows"]]
            preview_stem = stem
            preview_has_object = bool(scene["hasObject"])

        yield {
            "type": "progress",
            "current": i + 1,
            "total": count,
            "entry": entry,
            "outputDir": str(run_dir),
        }

    preview: dict[str, Any] | None = None
    if preview_points is not None:
        pts = _downsample_points(preview_points, 25000)
        preview = {
            "stem": preview_stem,
            "points": pts,
            "pointCount": len(preview_points),
            "classCounts": _class_counts_from_labeled(preview_points),
            "classLabels": {"0": "background", "1": "object"},
        }

    result = {
        "outputDir": str(run_dir),
        "baseDir": str(base),
        "runName": run_dir.name,
        "count": count,
        "seed": int(seed),
        "beamCount": beam_count,
        "lengthCount": length_count,
        "objectName": object_name,
        "objectScale": object_scale,
        "objectScaleIsMax": object_scale_is_max,
        "objectVertexCount": len(template),
        "classLabels": {"0": "background", "1": "object"},
        "stats": stats,
        "written": written,
        "preview": preview,
    }
    yield {"type": "done", "result": result}


def generate_dataset(
    *,
    count: int = 5,
    seed: int = 42,
    output_dir: str | Path = "sonar_dataset",
    object_points: list[list[float]],
    object_name: str | None = None,
    object_scale: float = 1.0,
    object_scale_is_max: bool = False,
    beam_count: int = 45,
    length_count: int | None = None,
) -> dict[str, Any]:
    """Generate `count` unique scenes into a new timestamped run folder under output_dir."""
    result: dict[str, Any] | None = None
    for event in iter_generate_dataset(
        count=count,
        seed=seed,
        output_dir=output_dir,
        object_points=object_points,
        object_name=object_name,
        object_scale=object_scale,
        object_scale_is_max=object_scale_is_max,
        beam_count=beam_count,
        length_count=length_count,
    ):
        if event.get("type") == "done":
            result = event["result"]
    if result is None:
        raise RuntimeError("Dataset generation produced no result.")
    return result