# lct.py

"""
Linear Canonical Transform (LCT)
===============================

Soummer-style two-sided MFT implementation of the Collins diffraction
integral for a general ABCD matrix:

    [[a, b],
     [c, d]]

We implement the 2D separable Collins kernel

    U'(x', y') =
        1 / (i λ b) *
        exp[i π d (x'^2 + y'^2) / (λ b)] *
        ∬ U(x, y)
            exp[i π a (x^2 + y^2) / (λ b)] *
            exp[-i 2π (x x' + y y') / (λ b)] dx dy

using a Soummer-style two-sided MFT:

  - pre-chirp in the input plane,
  - 1D MFT kernels along x and y,
  - post-chirp in the output plane.

This module now provides:

  - A spec-based kernel builder lct_kernels(...) that accepts the same
    CoordSpec conventions as mft/fraunhofer/asm.
  - A low-level "from kernels" application lct_kernel_prop(...).
  - A basic high-level Collins LCT lct_prop_basic(...) without curvature
    factoring.
  - A residual / curvature-aware LCT (lct_prop) that delegates curvature
    bookkeeping to curvature.py.
"""

from __future__ import annotations

import jax.numpy as np
from jax import Array

from .coords import unpack_coord_spec, pad_coords, pad_to
from .curvature import (
    r2_coords,
    apply_curv,
    remove_curv,
    factorise_curv,
    curvature_sampling_parameter,
)
from .mft import (
    mft,
    mft_sampling_parameter,
    mft_phase_parameter,
    mft_kernels,
)

__all__ = [
    "lct_kernels",
    "lct_kernel_prop",
    "lct_prop",
    "lct_fft_output_spec",
    "lct_fft_kernels",
    "lct_fft_kernel_prop",
    "lct_prop_fft",
    "lct_2ft_spec_freq",
    "lct_2ft_kernels",
    "lct_2ft_kernel_prop",
    "lct_prop_2ft",
    "lct_2ft_fft_kernels",
    "lct_2ft_fft_kernel_prop",
    "lct_prop_2ft_fft",
    "lct_prop_curv",
    "lct_sampling_parameter",
    "lct_phase_parameter",
    "lct_sampling_metrics",
    "lct_2ft_sampling_metrics",
    "lct_2ft_fft_spec_freq",
]

# ---------------------------------------------------------------------
# Fundamental LCT kernel builder (spec-based)
# ---------------------------------------------------------------------


def lct_kernels(
    spec_in: Array | tuple,
    spec_out: Array | tuple,
    lam: float,
    ABCD: Array,
) -> tuple:
    """
    Build LCT chirps and MFT kernels from coordinate specifications.

    This is the fundamental implementation of the Collins integral in
    dLux. Both uniform and explicit coordinate grids are supported via
    the CoordSpec API used throughout the Fresnel stack.

    Coordinate specs
    ----------------
    spec_in, spec_out can be:

      - (N, d) or ((Nx, Ny), (dx, dy)) for uniform centred grids, or
      - x (1D Array) for symmetric x=y grids, or
      - (x, y) for explicit per-axis coordinates.

    All arrays are expected to be 1D, monotonically increasing, and
    approximately equally spaced along each axis.

    Parameters
    ----------
    spec_in  : Array or tuple
        Input grid specification.
    spec_out : Array or tuple
        Output grid specification.
    lam : float
        Wavelength [m].
    ABCD : Array, shape (2, 2)
        Ray-transfer matrix [[a, b], [c, d]] for the segment.

    Returns
    -------
    pre  : Array
        Input-plane quadratic phase exp(i π a r^2 / (λ b)),
        shape (Ny_in, Nx_in).
    Kx   : Array
        1D LCT kernel along x, shape (Nx_out, Nx_in).
    Ky   : Array
        1D LCT kernel along y, shape (Ny_out, Ny_in).
    post : Array
        Output-plane quadratic phase exp(i π d r'^2 / (λ b)),
        shape (Ny_out, Nx_out).
    pref : complex
        LCT prefactor 1 / (i λ b).
    scale : float
        Pixel-scale ratio sqrt((dx_out * dy_out) / (dx_in * dy_in)) for
        discrete energy normalisation.

    Notes
    -----
    - We assume uniform, ascending coordinates along each axis when
      estimating dx_in, dy_in, dx_out, dy_out from the coordinate
      vectors implied by spec_in and spec_out.
    - If b is very small, the phase becomes extremely oscillatory and
      numerical sampling will generally be invalid; users should treat
      powered surfaces (b ≈ 0) as separate thin-lens phases.
    """
    x_in, y_in = unpack_coord_spec(spec_in)
    x_out, y_out = unpack_coord_spec(spec_out)

    a, b, c, d = ABCD.flatten()

    # r^2 grids
    r2_in = r2_coords((x_in, y_in))
    r2_out = r2_coords((x_out, y_out))

    # Quadratic chirps
    pre = np.exp(1j * np.pi * a * r2_in / (lam * b))
    post = np.exp(1j * np.pi * d * r2_out / (lam * b))

    # Spacings (uniform ascending assumed)
    dx_in = x_in[1] - x_in[0]
    dy_in = y_in[1] - y_in[0]
    dx_out = x_out[1] - x_out[0]
    dy_out = y_out[1] - y_out[0]

    # 1D MFT kernels: frequency-like factor alpha
    alpha = -2.0 * np.pi / (lam * b)

    # Build MFT kernels from specs with explicit per-axis weights
    Kx, Ky = mft_kernels(
        spec_in=spec_in,
        spec_out=spec_out,
        alpha=alpha,
        weight=(dx_in, dy_in),
    )

    pref = 1.0 / (1j * lam * b)
    scale = np.sqrt((dx_out * dy_out) / (dx_in * dy_in))

    return pre, Kx, Ky, post, pref, scale


# ---------------------------------------------------------------------
# Cached application
# ---------------------------------------------------------------------


