"""Curvilinear coordinate systems.

A coordinate system is a square local chart ``chi: U -> V`` on an ambient
space ``R^N``. Evaluating the chart at physical points ``x`` returns
``u = chi(x)``, the covariant tangent basis ``g_i = d x / d u_i``, and a
boolean mask marking the points where the chart is defined.

The numerical contract is:

- ``coords[..., i] = u_i(x)``
- ``basis[..., i, j] = d x_j / d u_i``
- ``valid[...]`` is ``True`` exactly on the domain of definition of the chart

The current implementation only supports square charts ``R^N -> R^N``.
Values returned outside ``valid == True`` are shape-correct but otherwise
unspecified; consumers are expected to honor ``valid`` explicitly.
"""

from __future__ import annotations

from abc import ABC, abstractmethod
from dataclasses import dataclass
from typing import Callable

import numpy as np

__all__ = [
    "CoordinateSampleND",
    "CoordinateSystemND",
    "CoordinateSystem2D",
    "CoordinateSystem3D",
    "EuclideanCoordinates",
    "AffineCoordinates",
    "CoordinateRemap",
    "RestrictedCoordinates",
    "CylindricalCoordinates",
    "PolarCoordinates",
    "AxisymmetricCoordinates",
]


def _validate_points(points, ndim: int) -> np.ndarray:
    """Validate that sampled ambient points have trailing dimension ``ndim``."""

    points = np.asarray(points)
    if points.shape[-1] != ndim:
        raise ValueError(f"Expected points with trailing dimension {ndim}, got {points.shape}.")
    return points


def _validate_square_matrix(matrix, ndim: int, name: str) -> np.ndarray:
    """Validate a square ``ndim x ndim`` matrix parameter."""

    matrix = np.asarray(matrix, dtype=np.float64)
    if matrix.shape != (ndim, ndim):
        raise ValueError(f"`{name}` must have shape ({ndim}, {ndim}), got {matrix.shape}.")
    return matrix


def _as_offset_vector(offset, ndim: int, name: str = "offset") -> np.ndarray:
    """Validate an affine offset vector in ``R^ndim``."""

    offset = np.asarray(offset, dtype=np.float64)
    if offset.shape != (ndim,):
        raise ValueError(f"`{name}` must have shape ({ndim},), got {offset.shape}.")
    return offset


def _identity_basis(shape: tuple[int, ...], ndim: int, dtype) -> np.ndarray:
    """Broadcast the Cartesian identity basis over a batch shape."""

    return np.broadcast_to(np.eye(ndim, dtype=dtype), shape + (ndim, ndim))


def _apply_inverse_coordinate_jacobian(basis: np.ndarray, inverse_jacobian: np.ndarray) -> np.ndarray:
    """Transform ``dx/du`` into ``dx/dv`` using the inverse coordinate Jacobian."""

    return np.einsum("...ia,...ij->...ja", basis, inverse_jacobian)


def _normalize_orientation(orientation: float) -> float:
    """Reduce an orientation sign to ``+1`` or ``-1``."""

    orientation = float(orientation)
    if orientation == 0:
        raise ValueError("`orientation` must be non-zero.")
    return 1.0 if orientation > 0 else -1.0


def _rotation_matrix_2d(angle: float) -> np.ndarray:
    """Return the standard 2D rotation matrix for ``angle``."""

    c = np.cos(angle)
    s = np.sin(angle)
    return np.array([[c, -s], [s, c]], dtype=np.float64)


