"""
Spherical wavefront curvature and quadratic phase utilities.

This module provides helpers for working with spherical-like wavefronts
parameterised by a curvature 1/R [1/m]:

    Q(x, y) = exp(i π curv (x² + y²) / λ),

as well as ABCD-based curvature propagation and simple sampling metrics.

All functions assume separable, rectilinear grids (equal spacing per axis,
but dx and dy may differ). Coordinates are handled via coords.unpack_*.
"""

from __future__ import annotations

import jax.numpy as np
from jax import Array
from .coords import unpack_coord_spec, unpack_size, unpack_scale, nd_coords

__all__ = [
    "r2_coords",
    "r2_from_N_d",
    "quad_phase",
    "apply_curv",
    "remove_curv",
    "curv_from_R",
    "residual_abcd",
    "residual_curv_cancel",
    "propagate_curv",
    "factorise_curv",
    "lct_1ft_curv",
    "lct_2ft_curv",
    "curvature_sampling_parameter",
]

# ---------------------------------------------------------------------
# r² helpers (single source of truth)
# ---------------------------------------------------------------------


def r2_coords(spec_in: Array | tuple) -> Array:
    """
    Return r² = x² + y² on a separable grid.

    Parameters
    ----------
    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.

    Returns
    -------
    Array
        r² grid with shape (Ny, Nx) in (y, x) ordering.
    """
    x, y = unpack_coord_spec(spec_in)
    return (y**2)[:, None] + (x**2)[None, :]


def r2_from_N_d(
    N: int | tuple[int, int],
    d: float | tuple,
) -> Array:
    """
    Build r² = x² + y² from grid size and sampling.

    This is a convenience wrapper around r2_coords that constructs
    x, y via dLux.utils.nd_coords for consistent centering.

    Parameters
    ----------
    N : int or (Nx, Ny)
        Grid size.
    d : float or (dx, dy)
        Sampling [m]. Integers are also accepted and cast to float.

    Returns
    -------
    Array
        r² grid with shape (Ny, Nx).
    """
    Nx, Ny = unpack_size(N)
    dx, dy = unpack_scale(d)

    x = nd_coords(Nx, dx)
    y = nd_coords(Ny, dy)
    return r2_coords((x, y))


# ---------------------------------------------------------------------
# Quadratic phase and curvature conversions
# ---------------------------------------------------------------------


def quad_phase(
    coords: Array | tuple,
    lam: float,
    curv: float | Array,
) -> Array:
    """
    Spherical-like quadratic phase:

        Q(x, y) = exp(i π curv (x² + y²) / λ),

    where curv = 1/R is the wavefront curvature [1/m].

    Parameters
    ----------
    coords : Array or (x, y)
        Coordinate specification for the grid.
    lam : float
        Wavelength [m].
    curv : float or Array
        Wavefront curvature 1/R [1/m].

    Returns
    -------
    Array
        Complex phase factor Q with shape (Ny, Nx).
    """
    r2 = r2_coords(coords)
    return np.exp(1j * np.pi * curv * r2 / lam)


def apply_curv(
    u: Array,
    coords: Array | tuple,
    lam: float,
    curv: float | Array,
) -> Array:
    """
    Apply a spherical-like quadratic phase of curvature 1/R.

    Parameters
    ----------
    u : Array
        Input field, shape (Ny, Nx).
    coords : Array or (x, y)
        Coordinate specification for the grid.
    lam : float
        Wavelength [m].
    curv : float or Array
        Wavefront curvature 1/R [1/m].

    Returns
    -------
    Array
        Physical field U = Q(curv) * u.
    """
    return u * quad_phase(coords, lam, curv)


def remove_curv(
    u: Array,
    coords: Array | tuple,
    lam: float,
    curv: float | Array,
) -> Array:
    """
    Remove a spherical-like quadratic phase of curvature 1/R.

    Parameters
    ----------
    u : Array
        Input field, shape (Ny, Nx).
    coords : Array or (x, y)
        Coordinate specification for the grid.
    lam : float
        Wavelength [m].
    curv : float or Array
        Wavefront curvature 1/R [1/m].

    Returns
    -------
    Array
        Residual field u_res = Q(curv)^* * u.
    """
    return u * quad_phase(coords, lam, curv).conj()


def curv_from_R(R: float | Array) -> Array:
    """
    Curvature 1/R [1/m] from radius of curvature R [m].

    For R = ∞, returns 0 (collimated). Implemented with np.where for
    JAX-compatibility.

    Parameters
    ----------
    R : float or Array
        Radius of curvature [m].

    Returns
    -------
    Array
        Curvature 1/R [1/m].
    """
    R_arr = np.asarray(R)
    return np.where(np.isinf(R_arr), 0.0, 1.0 / R_arr)


# ---------------------------------------------------------------------
# ABCD-based curvature propagation
# ---------------------------------------------------------------------


