"""CPU-only reconstruction of the LegacyCelerisLab scene geometry."""
from __future__ import annotations

from collections.abc import Sequence
from typing import Any

import numpy as np

from OID_analysis.configs import CENTER_Y, L0, NX, NY, get_scene

FIELD_AXES = ("x", "y")
COMPONENT_ORDER = ("ux", "uy")
_SUPPORTED_SCENES = ("karman_re100", "illusion_1.0L")


def _disk_points(center: tuple[float, float], radius: float):
    """Yield precisely the integer cells visited by FlowField.add_cylinder."""
    x_c, y_c = center
    for x in range(int(x_c - radius) - 1, int(x_c + radius) + 1):
        for y in range(int(y_c - radius) - 1, int(y_c + radius) + 1):
            if (x - x_c) ** 2 + (y - y_c) ** 2 < radius**2:
                yield x, y


def _declaration(kind: str, name: str, center: tuple[float, float], radius: float) -> dict[str, Any]:
    return {
        "kind": kind,
        "name": name,
        "center": [float(center[0]), float(center[1]), 0.0],
        "radius": float(radius),
    }


def build_legacy_geometry(scene: str) -> tuple[np.ndarray, dict[str, Any]]:
    """Return the Legacy fluid mask and deterministic geometry declaration.

    The mask follows collected velocity-field storage, ``(NX, NY)`` / ``(x, y)``.
    True means the Legacy ``FLUID`` bit remains set.  Sensor disks are declarations
    only because Legacy ``add_sensor`` ORs the SENSOR bit onto existing fluid flags.
    No CUDA or CFD module is imported or initialized here.
    """
    if scene not in _SUPPORTED_SCENES:
        raise ValueError(f"scene must be one of {_SUPPORTED_SCENES}")
    cfg = get_scene(scene)
    shape = (int(NX), int(NY))
    fluid = np.ones(shape, dtype=bool)
    fluid[0, :] = fluid[-1, :] = False
    fluid[:, 0] = fluid[:, -1] = False

    objects: list[dict[str, Any]] = []
    cylinders: list[tuple[str, tuple[float, float], float]] = []
    if cfg.get("has_disturbance", False):
        cylinders.append((
            "disturbance",
            (float(cfg.get("disturbance_x", 10.0)) * L0, float(CENTER_Y)),
            float(cfg.get("disturbance_radius", 1.0)) * L0,
        ))

    for name, y_offset in (("sensor_upper", 2.0), ("sensor_center", 0.0), ("sensor_lower", -2.0)):
        objects.append(_declaration(
            "sensor", name,
            (float(cfg["sensor_x"]) * L0, float(CENTER_Y + y_offset * L0)),
            L0 / 4.0,
        ))

    cylinders.extend((
        ("front", (float(cfg["pinball_front_x"]) * L0, float(CENTER_Y)), L0 / 2.0),
        ("upper", (float(cfg["pinball_rear_x"]) * L0, float(CENTER_Y + 0.75 * L0)), L0 / 2.0),
        ("lower", (float(cfg["pinball_rear_x"]) * L0, float(CENTER_Y - 0.75 * L0)), L0 / 2.0),
    ))

    cylinder_declarations = [_declaration("cylinder", name, center, radius)
                             for name, center, radius in cylinders]
    # Legacy object order is disturbance, sensors, then pinball cylinders.
    if cylinder_declarations and cylinder_declarations[0]["name"] == "disturbance":
        objects.insert(0, cylinder_declarations.pop(0))
    objects.extend(cylinder_declarations)

    for obj in objects:
        if obj["kind"] != "cylinder":
            continue
        x_c, y_c, _ = obj["center"]
        radius = obj["radius"]
        if (x_c - radius <= 0 or x_c + radius >= shape[0] - 1 or
                y_c - radius <= 0 or y_c + radius >= shape[1] - 1):
            raise ValueError(f"configured cylinder {obj['name']!r} is out of bounds")
        for x, y in _disk_points((x_c, y_c), radius):
            fluid[x, y] = False

    metadata = {
        "schema_id": "legacy-celerislab-geometry/v1",
        "scene": scene,
        "field_shape": [shape[0], shape[1]],
        "field_axes": list(FIELD_AXES),
        "fluid_mask_semantics": "true iff Legacy FLUID bit is set",
        "boundary_condition": "all four outer cell rows are SOLID",
        "cylinder_rasterization": "driver integer loop; squared distance < radius^2",
        "objects": objects,
    }
    return fluid, metadata


def crop_stride_metric(
    fluid_mask: np.ndarray,
    *,
    crop: tuple[slice, slice] | None = None,
    stride: int = 1,
    field_axes: Sequence[str] = FIELD_AXES,
    lattice_cell_area: float = 1.0,
) -> tuple[np.ndarray, np.ndarray, dict[str, Any]]:
    """Crop/stride an ``(x,y)`` mask and derive its fluid-cell quadrature."""
    mask = np.asarray(fluid_mask)
    if mask.ndim != 2 or mask.dtype != np.bool_:
        raise ValueError("fluid_mask must be a two-dimensional boolean array")
    if tuple(field_axes) != FIELD_AXES:
        raise ValueError(f"field_axes must be {FIELD_AXES}, matching collected fields")
    if type(stride) is not int or stride < 1:
        raise ValueError("stride must be a positive integer")
    if not np.isscalar(lattice_cell_area) or not np.isfinite(lattice_cell_area) or lattice_cell_area <= 0:
        raise ValueError("lattice_cell_area must be finite and positive")
    region = crop if crop is not None else (slice(None), slice(None))
    if not isinstance(region, tuple) or len(region) != 2 or any(not isinstance(part, slice) for part in region):
        raise ValueError("crop must contain exactly (x_slice, y_slice)")

    normalized: list[slice] = []
    bounds: list[list[int]] = []
    for axis, (part, size) in enumerate(zip(region, mask.shape)):
        if part.step not in (None, 1):
            raise ValueError("crop slices cannot have a step; use stride")
        start = 0 if part.start is None else part.start
        stop = size if part.stop is None else part.stop
        if type(start) is not int or type(stop) is not int or start < 0 or stop > size or start >= stop:
            raise ValueError(f"invalid crop on {FIELD_AXES[axis]} axis for size {size}")
        normalized.append(slice(start, stop))
        bounds.append([start, stop])

    sampled = mask[normalized[0], normalized[1]][::stride, ::stride].copy()
    represented_area = float(lattice_cell_area) * stride**2
    metric = sampled.astype(np.float64) * represented_area
    metadata = {
        "schema_id": "legacy-cell-metric/v1",
        "field_axes": list(FIELD_AXES),
        "crop": {"x": bounds[0], "y": bounds[1]},
        "stride": stride,
        "lattice_cell_area": float(lattice_cell_area),
        "sampled_cell_area": represented_area,
        "quadrature": "uniform lattice-cell area, zero on actual solid mask cells",
        "component_order": list(COMPONENT_ORDER),
    }
    return sampled, metric, metadata
