"""
ABCD matrices for paraxial optical systems.

Small helpers for building and inspecting 2×2 ray-transfer matrices:

    [ x_out ]   [ A  B ] [ x_in ]
    [ θ_out ] = [ C  D ] [ θ_in ]

All functions assume inputs are already valid JAX arrays / floats.
"""

from __future__ import annotations

import jax.numpy as np
from jax import Array

__all__ = [
    "abcd_surface_power",
    "abcd_lens",
    "abcd_mirror",
    "abcd_free_space",
    "abcd_fraunhofer",
    "compose_abcd",
    "is_surface",
    "is_free_space",
    "abcd_effective_focal_length",
    "abcd_paraxial_power",
    "abcd_geometric_magnification",
    "abcd_angular_magnification",
    "abcd_distance",
    "abcd_front_focal_length",
    "abcd_back_focal_length",
    "abcd_unimodularity",
]


# ---------------------------------------------------------------------
# Constructors
# ---------------------------------------------------------------------


def abcd_surface_power(power: float) -> Array:
    """
    Thin powered surface in air with optical power Φ.

    power = Φ [1/m]; for a thin lens Φ = 1/f, for a spherical mirror Φ = 2/R.
    """
    return np.array([[1.0, 0.0], [-power, 1.0]])


def abcd_lens(focal_length: float) -> Array:
    """
    Thin refractive lens in air with focal length f [m].
    """
    return abcd_surface_power(1.0 / focal_length)


def abcd_mirror(radius: float) -> Array:
    """
    Spherical mirror in air with radius of curvature R [m].

    Positive R corresponds to a concave (focusing) mirror.
    """
    power = 2.0 / radius
    return abcd_surface_power(power)


def abcd_free_space(z: float) -> Array:
    """
    Free-space propagation over distance z [m].
    """
    return np.array([[1.0, z], [0.0, 1.0]])


def abcd_fraunhofer(focal_length: float) -> Array:
    """
    ABCD matrix for an ideal Fraunhofer (Fourier) transform.
    """
    return np.array([[0.0, focal_length], [-1.0 / focal_length, 0.0]])


# ---------------------------------------------------------------------
# Composition and classification
# ---------------------------------------------------------------------


def compose_abcd(matrices: list | tuple) -> Array:
    """
    Compose a sequence of ABCD matrices by right-multiplication.

    Given [M1, M2, ..., MN], returns M_sys = MN · ... · M2 · M1.

    All ideal paraxial optical elements satisfy AD - BC = 1.
    The code does not enforce unimodularity, because non-unimodular matrices
    may arise from numerical drift or intentional scaling. Users should inspect
    `abcd_unimodularity(M)` when correctness matters.
    """
    M = np.eye(2)
    for Mi in matrices:
        M = Mi @ M
    return M


def is_surface(M: Array, tol: float = 1e-8) -> bool:
    """
    Return True if M is (numerically) a pure surface with B ≈ 0.
    """
    a, b, c, d = M.flatten()
    return np.abs(b) < tol


def is_free_space(M: Array, tol: float = 1e-8) -> bool:
    """
    Return True if M is (numerically) free-space: [[1, z], [0, 1]].

    Checks:
      - |A - 1| < tol
      - |C|     < tol
      - |D - 1| < tol
      - |B| ≥ tol  (otherwise it's effectively a surface)
    """
    a, b, c, d = M.flatten()
    if np.abs(b) < tol:
        return False
    return (np.abs(a - 1.0) < tol) and (np.abs(c) < tol) and (np.abs(d - 1.0) < tol)


# ---------------------------------------------------------------------
# Derived scalar properties
# ---------------------------------------------------------------------


def abcd_effective_focal_length(M: Array) -> float:
    """
    Effective focal length f_eff [m] for a system that focuses collimated input.

    For such a system C ≠ 0 and f_eff = -1 / C.
    """
    a, b, c, d = M.flatten()
    return -1.0 / c


def abcd_paraxial_power(M: Array) -> float:
    """
    Paraxial optical power Φ [1/m] of an ABCD system, defined via C = -Φ.

    For a thin lens: M = [[1, 0], [-1/f, 1]] ⇒ Φ = 1/f.
    """
    a, b, c, d = M.flatten()
    return -c


def abcd_geometric_magnification(M: Array) -> float:
    """
    Transverse geometric magnification m = A.

    Under imaging conditions, x_out = A x_in for θ_in = 0.
    """
    a, b, c, d = M.flatten()
    return a


def abcd_angular_magnification(M: Array) -> float:
    """Angular magnification m_θ = D."""
    a, b, c, d = M.flatten()
    return d


def abcd_distance(M: Array) -> float:
    """
    Effective propagation distance z_eff for an ABCD system.

    For pure free-space propagation, M = [[1, z], [0, 1]] ⇒ z_eff = B.
    """
    a, b, c, d = M.flatten()
    return b


def abcd_front_focal_length(M: Array) -> float:
    """Front focal length FFL = -A/C (for C ≠ 0)."""
    a, b, c, d = M.flatten()
    return -a / c


def abcd_back_focal_length(M: Array) -> float:
    """Back focal length BFL = -D/C (for C ≠ 0)."""
    a, b, c, d = M.flatten()
    return -d / c


def abcd_unimodularity(M: Array) -> float:
    """
    Return AD - BC (should be 1 for lossless paraxial systems). Useful for debugging composed matrices.
    """
    a, b, c, d = M.flatten()
    return a * d - b * c
