import numpy as np
import pytest
from numpy.testing import assert_allclose

import topovec as tv
from topovec import System
from topovec.ansatz import (
    ArtLoopFinger,
    AxisSection,
    BlochSection,
    CF1,
    CF2,
    CF3,
    CF4,
    CholestericSpiral3D,
    ConstantField3D,
    CoordinateSpiral3D,
    EuclideanCoordinates,
    Hopfion,
    LoopFinger,
    MaskND,
    SectionEmbedding3D,
    SpiralSection,
    StraightFinger,
    Twist,
)


def make_field_coords(size=(5, 4, 3)) -> np.ndarray:
    return System.cubic(size=size).spin_positions()


def make_section_coords(size=(5, 4, 3)) -> tuple[np.ndarray, np.ndarray]:
    xyz = make_field_coords(size=size)
    x, _, z = np.moveaxis(xyz, -1, 0)
    return x, z


def cartesian(aa, bb, nn):
    grid = np.meshgrid(
        *(np.linspace(a, b, n) for a, b, n in zip(aa, bb, nn)),
        indexing="ij",
    )
    return np.stack(grid, axis=-1)


def normalize(s, eps=1e-32):
    return s / np.sqrt(np.sum(s**2, axis=-1, keepdims=True) + eps)


def dist_to_surface(xyz, guard=1e-8, d=0.1, p=-2):
    z = xyz[..., 2]
    a = z - np.min(z)
    b = np.max(z) - z
    return (a / d + guard) ** p, (b / d + guard) ** p


def surface(xyz, **kwargs):
    a, b = dist_to_surface(xyz, **kwargs)
    s = np.zeros_like(xyz)
    s[..., 2] = a + b
    return s


def outer_cone_uniform(distance, c=1, shift=0):
    s = np.zeros(distance.shape + (3,), dtype=np.float64)
    s[..., 2] = np.exp(c * (distance - shift))
    return s


def cone_coordinates(phi: float):
    matrix = np.array(
        [
            [np.sin(phi), np.cos(phi)],
            [np.cos(phi), -np.sin(phi)],
        ],
        dtype=np.float64,
    )
    meridional = EuclideanCoordinates(ndim=2).mix(matrix=matrix)
    return meridional.axisymmetrize(axis=2, tol=1e-12, orientation=1.0)


def anticone_coordinates(phi: float):
    matrix = np.array(
        [
            [-np.sin(phi), np.cos(phi)],
            [-np.cos(phi), -np.sin(phi)],
        ],
        dtype=np.float64,
    )
    meridional = EuclideanCoordinates(ndim=2).mix(matrix=matrix)
    return meridional.axisymmetrize(axis=2, tol=1e-12, orientation=-1.0)


def apply_affine(points: np.ndarray, matrix: np.ndarray, offset: np.ndarray) -> np.ndarray:
    return np.asarray(points) @ np.asarray(matrix).T + np.asarray(offset)


def pushforward_rows(vectors: np.ndarray, matrix: np.ndarray) -> np.ndarray:
    return np.asarray(vectors) @ np.asarray(matrix).T


def section_output_basis_from_sample(sample) -> np.ndarray:
    basis = np.zeros(sample.coords.shape[:-1] + (3, 3), dtype=np.result_type(sample.basis, np.float64))
    basis[..., 0, 0] = sample.basis[..., 0, 0]
    basis[..., 0, 2] = sample.basis[..., 0, 1]
    basis[..., 1, 1] = 1.0
    basis[..., 2, 0] = sample.basis[..., 1, 0]
    basis[..., 2, 2] = sample.basis[..., 1, 1]
    return basis


def _mask_coordinates(ndim: int | None = None, coordinates=None):
    if coordinates is not None:
        return coordinates
    if ndim is None:
        raise TypeError("`ndim` is required when `coordinates` is not provided.")
    return EuclideanCoordinates(ndim=ndim)


def component_mask(index: int, *, ndim: int | None = None, coordinates=None) -> MaskND:
    coords = _mask_coordinates(ndim=ndim, coordinates=coordinates)
    return coords.mask(lambda *uu: uu[index])


