import numpy as np
from numpy.testing import assert_allclose

import topovec.ansatz.frame as frame_module
from topovec.ansatz import (
    AffineCoordinates,
    AxisymmetricCoordinates,
    CoordinateRemap,
    CoordinateSampleND,
    CylindricalCoordinates,
    EuclideanCoordinates,
    PolarCoordinates,
)


def finite_coordinate_jacobian(system, points, step=1e-6):
    points = np.asarray(points, dtype=np.float64)
    jacobian = np.zeros(points.shape[:-1] + (system.ndim, system.ndim), dtype=np.float64)
    for axis in range(system.ndim):
        delta = np.zeros(system.ndim, dtype=np.float64)
        delta[axis] = step
        plus = system.coords(points + delta)
        minus = system.coords(points - delta)
        jacobian[..., :, axis] = (plus - minus) / (2.0 * step)
    return jacobian


def assert_basis_is_inverse_of_coordinate_jacobian(system, points, atol=1e-5, rtol=1e-5):
    sample = system(points)
    valid = sample.valid
    jacobian = finite_coordinate_jacobian(system, points)
    inverse_identity = np.einsum("...ia,...ja->...ij", jacobian, sample.basis)
    expected = np.eye(system.ndim)
    assert np.any(valid)
    assert_allclose(
        inverse_identity[valid],
        np.broadcast_to(expected, inverse_identity[valid].shape),
        atol=atol,
        rtol=rtol,
    )


def test_coordinate_sample_shapes_and_validation():
    sample = CoordinateSampleND(
        coords=np.zeros((4, 3, 2)),
        basis=np.broadcast_to(np.eye(2), (4, 3, 2, 2)),
        valid=np.ones((4, 3), dtype=bool),
    )
    assert sample.ndim == 2
    assert sample.coords.shape == (4, 3, 2)
    assert sample.basis.shape == (4, 3, 2, 2)
    assert sample.valid.shape == (4, 3)


def test_euclidean_coordinates_shapes_and_basis():
    points2 = np.array([[[0.0, 1.0], [2.0, -1.0]]], dtype=np.float64)
    points3 = np.array([[[0.0, 1.0, 2.0], [2.0, -1.0, 0.5]]], dtype=np.float64)

    system2 = EuclideanCoordinates(ndim=2)
    system3 = EuclideanCoordinates(ndim=3)

    sample2 = system2(points2)
    sample3 = system3(points3)

    assert isinstance(system2.shift((1.0, 2.0)), AffineCoordinates)
    assert isinstance(system3.shift((1.0, 2.0, 3.0)), AffineCoordinates)

    assert_allclose(sample2.coords, points2)
    assert_allclose(sample3.coords, points3)
    assert_allclose(sample2.basis, np.broadcast_to(np.eye(2), sample2.basis.shape))
    assert_allclose(sample3.basis, np.broadcast_to(np.eye(3), sample3.basis.shape))
    assert np.all(sample2.valid)
    assert np.all(sample3.valid)


def test_shift_and_rotation_precompose_physical_space():
    points = np.array([[2.0, 3.0], [-1.0, 4.0]], dtype=np.float64)
    base = EuclideanCoordinates(ndim=2)

    shifted = base.shift((1.5, -2.0))
    rotated = base.rotate(angle=np.pi / 2)

    assert_allclose(shifted.coords(points), points - np.array([1.5, -2.0]))
    expected_rot = points @ np.array([[0.0, -1.0], [1.0, 0.0]])
    assert_allclose(rotated.coords(points), expected_rot, atol=1e-12)
    assert_basis_is_inverse_of_coordinate_jacobian(shifted, points)
    assert_basis_is_inverse_of_coordinate_jacobian(rotated, points)


def test_coordinate_system_mask_sugar_uses_curvilinear_coordinates():
    points2 = np.array([[2.0, 3.0], [-1.0, 4.0]], dtype=np.float64)
    shifted2 = EuclideanCoordinates(ndim=2).shift((1.5, -2.0))
    mask2 = shifted2.mask(lambda x, z: x - 2.0 * z)
    assert_allclose(mask2(points2), (points2[:, 0] - 1.5) - 2.0 * (points2[:, 1] + 2.0))

    points3 = np.array([[2.0, 3.0, 4.0], [-1.0, 4.0, 0.5]], dtype=np.float64)
    shifted3 = EuclideanCoordinates(ndim=3).shift((1.0, -2.0, 0.5))
    mask3 = shifted3.mask(lambda x, y, z: x + y - z)
    assert_allclose(mask3(points3), (points3[:, 0] - 1.0) + (points3[:, 1] + 2.0) - (points3[:, 2] - 0.5))


