"""
General Matrix Fourier Transform (MFT) utilities
================================================

This module provides low-level, unitless MFT kernels and application routines.
All physics-specific scaling (dx, λ, f, b, etc.) is applied externally by
Fraunhofer, ASM, LCT, or other Fresnel modules.

All grids are assumed *separable* and *rectilinear*:

  - separable: coordinates factor as x(i), y(j),
  - rectilinear: each axis has a single spacing (equally spaced samples),
  - anisotropic allowed: dx and dy may differ.

Coordinate specifications
-------------------------
We follow 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.
"""

from __future__ import annotations

import jax.numpy as np
from jax import Array

from .coords import unpack_coord_spec

__all__ = [
    "mft",
    "mft_kernels",
    "mft_phase_parameter",
    "mft_sampling_parameter",
]

# ---------------------------------------------------------------------
# 1D kernel (internal)
# ---------------------------------------------------------------------


def _mft_kernel_1d(
    x_in: Array,
    x_out: Array,
    alpha: float,
    weight: float,
) -> Array:
    """
    1D MFT kernel:

        K[j, k] = exp(i * alpha * x_in[k] * x_out[j]) * weight

    Parameters
    ----------
    x_in, x_out : Array
        1D coordinates in input and output planes.
    alpha : float
        Phase coefficient in exp(i * alpha * x * x').
    weight : float
        Input-sample weight (e.g. dx for a physical integral).
    """
    phase = alpha * (x_out[:, None] * x_in[None, :])
    return np.exp(1j * phase) * weight


# ---------------------------------------------------------------------
# 2D separable MFT application
# ---------------------------------------------------------------------


def mft(
    u: Array,
    Kx: Array,
    Ky: Array,
    left_conj: bool = False,
    right_conj: bool = False,
) -> Array:
    """
    Apply a separable 2D MFT:

        tmp = u @ (Kx or Kx.conj()).T
        out = (Ky or Ky.conj()) @ tmp

    Shape conventions
    -----------------
        u   : (Ny_in, Nx_in)
        Kx  : (Nx_out, Nx_in)   -- x-direction kernel
        Ky  : (Ny_out, Ny_in)   -- y-direction kernel
        out : (Ny_out, Nx_out)

    Parameters
    ----------
    u : Array
        Input field.
    Kx, Ky : Array
        1D kernels along x and y.
    left_conj : bool
        If True, conjugate Ky before applying (acts on the left).
    right_conj : bool
        If True, conjugate Kx before applying (acts on the right).

    Returns
    -------
    Array
        Output field on the MFT-defined grid.
    """
    Kx_eff = Kx.conj() if right_conj else Kx
    Ky_eff = Ky.conj() if left_conj else Ky

    # First: transform along x (columns)
    tmp = u @ Kx_eff.T

    # Second: transform along y (rows)
    out = Ky_eff @ tmp

    return out


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


def mft_kernels(
    spec_in: Array | tuple,
    spec_out: Array | tuple,
    alpha: float,
    weight: float | tuple = 1.0,
) -> tuple[Array, Array]:
    """
    Build MFT kernels (Kx, Ky) from coordinate specifications.

    This unifies the old "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.
    alpha : float
        Phase coefficient in exp(i * alpha * x * x').
    weight : float or (float, float)
        Input-sample weight. If scalar, applies to both axes. If tuple
        (wx, wy), wx is used for the x-kernel and wy for the y-kernel.

        For purely unitless kernels, use weight = 1.0.
        For physical Fresnel/LCT kernels, weight is typically dx_in
        (or (dx_in, dy_in)).

    Returns
    -------
    Kx, Ky : Array
        1D MFT kernels along x and y.
    """
    x_in, y_in = unpack_coord_spec(spec_in)
    x_out, y_out = unpack_coord_spec(spec_out)

    if isinstance(weight, (int, float)):
        wx = wy = float(weight)
    else:
        wx, wy = weight

    Kx = _mft_kernel_1d(x_in, x_out, alpha, wx)
    Ky = _mft_kernel_1d(y_in, y_out, alpha, wy)

    return Kx, Ky


# ---------------------------------------------------------------------
# MFT sampling diagnostics
# ---------------------------------------------------------------------


def mft_phase_parameter(
    x_in: Array,
    x_out: Array,
    alpha: float,
    eps: float = 1e-30,
) -> tuple[Array, Array]:
    """
    Phase-slope sampling parameters for the MFT kernel:

        K[j, k] = exp(i * alpha * x_in[k] * x_out[j]).

    For a step of one pixel in k (input index) at fixed j, the phase step is

        Δφ_in = alpha * dx_in * x_out.

    The worst case occurs at the largest |x_out|.

    Similarly, for a step in j (output index) at fixed k,

        Δφ_out = alpha * dx_out * x_in,

    with worst case at the largest |x_in|.

    We define Nyquist ratios (normalised by π):

        p_in  = π / max |Δφ_in|,
        p_out = π / max |Δφ_out|,

    so that:

        p >= 1  → safe (Nyquist-sampled),
        p <  1  → aliasing possible.

    This function is shift-safe: it only depends on max |x|, so grids
    that are shifted relative to zero are handled correctly.

    Parameters
    ----------
    x_in, x_out : Array
        1D input and output coordinates.
    alpha : float
        Phase coefficient in exp(i * alpha * x * x').
    eps : float
        Small term to avoid division by zero in flat cases.

    Returns
    -------
    p_in, p_out : Array
        Nyquist ratios for input and output directions.
    """
    dx_in = x_in[1] - x_in[0]
    dx_out = x_out[1] - x_out[0]

    xmax_in = np.max(np.abs(x_in))
    xmax_out = np.max(np.abs(x_out))

    dphi_in_max = np.abs(alpha) * xmax_out * dx_in
    dphi_out_max = np.abs(alpha) * xmax_in * dx_out

    # avoid division by zero in perfectly flat cases
    p_in = np.where(dphi_in_max > 0, np.pi / (dphi_in_max + eps), np.inf)
    p_out = np.where(dphi_out_max > 0, np.pi / (dphi_out_max + eps), np.inf)

    return p_in, p_out


def mft_sampling_parameter(
    x_in: Array,
    x_out: Array,
    alpha: float,
    eps: float = 1e-30,
) -> Array:
    """
    Single Nyquist ratio for the MFT kernel sampling.

    This is a summary metric built from mft_phase_parameter:

        p_in, p_out = mft_phase_parameter(...)

    and we return

        p = min(p_in, p_out),

    so that:

        p >= 1  → kernel is Nyquist-sampled in both input and output
                  directions;
        p <  1  → aliasing expected somewhere in the kernel.

    Parameters
    ----------
    x_in, x_out : Array
        1D input and output coordinates.
    alpha : float
        Phase coefficient in exp(i * alpha * x * x').
    eps : float
        Small term to avoid division by zero in flat cases.

    Returns
    -------
    p : Array
        Overall Nyquist ratio for the kernel.
    """
    p_in, p_out = mft_phase_parameter(x_in, x_out, alpha, eps=eps)
    return np.minimum(p_in, p_out)
