"""Factory functions for creating data structures.

Extended Summary
----------------
Factory functions for creating data structures with runtime type
checking.
All runtime validations use JAX safe conditional statements.

Routine Listings
----------------
make_lens_params : function
    Creates a LensParams instance with runtime type checking
make_grid_params : function
    Creates a GridParams instance with runtime type checking
make_optical_wavefront : function
    Creates an OpticalWavefront instance with runtime type checking
make_propagating_wavefront : function
    Creates a PropagatingWavefront instance with runtime type checking
optical2propagating : function
    Creates a PropagatingWavefront from a tuple of OpticalWavefronts
make_microscope_data : function
    Creates a MicroscopeData instance with runtime type checking
make_diffractogram : function
    Creates a Diffractogram instance with runtime type checking
make_sample_function : function
    Creates a SampleFunction instance with runtime type checking
make_sliced_material_function : function
    Creates a SlicedMaterialFunction instance with runtime type checking
make_optimizer_state : function
    Creates an OptimizerState instance with runtime type checking
make_ptychography_params : function
    Creates a PtychographyParams instance with runtime type checking
make_ptychography_reconstruction : function
    Creates a PtychographyReconstruction instance with runtime type checking

Notes
-----
Always use these factory functions instead of directly instantiating the
NamedTuple classes to ensure proper runtime type checking of the
contents.
"""

import jax.numpy as jnp
from beartype import beartype
from beartype.typing import Optional, Tuple, Union
from jax import lax
from jaxtyping import Array, Bool, Complex, Float, Int, Num, jaxtyped

from .common_types import (
    Diffractogram,
    EpieData,
    EpieParams,
    GridParams,
    LensParams,
    MicroscopeData,
    OpticalWavefront,
    OptimizerState,
    PropagatingWavefront,
    PtychographyParams,
    PtychographyReconstruction,
    SampleFunction,
    ScalarBool,
    ScalarComplex,
    ScalarFloat,
    ScalarInteger,
    ScalarNumeric,
    SlicedMaterialFunction,
)


@jaxtyped(typechecker=beartype)
def make_lens_params(
    focal_length: ScalarNumeric,
    diameter: ScalarNumeric,
    n: ScalarNumeric,
    center_thickness: ScalarNumeric,
    r1: ScalarNumeric,
    r2: ScalarNumeric,
) -> LensParams:
    """JAX-safe factory function for LensParams with data validation.

    Parameters
    ----------
    focal_length : ScalarNumeric
        Focal length of the lens in meters
    diameter : ScalarNumeric
        Diameter of the lens in meters
    n : ScalarNumeric
        Refractive index of the lens material
    center_thickness : ScalarNumeric
        Thickness at the center of the lens in meters
    r1 : ScalarNumeric
        Radius of curvature of the first surface in meters
        (positive for convex)
    r2 : ScalarNumeric
        Radius of curvature of the second surface in meters
        (positive for convex)

    Returns
    -------
    validated_lens_params : LensParams
        Validated lens parameters instance

    Raises
    ------
    ValueError
        If parameters are invalid or out of valid ranges

    Notes
    -----
    Algorithm:

    - Convert inputs to JAX arrays
    - Validate parameters:
        - Check focal_length is positive
        - Check diameter is positive
        - Check refractive index is positive
        - Check center_thickness is positive
        - Check radii are finite
    - Create and return LensParams instance
    """
    focal_length = jnp.asarray(focal_length, dtype=jnp.float64)
    diameter = jnp.asarray(diameter, dtype=jnp.float64)
    n = jnp.asarray(n, dtype=jnp.float64)
    center_thickness = jnp.asarray(center_thickness, dtype=jnp.float64)
    r1 = jnp.asarray(r1, dtype=jnp.float64)
    r2 = jnp.asarray(r2, dtype=jnp.float64)

    def validate_and_create() -> LensParams:
        def check_focal_length() -> Float[Array, " "]:
            return lax.cond(
                focal_length > 0,
                lambda: focal_length,
                lambda: lax.stop_gradient(
                    lax.cond(False, lambda: focal_length, lambda: focal_length)
                ),
            )

        def check_diameter() -> Float[Array, " "]:
            return lax.cond(
                diameter > 0,
                lambda: diameter,
                lambda: lax.stop_gradient(
                    lax.cond(False, lambda: diameter, lambda: diameter)
                ),
            )

        def check_refractive_index() -> Float[Array, " "]:
            return lax.cond(
                n > 0,
                lambda: n,
                lambda: lax.stop_gradient(
                    lax.cond(False, lambda: n, lambda: n)
                ),
            )

        def check_center_thickness() -> Float[Array, " "]:
            return lax.cond(
                center_thickness > 0,
                lambda: center_thickness,
                lambda: lax.stop_gradient(
                    lax.cond(
                        False,
                        lambda: center_thickness,
                        lambda: center_thickness,
                    )
                ),
            )

        def check_radii_finite() -> (
            Tuple[Float[Array, " "], Float[Array, " "]]
        ):
            return lax.cond(
                jnp.logical_and(jnp.isfinite(r1), jnp.isfinite(r2)),
                lambda: (r1, r2),
                lambda: lax.stop_gradient(
                    lax.cond(False, lambda: (r1, r2), lambda: (r1, r2))
                ),
            )

        check_focal_length()
        check_diameter()
        check_refractive_index()
        check_center_thickness()
        check_radii_finite()

        return LensParams(
            focal_length=focal_length,
            diameter=diameter,
            n=n,
            center_thickness=center_thickness,
            r1=r1,
            r2=r2,
        )

    validated_lens_params: LensParams = validate_and_create()
    return validated_lens_params