def lct_kernel_prop(
    u_in: Array,
    pre: Array,
    Kx: Array,
    Ky: Array,
    post: Array,
    pref: complex,
    scale: float,
    apply_out_curv: bool = True,
) -> Array:
    """
    Apply a precomputed 2D LCT.

    Implements:

        u_tmp = pre * u_in
        u_mft = MFT[u_tmp]    (via separable Kx, Ky)
        u_out = pref * scale * u_mft * post

    Parameters
    ----------
    u_in : Array
        Input complex field, shape (Ny_in, Nx_in).
    pre, post : Array
        Input/output quadratic chirps.
    Kx, Ky : Array
        1D LCT kernels along x and y.
    pref : complex
        LCT prefactor 1 / (i λ b).
    scale : float
        Pixel-scale ratio for discrete normalisation.

    Returns
    -------
    Array
        Output field, shape (Ny_out, Nx_out).
    """
    u_tmp = pre * u_in
    u_mft = mft(u_tmp, Kx, Ky)

    if apply_out_curv:
        return pref * scale * u_mft * post
    else:
        return pref * scale * u_mft


# ---------------------------------------------------------------------
# High-level Collins LCT (basic, no curvature factoring)
# ---------------------------------------------------------------------


def lct_prop(
    u_in: Array,
    spec_in: Array | tuple,
    spec_out: Array | tuple,
    lam: float,
    ABCD: Array,
    apply_out_curv: bool = True,
) -> Array:
    """
    High-level Collins LCT between arbitrary separable grids.

    This performs a pure Collins diffraction step with the given ABCD,
    without any curvature factoring or residualisation. It is the
    fundamental Fresnel engine; curvature-aware wrappers are built on
    top of this.

    Parameters
    ----------
    u_in : Array
        Input field, shape (Ny_in, Nx_in).
    spec_in : Array or tuple
        Input grid specification (CoordSpec). Typical examples:
          - (N, d)           → square uniform grid
          - ((Nx, Ny), (dx, dy))
          - x                → 1D vector, symmetric x=y
          - (x, y)           → explicit per-axis coords
        The lengths implied by spec_in must be compatible with
        u_in.shape; otherwise matrix multiplications will fail.
    spec_out : Array or tuple
        Output grid specification (same conventions).
    lam : float
        Wavelength [m].
    ABCD : Array, shape (2, 2)
        Ray-transfer matrix for this segment.

    Returns
    -------
    Array
        Output field on the grid defined by spec_out.

    Notes
    -----
    - This function does not check sampling validity. Use the LCT
      sampling helpers together with your coordinate specs to check
      whether the chosen (spec_in, spec_out, λ, ABCD) combination is
      Nyquist-sampled.
    """
    pre, Kx, Ky, post, pref, scale = lct_kernels(
        spec_in=spec_in,
        spec_out=spec_out,
        lam=lam,
        ABCD=ABCD,
    )
    return lct_kernel_prop(u_in, pre, Kx, Ky, post, pref, scale, apply_out_curv)


# ---------------------------------------------------------------------
# FFT LCT
# ---------------------------------------------------------------------


def lct_fft_output_spec(
    spec_in: Array | tuple,
    lam: float,
    ABCD: Array,
    npad: int | tuple[int, int] | None = None,
) -> tuple:
    """
    Compute the output coordinate specification for an FFT-based LCT.

    Parameters
    ----------
    spec_in : Array or tuple
        Input grid specification (CoordSpec).
    lam : float
        Wavelength [m].
    ABCD : Array, shape (2, 2)
        Ray-transfer matrix for this segment.
    npad : int or tuple[int, int] or None
        Padding factor for the input coordinates.

    Returns
    -------
    tuple
        Output coordinate specification (x_out, y_out).
    """
    a, b, c, d = ABCD.flatten()

    # Unpack input coordinates
    x_in, y_in = unpack_coord_spec(spec_in)

    # Apply padding if specified
    if npad is not None:
        x_fft, y_fft = pad_coords(x_in, y_in, npad)
    else:
        x_fft, y_fft = x_in, y_in

    # Get the frequency coordinates corresponding to the FFT grid
    fx = np.fft.fftshift(np.fft.fftfreq(x_fft.shape[0], d=x_fft[1] - x_fft[0]))
    fy = np.fft.fftshift(np.fft.fftfreq(y_fft.shape[0], d=y_fft[1] - y_fft[0]))

    # Output coordinates from the FFT-based LCT kernel
    return (lam * b * fx, lam * b * fy)


def lct_fft_kernels(
    spec_in: Array | tuple,
    lam: float,
    ABCD: Array,
    npad: int | tuple[int, int] | None = None,
):
    """
    Compute the FFT-based LCT kernels.

    Parameters
    ----------
    spec_in : Array or tuple
        Input grid specification (CoordSpec).
    lam : float
        Wavelength [m].
    ABCD : Array, shape (2, 2)
        Ray-transfer matrix for this segment.
    npad : int or tuple[int, int] or None
        Padding factor for the input coordinates.

    Returns
    -------
    tuple
        FFT-based LCT kernels
        (pre, fft_offset, post, pref, scale, dx_in, dy_in, spec_out).
    """
    a, b, c, d = ABCD.flatten()

    # Unpack input coordinates
    x_in, y_in = unpack_coord_spec(spec_in)

    # Apply padding if specified and get the FFT grid coordinates
    if npad is not None:
        x_fft, y_fft = pad_coords(x_in, y_in, npad)
    else:
        x_fft, y_fft = x_in, y_in

    # Get FFT grid spacings
    dx_in = x_fft[1] - x_fft[0]
    dy_in = y_fft[1] - y_fft[0]

    # Get the frequency coordinates corresponding to the FFT grid
    fx = np.fft.fftshift(np.fft.fftfreq(x_fft.shape[0], d=dx_in))
    fy = np.fft.fftshift(np.fft.fftfreq(y_fft.shape[0], d=dy_in))

    # Output coordinates from the FFT-based LCT kernel
    x_out = lam * b * fx
    y_out = lam * b * fy
    spec_out = (x_out, y_out)

    # Get the r^2 grids for the input and output coordinates
    r2_in = r2_coords((x_fft, y_fft))
    r2_out = r2_coords(spec_out)

    # Calculate the chirps
    pre = np.exp(1j * np.pi * a * r2_in / (lam * b))
    post = np.exp(1j * np.pi * d * r2_out / (lam * b))

    # FFT offset to account for non-centered starting coordinates in the input grid
    FX, FY = np.meshgrid(fx, fy, indexing="xy")
    fft_offset = np.exp(-2j * np.pi * (FX * x_fft[0] + FY * y_fft[0]))

    # LCT prefactor
    pref = 1.0 / (1j * lam * b)

    # Pixel-scale ratio for discrete normalisation
    dx_out = x_out[1] - x_out[0]
    dy_out = y_out[1] - y_out[0]
    scale = np.sqrt((dx_out * dy_out) / (dx_in * dy_in))

    # Return all the bits
    return pre, fft_offset, post, pref, scale, dx_in, dy_in, spec_out


