"""AZ-283 — DescriptorNormaliser helper AC tests.

Verifies the contract at ``_docs/02_document/contracts/shared_helpers/descriptor_normaliser.md`` v1.0.0.
"""

from __future__ import annotations

import ast
from pathlib import Path

import numpy as np
import pytest

from gps_denied_onboard.helpers import DescriptorNormaliser, DescriptorNormaliserError


def test_ac1_unit_vector_example() -> None:
    out = DescriptorNormaliser.l2_normalise(np.array([3.0, 4.0], dtype=np.float32))
    np.testing.assert_allclose(out, np.array([0.6, 0.8], dtype=np.float32), atol=1e-6)
    assert float(np.linalg.norm(out)) == pytest.approx(1.0, abs=1e-6)


def test_ac2_batch_normalisation() -> None:
    batch = np.array([[3.0, 4.0], [1.0, 0.0]], dtype=np.float32)
    out = DescriptorNormaliser.l2_normalise_batch(batch)
    np.testing.assert_allclose(out[0], np.array([0.6, 0.8], dtype=np.float32), atol=1e-6)
    np.testing.assert_allclose(out[1], np.array([1.0, 0.0], dtype=np.float32), atol=1e-6)
    for row in out:
        assert float(np.linalg.norm(row)) == pytest.approx(1.0, abs=1e-6)


def test_ac3_fp16_dtype_preservation() -> None:
    rng = np.random.default_rng(2026)
    x = rng.standard_normal(512).astype(np.float16)
    out = DescriptorNormaliser.l2_normalise(x)
    assert out.dtype == np.float16
    assert float(np.linalg.norm(out.astype(np.float32))) == pytest.approx(1.0, abs=1e-3)


def test_ac4_fp32_dtype_preservation() -> None:
    rng = np.random.default_rng(2026)
    x = rng.standard_normal(512).astype(np.float32)
    out = DescriptorNormaliser.l2_normalise(x)
    assert out.dtype == np.float32
    assert float(np.linalg.norm(out)) == pytest.approx(1.0, abs=1e-6)


def test_ac5_zero_vector_handling() -> None:
    zeros = np.zeros(128, dtype=np.float32)
    out = DescriptorNormaliser.l2_normalise(zeros)
    np.testing.assert_array_equal(out, zeros)
    assert not np.any(np.isnan(out))


def test_ac5b_zero_row_in_batch_remains_zero() -> None:
    batch = np.array([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]], dtype=np.float32)
    out = DescriptorNormaliser.l2_normalise_batch(batch)
    np.testing.assert_array_equal(out[0], np.zeros(3, dtype=np.float32))
    np.testing.assert_allclose(out[1], np.array([1.0, 0.0, 0.0], dtype=np.float32))


def test_ac6_idempotence_fp32() -> None:
    rng = np.random.default_rng(2026)
    x = rng.standard_normal(64).astype(np.float32)
    once = DescriptorNormaliser.l2_normalise(x)
    twice = DescriptorNormaliser.l2_normalise(once)
    assert once.tobytes() == twice.tobytes()


def test_ac7_idempotence_fp16_within_half_precision_tol() -> None:
    rng = np.random.default_rng(2026)
    x = rng.standard_normal(64).astype(np.float16)
    once = DescriptorNormaliser.l2_normalise(x)
    twice = DescriptorNormaliser.l2_normalise(once)
    np.testing.assert_allclose(twice.astype(np.float32), once.astype(np.float32), atol=1e-3)


def test_ac8_no_in_place_mutation() -> None:
    x = np.array([3.0, 4.0, 0.0], dtype=np.float32)
    snapshot = x.copy()
    _ = DescriptorNormaliser.l2_normalise(x)
    np.testing.assert_array_equal(x, snapshot)


def test_ac9_metric_is_inner_product_exact_string() -> None:
    assert DescriptorNormaliser.descriptor_metric() == "inner_product"


def test_ac10_float64_dtype_rejected() -> None:
    with pytest.raises(DescriptorNormaliserError, match=r"float16.*float32|float32.*float16"):
        DescriptorNormaliser.l2_normalise(np.array([1.0, 2.0], dtype=np.float64))


def test_ac11_shape_contract_single_rejects_2d() -> None:
    with pytest.raises(DescriptorNormaliserError, match=r"1-D"):
        DescriptorNormaliser.l2_normalise(np.zeros((2, 3), dtype=np.float32))


def test_ac11_shape_contract_batch_rejects_1d() -> None:
    with pytest.raises(DescriptorNormaliserError, match=r"2-D"):
        DescriptorNormaliser.l2_normalise_batch(np.zeros(128, dtype=np.float32))