def residual_abcd(
    ABCD: Array,
    curv_in: float,
    curv_out: float,
) -> Array:
    """
    Residual ABCD matrix after factoring input/output curvatures.

    We factor out quadratic phases

        Q_in(x)  = exp(i π curv_in  x² / λ),
        Q_out(x) = exp(i π curv_out x² / λ),

    from the Collins integral. The residual LCT has matrix:

        a_res = a + b * curv_in,
        d_res = d - b * curv_out,
        c_res = (a_res * d_res - 1) / b,

    with the same b. This preserves unimodularity exactly.
    """
    a, b, c, d = ABCD.flatten()

    a_res = a + b * curv_in
    d_res = d - b * curv_out
    c_res = (a_res * d_res - 1.0) / b

    return np.array([[a_res, b], [c_res, d_res]])


def residual_curv_cancel(ABCD: Array) -> tuple[float, float]:
    """
    Choose curvatures that cancel the LCT's intrinsic quadratics.

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

    which yields a_res = 0 and d_res = 0, i.e. the residual LCT has
    no quadratic phases and is purely kernel-like.

    This maximises curvature sampling headroom, but does not correspond
    to the physical beam curvature unless the system truly removes that
    curvature.
    """
    a, b, c, d = ABCD.flatten()
    curv_in = -a / b
    curv_out = d / b
    return float(curv_in), float(curv_out)


def propagate_curv(ABCD: Array, curv_in: float | Array) -> Array:
    """
    Propagate curvature 1/R through an ABCD segment.

    Uses the scalar curvature law:

        curv_out = (C + D * curv_in) / (A + B * curv_in),

    where ABCD = [[A, B], [C, D]].

    Parameters
    ----------
    ABCD : Array, shape (2, 2)
        Ray-transfer matrix.
    curv_in : float or Array
        Input curvature 1/R_in [1/m].

    Returns
    -------
    Array
        Output curvature 1/R_out [1/m].
    """
    a, b, c, d = ABCD.flatten()
    return (c + d * curv_in) / (a + b * curv_in)


def factorise_curv(
    ABCD: Array, curv_in: float | Array, curv_out: float | Array, mode: str = "physical"
) -> tuple:
    """
    Factorise input and output curvatures.

    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
    ----------
    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.

    Returns
    -------
    Tuple
        Factored ABCD matrix, and resulting input and output curvatures
    """
    a, b, c, d = ABCD.flatten()

    # Decide curvatures
    if mode == "physical":
        if curv_in is None:
            curv_in = 0.0
        if curv_out is None:
            curv_out = propagate_curv(ABCD, curv_in)
    elif mode == "cancel":
        curv_in, curv_out = residual_curv_cancel(ABCD)
    elif mode == "manual":
        if curv_in is None or curv_out is None:
            raise ValueError("mode='manual' requires curv_in and curv_out.")
    else:
        raise ValueError(
            f"Unknown mode '{mode}'. Use 'physical', 'cancel', or 'manual'."
        )

    # Build residual ABCD
    ABCD_res = residual_abcd(ABCD, curv_in, curv_out)

    return (ABCD_res, curv_in, curv_out)


# ---------------------------------------------------------------------
# Curvature factors for different propagations
# ---------------------------------------------------------------------
def lct_1ft_curv(
    ABCD: Array,
) -> tuple:
    a, b, c, d = ABCD.flatten()

    in_curv = a / b
    out_curv = d / b

    return in_curv, out_curv


def lct_2ft_curv(
    ABCD: Array,
) -> tuple:

    a, b, c, d = ABCD.flatten()
    fourier_curv = -b / a  # note that: when using "apply_curv" function in the fourier
    # domain you want lam to actually be inverse: 1/lam.
    # perhaps I will modify that function for Fourier bool = False something like that.
    out_curv = c / a

    return fourier_curv, out_curv


# ---------------------------------------------------------------------
# Curvature sampling diagnostics
# ---------------------------------------------------------------------


def curvature_sampling_parameter(
    x: Array,
    lam: float,
    curv: float | Array,
    eps: float = 1e-30,
) -> Array:
    """
    Nyquist ratio for a pure curvature phase on a 1D grid.

    For a phase

        φ(x) = π curv x² / λ,

    the edge phase increment is approximately

        Δφ_max ≈ 2 |curv| X dx * (π / λ),

    with X the half-range and dx the pixel spacing.

    This function returns

        p_curv = π / Δφ_max
               = λ / (2 |curv| X dx),

    so that p_curv >= 1 corresponds to Nyquist sampling of that curvature
    on the given grid.

    Parameters
    ----------
    x : Array
        1D coordinate vector along an axis.
    lam : float
        Wavelength [m].
    curv : float or Array
        Curvature 1/R [1/m].
    eps : float
        Small additive term to avoid division by zero in flat cases.

    Returns
    -------
    Array
        Nyquist ratio p_curv (>= 1 is safe).
    """
    X = 0.5 * (x[-1] - x[0])
    dx = x[1] - x[0]
    denom = 2 * np.abs(curv) * X * dx + eps
    return lam / denom
