"""Tests for terrain sheet operations."""

import numpy as np
import pytest
from numpy.testing import assert_allclose

from gri_utils.terrain import (
    get_sheet_resolution,
    interpolate_sheet,
    interpolate_sheet_batch,
    lla_to_xyz_sheet,
    xyz_to_lla_sheet,
)


def test_lla_to_xyz_sheet_flat_terrain() -> None:
    """Test conversion of flat terrain at sea level."""
    lats = np.array([39.0, 39.1, 39.2])
    lons = np.array([-105.0, -104.9])
    alts = np.zeros((3, 2))

    sheet = lla_to_xyz_sheet(lats, lons, alts)

    assert sheet.shape == (3, 2, 3)
    # All points should be roughly at Earth's surface
    radii = np.linalg.norm(sheet, axis=-1)
    assert np.all(radii > 6.3e6)
    assert np.all(radii < 6.4e6)


def test_lla_to_xyz_sheet_with_altitude() -> None:
    """Test conversion with non-zero altitude."""
    lats = np.array([0.0])
    lons = np.array([0.0])
    alts = np.array([[1000.0]])  # 1 km altitude

    sheet = lla_to_xyz_sheet(lats, lons, alts)

    # At equator, lon=0, the point should be at (ER + alt, 0, 0)
    assert sheet.shape == (1, 1, 3)
    assert sheet[0, 0, 0] > 6.378e6  # x > Earth radius
    assert abs(sheet[0, 0, 1]) < 1  # y ~ 0
    assert abs(sheet[0, 0, 2]) < 1  # z ~ 0


def test_xyz_to_lla_sheet_roundtrip() -> None:
    """Test round-trip conversion LLA -> XYZ -> LLA."""
    rng = np.random.default_rng(42)
    lats_orig = np.linspace(30.0, 40.0, 5)
    lons_orig = np.linspace(-110.0, -100.0, 6)
    alts_orig = rng.standard_normal((5, 6)) * 100 + 1000

    sheet = lla_to_xyz_sheet(lats_orig, lons_orig, alts_orig)
    lats_out, lons_out, alts_out = xyz_to_lla_sheet(sheet)

    assert_allclose(lats_out, lats_orig, rtol=1e-10)
    assert_allclose(lons_out, lons_orig, rtol=1e-10)
    assert_allclose(alts_out, alts_orig, rtol=1e-6)


def test_interpolate_sheet_center() -> None:
    """Test interpolation at grid center."""
    lats = np.array([0.0, 1.0])
    lons = np.array([0.0, 1.0])
    alts = np.array([[0.0, 0.0], [0.0, 0.0]])

    sheet = lla_to_xyz_sheet(lats, lons, alts)

    # Interpolate at center (0.5, 0.5)
    point = interpolate_sheet(sheet, 0.5, 0.5)

    assert point.shape == (3,)
    # Center point should be roughly the average of corners
    expected = np.mean(sheet.reshape(-1, 3), axis=0)
    assert_allclose(point, expected, rtol=0.01)


def test_interpolate_sheet_corners() -> None:
    """Test interpolation at exact grid corners."""
    lats = np.array([0.0, 1.0, 2.0])
    lons = np.array([0.0, 1.0])
    alts = np.zeros((3, 2))

    sheet = lla_to_xyz_sheet(lats, lons, alts)

    # Corner (0, 0)
    p00 = interpolate_sheet(sheet, 0.0, 0.0, order=1)
    assert_allclose(p00, sheet[0, 0], rtol=1e-10)

    # Corner (2, 1) -> indices (2, 1) in 3x2 grid
    p21 = interpolate_sheet(sheet, 2.0, 1.0, order=1)
    assert_allclose(p21, sheet[2, 1], rtol=1e-10)


def test_interpolate_sheet_batch() -> None:
    """Test batch interpolation."""
    lats = np.linspace(0.0, 1.0, 10)
    lons = np.linspace(0.0, 1.0, 10)
    alts = np.zeros((10, 10))

    sheet = lla_to_xyz_sheet(lats, lons, alts)

    row_fracs = np.array([0.0, 4.5, 9.0])
    col_fracs = np.array([0.0, 4.5, 9.0])

    # Test with linear interpolation for exact corner matches
    points = interpolate_sheet_batch(sheet, row_fracs, col_fracs, order=1)

    assert points.shape == (3, 3)
    # First point should match corner (use atol for near-zero values)
    assert_allclose(points[0], sheet[0, 0], rtol=1e-10, atol=1e-10)
    # Last point should match corner
    assert_allclose(points[2], sheet[9, 9], rtol=1e-10, atol=1e-10)


def test_get_sheet_resolution() -> None:
    """Test resolution computation."""
    # Create a 1 degree x 1 degree sheet at 0.1 degree spacing
    lats = np.linspace(39.0, 40.0, 11)  # 0.1 degree spacing
    lons = np.linspace(-105.0, -104.0, 11)
    alts = np.zeros((11, 11))

    sheet = lla_to_xyz_sheet(lats, lons, alts)
    row_res, col_res = get_sheet_resolution(sheet)

    # At 39 degrees latitude:
    # - 1 degree lat ~ 111 km
    # - 0.1 degree lat ~ 11.1 km
    # - 1 degree lon ~ 111 * cos(39) ~ 86 km
    # - 0.1 degree lon ~ 8.6 km
    assert 10000 < row_res < 12000  # ~11 km per row
    assert 7000 < col_res < 10000  # ~8.6 km per col


@pytest.mark.parametrize(
    "order",
    [0, 1, 3],
)
def test_interpolate_sheet_orders(order: int) -> None:
    """Test different interpolation orders."""
    lats = np.linspace(0.0, 1.0, 10)
    lons = np.linspace(0.0, 1.0, 10)
    alts = np.zeros((10, 10))

    sheet = lla_to_xyz_sheet(lats, lons, alts)

    # All orders should work
    point = interpolate_sheet(sheet, 5.0, 5.0, order=order)
    assert point.shape == (3,)
    assert np.all(np.isfinite(point))
