import numpy as np
import pytest

from he11lib.data import MeasurementPlane, ReconstructionResult, validate_planes


def test_measurement_plane_stores_fields():
    flux = np.ones((4, 4))
    plane = MeasurementPlane(flux=flux, z=0.3)

    assert plane.z == 0.3
    assert np.array_equal(plane.flux, flux)
    assert plane.z_tolerance == 0.0
    assert plane.label is None


def test_measurement_plane_stores_optional_fields():
    flux = np.ones((4, 4))
    plane = MeasurementPlane(flux=flux, z=0.4, z_tolerance=0.01, label="plane_40cm")

    assert plane.z_tolerance == 0.01
    assert plane.label == "plane_40cm"


def test_measurement_plane_rejects_non_2d_flux():
    with pytest.raises(ValueError, match="2D"):
        MeasurementPlane(flux=np.ones((4, 4, 3)), z=0.3)


def test_measurement_plane_rejects_non_positive_z():
    with pytest.raises(ValueError, match="positive"):
        MeasurementPlane(flux=np.ones((4, 4)), z=0.0)

    with pytest.raises(ValueError, match="positive"):
        MeasurementPlane(flux=np.ones((4, 4)), z=-0.1)


def test_measurement_plane_rejects_negative_z_tolerance():
    with pytest.raises(ValueError, match="z_tolerance"):
        MeasurementPlane(flux=np.ones((4, 4)), z=0.3, z_tolerance=-0.01)


def test_validate_planes_rejects_fewer_than_three():
    planes = [
        MeasurementPlane(flux=np.ones((4, 4)), z=0.3),
        MeasurementPlane(flux=np.ones((4, 4)), z=0.4),
    ]
    with pytest.raises(ValueError, match="[Aa]t least 3"):
        validate_planes(planes)


def test_validate_planes_rejects_mismatched_shapes():
    planes = [
        MeasurementPlane(flux=np.ones((4, 4)), z=0.3),
        MeasurementPlane(flux=np.ones((5, 5)), z=0.4),
        MeasurementPlane(flux=np.ones((4, 4)), z=0.5),
    ]
    with pytest.raises(ValueError, match="same shape"):
        validate_planes(planes)


def test_validate_planes_rejects_duplicate_z():
    planes = [
        MeasurementPlane(flux=np.ones((4, 4)), z=0.3),
        MeasurementPlane(flux=np.ones((4, 4)), z=0.3),
        MeasurementPlane(flux=np.ones((4, 4)), z=0.5),
    ]
    with pytest.raises(ValueError, match="distinct"):
        validate_planes(planes)


def test_validate_planes_accepts_valid_list():
    planes = [
        MeasurementPlane(flux=np.ones((4, 4)), z=0.3),
        MeasurementPlane(flux=np.ones((4, 4)), z=0.4),
        MeasurementPlane(flux=np.ones((4, 4)), z=0.5),
    ]
    validate_planes(planes)  # should not raise


def test_reconstruction_result_stores_fields():
    result = ReconstructionResult(
        purity={(0, 0): (1.0, 0.0)},
        reconstructed_field=np.ones((4, 4), dtype=complex),
        centers=[(0.0, 0.0)],
        pointing_angle_horizontal_deg=0.1,
        pointing_angle_vertical_deg=-0.2,
        geometry={"focal_length_px": 2000.0},
        residuals=[np.zeros((4, 4))],
        coefficient_uncertainty={(0, 0): 0.01},
        used_phase_retrieval=False,
    )

    assert result.purity[(0, 0)] == (1.0, 0.0)
    assert result.pointing_angle_horizontal_deg == 0.1
    assert result.pointing_angle_vertical_deg == -0.2
    assert result.used_phase_retrieval is False