"""
Angular Spectrum Method (ASM)
=============================

Paraxial (Fresnel) angular spectrum propagation for scalar fields in free
space on rectangular grids.

This module provides:

  - Sampling diagnostics for ASM (aliasing / padding heuristics).
  - Transfer-function (kernel) construction on the FFT grid.
  - Low-level FFT-based propagation engine on padded rectangular grids.
  - High-level user-facing propagation functions operating on grids
    specified by CoordSpec conventions (`spec_in`).

All propagation here is *paraxial*: the kernel corresponds to the Fresnel
approximation in spatial-frequency (fx, fy) space. Full non-paraxial ASM is
not included; if needed in future it should live in a dedicated module.

Conventions
-----------
- Input grids are **rectangular** and follow CoordSpec conventions
  (`spec_in`), describing the x and y coordinates of the sampled field:

      u.shape == (Ny, Nx)

- Pixel spacings are inferred from the coordinate specification.

- Padding is explicit via an integer or tuple `npad`, applied per axis
  when constructing the FFT propagation grid.

- The FFT-based ASM propagator is **unitary in ℓ²** on the padded grid
  (for |H| = 1), so it preserves discrete throughput up to truncation
  from cropping. No explicit scaling factor S is applied in the FFT-only
  path.

- The hybrid ASM→MFT path introduces a pixel-scale ratio between the FFT
  grid and the output spatial grid; this is corrected via a scalar S to
  keep the discrete energy consistent.
"""

from __future__ import annotations

from typing import Dict

import jax.numpy as np
from jax import Array

from .coords import unpack_coord_spec, pad_to, pad_coords, crop_to
from .mft import mft, mft_kernels
from .curvature import r2_coords

__all__ = [
    "asm_sampling_parameter",
    "asm_diagnostics",
    "asm_kernels",
    "asm_kernel_prop",
    "asm_prop",
    "asm_mft_kernels",
    "asm_mft_kernel_prop",
    "asm_prop_mft",
]

# ---------------------------------------------------------------------
# ASM sampling diagnostics
# ---------------------------------------------------------------------


def asm_sampling_parameter(
    N: int,
    d: float,
    lam: float,
    z: float,
) -> float:
    """
    ASM aliasing Nyquist ratio p_alias for a square grid.

    This is a simple heuristic derived from the alias-limited diameter
    D_alias ≈ sqrt(N * λ * z). For a grid with:

        N    : number of points per axis (square),
        d    : pixel spacing [m],
        D    : physical size = N * d,

    we define

        p_alias = D / D_alias = D / sqrt(N * λ * z).

    Interpretation
    --------------
      - p_alias >= 1  → grid is large enough (in principle) to avoid
                        strong wrap-around aliasing in ASM for distance z.
      - p_alias <  1  → aliasing likely; padding (larger N_fft) is
                        recommended.

    Parameters
    ----------
    N : int
        Unpadded grid size (N x N).
    d : float
        Pixel spacing [m].
    lam : float
        Wavelength [m].
    z : float
        Propagation distance [m].

    Returns
    -------
    float
        p_alias Nyquist ratio.
    """
    N_int = int(N)
    dx = float(d)
    lam = float(lam)
    z = float(z)

    D = N_int * dx
    D_alias = np.sqrt(N_int * lam * z)
    p_alias = D / D_alias

    return float(p_alias)


def asm_diagnostics(
    N: int,
    d: float,
    lam: float,
    z: float,
) -> Dict[str, float]:
    """
    Compute ASM-specific diagnostic quantities and suggested padding.

    This function focuses on FFT/ASM aliasing behaviour, not on generic
    beam-spread or Fresnel-number physics (those should live elsewhere).

    Parameters
    ----------
    N : int
        Unpadded grid size (N x N).
    d : float
        Pixel spacing [m].
    lam : float
        Wavelength [m].
    z : float
        Propagation distance [m].

    Returns
    -------
    dict
        {
          "N": int,
          "dx": float,
          "D": float,
          "f_max": float,
          "lambda_fmax": float,
          "D_alias": float,
          "p_alias": float,
          "Q_alias_float": float,
          "Q_alias_int": int,
        }

    Interpretation
    --------------
    - D_alias is the alias-limited characteristic size (~ sqrt(N λ z)).
    - p_alias = D / D_alias; p_alias >= 1 is numerically safer.
    - Q_alias_float = D_alias / D is the padding factor needed to
      reach p_alias ≈ 1; Q_alias_int = ceil(Q_alias_float).
    """
    N_int = int(N)
    dx = float(d)
    lam = float(lam)
    z = float(z)

    D = N_int * dx
    f_max = 1.0 / (2.0 * dx)
    lambda_fmax = lam * f_max

    D_alias = float(np.sqrt(N_int * lam * z))
    p_alias = float(D / D_alias)
    Q_alias_float = float(D_alias / D)
    Q_alias_int = int(np.ceil(Q_alias_float))

    return {
        "N": N_int,
        "dx": dx,
        "D": D,
        "f_max": f_max,
        "lambda_fmax": lambda_fmax,
        "D_alias": D_alias,
        "p_alias": p_alias,
        "Q_alias_float": Q_alias_float,
        "Q_alias_int": Q_alias_int,
    }


