"""Lab 4: why the router splits its FP32 weights into a BF16 "hi + lo" pair.

Run:   python3 lab4_router.py
Needs: Python 3 + numpy.

The router has FP32 weights [256, 8192] and must match an FP64 reference to 1e-4 relative L2 error
(1% for the other layers). Tensor cores multiply BF16, not FP32. submission.py (lines 713-717) does:
    hi = bf16(r)            the nearest BF16 value
    lo = bf16(r - hi)       the nearest BF16 value to what hi missed
and the kernels compute x @ hi.T + x @ lo.T with FP32 accumulation, in ONE pass over both matrices.

Things to try
  - Add a third term lo2 = bf16(r - hi - lo): how much does it still improve?
  - Make the inputs "outliers" like the judge (1% of the x columns scaled by 30): OUTLIERS = True.
"""
import numpy as np

OUTLIERS = False


def bf16(f):
    """Round FP32 -> nearest BF16 (round-to-nearest-even), returned as FP32 values."""
    u = np.asarray(f, dtype=np.float32).view(np.uint32).astype(np.uint64)
    u = ((u + 0x7FFF + ((u >> 16) & 1)) >> 16) << 16
    return u.astype(np.uint32).view(np.float32)


def rel_l2(a, ref):
    return float(np.linalg.norm(a - ref) / np.linalg.norm(ref))


def main():
    rng = np.random.default_rng(7)
    N, K, M = 256, 8192, 16
    r = (rng.standard_normal((N, K), dtype=np.float32) / np.sqrt(K)).astype(np.float32)   # FP32 weights
    x = bf16(rng.standard_normal((M, K), dtype=np.float32))                               # BF16 activations
    if OUTLIERS:
        cols = rng.permutation(K)[: K // 100]
        x[:, cols] = bf16(x[:, cols] * 30)
    ref = x.astype(np.float64) @ r.astype(np.float64).T          # the judge's FP64 reference

    hi = bf16(r)
    lo = bf16(r - hi)
    lo2 = bf16(r - hi - lo)
    print("how well the split represents the weights (relative error of the weight matrix itself):")
    print(f"  hi only          : {rel_l2(hi, r):.2e}")
    print(f"  hi + lo          : {rel_l2(hi.astype(np.float64) + lo, r):.2e}")
    print(f"  hi + lo + lo2    : {rel_l2(hi.astype(np.float64) + lo + lo2, r):.2e}")

    y_fp32 = (x @ r.T).astype(np.float32)                        # FP32 math (SGLang at <= 8 tokens)
    y_hi = (x @ hi.T).astype(np.float32)                         # BF16 weights only
    y_hilo = (x @ hi.T + x @ lo.T).astype(np.float32)            # what our kernels compute
    print("\noutput error vs the FP64 reference (router tolerance 1e-4):")
    for name, y in (("FP32 weights", y_fp32), ("BF16 hi only", y_hi), ("BF16 hi + lo", y_hilo)):
        e = rel_l2(y, ref)
        print(f"  {name:13s}: {e:.2e}  {'pass' if e <= 1e-4 else 'FAIL'}")

    # A judge-style probe row: x has one nonzero entry, a power of two, so the output is a weight column.
    c = 123
    probe = np.zeros((1, K), dtype=np.float32)
    probe[0, c] = 2.0 ** 5
    exact = probe.astype(np.float64) @ r.astype(np.float64).T
    got = (probe @ hi.T + probe @ lo.T).astype(np.float32)
    print(f"\nprobe row (one nonzero = 32 at column {c}): error {rel_l2(got, exact):.2e} (must be <= 1e-3)")


if __name__ == "__main__":
    main()