"""The regular grid that covers the bounding box of a plate, restricted to the
plate's outline by a mask. Arrays over the grid have shape ``(ny, nx)``, indexed
``[j, i]`` with ``i`` along x. Node coordinates are relative to the center of the
bounding box.

The edge is a staircase of whole cells: a cell belongs to the plate when its
center lies inside the outline, and a node when it is a corner of such a cell.
Deciding membership by cell centers keeps the covered area unbiased. A hole is
handled by the same rule.
"""

from __future__ import annotations

from dataclasses import dataclass
from functools import cached_property
from typing import Callable

import numpy as np

InsideFn = Callable[[np.ndarray, np.ndarray, np.ndarray, np.ndarray], np.ndarray]
"""``inside(x, y, r, theta) -> bool array``, evaluated on whole grid matrices."""


@dataclass(frozen=True)
class Grid:
    nx: int
    ny: int
    hx: float
    hy: float
    mask: np.ndarray
    """``(ny, nx)`` bool: True for the nodes that belong to the plate."""

    @property
    def n(self) -> int:
        """Number of grid nodes, inside or not."""
        return self.nx * self.ny

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

    @property
    def count(self) -> int:
        """Number of nodes in the plate."""
        return int(self.mask.sum())

    @cached_property
    def x(self) -> np.ndarray:
        """Node x coordinates ``(nx,)``, relative to the center of the bounding box."""
        return (np.arange(self.nx) - (self.nx - 1) / 2) * self.hx

    @cached_property
    def y(self) -> np.ndarray:
        return (np.arange(self.ny) - (self.ny - 1) / 2) * self.hy

    @cached_property
    def X(self) -> np.ndarray:
        """Node x coordinates on the grid, ``(ny, nx)``."""
        return np.broadcast_to(self.x[None, :], self.shape).copy()

    @cached_property
    def Y(self) -> np.ndarray:
        return np.broadcast_to(self.y[:, None], self.shape).copy()

    @cached_property
    def cell_mask(self) -> np.ndarray:
        """``(ny - 1, nx - 1)`` bool: cells whose four corners belong to the plate."""
        m = self.mask
        return m[:-1, :-1] & m[:-1, 1:] & m[1:, 1:] & m[1:, :-1]

    @cached_property
    def node_weight(self) -> np.ndarray:
        """Quadrature and lumped-mass weight of each node: the fraction of its
        four cells that belong to the plate (1 inside, 1/2 on an edge, 1/4 at a
        corner of a rectangle, 0 outside)."""
        return corners_to_nodes(self.cell_mask.astype(float)) / 4

    @cached_property
    def node_area(self) -> np.ndarray:
        """Area attributed to each node, hx hy times its weight [m^2]."""
        return self.hx * self.hy * self.node_weight

    @cached_property
    def area(self) -> float:
        """Area of the staircase plate [m^2]."""
        return float(self.node_area.sum())

    def cell_centers(self) -> tuple[np.ndarray, np.ndarray]:
        """Coordinates of the cell centers, two ``(ny - 1, nx - 1)`` arrays."""
        return np.meshgrid(self.x[:-1] + self.hx / 2, self.y[:-1] + self.hy / 2)

    def neighbor_masks(self) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
        """Whether the left, right, lower (j - 1) and upper (j + 1) neighbor of each node is in the plate."""
        m = self.mask
        left = np.zeros_like(m)
        left[:, 1:] = m[:, :-1]
        right = np.zeros_like(m)
        right[:, :-1] = m[:, 1:]
        down = np.zeros_like(m)
        down[1:, :] = m[:-1, :]
        up = np.zeros_like(m)
        up[:-1, :] = m[1:, :]
        return left & m, right & m, down & m, up & m


def corners_to_nodes(c: np.ndarray) -> np.ndarray:
    """Sums a per-cell array ``(ny - 1, nx - 1)`` onto the four corners of each cell, giving ``(ny, nx)``."""
    out = np.zeros((c.shape[0] + 1, c.shape[1] + 1), dtype=c.dtype)
    out[:-1, :-1] += c
    out[:-1, 1:] += c
    out[1:, 1:] += c
    out[1:, :-1] += c
    return out


def make_grid(Lx: float, Ly: float, spacing: float, inside: InsideFn | None = None) -> Grid:
    """The grid covering the box ``Lx`` by ``Ly`` with about the requested
    spacing (at least 6 nodes per side), restricted to ``inside`` when given."""
    nx = max(6, round(Lx / spacing) + 1)
    ny = max(6, round(Ly / spacing) + 1)
    g = Grid(nx, ny, Lx / (nx - 1), Ly / (ny - 1), np.ones((ny, nx), dtype=bool))
    if inside is None:
        return g
    XC, YC = g.cell_centers()
    cells = np.asarray(inside(XC, YC, np.hypot(XC, YC), np.arctan2(YC, XC)), dtype=bool)
    if cells.shape != (ny - 1, nx - 1):
        raise ValueError(f"inside returned shape {cells.shape} for a {ny - 1} x {nx - 1} array of cells")
    return with_cell_mask(g, cells)


def with_cell_mask(g: Grid, cells: np.ndarray) -> Grid:
    """Restricts a grid to the cells flagged in ``cells`` ``(ny - 1, nx - 1)``:
    a node belongs to the plate when it is a corner of an included cell. Any
    node left with no cell of its own would carry no mass and is removed."""
    mask = corners_to_nodes(cells.astype(np.int8)) > 0
    while True:
        gm = Grid(g.nx, g.ny, g.hx, g.hy, mask)
        bad = mask & (gm.node_weight == 0)
        if not bad.any():
            return gm
        mask = mask & ~bad