"""
Gaussian beams and q-parameter utilities.

Analytic helpers for tracking a Gaussian-like beam envelope and curvature
through ABCD systems using the complex q-parameter:

    1/q = 1/R - i λ / (π w²)

where:
  - w is the 1/e amplitude radius,
  - R is the wavefront radius of curvature (R = ∞ ⇒ collimated),
  - λ is the wavelength.

Quadratic phase / curvature utilities are provided in curvature.py and
re-exported here for convenience.
"""

from __future__ import annotations

import jax.numpy as np
from jax import Array

from .curvature import curv_from_R

__all__ = [
    "q_from_w_R",
    "w_R_from_q",
    "curv_from_q",
    "propagate_q",
    "gaussian_step",
    "propagate_q_chain",
    "w0_from_diameter",
    "rayleigh_range",
    "divergence_angle",
]


# ---------------------------------------------------------------------
# Basic conversions between (w, R) and q
# ---------------------------------------------------------------------


def q_from_w_R(w: float | Array, R: float | Array, lam: float) -> Array:
    """
    Construct q from 1/e amplitude radius w and curvature R.

    For R = ∞ we recover the collimated case:

        1/q = -i λ / (π w²)  ⇒  q = i π w² / λ

    Parameters
    ----------
    w : float or Array
        1/e amplitude radius [m].
    R : float or Array
        Radius of curvature [m]. R = ∞ → collimated.
    lam : float
        Wavelength [m].

    Returns
    -------
    Array
        Complex q-parameter.
    """
    w = np.asarray(w)
    R = np.asarray(R)

    collimated = np.isinf(R)
    q_coll = 1j * (np.pi * w**2) / lam
    q_gen = 1.0 / (1.0 / R - 1j * lam / (np.pi * w**2))
    return np.where(collimated, q_coll, q_gen)


def w_R_from_q(
    q: Array,
    lam: float,
    tol: float = 1e-8,
) -> tuple[Array, Array]:
    """
    Recover (w, R) from a q-parameter.

    Using:

        1/q = 1/R - i λ / (π w²)

    so:

        w² = -λ / (π Im(1/q)),
        1/R = Re(1/q)  (R = ∞ if Re(1/q) ≈ 0).

    Parameters
    ----------
    q : Array
        Complex q-parameter.
    lam : float
        Wavelength [m].
    tol : float
        Threshold on |Re(1/q)| for treating R as infinite.

    Returns
    -------
    w : Array
        1/e amplitude radius [m].
    R : Array
        Radius of curvature [m] (may be ∞).
    """
    invq = 1.0 / q

    re = np.real(invq)
    im = np.imag(invq)

    # Radius of curvature: R = 1 / Re(1/q), with R = ∞ if Re(1/q) ≈ 0
    R = np.where(np.abs(re) < tol, np.inf, 1.0 / re)

    # 1/e amplitude radius: w² = -λ / (π Im(1/q)), Im(1/q) < 0
    w2 = -lam / (np.pi * im)
    w = np.sqrt(w2)

    return w, R


def curv_from_q(q: Array, lam: float) -> Array:
    """
    Curvature 1/R from a q-parameter.

    Uses w_R_from_q(q, lam) and curvature.curv_from_R to remain consistent
    with existing Gaussian-beam and curvature utilities.
    """
    _, R = w_R_from_q(q, lam)
    return curv_from_R(R)


# ---------------------------------------------------------------------
# Propagation through ABCD
# ---------------------------------------------------------------------


def propagate_q(M: Array, q: Array) -> Array:
    """
    Propagate a q-parameter through a single ABCD segment:

        q' = (A q + B) / (C q + D).
    """
    a, b, c, d = M.flatten()
    return (a * q + b) / (c * q + d)


def gaussian_step(M: Array, lam: float, q_in: Array) -> dict:
    """
    Propagate a Gaussian-like envelope through one ABCD segment.

    Returns input/output (w, R, q) and the envelope magnification M_w.

    Parameters
    ----------
    M : Array
        ABCD matrix for this segment.
    lam : float
        Wavelength [m].
    q_in : Array
        Input q-parameter.

    Returns
    -------
    dict
        {
          "q_in", "q_out",
          "w_in", "w_out",
          "R_in", "R_out",
          "M_w",
        }
    """
    w_in, R_in = w_R_from_q(q_in, lam)
    q_out = propagate_q(M, q_in)
    w_out, R_out = w_R_from_q(q_out, lam)

    M_w = w_out / w_in

    return {
        "q_in": q_in,
        "q_out": q_out,
        "w_in": w_in,
        "w_out": w_out,
        "R_in": R_in,
        "R_out": R_out,
        "M_w": M_w,
    }


def propagate_q_chain(M_list: list | tuple, q0: Array) -> list:
    """
    Propagate q through a list of ABCD segments.

    M_list is a list or tuple of 2x2 ABCD matrices [M1, M2, ...].
    Returns [q0, q1, ..., qN].
    """
    q = q0
    out: list = [q0]
    for M in M_list:
        q = propagate_q(M, q)
        out.append(q)
    return out


# ---------------------------------------------------------------------
# Simple aperture-based initialisation and beam scalings
# ---------------------------------------------------------------------


def w0_from_diameter(diameter: float | Array) -> Array:
    """
    Estimate a beam "waist" w0 from an aperture diameter D.

    We assume the beam reaches 1/e amplitude at radius D/2:

        w0 = (D/2) / sqrt(2).
    """
    D = np.asarray(diameter)
    return 0.5 * D / np.sqrt(2.0)


def rayleigh_range(w0: float | Array, lam: float) -> Array:
    """
    Rayleigh range z_R for waist radius w0:

        z_R = π w0² / λ.
    """
    w0 = np.asarray(w0)
    return np.pi * w0**2 / lam


def divergence_angle(w0: float | Array, lam: float) -> Array:
    """
    Far-field half-angle divergence for a Gaussian waist w0:

        θ = λ / (π w0).
    """
    w0 = np.asarray(w0)
    return lam / (np.pi * w0)
