"""Tests for sheet stitching utilities."""

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

from gri_utils.terrain import lla_to_xyz_sheet
from gri_utils.terrain.source import TerrainBounds
from gri_utils.terrain.stitch import get_stitch_resolution, stitch_sheets


def _make_test_sheet(
    lat_min: float,
    lat_max: float,
    lon_min: float,
    lon_max: float,
    n_points: int = 11,
    altitude: float = 0.0,
) -> tuple[np.ndarray, TerrainBounds]:
    """Create a test sheet with given bounds."""
    lats = np.linspace(lat_min, lat_max, n_points)
    lons = np.linspace(lon_min, lon_max, n_points)
    alts = np.full((n_points, n_points), altitude)
    sheet = lla_to_xyz_sheet(lats, lons, alts)
    bounds = TerrainBounds(lat_min, lat_max, lon_min, lon_max)
    return sheet, bounds


def test_stitch_single_sheet() -> None:
    """Test stitching a single sheet returns it directly."""
    sheet, bounds = _make_test_sheet(38.0, 39.0, -106.0, -105.0)
    sheets = [(sheet, bounds)]

    result_sheet, result_bounds = stitch_sheets(sheets)

    assert result_sheet is sheet  # Should be same object, no copy
    assert result_bounds == bounds


def test_stitch_two_horizontal_tiles() -> None:
    """Test stitching two horizontally adjacent tiles."""
    sheet1, bounds1 = _make_test_sheet(38.0, 39.0, -106.0, -105.0, n_points=11)
    sheet2, bounds2 = _make_test_sheet(38.0, 39.0, -105.0, -104.0, n_points=11)

    result_sheet, result_bounds = stitch_sheets([(sheet1, bounds1), (sheet2, bounds2)])

    # Combined bounds
    assert result_bounds.lat_min == 38.0
    assert result_bounds.lat_max == 39.0
    assert result_bounds.lon_min == -106.0
    assert result_bounds.lon_max == -104.0

    # Combined shape: 11 rows, 21 cols (11 + 11 - 1 overlap)
    assert result_sheet.shape == (11, 21, 3)

    # Check corners match original sheets
    assert_allclose(result_sheet[0, 0], sheet1[0, 0], rtol=1e-10)
    assert_allclose(result_sheet[0, -1], sheet2[0, -1], rtol=1e-10)


def test_stitch_two_vertical_tiles() -> None:
    """Test stitching two vertically adjacent tiles."""
    sheet1, bounds1 = _make_test_sheet(38.0, 39.0, -106.0, -105.0, n_points=11)
    sheet2, bounds2 = _make_test_sheet(39.0, 40.0, -106.0, -105.0, n_points=11)

    result_sheet, result_bounds = stitch_sheets([(sheet1, bounds1), (sheet2, bounds2)])

    # Combined bounds
    assert result_bounds.lat_min == 38.0
    assert result_bounds.lat_max == 40.0
    assert result_bounds.lon_min == -106.0
    assert result_bounds.lon_max == -105.0

    # Combined shape: 21 rows (11 + 11 - 1 overlap), 11 cols
    assert result_sheet.shape == (21, 11, 3)


def test_stitch_2x2_grid() -> None:
    """Test stitching a 2x2 grid of tiles."""
    sheet_sw, bounds_sw = _make_test_sheet(38.0, 39.0, -106.0, -105.0, n_points=11)
    sheet_se, bounds_se = _make_test_sheet(38.0, 39.0, -105.0, -104.0, n_points=11)
    sheet_nw, bounds_nw = _make_test_sheet(39.0, 40.0, -106.0, -105.0, n_points=11)
    sheet_ne, bounds_ne = _make_test_sheet(39.0, 40.0, -105.0, -104.0, n_points=11)

    sheets = [
        (sheet_sw, bounds_sw),
        (sheet_se, bounds_se),
        (sheet_nw, bounds_nw),
        (sheet_ne, bounds_ne),
    ]

    result_sheet, result_bounds = stitch_sheets(sheets)

    # Combined bounds
    assert result_bounds.lat_min == 38.0
    assert result_bounds.lat_max == 40.0
    assert result_bounds.lon_min == -106.0
    assert result_bounds.lon_max == -104.0

    # Combined shape: 21x21
    assert result_sheet.shape == (21, 21, 3)


def test_stitch_with_gap_fills_nan() -> None:
    """Test that gaps between tiles are filled with NaN."""
    # Two tiles with a gap between them
    sheet1, bounds1 = _make_test_sheet(38.0, 39.0, -107.0, -106.0, n_points=11)
    sheet2, bounds2 = _make_test_sheet(38.0, 39.0, -105.0, -104.0, n_points=11)

    result_sheet, result_bounds = stitch_sheets([(sheet1, bounds1), (sheet2, bounds2)])

    # Combined bounds span from -107 to -104
    assert result_bounds.lon_min == -107.0
    assert result_bounds.lon_max == -104.0

    # Middle columns should be NaN
    mid_col = result_sheet.shape[1] // 2
    assert np.all(np.isnan(result_sheet[:, mid_col, :]))


def test_stitch_empty_list_raises() -> None:
    """Test that empty list raises ValueError."""
    with pytest.raises(ValueError, match="cannot be empty"):
        stitch_sheets([])


def test_stitch_incompatible_resolution_raises() -> None:
    """Test that incompatible resolutions raise ValueError."""
    sheet1, bounds1 = _make_test_sheet(38.0, 39.0, -106.0, -105.0, n_points=11)
    sheet2, bounds2 = _make_test_sheet(38.0, 39.0, -105.0, -104.0, n_points=21)

    with pytest.raises(ValueError, match="Incompatible"):
        stitch_sheets([(sheet1, bounds1), (sheet2, bounds2)])


def test_get_stitch_resolution() -> None:
    """Test resolution computation."""
    sheet, bounds = _make_test_sheet(38.0, 39.0, -106.0, -105.0, n_points=11)
    lat_res, lon_res = get_stitch_resolution([(sheet, bounds)])

    assert_allclose(lat_res, 0.1, rtol=1e-10)
    assert_allclose(lon_res, 0.1, rtol=1e-10)


def test_get_stitch_resolution_empty_raises() -> None:
    """Test resolution with empty list raises ValueError."""
    with pytest.raises(ValueError, match="cannot be empty"):
        get_stitch_resolution([])