def lct_fft_kernel_prop(
    u_pad: Array,
    pre: Array,
    fft_offset: Array,
    post: Array,
    pref: complex,
    scale: float,
    dx_in: float,
    dy_in: float,
    apply_out_curv: bool = True,
) -> Array:
    """
    Apply the FFT-based LCT kernels to a padded input field.

    Parameters
    ----------
    u_pad : Array
        Padded input field, shape (Ny_in_padded, Nx_in_padded).
    pre : Array
        Input-plane quadratic phase exp(i π a r^2 / (λ b)).
    fft_offset : Array
        FFT offset phase to account for non-zero starting coordinates.
    post : Array
        Output-plane quadratic phase exp(i π d r'^2 / (λ b)).
    pref : complex
        LCT prefactor 1 / (i λ b).
    scale : float
        Pixel-scale ratio for discrete normalisation.
    dx_in : float
        Input grid spacing along x.
    dy_in : float
        Input grid spacing along y.
    apply_out_curv : bool, optional
        Whether to apply the output quadratic phase term.

    Returns
    -------
    Array
        Output field after applying the FFT-based LCT kernels.
    """
    # Apply the pre-chirp to the padded input field
    u_tmp = pre * u_pad

    # Compute the FFT and apply the offset
    u_f = dx_in * dy_in * np.fft.fftshift(np.fft.fft2(u_tmp))
    u_f = u_f * fft_offset

    # Apply the post-chirp and prefactor if specified
    if apply_out_curv:
        return pref * scale * post * u_f
    else:
        return pref * scale * u_f


def lct_prop_fft(
    u_in: Array,
    spec_in: Array | tuple,
    lam: float,
    ABCD: Array,
    npad: int | tuple[int, int] | None = None,
    apply_out_curv: bool = True,
):
    """
    High-level LCT using an FFT-based implementation.

    Parameters
    ----------
    u_in : Array
        Input field, shape (Ny_in, Nx_in).
    spec_in : Array or tuple
        Input grid specification (CoordSpec).
    lam : float
        Wavelength [m].
    ABCD : Array, shape (2, 2)
        Ray-transfer matrix for this segment.
    npad : int or tuple[int, int] or None
        Padding factor for the input coordinates. If None, no padding is applied.
    apply_out_curv : bool, optional
        Whether to apply the output quadratic phase term.

    Returns
    -------
    tuple
        (u_out, spec_out) where:
        - u_out is the output field on the grid defined by spec_out.
        - spec_out is the output coordinate specification corresponding to the FFT grid.
    """
    # Compute the FFT-based LCT kernels and output spec
    pre, fft_offset, post, pref, scale, dx_in, dy_in, spec_out = lct_fft_kernels(
        spec_in=spec_in,
        lam=lam,
        ABCD=ABCD,
        npad=npad,
    )

    # Padding factor of the input coordinates.
    if npad is not None:
        u_pad = pad_to(u_in, npad)
    else:
        u_pad = u_in

    # Apply the FFT-based LCT kernels to the padded input field
    u_out = lct_fft_kernel_prop(
        u_pad=u_pad,
        pre=pre,
        fft_offset=fft_offset,
        post=post,
        pref=pref,
        scale=scale,
        dx_in=dx_in,
        dy_in=dy_in,
        apply_out_curv=apply_out_curv,
    )

    return u_out, spec_out


# ---------------------------------------------------------------------
# 2-FT Frequency Spec Generator
# ---------------------------------------------------------------------


def lct_2ft_spec_freq(spec_in: Array | tuple, freq_factor: float = 2.0) -> tuple:
    """
    Construct the intermediate Fourier-grid specification for the 2-FT LCT.

    The 2-FT factorisation uses an explicit spatial-frequency grid
    (fx, fy) for the intermediate Fourier-plane quadratic phase. This
    helper builds a frequency-grid CoordSpec from the input sampling,
    extending beyond the strict Nyquist sampling by an oversampling
    factor to reduce wrap-around / aliasing artefacts.

    Parameters
    ----------
    spec_in : Array or tuple
        Input grid specification in CoordSpec conventions.
    freq_factor : float, optional
        Oversampling factor applied to the Fourier grid relative to the
        Nyquist frequency implied by the input sampling. Values > 1
        reduce aliasing at the cost of larger intermediate arrays.
        Default is 2.0.

    Returns
    -------
    spec_freq : tuple
        Frequency-grid specification in CoordSpec conventions:
        ``((Nfx, Nfy), (dfx, dfy))`` where:
        - Nfx, Nfy are the number of frequency samples in x and y,
        - dfx, dfy are the frequency spacings in cycles per metre [1/m].

    Notes
    -----
    - The grid spans approximately [-f_Nyq, +f_Nyq] in each axis, with
      spacing reduced by `freq_factor`, increasing N accordingly.
    - This function does not validate downstream sampling of the
      Fourier-plane chirp; it only defines the intermediate grid size.
    """

    x_in, y_in = unpack_coord_spec(spec_in)

    Nx_in = x_in.shape[0]
    Ny_in = y_in.shape[0]

    dx_in = x_in[1] - x_in[0]
    dy_in = y_in[1] - y_in[0]

    # Nyquist sampling the grid
    fnyq_x = 1.0 / (2.0 * dx_in)
    fnyq_y = 1.0 / (2.0 * dy_in)

    dfx = 1.0 / (Nx_in * dx_in * freq_factor)  # to not have annoying aliasing
    dfy = 1.0 / (Ny_in * dy_in * freq_factor)

    Nfx = int(np.ceil(2 * fnyq_x / dfx + 1))
    Nfy = int(np.ceil(2 * fnyq_y / dfy + 1))

    spec_freq = ((Nfx, Nfy), (dfx, dfy))

    return spec_freq


# ---------------------------------------------------------------------
# 2-FT LCT Method Kernel Builder
# ---------------------------------------------------------------------