@jaxtyped(typechecker=beartype)
def make_grid_params(
    xx: Float[Array, " hh ww"],
    yy: Float[Array, " hh ww"],
    phase_profile: Float[Array, " hh ww"],
    transmission: Float[Array, " hh ww"],
) -> GridParams:
    """JAX-safe factory function for GridParams with data validation.

    Parameters
    ----------
    xx : Float[Array, " hh ww"]
        Spatial grid in the x-direction
    yy : Float[Array, " hh ww"]
        Spatial grid in the y-direction
    phase_profile : Float[Array, " hh ww"]
        Phase profile of the optical field
    transmission : Float[Array, " hh ww"]
        Transmission profile of the optical field

    Returns
    -------
    validated_grid_params : GridParams
        Validated grid parameters instance

    Raises
    ------
    ValueError
        If array shapes are inconsistent or data is invalid

    Notes
    -----
    Algorithm:

    - Convert inputs to JAX arrays
    - Validate array shapes:
        - Check all arrays are 2D
        - Check all arrays have the same shape
    - Validate data:
        - Ensure transmission values are between 0 and 1
        - Ensure phase values are finite
        - Ensure grid coordinates are finite
    - Create and return GridParams instance
    """
    xx = jnp.asarray(xx, dtype=jnp.float64)
    yy = jnp.asarray(yy, dtype=jnp.float64)
    phase_profile = jnp.asarray(phase_profile, dtype=jnp.float64)
    transmission = jnp.asarray(transmission, dtype=jnp.float64)

    def validate_and_create() -> GridParams:
        array_dims: int = 2
        hh: int
        ww: int
        hh, ww = xx.shape

        def check_2d_arrays() -> Tuple[
            Float[Array, " hh ww"],
            Float[Array, " hh ww"],
            Float[Array, " hh ww"],
            Float[Array, " hh ww"],
        ]:
            return lax.cond(
                jnp.logical_and(
                    jnp.logical_and(
                        xx.ndim == array_dims, yy.ndim == array_dims
                    ),
                    jnp.logical_and(
                        phase_profile.ndim == array_dims,
                        transmission.ndim == array_dims,
                    ),
                ),
                lambda: (xx, yy, phase_profile, transmission),
                lambda: lax.stop_gradient(
                    lax.cond(
                        False,
                        lambda: (xx, yy, phase_profile, transmission),
                        lambda: (xx, yy, phase_profile, transmission),
                    )
                ),
            )

        def check_same_shape() -> Tuple[
            Float[Array, " hh ww"],
            Float[Array, " hh ww"],
            Float[Array, " hh ww"],
            Float[Array, " hh ww"],
        ]:
            return lax.cond(
                jnp.logical_and(
                    jnp.logical_and(
                        xx.shape == (hh, ww), yy.shape == (hh, ww)
                    ),
                    jnp.logical_and(
                        phase_profile.shape == (hh, ww),
                        transmission.shape == (hh, ww),
                    ),
                ),
                lambda: (xx, yy, phase_profile, transmission),
                lambda: lax.stop_gradient(
                    lax.cond(
                        False,
                        lambda: (xx, yy, phase_profile, transmission),
                        lambda: (xx, yy, phase_profile, transmission),
                    )
                ),
            )

        def check_transmission_range() -> Float[Array, " hh ww"]:
            return lax.cond(
                jnp.logical_and(
                    jnp.all(transmission >= 0), jnp.all(transmission <= 1)
                ),
                lambda: transmission,
                lambda: lax.stop_gradient(
                    lax.cond(False, lambda: transmission, lambda: transmission)
                ),
            )

        def check_phase_finite() -> Float[Array, " hh ww"]:
            return lax.cond(
                jnp.all(jnp.isfinite(phase_profile)),
                lambda: phase_profile,
                lambda: lax.stop_gradient(
                    lax.cond(
                        False, lambda: phase_profile, lambda: phase_profile
                    )
                ),
            )

        def check_grid_finite() -> (
            Tuple[Float[Array, " hh ww"], Float[Array, " hh ww"]]
        ):
            return lax.cond(
                jnp.logical_and(
                    jnp.all(jnp.isfinite(xx)), jnp.all(jnp.isfinite(yy))
                ),
                lambda: (xx, yy),
                lambda: lax.stop_gradient(
                    lax.cond(False, lambda: (xx, yy), lambda: (xx, yy))
                ),
            )

        check_2d_arrays()
        check_same_shape()
        check_transmission_range()
        check_phase_finite()
        check_grid_finite()

        return GridParams(
            xx=xx,
            yy=yy,
            phase_profile=phase_profile,
            transmission=transmission,
        )

    validated_grid_params: GridParams = validate_and_create()
    return validated_grid_params


@jaxtyped(typechecker=beartype)
def make_optical_wavefront(
    field: Union[Complex[Array, " hh ww"], Complex[Array, " hh ww 2"]],
    wavelength: ScalarNumeric,
    dx: ScalarNumeric,
    z_position: ScalarNumeric,
    polarization: ScalarBool = False,
) -> OpticalWavefront:
    """JAX-safe factory function for OpticalWavefront with data
    validation.

    Parameters
    ----------
    field : Union[Complex[Array, " hh ww"], Complex[Array, " hh ww 2"]]
       Complex amplitude of the optical field. Should be 2D for scalar
       fields or 3D with last dimension 2 for polarized fields.
       Polarization is automatically detected from field dimensions.
    wavelength : ScalarNumeric
        Wavelength of the optical wavefront in meters
    dx : ScalarNumeric
        Spatial sampling interval (grid spacing) in meters
    z_position : ScalarNumeric
        Axial position of the wavefront in the propagation direction in
        meters.
    polarization : Bool[Array, " "]
        Boolean indicating whether the field is polarized.
        Default is False.

    Returns
    -------
    validated_optical_wavefront : OpticalWavefront
        Validated optical wavefront instance

    Raises
    ------
    ValueError
        If data is invalid or parameters are out of valid ranges

    Notes
    -----
    Algorithm:

    - Convert inputs to JAX arrays
    - Auto-detect polarization based on field dimensions (3D with last
      dimension 2 means polarized)
    - Validate field array:
        - Check it's 2D or 3D with last dimension 2
        - Ensure all values are finite
    - Validate parameters:
        - Check wavelength is positive
        - Check dx is positive
        - Check z_position is finite
    - Create and return OpticalWavefront instance
    """
    non_polar_dim: int = 2
    polar_dim: int = 3
    field: Complex[Array, " hh ww"] = jnp.asarray(field, dtype=jnp.complex128)
    wavelength: Float[Array, " "] = jnp.asarray(wavelength, dtype=jnp.float64)
    dx: Float[Array, " "] = jnp.asarray(dx, dtype=jnp.float64)
    z_position: Float[Array, " "] = jnp.asarray(z_position, dtype=jnp.float64)
    polarization: Bool[Array, " "] = jnp.asarray(polarization, dtype=jnp.bool_)

    polarization = jnp.where(
        field.ndim == polar_dim,
        jnp.asarray(field.shape[-1] == non_polar_dim, dtype=jnp.bool_),
        polarization,
    )

    def validate_and_create() -> OpticalWavefront:
        def check_field_dimensions() -> (
            Union[Complex[Array, " hh ww"], Complex[Array, " hh ww 2"]]
        ):
            non_polar_dimensions: int = 2
            polar_dimensions: int = 3

            def check_polarized() -> Complex[Array, " hh ww 2"]:
                return lax.cond(
                    jnp.logical_and(
                        field.ndim == polar_dimensions,
                        field.shape[-1] == non_polar_dimensions,
                    ),
                    lambda: field,
                    lambda: lax.stop_gradient(
                        lax.cond(False, lambda: field, lambda: field)
                    ),
                )

            def check_scalar() -> Complex[Array, " hh ww"]:
                return lax.cond(
                    field.ndim == non_polar_dimensions,
                    lambda: field,
                    lambda: lax.stop_gradient(
                        lax.cond(False, lambda: field, lambda: field)
                    ),
                )

            return lax.cond(
                polarization,
                check_polarized,
                check_scalar,
            )

        def check_field_finite() -> (
            Union[Complex[Array, " hh ww"], Complex[Array, " hh ww 2"]]
        ):
            return lax.cond(
                jnp.all(jnp.isfinite(field)),
                lambda: field,
                lambda: lax.stop_gradient(
                    lax.cond(False, lambda: field, lambda: field)
                ),
            )

        def check_wavelength() -> Float[Array, " "]:
            return lax.cond(
                wavelength > 0,
                lambda: wavelength,
                lambda: lax.stop_gradient(
                    lax.cond(False, lambda: wavelength, lambda: wavelength)
                ),
            )

        def check_dx() -> Float[Array, " "]:
            return lax.cond(
                dx > 0,
                lambda: dx,
                lambda: lax.stop_gradient(
                    lax.cond(False, lambda: dx, lambda: dx)
                ),
            )

        def check_z_position() -> Float[Array, " "]:
            return lax.cond(
                jnp.isfinite(z_position),
                lambda: z_position,
                lambda: lax.stop_gradient(
                    lax.cond(False, lambda: z_position, lambda: z_position)
                ),
            )

        check_field_dimensions()
        check_field_finite()
        check_wavelength()
        check_dx()
        check_z_position()

        return OpticalWavefront(
            field=field,
            wavelength=wavelength,
            dx=dx,
            z_position=z_position,
            polarization=polarization,
        )

    validated_optical_wavefront: OpticalWavefront = validate_and_create()
    return validated_optical_wavefront


