#!/usr/bin/env python3
"""Small, offline KV-cache laboratory. This is NOT a pretrained language model.

Run: python kv_lab.py --output results.json
Requires Python >=3.10 and NumPy. No downloads, credentials, or API requests.
The model is a randomly initialized, 3-layer causal decoder with GQA and RoPE.
Its purpose is to test state reuse, not language ability or production speed.
"""
from __future__ import annotations

import argparse
import hashlib
import json
import math
import platform
import sys
from collections import OrderedDict
from dataclasses import dataclass
from pathlib import Path
from typing import Any

import numpy as np

Array = np.ndarray
State = list[tuple[Array, Array]]


def softmax(x: Array) -> Array:
    e = np.exp(x - x.max(axis=-1, keepdims=True))
    return e / e.sum(axis=-1, keepdims=True)


def rmsnorm(x: Array) -> Array:
    return x / np.sqrt(np.mean(x * x, axis=-1, keepdims=True) + 1e-6)


def rotate(x: Array, positions: Array) -> Array:
    """x: [heads, tokens, head_dim]; interleaved even/odd RoPE pairs."""
    d = x.shape[-1]
    freq = 10000.0 ** (-np.arange(0, d, 2, dtype=np.float64) / d)
    angles = positions[:, None] * freq[None, :]
    c, s = np.cos(angles).astype(x.dtype), np.sin(angles).astype(x.dtype)
    y = np.empty_like(x)
    y[..., 0::2] = x[..., 0::2] * c - x[..., 1::2] * s
    y[..., 1::2] = x[..., 0::2] * s + x[..., 1::2] * c
    return y


@dataclass(frozen=True)
class Config:
    vocab: int = 41
    width: int = 32
    layers: int = 3
    q_heads: int = 4
    kv_heads: int = 2
    head_dim: int = 8
    seed: int = 7
    dtype: str = "float64"


