import numpy as np
import pytest

from he11lib.geometry import CameraModel
from he11lib.modes import LGBasis
from he11lib.synthetic import SyntheticBeamGenerator


W0 = 5e-3
Z0 = 0.5
WAVELENGTH = 1.76e-3
PIXEL_SCALE = 2e-4  # 0.2 mm/px, achieved at z=Z0
CAMERA_DISTANCE = 5.0  # camera stands 5 m upstream of the output window
IMAGE_SHAPE = (161, 161)  # odd so there's a well-defined center pixel


def make_camera(pixel_scale=PIXEL_SCALE, z0=Z0, camera_distance=CAMERA_DISTANCE):
    focal_length_px = (camera_distance + z0) / pixel_scale
    return CameraModel(
        focal_length_px=focal_length_px,
        position=(0.0, 0.0, -camera_distance),
        orientation_deg=(0.0, 0.0, 0.0),
    )


def make_generator():
    basis = LGBasis(w0=W0, z0=Z0, wavelength=WAVELENGTH)
    return SyntheticBeamGenerator(basis=basis, camera=make_camera())


def test_generate_returns_planes_with_requested_z():
    gen = make_generator()
    z_list = [0.3, 0.4, 0.5]
    planes = gen.generate(coefficients={(0, 0): 1 + 0j}, z_list=z_list, image_shape=IMAGE_SHAPE)

    assert [p.z for p in planes] == z_list
    assert all(p.flux.shape == IMAGE_SHAPE for p in planes)


def test_generate_pure_mode_peak_at_image_center_when_centered():
    gen = make_generator()
    planes = gen.generate(
        coefficients={(0, 0): 1 + 0j}, z_list=[Z0], image_shape=IMAGE_SHAPE, center=(0.0, 0.0)
    )
    flux = planes[0].flux

    peak_idx = np.unravel_index(np.argmax(flux), flux.shape)
    center_idx = (IMAGE_SHAPE[0] // 2, IMAGE_SHAPE[1] // 2)
    assert peak_idx == center_idx


def test_generate_applies_center_offset():
    gen = make_generator()
    offset_m = 20 * PIXEL_SCALE  # ~20 pixels at z=Z0
    planes = gen.generate(
        coefficients={(0, 0): 1 + 0j}, z_list=[Z0], image_shape=IMAGE_SHAPE, center=(offset_m, 0.0)
    )
    flux = planes[0].flux

    peak_idx = np.unravel_index(np.argmax(flux), flux.shape)
    center_row = IMAGE_SHAPE[0] // 2
    center_col = IMAGE_SHAPE[1] // 2
    assert peak_idx[0] == center_row
    assert peak_idx[1] == pytest.approx(center_col + 20, abs=1)


def test_generate_applies_pointing_angles_as_2d_linear_drift():
    gen = make_generator()
    z_list = [Z0, Z0 + 0.2]
    planes = gen.generate(
        coefficients={(0, 0): 1 + 0j},
        z_list=z_list,
        image_shape=IMAGE_SHAPE,
        center=(0.0, 0.0),
        pointing_angle_horizontal_deg=1.0,
        pointing_angle_vertical_deg=0.5,
    )

    peaks = []
    for plane in planes:
        peak_idx = np.unravel_index(np.argmax(plane.flux), plane.flux.shape)
        peaks.append(peak_idx)

    expected_shift_x_m = 0.2 * np.tan(np.deg2rad(1.0))
    expected_shift_y_m = 0.2 * np.tan(np.deg2rad(0.5))
    expected_shift_col_px = expected_shift_x_m / PIXEL_SCALE
    expected_shift_row_px = expected_shift_y_m / PIXEL_SCALE

    actual_shift_col_px = peaks[1][1] - peaks[0][1]
    actual_shift_row_px = peaks[1][0] - peaks[0][0]
    assert actual_shift_col_px == pytest.approx(expected_shift_col_px, abs=1)
    assert actual_shift_row_px == pytest.approx(expected_shift_row_px, abs=1)


def test_generate_noise_is_reproducible_with_seed():
    gen = make_generator()
    planes_a = gen.generate(
        coefficients={(0, 0): 1 + 0j}, z_list=[Z0], image_shape=IMAGE_SHAPE, noise_std=0.01, seed=42
    )
    planes_b = gen.generate(
        coefficients={(0, 0): 1 + 0j}, z_list=[Z0], image_shape=IMAGE_SHAPE, noise_std=0.01, seed=42
    )
    np.testing.assert_array_equal(planes_a[0].flux, planes_b[0].flux)


def test_generate_noise_std_matches_requested_level():
    gen = make_generator()
    noise_std = 0.02
    planes_noisy = gen.generate(
        coefficients={(0, 0): 1 + 0j}, z_list=[Z0], image_shape=IMAGE_SHAPE, noise_std=noise_std, seed=1
    )
    planes_clean = gen.generate(
        coefficients={(0, 0): 1 + 0j}, z_list=[Z0], image_shape=IMAGE_SHAPE, noise_std=0.0
    )

    diff = planes_noisy[0].flux - planes_clean[0].flux
    assert np.std(diff) == pytest.approx(noise_std, rel=0.15)


def test_generate_applies_z_tolerance_to_every_plane():
    gen = make_generator()
    planes = gen.generate(
        coefficients={(0, 0): 1 + 0j},
        z_list=[0.3, 0.4, 0.5],
        image_shape=IMAGE_SHAPE,
        z_tolerance=0.02,
    )
    assert all(p.z_tolerance == 0.02 for p in planes)


def test_generate_applies_nominal_z_offset_independent_of_true_z():
    gen = make_generator()
    true_z_list = [0.3, 0.4, 0.5]
    offsets = {0.3: 0.01, 0.4: -0.005, 0.5: 0.0}
    planes = gen.generate(
        coefficients={(0, 0): 1 + 0j},
        z_list=true_z_list,
        image_shape=IMAGE_SHAPE,
        nominal_z_offsets=offsets,
    )

    nominal_zs = [p.z for p in planes]
    assert nominal_zs == pytest.approx([0.31, 0.395, 0.5])
    # The flux is still rendered at each plane's *true* z (0.3, 0.4, 0.5),
    # not its offset nominal z -- verified indirectly in Task 7's
    # end-to-end tolerance-recovery test.