@jaxtyped(typechecker=beartype)
def make_propagating_wavefront(
    field: Union[Complex[Array, " zz hh ww"], Complex[Array, " zz hh ww 2"]],
    wavelength: ScalarNumeric,
    dx: ScalarNumeric,
    z_positions: Float[Array, " zz"],
    polarization: ScalarBool = False,
) -> PropagatingWavefront:
    """JAX-safe factory function for PropagatingWavefront with data
    validation.

    Parameters
    ----------
    field : Union[Complex[Array, " zz hh ww"], Complex[Array, " zz hh ww 2"]]
        Complex amplitude of the optical field. Should be 3D for scalar
        fields (Z, H, W) or 4D with last dimension 2 for polarized fields
        (Z, H, W, 2). Z represents slices along the propagation direction.
    wavelength : ScalarNumeric
        Wavelength of the optical wavefront in meters
    dx : ScalarNumeric
        Spatial sampling interval (grid spacing) in meters
    z_positions : Float[Array, " zz"]
        Axial positions of the wavefront slices along the propagation
        direction in meters.
    polarization : ScalarBool
        Boolean indicating whether the field is polarized.
        Default is False.

    Returns
    -------
    validated_propagating_wavefront : PropagatingWavefront
        Validated propagating wavefront instance

    Raises
    ------
    ValueError
        If data is invalid or parameters are out of valid ranges

    Notes
    -----
    Algorithm:

    - Convert inputs to JAX arrays
    - Auto-detect polarization based on field dimensions (4D with last
      dimension 2 means polarized)
    - Validate field array:
        - Check it's 3D or 4D with last dimension 2
        - Ensure all values are finite
    - Validate parameters:
        - Check wavelength is positive
        - Check dx is positive
        - Check z_positions are finite
        - Check z_positions length matches field's first dimension
    - Create and return PropagatingWavefront instance
    """
    non_polar_dim: int = 3
    polar_dim: int = 4
    polarization_components: int = 2
    field = jnp.asarray(field, dtype=jnp.complex128)
    wavelength_arr: Float[Array, " "] = jnp.asarray(
        wavelength, dtype=jnp.float64
    )
    dx_arr: Float[Array, " "] = jnp.asarray(dx, dtype=jnp.float64)
    z_positions_arr: Float[Array, " zz"] = jnp.asarray(
        z_positions, dtype=jnp.float64
    )
    polarization_arr: Bool[Array, " "] = jnp.asarray(
        polarization, dtype=jnp.bool_
    )

    polarization_arr = jnp.where(
        field.ndim == polar_dim,
        jnp.asarray(
            field.shape[-1] == polarization_components, dtype=jnp.bool_
        ),
        polarization_arr,
    )

    def validate_and_create() -> PropagatingWavefront:
        def check_field_dimensions() -> (
            Union[Complex[Array, " zz hh ww"], Complex[Array, " zz hh ww 2"]]
        ):
            def check_polarized() -> Complex[Array, " zz hh ww 2"]:
                return lax.cond(
                    jnp.logical_and(
                        field.ndim == polar_dim,
                        field.shape[-1] == polarization_components,
                    ),
                    lambda: field,
                    lambda: lax.stop_gradient(
                        lax.cond(False, lambda: field, lambda: field)
                    ),
                )

            def check_scalar() -> Complex[Array, " zz hh ww"]:
                return lax.cond(
                    field.ndim == non_polar_dim,
                    lambda: field,
                    lambda: lax.stop_gradient(
                        lax.cond(False, lambda: field, lambda: field)
                    ),
                )

            return lax.cond(
                polarization_arr,
                check_polarized,
                check_scalar,
            )

        def check_field_finite() -> (
            Union[Complex[Array, " zz hh ww"], Complex[Array, " zz hh ww 2"]]
        ):
            return lax.cond(
                jnp.all(jnp.isfinite(field)),
                lambda: field,
                lambda: lax.stop_gradient(
                    lax.cond(False, lambda: field, lambda: field)
                ),
            )

        def check_wavelength() -> Float[Array, " "]:
            return lax.cond(
                wavelength_arr > 0,
                lambda: wavelength_arr,
                lambda: lax.stop_gradient(
                    lax.cond(
                        False, lambda: wavelength_arr, lambda: wavelength_arr
                    )
                ),
            )

        def check_dx() -> Float[Array, " "]:
            return lax.cond(
                dx_arr > 0,
                lambda: dx_arr,
                lambda: lax.stop_gradient(
                    lax.cond(False, lambda: dx_arr, lambda: dx_arr)
                ),
            )

        def check_z_positions_finite() -> Float[Array, " zz"]:
            return lax.cond(
                jnp.all(jnp.isfinite(z_positions_arr)),
                lambda: z_positions_arr,
                lambda: lax.stop_gradient(
                    lax.cond(
                        False, lambda: z_positions_arr, lambda: z_positions_arr
                    )
                ),
            )

        def check_z_positions_length() -> Float[Array, " zz"]:
            return lax.cond(
                z_positions_arr.shape[0] == field.shape[0],
                lambda: z_positions_arr,
                lambda: lax.stop_gradient(
                    lax.cond(
                        False, lambda: z_positions_arr, lambda: z_positions_arr
                    )
                ),
            )

        check_field_dimensions()
        check_field_finite()
        check_wavelength()
        check_dx()
        check_z_positions_finite()
        check_z_positions_length()

        return PropagatingWavefront(
            field=field,
            wavelength=wavelength_arr,
            dx=dx_arr,
            z_positions=z_positions_arr,
            polarization=polarization_arr,
        )

    validated_propagating_wavefront: PropagatingWavefront = (
        validate_and_create()
    )
    return validated_propagating_wavefront


