"""CPU preparation for one fixed D2Q9 2:1 patch and optional interior circles.

Cell indices use y*nx+x; population slots use q*(nx*ny)+cell, matching
core/layout.cuh. Bounds are half-open (x0, y0, x1, y1) in coarse cells.
The four fine halo cells implement the pinned CPUREF volumetric geometry.
Base topology is independent of boundaries; prepare_fine_boundary checks them.
"""
from dataclasses import dataclass, replace
from numbers import Integral

import numpy as np

from ..config import LBMConfig
from .descriptors import D2Q9_EX, D2Q9_EY, FLUID, SOLID

ACTIVE, INNER_GHOST, OUTER_GHOST = 0, 1, 2
HALO = 4


def _readonly(values, dtype):
    array = np.ascontiguousarray(values, dtype=dtype)
    array.setflags(write=False)
    return array


@dataclass(frozen=True, eq=False)
class FixedPatch2D:
    """Layer metadata and compact exchange tables, constructed by from_config.

    Shapes are (ny, nx), including the fine halo. ``fine_origin`` is its lower
    cell edge in coarse physical units, with coarse dx=dt=1 and fine dx=dt=.5.
    ``ghost_fine`` and ``ghost_coarse`` pair fine cells with coarse donors.
    ``coalesce_coarse`` contains incoming population slots; each corresponding
    row of ``coalesce_fine`` contains four fine slots in (dy, dx) row order.
    Weights are all 1/4. Closed corner links are included even when the coarse
    donor is uncovered. These tables contain no DDF data and allocate no GPU.

    This is preparation, not permission to advance these levels with OneStep:
    its channel boundaries and fused collision do not implement patch coupling.
    """
    coarse_shape: tuple[int, int]
    bounds: tuple[int, int, int, int]
    tau_coarse: float
    tau_fine: float
    fine_kind: np.ndarray
    ghost_fine: np.ndarray
    ghost_coarse: np.ndarray
    coalesce_coarse: np.ndarray
    coalesce_fine: np.ndarray

    @classmethod
    def from_config(cls, cfg: LBMConfig, bounds) -> "FixedPatch2D":
        """Prepare maps without changing cfg or allocating coarse-sized tables.

        The initial backend contract is D2Q9/SRT/FP32/unshifted/double_buffer,
        without LES. Two coarse cells of margin keep all halo donors in bounds;
        this does NOT certify that those donors are fluid or periodic.
        """
        checked = replace(cfg)
        checked.validate()
        if (not checked.is_d2q9 or checked.collision != "SRT"
                or checked.streaming != "double_buffer"
                or checked.store_precision != "FP32"
                or checked.ddf_shifting or checked.les_enabled):
            raise ValueError("FixedPatch2D requires D2Q9/SRT/FP32/unshifted/double_buffer without LES")
        bounds = tuple(bounds)
        if len(bounds) != 4 or any(isinstance(v, (bool, np.bool_))
                                   or not isinstance(v, Integral) for v in bounds):
            raise ValueError("bounds must be four integers: (x0, y0, x1, y1)")
        x0, y0, x1, y1 = bounds = tuple(map(int, bounds))
        nx, ny = checked.nx, checked.ny
        if not (2 <= x0 < x1 <= nx - 2 and 2 <= y0 < y1 <= ny - 2):
            raise ValueError("patch must be nonempty with a two-coarse-cell margin")
        fx, fy = 2 * (x1 - x0) + 2 * HALO, 2 * (y1 - y0) + 2 * HALO
        if max(nx * ny, fx * fy) > np.iinfo(np.uint32).max:
            raise ValueError("layer cell counts exceed uint32 indexing")
        tau_c = 3.0 * checked.viscosity + 0.5
        tau_f = 6.0 * checked.viscosity + 0.5
        if not checked.omega_min <= 1.0 / tau_f <= checked.omega_max:
            raise ValueError("fine relaxation rate lies outside configured omega guards")

        kind = np.full((fy, fx), OUTER_GHOST, dtype=np.uint8)
        kind[2:-2, 2:-2] = INNER_GHOST
        kind[HALO:-HALO, HALO:-HALO] = ACTIVE
        gy, gx = np.nonzero(kind != ACTIVE)
        ghost_fine = gy * fx + gx
        ghost_coarse = (y0 + (gy - HALO) // 2) * nx + x0 + (gx - HALO) // 2

        targets, sources = [], []
        # Only the outside ring can receive fine-to-coarse populations.
        # Integer doubled midpoints include all eight diagonal corner passes.
        for y in range(y0 - 1, y1 + 1):
            xs = range(x0 - 1, x1 + 1) if y in (y0 - 1, y1) else (x0 - 1, x1)
            for x in xs:
                for q in range(1, 9):
                    mx, my = 2 * x + 1 - D2Q9_EX[q], 2 * y + 1 - D2Q9_EY[q]
                    if 2 * x0 <= mx <= 2 * x1 and 2 * y0 <= my <= 2 * y1:
                        targets.append(q * nx * ny + y * nx + x)
                        a, b = HALO + 2 * (x - x0), HALO + 2 * (y - y0)
                        sources.append([q * fx * fy + (b + dy) * fx + a + dx
                                        for dy in (0, 1) for dx in (0, 1)])
        return cls((ny, nx), bounds, tau_c, tau_f,
                   _readonly(kind, np.uint8),
                   _readonly(ghost_fine, np.uint32),
                   _readonly(ghost_coarse, np.uint32),
                   _readonly(targets, np.uint64),
                   _readonly(sources, np.uint64).reshape(-1, 4))

    @property
    def fine_shape(self) -> tuple[int, int]:
        return self.fine_kind.shape

    @property
    def fine_origin(self) -> tuple[float, float]:
        return self.bounds[0] - HALO / 2, self.bounds[1] - HALO / 2

    @property
    def table_bytes(self) -> int:
        return sum(a.nbytes for a in (self.fine_kind, self.ghost_fine,
                                     self.ghost_coarse, self.coalesce_coarse,
                                     self.coalesce_fine))

    def check_synchronization(self, coarse_steps: int, fine_steps: int, *, phase: str) -> None:
        """Check a proposed backend handoff; does not convert or collide data.

        Production storage is g=C(f), while CPUREF returns arrival_f. A future
        g-to-g coarse step must collide coarse cells once and fine cells twice;
        fine halo ghosts never collide. Coalescence must precede coarse collision.
        Converting reference f to g for comparison is an observation operation,
        not an additional production collision. Counts here are completed steps.
        """
        if phase != "postcollision_g":
            raise ValueError("patch backend handoff requires postcollision_g")
        for value in (coarse_steps, fine_steps):
            if isinstance(value, (bool, np.bool_)) or not isinstance(value, Integral) or value < 0:
                raise ValueError("step counts must be nonnegative integers")
        if fine_steps != 2 * coarse_steps:
            raise ValueError("levels are not synchronized: fine_steps must equal 2*coarse_steps")


@dataclass(frozen=True, eq=False)
class FineBoundary2D:
    """Fixed local projection plus auditable, finite transfer dependency sets.

    transfer_reads rows are (hop, direction, cell), traced from coalescence at
    hop 2 back to exploded G at hop 0. No previous coarse interval is traversed.
    Cell footprints also include the whole explosion write/donor tables.
    """
    objects: tuple
    flags: np.ndarray
    curved_host: tuple
    boundary_cells: np.ndarray
    transfer_fine_cells: np.ndarray
    transfer_coarse_cells: np.ndarray
    transfer_reads: np.ndarray


def _transfer_footprint(patch):
    ny, nx = patch.fine_shape
    cells = nx * ny
    pending = [(2, *divmod(int(slot), cells)) for slot in patch.coalesce_fine.ravel()]
    reads = set()
    while pending:
        hop, q, k = pending.pop()
        if (hop, q, k) in reads:
            continue
        reads.add((hop, q, k))
        if hop == 0:
            continue
        y, x = divmod(k, nx)
        kind = patch.fine_kind[y, x]
        if kind == OUTER_GHOST:
            raise ValueError("Transfer reads an invalid post-hop outer ghost")
        # A collided output depends on all nine arrivals; ghosts only stream.
        for direction in (range(9) if kind == ACTIVE else (q,)):
            sx, sy = x - D2Q9_EX[direction], y - D2Q9_EY[direction]
            if not (0 <= sx < nx and 0 <= sy < ny):
                raise ValueError("Transfer dependency leaves the fine allocation")
            pending.append((hop - 1, direction, sy * nx + sx))
    fine = set(map(int, patch.ghost_fine)) | {k for _, _, k in reads}
    coarse = set(map(int, patch.ghost_coarse)) | set(map(
        int, patch.coalesce_coarse % np.prod(patch.coarse_shape)))
    return fine, coarse, _readonly(sorted(reads), np.int64)


def prepare_fine_boundary(cfg, patch, objects, *, wall_aware=False):
    """Project fixed circles and reject boundary/transfer overlap before CUDA.

    Input centers/radii are in coarse physical coordinates (dx=dt=1).
    SurfaceActuation retains its lattice-velocity meaning: input coarse units
    are converted into independent fine-local objects with dx=dt=.5.
    Original objects are not retained or mutated. No runtime geometry updates.
    """
    if not objects:
        return None
    checked = replace(cfg)
    checked.validate()
    if (patch.coarse_shape != (cfg.ny, cfg.nx) or not checked.is_d2q9
            or checked.collision != "SRT" or checked.streaming != "double_buffer"
            or checked.store_precision != "FP32" or checked.ddf_shifting or checked.les_enabled
            or patch.tau_coarse != 3*checked.viscosity+.5):
        raise ValueError("Fine boundary requires the matching fixed-patch configuration")
    from ..body.geometry.circle import CircleGeometry
    from ..body.objects import SimObject
    from ..body.coupling.soa_packer import pack_cut_links_to_soa

    ny, nx = patch.fine_shape
    fine_cfg = replace(cfg, nx=nx, ny=ny, viscosity=2 * cfg.viscosity)
    flags = np.full(nx * ny, FLUID, np.uint16)
    active = patch.fine_kind.ravel() == ACTIVE
    local, links = [], []
    x0, y0, x1, y1 = patch.bounds
    dx = dt = .5
    for bid, obj in enumerate(objects):
        if (type(obj.geometry) is not CircleGeometry or obj.is_sensor or obj.is_force_region
                or any(getattr(obj.state, name) != 0 for name in ("vx", "vy", "omega"))
                or any(vars(obj.control).values())):
            raise ValueError("Fine boundaries require stationary circles without other roles")
        cx, cy = obj.geometry.center
        radius = obj.geometry._radius
        if (not np.isfinite([cx, cy, radius]).all() or radius <= 0
                or cx-radius < x0 or cx+radius > x1 or cy-radius < y0 or cy+radius > y1):
            raise ValueError("Entire fine circle must lie inside the physical patch")
        center = tuple((v-o)/dx-.5 for v, o in zip((cx, cy), patch.fine_origin))
        spec = obj.surface_actuation
        if spec is not None:
            spec = replace(spec, normal_velocity=spec.normal_velocity * dt / dx)
        item = SimObject(bid, geometry=CircleGeometry(*center, radius/dx),
                         center=center, radius=radius/dx, surface_actuation=spec)
        item.validate_surface_actuation(fine_cfg)
        mask = item.get_flag_mask(nx, ny)
        solid = (mask & SOLID) != 0
        if not solid.any() or np.any(solid & ~active) or np.any(flags[solid] & SOLID):
            raise ValueError("Fine solids must be resolved, disjoint and entirely active")
        flags[solid] = mask[solid]
        local.append(item)
    for item in local:
        cut = item.geometry.build_cut_links(nx, ny, domain_flags=flags)
        if not cut:
            raise ValueError("Fine circle has no resolved cut links")
        if item.surface_actuation is not None:
            _, aperture = item.surface_actuation.project(
                np.array([link.rx for link in cut]), np.array([link.ry for link in cut]))
            if not aperture.any():
                raise ValueError("Surface aperture hits no resolved cut links")
        for link in cut:
            link.body_id = item.obj_id
        links.extend(cut)

    boundary, writers, receivers = set(), set(), set()
    for link in links:
        r, q = link.fluid_idx, link.dir
        y, x = divmod(r, nx)
        sx, sy = x + D2Q9_EX[q], y + D2Q9_EY[q]
        s = sy * nx + sx
        if not (active[r] and flags[r] & FLUID and active[s] and flags[s] & SOLID):
            raise ValueError("Cut-link receiver/source must be active fluid/solid")
        writer = (s, q+1 if q % 2 else q-1)
        if writer in writers:
            raise ValueError("Duplicate fine boundary writer")
        writers.add(writer)
        receivers.add((r, q))
        boundary.update((r, s))
        if link.q < .5:
            dx_, dy_ = x - D2Q9_EX[q], y - D2Q9_EY[q]
            if not (0 <= dx_ < nx and 0 <= dy_ < ny and active[dy_*nx+dx_]):
                raise ValueError("q<.5 donor crosses the fine active interface; fallback cannot hide it")
            boundary.add(dy_*nx+dx_)
    # Reverse coverage: every fluid pull across a solid neighbor needs a writer.
    required = set()
    for r in np.flatnonzero(active & ((flags & FLUID) != 0)):
        for q in range(1, 9):
            if flags[r + D2Q9_EX[q] + nx*D2Q9_EY[q]] & SOLID:
                required.add((int(r), q))
    if required != receivers:
        raise ValueError("Fine solid pull coverage does not match the cut links")
    fine, coarse, reads = _transfer_footprint(patch)
    if fine & (boundary | set(map(int, np.flatnonzero(flags & SOLID)))):
        raise ValueError("Boundary stencil overlaps the two-hop transfer footprint")
    for k in coarse:
        y, x = divmod(k, cfg.nx)
        if x0 <= x < x1 and y0 <= y < y1:
            raise ValueError("Transfer coarse donor/target is not an uncovered bulk cell")
        if wall_aware and (y <= 2 or y >= cfg.ny-3):
            raise ValueError("Transfer footprint touches the coarse wall reconstruction stencil")
    return FineBoundary2D(tuple(local), _readonly(flags, np.uint16),
                          tuple(_readonly(a, a.dtype) for a in pack_cut_links_to_soa(links)),
                          _readonly(sorted(boundary), np.uint32),
                          _readonly(sorted(fine), np.uint32),
                          _readonly(sorted(coarse), np.uint32), reads)