import numpy as np
import pytest

from m3m_disease.features import FEATURE_NAMES, compute_indices, extract_patch_features, feature_vector


def test_multispectral_indices_match_hand_calculation():
    bands = {
        "green": np.array([[0.2, 0.3]], dtype=np.float32),
        "red": np.array([[0.1, 0.2]], dtype=np.float32),
        "red_edge": np.array([[0.3, 0.4]], dtype=np.float32),
        "nir": np.array([[0.7, 0.8]], dtype=np.float32),
    }
    indices = compute_indices(bands)
    np.testing.assert_allclose(indices["ndvi"], [[0.75, 0.6]], atol=1e-6)
    np.testing.assert_allclose(indices["ndre"], [[0.4, 1 / 3]], atol=1e-6)
    np.testing.assert_allclose(indices["gndvi"], [[5 / 9, 5 / 11]], atol=1e-6)


def test_feature_contract_uses_rgb_and_all_mavic_3m_bands():
    rgb = np.zeros((8, 8, 3), dtype=np.float32)
    rgb[:, :, 0] = 0.2
    rgb[:, :, 1] = 0.5
    rgb[:, :, 2] = 0.1
    rgb[::2, ::2, :] += 0.02
    bands = {
        "green": np.full((8, 8), 0.2, dtype=np.float32),
        "red": np.full((8, 8), 0.1, dtype=np.float32),
        "red_edge": np.full((8, 8), 0.3, dtype=np.float32),
        "nir": np.full((8, 8), 0.7, dtype=np.float32),
    }
    features = extract_patch_features(rgb, bands)
    assert tuple(features) == FEATURE_NAMES
    assert features["rgb_exg_mean"] > 0
    assert features["ndvi_mean"] == pytest.approx(0.75, abs=1e-6)
    assert all(f"ms_{name}_mean" in features for name in ("green", "red", "red_edge", "nir"))


def test_feature_extraction_rejects_unregistered_shapes():
    rgb = np.zeros((8, 8, 3), dtype=np.float32)
    bands = {name: np.zeros((8, 7), dtype=np.float32) for name in ("green", "red", "red_edge", "nir")}
    with pytest.raises(ValueError, match="registered"):
        extract_patch_features(rgb, bands)


def test_feature_vector_rejects_missing_or_extra_features():
    with pytest.raises(ValueError, match="contract"):
        feature_vector({"ndvi_mean": 0.5})
