# /// script
# requires-python = ">=3.11"
# dependencies = ["numpy==2.2.6"]
# ///
"""Write the golden of the text Dtype.pp_value prints for chosen floats.

    uv run dev/nx2/test/array/gen/float_text.py [--check]

The text of a float is the fewest significant digits that round to it in
its format, with the exponent unbounded above, and of those the nearest,
ties to even. Two oracles state it. numpy's shortest repr (Dragon4,
unique=True) covers float16, float32 and float64. The other formats, which
numpy lacks, take an exact search: the shortest decimal inside the interval
of values that round to the float, in rationals. The search runs on every
value written and must agree with numpy where numpy has the format.

The golden holds, under a line naming each dtype, one line per float as its
bits in hexadecimal and its text:

- the positive values of each format of 16 bits or less whose decimal
  search meets a tie: the value is halfway between two decimals of as many
  digits as the search tries, which C's printf must round to even;
- every positive power of two of float32 and float64.

These hold the values a C library prints or reads wrong when it rounds a
tie away from zero, rounds twice, or reads a decimal near a power of two
into the wrong float. Without --check the golden is written; with it, the
run fails if the golden would change.
"""

import math
import sys
from fractions import Fraction
from pathlib import Path

import numpy as np

GOLDEN = Path(__file__).resolve().parent.parent / "golden" / "float_text.txt"
HEADER = "# generated by dev/nx2/test/array/gen/float_text.py; do not edit\n"


# A format: [exp] exponent bits, [mant] stored fraction bits, and what its
# all-ones exponent holds: infinities and NaNs ("ieee"), finite values and one
# NaN code ("nan_only"), or finite values ("finite").
class Format:
    def __init__(self, name, exp, mant, top):
        self.name, self.exp, self.mant, self.top = name, exp, mant, top
        self.bias = (1 << (exp - 1)) - 1
        self.width = 1 + exp + mant

    def value(self, code):
        e = (code >> self.mant) & ((1 << self.exp) - 1)
        f = code & ((1 << self.mant) - 1)
        if e == 0:
            return Fraction(f, 1 << self.mant) * Fraction(2) ** (1 - self.bias)
        return Fraction((1 << self.mant) + f, 1 << self.mant) * Fraction(2) ** (
            e - self.bias
        )

    def finite(self, code):
        if code >= 1 << (self.width - 1):
            return False
        all_ones = (1 << (self.exp + self.mant)) - 1
        if self.top == "ieee":
            return code >> self.mant != (1 << self.exp) - 1
        if self.top == "nan_only":
            return code != all_ones
        return True

    def positives(self):
        return [c for c in range(1, 1 << (self.width - 1)) if self.finite(c)]


F16 = Format("float16", 5, 10, "ieee")
BF16 = Format("bfloat16", 8, 7, "ieee")
E4M3 = Format("float8_e4m3fn", 4, 3, "nan_only")
E5M2 = Format("float8_e5m2", 5, 2, "ieee")
E2M1 = Format("float4_e2m1fn", 2, 1, "finite")
F32 = Format("float32", 8, 23, "ieee")
F64 = Format("float64", 11, 52, "ieee")
NUMPY = {"float16": np.float16, "float32": np.float32, "float64": np.float64}


def interval(fmt, code):
    """The values that round to [code], to nearest with ties to even: its
    bounds and whether they belong. Past the largest finite value, the values
    within half a step of it round to it."""
    x = fmt.value(code)
    prev = fmt.value(code - 1) if code > 1 else Fraction(0)
    nxt = fmt.value(code + 1) if fmt.finite(code + 1) else x + (x - prev)
    even = code % 2 == 0
    return (prev + x) / 2, (x + nxt) / 2, even