def radial_distance_mask(
    *,
    center=(0.0, 0.0),
    axes=(0, 1),
    tol: float = 0.0,
    ndim: int | None = None,
    coordinates=None,
) -> MaskND:
    coords = _mask_coordinates(ndim=ndim, coordinates=coordinates)
    center = tuple(float(value) for value in center)
    axes = tuple(int(axis) for axis in axes)

    def fn(*uu):
        radius2 = 0.0
        for axis, value in zip(axes, center):
            radius2 = radius2 + (uu[axis] - value) ** 2
        if tol > 0:
            radius2 = np.maximum(radius2, tol)
        return np.sqrt(radius2)

    return MaskND.from_function(fn, coordinates=coords)


def gaussian_mask(
    *,
    center=(0.0, 0.0),
    sigma=1.0,
    axes=(0, 1),
    ndim: int | None = None,
    coordinates=None,
) -> MaskND:
    coords = _mask_coordinates(ndim=ndim, coordinates=coordinates)
    center = tuple(float(value) for value in center)
    axes = tuple(int(axis) for axis in axes)
    sigma = np.asarray(sigma, dtype=np.float64)
    if sigma.ndim == 0:
        sigma = np.broadcast_to(sigma, (len(axes),))

    def fn(*uu):
        radius2 = 0.0
        for axis, value, scale in zip(axes, center, sigma):
            radius2 = radius2 + ((uu[axis] - value) / scale) ** 2
        return np.exp(-0.5 * radius2)

    return MaskND.from_function(fn, coordinates=coords)


def boundary_distance_mask(
    *,
    axis: int = 0,
    side: str = "min",
    ndim: int | None = None,
    coordinates=None,
) -> MaskND:
    coords = _mask_coordinates(ndim=ndim, coordinates=coordinates)

    def fn(*uu):
        coord = uu[axis]
        if side == "min":
            return coord - np.min(coord)
        return np.max(coord) - coord

    return MaskND.from_function(fn, coordinates=coords)


def boundary_power_mask(
    *,
    axis: int = 0,
    scale: float = 1.0,
    exponent: float = -2.0,
    guard: float = 1e-8,
    sides="both",
    ndim: int | None = None,
    coordinates=None,
) -> MaskND:
    coords = _mask_coordinates(ndim=ndim, coordinates=coordinates)
    if sides == "both":
        sides = ("min", "max")
    elif isinstance(sides, str):
        sides = (sides,)
    sides = tuple(sides)

    def fn(*uu):
        coord = uu[axis]
        result = 0.0
        if "min" in sides:
            result = result + ((coord - np.min(coord)) / scale + guard) ** exponent
        if "max" in sides:
            result = result + ((np.max(coord) - coord) / scale + guard) ** exponent
        return result

    return MaskND.from_function(fn, coordinates=coords)


def sigmoid_mask(source: MaskND, *, center: float = 0.0, width: float = 1.0, inverted: bool = False) -> MaskND:
    value = (source - center) / width
    if inverted:
        value = -value
    return 1.0 / (1.0 + (-value).exp())


def test_section_shapes_and_coordinate_binding():
    x, z = make_section_coords()
    sections = [
        AxisSection(axis=(0, 0, 1)),
        BlochSection(),
        SpiralSection(),
        CF1(),
        CF2(),
        CF3(),
        CF4(),
        Twist(),
    ]

    for section in sections:
        values = section(x, z)
        assert values.shape == x.shape + (3,)

    base = Twist()
    shifted = base.shift(x=1.25, z=0.75)
    explicit = Twist(coordinates=EuclideanCoordinates(ndim=2).shift((1.25, 0.75)))

    assert_allclose(shifted(x, z), base(x - 1.25, z - 0.75))
    assert_allclose(explicit(x, z), shifted(x, z))


