"""Round weights to fewer bits and measure what that costs in error, size and decoding speed.

numpy only. The method is round-to-nearest with a symmetric (absmax) scale: each group of weights
shares one 16-bit scale, and every weight becomes a small signed whole number times that scale.
The model shape comes from llama31_8b_config.json, a copy of the config.json of Llama 3.1 8B on
Hugging Face (read 2026-10-19). Sizes are in bytes unless a name says GiB; bandwidths are in
bytes per second (GB and TB mean 10^9 and 10^12 bytes).
"""
import json
from pathlib import Path

import numpy as np

HERE = Path(__file__).parent
GIB = 1024**3

BITS = (16, 8, 4, 3)      # bits stored per weight the lesson compares; 16 is the unquantized reference
FULL_BITS = 16            # the token table and output matrix stay at 16 bits in published 8-bit and 4-bit files
SCALE_BITS = 16           # each scale is stored as a 16-bit float
GROUP = 128               # weights that share one scale in the usual 4-bit files

# Graphics-card memory bandwidth, bytes per second. Sources (read 2026-10-19):
#   H100 SXM  3.35 TB/s  NVIDIA H100 page, specification table
#   H200      4.8 TB/s   NVIDIA H200 page
#   A100 80 GB SXM  2,039 GB/s  NVIDIA A100 data sheet (June 2021 revision), same figure as lesson 19
BANDWIDTHS = {"H100 SXM": 3.35e12, "H200": 4.8e12, "A100 80 GB SXM": 2.039e12}

# Sixteen example weights for the widget: seed, size, and where the optional outlier goes.
SAMPLE_SEED = 18
SAMPLE_GROUP = 4          # the picture shows groups of 4; real files use 128
SAMPLE_OUTLIER_SLOT = 5
SAMPLE_OUTLIER = 1.0


def load_config() -> dict:
    return json.loads((HERE / "llama31_8b_config.json").read_text())


# ---------- the quantizer

def qmax(bits: int) -> int:
    """Largest whole number a signed `bits`-bit code uses: 127 at 8 bits, 7 at 4, 3 at 3.
    The most negative code, -2**(bits-1), goes unused so that the range is symmetric around zero."""
    if not 2 <= bits <= 8:
        raise ValueError("bits must be between 2 and 8")
    return 2 ** (bits - 1) - 1


