"""Print every number the lesson quotes. Run: python3 show_numbers.py"""
import numpy as np

import quantize as qz

CFG = qz.load_config()
OTHERS = qz.param_count(CFG) - qz.embedding_params(CFG)

# Sizes of the .safetensors files of published copies of Llama 3.1 8B Instruct (Hugging Face API, blobs=true, read 2026-10-19).
FILES = (
    ("16-bit", "NousResearch/Meta-Llama-3.1-8B", 16_060_556_376),
    ("8-bit, one scale per output row (w8a16)", "RedHatAI/Meta-Llama-3.1-8B-Instruct-quantized.w8a16", 9_084_040_952),
    ("4-bit AWQ, group 128", "hugging-quants/Meta-Llama-3.1-8B-Instruct-AWQ-INT4", 5_727_938_576),
    ("4-bit GPTQ, group 128", "hugging-quants/Meta-Llama-3.1-8B-Instruct-GPTQ-INT4", 5_735_719_888),
    ("4-bit GPTQ, group 128 (w4a16)", "RedHatAI/Meta-Llama-3.1-8B-Instruct-quantized.w4a16", 5_735_720_552),
)

# Perplexity on WikiText-2, lower is better (AWQ paper, arXiv 2306.00978v6, Table 4; the RTN rows are round-to-nearest, group 128).
PERPLEXITY = (
    ("Llama-2 7B", 5.47, 5.73, 6.66),
    ("Llama-2 13B", 4.88, 4.98, 5.52),
    ("Llama-2 70B", 3.32, 3.46, 3.98),
)

# Model-card results, Llama 3.1 8B Instruct, 16-bit baseline against 4-bit GPTQ (w4a16, group 128): each card reports its own baseline run.
CARD_W4 = (("OpenLLM v1 average", 74.3, 73.5), ("MMLU 5-shot", 68.3, 66.9), ("GSM-8K", 82.8, 82.9),
           ("TruthfulQA", 54.5, 52.8), ("MMLU-Pro", 30.8, 28.8), ("IFEval", 77.9, 76.3))
CARD_W8 = (("OpenLLM v1 average", 74.05, 74.12), ("Hellaswag", 80.47, 80.48), ("Winogrande", 78.06, 77.51), ("TruthfulQA", 54.48, 54.41))


def gib(n):
    return f"{n / qz.GIB:.2f}"


