"""
tiny_engine.py — one request, end to end, in pure NumPy.

A GPT-2-small-shaped decoder (12 layers, d_model=768, 12 heads) with random
weights. The text it produces is noise; the *timing* is real, because the
work per step is exactly the work a trained model of this shape does.

Stages: tokenize -> prefill -> decode (with KV cache) -> stream (detokenize).
"""
import codecs
import time
import numpy as np

# ---------------------------------------------------------------- tokenizer
class ByteTokenizer:
    """Byte-level tokenizer: every UTF-8 byte is one token (vocab = 256)."""
    vocab_size = 256

    def encode(self, text: str) -> list[int]:
        return list(text.encode("utf-8"))

    def stream_decoder(self):
        # Incremental decoder: holds back a partial multi-byte character
        # until the rest of it arrives, exactly like a real detokenizer.
        dec = codecs.getincrementaldecoder("utf-8")(errors="replace")
        return lambda tok: dec.decode(bytes([tok]))


# -------------------------------------------------------------------- model
class TinyDecoder:
    def __init__(self, n_layers=12, d=768, n_heads=12, d_ff=3072,
                 vocab=256, max_seq=4096, seed=0):
        rng = np.random.default_rng(seed)
        w = lambda *s: (rng.standard_normal(s) * 0.02).astype(np.float32)
        self.L, self.d, self.H, self.hd = n_layers, d, n_heads, d // n_heads
        self.tok_emb, self.pos_emb = w(vocab, d), w(max_seq, d)
        self.layers = [dict(qkv=w(d, 3 * d), o=w(d, d),
                            up=w(d, d_ff), down=w(d_ff, d))
                       for _ in range(n_layers)]
        self.lm_head = w(d, vocab)
        # KV cache: one K and one V tensor per layer, preallocated.
        self.k = np.zeros((n_layers, n_heads, max_seq, self.hd), np.float32)
        self.v = np.zeros_like(self.k)
        self.n_cached = 0

    @staticmethod
    def _norm(x):
        return (x - x.mean(-1, keepdims=True)) / np.sqrt(x.var(-1, keepdims=True) + 1e-5)

    def forward(self, ids: np.ndarray) -> np.ndarray:
        """Run len(ids) new tokens through the model, append their K/V to the
        cache, and return logits for the LAST position only.
        Prefill calls this once with the whole prompt; decode calls it with 1 id."""
        T, start = len(ids), self.n_cached
        x = self.tok_emb[ids] + self.pos_emb[start:start + T]          # (T, d)
        for i, p in enumerate(self.layers):
            q, k, v = np.split(self._norm(x) @ p["qkv"], 3, axis=-1)
            q, k, v = (t.reshape(T, self.H, self.hd).transpose(1, 0, 2) for t in (q, k, v))
            self.k[i, :, start:start + T], self.v[i, :, start:start + T] = k, v
            K, V = self.k[i, :, :start + T], self.v[i, :, :start + T]   # all past + new
            att = q @ K.transpose(0, 2, 1) / np.sqrt(self.hd)           # (H, T, start+T)
            mask = np.triu(np.ones((T, start + T), bool), k=start + 1)   # causal
            att = np.where(mask, -1e9, att)
            att = np.exp(att - att.max(-1, keepdims=True))
            att /= att.sum(-1, keepdims=True)
            y = (att @ V).transpose(1, 0, 2).reshape(T, self.d)
            x = x + y @ p["o"]
            h = self._norm(x) @ p["up"]
            x = x + (0.5 * h * (1 + np.tanh(0.79788456 * (h + 0.044715 * h**3)))) @ p["down"]
        self.n_cached += T
        return self._norm(x[-1]) @ self.lm_head                          # (vocab,)


# ------------------------------------------------------------------ engine
def generate(model, tok, prompt: str, max_new_tokens: int):
    """Yield (event, text, timestamp) for one request, end to end."""
    t0 = time.perf_counter()
    model.n_cached = 0
    ids = np.array(tok.encode(prompt))                     # 1. tokenize
    yield "tokenized", "", time.perf_counter()
    logits = model.forward(ids)                            # 2. prefill (whole prompt)
    yield "prefilled", "", time.perf_counter()
    detok = tok.stream_decoder()
    for _ in range(max_new_tokens):                        # 3. decode loop
        nxt = int(np.argmax(logits))                       #    sample (greedy)
        yield "token", detok(nxt), time.perf_counter()     # 4. stream
        logits = model.forward(np.array([nxt]))            #    one token in, one out


def measure(model, tok, prompt, max_new_tokens):
    t_start = time.perf_counter()
    token_times = []
    for event, text, t in generate(model, tok, prompt, max_new_tokens):
        if event == "tokenized": t_tok = t
        elif event == "prefilled": t_pre = t
        else: token_times.append(t)
    itl = np.diff(token_times) * 1e3
    e2e = token_times[-1] - t_start
    return dict(prompt_tokens=len(tok.encode(prompt)),
                tokenize_ms=(t_tok - t_start) * 1e3,
                prefill_ms=(t_pre - t_tok) * 1e3,
                ttft_ms=(token_times[0] - t_start) * 1e3,
                itl_p50_ms=float(np.percentile(itl, 50)),
                itl_p99_ms=float(np.percentile(itl, 99)),
                e2e_s=e2e,
                decode_tok_s=(len(token_times) - 1) / (token_times[-1] - token_times[0]))


if __name__ == "__main__":
    np.random.seed(0)
    tok, model = ByteTokenizer(), TinyDecoder()
    base = "Explain how a large language model answers a question, step by step. "
    measure(model, tok, base, 8)                              # warm-up
    print(f"{'prompt':>7} {'tok ms':>7} {'prefill ms':>10} {'TTFT ms':>8} "
          f"{'ITL p50':>8} {'ITL p99':>8} {'E2E s':>6} {'decode tok/s':>12}")
    for n in (64, 256, 1024):
        prompt = (base * 40)[:n]                              # n bytes = n tokens
        r = measure(model, tok, prompt, 64)
        print(f"{r['prompt_tokens']:>7} {r['tokenize_ms']:>7.2f} {r['prefill_ms']:>10.1f} "
              f"{r['ttft_ms']:>8.1f} {r['itl_p50_ms']:>8.1f} {r['itl_p99_ms']:>8.1f} "
              f"{r['e2e_s']:>6.2f} {r['decode_tok_s']:>12.1f}")