def test_ansatz2d_is_covariant_and_eval_normalizes():
    points = cartesian((-1.0, -0.6), (0.8, 0.7), (6, 5))
    matrix = np.array([[2.0, 0.3], [-0.1, 0.7]], dtype=np.float64)
    offset = np.array([0.25, -0.4], dtype=np.float64)
    deformed_points = apply_affine(points, matrix, offset)
    coordinates = EuclideanCoordinates(ndim=2).affine(matrix=matrix, offset=offset)

    base = Twist()
    deformed = Twist(coordinates=coordinates)
    sample = coordinates(deformed_points)
    expected = np.einsum(
        "...i,...ij->...j",
        base.call(points),
        section_output_basis_from_sample(sample),
    )

    raw = deformed.call(deformed_points)
    values = deformed.eval(deformed_points)

    assert_allclose(raw, expected, atol=1e-12)
    assert np.any(np.abs(np.linalg.norm(raw, axis=-1) - 1.0) > 1e-3)
    assert_allclose(np.linalg.norm(values, axis=-1), 1.0, atol=1e-6)
    assert_allclose(deformed(deformed_points), values)


def test_field_shapes_and_eval_alias():
    xyz = make_field_coords()
    fields = [
        Hopfion(size=10),
        StraightFinger(Twist()),
        LoopFinger(Twist()),
        ArtLoopFinger(Twist(), rad=2.0),
        Hopfion(size=10).shift(x=1.0, y=-0.5, z=0.25),
        ConstantField3D((0, 0, 1)),
        CholestericSpiral3D(period=4.0),
        ConstantField3D((0, 0, 1)) + CholestericSpiral3D(period=4.0),
        CholestericSpiral3D(period=4.0) * gaussian_mask(center=(0.0, 0.0), sigma=5.0, ndim=3),
    ]

    for field in fields:
        raw = field.call(xyz)
        values = field.eval(xyz)
        assert raw.shape == xyz.shape
        assert values.shape == xyz.shape
        assert_allclose(field(xyz), values)


def test_ansatz3d_call_is_raw_and_eval_normalizes():
    xyz = make_field_coords(size=(9, 3, 7))
    field = StraightFinger(Twist(), uni=1.0)

    raw = field.call(xyz)
    evaluated = field.eval(xyz)

    raw_norms = np.linalg.norm(raw, axis=-1)
    eval_norms = np.linalg.norm(evaluated, axis=-1)

    assert np.any(np.abs(raw_norms - 1.0) > 1e-3)
    assert_allclose(eval_norms, 1.0, atol=1e-6)
    assert_allclose(field(xyz), evaluated)


def test_constant_field_raw_pushforward_matches_affine_covariance():
    xyz = cartesian((-1.2, -0.8, -0.4), (0.7, 0.6, 0.5), (5, 4, 3))
    matrix = np.array(
        [
            [1.7, 0.2, 0.0],
            [0.0, 0.9, -0.1],
            [0.1, 0.0, 1.3],
        ],
        dtype=np.float64,
    )
    offset = np.array([0.4, -0.5, 0.25], dtype=np.float64)
    deformed_xyz = apply_affine(xyz, matrix, offset)

    base = ConstantField3D((0.5, -1.0, 2.0))
    deformed = ConstantField3D(
        (0.5, -1.0, 2.0),
        coordinates=EuclideanCoordinates(ndim=3).affine(matrix=matrix, offset=offset),
    )

    assert_allclose(deformed.call(deformed_xyz), pushforward_rows(base.call(xyz), matrix))


def test_hopfion_and_straight_finger_are_covariant_under_affine_coordinates():
    xyz = cartesian((-1.0, -0.7, -0.6), (0.8, 0.9, 0.5), (6, 5, 4))
    matrix = np.array(
        [
            [1.4, 0.1, 0.0],
            [0.0, 0.8, -0.2],
            [0.15, 0.0, 1.1],
        ],
        dtype=np.float64,
    )
    offset = np.array([-0.3, 0.2, 0.4], dtype=np.float64)
    deformed_xyz = apply_affine(xyz, matrix, offset)
    coordinates = EuclideanCoordinates(ndim=3).affine(matrix=matrix, offset=offset)

    base_hopfion = Hopfion(size=3.5)
    base_finger = StraightFinger(Twist(), uni=0.25)
    deformed_hopfion = Hopfion(size=3.5, coordinates=coordinates)
    deformed_finger = StraightFinger(Twist(), uni=0.25, coordinates=coordinates)

    assert_allclose(
        deformed_hopfion.call(deformed_xyz),
        pushforward_rows(base_hopfion.call(xyz), matrix),
        atol=1e-12,
    )
    assert_allclose(
        deformed_finger.call(deformed_xyz),
        pushforward_rows(base_finger.call(xyz), matrix),
        atol=1e-12,
    )