def lct_2ft_kernels(
    spec_in: Array | tuple,
    spec_out: Array | tuple,
    spec_freq: Array | tuple,
    lam: float,
    ABCD: Array,
) -> tuple:
    """
    Build kernels for the 2-FT factorisation of the LCT.

    This implements the two-Fourier-transform (2FT) decomposition of the
    Linear Canonical Transform as described in the Fienup (2024) tutorial.
    The Collins integral is factorised as:

        FT  →  quadratic phase in frequency  →  inverse FT  →  output chirp

    The intermediate Fourier-domain quadratic phase corresponds to the
    LCT kernel expressed in spatial-frequency coordinates.

    Parameters
    ----------
    spec_in : Array or tuple
        Input grid specification (CoordSpec conventions).
    spec_out : Array or tuple
        Output grid specification (CoordSpec conventions).
    spec_freq : Array or tuple
        Intermediate frequency-grid specification (CoordSpec conventions),
        for example produced by `lct_2ft_spec_freq(spec_in, ...)`.
        This defines the sampling of the Fourier-plane chirp.
    lam : float
        Wavelength [m].
    ABCD : Array, shape (2, 2)
        Ray-transfer matrix [[a, b], [c, d]].
        This implementation assumes a ≠ 0.

    Returns
    -------
    Kx_f, Ky_f : Array
        Forward MFT kernels mapping (x, y) → (fx, fy).
    H : Array
        Fourier-plane quadratic phase exp[-i π (λ b / a)(fx² + fy²)].
    Kx_i, Ky_i : Array
        Inverse MFT kernels mapping (fx, fy) → (x_out / a, y_out / a).
    post : Array
        Output-plane quadratic phase exp[i π c r'^2 / (λ a)].
    pref : float
        Overall scalar prefactor 1 / a (global phase omitted).
    scale : float
        Pixel-scale ratio for discrete energy normalisation.

    Notes
    -----
    - This factorisation is often numerically advantageous when the
      Fourier-domain quadratic phase is easier to sample than the
      spatial-domain chirps in the basic Collins form.
    - If a ≈ 0 the transform becomes ill-conditioned in this form and
      the standard Collins (b-based) factorisation should be used.
    """
    a, b, c, _ = ABCD.flatten()

    x_in, y_in = unpack_coord_spec(spec_in)
    x_out, y_out = unpack_coord_spec(spec_out)
    x_freq, y_freq = unpack_coord_spec(spec_freq)

    dx_in = x_in[1] - x_in[0]
    dy_in = y_in[1] - y_in[0]
    dx_out = x_out[1] - x_out[0]
    dy_out = y_out[1] - y_out[0]
    dx_freq = x_freq[1] - x_freq[0]
    dy_freq = y_freq[1] - y_freq[0]

    r2_freq = r2_coords(spec_freq)

    # Fourier-plane LCT chirp
    H = np.exp(-1j * np.pi * (lam * b / a) * (r2_freq))

    # Include physical Riemann weights dx,dy here.
    Kx_f, Ky_f = mft_kernels(
        spec_in=spec_in,
        spec_out=spec_freq,
        alpha=-2.0 * np.pi,
        weight=(dx_in, dy_in),
    )

    # Inverse MFT: (fx,fy) -> (x/A, y/A)
    # Include physical Riemann weights dfx,dfy here.
    Kx_i, Ky_i = mft_kernels(
        spec_in=spec_freq,
        spec_out=(x_out / a, y_out / a),
        alpha=2.0 * np.pi,
        weight=(dx_freq, dy_freq),
    )

    # output chirp
    r2_out = r2_coords(spec_out)

    post = np.exp(
        1j * np.pi * c * r2_out / (lam * a)
    )  # could also use "quad_phase" here?

    pref = 1 / a  # ignoring the e^ikL0 constant phase factor

    scale = np.sqrt((dx_out * dy_out) / (dx_in * dy_in))

    # returning in the order that operations are applied (Fourier chirp in the middle)
    return Kx_f, Ky_f, H, Kx_i, Ky_i, post, pref, scale


# ---------------------------------------------------------------------
# Cached application
# ---------------------------------------------------------------------


def lct_2ft_kernel_prop(
    u_in: Array,
    Kx_f: Array,
    Ky_f: Array,
    H: Array,  # Fourier domain chirp
    Kx_i: Array,
    Ky_i: Array,
    post: Array,  # Output chirp
    pref: float,
    scale: float,
    apply_out_curv: bool = True,  # perhaps "post" should be an option if this is also an option...
) -> Array:
    """
    Apply a precomputed 2-FT LCT.

    Implements the sequence:

        u_f   = MFT_forward[u_in]
        u_fc  = u_f * H
        u_out = MFT_inverse[u_fc]
        result = pref * scale * post * u_out

    Parameters
    ----------
    u_in : Array
        Input complex field, shape (Ny_in, Nx_in).
    Kx_f, Ky_f : Array
        Forward MFT kernels mapping the input grid to the Fourier grid.
    H : Array
        Fourier-domain quadratic phase.
    Kx_i, Ky_i : Array
        Inverse MFT kernels mapping the Fourier grid to the scaled
        output coordinates.
    post : Array
        Output-plane quadratic phase.
    pref : float
        Scalar prefactor (1 / a in this implementation).
    scale : float
        Pixel-scale ratio for discrete normalisation.
    apply_out_curv : bool, optional
        If True, applies the output quadratic phase `post`. If False,
        returns the field prior to this curvature term.

    Returns
    -------
    Array
        Output field on the grid defined by spec_out.

    Notes
    -----
    - This function assumes all kernels were constructed consistently
      using `lct_2ft_kernels`.
    - No sampling checks are performed; aliasing control depends on the
      chosen spec_freq during kernel construction.
    """
    u_f = mft(u_in, Kx_f, Ky_f)  # go into Fourier domain

    u_fc = u_f * H  # apply the chirp

    u_2ft = mft(u_fc, Kx_i, Ky_i)  # go back to the physical domain

    if apply_out_curv:
        return pref * scale * post * u_2ft
    else:
        return pref * scale * u_2ft


# ---------------------------------------------------------------------
# High-level 2-FT LCT application
# ---------------------------------------------------------------------


