"""Linear subdivision for simplicial meshes.

Linear subdivision is the simplest subdivision scheme: each edge is split at
its midpoint, and each n-simplex is divided into 2^n smaller simplices.
This is an interpolating scheme - original vertices remain unchanged.

Works for any manifold dimension and any spatial dimension (including higher
codimensions like curves in 3D or surfaces in 4D).
"""

from typing import TYPE_CHECKING

import torch

from torchmesh.subdivision._data import (
    interpolate_point_data_to_edges,
    propagate_cell_data_to_children,
)
from torchmesh.subdivision._topology import (
    extract_unique_edges,
    generate_child_cells,
    get_subdivision_pattern,
)

if TYPE_CHECKING:
    from torchmesh.mesh import Mesh


def subdivide_linear(mesh: "Mesh") -> "Mesh":
    """Perform one level of linear subdivision on the mesh.

    Linear subdivision splits each n-simplex into 2^n child simplices by:
    1. Adding new vertices at edge midpoints
    2. Connecting vertices according to a subdivision pattern

    This is an interpolating scheme: original vertices keep their positions,
    and new vertices are placed exactly at edge midpoints.

    Properties:
    - Preserves manifold dimension and spatial dimension
    - Increases mesh resolution uniformly
    - Point data is interpolated to new vertices (averaged from endpoints)
    - Cell data is propagated to children (each child inherits parent's data)
    - Global data is preserved unchanged

    Args:
        mesh: Input mesh to subdivide (any manifold/spatial dimension)

    Returns:
        Subdivided mesh with:
        - n_points = original_n_points + n_edges
        - n_cells = original_n_cells * 2^n_manifold_dims

    Example:
        >>> # Triangle mesh: 2 triangles -> 8 triangles
        >>> mesh = create_triangle_mesh()
        >>> subdivided = subdivide_linear(mesh)
        >>> assert subdivided.n_cells == mesh.n_cells * 4  # 2^2 for 2D

        >>> # Tetrahedral mesh: 1 tet -> 8 tets
        >>> tet_mesh = create_tet_mesh()
        >>> subdivided = subdivide_linear(tet_mesh)
        >>> assert subdivided.n_cells == tet_mesh.n_cells * 8  # 2^3 for 3D
    """
    from torchmesh.mesh import Mesh

    ### Handle empty mesh
    if mesh.n_cells == 0:
        return mesh

    ### Extract unique edges from mesh
    unique_edges, edge_inverse = extract_unique_edges(mesh)
    n_original_points = mesh.n_points

    ### Compute edge midpoints
    # Shape: (n_edges, n_spatial_dims)
    edge_vertices = mesh.points[unique_edges]  # (n_edges, 2, n_spatial_dims)
    edge_midpoints = edge_vertices.mean(dim=1)  # Average the two endpoints

    ### Create new points array: original + midpoints
    # Shape: (n_original_points + n_edges, n_spatial_dims)
    new_points = torch.cat([mesh.points, edge_midpoints], dim=0)

    ### Interpolate point_data to edge midpoints
    new_point_data = interpolate_point_data_to_edges(
        point_data=mesh.point_data,
        edges=unique_edges,
        n_original_points=n_original_points,
    )

    ### Get subdivision pattern for this manifold dimension
    subdivision_pattern = get_subdivision_pattern(mesh.n_manifold_dims)
    subdivision_pattern = subdivision_pattern.to(mesh.cells.device)

    ### Generate child cells from parents
    child_cells, parent_indices = generate_child_cells(
        parent_cells=mesh.cells,
        edge_inverse=edge_inverse,
        n_original_points=n_original_points,
        subdivision_pattern=subdivision_pattern,
    )

    ### Propagate cell_data from parents to children
    new_cell_data = propagate_cell_data_to_children(
        cell_data=mesh.cell_data,
        parent_indices=parent_indices,
        n_total_children=len(child_cells),
    )

    ### Create and return subdivided mesh
    return Mesh(
        points=new_points,
        cells=child_cells,
        point_data=new_point_data,
        cell_data=new_cell_data,
        global_data=mesh.global_data,  # Preserved unchanged
    )