def test_section_embedding_identity_lift_reproduces_straight_finger():
    xyz = make_field_coords(size=(8, 5, 6))
    section = Twist(
        coordinates=EuclideanCoordinates(ndim=2).affine(
            matrix=np.array([[1.3, 0.2], [0.1, 0.9]], dtype=np.float64),
            offset=(0.2, -0.3),
        )
    )

    generic = SectionEmbedding3D(
        section,
        section_axes=(0, 2),
        background_components=(0.0, 0.0, 0.5),
    )
    straight = StraightFinger(section, uni=0.5)

    assert_allclose(generic.call(xyz), straight.call(xyz))


def test_section_embedding_cylindrical_lift_reproduces_loop_finger():
    xyz = make_field_coords(size=(8, 7, 5))
    section = Twist()

    generic = SectionEmbedding3D(
        section,
        lift_builder=lambda coords: coords.cylindrize(
            radius=0,
            angle=1,
            angle_min=0.0,
            angle_max=2.0 * np.pi,
            tol=1e-15,
            orientation=1.0,
        ),
        section_input_map=lambda sample: (sample.coords[..., 0], sample.coords[..., 2]),
        background_components=(0.0, 0.0, 0.25),
    )
    loop = LoopFinger(section, uni=0.25)

    assert_allclose(generic.call(xyz), loop.call(xyz))


def test_section_embedding_cylindrical_shift_and_normalize_reproduces_art_loop():
    xyz = make_field_coords(size=(8, 7, 5))
    section = AxisSection(axis=(0.0, 2.0, 1.0))

    generic = SectionEmbedding3D(
        section,
        lift_builder=lambda coords: coords.cylindrize(
            radius=0,
            angle=1,
            angle_min=0.0,
            angle_max=2.0 * np.pi,
            tol=1e-15,
            orientation=1.0,
        ),
        section_input_map=lambda sample: (sample.coords[..., 0] - 2.0, sample.coords[..., 2]),
        normalize_section=True,
    )
    art = ArtLoopFinger(section, rad=2.0)

    assert_allclose(generic.call(xyz), art.call(xyz))


def test_raw_pushforward_is_not_orthonormalized_under_anisotropic_affine_maps():
    matrix = np.diag([2.0, 3.0, 1.0]).astype(np.float64)
    field = ConstantField3D((1.0, 1.0, 0.0), coordinates=EuclideanCoordinates(ndim=3).affine(matrix=matrix))
    point = np.zeros((1, 1, 1, 3), dtype=np.float64)

    raw = field.call(point)

    assert_allclose(raw, np.array([[[[2.0, 3.0, 0.0]]]]))
    assert not np.allclose(raw, np.array([[[[1.0, 1.0, 0.0]]]]))


def test_coordinate_binding_for_fields_and_masks():
    xyz = make_field_coords(size=(7, 6, 5))

    base_field = StraightFinger(Twist(), uni=0.5)
    shifted_field = base_field.shift(x=1.0, y=-0.5, z=0.25)
    field_delta = np.array([1.0, -0.5, 0.25])
    assert_allclose(shifted_field.call(xyz), base_field.call(xyz - field_delta))

    base_mask = radial_distance_mask(center=(0.0, 0.0), axes=(0, 1), ndim=3)
    shifted_mask = base_mask.shift(x=1.0, y=-2.0, z=0.0)
    mask_delta = np.array([1.0, -2.0, 0.0])
    assert_allclose(shifted_mask(xyz), base_mask(xyz - mask_delta))

    explicit_field = Hopfion(
        size=7.0,
        coordinates=EuclideanCoordinates(ndim=3).shift((0.25, -0.75, 0.5)),
    )
    reference_field = Hopfion(size=7.0)
    assert_allclose(explicit_field.call(xyz), reference_field.call(xyz - np.array([0.25, -0.75, 0.5])))

    base_loop = LoopFinger(Twist(), uni=0.5)
    shifted_loop = base_loop.shift(x=1.0, y=-0.5, z=0.25)
    assert isinstance(shifted_loop, LoopFinger)
    assert_allclose(shifted_loop.call(xyz), base_loop.call(xyz - field_delta))


