import warnings

import numpy as np
import pytest

from he11lib.data import validate_planes
from he11lib.fitting import ModalFitter, generate_mode_shells
from he11lib.geometry import CameraModel, CameraModelTolerance
from he11lib.modes import LGBasis
from he11lib.synthetic import SyntheticBeamGenerator

W0 = 5e-3
Z0 = 0.5
WAVELENGTH = 1.76e-3
PIXEL_SCALE = 4e-4
CAMERA_DISTANCE = 5.0
IMAGE_SHAPE = (61, 61)
Z_LIST = [0.35, 0.5, 0.65, 0.8]


def make_basis():
    return LGBasis(w0=W0, z0=Z0, wavelength=WAVELENGTH)


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


def zero_tolerance():
    return CameraModelTolerance(
        focal_length_px=0.0, position=(0.0, 0.0, 0.0), orientation_deg=(0.0, 0.0, 0.0)
    )


def make_generator(basis, camera):
    return SyntheticBeamGenerator(basis=basis, camera=camera)


def test_generate_mode_shells_orders_by_2p_plus_abs_l():
    shells = generate_mode_shells(max_order=2)
    assert shells[0] == [(0, 0)]
    assert set(shells[1]) == {(0, 1), (0, -1)}
    assert set(shells[2]) == {(0, 2), (0, -2), (1, 0)}


def test_fit_recovers_pure_fundamental_mode():
    basis = make_basis()
    camera = make_camera()
    gen = make_generator(basis, camera)
    planes = gen.generate(
        coefficients={(0, 0): 1.0 + 0j}, z_list=Z_LIST, image_shape=IMAGE_SHAPE, noise_std=1e-4, seed=0
    )

    fitter = ModalFitter(basis)
    result = fitter.fit(planes, modes=[(0, 0)], camera=camera, camera_tolerance=zero_tolerance())

    power_fraction, _ = result.purity[(0, 0)]
    assert power_fraction == pytest.approx(1.0, abs=1e-6)
    for cx, cy in result.centers:
        assert cx == pytest.approx(0.0, abs=2 * PIXEL_SCALE)
        assert cy == pytest.approx(0.0, abs=2 * PIXEL_SCALE)


def test_fit_recovers_two_mode_purity_ratio():
    basis = make_basis()
    camera = make_camera()
    gen = make_generator(basis, camera)
    true_coeffs = {(0, 0): 0.9 + 0j, (1, 0): 0.3 + 0.1j}
    planes = gen.generate(
        coefficients=true_coeffs, z_list=Z_LIST, image_shape=IMAGE_SHAPE, noise_std=1e-4, seed=1
    )

    fitter = ModalFitter(basis)
    result = fitter.fit(
        planes, modes=list(true_coeffs.keys()), camera=camera, camera_tolerance=zero_tolerance()
    )

    true_total = sum(abs(c) ** 2 for c in true_coeffs.values())
    for mode, c in true_coeffs.items():
        expected_fraction = abs(c) ** 2 / true_total
        recovered_fraction, _ = result.purity[mode]
        assert recovered_fraction == pytest.approx(expected_fraction, abs=0.03)


def test_fit_recovers_center_offset():
    basis = make_basis()
    camera = make_camera()
    gen = make_generator(basis, camera)
    true_center = (10 * PIXEL_SCALE, -5 * PIXEL_SCALE)
    planes = gen.generate(
        coefficients={(0, 0): 1.0 + 0j},
        z_list=Z_LIST,
        image_shape=IMAGE_SHAPE,
        center=true_center,
        noise_std=1e-4,
        seed=2,
    )

    fitter = ModalFitter(basis)
    result = fitter.fit(
        planes,
        modes=[(0, 0)],
        camera=camera,
        camera_tolerance=zero_tolerance(),
        initial_center=true_center,
    )

    for cx, cy in result.centers:
        assert cx == pytest.approx(true_center[0], abs=2 * PIXEL_SCALE)
        assert cy == pytest.approx(true_center[1], abs=2 * PIXEL_SCALE)