def lct_prop_2ft(
    u_in: Array,
    spec_in: Array | tuple,
    spec_out: Array | tuple,
    spec_freq: Array | tuple,
    lam: float,
    ABCD: Array,
    apply_out_curv: bool = True,
) -> Array:
    """
    High-level 2-FT Linear Canonical Transform.

    Performs an LCT using the two-Fourier-transform factorisation
    described in the Fienup (2024) tutorial. This provides an
    alternative to `lct_prop_basic`, expressing the transform as:

        MFT → quadratic phase (frequency domain) → inverse MFT
        → output chirp.

    Parameters
    ----------
    u_in : Array
        Input field, shape (Ny_in, Nx_in).
    spec_in : Array or tuple
        Input grid specification (CoordSpec conventions).
    spec_out : Array or tuple
        Output grid specification (CoordSpec conventions).
    spec_freq : Array or tuple
        Intermediate frequency-grid specification (CoordSpec conventions),
        for example produced by `lct_2ft_spec_freq(spec_in, ...)`.
        This defines the sampling of the Fourier-plane chirp.
    lam : float
        Wavelength [m].
    ABCD : Array, shape (2, 2)
        Ray-transfer matrix for the propagation segment.
        This method assumes a ≠ 0.
    apply_out_curv : bool, optional
        Whether to apply the output quadratic phase term.

    Returns
    -------
    Array
        Output field on the grid defined by spec_out.

    Notes
    -----
    - Compared to the basic Collins implementation, this factorisation
      moves the dominant quadratic phase into the Fourier domain, which
      can be advantageous for certain sampling regimes.
    - No automatic validation of LCT sampling conditions is performed.
      Users should verify that both spatial and frequency grids are
      adequately sampled for their chosen ABCD segment.
    """
    Kx_f, Ky_f, H, Kx_i, Ky_i, post, pref, scale = lct_2ft_kernels(
        spec_in=spec_in,
        spec_out=spec_out,
        spec_freq=spec_freq,
        lam=lam,
        ABCD=ABCD,
    )  # Not sure if this is the best way to do it but it makes sense for now.

    u_out = lct_2ft_kernel_prop(
        u_in=u_in,
        Kx_f=Kx_f,
        Ky_f=Ky_f,
        H=H,
        Kx_i=Kx_i,
        Ky_i=Ky_i,
        post=post,
        pref=pref,
        scale=scale,
        apply_out_curv=apply_out_curv,
    )

    return u_out


# ---------------------------------------------------------------------
# 2-FT LCT: FFT --> MFT propagator
# ---------------------------------------------------------------------

# ---------------------------------------------------------------------
# Generating kernels
# ---------------------------------------------------------------------


def lct_2ft_fft_kernels(
    spec_in: Array | tuple,
    spec_out: Array | tuple,
    lam: float,
    ABCD: Array,
    npad: int | tuple[int, int] | None = None,
) -> tuple:
    """
    Build kernels for the FFT --> MFT factorisation of the LCT.

    This implements a mixed two-transform decomposition of the
    Linear Canonical Transform, where the forward transform from the
    input plane to the Fourier plane is performed with a discrete FFT,
    and the return from the Fourier plane to the output plane is
    performed with an MFT. The factorisation is:

        FFT  →  quadratic phase in frequency  →  inverse MFT  →  output chirp

    The intermediate Fourier-domain quadratic phase corresponds to the
    LCT kernel expressed in spatial-frequency coordinates, with the
    Fourier grid fixed by the native FFT frequencies of the input grid.

    Parameters
    ----------
    spec_in : Array or tuple
        Input grid specification (CoordSpec conventions).
    spec_out : Array or tuple
        Output grid specification (CoordSpec conventions).
    lam : float
        Wavelength [m].
    ABCD : Array, shape (2, 2)
        Ray-transfer matrix [[a, b], [c, d]].
        This implementation assumes a ≠ 0.

    Returns
    -------
    H : Array
        Fourier-plane quadratic phase
        exp[-i π (λ b / a)(fx² + fy²)] evaluated on the FFT frequency grid.
    Kx_i, Ky_i : Array
        Inverse MFT kernels mapping (fx, fy) → (x_out / a, y_out / a).
    post : Array
        Output-plane quadratic phase exp[i π c r'^2 / (λ a)].
    pref : float
        Overall scalar prefactor 1 / a (global phase omitted).
    scale : float
        Pixel-scale ratio for discrete energy normalisation.

    Notes
    -----
    - Unlike `lct_2ft_kernels`, this routine does not require an explicit
      intermediate frequency-grid specification. The Fourier grid is set
      directly by the FFT sampling implied by `spec_in`.
    - This can be more efficient than a forward MFT when the input grid is
      uniformly sampled and the FFT frequency sampling is adequate for the
      Fourier-plane chirp.
    - If a ≈ 0 the transform becomes ill-conditioned in this form and
      the standard Collins (b-based) factorisation should be used.
    """

    a, b, c, _ = ABCD.flatten()

    x_in, y_in = unpack_coord_spec(spec_in)
    x_out, y_out = unpack_coord_spec(spec_out)

    # Padding factor of the input coordinates.
    if npad is not None:
        x_fft, y_fft = pad_coords(x_in, y_in, npad)
    else:
        x_fft, y_fft = x_in, y_in

    dx_in = x_fft[1] - x_fft[0]
    dy_in = y_fft[1] - y_fft[0]
    dx_out = x_out[1] - x_out[0]
    dy_out = y_out[1] - y_out[0]

    Nx_in = x_fft.shape[0]
    Ny_in = y_fft.shape[0]

    fx = np.fft.fftshift(np.fft.fftfreq(Nx_in, d=dx_in))
    fy = np.fft.fftshift(np.fft.fftfreq(Ny_in, d=dy_in))
    spec_freq = (fx, fy)

    dfx = fx[1] - fx[0]
    dfy = fy[1] - fy[0]

    r2_freq = r2_coords(spec_freq)

    # Quadratic phase kernel: Also include offset shifted FFT!!
    # construct grid
    # I'm 99.9% sure this is the correct way to do this.
    # otherwise you cook yourself with an incorrect/uncentered fourier transform.
    FX, FY = np.meshgrid(fx, fy, indexing="xy")
    fft_offset = np.exp(-2j * np.pi * (FX * x_fft[0] + FY * y_fft[0]))

    H = np.exp(-1j * np.pi * (lam * b / a) * r2_freq) * fft_offset

    Kx_i, Ky_i = mft_kernels(
        spec_in=spec_freq,
        spec_out=(x_out / a, y_out / a),
        alpha=2.0 * np.pi,
        weight=(dx_in * dfx, dy_in * dfy),  # for scaling with FFT
    )

    r2_out = r2_coords(spec_out)
    post = np.exp(1j * np.pi * c * r2_out / (lam * a))

    pref = 1 / a
    scale = np.sqrt((dx_out * dy_out) / (dx_in * dy_in))

    return H, Kx_i, Ky_i, post, pref, scale


