"""Bit-exact GGML numerical primitives.

The C engine relies on two non-obvious primitives:

* ``gguf_nearest_int`` implements round-to-nearest-even with the classic
  ``float + 1.5*2**23`` mantissa trick rather than ``rint``.  The trick is
  reproduced here so that every quantized code is bit-identical to the
  reference implementation.
* ``gguf_fp32_to_fp16`` / ``gguf_fp16_to_fp32`` are the emitted parameter
  domain.  Every solver in V2 may optimise in float32, but it reports and
  stores only fp16-rounded values.
"""

from __future__ import annotations

import numpy as np

_MAGIC = np.float32(12582912.0)
_F16_MAX_BITS = 0x7BFF


def nearest_int(value) -> np.ndarray:
    """GGML-compatible nearest-integer rounding (ties to even)."""
    arr = np.asarray(value, dtype=np.float32)
    shifted = arr + _MAGIC
    bits = shifted.view(np.int32)
    return ((bits & 0x007FFFFF) - 0x00400000).astype(np.int64)


def to_fp16_bits(value) -> np.ndarray:
    arr = np.asarray(value, dtype=np.float32)
    return arr.astype(np.float16).view(np.uint16).astype(np.uint16)


def from_fp16_bits(bits) -> np.ndarray:
    arr = np.asarray(bits, dtype=np.uint16)
    return arr.view(np.float16).astype(np.float32)


def round_fp16(value) -> np.ndarray:
    return from_fp16_bits(to_fp16_bits(value))


def fp16_ulp_neighbors(bits, radius: int) -> np.ndarray:
    """All non-negative finite fp16 bit patterns within ``radius`` ULPs."""
    if radius < 0:
        raise ValueError("radius must be non-negative")
    b = int(bits)
    lo = max(0, b - radius)
    hi = min(_F16_MAX_BITS, b + radius)
    return np.arange(lo, hi + 1, dtype=np.uint16)


def epsilon_for(scale: np.ndarray, relative: float = 1.0e-7) -> np.ndarray:
    return relative * np.abs(scale) + np.float32(1.0e-30)