def _rotation_matrix_3d(axis: int | np.ndarray | tuple[float, float, float], angle: float) -> np.ndarray:
    """Return the 3D rotation matrix around a Cartesian or explicit axis."""

    if np.isscalar(axis):
        axis_index = int(axis)
        if axis_index not in (0, 1, 2):
            raise ValueError("Integer `axis` must be one of 0, 1, 2.")
        axis_vec = np.zeros(3, dtype=np.float64)
        axis_vec[axis_index] = 1.0
    else:
        axis_vec = np.asarray(axis, dtype=np.float64)
        if axis_vec.shape != (3,):
            raise ValueError(f"`axis` must have shape (3,), got {axis_vec.shape}.")
        norm = np.linalg.norm(axis_vec)
        if norm == 0:
            raise ValueError("`axis` must be non-zero.")
        axis_vec = axis_vec / norm

    x, y, z = axis_vec
    c = np.cos(angle)
    s = np.sin(angle)
    cc = 1.0 - c
    return np.array(
        [
            [c + x * x * cc, x * y * cc - z * s, x * z * cc + y * s],
            [y * x * cc + z * s, c + y * y * cc, y * z * cc - x * s],
            [z * x * cc - y * s, z * y * cc + x * s, c + z * z * cc],
        ],
        dtype=np.float64,
    )


@dataclass(slots=True)
class CoordinateSampleND:
    """Evaluated square local chart.

    Parameters
    ----------
    coords:
        Array with shape ``(..., ndim)`` storing the curvilinear coordinates
        ``u = chi(x)``.
    basis:
        Array with shape ``(..., ndim, ndim)`` storing the covariant tangent
        basis vectors. By convention ``basis[..., i, :] = d x / d u_i``.
    valid:
        Boolean array with shape ``(...)`` marking where the chart is defined.

    Notes
    -----
    This object is purely structural. It does not attempt to sanitize values
    outside ``valid == True``; such values remain unspecified by contract.
    """

    coords: np.ndarray
    basis: np.ndarray
    valid: np.ndarray

    def __post_init__(self):
        self.coords = np.asarray(self.coords)
        self.basis = np.asarray(self.basis)
        self.valid = np.asarray(self.valid, dtype=bool)

        if self.coords.ndim < 1:
            raise ValueError("`coords` must have at least one dimension.")
        ndim = self.coords.shape[-1]
        if ndim <= 0:
            raise ValueError("The trailing coordinate dimension must be positive.")
        expected_basis_shape = self.coords.shape[:-1] + (ndim, ndim)
        if self.basis.shape != expected_basis_shape:
            raise ValueError(
                f"`basis` must have shape {expected_basis_shape}, got {self.basis.shape}."
            )
        if self.valid.shape != self.coords.shape[:-1]:
            raise ValueError(
                f"`valid` must have shape {self.coords.shape[:-1]}, got {self.valid.shape}."
            )

    @property
    def ndim(self) -> int:
        """Coordinate dimension ``N`` of the underlying square chart."""

        return int(self.coords.shape[-1])