@jaxtyped(typechecker=beartype)
def optical2propagating(
    wavefronts: Tuple[OpticalWavefront, ...],
) -> PropagatingWavefront:
    """Create a PropagatingWavefront from a tuple of OpticalWavefronts.

    Parameters
    ----------
    wavefronts : Tuple[OpticalWavefront, ...]
        Tuple of OpticalWavefront instances. All wavefronts must have the
        same wavelength, dx, polarization, and field shape (H, W).

    Returns
    -------
    propagating_wavefront : PropagatingWavefront
        A PropagatingWavefront containing all input wavefronts stacked
        along the z dimension.

    Raises
    ------
    ValueError
        If wavefronts tuple is empty, or if wavefronts have inconsistent
        wavelength, dx, polarization, or field shapes.

    Notes
    -----
    Algorithm:

    - Extract fields from all wavefronts and stack along axis 0
    - Extract z_positions from each wavefront
    - Validate all wavefronts have consistent wavelength, dx, and
      polarization
    - Create PropagatingWavefront using the factory function
    """
    if len(wavefronts) == 0:
        raise ValueError("wavefronts tuple cannot be empty")

    stacked_fields = jnp.stack([wf.field for wf in wavefronts], axis=0)
    z_positions = jnp.array([wf.z_position for wf in wavefronts])
    wavelength = wavefronts[0].wavelength
    dx = wavefronts[0].dx
    polarization = wavefronts[0].polarization

    return make_propagating_wavefront(
        field=stacked_fields,
        wavelength=wavelength,
        dx=dx,
        z_positions=z_positions,
        polarization=polarization,
    )


@jaxtyped(typechecker=beartype)
def make_microscope_data(
    image_data: Union[Float[Array, " pp hh ww"], Float[Array, " xx yy hh ww"]],
    positions: Num[Array, " pp 2"],
    wavelength: ScalarNumeric,
    dx: ScalarNumeric,
) -> MicroscopeData:
    """JAX-safe factory function for MicroscopeData with data
    validation.

    Parameters
    ----------
    image_data :
        Union[Float[Array, " pp hh ww"], Float[Array, " xx yy hh ww"]]
        3D or 4D image data representing the optical field
    positions : Num[Array, " pp 2"]
        Positions of the images during collection
    wavelength : ScalarNumeric
        Wavelength of the optical wavefront in meters
    dx : ScalarNumeric
        Spatial sampling interval (grid spacing) in meters

    Returns
    -------
    validated_microscope_data : MicroscopeData
        Validated microscope data instance

    Raises
    ------
    ValueError
        If data is invalid or parameters are out of valid ranges

    Notes
    -----
    Algorithm:

    - Convert inputs to JAX arrays
    - Validate image_data:
        - Check it's 3D or 4D
        - Ensure all values are finite and non-negative
    - Validate positions:
        - Check it's 2D with shape (pp, 2)
        - Ensure all values are finite
    - Validate parameters:
        - Check wavelength is positive
        - Check dx is positive
    - Validate consistency:
        - Check P matches between image_data and positions
    - Create and return MicroscopeData instance
    """
    image_data: Union[
        Float[Array, " pp hh ww"], Float[Array, " xx yy hh ww"]
    ] = jnp.asarray(image_data, dtype=jnp.float64)
    positions: Num[Array, " pp 2"] = jnp.asarray(positions, dtype=jnp.float64)
    wavelength: Float[Array, " "] = jnp.asarray(wavelength, dtype=jnp.float64)
    dx: Float[Array, " "] = jnp.asarray(dx, dtype=jnp.float64)
    expected_image_dim = 2
    expected_diffractogram_dim_3d: int = 3
    expected_diffractogram_dim_4d: int = 4

    def validate_and_create() -> MicroscopeData:
        def check_image_dimensions() -> (
            Union[Float[Array, " P H W"], Float[Array, " X Y H W"]]
        ):
            return lax.cond(
                jnp.logical_or(
                    image_data.ndim == expected_diffractogram_dim_3d,
                    image_data.ndim == expected_diffractogram_dim_4d,
                ),
                lambda: image_data,
                lambda: lax.stop_gradient(
                    lax.cond(False, lambda: image_data, lambda: image_data)
                ),
            )

        def check_image_finite() -> (
            Union[Float[Array, " P H W"], Float[Array, " X Y H W"]]
        ):
            return lax.cond(
                jnp.all(jnp.isfinite(image_data)),
                lambda: image_data,
                lambda: lax.stop_gradient(
                    lax.cond(False, lambda: image_data, lambda: image_data)
                ),
            )

        def check_image_nonnegative() -> (
            Union[Float[Array, " P H W"], Float[Array, " X Y H W"]]
        ):
            return lax.cond(
                jnp.all(image_data >= 0),
                lambda: image_data,
                lambda: lax.stop_gradient(
                    lax.cond(False, lambda: image_data, lambda: image_data)
                ),
            )

        def check_positions_shape() -> Num[Array, " P 2"]:
            return lax.cond(
                positions.shape[1] == expected_image_dim,
                lambda: positions,
                lambda: lax.stop_gradient(
                    lax.cond(False, lambda: positions, lambda: positions)
                ),
            )

        def check_positions_finite() -> Num[Array, " P 2"]:
            return lax.cond(
                jnp.all(jnp.isfinite(positions)),
                lambda: positions,
                lambda: lax.stop_gradient(
                    lax.cond(False, lambda: positions, lambda: positions)
                ),
            )

        def check_wavelength() -> Float[Array, " "]:
            return lax.cond(
                wavelength > 0,
                lambda: wavelength,
                lambda: lax.stop_gradient(
                    lax.cond(False, lambda: wavelength, lambda: wavelength)
                ),
            )

        def check_dx() -> Float[Array, " "]:
            return lax.cond(
                dx > 0,
                lambda: dx,
                lambda: lax.stop_gradient(
                    lax.cond(False, lambda: dx, lambda: dx)
                ),
            )

        def check_consistency() -> Tuple[
            Union[Float[Array, " P H W"], Float[Array, " X Y H W"]],
            Num[Array, " P 2"],
        ]:
            pp = positions.shape[0]

            def check_3d_consistency() -> Tuple[
                Union[Float[Array, " pp H W"], Float[Array, " X Y H W"]],
                Num[Array, " pp 2"],
            ]:
                return lax.cond(
                    image_data.shape[0] == pp,
                    lambda: (image_data, positions),
                    lambda: lax.stop_gradient(
                        lax.cond(
                            False,
                            lambda: (image_data, positions),
                            lambda: (image_data, positions),
                        )
                    ),
                )

            def check_4d_consistency() -> Tuple[
                Union[Float[Array, " P H W"], Float[Array, " X Y H W"]],
                Num[Array, " P 2"],
            ]:
                return lax.cond(
                    image_data.shape[0] * image_data.shape[1] == pp,
                    lambda: (image_data, positions),
                    lambda: lax.stop_gradient(
                        lax.cond(
                            False,
                            lambda: (image_data, positions),
                            lambda: (image_data, positions),
                        )
                    ),
                )

            return lax.cond(
                image_data.ndim == expected_image_dim,
                check_3d_consistency,
                check_4d_consistency,
            )

        check_image_dimensions()
        check_image_finite()
        check_image_nonnegative()
        check_positions_shape()
        check_positions_finite()
        check_wavelength()
        check_dx()
        check_consistency()

        return MicroscopeData(
            image_data=image_data,
            positions=positions,
            wavelength=wavelength,
            dx=dx,
        )

    validated_microscope_data: MicroscopeData = validate_and_create()
    return validated_microscope_data


