import numpy as np
import pytest

from he11lib.data import MeasurementPlane
from he11lib.geometry import GeometryCalibration


def test_pixel_scale_known_reflects_plane():
    plane_known = MeasurementPlane(flux=np.ones((5, 5)), z=0.3, pixel_scale=1e-4)
    plane_unknown = MeasurementPlane(flux=np.ones((5, 5)), z=0.3)

    assert GeometryCalibration(plane_known).pixel_scale_known is True
    assert GeometryCalibration(plane_unknown).pixel_scale_known is False


def test_viewing_angle_known_reflects_plane():
    plane_known = MeasurementPlane(flux=np.ones((5, 5)), z=0.3, viewing_angle_deg=10.0)
    plane_unknown = MeasurementPlane(flux=np.ones((5, 5)), z=0.3)

    assert GeometryCalibration(plane_known).viewing_angle_known is True
    assert GeometryCalibration(plane_unknown).viewing_angle_known is False


def test_physical_coordinates_uses_known_calibration():
    plane = MeasurementPlane(
        flux=np.ones((5, 5)), z=0.3, pixel_scale=2e-4, viewing_angle_deg=0.0
    )
    calib = GeometryCalibration(plane)
    x, y = calib.physical_coordinates()

    row_idx = np.arange(5) - 2
    col_idx = np.arange(5) - 2
    expected_x = col_idx * 2e-4
    expected_y = row_idx * 2e-4
    np.testing.assert_allclose(x[2, :], expected_x)
    np.testing.assert_allclose(y[:, 2], expected_y)


def test_physical_coordinates_compresses_x_for_viewing_angle():
    plane = MeasurementPlane(
        flux=np.ones((5, 5)), z=0.3, pixel_scale=2e-4, viewing_angle_deg=60.0
    )
    calib = GeometryCalibration(plane)
    x, y = calib.physical_coordinates()

    col_idx = np.arange(5) - 2
    expected_x = col_idx * 2e-4 / np.cos(np.deg2rad(60.0))
    np.testing.assert_allclose(x[2, :], expected_x)


def test_physical_coordinates_raises_without_calibration_or_override():
    plane = MeasurementPlane(flux=np.ones((5, 5)), z=0.3)
    calib = GeometryCalibration(plane)

    with pytest.raises(ValueError, match="pixel_scale"):
        calib.physical_coordinates()


def test_physical_coordinates_accepts_override_for_unknown_values():
    plane = MeasurementPlane(flux=np.ones((5, 5)), z=0.3)
    calib = GeometryCalibration(plane)

    x, y = calib.physical_coordinates(pixel_scale=1e-4, viewing_angle_deg=0.0)
    col_idx = np.arange(5) - 2
    np.testing.assert_allclose(x[2, :], col_idx * 1e-4)


def test_known_calibration_takes_precedence_over_override():
    plane = MeasurementPlane(
        flux=np.ones((5, 5)), z=0.3, pixel_scale=2e-4, viewing_angle_deg=0.0
    )
    calib = GeometryCalibration(plane)

    # override should be ignored since plane already specifies calibration
    x, _ = calib.physical_coordinates(pixel_scale=999.0, viewing_angle_deg=45.0)
    col_idx = np.arange(5) - 2
    np.testing.assert_allclose(x[2, :], col_idx * 2e-4)