def shortest(fmt, code):
    """The digits and the decimal exponent of the first digit of the shortest
    decimal that rounds to [code], the nearest of them on a tie in length."""
    x = fmt.value(code)
    lo, hi, closed = interval(fmt, code)
    inside = (lambda v: lo <= v <= hi) if closed else (lambda v: lo < v < hi)
    k0 = math.floor(math.log10(x))
    for p in range(1, 40):
        found = []
        for k in (k0 - 1, k0, k0 + 1):
            unit = Fraction(10) ** (k - p + 1)
            first = max(math.ceil(lo / unit), 10 ** (p - 1))
            last = min(math.floor(hi / unit), 10**p - 1)
            for m in range(first, last + 1):
                if inside(m * unit):
                    found.append((abs(m * unit - x), m, k))
        if found:
            # The nearest, and of two as near the one whose last digit is even.
            _, m, k = min(found, key=lambda c: (c[0], c[1] % 2))
            return str(m).rstrip("0"), k
    raise SystemExit(f"{fmt.name} {code:x}: no decimal found")


def numpy_shortest(fmt, code):
    v = np.frombuffer(code.to_bytes(fmt.width // 8, "little"), dtype=NUMPY[fmt.name])
    s = np.format_float_scientific(v[0], unique=True, trim="-")
    mantissa, exponent = s.split("e")
    return mantissa.replace(".", "").rstrip("0"), int(exponent)


def oracle(fmt, code):
    d = shortest(fmt, code)
    if fmt.name in NUMPY and d != numpy_shortest(fmt, code):
        raise SystemExit(f"{fmt.name} {code:x}: numpy and the search differ")
    return d


def exact_digits(x):
    """The significant digits of the exact decimal expansion of the dyadic
    [x]."""
    n, d = x.numerator, x.denominator
    k = d.bit_length() - 1
    return str(n * 5**k).rstrip("0")


def meets_tie(fmt, code, digits):
    """Whether the search for the text [digits] of [code] tries a number of
    digits at which the value is halfway between two decimals: its exact
    expansion has [q] digits ending in 5, so it is a tie at [q - 1] digits,
    and the search tries every count up to the length of its text."""
    e = exact_digits(fmt.value(code))
    return e.endswith("5") and len(e) - 1 <= len(digits)


def layout(digits, exp):
    """The text printed for the digits [d0 d1 ...] times 10^exp at the first
    digit: in full for exponents from -4 to 15, with an exponent beyond."""
    p = len(digits)
    if exp < -4 or exp >= 16:
        fraction = "." + digits[1:] if p > 1 else ""
        return f"{digits[0]}{fraction}e{'-' if exp < 0 else '+'}{abs(exp):02d}"
    if exp >= p - 1:
        return digits + "0" * (exp - p + 1)
    if exp >= 0:
        return digits[: exp + 1] + "." + digits[exp + 1 :]
    return "0." + "0" * (-exp - 1) + digits


def power_codes(fmt):
    """The codes of the positive powers of two of a format: the subnormals
    with one bit set and the normals with a zero fraction."""
    subnormal = [1 << i for i in range(fmt.mant)]
    normal = [e << fmt.mant for e in range(1, (1 << fmt.exp) - 1)]
    return subnormal + normal


def rows():
    out = [HEADER]
    for fmt in (F16, BF16, E4M3, E5M2, E2M1):
        out.append(fmt.name + "\n")
        for code in fmt.positives():
            d = oracle(fmt, code)
            if meets_tie(fmt, code, d[0]):
                out.append(f"{code:x} {layout(*d)}\n")
    for fmt in (F32, F64):
        out.append(fmt.name + "\n")
        for code in power_codes(fmt):
            out.append(f"{code:x} {layout(*oracle(fmt, code))}\n")
    return "".join(out)


def main():
    text = rows()
    if "--check" in sys.argv[1:]:
        if not GOLDEN.exists() or GOLDEN.read_text() != text:
            raise SystemExit(f"{GOLDEN} differs from the oracles")
        return
    GOLDEN.parent.mkdir(exist_ok=True)
    GOLDEN.write_text(text)


if __name__ == "__main__":
    main()