@jaxtyped(typechecker=beartype)
def make_diffractogram(
    image: Float[Array, " hh ww"],
    wavelength: ScalarNumeric,
    dx: ScalarNumeric,
) -> Diffractogram:
    """JAX-safe factory function for Diffractogram with data validation.

    Parameters
    ----------
    image : Float[Array, " hh ww"]
        Image data
    wavelength : ScalarNumeric
        Wavelength of the optical wavefront in meters
    dx : ScalarNumeric
        Spatial sampling interval (grid spacing) in meters

    Returns
    -------
    validated_diffractogram : Diffractogram
        Validated diffractogram instance

    Raises
    ------
    ValueError
        If data is invalid or parameters are out of valid ranges

    Notes
    -----
    Algorithm:

    - Convert inputs to JAX arrays
    - Validate image array:
        - Check it's 2D
        - Ensure all values are finite and non-negative
    - Validate parameters:
        - Check wavelength is positive
        - Check dx is positive
    - Create and return Diffractogram instance
    """
    image: Float[Array, " H W"] = jnp.asarray(image, dtype=jnp.float64)
    wavelength: Float[Array, " "] = jnp.asarray(wavelength, dtype=jnp.float64)
    dx: Float[Array, " "] = jnp.asarray(dx, dtype=jnp.float64)
    expected_sample_dim: int = 2

    def validate_and_create() -> Diffractogram:
        def check_2d_image() -> Float[Array, " H W"]:
            return lax.cond(
                image.ndim == expected_sample_dim,
                lambda: image,
                lambda: lax.stop_gradient(
                    lax.cond(False, lambda: image, lambda: image)
                ),
            )

        def check_image_finite() -> Float[Array, " H W"]:
            return lax.cond(
                jnp.all(jnp.isfinite(image)),
                lambda: image,
                lambda: lax.stop_gradient(
                    lax.cond(False, lambda: image, lambda: image)
                ),
            )

        def check_image_nonnegative() -> Float[Array, " H W"]:
            return lax.cond(
                jnp.all(image >= 0),
                lambda: image,
                lambda: lax.stop_gradient(
                    lax.cond(False, lambda: image, lambda: image)
                ),
            )

        def check_wavelength() -> Float[Array, " "]:
            return lax.cond(
                wavelength > 0,
                lambda: wavelength,
                lambda: lax.stop_gradient(
                    lax.cond(False, lambda: wavelength, lambda: wavelength)
                ),
            )

        def check_dx() -> Float[Array, " "]:
            return lax.cond(
                dx > 0,
                lambda: dx,
                lambda: lax.stop_gradient(
                    lax.cond(False, lambda: dx, lambda: dx)
                ),
            )

        check_2d_image()
        check_image_finite()
        check_image_nonnegative()
        check_wavelength()
        check_dx()

        return Diffractogram(
            image=image,
            wavelength=wavelength,
            dx=dx,
        )

    validated_diffractogram: Diffractogram = validate_and_create()
    return validated_diffractogram


@jaxtyped(typechecker=beartype)
def make_sample_function(
    sample: Num[Array, " hh ww"],
    dx: ScalarNumeric,
) -> SampleFunction:
    """JAX-safe factory function for SampleFunction with data
    validation.

    Parameters
    ----------
    sample : Num[Array, " hh ww"]
        The sample function. Will be converted to complex if real.
    dx : ScalarNumeric
        Spatial sampling interval (grid spacing) in meters

    Returns
    -------
    validated_sample_function : SampleFunction
        Validated sample function instance

    Raises
    ------
    ValueError
        If data is invalid or parameters are out of valid ranges

    Notes
    -----
    Algorithm:

    - Convert inputs to JAX arrays
    - Validate sample array:
        - Check it's 2D
        - Ensure all values are finite
    - Validate parameters:
        - Check dx is positive
    - Create and return SampleFunction instance
    """
    sample: Complex[Array, " hh ww"] = jnp.asarray(
        sample, dtype=jnp.complex128
    )
    dx: Float[Array, " "] = jnp.asarray(dx, dtype=jnp.float64)
    expected_sample_dim: int = 2

    def validate_and_create() -> SampleFunction:
        def check_2d_sample() -> Complex[Array, " hh ww"]:
            return lax.cond(
                sample.ndim == expected_sample_dim,
                lambda: sample,
                lambda: lax.stop_gradient(
                    lax.cond(False, lambda: sample, lambda: sample)
                ),
            )

        def check_sample_finite() -> Complex[Array, " hh ww"]:
            return lax.cond(
                jnp.all(jnp.isfinite(sample)),
                lambda: sample,
                lambda: lax.stop_gradient(
                    lax.cond(False, lambda: sample, lambda: sample)
                ),
            )

        def check_dx() -> Float[Array, " "]:
            return lax.cond(
                dx > 0,
                lambda: dx,
                lambda: lax.stop_gradient(
                    lax.cond(False, lambda: dx, lambda: dx)
                ),
            )

        check_2d_sample()
        check_sample_finite()
        check_dx()

        return SampleFunction(
            sample=sample,
            dx=dx,
        )

    validated_sample_function: SampleFunction = validate_and_create()
    return validated_sample_function