# ---------------------------------------------------------------------
# Cached application
# ---------------------------------------------------------------------


def lct_2ft_fft_kernel_prop(
    u_pad: Array,
    H: Array,  # Fourier domain chirp
    Kx_i: Array,
    Ky_i: Array,
    post: Array,  # Output chirp
    pref: float,
    scale: float,
    apply_out_curv: bool = True,
) -> Array:
    """
    Apply a precomputed FFT --> MFT 2-FT LCT.

    Implements the sequence:

        u_f   = FFT[u_in]
        u_fc  = u_f * H
        u_out = MFT_inverse[u_fc]
        result = pref * scale * post * u_out

    Parameters
    ----------
    u_in : Array
        Input complex field, shape (Ny_in, Nx_in).
    H : Array
        Fourier-domain quadratic phase.
    Kx_i, Ky_i : Array
        Inverse MFT kernels mapping the FFT frequency grid to the scaled
        output coordinates.
    post : Array
        Output-plane quadratic phase.
    pref : float
        Scalar prefactor (1 / a in this implementation).
    scale : float
        Pixel-scale ratio for discrete normalisation.
    apply_out_curv : bool, optional
        If True, applies the output quadratic phase `post`. If False,
        returns the field prior to this curvature term.

    Returns
    -------
    Array
        Output field on the grid defined by `spec_out`.

    Notes
    -----
    - This function assumes all kernels were constructed consistently
      using `lct_2ft_fft_kernels`.
    - The forward transform uses the native FFT frequency sampling of
      the input grid, so aliasing control depends on whether that FFT
      grid is sufficient for the Fourier-plane chirp.
    """
    u_f = np.fft.fftshift(np.fft.fft2(u_pad))
    u_fc = u_f * H
    u_2ft = mft(u_fc, Kx_i, Ky_i)

    if apply_out_curv:
        return pref * scale * post * u_2ft
    else:
        return pref * scale * u_2ft


# ---------------------------------------------------------------------
# High-level 2-FT FFT --> MFT LCT application
# ---------------------------------------------------------------------


def lct_prop_2ft_fft(
    u_in: Array,
    spec_in: Array | tuple,
    spec_out: Array | tuple,
    lam: float,
    ABCD: Array,
    npad: int | tuple[int, int] | None = None,
    apply_out_curv: bool = True,
) -> Array:
    """
    High-level 2-FT Linear Canonical Transform (FFT --> MFT version).

    Performs an LCT using a mixed two-transform factorisation, in which
    the forward transform to the Fourier plane is evaluated with an FFT
    and the return to the output plane is evaluated with an MFT. This
    provides an alternative to `lct_prop_basic`, expressing the
    transform as:

        FFT → quadratic phase (frequency domain) → inverse MFT
        → output chirp.

    Parameters
    ----------
    u_in : Array
        Input field, shape (Ny_in, Nx_in).
    spec_in : Array or tuple
        Input grid specification (CoordSpec conventions).
    spec_out : Array or tuple
        Output grid specification (CoordSpec conventions).
    lam : float
        Wavelength [m].
    ABCD : Array, shape (2, 2)
        Ray-transfer matrix for the propagation segment.
        This method assumes a ≠ 0.
    apply_out_curv : bool, optional
        Whether to apply the output quadratic phase term.

    Returns
    -------
    Array
        Output field on the grid defined by `spec_out`.

    Notes
    -----
    - Compared to the basic Collins implementation, this factorisation
      moves the dominant quadratic phase into the Fourier domain, which
      can be advantageous for certain sampling regimes.
    - Compared to `lct_prop_2ft`, this version fixes the intermediate
      Fourier grid to the native FFT frequencies of the input sampling,
      avoiding the need for an explicit `spec_freq`.
    - No automatic validation of LCT sampling conditions is performed.
      Users should verify that both the FFT frequency grid and the
      output-plane sampling are adequate for their chosen ABCD segment.
    """
    H, Kx_i, Ky_i, post, pref, scale = lct_2ft_fft_kernels(
        spec_in=spec_in,
        spec_out=spec_out,
        lam=lam,
        ABCD=ABCD,
        npad=npad,
    )

    # Let us pad the actual field in the propagator.
    u_pad = pad_to(u_in, npad)

    u_out = lct_2ft_fft_kernel_prop(
        u_pad=u_pad,
        H=H,
        Kx_i=Kx_i,
        Ky_i=Ky_i,
        post=post,
        pref=pref,
        scale=scale,
        apply_out_curv=apply_out_curv,
    )

    return u_out


# ---------------------------------------------------------------------
# Residual LCT: curvature factoring + high-level interface
# ---------------------------------------------------------------------


