#!/usr/bin/env python3
"""float_census.py -- classify every float leaf in a set of saves.

This is the evidence behind the float-parity policy in STATE_CHECKSUM.md: it
says whether the corpus actually contains the values where a `bits` policy and
a `canonical` policy disagree (signed zero, NaN payloads), and whether it
contains the values where an x87 -> SSE port is most likely to drift
(subnormals, values at the edge of float32 range).

    uv run python3 float_census.py SAVE...
"""
from __future__ import annotations

import collections
import math
import os
import struct
import sys

_HERE = os.path.dirname(os.path.abspath(__file__))
if _HERE not in sys.path:
    sys.path.insert(0, _HERE)

import state_checksum as sc  # noqa: E402

FLT_MIN_NORMAL = 1.1754943508222875e-38
FLT_MAX = 3.4028234663852886e38


def classify(raw: bytes) -> str:
    (bits,) = struct.unpack("<I", raw)
    (v,) = struct.unpack("<f", raw)
    if bits == 0x80000000:
        return "negative zero"
    if v == 0.0:
        return "positive zero"
    if math.isnan(v):
        return "NaN"
    if math.isinf(v):
        return "infinity"
    if bits == 0x7F7FFFFF:
        return "FLT_MAX (0x7f7fffff)"
    if bits == 0xFF7FFFFF:
        return "-FLT_MAX"
    if abs(v) < FLT_MIN_NORMAL:
        return "subnormal"
    if v == int(v) and abs(v) < 2 ** 24:
        return "exact integer"
    return "ordinary"


def main(argv=None) -> int:
    argv = list(sys.argv[1:] if argv is None else argv)
    if not argv:
        print("usage: float_census.py SAVE...", file=sys.stderr)
        return 2
    seen: set = set()
    total = collections.Counter()
    print("# float census -- classes that separate the 'bits' and 'canonical'")
    print("# policies are: negative zero, NaN.  Classes an x87 -> SSE port is")
    print("# most likely to move are: subnormal, and anything near FLT_MAX.")
    print()
    for p in argv:
        with open(p, "rb") as f:
            data = f.read()
        key = hash(data)
        if key in seen:
            continue
        seen.add(key)
        ck = sc.checksum_bytes(data, path=p, audit=False)
        c = collections.Counter()
        examples: dict = {}
        for n in ck.root.walk():
            if n.kind != "float":
                continue
            k = classify(n.raw)
            c[k] += 1
            examples.setdefault(k, n.path)
        total += c
        print(f"{os.path.basename(p)}: {sum(c.values())} float leaves")
        for k, v in sorted(c.items(), key=lambda kv: -kv[1]):
            print(f"   {v:>6}  {k:<24} e.g. {examples[k]}")
        print()
    print("all distinct saves combined:")
    for k, v in sorted(total.items(), key=lambda kv: -kv[1]):
        print(f"   {v:>6}  {k}")
    print()
    risky = total["negative zero"] + total["NaN"]
    print(f"'canonical' would change {risky} leaf/leaves across the corpus "
          f"({'a no-op today' if risky == 0 else 'NOT a no-op'}).")
    print(f"subnormals: {total['subnormal']}  (each one is an x87/SSE parity risk)")
    return 0


if __name__ == "__main__":
    sys.exit(main())