"""Proof: the field guide's numbers, derived by hand from each
format's two counts (E exponent bits, M mantissa bits), then checked
against torch. Nothing on the cards is typed from memory:
  bias        = 2^(E-1) - 1
  ceiling     = (2 - 2^-M) * 2^bias
  floor       = 2^(1 - bias)          (smallest normal)
  ramp floor  = 2^(1 - bias - M)      (smallest subnormal)
  step at 1   = 2^-M                  (the ulp of 1.0; finfo calls it eps)
  cells       = 2^E - 2               (power-of-two intervals)
  points/cell = 2^M
Also prints the torch defaults the float32 section quotes."""
import math
import torch

FORMATS = [
    ("float64", 11, 52, torch.float64, torch.int64),
    ("float32", 8, 23, torch.float32, torch.int32),
    ("float16", 5, 10, torch.float16, torch.int16),
    ("bfloat16", 8, 7, torch.bfloat16, torch.int16),
]

for name, E, M, dt, it in FORMATS:
    bias = 2 ** (E - 1) - 1
    ceiling = (2 - 2.0 ** -M) * 2.0 ** bias
    floor = 2.0 ** (1 - bias)
    ramp_floor = 2.0 ** (1 - bias - M)
    step_at_1 = 2.0 ** -M
    cells = 2 ** E - 2
    digits = (M + 1) * math.log10(2)

    fi = torch.finfo(dt)
    # smallest subnormal: bit pattern 0...01 viewed as this dtype
    sub = torch.tensor([1], dtype=it).view(dt).double().item()

    assert ceiling == fi.max, (name, ceiling, fi.max)
    assert floor == fi.tiny, (name, floor, fi.tiny)
    assert step_at_1 == fi.eps, (name, step_at_1, fi.eps)
    assert ramp_floor == sub, (name, ramp_floor, sub)

    print(f"{name:9s} E={E:2d} M={M:2d}  bias {bias:4d}   "
          f"ceiling {ceiling:.4g}")
    print(f"          floor {floor:.4g}   ramp floor {ramp_floor:.4g}   "
          f"step at 1 = {step_at_1:.4g}")
    print(f"          {cells} cells of {2**M:,} points; "
          f"about {digits:.1f} decimal digits   "
          f"(torch.finfo agrees on all four numbers)")

# absorption never leaves; it only moves. float64's step at 1e16 is 2.
x = torch.tensor(1e16, dtype=torch.float64)
print(f"\nfloat64: (1e16 + 1) - 1e16 = {((x + 1) - x).item():g}"
      f"   (the step at 1e16 is 2; the 1 is below half of it)")

# the defaults the float32 section quotes, read from this install
print(f"\ntorch {torch.__version__} defaults:")
print(f"  default dtype                  {torch.get_default_dtype()}")
print(f"  float32_matmul_precision       "
      f"{torch.get_float32_matmul_precision()!r}")
print(f"  backends.cuda.matmul.allow_tf32  "
      f"{torch.backends.cuda.matmul.allow_tf32}")
print(f"  backends.cudnn.allow_tf32        "
      f"{torch.backends.cudnn.allow_tf32}")
print(f"  autocast dtype on cpu          "
      f"{torch.get_autocast_dtype('cpu')}")