# ---------------------------------------------------------------------
# FFT-grid kernel builder
# ---------------------------------------------------------------------


def asm_kernels(
    spec_in: Array | tuple,
    lam: float,
    z: float,
    npad: int | tuple[int, int] | None = None,
) -> tuple:
    """
    Build the paraxial ASM propagation kernel on a shifted FFT grid.

    The input spatial coordinate specification is optionally padded, and
    the corresponding shifted FFT frequency coordinates are constructed
    along each axis with `fftfreq(...)`. The returned kernel is
    the Fresnel angular-spectrum transfer function evaluated on this
    shifted frequency grid,

        H(fx, fy) = exp[-i π λ z (fx² + fy²)].

    Parameters
    ----------
    spec_in : Array or tuple
        Input coordinate specification, interpreted by
        `unpack_coord_spec`, describing the x- and y-coordinates of the
        input field.
    lam : float
        Wavelength, in metres.
    z : float
        Propagation distance, in metres.
    npad : int or tuple[int, int] or None, optional
        Target padded size passed to `pad_coords`. If `None`, no padding
        is applied and the input coordinates are used directly.

    Returns
    -------
    H : Array
        Complex ASM kernel sampled on the shifted FFT grid, with shape
        `(Ny_fft, Nx_fft)`.
    Nx_in : int
        Number of x-samples in the input propagation grid.
    Ny_in : int
        Number of y-samples in the input propagation grid.
    """
    x_in, y_in = unpack_coord_spec(spec_in)

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

    # 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]

    Nx_fft = x_fft.shape[0]
    Ny_fft = y_fft.shape[0]

    # We don't need to shift the coordinates here
    fx = np.fft.fftfreq(Nx_fft, d=dx_in)
    fy = np.fft.fftfreq(Ny_fft, d=dy_in)
    spec_freq = (fx, fy)

    r2_freq = r2_coords(spec_freq)

    # Here we simply have the quadratic phase in frequency
    # There is no need for a phase slope in frequency -- the output
    # by definition lives on the same coordiantes as the input
    # (We cannot independently specify a spec_out for the 2fft propagation).
    H = np.exp(-1j * np.pi * lam * z * r2_freq)

    return H, Nx_in, Ny_in


# ---------------------------------------------------------------------
# Low-level FFT engine
# ---------------------------------------------------------------------


def asm_kernel_prop(
    u_pad: Array,
    H: Array,
    Nx_in: int,
    Ny_in: int,
    crop: bool = True,
) -> Array:
    """
    Propagate a field using a precomputed shifted-grid ASM kernel.

    The input field is assumed to already be sampled on the padded spatial
    grid associated with `H`. Propagation is performed by taking a 2D FFT,
    multiplying by the precomputed kernel, and inverse transforming back
    to the spatial domain:

        fft2(u_pad)
        -> multiply by H
        -> ifft2(...)

    If requested, the propagated field is then centrally cropped back to
    the specified unpadded size.

    Parameters
    ----------
    u_pad : Array
        Input complex field on the padded spatial grid, with shape
        `(Ny_fft, Nx_fft)`.
    H : Array
        Complex ASM propagation kernel on the correspondingly shifted FFT
        grid, with shape `(Ny_fft, Nx_fft)`.
    Nx_in : int
        Target cropped size along the x-axis.
    Ny_in : int
        Target cropped size along the y-axis.
    crop : bool, optional
        If `True`, centrally crop the propagated field to
        `(Ny_in, Nx_in)`. If `False`, return the full padded propagated
        field.

    Returns
    -------
    Array
        Propagated complex field, either cropped to `(Ny_in, Nx_in)` or
        returned on the full padded grid depending on `crop`.
    """
    u_fft = np.fft.fft2(u_pad)
    u_fftc = u_fft * H
    u_2fft = np.fft.ifft2(u_fftc)

    if crop:
        return crop_to(u_2fft, (Nx_in, Ny_in))
    else:
        return u_2fft


