"""
Module: janssen.simul.apertures
-------------------------------
Aperture and apodizer elements for shaping optical wavefronts.

Functions
---------
circular_aperture
    Applies a circular aperture (optionally offset) with uniform transmittivity.
rectangular_aperture
    Applies an axis-aligned rectangular aperture with uniform transmittivity.
annular_aperture
    Applies a concentric ring (donut) aperture between inner/outer diameters.
variable_transmission_aperture
    Applies an arbitrary transmission mask (array or callable), including
    common apodizers such as Gaussian or super-Gaussian.
gaussian_apodizer
    Applies a Gaussian apodizer (smooth transmission mask) to the wavefront.
supergaussian_apodizer
    Applies a super-Gaussian apodizer (smooth transmission mask) to the wavefront.
gaussian_apodizer_elliptical
    Applies a Gaussian apodizer (smooth transmission mask) to the wavefront.
supergaussian_apodizer_elliptical
    Applies a super-Gaussian apodizer (smooth transmission mask) to the wavefront.
"""

import jax
import jax.numpy as jnp
from beartype.typing import Optional, Tuple, Union
from jaxtyping import Array, Bool, Float

from janssen.utils import (
    OpticalWavefront,
    beartype,
    jaxtyped,
    make_optical_wavefront,
    scalar_float,
    scalar_numeric,
)

jax.config.update("jax_enable_x64", True)


