import numpy as np
import pytest

from he11lib.modes import LGBasis


W0 = 5e-3  # 5 mm waist
Z0 = 0.5  # waist at 0.5 m
WAVELENGTH = 1.76e-3  # ~170 GHz gyrotron


def make_grid(w, half_widths_in_w=6.0, n=300):
    extent = half_widths_in_w * w
    coords = np.linspace(-extent, extent, n)
    dx = coords[1] - coords[0]
    x, y = np.meshgrid(coords, coords)
    return x, y, dx


def test_beam_radius_at_waist_equals_w0():
    basis = LGBasis(w0=W0, z0=Z0, wavelength=WAVELENGTH)
    assert basis.beam_radius(Z0) == pytest.approx(W0)


def test_beam_radius_at_one_rayleigh_range():
    basis = LGBasis(w0=W0, z0=Z0, wavelength=WAVELENGTH)
    zR = np.pi * W0**2 / WAVELENGTH
    assert basis.beam_radius(Z0 + zR) == pytest.approx(W0 * np.sqrt(2))


def test_gouy_phase_zero_at_waist():
    basis = LGBasis(w0=W0, z0=Z0, wavelength=WAVELENGTH)
    assert basis.gouy_phase(Z0, p=0, l=0) == pytest.approx(0.0)
    assert basis.gouy_phase(Z0, p=1, l=2) == pytest.approx(0.0)


def test_gouy_phase_at_one_rayleigh_range_for_fundamental():
    basis = LGBasis(w0=W0, z0=Z0, wavelength=WAVELENGTH)
    zR = np.pi * W0**2 / WAVELENGTH
    # order (2p+|l|+1) = 1 for p=0,l=0; atan(1) = pi/4
    assert basis.gouy_phase(Z0 + zR, p=0, l=0) == pytest.approx(np.pi / 4)


def test_fundamental_mode_matches_analytic_gaussian_at_waist():
    basis = LGBasis(w0=W0, z0=Z0, wavelength=WAVELENGTH)
    x, y, _ = make_grid(W0, n=50)
    r2 = x**2 + y**2

    field = basis.field(x, y, Z0, p=0, l=0)
    expected_intensity_shape = np.exp(-2 * r2 / W0**2)
    intensity = np.abs(field) ** 2

    # shapes should match up to a constant normalization factor
    ratio = intensity / expected_intensity_shape
    ratio_center = ratio[len(ratio) // 2, len(ratio) // 2]
    np.testing.assert_allclose(ratio / ratio_center, 1.0, atol=1e-6)


def test_mode_is_normalized_at_waist():
    basis = LGBasis(w0=W0, z0=Z0, wavelength=WAVELENGTH)
    x, y, dx = make_grid(W0, n=400)

    field = basis.field(x, y, Z0, p=0, l=0)
    total_power = np.sum(np.abs(field) ** 2) * dx**2
    assert total_power == pytest.approx(1.0, rel=2e-3)


def test_mode_is_normalized_away_from_waist():
    basis = LGBasis(w0=W0, z0=Z0, wavelength=WAVELENGTH)
    z = Z0 + 0.2
    w_z = basis.beam_radius(z)
    x, y, dx = make_grid(w_z, n=400)

    field = basis.field(x, y, z, p=1, l=2)
    total_power = np.sum(np.abs(field) ** 2) * dx**2
    assert total_power == pytest.approx(1.0, rel=2e-3)


def test_modes_are_orthogonal():
    basis = LGBasis(w0=W0, z0=Z0, wavelength=WAVELENGTH)
    z = Z0 + 0.1
    w_z = basis.beam_radius(z)
    x, y, dx = make_grid(w_z, n=400)

    field_a = basis.field(x, y, z, p=0, l=0)
    field_b = basis.field(x, y, z, p=1, l=0)
    inner_product = np.sum(field_a * np.conj(field_b)) * dx**2
    assert abs(inner_product) < 1e-3


def test_project_recovers_known_coefficients():
    basis = LGBasis(w0=W0, z0=Z0, wavelength=WAVELENGTH)
    z = Z0 + 0.15
    w_z = basis.beam_radius(z)
    x, y, dx = make_grid(w_z, n=400)

    true_coeffs = {(0, 0): 0.8 + 0.1j, (1, 0): 0.2 - 0.3j}
    field = basis.field_superposition(x, y, z, true_coeffs)

    recovered = basis.project(field, x, y, dx, z, modes=list(true_coeffs.keys()))

    for mode, coeff in true_coeffs.items():
        assert recovered[mode] == pytest.approx(coeff, abs=5e-3)