# ---------------------------------------------------------------------
# High-level FFT-only ASM propagator
# ---------------------------------------------------------------------


def asm_prop(
    u_in: Array,
    spec_in: Array,
    lam: float,
    z: float,
    npad: int | tuple[int, int] | None = None,
    crop: bool = True,
) -> Array:
    """
    Propagate a field with the paraxial angular spectrum method.

    This is the high-level interface to the FFT-based ASM propagator. The
    input coordinate specification is used to construct the shifted
    frequency-domain propagation kernel, optionally on a padded grid.
    The input field is padded to match that propagation grid, propagated
    with `asm_kernel_prop`, and optionally cropped back to the original
    field size.

    Parameters
    ----------
    u_in : Array
        Input complex field sampled on the coordinates described by
        `spec_in`, with shape `(Ny_in, Nx_in)`.
    spec_in : Array
        Input coordinate specification, interpreted by
        `unpack_coord_spec`.
    lam : float
        Wavelength, in metres.
    z : float
        Propagation distance, in metres.
    npad : int or tuple[int, int] or None, optional
        Target padded size passed to both the coordinate-padding and
        array-padding routines. If `None`, no padding is applied.
    crop : bool, optional
        If `True`, crop the propagated field back to the original input
        size. If `False`, return the propagated field on the padded grid.

    Returns
    -------
    Array
        Propagated complex field, either on the original input grid or on
        the padded grid depending on `crop`.
    """
    H, Nx_in, Ny_in = asm_kernels(
        spec_in=spec_in,
        lam=lam,
        z=z,
        npad=npad,
    )

    u_pad = pad_to(u_in, npad)

    u_out = asm_kernel_prop(
        u_pad=u_pad,
        H=H,
        Nx_in=Nx_in,
        Ny_in=Ny_in,
        crop=crop,
    )

    return u_out


# ---------------------------------------------------------------------
# ASM FFT --> MFT kernel builder
# ---------------------------------------------------------------------


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

    The factorisation is:

        FFT  →  quadratic phase in frequency  →  inverse MFT

    The intermediate Fourier-domain quadratic phase corresponds to the
    ASM 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].
    z : float
        Propagation distance [m].

    Returns
    -------
    H : Array
        Fourier-plane quadratic phase
        exp[-i π λ z(fx² + fy²)] evaluated on the FFT frequency grid.
    Kx_i, Ky_i : Array
        Inverse MFT kernels mapping (fx, fy) → (x_out, y_out).
    scale : float
        Pixel-scale ratio for discrete energy normalisation.
    """

    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 * z * r2_freq) * fft_offset

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

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

    return H, Kx_i, Ky_i, scale


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


def asm_mft_kernel_prop(
    u_pad: Array,
    H: Array,  # Fourier domain chirp
    Kx_i: Array,
    Ky_i: Array,
    scale: float,
) -> Array:
    """
    Apply a precomputed FFT --> MFT ASM propagation.

    Implements the sequence:

        u_f   = FFT[u_in]
        u_fc  = u_f * H
        u_out = MFT_inverse[u_fc]
        result = scale * 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.
    scale : float
        Pixel-scale ratio for discrete normalisation.

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

    Notes
    -----
    - This function assumes all kernels were constructed consistently
      using `asm_mft_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)

    return scale * u_2ft


# ---------------------------------------------------------------------
# High-level FFT --> MFT ASM application
# ---------------------------------------------------------------------


def asm_prop_mft(
    u_in: Array,
    spec_in: Array | tuple,
    spec_out: Array | tuple,
    lam: float,
    z: float,
    npad: int | tuple[int, int] | None = None,
) -> Array:
    """
    High-level ASM propagation (FFT --> MFT version).

    Performs an ASM 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 `asm_prop_fft`, expressing the
    transform as:

        FFT → quadratic phase (frequency domain) → inverse MFT

    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].
    z : float
        Propagation distance [m].

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

    Notes
    -----
    - No automatic validation of ASM sampling conditions is performed.
      Users should verify that the FFT frequency grid is adequate for
      their chosen propagation distance.
    """
    H, Kx_i, Ky_i, scale = asm_mft_kernels(
        spec_in=spec_in,
        spec_out=spec_out,
        lam=lam,
        z=z,
        npad=npad,
    )

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

    u_out = asm_mft_kernel_prop(
        u_pad=u_pad,
        H=H,
        Kx_i=Kx_i,
        Ky_i=Ky_i,
        scale=scale,
    )

    return u_out
