"""A design: the material, the bounding box, the outline, the thickness field and
the stress-free profile of a plate, as numbers or vectorized functions of position.

The functions take four ``(ny, nx)`` arrays, ``(x, y, r, theta)``: the point
relative to the center of the bounding box, its distance from the center and
its angle from the x axis, all in meters and radians, and return an array of
the same shape (or a scalar). ``Design.evaluate`` samples them on a grid.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Callable, Union

import numpy as np

from .constants import ROOM_TEMPERATURE
from .grid import InsideFn, make_grid
from .material import Material
from .plate import Plate, remove_plane

Field = Union[float, Callable[[np.ndarray, np.ndarray, np.ndarray, np.ndarray], np.ndarray]]
"""A number, or a function ``f(x, y, r, theta)`` of grid matrices."""


class DesignError(ValueError):
    """A design that cannot be evaluated."""


@dataclass
class Design:
    """A plate design. Lengths are meters.

    Example::

        inch, mm = 0.0254, 1e-3
        R = 8 * inch
        crash = Design(
            name="16 inch flat crash",
            material=BRONZE,
            width=2 * R, height=2 * R,
            inside=lambda x, y, r, theta: r <= R,
            thickness=lambda x, y, r, theta: 1.2 * mm + (0.7 * mm - 1.2 * mm) * np.clip((r - 0.04) / (R - 0.04), 0, 1),
        )
        plate = crash.evaluate(spacing=0.008)
    """

    name: str
    material: Material
    width: float
    """Bounding box [m]."""
    height: float
    thickness: Field
    """Thickness [m]: a number or a function of ``(x, y, r, theta)``; must be positive on the plate."""
    inside: InsideFn | None = None
    """Outline: ``inside(x, y, r, theta) -> bool``; ``None`` for the whole bounding box."""
    profile: Field = 0.0
    """Stress-free height [m]: 0 for a flat plate, or a function for a bow or a bell. Its mean plane is removed."""
    temperature: float = ROOM_TEMPERATURE
    """Plate temperature [K], for the thermoelastic damping."""

    def __post_init__(self) -> None:
        if not (self.width >= 1e-3 and self.height >= 1e-3):
            raise DesignError("width and height must be at least 1 mm")
        if not callable(self.thickness) and not float(self.thickness) > 0:
            raise DesignError("thickness must be positive")

    def evaluate(self, spacing: float = 0.008) -> Plate:
        """Samples the design on a grid of about the given spacing: the mask from
        the outline at the cell centers, the thickness and the profile at the
        nodes. Only the slopes of the profile enter the physics, so its mean is
        immaterial, but its tilt is not (the von Karman strains are not invariant
        under rotation), so the mass-weighted least-squares plane is removed."""
        g = make_grid(self.width, self.height, spacing, self.inside)
        if g.count == 0:
            raise DesignError("the outline contains no grid cells; enlarge it or reduce the grid spacing")
        X, Y = g.X, g.Y
        R, T = np.hypot(X, Y), np.arctan2(Y, X)
        h = sample_field(self.thickness, "thickness", X, Y, R, T)
        w0 = sample_field(self.profile, "profile", X, Y, R, T)
        m = g.mask
        bad = m & ~np.isfinite(h)
        if bad.any():
            j, i = np.argwhere(bad)[0]
            raise DesignError(f"thickness is not finite at ({X[j, i]:.4f}, {Y[j, i]:.4f})")
        bad = m & (h <= 0)
        if bad.any():
            j, i = np.argwhere(bad)[0]
            raise DesignError(f"thickness must be positive; it is {h[j, i]} at ({X[j, i]:.4f}, {Y[j, i]:.4f})")
        bad = m & ~np.isfinite(w0)
        if bad.any():
            j, i = np.argwhere(bad)[0]
            raise DesignError(f"profile is not finite at ({X[j, i]:.4f}, {Y[j, i]:.4f})")
        w0 = np.where(m, remove_plane(np.where(m, w0, 0.0), g), 0.0)
        return Plate(g, self.material, np.where(m, h, 0.0), w0, self.width, self.height, self.temperature, self.name)


def sample_field(f: Field, what: str, X: np.ndarray, Y: np.ndarray, R: np.ndarray, T: np.ndarray) -> np.ndarray:
    """Evaluates a field (a number or a function of ``(x, y, r, theta)``) on grid matrices."""
    if callable(f):
        try:
            v = np.asarray(f(X, Y, R, T), dtype=float)
        except Exception as e:  # noqa: BLE001
            raise DesignError(f"{what}: {e}") from e
        if v.ndim == 0:
            return np.full(X.shape, float(v))
        if v.shape != X.shape:
            raise DesignError(f"{what} returned shape {v.shape} for a {X.shape[0]} x {X.shape[1]} grid; use elementwise numpy operations")
        return v
    try:
        return np.full(X.shape, float(f))
    except (TypeError, ValueError) as e:
        raise DesignError(f"{what} must be a number or a function of (x, y, r, theta)") from e