class CoordinateSystemND(ABC):
    """Abstract square local chart on ``R^N``.

    A concrete coordinate system represents a local chart ``chi: U -> V`` on a
    square ambient space ``R^N``. Evaluating the chart at physical points ``x``
    returns:

    - coordinates ``u = chi(x)``
    - the covariant tangent basis ``g_i = d x / d u_i``
    - a boolean mask describing the domain of definition ``x in U``

    Implementations must respect the numerical contract:

    - input points have shape ``(..., ndim)``
    - ``coords`` has shape ``(..., ndim)``
    - ``basis`` has shape ``(..., ndim, ndim)``
    - ``basis[..., i, :] = d x / d u_i``
    - ``valid`` has shape ``(...)``

    The chart is allowed to be defined only on a patch. Mixed arrays of valid
    and invalid points are supported by construction through the ``valid`` mask.
    """

    _fixed_ndim: int | None = None

    def __init__(self, ndim: int | None = None):
        fixed_ndim = self._fixed_ndim
        if ndim is None:
            if fixed_ndim is None:
                raise TypeError("`ndim` must be provided for generic coordinate systems.")
            ndim = fixed_ndim
        ndim = int(ndim)
        if fixed_ndim is not None and ndim != fixed_ndim:
            raise ValueError(f"{self.__class__.__name__} requires ndim={fixed_ndim}, got {ndim}.")
        if ndim <= 0:
            raise ValueError("`ndim` must be positive.")
        self.ndim = ndim

    @abstractmethod
    def call(self, points: np.ndarray) -> CoordinateSampleND:
        """Evaluate the chart on physical points."""

        raise NotImplementedError

    def eval(self, points) -> CoordinateSampleND:
        """Validate points and return a checked :class:`CoordinateSampleND`."""

        points = _validate_points(points, self.ndim)
        sample = self.call(points)
        if not isinstance(sample, CoordinateSampleND):
            raise TypeError(
                f"{self.__class__.__name__}.call() must return CoordinateSampleND, got {type(sample).__name__}."
            )
        if sample.ndim != self.ndim:
            raise ValueError(
                f"{self.__class__.__name__} returned ndim={sample.ndim}, expected {self.ndim}."
            )
        expected_coords = points.shape
        expected_basis = points.shape[:-1] + (self.ndim, self.ndim)
        expected_valid = points.shape[:-1]
        if sample.coords.shape != expected_coords:
            raise ValueError(
                f"{self.__class__.__name__} returned coords shape {sample.coords.shape}, expected {expected_coords}."
            )
        if sample.basis.shape != expected_basis:
            raise ValueError(
                f"{self.__class__.__name__} returned basis shape {sample.basis.shape}, expected {expected_basis}."
            )
        if sample.valid.shape != expected_valid:
            raise ValueError(
                f"{self.__class__.__name__} returned valid shape {sample.valid.shape}, expected {expected_valid}."
            )
        return sample

    def __call__(self, points) -> CoordinateSampleND:
        """Alias for :meth:`eval`."""

        return self.eval(points)

    def coords(self, points) -> np.ndarray:
        """Return only the coordinate values ``u = chi(x)``."""

        return self.eval(points).coords

    def basis(self, points) -> np.ndarray:
        """Return only the covariant tangent basis ``d x / d u_i``."""

        return self.eval(points).basis

    def valid(self, points) -> np.ndarray:
        """Return only the domain-of-definition mask."""

        return self.eval(points).valid

    def mask(self, fn):
        """Build a scalar mask evaluated in this coordinate system.

        The callable is invoked as ``fn(*coords)`` where ``coords[i]`` is the
        ``i``-th curvilinear coordinate component returned by this chart. This
        is a convenience wrapper over ``MaskND.from_function(..., coordinates=self)``.
        """

        from .base import MaskND

        return MaskND.from_function(fn, coordinates=self)

    def affine(self, matrix, offset=None) -> "AffineCoordinates":
        """Precompose the chart with an ambient affine map."""

        return AffineCoordinates(self, matrix=matrix, offset=offset)

    def shift(self, offset) -> "AffineCoordinates":
        """Precompose the chart with an ambient translation.

        This is a convenience wrapper over :meth:`affine` with the identity
        ambient matrix and the requested translation vector.
        """

        offset = _as_offset_vector(offset, self.ndim, name="offset")
        return self.affine(np.eye(self.ndim, dtype=np.float64), offset=offset)

    def remap(
        self,
        map_fn: Callable[[np.ndarray], np.ndarray],
        inverse_jacobian_fn: Callable[[np.ndarray, np.ndarray], np.ndarray],
        valid_fn: Callable[[np.ndarray, np.ndarray], np.ndarray] | None = None,
    ) -> "CoordinateRemap":
        """Postcompose the chart with a coordinate-space map."""

        return CoordinateRemap(
            self,
            map_fn=map_fn,
            inverse_jacobian_fn=inverse_jacobian_fn,
            valid_fn=valid_fn,
        )

    def mix(self, matrix, offset=None) -> "CoordinateRemap":
        """Postcompose the chart with an invertible affine coordinate mix.

        If the base chart returns ``u``, the new coordinates are
        ``v = M u + b``. This is a convenience wrapper over :meth:`remap`
        rather than a separate public class.
        """

        matrix = _validate_square_matrix(matrix, self.ndim, "matrix")
        inverse_matrix = np.linalg.inv(matrix)
        if offset is None:
            offset = np.zeros(self.ndim, dtype=np.float64)
        offset = _as_offset_vector(offset, self.ndim)

        def map_fn(u: np.ndarray) -> np.ndarray:
            return u @ matrix.T + offset

        def inverse_jacobian_fn(u: np.ndarray, v: np.ndarray) -> np.ndarray:
            return np.broadcast_to(inverse_matrix, u.shape[:-1] + (self.ndim, self.ndim))

        return self.remap(map_fn=map_fn, inverse_jacobian_fn=inverse_jacobian_fn)

    def restrict(
        self, predicate: Callable[[np.ndarray, CoordinateSampleND], np.ndarray]
    ) -> "RestrictedCoordinates":
        """Restrict the chart to a smaller patch of its domain."""

        return RestrictedCoordinates(self, predicate=predicate)

    def cylindrize(
        self,
        radius: int,
        angle: int,
        angle_min: float = 0.0,
        angle_max: float = 2.0 * np.pi,
        tol: float = 1e-32,
        orientation: float = 1.0,
    ) -> "CylindricalCoordinates":
        """Postcompose the current chart with a cylindrical coordinate remap.

        The selected coordinate components are reinterpreted as a radius-like
        and angle-like pair, while all remaining coordinates are passed through
        unchanged.
        """

        return CylindricalCoordinates(
            self,
            radius=radius,
            angle=angle,
            angle_min=angle_min,
            angle_max=angle_max,
            tol=tol,
            orientation=orientation,
        )

    def rotate(
        self,
        angle: float,
        axis: int | np.ndarray | tuple[float, float, float] = 2,
        origin=None,
    ) -> "AffineCoordinates":
        """Rotate the physical input space when ``ndim`` is 2 or 3.

        This method exists on the generic base for usability. It is defined
        only for ambient dimensions 2 and 3 and raises ``NotImplementedError``
        otherwise.
        """

        if self.ndim == 2:
            if origin is None:
                origin = (0.0, 0.0)
            origin = _as_offset_vector(origin, 2, name="origin")
            rotation = _rotation_matrix_2d(float(angle))
        elif self.ndim == 3:
            if origin is None:
                origin = (0.0, 0.0, 0.0)
            origin = _as_offset_vector(origin, 3, name="origin")
            rotation = _rotation_matrix_3d(axis, float(angle))
        else:
            raise NotImplementedError("Ambient rotation is only implemented for ndim=2 and ndim=3.")
        offset = origin - origin @ rotation.T
        return self.affine(rotation, offset=offset)

    def axisymmetrize(
        self,
        axis: int = 2,
        angle_min: float = 0.0,
        angle_max: float = 2.0 * np.pi,
        tol: float = 1e-32,
        orientation: float = 1.0,
    ) -> "AxisymmetricCoordinates":
        """Lift a 2D meridional chart to a 3D axisymmetric chart.

        This method is only defined for ``ndim = 2`` and raises
        ``NotImplementedError`` otherwise.
        """

        if self.ndim != 2:
            raise NotImplementedError("Axisymmetric lifting is only implemented for ndim=2.")
        return AxisymmetricCoordinates(
            self,
            axis=axis,
            angle_min=angle_min,
            angle_max=angle_max,
            tol=tol,
            orientation=orientation,
        )