class TinyDecoder:
    def __init__(self, cfg: Config = Config()):
        if (cfg.width != cfg.q_heads * cfg.head_dim or
            cfg.q_heads % cfg.kv_heads or cfg.head_dim % 2):
            raise ValueError("Invalid head dimensions or GQA grouping")
        self.cfg = cfg
        dtype = np.dtype(cfg.dtype)
        if dtype not in (np.dtype('float32'), np.dtype('float64')):
            raise ValueError("This lab supports float32 and float64 only")
        rng = np.random.default_rng(cfg.seed)
        self.embed = (rng.normal(size=(cfg.vocab, cfg.width)) * 0.3).astype(dtype)
        shapes = [(cfg.width, cfg.q_heads * cfg.head_dim),
                  (cfg.width, cfg.kv_heads * cfg.head_dim),
                  (cfg.width, cfg.kv_heads * cfg.head_dim),
                  (cfg.width, cfg.width), (cfg.width, cfg.width * 2),
                  (cfg.width * 2, cfg.width)]
        self.weights = [[(rng.normal(size=sh) / math.sqrt(sh[0])).astype(dtype)
                         for sh in shapes] for _ in range(cfg.layers)]
        self.lm_head = (rng.normal(size=(cfg.width, cfg.vocab)) /
                        math.sqrt(cfg.width)).astype(dtype)
        h = hashlib.sha256(json.dumps(cfg.__dict__, sort_keys=True).encode())
        for a in [self.embed, self.lm_head] + [w for row in self.weights for w in row]:
            h.update(a.tobytes())
        # Includes actual weights, not just a model name. RoPE rule is fixed by code.
        self.identity = 'tiny-gqa-rope-v1:' + h.hexdigest()
        self.processed_positions = 0

    def empty_state(self) -> State:
        c = self.cfg
        return [(np.empty((c.kv_heads, 0, c.head_dim), dtype=c.dtype),
                 np.empty((c.kv_heads, 0, c.head_dim), dtype=c.dtype))
                for _ in range(c.layers)]

    def forward(self, tokens: list[int], past: State | None = None) -> tuple[Array, State]:
        """Processes ONLY tokens supplied here; past has already been computed.

        Token indices in the causal mask are ABSOLUTE, including the cached offset.
        This function returns new arrays; it never mutates a supplied cache.
        """
        c = self.cfg
        if not tokens or any(type(t) is not int or not 0 <= t < c.vocab for t in tokens):
            raise ValueError("Provide a nonempty list of integer IDs within the vocabulary")
        past = self.empty_state() if past is None else past
        if len(past) != c.layers:
            raise ValueError("Cache layer count mismatch")
        start = past[0][0].shape[1]
        for k, v in past:
            expected = (c.kv_heads, start, c.head_dim)
            if k.shape != expected or v.shape != expected:
                raise ValueError("Inconsistent cache shapes")
            if k.dtype != np.dtype(c.dtype) or v.dtype != np.dtype(c.dtype):
                raise ValueError("Cache dtype mismatch")
        pos = np.arange(start, start + len(tokens))
        x = self.embed[tokens].copy()
        out_state = []
        self.processed_positions += len(tokens)
        for layer, (wq, wk, wv, wo, w1, w2) in enumerate(self.weights):
            h = rmsnorm(x)
            q = (h @ wq).reshape(-1, c.q_heads, c.head_dim).transpose(1, 0, 2)
            k = (h @ wk).reshape(-1, c.kv_heads, c.head_dim).transpose(1, 0, 2)
            v = (h @ wv).reshape(-1, c.kv_heads, c.head_dim).transpose(1, 0, 2)
            q, k = rotate(q, pos), rotate(k, pos)
            k = np.concatenate([past[layer][0], k], axis=1)
            v = np.concatenate([past[layer][1], v], axis=1)
            out_state.append((k, v))  # Persist KV heads, not repeated Q heads.
            kr = np.repeat(k, c.q_heads // c.kv_heads, axis=0)
            vr = np.repeat(v, c.q_heads // c.kv_heads, axis=0)
            scores = (q @ kr.transpose(0, 2, 1)) / math.sqrt(c.head_dim)
            allowed = np.arange(k.shape[1])[None, :] <= pos[:, None]
            scores = np.where(allowed[None, :, :], scores, -np.inf)
            a = (softmax(scores) @ vr).transpose(1, 0, 2).reshape(-1, c.width)
            x = x + a @ wo
            z = rmsnorm(x) @ w1
            x = x + (z / (1 + np.exp(-z))) @ w2
        return rmsnorm(x) @ self.lm_head, out_state


@dataclass
class Entry:
    tokens: tuple[int, ...]  # Collision guard; each entry stores a full block only.
    state: State            # This block's K/V across all layers, never a whole prefix.


class PrefixCache:
    """Serial educational prefix cache. Hits are copied before model execution.

    Source blocks cannot be evicted midway because this implementation is serial.
    Production async transfers need ownership protection until the copy completes.
    This class stores neither final logits nor hidden states, so it leaves >=1 input.
    """
    def __init__(self, model: TinyDecoder, block_size: int = 16,
                 capacity_blocks: int = 128, namespace: str = "local-demo"):
        if block_size < 1 or capacity_blocks < 1:
            raise ValueError("Block size and capacity must be positive")
        self.model = model
        self.block_size = block_size
        self.capacity = capacity_blocks
        self.base = hashlib.sha256((model.identity + ':' + namespace).encode()).digest()
        self.blocks: OrderedDict[bytes, Entry] = OrderedDict()

    def keys(self, tokens: list[int]):
        parent = self.base
        b = self.block_size
        for end in range(b, len(tokens) + 1, b):
            block = tuple(tokens[end-b:end])
            raw = np.asarray(block, dtype='<i8').tobytes()
            key = hashlib.sha256(parent + raw).digest()
            yield end, key, block
            parent = key

    def lookup(self, tokens: list[int]) -> tuple[int, State | None]:
        rows: list[State] = []
        hit = 0
        for end, key, block in self.keys(tokens):
            if end >= len(tokens):
                break  # Need a final hidden state, not only K/V.
            entry = self.blocks.get(key)
            if entry is None or entry.tokens != block:
                break
            self.blocks.move_to_end(key)
            rows.append(entry.state)
            hit = end
        if not rows:
            return 0, None
        state = [(np.concatenate([r[i][0] for r in rows], axis=1).copy(),
                  np.concatenate([r[i][1] for r in rows], axis=1).copy())
                 for i in range(self.model.cfg.layers)]
        return hit, state

    def store(self, tokens: list[int], state: State) -> None:
        # Reverse insertion makes the beginning newer than the tail in this toy LRU.
        # It preserves a useful prefix when one old block must be evicted.
        for end, key, block in reversed(list(self.keys(tokens))):
            if key not in self.blocks:
                if len(self.blocks) >= self.capacity:
                    self.blocks.popitem(last=False)
                sliced = [(k[:, end-self.block_size:end].copy(),
                           v[:, end-self.block_size:end].copy()) for k, v in state]
                self.blocks[key] = Entry(block, sliced)
            self.blocks.move_to_end(key)

    def request(self, tokens: list[int]) -> tuple[Array, State, dict[str, int]]:
        hit, state = self.lookup(tokens)
        logits, state = self.model.forward(tokens[hit:], state)
        self.store(tokens, state)
        return logits[-1], state, {"input_tokens": len(tokens), "hit_tokens": hit,
                                   "computed_tokens": len(tokens)-hit}


def online_attention(q: Array, k: Array, v: Array, block: int = 7) -> Array:
    """One query, all keys visible. Online global softmax; not a GPU kernel."""
    m = -np.inf
    l = 0.0
    acc = np.zeros(v.shape[-1], dtype=np.float64)
    for i in range(0, len(k), block):
        s = k[i:i+block] @ q / math.sqrt(q.size)
        new_m = max(m, float(s.max()))
        alpha = math.exp(m-new_m)
        p = np.exp(s-new_m)
        acc = alpha * acc + p @ v[i:i+block]
        l = alpha * l + float(p.sum())
        m = new_m
    return acc / l


def pack_int4(x: Array, group: int = 16) -> tuple[Array, Array, tuple[int, ...]]:
    """Uniform signed INT4, NOT NVFP4. One FP32 scale per group, no zero point.
    Codes use [-7,7]; +8 offset enables two codes per uint8. Pad zeros at the tail.
    """
    if group < 2 or group % 2:
        raise ValueError("Group must be positive and even")
    f = np.asarray(x, dtype=np.float32).reshape(-1)
    if not np.isfinite(f).all():
        raise ValueError("Quantization requires finite input")
    n = ((len(f) + group - 1) // group) * group
    a = np.pad(f, (0, n-len(f))).reshape(-1, group)
    scales = np.max(np.abs(a), axis=1).astype(np.float32) / 7
    scales[scales == 0] = 1
    codes = (np.rint(a / scales[:, None]).clip(-7, 7).astype(np.int8) + 8).astype(np.uint8)
    codes = codes.reshape(-1)
    packed = codes[0::2] | (codes[1::2] << 4)
    return packed, scales, x.shape


def unpack_int4(packed: Array, scales: Array, shape: tuple[int, ...], group: int = 16) -> Array:
    codes = np.empty(packed.size * 2, dtype=np.int16)
    codes[0::2] = packed & 15
    codes[1::2] = packed >> 4
    a = ((codes - 8).reshape(-1, group) * scales[:, None]).reshape(-1)
    return a[:math.prod(shape)].reshape(shape)


class PagePool:
    """Small single-threaded page-ownership simulation (not model execution)."""
    def __init__(self, capacity: int = 4):
        self.capacity = capacity
        self.pages: dict[int, Array] = {}
        self.refs: dict[int, int] = {}

    def allocate(self, data: Array) -> int:
        for p in range(self.capacity):
            if self.refs.get(p, 0) == 0:
                self.pages[p] = data.copy()
                self.refs[p] = 1
                return p
        raise MemoryError("All pages are held by readers")

    def retain(self, p: int) -> None:
        if self.refs.get(p, 0) <= 0:
            raise RuntimeError("Cannot retain a free page")
        self.refs[p] += 1

    def release(self, p: int) -> None:
        if self.refs.get(p, 0) <= 0:
            raise RuntimeError("Double release or unknown page")
        self.refs[p] -= 1

    def private_for_write(self, p: int) -> int:
        if self.refs.get(p, 0) <= 0:
            raise RuntimeError("Cannot write a free page")
        if self.refs[p] == 1:
            return p
        # Allocate before dropping the old reference, so failure is safe.
        new_p = self.allocate(self.pages[p])
        self.release(p)
        return new_p


def max_error(a: Array, b: Array) -> float:
    return float(np.max(np.abs(a-b)))


def experiments() -> dict[str, Any]:
    rng = np.random.default_rng(20261009)
    prefix = rng.integers(0, 41, size=42).tolist()
    request_a = prefix + [4, 5, 6]
    request_b = prefix + [7, 8, 9, 10, 11]
    rows = []
    for dtype in ('float64', 'float32'):
        model = TinyDecoder(Config(dtype=dtype))
        cache = PrefixCache(model)
        cache.request(request_a)
        cached, _, usage = cache.request(request_b)
        full, _ = model.forward(request_b)
        _, pre = model.forward(request_b[:usage['hit_tokens']])
        split, _ = model.forward(request_b[usage['hit_tokens']:], pre)
        tolerance = 1e-10 if dtype == 'float64' else 2e-5
        assert max_error(cached, full[-1]) < tolerance
        assert max_error(cached, split[-1]) < tolerance
        rows.append({'dtype': dtype, **usage,
                     'cached_vs_full_max_abs': max_error(cached, full[-1]),
                     'split_vs_full_max_abs': max_error(split[-1], full[-1]),
                     'cached_vs_split_max_abs': max_error(cached, split[-1])})

    model = TinyDecoder()
    full, state = model.forward(request_b)
    _, old = model.forward([1, 2, 3, 4] + prefix)
    _, new = model.forward([5, 6, 7, 8] + prefix)
    first = max_error(old[0][0][:,4:], new[0][0][:,4:])
    higher = max_error(old[1][0][:,4:], new[1][0][:,4:])
    assert first < 1e-10 and higher > 1e-6
    # Same-length prefixes: isolate content dependency without a RoPE-offset change.
    _, lead = model.forward([5, 6, 7, 8])
    wrong = [(np.concatenate([lead[i][0], old[i][0][:,4:]], axis=1),
              np.concatenate([lead[i][1], old[i][1][:,4:]], axis=1))
             for i in range(model.cfg.layers)]
    wrong_logits, _ = model.forward([9], wrong)
    right_logits, _ = model.forward([9], new)
    wrong_error = max_error(wrong_logits, right_logits)
    assert wrong_error > 1e-6
    # Replay document under its new prefix: exact repair baseline, not CacheBlend.
    repaired_logits, repaired = model.forward(prefix, lead)
    repaired_next, _ = model.forward([9], repaired)
    assert max_error(repaired_next, right_logits) < 1e-10

    q = rng.normal(size=16)
    k = rng.normal(size=(37,16))
    v = rng.normal(size=(37,12))
    dense = softmax(k @ q / 4) @ v
    online = online_attention(q,k,v)
    online_error = max_error(dense, online)
    assert online_error < 1e-12
    data = rng.normal(size=(4,32,16)).astype(np.float32)
    packed, scales, shape = pack_int4(data)
    recovered = unpack_int4(packed, scales, shape)
    # Quantized model state errors are measured, never interpreted as task quality.
    quant_state = []
    for kk,vv in state:
        quant_state.append(tuple(unpack_int4(*pack_int4(a)).astype('float64') for a in (kk,vv)))
    ref_next,_ = model.forward([12],state)
    quant_next,_ = model.forward([12],quant_state)
    return {'lab_version':'1.0', 'run_date':'2026-10-09',
            'environment': {'python':platform.python_version(), 'numpy':np.__version__,
                            'platform':platform.platform(), 'device':'CPU'},
            'scope':'Random tiny model; numerical and state-management tests only. No API or pretrained-model benchmark.',
            'config':model.cfg.__dict__, 'model_identity':model.identity,
            'prefix_tests':rows,
            'context_dependency': {'first_layer_document_k_max_abs': first,
                                    'second_layer_document_k_max_abs': higher,
                                    'wrong_splice_next_logits_max_abs': wrong_error,
                                    'full_replay_repair_max_abs': max_error(repaired_next,right_logits)},
            'online_softmax_max_abs': online_error,
            'quantization': {'format':'uniform signed INT4; FP32 scale per 16 values',
                             'original_fp32_bytes':int(data.nbytes),
                             'hypothetical_fp16_bytes':int(data.size * 2),
                             'packed_data_bytes':int(packed.nbytes),
                             'scale_bytes':int(scales.nbytes),
                             'total_payload_bytes':int(packed.nbytes + scales.nbytes),
                             'reconstruction_max_abs':max_error(data,recovered),
                             'next_logits_max_abs':max_error(ref_next,quant_next)},
            'timing_note':'No latency benchmark: toy CPU timing is not evidence of GPU speedup.'}


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--output', type=Path, default=Path('results.json'))
    args = parser.parse_args()
    result = experiments()
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(result, ensure_ascii=False, indent=2)+'\n',encoding='utf-8')
    print(json.dumps(result, ensure_ascii=False, indent=2))


if __name__ == '__main__':
    main()