def test_linear_coordinate_mix_and_coordinate_remap():
    points = np.array([[1.0, 2.0], [3.0, -4.0]], dtype=np.float64)
    base = EuclideanCoordinates(ndim=2)

    matrix = np.array([[2.0, 1.0], [1.0, 1.0]], dtype=np.float64)
    mix = base.mix(matrix=matrix, offset=(0.5, -1.0))
    sample = mix(points)
    expected_coords = points @ matrix.T + np.array([0.5, -1.0])
    expected_basis = np.broadcast_to(np.linalg.inv(matrix).T, sample.basis.shape)
    assert_allclose(sample.coords, expected_coords)
    assert_allclose(sample.basis, expected_basis)
    assert_basis_is_inverse_of_coordinate_jacobian(mix, points)

    remapped = CoordinateRemap(
        base,
        map_fn=lambda u: np.stack([u[..., 0] + 2.0 * u[..., 1], u[..., 1]], axis=-1),
        inverse_jacobian_fn=lambda u, v: np.broadcast_to(
            np.array([[1.0, -2.0], [0.0, 1.0]], dtype=np.float64),
            u.shape[:-1] + (2, 2),
        ),
    )
    remapped_sample = remapped(points)
    assert_allclose(remapped_sample.coords, np.stack([points[:, 0] + 2.0 * points[:, 1], points[:, 1]], axis=-1))
    assert_allclose(
        remapped_sample.basis,
        np.broadcast_to(np.array([[1.0, 0.0], [-2.0, 1.0]], dtype=np.float64), remapped_sample.basis.shape),
    )
    assert_basis_is_inverse_of_coordinate_jacobian(remapped, points)


def test_cylindrical_coordinates_2d_contract_and_singularity():
    points = np.array([[3.0, 4.0], [0.0, 0.0]], dtype=np.float64)
    cylindrical = CylindricalCoordinates(
        EuclideanCoordinates(ndim=2),
        radius=0,
        angle=1,
        angle_min=0.0,
        angle_max=2.0 * np.pi,
        tol=1e-12,
    )
    sample = cylindrical(points)

    assert_allclose(sample.coords[0], np.array([5.0, np.arctan2(4.0, 3.0)]))
    assert_allclose(sample.basis[0, 0], np.array([3.0 / 5.0, 4.0 / 5.0]))
    assert_allclose(sample.basis[0, 1], np.array([-4.0, 3.0]))
    assert bool(sample.valid[0])
    assert not bool(sample.valid[1])
    assert np.all(np.isfinite(sample.coords))
    assert np.all(np.isfinite(sample.basis))
    assert_basis_is_inverse_of_coordinate_jacobian(cylindrical, points[:1])


def test_polar_coordinates_is_convenience_specialization():
    points = np.array([[3.0, 4.0], [5.0, -12.0]], dtype=np.float64)
    polar = PolarCoordinates(tol=1e-12)
    cylindrical = CylindricalCoordinates(
        EuclideanCoordinates(ndim=2),
        radius=0,
        angle=1,
        angle_min=0.0,
        angle_max=2.0 * np.pi,
        tol=1e-12,
    )
    polar_sample = polar(points)
    cylindrical_sample = cylindrical(points)

    assert_allclose(polar_sample.coords, cylindrical_sample.coords)
    assert_allclose(polar_sample.basis, cylindrical_sample.basis)
    assert np.array_equal(polar_sample.valid, cylindrical_sample.valid)


def test_cylindrical_coordinates_3d_only_transform_selected_plane():
    points = np.array([[3.0, 4.0, 7.0], [5.0, -12.0, -1.0]], dtype=np.float64)
    cylindrical = CylindricalCoordinates(
        EuclideanCoordinates(ndim=3),
        radius=0,
        angle=1,
        angle_min=-1.0,
        angle_max=3.0,
        tol=1e-12,
    )
    sample = cylindrical(points)

    radius = np.sqrt(points[:, 0] ** 2 + points[:, 1] ** 2)
    psi = np.mod(np.arctan2(points[:, 1], points[:, 0]), 2.0 * np.pi)
    expected_angle = -1.0 + 4.0 * psi / (2.0 * np.pi)
    assert_allclose(sample.coords[:, 0], radius)
    assert_allclose(sample.coords[:, 1], expected_angle)
    assert_allclose(sample.coords[:, 2], points[:, 2])
    assert_basis_is_inverse_of_coordinate_jacobian(cylindrical, points)