def test_callable_mask_sugar_for_2d_and_3d_ansatze():
    x, z = make_section_coords(size=(8, 1, 6))
    xyz = make_field_coords(size=(8, 7, 6))

    section = Twist().shift(x=0.5, z=-0.25)
    field = Hopfion(size=5.0).shift(x=0.5, y=-0.25, z=0.75)

    section_masked = section.mask(lambda x, z: np.exp(-x))
    field_masked = field.mask(lambda x, y, z: np.exp(-x))

    explicit_section_mask = MaskND.from_function(lambda x, z: np.exp(-x), coordinates=section.coordinates)
    explicit_field_mask = MaskND.from_function(
        lambda x, y, z: np.exp(-x),
        coordinates=field.coordinates,
    )

    assert_allclose(section_masked(x, z), (section * explicit_section_mask)(x, z))
    assert_allclose(field_masked.call(xyz), (field * explicit_field_mask).call(xyz))


def test_callable_mask_sugar_chains_and_sees_curvilinear_coordinates():
    xyz = make_field_coords(size=(7, 6, 5))
    field = Hopfion(size=4.0, coordinates=EuclideanCoordinates(ndim=3).shift((1.0, -0.5, 0.0)))

    f = lambda x, y, z: np.exp(-x)
    g = lambda x, y, z: np.exp(x - 1.0)

    chained = field.mask(f).mask(g)
    explicit = field * MaskND.from_function(f, coordinates=field.coordinates) * MaskND.from_function(
        g,
        coordinates=field.coordinates,
    )

    assert_allclose(chained.call(xyz), explicit.call(xyz))
    assert_allclose(field.mask(lambda x, y, z: x).call(xyz), field.call(xyz) * (xyz[..., 0] - 1.0)[..., None])


def test_callable_mask_construction_supports_scalar_array_and_broadcast():
    xyz = make_field_coords(size=(6, 5, 4))
    coords = EuclideanCoordinates(ndim=3)

    scalar = MaskND.from_function(lambda x, y, z: 2.0, coordinates=coords)
    array = MaskND.from_function(lambda x, y, z: x + z, coordinates=coords)
    broadcast = MaskND.from_function(
        lambda x, y, z: np.ones((xyz.shape[0], 1, xyz.shape[2], 1)),
        coordinates=coords,
    )

    assert_allclose(scalar(xyz), 2.0)
    assert_allclose(array(xyz), xyz[..., 0] + xyz[..., 2])
    assert_allclose(broadcast(xyz), 1.0)


def test_callable_mask_sugar_is_strict_on_ambiguous_compositions():
    left = Hopfion(size=4.0, coordinates=EuclideanCoordinates(ndim=3).shift((1.0, 0.0, 0.0)))
    right = Hopfion(size=4.0, coordinates=EuclideanCoordinates(ndim=3).shift((0.0, 1.0, 0.0)))

    with pytest.raises(ValueError, match="unambiguous associated coordinate system"):
        (left + right).mask(lambda x, y, z: np.exp(-x))


def test_mask_shapes_in_2d_and_3d():
    x, z = make_section_coords(size=(6, 1, 4))
    points2 = np.stack([x[:, 0, :, 0], z[:, 0, :, 0]], axis=-1)
    xyz = make_field_coords(size=(6, 5, 4))
    coords2 = EuclideanCoordinates(ndim=2)
    coords3 = EuclideanCoordinates(ndim=3)

    masks2 = [
        MaskND.from_function(lambda x, z: 1.0, coordinates=coords2),
        coords2.mask(lambda x, z: x),
        radial_distance_mask(center=(0.0, 0.0), axes=(0, 1), coordinates=coords2),
        gaussian_mask(center=(0.0, 0.0), sigma=2.0, axes=(0, 1), coordinates=coords2),
    ]
    for mask in masks2:
        assert mask(points2).shape == points2.shape[:-1]

    masks3 = [
        MaskND.from_function(lambda x, y, z: 1.0, coordinates=coords3),
        coords3.mask(lambda x, y, z: x),
        radial_distance_mask(center=(0.0, 0.0), axes=(0, 1), coordinates=coords3),
        gaussian_mask(center=(0.0, 0.0), sigma=2.0, axes=(0, 1), coordinates=coords3),
        boundary_distance_mask(axis=2, side="min", coordinates=coords3),
        boundary_power_mask(axis=2, scale=1.0, exponent=-2, coordinates=coords3),
        sigmoid_mask(
            radial_distance_mask(center=(0.0, 0.0), axes=(0, 1), coordinates=coords3),
            center=2.0,
            width=0.5,
        ),
    ]
    for mask in masks3:
        assert mask(xyz).shape == xyz.shape[:-1]