class CoordinateSystem2D(CoordinateSystemND):
    """Typed specialization of :class:`CoordinateSystemND` for ``ndim = 2``.

    This class is intentionally lightweight. It exists for intrinsically 2D
    charts and for future type-level constraints in ansatz code, not to mirror
    every generic coordinate transformation with a dedicated ``*2D`` wrapper.
    """

    _fixed_ndim = 2


class CoordinateSystem3D(CoordinateSystemND):
    """Typed specialization of :class:`CoordinateSystemND` for ``ndim = 3``.

    This class is intentionally lightweight. It exists for intrinsically 3D
    charts and for future type-level constraints in ansatz code, not to mirror
    every generic coordinate transformation with a dedicated ``*3D`` wrapper.
    """

    _fixed_ndim = 3


class EuclideanCoordinates(CoordinateSystemND):
    """Identity chart on ``R^N``.

    The returned coordinates coincide with the ambient Euclidean coordinates,
    the basis is the constant Cartesian basis, and ``valid`` is identically
    ``True``.
    """

    def call(self, points: np.ndarray) -> CoordinateSampleND:
        """Return the identity chart evaluated on Euclidean points."""

        points = _validate_points(points, self.ndim)
        dtype = np.result_type(points, np.float64)
        coords = np.asarray(points, dtype=dtype)
        basis = _identity_basis(points.shape[:-1], self.ndim, dtype)
        valid = np.ones(points.shape[:-1], dtype=bool)
        return CoordinateSampleND(coords=coords, basis=basis, valid=valid)


