"""Proof: the whole-family sheet, derived from each format's bit
counts and checked against torch wherever torch has the type.
For every format: ceiling, smallest positive value, the number of
doublings between them (log2 of their ratio), step at 1.0, and
decimal digits. kind says how the top exponent codes are used:
  ieee  top exponent code reserved for inf and NaN
  fn    no inf; only the top mantissa pattern of the top code is NaN
  all   no inf and no NaN; every code is a number (FP4 E2M1)
"""
import math
import torch

def family(E, M, bias, kind):
    if kind == "ieee":
        emax = 2 ** E - 2 - bias
        ceiling = (2 - 2.0 ** -M) * 2.0 ** emax
    elif kind == "fn":
        emax = 2 ** E - 1 - bias
        ceiling = (2 - 2.0 ** (1 - M)) * 2.0 ** emax
    else:  # all
        emax = 2 ** E - 1 - bias
        ceiling = (2 - 2.0 ** -M) * 2.0 ** emax
    smallest = 2.0 ** (1 - bias - M)          # smallest subnormal
    return {
        "ceiling": ceiling,
        "smallest": smallest,
        "doublings": math.log2(ceiling) - math.log2(smallest),
        "step1": 2.0 ** -M,
        "digits": (M + 1) * math.log10(2),
    }

ROWS = [
    ("float64", 11, 52, 1023, "ieee", torch.float64),
    ("float32", 8, 23, 127, "ieee", torch.float32),
    ("TF32", 8, 10, 127, "ieee", None),
    ("bfloat16", 8, 7, 127, "ieee", torch.bfloat16),
    ("float16", 5, 10, 15, "ieee", torch.float16),
    ("FP8 E5M2", 5, 2, 15, "ieee", torch.float8_e5m2),
    ("FP8 E4M3", 4, 3, 7, "fn", torch.float8_e4m3fn),
    ("FP6 E3M2", 3, 2, 3, "all", None),
    ("FP6 E2M3", 2, 3, 1, "all", None),
    ("FP4 E2M1", 2, 1, 1, "all", None),
]

print(f"{'format':10s} {'ceiling':>10s} {'smallest':>10s} "
      f"{'doublings':>9s} {'step at 1':>10s} {'digits':>6s}")
for name, E, M, bias, kind, dt in ROWS:
    f = family(E, M, bias, kind)
    print(f"{name:10s} {f['ceiling']:>10.4g} {f['smallest']:>10.3g} "
          f"{f['doublings']:>9.1f} {f['step1']:>10.4g} "
          f"{f['digits']:>6.1f}")
    if dt is not None:
        fi = torch.finfo(dt)
        assert f["ceiling"] == fi.max, (name, f["ceiling"], fi.max)
        assert f["step1"] == fi.eps, (name, f["step1"], fi.eps)
print("\ntorch.finfo agrees on ceiling and step for every format "
      "torch carries\n(float64/32/16, bfloat16, both FP8 types).")
print("FP6 sheets from the same formulas; Blackwell tensor cores "
      "run them,\ntorch has no storage type for them yet.")