def test_callable_masks_reproduce_removed_mask_recipes():
    xyz = make_field_coords(size=(6, 5, 4))
    coords = EuclideanCoordinates(ndim=3)

    component = component_mask(0, coordinates=coords)
    radial = radial_distance_mask(center=(1.0, -0.5), axes=(0, 1), coordinates=coords)
    gaussian = gaussian_mask(center=(1.0, -0.5), sigma=(2.0, 3.0), axes=(0, 1), coordinates=coords)
    boundary = boundary_power_mask(axis=2, scale=1.1, exponent=-2.0, guard=1e-8, coordinates=coords)

    expected_radial = np.sqrt((xyz[..., 0] - 1.0) ** 2 + (xyz[..., 1] + 0.5) ** 2)
    expected_gaussian = np.exp(-0.5 * (((xyz[..., 0] - 1.0) / 2.0) ** 2 + ((xyz[..., 1] + 0.5) / 3.0) ** 2))
    z = xyz[..., 2]
    expected_boundary = ((z - np.min(z)) / 1.1 + 1e-8) ** -2.0 + ((np.max(z) - z) / 1.1 + 1e-8) ** -2.0

    assert_allclose(component(xyz), xyz[..., 0])
    assert_allclose(radial(xyz), expected_radial)
    assert_allclose(gaussian(xyz), expected_gaussian)
    assert_allclose(boundary(xyz), expected_boundary)


def test_additive_masked_composition_reproduces_notebook_recipe():
    size = np.array([12.0, 12.0, 4.0])
    xyz = cartesian(-size / 2, size / 2, (20, 20, 9))
    thickness = size[2]

    pitch = 7.5
    arm = 6.0
    outer_radius = arm / 2
    angle = np.arctan2(thickness, arm)
    shift0 = 0.5
    radial_decay = 0.05
    cut_slope = 0.7
    w = np.sqrt(thickness**2 + arm**2)

    radial = np.exp(-radial_decay * np.sqrt(xyz[..., 0] ** 2 + xyz[..., 1] ** 2))[..., None]

    cone0 = cone_coordinates(angle)
    anticone = anticone_coordinates(angle)
    cone2 = cone_coordinates(-angle)

    sample0 = cone0(xyz)
    sample1 = anticone(xyz)
    sample2 = cone2(xyz)

    phase0 = ((sample0.coords[..., 0] + shift0) * 2.0 * np.pi / pitch)[..., None]
    phase1 = ((sample1.coords[..., 0] + shift0) * 2.0 * np.pi / pitch)[..., None]
    spiral0 = (
        np.cos(phase0) * sample0.basis[..., 1, :]
        + np.sin(phase0) * sample0.basis[..., 2, :]
    )
    spiral1 = (
        np.cos(phase1) * sample1.basis[..., 1, :]
        + np.sin(phase1) * sample1.basis[..., 2, :]
    )
    reference = normalize(
        spiral0 * radial
        + spiral1
        + outer_cone_uniform(sample2.coords[..., 0], c=-cut_slope, shift=-(thickness * outer_radius) / w)
        + outer_cone_uniform(sample0.coords[..., 0], c=cut_slope, shift=(thickness * outer_radius) / w)
        + surface(xyz, p=-2, d=1.1)
    )

    z_axis = ConstantField3D((0.0, 0.0, 1.0))
    a0_mask = component_mask(0, coordinates=cone0)
    a2_mask = component_mask(0, coordinates=cone2)
    radial_mask = (-radial_decay * radial_distance_mask(center=(0.0, 0.0), axes=(0, 1), ndim=3)).exp()
    soliton = (
        CholestericSpiral3D(
            phase_axis=0,
            plane_axes=(1, 2),
            period=pitch,
            shift=shift0,
            coordinates=cone0,
        )
        * radial_mask
        + CholestericSpiral3D(
            phase_axis=0,
            plane_axes=(1, 2),
            period=pitch,
            shift=shift0,
            coordinates=anticone,
        )
        + z_axis * (-cut_slope * (a2_mask + (thickness * outer_radius) / w)).exp()
        + z_axis * (cut_slope * (a0_mask - (thickness * outer_radius) / w)).exp()
        + z_axis * boundary_power_mask(axis=2, scale=1.1, exponent=-2, ndim=3)
    )

    candidate = soliton(xyz)
    assert_allclose(candidate, reference, atol=1e-10, rtol=1e-10)