def test_fit_recovers_pointing_angles_independently():
    basis = make_basis()
    camera = make_camera()
    gen = make_generator(basis, camera)
    planes = gen.generate(
        coefficients={(0, 0): 1.0 + 0j},
        z_list=Z_LIST,
        image_shape=IMAGE_SHAPE,
        pointing_angle_horizontal_deg=0.3,
        pointing_angle_vertical_deg=-0.15,
        noise_std=1e-4,
        seed=6,
    )

    fitter = ModalFitter(basis)
    result = fitter.fit(planes, modes=[(0, 0)], camera=camera, camera_tolerance=zero_tolerance())

    assert result.pointing_angle_horizontal_deg == pytest.approx(0.3, abs=0.05)
    assert result.pointing_angle_vertical_deg == pytest.approx(-0.15, abs=0.05)


def test_fit_holds_zero_tolerance_camera_field_fixed_at_wrong_nominal():
    # A tolerance=0 field must stay exactly at its (deliberately wrong)
    # nominal value rather than being corrected.
    basis = make_basis()
    true_camera = make_camera()
    gen = make_generator(basis, true_camera)
    planes = gen.generate(
        coefficients={(0, 0): 1.0 + 0j}, z_list=Z_LIST, image_shape=IMAGE_SHAPE, noise_std=1e-4, seed=7
    )

    wrong_focal_length = true_camera.focal_length_px * 1.2
    nominal_camera = CameraModel(
        focal_length_px=wrong_focal_length,
        position=true_camera.position,
        orientation_deg=true_camera.orientation_deg,
    )

    fitter = ModalFitter(basis)
    result = fitter.fit(
        planes, modes=[(0, 0)], camera=nominal_camera, camera_tolerance=zero_tolerance()
    )

    assert result.geometry["focal_length_px"] == wrong_focal_length


def test_fit_recovers_offset_camera_field_within_tolerance():
    # A tolerance>0 field recovers a ground-truth offset from nominal, but
    # within its band.
    basis = make_basis()
    true_camera = make_camera()
    gen = make_generator(basis, true_camera)
    planes = gen.generate(
        coefficients={(0, 0): 1.0 + 0j}, z_list=Z_LIST, image_shape=IMAGE_SHAPE, noise_std=1e-4, seed=8
    )

    offset = true_camera.focal_length_px * 0.02  # 2% off nominal
    nominal_camera = CameraModel(
        focal_length_px=true_camera.focal_length_px + offset,
        position=true_camera.position,
        orientation_deg=true_camera.orientation_deg,
    )
    tolerance = CameraModelTolerance(
        focal_length_px=true_camera.focal_length_px * 0.05,  # +/-5% band
        position=(0.0, 0.0, 0.0),
        orientation_deg=(0.0, 0.0, 0.0),
    )

    fitter = ModalFitter(basis)
    result = fitter.fit(planes, modes=[(0, 0)], camera=nominal_camera, camera_tolerance=tolerance)

    assert result.geometry["focal_length_px"] == pytest.approx(
        true_camera.focal_length_px, rel=0.02
    )


def test_fit_clips_out_of_band_ground_truth_to_bound():
    # A ground truth placed outside a deliberately too-tight band is
    # clipped to the bound rather than escaping it.
    basis = make_basis()
    true_camera = make_camera()
    gen = make_generator(basis, true_camera)
    planes = gen.generate(
        coefficients={(0, 0): 1.0 + 0j}, z_list=Z_LIST, image_shape=IMAGE_SHAPE, noise_std=1e-4, seed=9
    )

    # nominal is 10% off true, but the band only allows +/-1%.
    nominal_focal_length = true_camera.focal_length_px * 1.10
    nominal_camera = CameraModel(
        focal_length_px=nominal_focal_length,
        position=true_camera.position,
        orientation_deg=true_camera.orientation_deg,
    )
    tight_tolerance = CameraModelTolerance(
        focal_length_px=nominal_focal_length * 0.01,
        position=(0.0, 0.0, 0.0),
        orientation_deg=(0.0, 0.0, 0.0),
    )

    fitter = ModalFitter(basis)
    result = fitter.fit(
        planes, modes=[(0, 0)], camera=nominal_camera, camera_tolerance=tight_tolerance
    )

    lower_bound = nominal_focal_length - tight_tolerance.focal_length_px
    assert result.geometry["focal_length_px"] == pytest.approx(lower_bound, rel=1e-3)


