"""L2 descriptor normaliser aligning cosine similarity to FAISS inner-product (AZ-283).

Public surface frozen by
``_docs/02_document/contracts/shared_helpers/descriptor_normaliser.md`` v1.0.0.

Used on both the corpus side (C10 index build) and the query side (C2 runtime
lookup). The two sides MUST go through the same helper so the FAISS HNSW
search returns useful neighbours.
"""

from __future__ import annotations

from typing import Final

import numpy as np

__all__ = [
    "ALLOWED_DTYPES",
    "DescriptorNormaliser",
    "DescriptorNormaliserError",
]

# Allowed input dtypes; anything else is rejected to keep the FAISS index and
# query path on the same precision.
ALLOWED_DTYPES: Final[tuple[np.dtype, ...]] = (
    np.dtype(np.float16),
    np.dtype(np.float32),
)

_METRIC_VALUE: Final[str] = "inner_product"


class DescriptorNormaliserError(ValueError):
    """Raised on shape / dtype violations (AZ-283)."""


def _validate_dtype(arr: np.ndarray, label: str) -> None:
    if arr.dtype not in ALLOWED_DTYPES:
        raise DescriptorNormaliserError(
            f"{label}: dtype {arr.dtype} not in allowed set (float16, float32)"
        )


class DescriptorNormaliser:
    """Stateless L2-normalisation helper; dtype-preserving; zero-norm safe."""

    @staticmethod
    def l2_normalise(descriptor: np.ndarray) -> np.ndarray:
        if not isinstance(descriptor, np.ndarray):
            raise DescriptorNormaliserError(
                f"l2_normalise: expected np.ndarray; got {type(descriptor).__name__}"
            )
        if descriptor.ndim != 1:
            raise DescriptorNormaliserError(
                f"l2_normalise: expected 1-D shape (D,); got shape {descriptor.shape}"
            )
        if descriptor.shape[0] < 1:
            raise DescriptorNormaliserError(
                f"l2_normalise: dimension must be >= 1; got shape {descriptor.shape}"
            )
        _validate_dtype(descriptor, "l2_normalise")
        in_dtype = descriptor.dtype
        # Compute norm in float32 to stabilise float16 inputs against overflow /
        # underflow; cast back to the caller dtype so we never silently up-cast.
        as_f32 = descriptor.astype(np.float32, copy=False)
        norm = float(np.linalg.norm(as_f32))
        if norm == 0.0:
            return np.zeros_like(descriptor)
        normalised_f32 = as_f32 / norm
        return normalised_f32.astype(in_dtype, copy=False)

    @staticmethod
    def l2_normalise_batch(descriptors: np.ndarray) -> np.ndarray:
        if not isinstance(descriptors, np.ndarray):
            raise DescriptorNormaliserError(
                f"l2_normalise_batch: expected np.ndarray; got {type(descriptors).__name__}"
            )
        if descriptors.ndim != 2:
            raise DescriptorNormaliserError(
                f"l2_normalise_batch: expected 2-D shape (N, D); got shape {descriptors.shape}"
            )
        if descriptors.shape[0] < 1 or descriptors.shape[1] < 1:
            raise DescriptorNormaliserError(
                f"l2_normalise_batch: N and D must be >= 1; got shape {descriptors.shape}"
            )
        _validate_dtype(descriptors, "l2_normalise_batch")
        in_dtype = descriptors.dtype
        as_f32 = descriptors.astype(np.float32, copy=False)
        norms = np.linalg.norm(as_f32, axis=1, keepdims=True)
        # Avoid division-by-zero: leave zero rows as zero.
        safe = np.where(norms == 0.0, 1.0, norms)
        normalised_f32 = np.where(norms == 0.0, 0.0, as_f32 / safe)
        return normalised_f32.astype(in_dtype, copy=False)

    @staticmethod
    def descriptor_metric() -> str:
        return _METRIC_VALUE