def _blocks(a: np.ndarray, group: int | None) -> np.ndarray:
    """View `a` as (..., groups, weights per group). group=None makes the whole tensor one group."""
    if group is None:
        return a.reshape(1, -1)
    width = a.shape[-1]
    if width % group:
        raise ValueError(f"group {group} does not divide the row length {width}")
    return a.reshape(*a.shape[:-1], width // group, group)


def quantize(w, bits: int, group: int | None = None) -> tuple[np.ndarray, np.ndarray]:
    """Return whole-number codes (same shape as `w`) and one 16-bit scale per group.

    scale = largest |weight| in the group / qmax, stored as float16. code = round(weight / scale), ties to even,
    kept inside [-qmax, qmax] (the 16-bit scale can round down, which would push the largest weight past qmax).
    group=None uses one scale for the whole tensor; otherwise a group is `group` neighbours along the last axis."""
    w = np.asarray(w, dtype=np.float64)
    top = qmax(bits)
    blocks = _blocks(w, group)
    scales = (np.abs(blocks).max(axis=-1, keepdims=True) / top).astype(np.float16)
    step = scales.astype(np.float64)
    step[step == 0] = 1.0   # a group of zeros has codes of 0 whatever the step is
    codes = np.clip(np.rint(blocks / step), -top, top).astype(np.int32)
    return codes.reshape(w.shape), scales


def dequantize(codes: np.ndarray, scales: np.ndarray, group: int | None = None) -> np.ndarray:
    """The weights the model computes with: code times its group's scale."""
    out = _blocks(codes, group) * scales.astype(np.float64)
    return out.reshape(codes.shape)


def roundtrip(w, bits: int, group: int | None = None) -> np.ndarray:
    codes, scales = quantize(w, bits, group)
    return dequantize(codes, scales, group)


def rms(a) -> float:
    return float(np.sqrt(np.mean(np.asarray(a, dtype=np.float64) ** 2)))


def relative_rms_error(w, w_hat, typical: float | None = None) -> float:
    """Root-mean-square gap between the weights and their rounded versions, as a share of a typical weight's size
    (by default the weights' own root-mean-square size)."""
    w = np.asarray(w, dtype=np.float64)
    return rms(w - w_hat) / (rms(w) if typical is None else typical)


def error(w, bits: int, group: int | None = None) -> float:
    return 0.0 if bits == FULL_BITS else relative_rms_error(w, roundtrip(w, bits, group))


def bits_per_weight(bits: int, group: int | None = GROUP) -> float:
    """Stored bits per weight including the scales: 4 bits and a 16-bit scale per 128 weights is 4.125."""
    if bits == FULL_BITS or group is None:
        return float(bits)   # one scale per whole tensor is a few bytes in all
    return bits + SCALE_BITS / group


# ---------- the data

def synthetic_weights(rows: int = 256, cols: int = 1024, seed: int = 0, outlier: float = 0.0) -> np.ndarray:
    """Bell-curve weights with standard deviation 0.02 (the initializer_range of the Llama config; a stand-in,
    not a trained matrix). outlier > 0 sets the top-left weight to that value."""
    w = np.random.default_rng(seed).normal(0.0, 0.02, size=(rows, cols))
    if outlier:
        w[0, 0] = outlier
    return w


def sample_weights() -> list[float]:
    """Sixteen weights of the widget, three decimals, none larger than 0.3 in size."""
    w = np.random.default_rng(SAMPLE_SEED).normal(0.0, 0.1, size=16)
    return [float(x) for x in np.round(np.clip(w, -0.3, 0.3), 3)]


def sample_with_outlier() -> list[float]:
    w = sample_weights()
    w[SAMPLE_OUTLIER_SLOT] = SAMPLE_OUTLIER
    return w


def sample_error(bits: int, group: int | None, outlier: bool) -> float:
    """The widget's error figure. With the outlier in, one weight of 1.0 would dominate the root-mean-square of
    sixteen, so the gap is always compared with the size of the ordinary sample weights."""
    w = sample_with_outlier() if outlier else sample_weights()
    if bits == FULL_BITS:
        return 0.0
    return relative_rms_error(w, roundtrip(w, bits, group), typical=rms(sample_weights()))


# ---------- the model

def head_dim(cfg: dict) -> int:
    return cfg.get("head_dim") or cfg["hidden_size"] // cfg["num_attention_heads"]


def param_count(cfg: dict) -> int:
    """Weights of a Llama-style model from its shape (as in lesson 18): attention, feed-forward and two norms per layer,
    the token table, an output matrix the size of the table unless tied, and the final norm."""
    h, ffn = cfg["hidden_size"], cfg["intermediate_size"]
    kv_width = cfg["num_key_value_heads"] * head_dim(cfg)
    per_layer = 2 * h * h + 2 * h * kv_width + 3 * h * ffn + 2 * h
    return cfg["num_hidden_layers"] * per_layer + embedding_params(cfg) + h


def table_params(cfg: dict) -> int:
    return cfg["vocab_size"] * cfg["hidden_size"]


def embedding_params(cfg: dict) -> int:
    """The token table, plus an output matrix of the same size unless the model reuses the table ("tied")."""
    return table_params(cfg) * (1 if cfg.get("tie_word_embeddings") else 2)


def weight_bytes(cfg: dict, bits: int, group: int | None = GROUP) -> int:
    """Embeddings stay at 16 bits; every other weight uses `bits`, plus one 16-bit scale per `group` of them.
    16 bits means no quantization, so no scales. group=None adds none (a scale per matrix is a few hundred bytes)."""
    if bits not in BITS:
        raise ValueError(f"bits must be one of {BITS}")
    embeddings = embedding_params(cfg)
    others = param_count(cfg) - embeddings
    total = embeddings * FULL_BITS // 8 + others * bits // 8
    if bits != FULL_BITS and group is not None:
        if others % group:
            raise ValueError(f"group {group} does not divide {others}")
        total += others // group * SCALE_BITS // 8
    return total


def bytes_per_token(cfg: dict, bits: int, group: int | None = GROUP) -> int:
    """Bytes read from memory to produce one token for one user.
    Every weight is read once, except that an untied model looks up one row of its token table (hidden x 2 bytes)
    instead of reading the whole table. The attention cache and working memory are not counted."""
    total = weight_bytes(cfg, bits, group)
    if cfg.get("tie_word_embeddings"):
        return total
    return total - table_params(cfg) * FULL_BITS // 8 + cfg["hidden_size"] * FULL_BITS // 8


def tokens_per_second_bound(bytes_read: int, bandwidth: float) -> float:
    """Ceiling for one user: the card cannot read the weights faster than its memory bandwidth."""
    return bandwidth / bytes_read