def test_ac12_no_upward_imports_to_components() -> None:
    module_path = (
        Path(__file__).resolve().parents[2]
        / "src"
        / "gps_denied_onboard"
        / "helpers"
        / "descriptor_normaliser.py"
    )
    tree = ast.parse(module_path.read_text(encoding="utf-8"))
    bad: list[str] = []
    for node in ast.walk(tree):
        if isinstance(node, ast.ImportFrom) and node.module:
            if node.module.startswith("gps_denied_onboard.components"):
                bad.append(node.module)
        elif isinstance(node, ast.Import):
            for alias in node.names:
                if alias.name.startswith("gps_denied_onboard.components"):
                    bad.append(alias.name)
    assert not bad, f"descriptor_normaliser must not import components.*; found: {bad}"


# AZ-338 — DescriptorNormaliser v1.1.0: intra_cluster_normalise


def test_intra_cluster_normalise_per_cluster_unit_norm() -> None:
    # Arrange — 4 clusters of dim 3, residuals not yet normalised
    raw = np.array(
        [3.0, 4.0, 0.0, 1.0, 0.0, 0.0, 0.0, 5.0, 12.0, 2.0, 2.0, 1.0],
        dtype=np.float32,
    )
    # Act
    out = DescriptorNormaliser.intra_cluster_normalise(raw, num_clusters=4)
    # Assert — each cluster sub-vector is unit norm in the intra-cluster sense
    reshaped = out.reshape(4, 3)
    for k in range(4):
        assert float(np.linalg.norm(reshaped[k])) == pytest.approx(1.0, abs=1e-6)


def test_intra_cluster_normalise_dtype_preserved_fp16() -> None:
    rng = np.random.default_rng(2026)
    descriptor = rng.standard_normal(64 * 32).astype(np.float16)
    out = DescriptorNormaliser.intra_cluster_normalise(descriptor, num_clusters=64)
    assert out.dtype == np.float16
    reshaped = out.astype(np.float32).reshape(64, 32)
    for k in range(64):
        assert float(np.linalg.norm(reshaped[k])) == pytest.approx(1.0, abs=1e-3)


def test_intra_cluster_normalise_zero_cluster_returns_zero() -> None:
    # 2 clusters of dim 3; first is all zeros, second is unit-normalisable
    raw = np.array([0.0, 0.0, 0.0, 0.0, 3.0, 4.0], dtype=np.float32)
    out = DescriptorNormaliser.intra_cluster_normalise(raw, num_clusters=2)
    np.testing.assert_array_equal(out[:3], np.zeros(3, dtype=np.float32))
    np.testing.assert_allclose(
        out[3:], np.array([0.0, 0.6, 0.8], dtype=np.float32), atol=1e-6
    )


def test_intra_cluster_normalise_rejects_non_divisible_length() -> None:
    raw = np.zeros(7, dtype=np.float32)
    with pytest.raises(DescriptorNormaliserError, match=r"not divisible"):
        DescriptorNormaliser.intra_cluster_normalise(raw, num_clusters=3)


def test_intra_cluster_normalise_rejects_2d_input() -> None:
    raw = np.zeros((4, 3), dtype=np.float32)
    with pytest.raises(DescriptorNormaliserError, match=r"1-D"):
        DescriptorNormaliser.intra_cluster_normalise(raw, num_clusters=4)


def test_intra_cluster_normalise_rejects_zero_num_clusters() -> None:
    raw = np.zeros(12, dtype=np.float32)
    with pytest.raises(DescriptorNormaliserError, match=r">= 1"):
        DescriptorNormaliser.intra_cluster_normalise(raw, num_clusters=0)


def test_intra_cluster_normalise_rejects_bool_num_clusters() -> None:
    raw = np.zeros(12, dtype=np.float32)
    with pytest.raises(DescriptorNormaliserError, match=r"non-bool"):
        DescriptorNormaliser.intra_cluster_normalise(raw, num_clusters=True)  # type: ignore[arg-type]


def test_intra_cluster_normalise_rejects_float64() -> None:
    raw = np.zeros(12, dtype=np.float64)
    with pytest.raises(DescriptorNormaliserError, match=r"float16.*float32|float32.*float16"):
        DescriptorNormaliser.intra_cluster_normalise(raw, num_clusters=4)


def test_intra_cluster_normalise_no_in_place_mutation() -> None:
    raw = np.array(
        [3.0, 4.0, 0.0, 1.0, 0.0, 0.0, 0.0, 5.0, 12.0, 2.0, 2.0, 1.0],
        dtype=np.float32,
    )
    snapshot = raw.copy()
    _ = DescriptorNormaliser.intra_cluster_normalise(raw, num_clusters=4)
    np.testing.assert_array_equal(raw, snapshot)