class AffineCoordinates(CoordinateSystemND):
    """Precompose a chart with an invertible affine ambient transformation.

    If the base chart is ``chi(x)`` and the affine ambient map is
    ``y = A x + b``, then this wrapper defines the new chart

    ``chi'(y) = chi(A^{-1} (y - b))``.

    The covariant tangent basis transforms as

    ``g'_i(y) = A g_i(A^{-1}(y - b))``.
    """

    def __init__(self, system: CoordinateSystemND, matrix, offset=None):
        self.system = system
        self.matrix = _validate_square_matrix(matrix, system.ndim, "matrix")
        self.inverse_matrix = np.linalg.inv(self.matrix)
        if offset is None:
            offset = np.zeros(system.ndim, dtype=np.float64)
        self.offset = _as_offset_vector(offset, system.ndim)
        super().__init__(system.ndim)

    def call(self, points: np.ndarray) -> CoordinateSampleND:
        """Evaluate the wrapped chart after applying the inverse affine map."""

        points = _validate_points(points, self.ndim)
        source_points = (np.asarray(points, dtype=np.result_type(points, np.float64)) - self.offset) @ self.inverse_matrix.T
        sample = self.system.eval(source_points)
        basis = sample.basis @ self.matrix.T
        return CoordinateSampleND(coords=sample.coords, basis=basis, valid=sample.valid)


class CoordinateRemap(CoordinateSystemND):
    """Postcompose a chart with a coordinate-space map.

    If the base chart is ``u = chi(x)`` and the new coordinates are defined by
    ``v = F(u)``, then this wrapper returns ``v`` as the coordinate values.

    The caller must provide the inverse differential
    ``d u / d v``. The tangent basis is updated as

    ``d x / d v_j = sum_i (d x / d u_i) (d u_i / d v_j)``.

    Parameters
    ----------
    map_fn:
        Function taking ``u`` with shape ``(..., ndim)`` and returning
        ``v = F(u)`` with the same shape.
    inverse_jacobian_fn:
        Function returning ``d u / d v`` with shape ``(..., ndim, ndim)``.
    valid_fn:
        Optional extra validity mask in the new coordinates. The final validity
        mask is ``base.valid & valid_fn(u, v)``.
    """

    def __init__(
        self,
        system: CoordinateSystemND,
        map_fn: Callable[[np.ndarray], np.ndarray],
        inverse_jacobian_fn: Callable[[np.ndarray, np.ndarray], np.ndarray],
        valid_fn: Callable[[np.ndarray, np.ndarray], np.ndarray] | None = None,
    ):
        self.system = system
        self.map_fn = map_fn
        self.inverse_jacobian_fn = inverse_jacobian_fn
        self.valid_fn = valid_fn
        super().__init__(system.ndim)

    def call(self, points: np.ndarray) -> CoordinateSampleND:
        """Evaluate the wrapped chart and then remap its coordinate values."""

        sample = self.system.eval(points)
        mapped_coords = np.asarray(self.map_fn(sample.coords))
        if mapped_coords.shape != sample.coords.shape:
            raise ValueError(
                f"`map_fn` must return shape {sample.coords.shape}, got {mapped_coords.shape}."
            )
        inverse_jacobian = np.asarray(self.inverse_jacobian_fn(sample.coords, mapped_coords))
        expected_jacobian = sample.coords.shape[:-1] + (self.ndim, self.ndim)
        if inverse_jacobian.shape != expected_jacobian:
            raise ValueError(
                f"`inverse_jacobian_fn` must return shape {expected_jacobian}, got {inverse_jacobian.shape}."
            )
        valid = sample.valid
        if self.valid_fn is not None:
            valid_extra = np.asarray(self.valid_fn(sample.coords, mapped_coords), dtype=bool)
            if valid_extra.shape != sample.valid.shape:
                raise ValueError(
                    f"`valid_fn` must return shape {sample.valid.shape}, got {valid_extra.shape}."
                )
            valid = valid & valid_extra
        basis = _apply_inverse_coordinate_jacobian(sample.basis, inverse_jacobian)
        return CoordinateSampleND(coords=mapped_coords, basis=basis, valid=valid)


