"""Quadrilateral surface meshers for closed cylinder geometries.

A cylinder is assembled from three parts:

1. **Bottom cap** — disk mesh placed at ``z = -height/2``, normals down.
2. **Top cap**    — same disk mesh with reversed normals at ``z = +height/2``.
3. **Lateral surface** — generated by revolving the side-wall curve around
   the Z-axis with :func:`~surfmesh.operations.revolve.circular_revolve`.

Two cap variants are provided, matching the two disk strategies:

* :func:`cylinder_mesher_radial` — purely polar caps.
* :func:`cylinder_mesher_square_centered` — square-centred caps for
  better element quality.
"""

from __future__ import annotations

import numpy as np

from surfmesh.core.validate import validate_positive, validate_positive_int
from surfmesh.operations.revolve import circular_revolve
from surfmesh.operations.transform import convert_2d_face_to_3d
from surfmesh.primitives.disk import disk_mesher_radial, disk_mesher_square_centered

__all__ = [
    "cylinder_mesher_radial",
    "cylinder_mesher_square_centered",
]


def cylinder_mesher_radial(
    radius: float,
    height: float,
    radial_resolution: int,
    segment_resolution: int,
    height_resolution: int,
) -> np.ndarray:
    """Generate a closed cylinder mesh with radial disk caps.

    Parameters
    ----------
    radius:
        Cylinder radius.  Must be positive.
    height:
        Cylinder height (total, symmetric about z = 0).  Must be positive.
    radial_resolution:
        Number of radial quad rings on each disk cap.  Must be >= 1.
    segment_resolution:
        Number of angular segments around the cylinder.  Must be >= 1.
    height_resolution:
        Number of quad layers along the cylinder axis.  Must be >= 1.

    Returns
    -------
    ndarray of shape ``(n_panels, 4, 3)``
        Closed cylinder mesh: bottom cap + top cap + lateral surface.

    Raises
    ------
    ValueError
        If any parameter violates its constraint.

    Examples
    --------
    >>> from surfmesh.primitives.cylinder import cylinder_mesher_radial
    >>> mesh = cylinder_mesher_radial(
    ...     radius=1.0, height=2.0,
    ...     radial_resolution=4, segment_resolution=8, height_resolution=4,
    ... )
    >>> mesh.shape
    (96, 4, 3)
    """
    validate_positive(radius, "radius")
    validate_positive(height, "height")
    validate_positive_int(radial_resolution, "radial_resolution")
    validate_positive_int(segment_resolution, "segment_resolution")
    validate_positive_int(height_resolution, "height_resolution")

    # Disk caps (2-D), then lifted into 3-D.
    disk_2d = disk_mesher_radial(radius, radial_resolution, segment_resolution)
    disk_2d_flipped = np.flip(disk_2d, axis=1)

    top_cap    = convert_2d_face_to_3d(disk_2d,        axis=2, offset= height / 2.0)
    bottom_cap = convert_2d_face_to_3d(disk_2d_flipped, axis=2, offset=-height / 2.0)

    # Lateral surface: revolve a vertical line segment.
    h_coords = height / 2.0 * np.linspace(-1.0, 1.0, height_resolution + 1)
    r_coords = radius * np.ones_like(h_coords)
    lateral_curve = np.stack([r_coords, h_coords], axis=1)
    lateral = circular_revolve(lateral_curve, segment_resolution)

    return np.vstack([bottom_cap, top_cap, lateral])


def cylinder_mesher_square_centered(
    radius: float,
    height: float,
    radial_resolution: int,
    half_square_resolution: int,
    height_resolution: int,
) -> np.ndarray:
    """Generate a closed cylinder mesh with square-centred disk caps.

    The disk caps use :func:`~surfmesh.primitives.disk.disk_mesher_square_centered`
    so the central region has well-conditioned quad elements.  The lateral
    surface is identical to :func:`cylinder_mesher_radial`.

    Parameters
    ----------
    radius:
        Cylinder radius.  Must be positive.
    height:
        Cylinder height (total, symmetric about z = 0).  Must be positive.
    radial_resolution:
        Number of radial transition layers on each cap.  Must be >= 1.
    half_square_resolution:
        Number of quad cells along **half** a square-core side edge.
        The full square resolution is ``2 * half_square_resolution``.
        Must be >= 1.
    height_resolution:
        Number of quad layers along the cylinder axis.  Must be >= 1.

    Returns
    -------
    ndarray of shape ``(n_panels, 4, 3)``
        Closed cylinder mesh: bottom cap + top cap + lateral surface.

    Raises
    ------
    ValueError
        If any parameter violates its constraint.

    Examples
    --------
    >>> from surfmesh.primitives.cylinder import cylinder_mesher_square_centered
    >>> mesh = cylinder_mesher_square_centered(
    ...     radius=1.0, height=2.0,
    ...     radial_resolution=4, half_square_resolution=4, height_resolution=4,
    ... )
    >>> mesh.shape
    (512, 4, 3)
    """
    validate_positive(radius, "radius")
    validate_positive(height, "height")
    validate_positive_int(radial_resolution, "radial_resolution")
    validate_positive_int(half_square_resolution, "half_square_resolution")
    validate_positive_int(height_resolution, "height_resolution")

    square_resolution = 2 * half_square_resolution

    # Disk caps (2-D), then lifted into 3-D.
    disk_2d = disk_mesher_square_centered(radius, square_resolution, radial_resolution)
    disk_2d_flipped = np.flip(disk_2d, axis=1)

    top_cap    = convert_2d_face_to_3d(disk_2d,         axis=2, offset= height / 2.0)
    bottom_cap = convert_2d_face_to_3d(disk_2d_flipped, axis=2, offset=-height / 2.0)

    # Lateral surface: number of segments must match the square-centred circumference.
    segment_resolution = square_resolution * 4
    h_coords = height / 2.0 * np.linspace(-1.0, 1.0, height_resolution + 1)
    r_coords = radius * np.ones_like(h_coords)
    lateral_curve = np.stack([r_coords, h_coords], axis=1)
    lateral = circular_revolve(lateral_curve, segment_resolution)

    return np.vstack([bottom_cap, top_cap, lateral])