def test_loop_finger_uses_local_cylindrical_basis_under_affine_deformation():
    matrix = np.diag([2.0, 3.0, 1.0]).astype(np.float64)
    coordinates = EuclideanCoordinates(ndim=3).affine(matrix=matrix)

    point_y_axis = np.array([[[[0.0, 3.0, 0.0]]]], dtype=np.float64)
    radial_section = AxisSection(axis=(1.0, 0.0, 0.0))
    azimuthal_section = AxisSection(axis=(0.0, 1.0, 0.0))

    radial_field = LoopFinger(radial_section, uni=0.0, coordinates=coordinates)
    azimuthal_field = LoopFinger(azimuthal_section, uni=0.0, coordinates=coordinates)
    radial_back = LoopFinger(radial_section, uni=0.0, back=True, coordinates=coordinates)

    assert_allclose(radial_field.call(point_y_axis), np.array([[[[0.0, 9.0, 0.0]]]]), atol=1e-12)
    assert_allclose(azimuthal_field.call(point_y_axis), np.array([[[[-6.0, 0.0, 0.0]]]]), atol=1e-12)
    assert_allclose(radial_back.call(point_y_axis), np.array([[[[0.0, -9.0, 0.0]]]]), atol=1e-12)


def test_loop_finger_masking_uses_public_ambient_coordinates():
    xyz = make_field_coords(size=(7, 6, 5))
    field = LoopFinger(Twist(), coordinates=EuclideanCoordinates(ndim=3).shift((1.0, 0.0, 0.0)))

    masked = field.mask(lambda x, y, z: x)

    assert_allclose(masked.call(xyz), field.call(xyz) * (xyz[..., 0] - 1.0)[..., None])


def test_art_loop_finger_uses_local_cylindrical_basis_under_affine_deformation():
    matrix = np.diag([2.0, 3.0, 1.0]).astype(np.float64)
    coordinates = EuclideanCoordinates(ndim=3).affine(matrix=matrix)

    point_y_axis = np.array([[[[0.0, 3.0, 0.0]]]], dtype=np.float64)
    azimuthal_section = AxisSection(axis=(0.0, 1.0, 0.0))

    front = ArtLoopFinger(azimuthal_section, rad=1.0, coordinates=coordinates)
    back = ArtLoopFinger(azimuthal_section, rad=1.0, back=True, coordinates=coordinates)

    assert_allclose(front.call(point_y_axis), np.array([[[[-2.0, 0.0, 0.0]]]]), atol=1e-12)
    assert_allclose(back.call(point_y_axis), np.array([[[[2.0, 0.0, 0.0]]]]), atol=1e-12)


