"""
Fraunhofer propagation via MFT
==============================

This module provides Fraunhofer (far-field) propagation using the unitless MFT
kernels defined in :mod:`mft`. All normalisation is handled through a single
physics-derived scaling factor S returned together with the kernels.

Coordinate specifications
-------------------------
Both input and output grids are specified using the same conventions as
coords.unpack_coord_spec. For any "spec" argument (spec_in, spec_out) we accept:

1) Uniform centred grid:

       spec = (N, d)
       spec = ((Nx, Ny), (dx, dy))

   where N is int or (int, int) and d is float/int or (float/int, float/int).
   Coordinates are generated with dLux.utils.nd_coords for each axis.

2) Single 1D coordinate vector:

       spec = x    # Array, shape (N,)

   interpreted as symmetric 2D grid with x and y identical.

3) Explicit per-axis coordinates:

       spec = (x, y)  # tuple of Arrays, each shape (N,)

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

The same spec API works for both uniform and non-uniform grids. The
Fraunhofer scaling S is computed from the effective spacings inferred
from the coordinates (dx = x[1] - x[0], etc.).
"""

from __future__ import annotations

import jax.numpy as np
from jax import Array

from .coords import unpack_coord_spec
from .mft import mft, mft_kernels


__all__ = [
    "fraunhofer_kernels",
    "fraunhofer_kernel_prop",
    "fraunhofer_prop",
]

# ---------------------------------------------------------------------
# Internal helpers
# ---------------------------------------------------------------------


def _log_safe_mul(*vals: float) -> float:
    """Compute product(vals) as exp(sum(log(vals))) for numerical stability."""
    arr = np.array(vals, dtype=float)
    return np.exp(np.sum(np.log(arr)))


def _fraunhofer_scale_from_coords(
    x_in: Array,
    y_in: Array,
    x_out: Array,
    y_out: Array,
    lam: float,
    f: float,
) -> float:
    """
    Compute the discrete Fraunhofer scaling factor S from coordinates:

        S = sqrt(dx_in * dx_out * dy_in * dy_out) / (lam * f)

    where dx_in, dy_in, dx_out, dy_out are inferred as first differences of
    the 1D coordinate vectors.

    Parameters
    ----------
    x_in, y_in : Array
        1D input coordinates along x and y.
    x_out, y_out : Array
        1D output coordinates along x and y.
    lam : float
        Wavelength [m].
    f : float
        Focal length [m].

    Returns
    -------
    float
        Scaling factor S.
    """
    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]

    num = (_log_safe_mul(dx_in, dx_out) * _log_safe_mul(dy_in, dy_out)) ** 0.5
    den = lam * f
    return num / den


# ---------------------------------------------------------------------
# Kernel builder (spec-based, unified)
# ---------------------------------------------------------------------


def fraunhofer_kernels(
    spec_in: Array | tuple,
    spec_out: Array | tuple,
    lam: float,
    f: float,
) -> tuple[float, Array, Array]:
    """
    Build Fraunhofer (S, Kx, Ky) for arbitrary separable grids.

    This unifies the previous "uniform" and "_from_coords" variants. The
    input and output specs 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.

    Parameters
    ----------
    spec_in : Array or tuple
        Input grid specification, see module docstring.
    spec_out : Array or tuple
        Output grid specification, see module docstring.
    lam : float
        Wavelength [m].
    f : float
        Focal length [m].

    Returns
    -------
    S  : float
        Discrete throughput-preserving scaling factor.
    Kx : Array
    Ky : Array
        Unitless MFT kernels.
    """
    x_in, y_in = unpack_coord_spec(spec_in)
    x_out, y_out = unpack_coord_spec(spec_out)

    alpha = -2.0 * np.pi / (lam * f)
    weight = 1.0  # unitless kernels; all physics in S

    # Build MFT kernels from specs (re-uses the same CoordSpec API)
    Kx, Ky = mft_kernels(
        spec_in=spec_in,
        spec_out=spec_out,
        alpha=alpha,
        weight=weight,
    )

    S = _fraunhofer_scale_from_coords(
        x_in=x_in,
        y_in=y_in,
        x_out=x_out,
        y_out=y_out,
        lam=lam,
        f=f,
    )

    return S, Kx, Ky


# ---------------------------------------------------------------------
# Propagation
# ---------------------------------------------------------------------


def fraunhofer_kernel_prop(
    u_pupil: Array,
    S: float,
    Kx: Array,
    Ky: Array,
) -> Array:
    """
    Fraunhofer propagation using precomputed (S, Kx, Ky).

    Parameters
    ----------
    u_pupil : Array
        Pupil-plane field (Ny_in, Nx_in).
    S : float
        Discrete Fraunhofer scaling factor.
    Kx, Ky : Array
        Unitless MFT kernels.

    Returns
    -------
    Array
        Focal-plane field.
    """
    u_fft = mft(u_pupil, Kx, Ky)
    u_focal = S * u_fft
    return u_focal


def fraunhofer_prop(
    u_pupil: Array,
    spec_in: Array | tuple,
    spec_out: Array | tuple,
    lam: float,
    f: float,
) -> Array:
    """
    High-level Fraunhofer propagation between arbitrary separable grids.

    Parameters
    ----------
    u_pupil : Array
        Pupil-plane field, shape (Ny_in, Nx_in).
    spec_in : Array or tuple
        Input grid specification for the pupil plane. 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 of x_in and y_in implied by spec_in must be compatible
        with u_pupil.shape; otherwise matrix multiplications will fail.
    spec_out : Array or tuple
        Output grid specification for the focal plane (same conventions).
    lam : float
        Wavelength [m].
    f : float
        Focal length [m].

    Returns
    -------
    Array
        Focal-plane field on the grid defined by spec_out.

    Notes
    -----
    - No sampling checks are performed here. Use the MFT sampling
      diagnostics together with your coordinate specs to check whether
      the chosen grids are Nyquist-sampled.
    - Fraunhofer is strictly valid in the far-field / small-angle regime.
      For more general Fresnel propagation, use the LCT/ASM modules.
    """
    S, Kx, Ky = fraunhofer_kernels(
        spec_in=spec_in,
        spec_out=spec_out,
        lam=lam,
        f=f,
    )

    return fraunhofer_kernel_prop(
        u_pupil=u_pupil,
        S=S,
        Kx=Kx,
        Ky=Ky,
    )