class RestrictedCoordinates(CoordinateSystemND):
    """Restrict an existing chart to an explicit atlas patch.

    The wrapped chart is evaluated unchanged, but the final validity mask is

    ``valid' = valid & predicate(x, sample)``.
    """

    def __init__(
        self,
        system: CoordinateSystemND,
        predicate: Callable[[np.ndarray, CoordinateSampleND], np.ndarray],
    ):
        self.system = system
        self.predicate = predicate
        super().__init__(system.ndim)

    def call(self, points: np.ndarray) -> CoordinateSampleND:
        """Evaluate the wrapped chart and apply the extra patch predicate."""

        sample = self.system.eval(points)
        valid_extra = np.asarray(self.predicate(np.asarray(points), sample), dtype=bool)
        if valid_extra.shape != sample.valid.shape:
            raise ValueError(
                f"`predicate` must return shape {sample.valid.shape}, got {valid_extra.shape}."
            )
        return CoordinateSampleND(coords=sample.coords, basis=sample.basis, valid=sample.valid & valid_extra)


class CylindricalCoordinates(CoordinateSystemND):
    """Coordinate-space cylindrical remap of an existing chart.

    Let the wrapped chart return coordinates ``u`` and choose two distinct
    coordinate components ``u[radius]`` and ``u[angle]``. This wrapper defines
    new coordinates ``v`` by

    ``v[radius] = sqrt(u[radius]^2 + u[angle]^2)``

    ``v[angle] = angle_min + (angle_max - angle_min) * phi / (2π)``

    where ``phi = mod(orientation * atan2(u[angle], u[radius]), 2π)``. All
    remaining coordinate components are copied unchanged.

    The covariant tangent basis is transformed by the inverse differential
    ``d u / d v`` of the cylindrical remap, so ``basis[..., i, :]`` continues
    to store ``d x / d v_i`` in the ambient space of the wrapped chart.
    """

    def __init__(
        self,
        system: CoordinateSystemND,
        radius: int,
        angle: int,
        angle_min: float = 0.0,
        angle_max: float = 2.0 * np.pi,
        tol: float = 1e-32,
        orientation: float = 1.0,
    ):
        self.system = system
        self.radius = int(radius)
        self.angle = int(angle)
        if self.radius == self.angle:
            raise ValueError("`radius` and `angle` must be different indices.")
        if self.radius < 0 or self.radius >= system.ndim:
            raise ValueError(f"`radius` must be in [0, {system.ndim}), got {self.radius}.")
        if self.angle < 0 or self.angle >= system.ndim:
            raise ValueError(f"`angle` must be in [0, {system.ndim}), got {self.angle}.")
        self.angle_min = float(angle_min)
        self.angle_max = float(angle_max)
        if not self.angle_max > self.angle_min:
            raise ValueError("`angle_max` must be strictly greater than `angle_min`.")
        self.angle_period = self.angle_max - self.angle_min
        self.angle_scale = _normalize_orientation(orientation) * (2.0 * np.pi / self.angle_period)
        self.orientation = 1.0 if self.angle_scale > 0 else -1.0
        self.tol = float(tol)
        if self.tol < 0:
            raise ValueError("`tol` must be non-negative.")
        super().__init__(system.ndim)

    def call(self, points: np.ndarray) -> CoordinateSampleND:
        """Evaluate the cylindrical remap and its transformed tangent basis."""

        points = _validate_points(points, self.ndim)
        source_sample = self.system.eval(points)
        dtype = np.result_type(source_sample.coords, np.float64)
        source_coords = np.asarray(source_sample.coords, dtype=dtype)

        radial_x = np.asarray(source_coords[..., self.radius], dtype=dtype)
        radial_y = np.asarray(source_coords[..., self.angle], dtype=dtype)
        radius = np.sqrt(radial_x**2 + radial_y**2)
        psi = np.mod(np.arctan2(radial_y, radial_x), 2.0 * np.pi)
        angle_coord = self.angle_min + (self.angle_period / (2.0 * np.pi)) * np.mod(
            self.orientation * psi,
            2.0 * np.pi,
        )

        coords = np.array(source_coords, copy=True)
        coords[..., self.radius] = radius
        coords[..., self.angle] = angle_coord

        phi = self.angle_scale * (angle_coord - self.angle_min)
        cos_phi = np.cos(phi)
        sin_phi = np.sin(phi)

        inverse_jacobian = np.array(_identity_basis(points.shape[:-1], self.ndim, dtype), copy=True)
        inverse_jacobian[..., self.radius, self.radius] = cos_phi
        inverse_jacobian[..., self.angle, self.radius] = sin_phi
        inverse_jacobian[..., self.radius, self.angle] = -radius * sin_phi * self.angle_scale
        inverse_jacobian[..., self.angle, self.angle] = radius * cos_phi * self.angle_scale

        basis = _apply_inverse_coordinate_jacobian(source_sample.basis, inverse_jacobian)
        valid = source_sample.valid & (radius > self.tol)
        return CoordinateSampleND(coords=coords, basis=basis, valid=valid)