def lct_prop_curv(
    u_in: Array,
    spec_in: Array | tuple,
    spec_out: Array | tuple,
    lam: float,
    ABCD: Array,
    curv_in: float | None = None,
    curv_out: float | None = None,
    mode: str = "physical",
    strip_input: bool = True,
    return_residual: bool = False,
) -> Array:
    """
    High-level LCT with optional curvature factoring.

    This wraps lct_prop_basic(...) and allows you to factor out
    spherical-like quadratic phases at the input and output planes.

    Modes
    -----
    mode = "physical"
        Treat curv_in as a physical input curvature 1/R_in. If curv_in
        is None, assume a flat wavefront (curv_in = 0). The output
        curvature is then propagated via:

            curv_out = propagate_curv(ABCD, curv_in).

    mode = "cancel"
        Ignore any user-supplied curvatures and choose:

            curv_in  = -a / b,
            curv_out =  d / b,

        so that the residual LCT has no quadratic phases (a_res = d_res = 0).
        This maximises curvature sampling headroom, but the residual
        field's phase no longer directly matches the physical wavefront.

    mode = "manual"
        Use the user-supplied curv_in and curv_out directly.

    Parameters
    ----------
    u_in : Array
        Input field. If strip_input=True, this is interpreted as the
        physical field U_in and we build a residual u_in by removing
        the input curvature. If strip_input=False, u_in is assumed to
        already be residual and is passed into the LCT directly.
    spec_in : Array or tuple
        Input grid specification (CoordSpec). Typical examples:
          - (N, d)           → square uniform grid
          - ((Nx, Ny), (dx, dy))
          - x                → 1D vector, symmetric x=y
          - (x, y)           → explicit per-axis coords
        The lengths implied by spec_in must be compatible with
        u_in.shape; otherwise matrix multiplications will fail.
    lam : float
        Wavelength [m].
    ABCD : Array, shape (2, 2)
        Physical ABCD matrix.
    curv_in, curv_out : float or None
        Input/output curvatures 1/R [1/m], depending on mode.
    mode : {"physical", "cancel", "manual"}
        Curvature handling strategy.
    strip_input : bool
        If True, remove Q_in from the input before applying the
        residual LCT. If False, assume u_in is already residual.
    return_residual : bool
        If True, return the residual output field u_out (without
        applying Q_out). If False, return the physical field
        U_out = Q_out * u_out.

    Returns
    -------
    Array
        Output field, either residual or physical depending on
        return_residual.
    """

    # factorise curvatures and build residual ABCD matrix
    ABCD_res, curv_in, curv_out = factorise_curv(ABCD, curv_in, curv_out, mode)

    # Build residual input field (if needed)
    if strip_input and curv_in != 0.0:
        u_res_in = remove_curv(u_in, spec_in, lam, curv_in)
    else:
        u_res_in = u_in

    # Apply core LCT on residual ABCD
    u_res_out = lct_prop(
        u_in=u_res_in,
        spec_in=spec_in,
        spec_out=spec_out,
        lam=lam,
        ABCD=ABCD_res,
    )

    if return_residual:
        return u_res_out

    # Re-apply output curvature to get physical field
    if curv_out != 0.0:
        return apply_curv(u_res_out, spec_out, lam, curv_out)

    return u_res_out


# ---------------------------------------------------------------------
# Sampling helpers (to be refined to p≥1 semantics later)
# ---------------------------------------------------------------------


def lct_sampling_parameter(x_in: Array, x_out: Array, lam: float, b):
    """Soummer-style sampling measure for the LCT kernel via mft_sampling_parameter."""
    alpha = -2.0 * np.pi / (lam * b)
    return mft_sampling_parameter(x_in, x_out, alpha)


def lct_phase_parameter(x_in: Array, x_out: Array, lam: float, b):
    """Phase-slope sampling parameters (p_in, p_out) for the LCT kernel via mft_phase_parameter."""
    alpha = -2.0 * np.pi / (lam * b)
    return mft_phase_parameter(x_in, x_out, alpha)


def lct_sampling_metrics(
    lam: float,
    ABCD: Array,
    spec_in,
    spec_out,
    curv_in=None,
    curv_out=None,
    mode=None,
) -> dict:
    """
    Sampling diagnostics for the basic (Collins-form) LCT.

    Computes sampling margins for both the separable Fourier/MFT kernel and
    the quadratic phase (chirp) factors appearing in the standard Collins
    factorisation:

        pre-chirp → MFT kernel → post-chirp

    Two types of metrics are reported:

    - s_* : Fourier kernel sampling factors.
        These come from `mft_sampling_parameter` and measure the sampling
        adequacy of the oscillatory MFT kernel exp[-i α x x'] on the chosen
        (x_in, x_out) and (y_in, y_out) grids. Values greater than 1
        indicate adequate sampling.

    - p_* : Quadratic phase sampling factors.
        These measure how well the quadratic chirps are sampled, using the
        maximum phase slope at the grid edge. For a 1D quadratic phase
            φ(x) = π * curv * x^2 / λ,
        the analytic derivative is
            dφ/dx = 2π * curv * x / λ.
        Evaluating at x = X_max and multiplying by Δx gives the maximum
        phase step per pixel at the edge:
            Δφ_max ≈ (2π * curv * X_max / λ) * Δx.
        The reported sampling factor is the Nyquist margin
            p = π / Δφ_max
            = λ / (2 * curv * X_max * Δx).
        With this convention, p ≥ 1 indicates the quadratic phase is
        adequately Nyquist-sampled (larger is better).

    Parameters
    ----------
    lam : float
        Wavelength [m].
    ABCD : Array, shape (2, 2)
        LCT matrix [[a, b], [c, d]].
    spec_in : CoordSpec
        Input grid specification.
    spec_out : CoordSpec
        Output grid specification.
    curv_in, curv_out : float or None, optional
        Input/output curvatures for residualised propagation. If `mode` is
        provided, the ABCD matrix is first factorised via `factorise_curv`
        before computing sampling metrics.
    mode : str or None, optional
        Curvature factorisation mode (see curvature.py).

    Returns
    -------
    dict
        {
        "s_x", "s_y",          # Fourier kernel sampling
        "p_in_x", "p_out_x",   # input/output chirp sampling (x)
        "p_in_y", "p_out_y",   # input/output chirp sampling (y)
        }

    Notes
    -----
    - Kernel sampling uses α = 2π / (λ b) (sign does not affect adequacy).
    - Chirp sampling uses effective Collins curvatures:
        curv_in  = a / b,
        curv_out = d / b.
    - Interpretation: s > 1 and p > 1 indicate adequate sampling margins.
    """

    if mode is not None:
        ABCD, _, _ = factorise_curv(ABCD, curv_in, curv_out, mode)

    a, b, c, d = ABCD.flatten()

    x_in, y_in = unpack_coord_spec(spec_in)
    x_out, y_out = unpack_coord_spec(spec_out)

    # Kernel sampling (Soummer-style, via MFT helper)
    alpha = 2.0 * np.pi / (lam * b)  # sign irrelevant for magnitude
    s_x = mft_sampling_parameter(x_in, x_out, alpha)
    s_y = mft_sampling_parameter(y_in, y_out, alpha)

    # Quadratic chirps
    # we can use curvature sampling parameter here
    # this was wrong algebraically, now it should be correct
    input_curv = a / b
    output_curv = d / b

    p_in_x = curvature_sampling_parameter(x_in, lam, input_curv)
    p_in_y = curvature_sampling_parameter(y_in, lam, input_curv)

    p_out_x = curvature_sampling_parameter(x_out, lam, output_curv)
    p_out_y = curvature_sampling_parameter(y_out, lam, output_curv)

    return {
        "s_x": s_x,
        "s_y": s_y,
        "p_in_x": p_in_x,
        "p_out_x": p_out_x,
        "p_in_y": p_in_y,
        "p_out_y": p_out_y,
    }