def main():
    print("== hand-worked examples")
    for bits, w in ((4, [0.32, -0.10, 0.04, -0.70]), (3, [0.32, -0.10, 0.04, -0.70])):
        codes, scales = qz.quantize(w, bits, group=4)
        back = qz.dequantize(codes, scales, 4)
        print(f"{bits} bits: weights {w} scale {float(scales[0, 0]):.7f} codes {codes.tolist()} back {back.round(3).tolist()} "
              f"gaps {(np.array(w) - back).round(3).tolist()}")

    print("== synthetic matrix 256 x 1024, bell curve with std 0.02, seed 0: relative root-mean-square error (%)")
    clean = qz.synthetic_weights()
    print(f"weights {clean.size:,}; groups of 128 per row: {clean.size // 128:,}; largest |weight| {abs(clean).max():.4f} "
          f"= {abs(clean).max() / 0.02:.1f} std")
    for label, w in (("no outlier", clean), ("one weight of 1.0", qz.synthetic_weights(outlier=1.0))):
        for bits in (8, 4, 3):
            tensor, group = qz.error(w, bits), qz.error(w, bits, qz.GROUP)
            zero = np.mean(qz.roundtrip(w, bits) == 0)
            print(f"{label}, {bits} bits: one scale {tensor:.1%} (weights rounded to zero {zero:.1%}); "
                  f"one per 128 {group:.1%} (zero {np.mean(qz.roundtrip(w, bits, qz.GROUP) == 0):.1%})")

    print("== outlier size sweep, 4 bits (std 0.02)")
    for size in (0.0, 0.1, 0.2, 0.5, 1.0):
        w = qz.synthetic_weights(outlier=size)
        print(f"outlier {size:.1f} ({size / 0.02:.0f} std): one scale {qz.error(w, 4):.1%}, one per 128 {qz.error(w, 4, qz.GROUP):.1%}")

    print("== bits per weight including scales (group 128)")
    print("  ".join(f"{b}-bit {qz.bits_per_weight(b):.3f}" for b in qz.BITS))

    print("== Llama 3.1 8B")
    n, e = qz.param_count(CFG), qz.embedding_params(CFG)
    print(f"weights {n:,}; embeddings (table + output) {e:,} = {e * 2:,} B = {gib(e * 2)} GiB kept at 16 bits; others {OTHERS:,}")
    base = qz.weight_bytes(CFG, 16)
    for b in qz.BITS:
        none, grp = qz.weight_bytes(CFG, b, None), qz.weight_bytes(CFG, b)
        print(f"{b:>2} bits: no scales {none:,} B = {gib(none)} GiB; with group-128 scales {grp:,} B = {gib(grp)} GiB "
              f"({base / grp:.2f}x smaller); scales alone {grp - none:,} B = {(grp - none) / 1e6:.1f} MB = {(grp - none) / grp:.1%} of the file")

    print("== decoding ceiling for one user: tokens per second = bandwidth / bytes read per token")
    print(f"token table {qz.table_params(CFG) * 2:,} B; one row {CFG['hidden_size'] * 2:,} B")
    for b in qz.BITS:
        read = qz.bytes_per_token(CFG, b)
        row = "  ".join(f"{name} {qz.tokens_per_second_bound(read, bw):,.0f}" for name, bw in qz.BANDWIDTHS.items())
        print(f"{b:>2} bits: reads {read:,} B ({read / 1e9:.2f} GB) per token; {row}")
    h100 = qz.BANDWIDTHS["H100 SXM"]
    ref = qz.tokens_per_second_bound(qz.bytes_per_token(CFG, 16), h100)
    print("gain over 16 bits on H100 SXM: " + "  ".join(
        f"{b}-bit {qz.tokens_per_second_bound(qz.bytes_per_token(CFG, b), h100) / ref:.2f}x" for b in qz.BITS))
    whole = [qz.weight_bytes(CFG, b) for b in qz.BITS]
    print("if the whole file were read (the simpler whole-file rule): "
          + "  ".join(f"{b}-bit {h100 / w:,.0f}" for b, w in zip(qz.BITS, whole)))

    print("== size formula against published files (bytes)")
    zero_points = OTHERS // qz.GROUP // 2
    for label, repo, size in FILES:
        print(f"{label}: file {size:,} ({repo})")
    f16, w8, awq, gptq, w4 = (f[2] for f in FILES)
    print(f"16-bit: calculator {qz.weight_bytes(CFG, 16):,}; gap {f16 - qz.weight_bytes(CFG, 16):,}")
    h, kv, ffn = CFG["hidden_size"], CFG["num_key_value_heads"] * qz.head_dim(CFG), CFG["intermediate_size"]
    rows = CFG["num_hidden_layers"] * (2 * h + 2 * kv + 2 * ffn + h)   # output rows of q, o, k, v, gate, up, down
    print(f"w8a16: no scales {qz.weight_bytes(CFG, 8, None):,}; group-128 scales {qz.weight_bytes(CFG, 8):,}; file {w8:,}; "
          f"above no-scales {w8 - qz.weight_bytes(CFG, 8, None):,}; one 16-bit scale per output row: {rows:,} scales = {rows * 2:,} B")
    print(f"AWQ: formula + zero points {qz.weight_bytes(CFG, 4) + zero_points:,}; file {awq:,}; gap {awq - qz.weight_bytes(CFG, 4) - zero_points:,} "
          f"({(awq - qz.weight_bytes(CFG, 4) - zero_points) / awq:.4%})")
    columns = CFG["num_hidden_layers"] * (6 * CFG["hidden_size"] + CFG["intermediate_size"])
    rest = gptq - qz.weight_bytes(CFG, 4) - zero_points - columns * 4
    print(f"GPTQ: file above group-128 formula {gptq - qz.weight_bytes(CFG, 4):,}; zero points {zero_points:,}; "
          f"column index 4 B x {columns:,} = {columns * 4:,}; unitemised remainder {rest:,} ({rest / gptq:.3%} of the file)")
    print(f"Red Hat w4a16 vs hugging-quants GPTQ: {w4 - gptq:+,} B")

    print("== perplexity on WikiText-2 (AWQ paper), 16-bit | 4-bit RTN g128 | 3-bit RTN g128; rise over 16-bit")
    for name, f, i4, i3 in PERPLEXITY:
        print(f"{name}: {f:.2f} | {i4:.2f} (+{i4 - f:.2f}, +{(i4 - f) / f:.1%}) | {i3:.2f} (+{i3 - f:.2f}, +{(i3 - f) / f:.1%})")
    print("OPT-6.7B 3-bit g128 (Table 1): 16-bit 10.86, RTN 23.54, keep 1% of channels in 16 bits chosen by activations 11.39, by weight size 22.37, at random 24.23")
    print("Llama-2 7B GSM8K, 4-bit g128 (Table 8): 16-bit 13.87, RTN 11.07, GPTQ 12.13, AWQ 13.57")

    print("== model cards (Llama 3.1 8B Instruct; each card has its own 16-bit run)")
    for title, rows in (("4-bit GPTQ, group 128", CARD_W4), ("8-bit, one scale per output row", CARD_W8)):
        print(title)
        for name, a, b in rows:
            print(f"  {name}: {a} -> {b} ({b - a:+.2f} points, recovery {b / a:.1%})")

    print("== widget sample: 16 weights, groups of 4; error as a share of a typical weight (%)")
    print("weights", qz.sample_weights(), "typical", f"{qz.rms(qz.sample_weights()):.4f}")
    for outlier in (False, True):
        for group, label in ((None, "one scale"), (qz.SAMPLE_GROUP, "one per 4")):
            print(f"outlier {outlier}, {label}: " + "  ".join(f"{b} bits {qz.sample_error(b, group, outlier):.1%}" for b in (8, 4, 3)))


if __name__ == "__main__":
    main()