@jaxtyped(typechecker=beartype)
def make_sliced_material_function(
    material: Num[Array, " hh ww zz"],
    dx: ScalarNumeric,
    tz: ScalarNumeric,
) -> SlicedMaterialFunction:
    """JAX-safe validated factory function for SlicedMaterialFunction.

    Parameters
    ----------
    material : Num[Array, " hh ww zz"]
        3D array of complex refractive indices. The real part represents
        the refractive index n, and the imaginary part represents the
        extinction coefficient κ (absorption). Will be converted to complex
        if real.
    dx : ScalarNumeric
        Spatial sampling interval (pixel spacing) within each slice in meters
    tz : ScalarNumeric
        Interslice distance (spacing between slices) in the z-direction in
        meters.

    Returns
    -------
    validated_sliced_material : SlicedMaterialFunction
        Validated sliced material function instance

    Raises
    ------
    ValueError
        If data is invalid or parameters are out of valid ranges

    Notes
    -----
    Algorithm:

    - Convert inputs to JAX arrays
    - Validate material array:
        - Check it's 3D
        - Ensure all values are finite
    - Validate parameters:
        - Check dx is positive
        - Check tz is positive
    - Create and return SlicedMaterialFunction instance
    """
    material_array: Complex[Array, " hh ww zz"] = jnp.asarray(
        material, dtype=jnp.complex128
    )
    dx_array: Float[Array, " "] = jnp.asarray(dx, dtype=jnp.float64)
    tz_array: Float[Array, " "] = jnp.asarray(tz, dtype=jnp.float64)
    expected_material_dim: int = 3

    def validate_and_create() -> SlicedMaterialFunction:
        def check_3d_material() -> Complex[Array, " hh ww zz"]:
            return lax.cond(
                material_array.ndim == expected_material_dim,
                lambda: material_array,
                lambda: lax.stop_gradient(
                    lax.cond(
                        False, lambda: material_array, lambda: material_array
                    )
                ),
            )

        def check_material_finite() -> Complex[Array, " hh ww zz"]:
            return lax.cond(
                jnp.all(jnp.isfinite(material_array)),
                lambda: material_array,
                lambda: lax.stop_gradient(
                    lax.cond(
                        False, lambda: material_array, lambda: material_array
                    )
                ),
            )

        def check_dx() -> Float[Array, " "]:
            return lax.cond(
                dx_array > 0,
                lambda: dx_array,
                lambda: lax.stop_gradient(
                    lax.cond(False, lambda: dx_array, lambda: dx_array)
                ),
            )

        def check_tz() -> Float[Array, " "]:
            return lax.cond(
                tz_array > 0,
                lambda: tz_array,
                lambda: lax.stop_gradient(
                    lax.cond(False, lambda: tz_array, lambda: tz_array)
                ),
            )

        check_3d_material()
        check_material_finite()
        check_dx()
        check_tz()

        return SlicedMaterialFunction(
            material=material_array,
            dx=dx_array,
            tz=tz_array,
        )

    validated_sliced_material: SlicedMaterialFunction = validate_and_create()
    return validated_sliced_material


@jaxtyped(typechecker=beartype)
def make_optimizer_state(
    shape: Tuple,
    m: Optional[Union[Complex[Array, " ..."], ScalarComplex]] = 1j,
    v: Optional[Union[Float[Array, " ..."], ScalarFloat]] = 0.0,
    step: Optional[ScalarInteger] = 0,
) -> OptimizerState:
    """JAX-safe factory function for OptimizerState with data
    validation.

    Parameters
    ----------
    shape : Tuple
        Shape of the parameters to be optimized
    m : Optional[Complex[Array, "..."]], optional
        First moment estimate. If None, initialized to zeros with given
        shape.
        Default is 1j.
    v : Optional[Float[Array, "..."]], optional
        Second moment estimate. If None, initialized to zeros with given
        shape.
        Default is 0.0.
    step : Optional[ScalarInteger], optional
        Step count. Default is 0.

    Returns
    -------
    validated_optimizer_state : OptimizerState
        Validated optimizer state instance

    Raises
    ------
    ValueError
        If arrays have incompatible shapes with the given shape
        parameter

    Notes
    -----
    Algorithm:

    - Convert all inputs to JAX arrays with appropriate dtypes
    - Always broadcast m and v to the target shape (if already the
      right shape, broadcast_to is a no-op)
    - Validate arrays have compatible shapes
    - Create and return OptimizerState instance
    """
    m_input = jnp.asarray(m, dtype=jnp.complex128)
    v_input = jnp.asarray(v, dtype=jnp.float64)
    step_input = jnp.asarray(step, dtype=jnp.int32)

    m_array = jnp.broadcast_to(m_input, shape).astype(jnp.complex128)
    v_array = jnp.broadcast_to(v_input, shape).astype(jnp.float64)

    step_array = step_input

    def validate_and_create() -> OptimizerState:
        def check_m_shape() -> Complex[Array, " ..."]:
            return lax.cond(
                m_array.shape == shape,
                lambda: m_array,
                lambda: lax.stop_gradient(
                    lax.cond(False, lambda: m_array, lambda: m_array)
                ),
            )

        def check_v_shape() -> Float[Array, " ..."]:
            return lax.cond(
                v_array.shape == shape,
                lambda: v_array,
                lambda: lax.stop_gradient(
                    lax.cond(False, lambda: v_array, lambda: v_array)
                ),
            )

        def check_step_scalar() -> Int[Array, " "]:
            return lax.cond(
                step_array.ndim == 0,
                lambda: step_array,
                lambda: lax.stop_gradient(
                    lax.cond(False, lambda: step_array, lambda: step_array)
                ),
            )

        check_m_shape()
        check_v_shape()
        check_step_scalar()

        return OptimizerState(
            m=m_array,
            v=v_array,
            step=step_array,
        )

    validated_optimizer_state: OptimizerState = validate_and_create()
    return validated_optimizer_state


