"""Flat plane in 3D space.

Dimensional: 2D manifold in 3D space (has boundary).
"""

import torch

from torchmesh.mesh import Mesh


def load(
    size: float = 2.0,
    n_subdivisions: int = 10,
    normal: tuple[float, float, float] = (0.0, 0.0, 1.0),
    device: str = "cpu",
) -> Mesh:
    """Create a flat triangulated plane in 3D space.

    Args:
        size: Size of the plane (length of each side)
        n_subdivisions: Number of subdivisions per edge
        normal: Normal vector to the plane (will be normalized)
        device: Compute device ('cpu' or 'cuda')

    Returns:
        Mesh with n_manifold_dims=2, n_spatial_dims=3
    """
    if n_subdivisions < 1:
        raise ValueError(f"n_subdivisions must be at least 1, got {n_subdivisions=}")

    n = n_subdivisions + 1

    # Create grid of points in xy-plane
    x = torch.linspace(-size / 2, size / 2, n, device=device)
    y = torch.linspace(-size / 2, size / 2, n, device=device)
    xx, yy = torch.meshgrid(x, y, indexing="ij")

    points_flat = torch.stack(
        [xx.flatten(), yy.flatten(), torch.zeros_like(xx.flatten())], dim=1
    )

    # Rotate to align with normal if not (0, 0, 1)
    normal_t = torch.tensor(normal, dtype=torch.float32, device=device)
    normal_t = normal_t / torch.norm(normal_t)

    if not torch.allclose(normal_t, torch.tensor([0.0, 0.0, 1.0], device=device)):
        # Find rotation axis and angle
        z_axis = torch.tensor([0.0, 0.0, 1.0], dtype=torch.float32, device=device)
        axis = torch.cross(z_axis, normal_t)
        axis_norm = torch.norm(axis)

        if axis_norm > 1e-6:
            axis = axis / axis_norm
            angle = torch.acos(torch.dot(z_axis, normal_t))

            # Rodrigues' rotation formula
            K = torch.tensor(
                [
                    [0, -axis[2], axis[1]],
                    [axis[2], 0, -axis[0]],
                    [-axis[1], axis[0], 0],
                ],
                dtype=torch.float32,
                device=device,
            )
            R = (
                torch.eye(3, device=device)
                + torch.sin(angle) * K
                + (1 - torch.cos(angle)) * torch.mm(K, K)
            )
            points = torch.mm(points_flat, R.T)
        else:
            points = points_flat
    else:
        points = points_flat

    # Create triangular cells
    cells = []
    for i in range(n_subdivisions):
        for j in range(n_subdivisions):
            idx = i * n + j
            # Two triangles per quad
            cells.append([idx, idx + 1, idx + n])
            cells.append([idx + 1, idx + n + 1, idx + n])

    cells = torch.tensor(cells, dtype=torch.int64, device=device)
    return Mesh(points=points, cells=cells)