def test_fit_recovers_offset_z_within_tolerance():
    basis = make_basis()
    camera = make_camera()
    gen = make_generator(basis, camera)
    true_z_list = [0.35, 0.5, 0.65, 0.8]
    offsets = {z: 0.01 for z in true_z_list}  # nominal is 1 cm off true
    planes = gen.generate(
        coefficients={(0, 0): 1.0 + 0j},
        z_list=true_z_list,
        image_shape=IMAGE_SHAPE,
        nominal_z_offsets=offsets,
        z_tolerance=0.03,
        noise_std=1e-4,
        seed=10,
    )

    fitter = ModalFitter(basis)
    result = fitter.fit(planes, modes=[(0, 0)], camera=camera, camera_tolerance=zero_tolerance())

    for i, true_z in enumerate(true_z_list):
        assert result.geometry[f"z_{i}"] == pytest.approx(true_z, abs=0.005)


def test_fit_auto_does_not_add_modes_for_pure_fundamental():
    basis = make_basis()
    camera = make_camera()
    gen = make_generator(basis, camera)
    planes = gen.generate(
        coefficients={(0, 0): 1.0 + 0j}, z_list=Z_LIST, image_shape=IMAGE_SHAPE, noise_std=1e-4, seed=4
    )

    fitter = ModalFitter(basis)
    result = fitter.fit_auto(planes, camera=camera, camera_tolerance=zero_tolerance(), max_order=2)

    assert set(result.purity.keys()) == {(0, 0)}


def test_fit_auto_grows_to_include_second_mode():
    basis = make_basis()
    camera = make_camera()
    gen = make_generator(basis, camera)
    true_coeffs = {(0, 0): 0.9 + 0j, (0, 1): 0.4 + 0j}
    planes = gen.generate(
        coefficients=true_coeffs, z_list=Z_LIST, image_shape=IMAGE_SHAPE, noise_std=1e-4, seed=5
    )

    fitter = ModalFitter(basis)
    result = fitter.fit_auto(planes, camera=camera, camera_tolerance=zero_tolerance(), max_order=2)

    assert (0, 1) in result.purity or (0, -1) in result.purity


def test_fit_auto_warns_when_free_geometry_params_exceed_plane_count():
    basis = make_basis()
    camera = make_camera()
    gen = make_generator(basis, camera)
    planes = gen.generate(
        coefficients={(0, 0): 1.0 + 0j},
        z_list=Z_LIST,  # 4 planes
        image_shape=IMAGE_SHAPE,
        z_tolerance=0.05,  # +4 free z params
        noise_std=1e-4,
        seed=11,
    )

    # +7 free camera params (all but the 2 principal_point components) +
    # 4 free z params = 11 free geometry params > 4 planes.
    generous_tolerance = CameraModelTolerance(
        focal_length_px=camera.focal_length_px * 0.05,
        position=(0.01, 0.01, 0.01),
        orientation_deg=(2.0, 2.0, 2.0),
        principal_point=(0.0, 0.0),
    )

    fitter = ModalFitter(basis)
    with pytest.warns(UserWarning, match="free camera/z geometry parameters"):
        fitter.fit_auto(planes, camera=camera, camera_tolerance=generous_tolerance, max_order=1)


def test_fit_auto_does_not_warn_when_geometry_fully_fixed():
    basis = make_basis()
    camera = make_camera()
    gen = make_generator(basis, camera)
    planes = gen.generate(
        coefficients={(0, 0): 1.0 + 0j}, z_list=Z_LIST, image_shape=IMAGE_SHAPE, noise_std=1e-4, seed=12
    )

    fitter = ModalFitter(basis)
    with warnings.catch_warnings():
        warnings.simplefilter("error", UserWarning)
        fitter.fit_auto(planes, camera=camera, camera_tolerance=zero_tolerance(), max_order=1)