@jaxtyped(typechecker=beartype)
def make_ptychography_params(
    camera_pixel_size: ScalarNumeric,
    num_iterations: ScalarInteger,
    learning_rate: ScalarNumeric = 1e-3,
    loss_type: ScalarInteger = 0,
    optimizer_type: ScalarInteger = 0,
    zoom_factor_bounds: Optional[Float[Array, " 2"]] = None,
    aperture_diameter_bounds: Optional[Float[Array, " 2"]] = None,
    travel_distance_bounds: Optional[Float[Array, " 2"]] = None,
    aperture_center_bounds: Optional[Float[Array, " 2 2"]] = None,
) -> PtychographyParams:
    """Create a PtychographyParams PyTree with validated parameters.

    Parameters
    ----------
    camera_pixel_size : ScalarNumeric
        Camera pixel size in meters (must be positive)
    num_iterations : ScalarInteger
        Number of optimization iterations per call (must be positive)
    learning_rate : ScalarNumeric, optional
        Learning rate for optimization. Default is 1e-3.
    loss_type : ScalarInteger, optional
        Loss function type (0=mse, 1=mae, 2=poisson). Default is 0 (mse).
    optimizer_type : ScalarInteger, optional
        Optimizer type (0=adam, 1=adagrad, 2=rmsprop, 3=sgd). Default is 0.
    zoom_factor_bounds : Float[Array, " 2"], optional
        Lower and upper bounds for zoom factor [lower, upper].
        Default is [-inf, inf] (no bounds).
    aperture_diameter_bounds : Float[Array, " 2"], optional
        Lower and upper bounds for aperture diameter [lower, upper].
        Default is [-inf, inf] (no bounds).
    travel_distance_bounds : Float[Array, " 2"], optional
        Lower and upper bounds for travel distance [lower, upper].
        Default is [-inf, inf] (no bounds).
    aperture_center_bounds : Float[Array, " 2 2"], optional
        Lower and upper bounds for aperture center
        [[lower_x, lower_y], [upper_x, upper_y]].
        Default is [[-inf, -inf], [inf, inf]] (no bounds).

    Returns
    -------
    PtychographyParams
        Validated ptychography parameters as a PyTree

    Notes
    -----
    This function performs runtime validation to ensure all parameters
    are properly formatted and within valid ranges before creating the
    PtychographyParams PyTree. All scalar inputs are converted to JAX
    arrays.

    Loss types: 0=mse, 1=mae, 2=poisson
    Optimizer types: 0=adam, 1=adagrad, 2=rmsprop, 3=sgd
    """
    camera_pixel_size_arr: Float[Array, " "] = jnp.asarray(
        camera_pixel_size, dtype=jnp.float64
    )
    num_iterations_arr: Int[Array, " "] = jnp.asarray(
        num_iterations, dtype=jnp.int64
    )
    learning_rate_arr: Float[Array, " "] = jnp.asarray(
        learning_rate, dtype=jnp.float64
    )
    loss_type_arr: Int[Array, " "] = jnp.asarray(loss_type, dtype=jnp.int64)
    optimizer_type_arr: Int[Array, " "] = jnp.asarray(
        optimizer_type, dtype=jnp.int64
    )

    inf: float = float("inf")
    zoom_factor_bounds_arr: Float[Array, " 2"] = (
        jnp.asarray(zoom_factor_bounds, dtype=jnp.float64)
        if zoom_factor_bounds is not None
        else jnp.array([-inf, inf], dtype=jnp.float64)
    )
    aperture_diameter_bounds_arr: Float[Array, " 2"] = (
        jnp.asarray(aperture_diameter_bounds, dtype=jnp.float64)
        if aperture_diameter_bounds is not None
        else jnp.array([-inf, inf], dtype=jnp.float64)
    )
    travel_distance_bounds_arr: Float[Array, " 2"] = (
        jnp.asarray(travel_distance_bounds, dtype=jnp.float64)
        if travel_distance_bounds is not None
        else jnp.array([-inf, inf], dtype=jnp.float64)
    )
    aperture_center_bounds_arr: Float[Array, " 2 2"] = (
        jnp.asarray(aperture_center_bounds, dtype=jnp.float64)
        if aperture_center_bounds is not None
        else jnp.array([[-inf, -inf], [inf, inf]], dtype=jnp.float64)
    )

    return PtychographyParams(
        camera_pixel_size=camera_pixel_size_arr,
        num_iterations=num_iterations_arr,
        learning_rate=learning_rate_arr,
        loss_type=loss_type_arr,
        optimizer_type=optimizer_type_arr,
        zoom_factor_bounds=zoom_factor_bounds_arr,
        aperture_diameter_bounds=aperture_diameter_bounds_arr,
        travel_distance_bounds=travel_distance_bounds_arr,
        aperture_center_bounds=aperture_center_bounds_arr,
    )


@jaxtyped(typechecker=beartype)
def make_epie_params(
    effective_dx: ScalarNumeric,
    num_iterations: ScalarInteger,
    alpha: ScalarNumeric = 1.0,
    beta: ScalarNumeric = 1.0,
    padding: ScalarInteger = 32,
) -> EpieParams:
    """Create an EpieParams PyTree with validated parameters.

    Parameters
    ----------
    effective_dx : ScalarNumeric
        Desired pixel size in meters for sample/probe reconstruction.
        This is a user input that determines resolution.
    num_iterations : ScalarInteger
        Number of ePIE iterations (sweeps over all scan positions).
    alpha : ScalarNumeric, optional
        Object update step size. Default is 1.0. Typical range is 0.5-1.0.
    beta : ScalarNumeric, optional
        Probe update step size. Default is 1.0. Set to 0 to freeze the
        probe and only update the object. Typical range is 0.0-1.0.
    padding : ScalarInteger, optional
        Additional padding in pixels around the scanned region. Default is 32.
        Controls the reconstructed FOV beyond the scan area.

    Returns
    -------
    EpieParams
        Validated ePIE parameters as a PyTree

    Notes
    -----
    ePIE parameters differ from gradient-based optimization:

    - alpha/beta are step sizes for ePIE update formulas, not learning rates
    - Setting beta=0 freezes the probe (useful when probe is well-known)
    - effective_dx is user-specified, not derived from optical parameters
    - The sample is initialized to 1+0j (unity transmission)

    Examples
    --------
    >>> params = make_epie_params(
    ...     effective_dx=1.6e-6,  # 1.6 µm pixels
    ...     num_iterations=100,
    ...     alpha=0.8,
    ...     beta=0.5,  # Update probe more slowly than object
    ... )
    """
    effective_dx_arr: Float[Array, " "] = jnp.asarray(
        effective_dx, dtype=jnp.float64
    )
    num_iterations_arr: Int[Array, " "] = jnp.asarray(
        num_iterations, dtype=jnp.int64
    )
    alpha_arr: Float[Array, " "] = jnp.asarray(alpha, dtype=jnp.float64)
    beta_arr: Float[Array, " "] = jnp.asarray(beta, dtype=jnp.float64)
    padding_arr: Int[Array, " "] = jnp.asarray(padding, dtype=jnp.int64)

    return EpieParams(
        effective_dx=effective_dx_arr,
        num_iterations=num_iterations_arr,
        alpha=alpha_arr,
        beta=beta_arr,
        padding=padding_arr,
    )


