gps-denied-onboard/tests/test_metric.py

"""Tests for Metric Refinement (F09)."""

import numpy as np
import pytest

from gps_denied.core.metric import MetricRefinement
from gps_denied.core.models import ModelManager
from gps_denied.schemas import GPSPoint
from gps_denied.schemas.metric import AlignmentResult, ChunkAlignmentResult
from gps_denied.schemas.satellite import TileBounds


@pytest.fixture
def metric():
    manager = ModelManager()
    return MetricRefinement(manager)

@pytest.fixture
def bounds():
    # Covers precisely 1 degree lat and lon around 49, 32
    return TileBounds(
        nw=GPSPoint(lat=50.0, lon=32.0),
        ne=GPSPoint(lat=50.0, lon=33.0),
        sw=GPSPoint(lat=49.0, lon=32.0),
        se=GPSPoint(lat=49.0, lon=33.0),
        center=GPSPoint(lat=49.5, lon=32.5),
        gsd=1.0 # dummy
    )

def test_extract_gps_from_alignment(metric, bounds):
    # Homography is identity -> map center to center
    H = np.eye(3, dtype=np.float64)
    # The image is 256x256 in our mock
    # Center pixel is 128, 128
    gps = metric.extract_gps_from_alignment(H, bounds, (128, 128))
    
    # 128 is middle -> should be EXACTLY at 49.5 lat and 32.5 lon
    assert np.isclose(gps.lat, 49.5)
    assert np.isclose(gps.lon, 32.5)

def test_align_to_satellite(metric, bounds, monkeypatch):
    def mock_infer(*args, **kwargs):
        H = np.eye(3, dtype=np.float64)
        return {"homography": H, "inlier_count": 80, "total_correspondences": 100, "confidence": 0.8, "reprojection_error": 1.0}

    engine = metric.model_manager.get_inference_engine("LiteSAM")
    monkeypatch.setattr(engine, "infer", mock_infer)

    uav = np.zeros((256, 256, 3))
    sat = np.zeros((256, 256, 3))

    res = metric.align_to_satellite(uav, sat, bounds)
    assert res is not None
    assert isinstance(res, AlignmentResult)
    assert res.matched is True
    assert res.inlier_count == 80
    # SAT-04: confidence = inlier_ratio
    assert np.isclose(res.confidence, 80 / 100)


# ---------------------------------------------------------------
# SAT-03: GSD normalization
# ---------------------------------------------------------------

def test_normalize_gsd_downsamples(metric):
    """UAV frame at 0.16 m/px downsampled to satellite 0.6 m/px."""
    uav = np.zeros((480, 640, 3), dtype=np.uint8)
    out = metric.normalize_gsd(uav, uav_gsd_mpp=0.16, sat_gsd_mpp=0.6)
    # Should be roughly 640 * (0.16/0.6) ≈ 170 wide
    assert out.shape[1] < 640
    assert out.shape[0] < 480


def test_normalize_gsd_no_downscale_needed(metric):
    """UAV GSD already coarser than satellite → image unchanged."""
    uav = np.zeros((256, 256, 3), dtype=np.uint8)
    out = metric.normalize_gsd(uav, uav_gsd_mpp=0.8, sat_gsd_mpp=0.6)
    assert out.shape == uav.shape


def test_normalize_gsd_zero_args(metric):
    """Zero GSD args → image returned unchanged (guard against divide-by-zero)."""
    uav = np.zeros((100, 100, 3), dtype=np.uint8)
    out = metric.normalize_gsd(uav, uav_gsd_mpp=0.0, sat_gsd_mpp=0.6)
    assert out.shape == uav.shape


# ---------------------------------------------------------------
# SAT-04: confidence = inlier ratio via align_to_satellite
# ---------------------------------------------------------------

def test_align_confidence_is_inlier_ratio(metric, bounds, monkeypatch):
    """SAT-04: returned confidence must equal inlier_count / total_correspondences."""
    def mock_infer(*args, **kwargs):
        H = np.eye(3, dtype=np.float64)
        return {"homography": H, "inlier_count": 60, "total_correspondences": 150,
                "confidence": 0.4, "reprojection_error": 1.0}

    engine = metric.model_manager.get_inference_engine("LiteSAM")
    monkeypatch.setattr(engine, "infer", mock_infer)

    res = metric.align_to_satellite(np.zeros((256, 256, 3)), np.zeros((256, 256, 3)), bounds)
    if res is not None:
        assert np.isclose(res.confidence, 60 / 150)

def test_align_chunk_to_satellite(metric, bounds, monkeypatch):
    def mock_infer(*args, **kwargs):
        H = np.eye(3, dtype=np.float64)
        return {"homography": H, "inlier_count": 80, "confidence": 0.8}
    
    engine = metric.model_manager.get_inference_engine("LiteSAM")
    monkeypatch.setattr(engine, "infer", mock_infer)
    
    uavs = [np.zeros((256, 256, 3)) for _ in range(5)]
    sat = np.zeros((256, 256, 3))
    
    res = metric.align_chunk_to_satellite(uavs, sat, bounds)
    assert res is not None
    assert isinstance(res, ChunkAlignmentResult)
    assert res.matched is True
    assert res.chunk_id == "chunk1"