def test_field_algebra_helpers_preserve_local_coordinates():
    xyz = make_field_coords(size=(7, 7, 5))
    base = ConstantField3D((0.0, 0.0, 2.0))
    spiral = CholestericSpiral3D(
        period=5.0,
        coordinates=EuclideanCoordinates(ndim=3).rotate(angle=np.pi / 6, axis=2),
    )
    mask = gaussian_mask(
        center=(1.0, -0.5),
        sigma=5.0,
        axes=(0, 1),
        coordinates=EuclideanCoordinates(ndim=3).shift((1.0, -0.5, 0.0)),
    )

    weighted_left = spiral * mask
    weighted_right = mask * spiral
    combined = base + weighted_left
    negated = -base

    assert_allclose(weighted_left.call(xyz), weighted_right.call(xyz))
    assert_allclose(weighted_left.call(xyz), spiral.call(xyz) * mask(xyz)[..., None])
    assert_allclose(combined.call(xyz), base.call(xyz) + weighted_left.call(xyz))
    assert_allclose(negated.call(xyz), -base.call(xyz))
    assert_allclose(combined(xyz), combined.eval(xyz))
    assert_allclose(np.linalg.norm(combined(xyz), axis=-1), 1.0, atol=1e-6)


def test_art_loop_finger_handles_unnormalized_section():
    xyz = make_field_coords(size=(9, 9, 5))
    field = ArtLoopFinger(AxisSection(axis=(0.0, 0.0, 3.0)), rad=2.0)

    raw = field.call(xyz)
    values = field(xyz)

    assert raw.shape == xyz.shape
    assert values.shape == xyz.shape
    assert np.all(np.isfinite(raw))
    assert np.all(np.isfinite(values))
    assert_allclose(np.linalg.norm(values, axis=-1), 1.0, atol=1e-6)


def test_public_exports_drop_frame_and_shift_symbols():
    assert hasattr(tv.ansatz, "AnsatzND")
    assert hasattr(tv.ansatz, "MaskND")
    assert not hasattr(tv.ansatz, "Mask2D")
    assert not hasattr(tv.ansatz, "Mask3D")
    assert hasattr(tv.ansatz, "CholestericSpiral3D")
    assert hasattr(tv.ansatz, "CoordinateSpiral3D")
    assert hasattr(tv.ansatz, "SectionEmbedding3D")
    assert not hasattr(tv.ansatz, "ConstantMaskND")
    assert not hasattr(tv.ansatz, "CoordinateComponentMaskND")
    assert not hasattr(tv.ansatz, "RadialDistanceMaskND")
    assert not hasattr(tv.ansatz, "GaussianMaskND")
    assert not hasattr(tv.ansatz, "BoundaryDistanceMaskND")
    assert not hasattr(tv.ansatz, "BoundaryPowerMaskND")
    assert not hasattr(tv.ansatz, "SmoothStepMaskND")
    assert not hasattr(tv.ansatz, "Frame3D")
    assert not hasattr(tv.ansatz, "FrameSpiral3D")
    assert not hasattr(tv.ansatz, "FrameCoordinateMask3D")
    assert not hasattr(tv.ansatz, "Shift2D")
    assert not hasattr(tv.ansatz, "Shift3D")


def test_smoke_default_public_classes():
    x, z = make_section_coords()
    xyz = make_field_coords()

    AxisSection(axis=(0, 0, 1))(x, z)
    BlochSection()(x, z)
    SpiralSection()(x, z)
    CF1()(x, z)
    CF2()(x, z)
    CF3()(x, z)
    CF4()(x, z)
    Twist()(x, z)
    Hopfion()(xyz)
    SectionEmbedding3D(Twist(), section_axes=(0, 2))(xyz)
    StraightFinger(Twist())(xyz)
    LoopFinger(Twist())(xyz)
    ArtLoopFinger(Twist(), rad=2.0)(xyz)
    Twist().shift(x=0.5, z=0.5)(x, z)
    Hopfion().shift(x=0.5, y=0.5, z=0.5)(xyz)
    ConstantField3D((0, 0, 1))(xyz)
    CholestericSpiral3D(period=4.0)(xyz)
    (
        ConstantField3D((0, 0, 1))
        + CholestericSpiral3D(period=4.0)
        + ConstantField3D((0, 0, 1)) * 0.5
    )(xyz)


def test_cholesteric_spiral_alias_matches_legacy_name():
    xyz = make_field_coords(size=(6, 5, 4))
    new = CholestericSpiral3D(period=4.0)
    old = CoordinateSpiral3D(period=4.0)

    assert_allclose(new.call(xyz), old.call(xyz))
    assert_allclose(new(xyz), old(xyz))