class PolarCoordinates(CylindricalCoordinates, CoordinateSystem2D):
    """Convenience specialization of :class:`CylindricalCoordinates` on ``R^2``.

    This class keeps the ergonomic public name ``PolarCoordinates`` but does
    not implement separate mathematics. It is exactly
    ``CylindricalCoordinates(EuclideanCoordinates(ndim=2), radius=0, angle=1, ...)``.
    """

    def __init__(
        self,
        angle_min: float = 0.0,
        angle_max: float = 2.0 * np.pi,
        tol: float = 1e-32,
        orientation: float = 1.0,
    ):
        super().__init__(
            EuclideanCoordinates(ndim=2),
            radius=0,
            angle=1,
            angle_min=angle_min,
            angle_max=angle_max,
            tol=tol,
            orientation=orientation,
        )


class AxisymmetricCoordinates(CoordinateSystem3D):
    """Lift a meridional 2D chart to a 3D axisymmetric chart.

    The wrapped meridional chart is evaluated on physical meridional points
    ``(rho, z)`` where ``rho`` is the Euclidean distance to the chosen rotation
    axis and ``z`` is the ambient coordinate along that axis.

    If the meridional chart returns coordinates ``(u_1, u_2)``, this class
    returns the 3D coordinates ``(u_1, u_2, a)``, where ``a`` is an angular
    source coordinate whose physical azimuth is

    ``phi(a) = orientation * 2π * (a - angle_min) / (angle_max - angle_min)``.

    The tangent basis is

    - ``d x / d u_k = (d rho / d u_k) e_r + (d z / d u_k) e_axis`` for
      ``k = 1, 2``
    - ``d x / d a = rho * d phi / d a * e_phi``

    The chart is singular on the rotation axis ``rho <= tol`` and those points
    are marked invalid via ``valid == False``.
    """

    def __init__(
        self,
        meridional: CoordinateSystemND,
        axis: int = 2,
        angle_min: float = 0.0,
        angle_max: float = 2.0 * np.pi,
        tol: float = 1e-32,
        orientation: float = 1.0,
    ):
        super().__init__()
        if meridional.ndim != 2:
            raise ValueError(
                f"`meridional` must have ndim=2, got {meridional.ndim} from {meridional.__class__.__name__}."
            )
        self.meridional = meridional
        self.axis = int(axis)
        if self.axis not in (0, 1, 2):
            raise ValueError("`axis` must be one of 0, 1, 2.")
        self.angle_min = float(angle_min)
        self.angle_max = float(angle_max)
        if not self.angle_max > self.angle_min:
            raise ValueError("`angle_max` must be strictly greater than `angle_min`.")
        self.angle_period = self.angle_max - self.angle_min
        self.angle_scale = _normalize_orientation(orientation) * (2.0 * np.pi / self.angle_period)
        self.orientation = 1.0 if self.angle_scale > 0 else -1.0
        self.tol = float(tol)
        if self.tol < 0:
            raise ValueError("`tol` must be non-negative.")
        self.transverse_axes = tuple(index for index in range(3) if index != self.axis)

    def call(self, points: np.ndarray) -> CoordinateSampleND:
        """Evaluate the lifted 3D axisymmetric chart on Euclidean points."""

        points = _validate_points(points, self.ndim)
        dtype = np.result_type(points, np.float64)
        transverse_x = np.asarray(points[..., self.transverse_axes[0]], dtype=dtype)
        transverse_y = np.asarray(points[..., self.transverse_axes[1]], dtype=dtype)
        axial = np.asarray(points[..., self.axis], dtype=dtype)

        rho = np.sqrt(transverse_x**2 + transverse_y**2)
        psi = np.mod(np.arctan2(transverse_y, transverse_x), 2.0 * np.pi)
        angle_coord = self.angle_min + (self.angle_period / (2.0 * np.pi)) * np.mod(
            self.orientation * psi,
            2.0 * np.pi,
        )
        meridional_points = np.stack([rho, axial], axis=-1)
        meridional_sample = self.meridional.eval(meridional_points)
        phi = self.angle_scale * (angle_coord - self.angle_min)

        er = np.zeros(points.shape, dtype=dtype)
        er[..., self.transverse_axes[0]] = np.cos(phi)
        er[..., self.transverse_axes[1]] = np.sin(phi)

        ephi = np.zeros(points.shape, dtype=dtype)
        ephi[..., self.transverse_axes[0]] = -np.sin(phi)
        ephi[..., self.transverse_axes[1]] = np.cos(phi)

        eaxis = np.zeros(points.shape, dtype=dtype)
        eaxis[..., self.axis] = 1.0

        basis = np.zeros(points.shape[:-1] + (3, 3), dtype=dtype)
        basis[..., 0, :] = (
            meridional_sample.basis[..., 0, 0, None] * er
            + meridional_sample.basis[..., 0, 1, None] * eaxis
        )
        basis[..., 1, :] = (
            meridional_sample.basis[..., 1, 0, None] * er
            + meridional_sample.basis[..., 1, 1, None] * eaxis
        )
        basis[..., 2, :] = (rho * self.angle_scale)[..., None] * ephi

        coords = np.concatenate([meridional_sample.coords, angle_coord[..., None]], axis=-1)
        valid = meridional_sample.valid & (rho > self.tol)
        return CoordinateSampleND(coords=coords, basis=basis, valid=valid)