@jaxtyped(typechecker=beartype)
def make_ptychography_reconstruction(  # noqa: PLR0913
    sample: SampleFunction,
    lightwave: OpticalWavefront,
    translated_positions: Float[Array, " N 2"],
    zoom_factor: ScalarNumeric,
    aperture_diameter: ScalarNumeric,
    aperture_center: Optional[Float[Array, " 2"]],
    travel_distance: ScalarNumeric,
    intermediate_samples: Complex[Array, " Hs Ws S"],
    intermediate_lightwaves: Complex[Array, " Hp Wp S"],
    intermediate_zoom_factors: Float[Array, " S"],
    intermediate_aperture_diameters: Float[Array, " S"],
    intermediate_aperture_centers: Float[Array, " 2 S"],
    intermediate_travel_distances: Float[Array, " S"],
    losses: Float[Array, " L 2"],
) -> PtychographyReconstruction:
    """Create a PtychographyReconstruction PyTree with validated results.

    Parameters
    ----------
    sample : SampleFunction
        Final reconstructed sample covering the scanned FOV
    lightwave : OpticalWavefront
        Final reconstructed probe/lightwave
    translated_positions : Float[Array, " N 2"]
        Scan positions translated to FOV coordinates (in meters)
    zoom_factor : ScalarNumeric
        Final optimized zoom factor
    aperture_diameter : ScalarNumeric
        Final optimized aperture diameter in meters
    aperture_center : Optional[Float[Array, " 2"]]
        Final optimized aperture center position (x, y)
    travel_distance : ScalarNumeric
        Final optimized light propagation distance in meters
    intermediate_samples : Complex[Array, " Hs Ws S"]
        Intermediate sample reconstructions during optimization
    intermediate_lightwaves : Complex[Array, " Hp Wp S"]
        Intermediate probe reconstructions during optimization
    intermediate_zoom_factors : Float[Array, " S"]
        Intermediate zoom factors during optimization
    intermediate_aperture_diameters : Float[Array, " S"]
        Intermediate aperture diameters during optimization
    intermediate_aperture_centers : Float[Array, " 2 S"]
        Intermediate aperture centers during optimization
    intermediate_travel_distances : Float[Array, " S"]
        Intermediate travel distances during optimization
    losses : Float[Array, " L 2"]
        Loss history with columns [iteration, loss_value]. L is the number
        of recorded iterations (may differ from number of positions N).

    Returns
    -------
    PtychographyReconstruction
        Validated ptychography reconstruction results as a PyTree

    Notes
    -----
    This function performs runtime validation to ensure all results
    are properly formatted before creating the PtychographyReconstruction
    PyTree. Scalar inputs are converted to JAX arrays.
    """
    zoom_factor_array = jnp.asarray(zoom_factor, dtype=jnp.float64)
    aperture_diameter_array = jnp.asarray(aperture_diameter, dtype=jnp.float64)
    travel_distance_array = jnp.asarray(travel_distance, dtype=jnp.float64)
    aperture_center_array = (
        jnp.asarray(aperture_center, dtype=jnp.float64)
        if aperture_center is not None
        else None
    )

    intermediate_samples_array = jnp.asarray(
        intermediate_samples, dtype=jnp.complex128
    )
    intermediate_lightwaves_array = jnp.asarray(
        intermediate_lightwaves, dtype=jnp.complex128
    )
    intermediate_zoom_factors_array = jnp.asarray(
        intermediate_zoom_factors, dtype=jnp.float64
    )
    intermediate_aperture_diameters_array = jnp.asarray(
        intermediate_aperture_diameters, dtype=jnp.float64
    )
    intermediate_aperture_centers_array = jnp.asarray(
        intermediate_aperture_centers, dtype=jnp.float64
    )
    intermediate_travel_distances_array = jnp.asarray(
        intermediate_travel_distances, dtype=jnp.float64
    )
    losses_array = jnp.asarray(losses, dtype=jnp.float64)

    translated_positions_array = jnp.asarray(
        translated_positions, dtype=jnp.float64
    )

    return PtychographyReconstruction(
        sample=sample,
        lightwave=lightwave,
        translated_positions=translated_positions_array,
        zoom_factor=zoom_factor_array,
        aperture_diameter=aperture_diameter_array,
        aperture_center=aperture_center_array,
        travel_distance=travel_distance_array,
        intermediate_samples=intermediate_samples_array,
        intermediate_lightwaves=intermediate_lightwaves_array,
        intermediate_zoom_factors=intermediate_zoom_factors_array,
        intermediate_aperture_diameters=intermediate_aperture_diameters_array,
        intermediate_aperture_centers=intermediate_aperture_centers_array,
        intermediate_travel_distances=intermediate_travel_distances_array,
        losses=losses_array,
    )


@jaxtyped(typechecker=beartype)
def make_epie_data(
    diffraction_patterns: Float[Array, " N H W"],
    probe: Complex[Array, " H W"],
    sample: Complex[Array, " Hs Ws"],
    positions: Float[Array, " N 2"],
    effective_dx: ScalarNumeric,
    wavelength: ScalarNumeric,
    original_camera_pixel_size: ScalarNumeric,
    zoom_factor: ScalarNumeric,
) -> EpieData:
    """Create an EpieData PyTree for FFT-compatible ePIE reconstruction.

    This factory function creates the preprocessed data structure needed
    for running ePIE with pure FFT-based propagation. All physical quantities
    are scaled so that the FFT naturally produces the correct far-field
    coordinates.

    Parameters
    ----------
    diffraction_patterns : Float[Array, " N H W"]
        Preprocessed diffraction patterns scaled to FFT-natural pixel size.
        Shape (N, H, W) where N is number of positions, H and W are probe size.
    probe : Complex[Array, " H W"]
        Initial probe field (typically plane wave with aperture applied).
        Same shape as diffraction patterns (H, W).
    sample : Complex[Array, " Hs Ws"]
        Initial sample estimate covering the full FOV.
    positions : Float[Array, " N 2"]
        Scan positions in pixels (in the effective coordinate system).
    effective_dx : ScalarNumeric
        Effective pixel size at sample plane: camera_pixel_size / zoom_factor.
    wavelength : ScalarNumeric
        Wavelength of light in meters.
    original_camera_pixel_size : ScalarNumeric
        Original camera pixel size before preprocessing (for reference).
    zoom_factor : ScalarNumeric
        Original zoom factor (for reference/postprocessing).

    Returns
    -------
    EpieData
        Validated EpieData PyTree ready for FFT-based ePIE reconstruction.

    Notes
    -----
    The key insight is that optical zoom just scales all physical dimensions.
    By dividing pixel sizes and aperture diameter by zoom_factor, we get an
    equivalent problem where the FFT naturally gives the correct far-field
    coordinates without needing scaled Fraunhofer propagation.
    """
    diffraction_patterns_arr = jnp.asarray(
        diffraction_patterns, dtype=jnp.float64
    )
    probe_arr = jnp.asarray(probe, dtype=jnp.complex128)
    sample_arr = jnp.asarray(sample, dtype=jnp.complex128)
    positions_arr = jnp.asarray(positions, dtype=jnp.float64)
    effective_dx_arr = jnp.asarray(effective_dx, dtype=jnp.float64)
    wavelength_arr = jnp.asarray(wavelength, dtype=jnp.float64)
    original_camera_pixel_size_arr = jnp.asarray(
        original_camera_pixel_size, dtype=jnp.float64
    )
    zoom_factor_arr = jnp.asarray(zoom_factor, dtype=jnp.float64)

    return EpieData(
        diffraction_patterns=diffraction_patterns_arr,
        probe=probe_arr,
        sample=sample_arr,
        positions=positions_arr,
        effective_dx=effective_dx_arr,
        wavelength=wavelength_arr,
        original_camera_pixel_size=original_camera_pixel_size_arr,
        zoom_factor=zoom_factor_arr,
    )