@jaxtyped(typechecker=beartype)
def _xy_grids(
    nx: int, ny: int, dx: float
) -> Tuple[Float[Array, " H W"], Float[Array, " H W"]]:
    """
    Internal helper to create centered spatial coordinate grids (in meters).

    Parameters
    ----------
    nx : int
        Number of grid points along x-axis.
    ny : int
        Number of grid points along y-axis.
    dx : float
        Grid spacing in meters.

    Returns
    -------
    xx : Float[Array, "H W"]
        X coordinate grid in meters.
    yy : Float[Array, "H W"]
        Y coordinate grid in meters.
    """
    x: Float[Array, " W"] = jnp.arange(-nx // 2, nx // 2) * dx
    y: Float[Array, " H"] = jnp.arange(-ny // 2, ny // 2) * dx
    xx: Float[Array, "H W"]
    yy: Float[Array, "H W"]
    xx, yy = jnp.meshgrid(x, y)
    return (xx, yy)


@jaxtyped(typechecker=beartype)
def circular_aperture(
    incoming: OpticalWavefront,
    diameter: scalar_float,
    center: Optional[Float[Array, " 2"]] = jnp.zeros(2),
    transmittivity: Optional[scalar_float] = 1.0,
) -> OpticalWavefront:
    """
    Apply a circular aperture to the incoming wavefront.

    The aperture is defined by its physical diameter and (optional) center.

    Parameters
    ----------
    incoming : OpticalWavefront
        Input wavefront PyTree.
    diameter : scalar_float
        Aperture diameter in meters.
    center : Optional[Float[Array, " 2"]], optional
        Physical center [x0, y0] of the aperture in meters, by default [0, 0].
    transmittivity : Optional[scalar_float], optional
        Uniform transmittivity inside the aperture (0..1), by default 1.0.

    Returns
    -------
    apertured : OpticalWavefront
        Wavefront after applying the circular aperture.

    Notes
    -----
    - Build centered (x, y) grids in meters.
    - Compute radial distance from the specified center.
    - Create a binary mask for r <= diameter/2.
    - Multiply by transmittivity (clipped to [0, 1]).
    - Apply to the complex field and return.
    """
    ny: int = incoming.field.shape[0]
    nx: int = incoming.field.shape[1]
    xx, yy = _xy_grids(nx, ny, float(incoming.dx))
    x0, y0 = center[0], center[1]
    r: Float[Array, " H W"] = jnp.sqrt((xx - x0) ** 2 + (yy - y0) ** 2)
    inside: Bool[Array, " H W"] = r <= (diameter / 2.0)
    t = jnp.clip(jnp.asarray(transmittivity, dtype=float), 0.0, 1.0)
    transmission: Float[Array, " H W"] = inside.astype(float) * t
    apertured: OpticalWavefront = make_optical_wavefront(
        field=incoming.field * transmission,
        wavelength=incoming.wavelength,
        dx=incoming.dx,
        z_position=incoming.z_position,
    )
    return apertured


@jaxtyped(typechecker=beartype)
def rectangular_aperture(
    incoming: OpticalWavefront,
    width: scalar_float,
    height: scalar_float,
    center: Optional[Float[Array, " 2"]] = jnp.zeros(2),
    transmittivity: Optional[scalar_float] = 1.0,
) -> OpticalWavefront:
    """
    Apply an axis-aligned rectangular aperture to the incoming wavefront.

    Parameters
    ----------
    incoming : OpticalWavefront
        Input wavefront PyTree.
    width : scalar_float
        Rectangle width along x in meters.
    height : scalar_float
        Rectangle height along y in meters.
    center : Optional[Float[Array, " 2"]], optional
        Rectangle center [x0, y0] in meters, by default [0, 0].
    transmittivity : Optional[scalar_float], optional
        Uniform transmittivity inside the rectangle (0..1), by default 1.0.

    Returns
    -------
    apertured : OpticalWavefront
        Wavefront after applying the rectangular aperture.

    Notes
    -----
    - Build centered (x, y) grids in meters.
    - Compute half-width/half-height and an inside-rectangle mask.
    - Multiply by transmittivity (clipped).
    - Apply to the complex field and return.
    """
    ny: int = incoming.field.shape[0]
    nx: int = incoming.field.shape[1]
    xx, yy = _xy_grids(nx, ny, float(incoming.dx))
    x0, y0 = center[0], center[1]
    hx = width / 2.0
    hy = height / 2.0
    inside_x: Bool[Array, " H W"] = ((x0 - hx) <= xx) & ((x0 + hx) >= xx)
    inside_y: Bool[Array, " H W"] = ((y0 - hy) <= yy) & ((y0 + hy) >= yy)
    inside: Bool[Array, " H W"] = inside_x & inside_y
    t = jnp.clip(jnp.asarray(transmittivity, dtype=float), 0.0, 1.0)
    transmission: Float[Array, " H W"] = inside.astype(float) * t
    apertured: OpticalWavefront = make_optical_wavefront(
        field=incoming.field * transmission,
        wavelength=incoming.wavelength,
        dx=incoming.dx,
        z_position=incoming.z_position,
    )
    return apertured


@jaxtyped(typechecker=beartype)
def annular_aperture(
    incoming: OpticalWavefront,
    inner_diameter: scalar_float,
    outer_diameter: scalar_float,
    center: Optional[Float[Array, " 2"]] = jnp.zeros(2),
    transmittivity: Optional[scalar_float] = 1.0,
) -> OpticalWavefront:
    """
    Apply an annular (ring) aperture with inner and outer diameters.

    Parameters
    ----------
    incoming : OpticalWavefront
        Input wavefront PyTree.
    inner_diameter : scalar_float
        Inner blocked diameter in meters.
    outer_diameter : scalar_float
        Outer clear aperture diameter in meters.
    center : Optional[Float[Array, " 2"]], optional
        Ring center [x0, y0] in meters, by default [0, 0].
    transmittivity : Optional[scalar_float], optional
        Uniform transmittivity in the ring (0..1), by default 1.0.

    Returns
    -------
    apertured : OpticalWavefront
        Wavefront after applying the annular aperture.

    Notes
    -----
    - Build centered (x, y) grids in meters.
    - Compute radial distance from center.
    - Create mask for inner_radius < r <= outer_radius.
    - Multiply by transmittivity (clipped), apply, and return.
    """
    ny: int = incoming.field.shape[0]
    nx: int = incoming.field.shape[1]
    xx, yy = _xy_grids(nx, ny, float(incoming.dx))
    x0, y0 = center[0], center[1]
    r: Float[Array, " H W"] = jnp.sqrt((xx - x0) ** 2 + (yy - y0) ** 2)
    r_in = inner_diameter / 2.0
    r_out = outer_diameter / 2.0
    ring: Bool[Array, " H W"] = (r > r_in) & (r <= r_out)
    t = jnp.clip(jnp.asarray(transmittivity, dtype=float), 0.0, 1.0)
    transmission: Float[Array, " H W"] = ring.astype(float) * t
    apertured: OpticalWavefront = make_optical_wavefront(
        field=incoming.field * transmission,
        wavelength=incoming.wavelength,
        dx=incoming.dx,
        z_position=incoming.z_position,
    )
    return apertured


@jaxtyped(typechecker=beartype)
def variable_transmission_aperture(
    incoming: OpticalWavefront,
    transmission: Union[scalar_float, Float[Array, " ..."]],
) -> OpticalWavefront:
    """
    Apply an arbitrary (spatially varying) transmission to the wavefront.

    Parameters
    ----------
    incoming : OpticalWavefront
        Input wavefront PyTree.
    transmission : Union[scalar_float, Float[Array, " H W"]]
        Precomputed transmission map (0..1) with shape "H W", or a scalar
        attenuation factor for uniform transmission.

    Returns
    -------
    apertured : OpticalWavefront
        Wavefront after applying the transmission.

    Examples
    --------
    Uniform attenuation::

        >>> wf2 = variable_transmission_aperture(wf, 0.5)  # 50% transmission

    Spatially varying transmission::

        >>> tmap = create_transmission_map(...)  # Shape (H, W)
        >>> wf2 = variable_transmission_aperture(wf, tmap)

    Notes
    -----
    - For scalar transmission: applies uniform attenuation.
    - For array transmission: applies spatially varying transmission map.
    - Transmission values are clipped to [0, 1].
    - This function is fully JAX-compatible and uses jax.lax.cond.
    """
    trans = jnp.asarray(transmission, dtype=float)

    def apply_scalar_transmission() -> OpticalWavefront:
        t = jnp.clip(trans, 0.0, 1.0)
        return make_optical_wavefront(
            field=incoming.field * t,
            wavelength=incoming.wavelength,
            dx=incoming.dx,
            z_position=incoming.z_position,
        )

    def apply_array_transmission() -> OpticalWavefront:
        tmap = jnp.clip(trans, 0.0, 1.0)
        return make_optical_wavefront(
            field=incoming.field * tmap,
            wavelength=incoming.wavelength,
            dx=incoming.dx,
            z_position=incoming.z_position,
        )

    return jax.lax.cond(
        trans.ndim == 0, apply_scalar_transmission, apply_array_transmission
    )


@jaxtyped(typechecker=beartype)
def gaussian_apodizer(
    incoming: OpticalWavefront,
    sigma: scalar_float,
    center: Optional[Float[Array, " 2"]] = jnp.zeros(2),
    peak_transmittivity: Optional[scalar_float] = 1.0,
) -> OpticalWavefront:
    """
    Apply a Gaussian apodizer (smooth transmission mask) to the wavefront.

    Parameters
    ----------
    incoming : OpticalWavefront
        Input optical wavefront.
    sigma : scalar_float
        Gaussian width parameter in meters.
    center : Optional[Float[Array, " 2"]], optional
        Physical center [x0, y0] of the Gaussian in meters, by default [0, 0].
    peak_transmittivity : Optional[scalar_float], optional
        Maximum transmission at the Gaussian center, by default 1.0.

    Returns
    -------
    apertured : OpticalWavefront
        Wavefront after applying Gaussian apodization.

    Notes
    -----
    - Build centered (x, y) grids.
    - Compute squared radial distance from center.
    - Evaluate Gaussian exp(-r^2 / (2*sigma^2)).
    - Scale by peak transmittivity, clip to [0,1].
    - Multiply with incoming field and return.
    """
    ny: int = incoming.field.shape[0]
    nx: int = incoming.field.shape[1]
    xx, yy = _xy_grids(nx, ny, float(incoming.dx))
    x0, y0 = center[0], center[1]
    r2: Float[Array, " H W"] = (xx - x0) ** 2 + (yy - y0) ** 2
    gauss: Float[Array, " H W"] = jnp.exp(-r2 / (2.0 * sigma**2))
    tmap: Float[Array, " H W"] = jnp.clip(gauss * peak_transmittivity, 0.0, 1.0)
    apertured: OpticalWavefront = make_optical_wavefront(
        field=incoming.field * tmap,
        wavelength=incoming.wavelength,
        dx=incoming.dx,
        z_position=incoming.z_position,
    )
    return apertured


@jaxtyped(typechecker=beartype)
def supergaussian_apodizer(
    incoming: OpticalWavefront,
    sigma: scalar_float,
    m: scalar_numeric,
    center: Optional[Float[Array, " 2"]] = jnp.zeros(2),
    peak_transmittivity: Optional[scalar_float] = 1.0,
) -> OpticalWavefront:
    """
    Apply a super-Gaussian apodizer to the wavefront.

    Transmission profile: exp(- (r^2 / sigma^2)^m ).

    Parameters
    ----------
    incoming : OpticalWavefront
        Input optical wavefront.
    sigma : scalar_float
        Width parameter in meters (sets the roll-off scale).
    m : scalar_numeric
        Super-Gaussian order (m=1 → Gaussian, m>1 → flatter top).
    center : Optional[Float[Array, " 2"]], optional
        Physical center [x0, y0] of the profile, by default [0, 0].
    peak_transmittivity : Optional[scalar_float], optional
        Maximum transmission at the center, by default 1.0.

    Returns
    -------
    apertured : OpticalWavefront
        Wavefront after applying super-Gaussian apodization.

    Notes
    -----
    - Build centered (x, y) grids.
    - Compute squared radial distance from center.
    - Evaluate exp(- (r^2 / sigma^2)^m ).
    - Scale by peak transmittivity, clip to [0,1].
    - Multiply with incoming field and return.
    """
    ny: int = incoming.field.shape[0]
    nx: int = incoming.field.shape[1]
    xx, yy = _xy_grids(nx, ny, float(incoming.dx))
    x0, y0 = center[0], center[1]
    r2: Float[Array, " H W"] = (xx - x0) ** 2 + (yy - y0) ** 2
    super_gauss: Float[Array, " H W"] = jnp.exp(-((r2 / (sigma**2)) ** m))
    tmap: Float[Array, " H W"] = jnp.clip(super_gauss * peak_transmittivity, 0.0, 1.0)
    apertured: OpticalWavefront = make_optical_wavefront(
        field=incoming.field * tmap,
        wavelength=incoming.wavelength,
        dx=incoming.dx,
        z_position=incoming.z_position,
    )
    return apertured


@jaxtyped(typechecker=beartype)
def gaussian_apodizer_elliptical(
    incoming: OpticalWavefront,
    sigma_x: scalar_float,
    sigma_y: scalar_float,
    theta: Optional[scalar_float] = 0.0,
    center: Optional[Float[Array, " 2"]] = jnp.zeros(2),
    peak_transmittivity: Optional[scalar_float] = 1.0,
) -> OpticalWavefront:
    """
    Apply an elliptical Gaussian apodizer to the wavefront, with optional rotation.

    Parameters
    ----------
    incoming : OpticalWavefront
        Input optical wavefront.
    sigma_x : scalar_float
        Gaussian width along the x'-axis (meters) after rotation by `theta`.
    sigma_y : scalar_float
        Gaussian width along the y'-axis (meters) after rotation by `theta`.
    theta : Optional[scalar_float], optional
        Rotation angle in radians (counter-clockwise), by default 0.0.
    center : Optional[Float[Array, " 2"]], optional
        Physical center [x0, y0] in meters, by default [0, 0].
    peak_transmittivity : Optional[scalar_float], optional
        Maximum transmission at the center, by default 1.0.

    Returns
    -------
    apertured : OpticalWavefront
        Wavefront after applying elliptical Gaussian apodization.

    Notes
    -----
    - Build centered (x, y) grids.
    - Translate by `center`, rotate by `theta` → (x', y').
    - Evaluate exp(-0.5 * ( (x'/sigma_x)^2 + (y'/sigma_y)^2 )).
    - Scale by `peak_transmittivity`, clip to [0, 1].
    - Multiply with incoming field and return.
    """
    ny: int = incoming.field.shape[0]
    nx: int = incoming.field.shape[1]
    xx, yy = _xy_grids(nx, ny, float(incoming.dx))
    x0, y0 = center[0], center[1]
    xc = xx - x0
    yc = yy - y0
    ct = jnp.cos(theta)
    st = jnp.sin(theta)
    xp = ct * xc + st * yc
    yp = -st * xc + ct * yc
    arg = (xp / sigma_x) ** 2 + (yp / sigma_y) ** 2
    gauss = jnp.exp(-0.5 * arg)
    tmap = jnp.clip(gauss * peak_transmittivity, 0.0, 1.0)
    apertured: OpticalWavefront = make_optical_wavefront(
        field=incoming.field * tmap,
        wavelength=incoming.wavelength,
        dx=incoming.dx,
        z_position=incoming.z_position,
    )
    return apertured


@jaxtyped(typechecker=beartype)
def supergaussian_apodizer_elliptical(
    incoming: OpticalWavefront,
    sigma_x: scalar_float,
    sigma_y: scalar_float,
    m: scalar_numeric,
    theta: Optional[scalar_float] = 0.0,
    center: Optional[Float[Array, " 2"]] = jnp.zeros(2),
    peak_transmittivity: Optional[scalar_float] = 1.0,
) -> OpticalWavefront:
    """
    Apply an elliptical super-Gaussian apodizer with optional rotation.

    Transmission profile: exp( - ( (x'/sigma_x)^2 + (y'/sigma_y)^2 )^m ).

    Parameters
    ----------
    incoming : OpticalWavefront
        Input optical wavefront.
    sigma_x : scalar_float
        Width along x' (meters) after rotation by `theta`.
    sigma_y : scalar_float
        Width along y' (meters) after rotation by `theta`.
    m : scalar_numeric
        Super-Gaussian order (m=1 → Gaussian; m>1 → flatter top, sharper edges).
    theta : Optional[scalar_float], optional
        Rotation angle in radians (counter-clockwise), by default 0.0.
    center : Optional[Float[Array, " 2"]], optional
        Physical center [x0, y0] in meters, by default [0, 0].
    peak_transmittivity : Optional[scalar_float], optional
        Maximum transmission at the center, by default 1.0.

    Returns
    -------
    apertured : OpticalWavefront
        Wavefront after applying elliptical super-Gaussian apodization.

    Notes
    -----
    - Build centered (x, y) grids.
    - Translate by `center`, rotate by `theta` → (x', y').
    - Evaluate exp( - ( (x'/sigma_x)^2 + (y'/sigma_y)^2 )^m ).
    - Scale by `peak_transmittivity`, clip to [0, 1].
    - Multiply with incoming field and return.
    """
    ny: int = incoming.field.shape[0]
    nx: int = incoming.field.shape[1]
    xx, yy = _xy_grids(nx, ny, float(incoming.dx))
    x0, y0 = center[0], center[1]
    xc = xx - x0
    yc = yy - y0
    ct = jnp.cos(theta)
    st = jnp.sin(theta)
    xp = ct * xc + st * yc
    yp = -st * xc + ct * yc
    base = (xp / sigma_x) ** 2 + (yp / sigma_y) ** 2
    super_gauss = jnp.exp(-(base**m))
    tmap = jnp.clip(super_gauss * peak_transmittivity, 0.0, 1.0)
    apertured: OpticalWavefront = make_optical_wavefront(
        field=incoming.field * tmap,
        wavelength=incoming.wavelength,
        dx=incoming.dx,
        z_position=incoming.z_position,
    )
    return apertured