def test_axisymmetric_coordinates_lift_meridional_chart():
    points = np.array(
        [
            [3.0, 4.0, 2.0],
            [5.0, 12.0, -1.0],
            [0.0, 0.0, 7.0],
        ],
        dtype=np.float64,
    )
    system = EuclideanCoordinates(ndim=2).axisymmetrize(
        axis=2,
        angle_min=1.0,
        angle_max=5.0,
        tol=1e-12,
    )
    sample = system(points)

    rho = np.sqrt(points[:, 0] ** 2 + points[:, 1] ** 2)
    psi = np.mod(np.arctan2(points[:, 1], points[:, 0]), 2.0 * np.pi)
    expected_angle = 1.0 + 4.0 * psi / (2.0 * np.pi)
    expected_coords = np.stack([rho, points[:, 2], expected_angle], axis=-1)
    assert_allclose(sample.coords, expected_coords)
    assert_allclose(sample.basis[0, 0], np.array([3.0 / 5.0, 4.0 / 5.0, 0.0]))
    assert_allclose(sample.basis[0, 1], np.array([0.0, 0.0, 1.0]))
    assert_allclose(
        sample.basis[0, 2],
        (5.0 * (2.0 * np.pi / 4.0)) * np.array([-4.0 / 5.0, 3.0 / 5.0, 0.0]),
    )
    assert bool(sample.valid[0])
    assert not bool(sample.valid[2])
    assert_basis_is_inverse_of_coordinate_jacobian(system, points[:2])


def test_restricted_patch_and_valid_propagation():
    base = PolarCoordinates(tol=1e-12)
    patch = base.restrict(predicate=lambda points, sample: sample.coords[..., 1] <= np.pi)
    points = np.array([[1.0, 1.0], [1.0, -1.0], [0.0, 0.0]], dtype=np.float64)
    valid = patch.valid(points)
    assert valid.tolist() == [True, False, False]


def test_cone_like_coordinates_from_meridional_mix_smoke():
    alpha = 0.37
    meridional = EuclideanCoordinates(ndim=2).mix(
        matrix=np.array(
            [
                [np.sin(alpha), np.cos(alpha)],
                [np.cos(alpha), -np.sin(alpha)],
            ],
            dtype=np.float64,
        )
    )
    cone_like = meridional.axisymmetrize(axis=2, tol=1e-12)
    points = np.array([[3.0, 4.0, 2.0], [5.0, 12.0, -1.0]], dtype=np.float64)
    sample = cone_like(points)

    rho = np.sqrt(points[:, 0] ** 2 + points[:, 1] ** 2)
    expected_first = np.sin(alpha) * rho + np.cos(alpha) * points[:, 2]
    assert_allclose(sample.coords[:, 0], expected_first)
    assert np.all(sample.valid)
    assert np.all(np.isfinite(sample.basis))


def test_dimension_specific_methods_use_runtime_dimension_checks():
    system2 = EuclideanCoordinates(ndim=2)
    system3 = EuclideanCoordinates(ndim=3)
    system4 = EuclideanCoordinates(ndim=4)

    assert isinstance(system2.rotate(np.pi / 3), AffineCoordinates)
    assert isinstance(system3.rotate(np.pi / 3, axis=1), AffineCoordinates)
    assert isinstance(system2.axisymmetrize(axis=2), AxisymmetricCoordinates)

    try:
        system4.rotate(0.1)
    except NotImplementedError:
        pass
    else:
        raise AssertionError("rotate() should reject dimensions other than 2 and 3.")

    try:
        system3.axisymmetrize(axis=2)
    except NotImplementedError:
        pass
    else:
        raise AssertionError("axisymmetrize() should reject dimensions other than 2.")


def test_public_coordinate_docstrings_present():
    assert frame_module.__doc__ is not None and "local chart" in frame_module.__doc__
    for cls in (
        frame_module.CoordinateSystemND,
        frame_module.AffineCoordinates,
        frame_module.CoordinateRemap,
        frame_module.RestrictedCoordinates,
        frame_module.CylindricalCoordinates,
        frame_module.PolarCoordinates,
        frame_module.AxisymmetricCoordinates,
    ):
        assert cls.__doc__ is not None
        assert len(cls.__doc__.strip()) > 40