def lct_2ft_sampling_metrics(
    lam: float,
    ABCD: Array,
    spec_in: Array | tuple,
    spec_out: Array | tuple,
    spec_freq: Array | tuple,
) -> dict:
    """
    Sampling diagnostics for the 2-FT LCT factorisation.

    Evaluates sampling margins for each stage of the two-Fourier-transform
    decomposition:

        FT → Fourier-domain quadratic phase → inverse FT → output chirp

    Metrics returned:

    - s_f_* : Sampling factors for the forward Fourier transform kernel.
    - p_f_* : Sampling factors for the Fourier-domain quadratic phase.
    - s_i_* : Sampling factors for the inverse Fourier transform kernel.
    - p_out_* : Sampling factors for the final spatial quadratic phase.

    Definitions
    -----------
    s parameters:
        Fourier/MFT sampling factors from `mft_sampling_parameter`, which
        quantify adequacy of sampling of exp[-i 2π x f]-type kernels on the
        corresponding coordinate pairs. Values greater than 1 indicate
        adequate sampling.

    p parameters:
        Quadratic phase sampling factors computed from the Nyquist margin
            p = π / Δφ_max,
        where Δφ_max is the maximum phase increment per pixel at the grid
        edge estimated from the analytic derivative. For
            φ(u) = π * curv * u^2 / λ_eff,
        with u = x or f, the closed-form edge expression is
            p = λ_eff / (2 * curv * U_max * Δu).
        With this convention, p ≥ 1 indicates adequate Nyquist sampling of
        the quadratic phase (larger is better).

    Parameters
    ----------
    lam : float
        Wavelength [m].
    ABCD : Array, shape (2, 2)
        LCT matrix [[a, b], [c, d]] (assumes a ≠ 0).
    spec_in : CoordSpec
        Input spatial grid specification.
    spec_out : CoordSpec
        Output spatial grid specification.
    spec_freq: CoordSpec
        Frequency spatial grid specification.

    Returns
    -------
    dict
        {
        "s_f_x", "s_f_y",      # forward FT sampling
        "p_f_x", "p_f_y",      # Fourier-domain chirp sampling
        "s_i_x", "s_i_y",      # inverse FT sampling
        "p_out_x", "p_out_y",  # final spatial chirp sampling
        }

    Notes
    -----
    - The Fourier-domain quadratic phase uses effective curvature
        curv_f = -b / a
    and effective wavelength
        λ_eff = 1 / λ
    (consistent with your implementation via `curvature_sampling_parameter`).
    - The final spatial chirp uses curvature curv_out = c / a.
    - Interpretation: s > 1 and p > 1 indicate adequate sampling margins.
    """
    a, b, c, _ = ABCD.flatten()

    x_in, y_in = unpack_coord_spec(spec_in)
    x_out, y_out = unpack_coord_spec(spec_out)

    x_freq, y_freq = unpack_coord_spec(spec_freq)

    alpha_f = -2.0 * np.pi

    s_f_x = mft_sampling_parameter(x_in, x_freq, alpha_f)
    s_f_y = mft_sampling_parameter(y_in, y_freq, alpha_f)

    # Now the curvature in the frequency domain
    # Basically we have a curvature given by -b/a
    # But our wavelength should be inverse (1/lam)
    # then the algebra is correct

    curv_f = -b / a

    p_f_x = curvature_sampling_parameter(x_freq, 1 / lam, curv_f)
    p_f_y = curvature_sampling_parameter(y_freq, 1 / lam, curv_f)

    # Now the sampling of the inverse fourier transform
    alpha_i = 2.0 * np.pi

    s_i_x = mft_sampling_parameter(x_freq, x_out, alpha_i)
    s_i_y = mft_sampling_parameter(y_freq, y_out, alpha_i)

    # now the sampling of the final chirp
    curv_out = c / a

    p_out_x = curvature_sampling_parameter(x_out, lam, curv_out)
    p_out_y = curvature_sampling_parameter(y_out, lam, curv_out)

    return {
        "s_f_x": s_f_x,
        "s_f_y": s_f_y,
        "p_f_x": p_f_x,
        "p_f_y": p_f_y,
        "s_i_x": s_i_x,
        "s_i_y": s_i_y,
        "p_out_x": p_out_x,
        "p_out_y": p_out_y,
    }


# Building the FFT coord spec to be used in the sampling metrics
# Useful in running lct_2ft_sampling_metrics for the FFT version
def lct_2ft_fft_spec_freq(
    spec_in: Array | tuple,
) -> tuple:
    """
    Construct the FFT frequency-grid specification for the 2-FT FFT method.

    Generates the spatial-frequency coordinate specification corresponding
    to the native FFT frequencies of the input grid. This grid is used as
    the intermediate Fourier plane in the FFT → MFT implementation of the
    two-Fourier-transform LCT factorisation.

    The returned coordinates follow the same centred convention used by
    the propagation routines:

        fx = fftshift(fftfreq(Nx, dx))
        fy = fftshift(fftfreq(Ny, dy))

    Parameters
    ----------
    spec_in : Array or tuple
        Input grid specification (CoordSpec conventions). The coordinate
        spacing and array lengths determine the FFT frequency sampling.

    Returns
    -------
    tuple
        Frequency-grid specification `(fx, fy)` compatible with the
        CoordSpec convention, representing the spatial-frequency
        coordinates associated with the FFT of the input grid.

    Notes
    -----
    - This helper is primarily intended for use with
      `lct_2ft_sampling_metrics` when evaluating sampling diagnostics for
      the FFT → MFT propagator.
    - The resulting grid corresponds exactly to the discrete frequencies
      produced by `fftshift(fft2(ifftshift(...)))`.
    """
    x_in, y_in = unpack_coord_spec(spec_in)

    dx_in = x_in[1] - x_in[0]
    dy_in = y_in[1] - y_in[0]

    Nx_in = x_in.shape[0]
    Ny_in = y_in.shape[0]

    fx = np.fft.fftshift(np.fft.fftfreq(Nx_in, d=dx_in))
    fy = np.fft.fftshift(np.fft.fftfreq(Ny_in, d=dy_in))
    spec_freq = (fx, fy)

    return spec_freq
