# examples/pagedattention.py
"""Authored offline CPU KV ownership mechanics for PagedAttention.

Grounding: arxiv.org/html/2309.06180v1, Sections 4.1-4.4 and Figure 6.
This serial, single-head NumPy reference is not a paged GPU kernel, language
model, modern prefix hash cache, scheduler, or serving benchmark. All indexing
uses zero-based token positions. No upstream implementation is ported.
"""

import argparse
import hashlib
import json
from pathlib import Path

import numpy as np


def positive_integer(value, name):
    """Validate dimensions before allocating or changing any state."""
    if type(value) is not int or value <= 0:
        raise ValueError(f"{name} must be a positive integer")
    return value


def identity(value):
    """Require an unambiguous, nonempty sequence name."""
    if not isinstance(value, str) or not value.strip():
        raise ValueError("sequence identity must be a nonempty string")
    return value


def cache_context(value):
    """Validate a declared (model, cache configuration, position context)."""
    if (not isinstance(value, tuple) or len(value) != 3
            or any(not isinstance(x, str) or not x.strip() for x in value)):
        raise ValueError("context must contain three nonempty strings")
    return value


def vector(value, width):
    """Copy one finite, real KV or query vector without casting bad dtypes."""
    array = np.asarray(value)
    if (array.shape != (width,) or array.dtype.kind not in "fiu"
            or not np.isfinite(array).all()):
        raise ValueError(f"expected a finite real vector of shape ({width},)")
    result = array.astype(np.float64, copy=True)
    if not np.isfinite(result).all():
        raise ValueError("vector cannot be represented in float64")
    return result


class KVPool:
    """Own a finite physical pool and append-only logical sequence tables.

    Callers must use the methods for writes; public arrays expose numerical
    fixtures for inspection. Context declarations are an API precondition, not
    an automatic proof that arbitrary text has equal computed KV states.
    """

    # region: physical_pool
    def __init__(self, num_blocks, block_size, key_dim, value_dim,
                 allocation_order=None):
        self.num_blocks = positive_integer(num_blocks, "num_blocks")
        self.block_size = positive_integer(block_size, "block_size")
        self.key_dim = positive_integer(key_dim, "key_dim")
        self.value_dim = positive_integer(value_dim, "value_dim")
        try:
            order = (list(range(num_blocks)) if allocation_order is None
                     else list(allocation_order))
        except TypeError as error:
            raise ValueError(
                "allocation_order must be a permutation") from error
        if (any(type(x) is not int for x in order)
                or sorted(order) != list(range(num_blocks))):
            raise ValueError("allocation_order must be a full permutation")
        self.order = tuple(order)
        self.keys = np.zeros((num_blocks, block_size, key_dim))
        self.values = np.zeros((num_blocks, block_size, value_dim))
        self.refs = np.zeros(num_blocks, dtype=np.int64)
        self.free = set(range(num_blocks))
        self.tables, self.lengths, self.contexts = {}, {}, {}

    def create(self, sequence, context):
        """Create an empty table; do not reserve a per-request block."""
        self._new_identity(sequence)
        context = cache_context(context)
        self.tables[sequence] = []
        self.lengths[sequence] = 0
        self.contexts[sequence] = context
    # endregion: physical_pool

    def _new_identity(self, sequence):
        identity(sequence)
        if sequence in self.tables:
            raise ValueError("duplicate sequence identity")

    def _existing(self, sequence):
        identity(sequence)
        if sequence not in self.tables:
            raise ValueError("unknown or released sequence")

    def _allocate(self):
        """Claim the first free ID in the authored deterministic order."""
        if not self.free:
            raise MemoryError("physical KV pool exhausted")
        physical = next(x for x in self.order if x in self.free)
        self.free.remove(physical)
        self.refs[physical] = 1
        self.keys[physical].fill(0)
        self.values[physical].fill(0)
        return physical

    # region: logical_address
    def address(self, sequence, token):
        """Translate only live, zero-based token indices."""
        self._existing(sequence)
        if (type(token) is not int
                or not 0 <= token < self.lengths[sequence]):
            raise ValueError("token index must address a live token")
        logical, offset = divmod(token, self.block_size)
        return self.tables[sequence][logical], offset

    def gather(self, sequence):
        """Gather live K/V copies in logical order, excluding unused tails."""
        self._existing(sequence)
        addresses = [self.address(sequence, i)
                     for i in range(self.lengths[sequence])]
        keys = np.empty((len(addresses), self.key_dim))
        values = np.empty((len(addresses), self.value_dim))
        for i, (physical, offset) in enumerate(addresses):
            keys[i] = self.keys[physical, offset]
            values[i] = self.values[physical, offset]
        return keys, values
    # endregion: logical_address

    # region: append_or_allocate
    def append(self, sequence, key, value):
        """Validate first; allocate on a boundary or copy a shared partial tail.

        A single-token append needs at most one free block. Exhaustion is
        checked before changing tables, refcounts, data or the free set.
        """
        self._existing(sequence)
        key = vector(key, self.key_dim)
        value = vector(value, self.value_dim)
        length = self.lengths[sequence]
        offset = length % self.block_size
        if offset == 0:
            physical = self._allocate()
            self.tables[sequence].append(physical)
        else:
            physical = self._private_tail(sequence)
        self.keys[physical, offset] = key
        self.values[physical, offset] = value
        self.lengths[sequence] = length + 1
    # endregion: append_or_allocate

    # region: fork_prefix
    def fork(self, parent, child, context=None):
        """Fork actual KV data in the same model/cache/position context."""
        self._existing(parent)
        self._new_identity(child)
        inherited = self.contexts[parent]
        if context is not None and cache_context(context) != inherited:
            raise ValueError("fork requires identical cache context")
        table = self.tables[parent].copy()
        self.tables[child] = table
        self.lengths[child] = self.lengths[parent]
        self.contexts[child] = inherited
        for physical in table:
            self.refs[physical] += 1
    # endregion: fork_prefix

    # region: copy_on_write
    def _private_tail(self, sequence):
        """Copy a shared partial tail; leave completed prefix blocks shared."""
        old = self.tables[sequence][-1]
        if self.refs[old] == 1:
            return old
        # Allocate first: MemoryError leaves every ownership field unchanged.
        new = self._allocate()
        self.keys[new] = self.keys[old]
        self.values[new] = self.values[old]
        self.tables[sequence][-1] = new
        self.refs[old] -= 1
        return new
    # endregion: copy_on_write

    # region: paged_attention
    def attention(self, sequence, query):
        """Gather live rows and compute Eq. 3, not the blockwise CUDA kernel.

        The sole query represents the next decoding position and can attend
        to the entire existing cache. The authored empty-cache result is zero.
        """
        self._existing(sequence)
        query = vector(query, self.key_dim)
        keys, values = self.gather(sequence)
        if not len(keys):
            return np.zeros(self.value_dim)
        with np.errstate(over="ignore", invalid="ignore"):
            scores = keys @ query / np.sqrt(self.key_dim)
        if not np.isfinite(scores).all():
            raise ValueError("attention scores overflow float64")
        # Finite opposite-sign scores may subtract to -inf: exp(-inf)=0.
        with np.errstate(over="ignore"):
            shifted = scores - scores.max()
        weights = np.exp(shifted)
        weights /= weights.sum()
        output = weights @ values
        if not np.isfinite(output).all():
            raise ValueError("attention output overflow float64")
        return output
    # endregion: paged_attention

    # region: release_sequence
    def release(self, sequence):
        """Free only blocks with no remaining owners; reject double release."""
        self._existing(sequence)
        for physical in self.tables[sequence]:
            self.refs[physical] -= 1
            if self.refs[physical] == 0:
                self.free.add(physical)
                self.keys[physical].fill(0)
                self.values[physical].fill(0)
        del self.tables[sequence]
        del self.lengths[sequence]
        del self.contexts[sequence]
    # endregion: release_sequence

    # region: ownership_trace
    def snapshot(self, event):
        """Measure occupied slots once per physical block, even with sharing."""
        occupied = {}
        live = {}
        for sequence, length in self.lengths.items():
            keys, values = self.gather(sequence)
            live[sequence] = {"keys": keys.tolist(), "values": values.tolist()}
            for logical, physical in enumerate(self.tables[sequence]):
                filled = min(self.block_size,
                             length - logical * self.block_size)
                occupied[physical] = max(occupied.get(physical, 0), filled)
        unique = len(occupied)
        live_slots = sum(occupied.values())
        return {"event": event,
                "tables": {s: t.copy() for s, t in self.tables.items()},
                "lengths": self.lengths.copy(), "refs": self.refs.tolist(),
                "free": sorted(self.free), "live_data": live,
                "logical_live_slots": sum(self.lengths.values()),
                "unique_live_slots": live_slots, "unique_blocks": unique,
                "tail_waste_slots": unique * self.block_size - live_slots}


def ownership_trace(seed):
    """Execute a same-context fork, independent appends and both releases."""
    rng = np.random.default_rng(seed)
    keys = rng.normal(size=(10, 3))
    values = rng.normal(size=(10, 2))
    pool = KVPool(8, 4, 3, 2, [5, 2, 7, 0, 3, 6, 1, 4])
    pool.create("A", ("synthetic-model", "single-head-float64", "origin-0"))
    trace = [pool.snapshot("empty")]
    for key, value in zip(keys[:7], values[:7]):
        pool.append("A", key, value)
    trace.append(pool.snapshot("prompt-7"))
    pool.fork("A", "B")
    trace.append(pool.snapshot("fork"))
    pool.append("B", keys[7], values[7])
    trace.append(pool.snapshot("B-copy-on-write-8"))
    pool.append("A", keys[8], values[8])
    trace.append(pool.snapshot("A-private-8"))
    pool.append("B", keys[9], values[9])
    trace.append(pool.snapshot("B-boundary-9"))
    query = rng.normal(size=3)
    outputs = {s: pool.attention(s, query).tolist() for s in ("A", "B")}
    pool.release("A")
    trace.append(pool.snapshot("release-A"))
    pool.release("B")
    trace.append(pool.snapshot("release-B"))
    return trace, outputs
# endregion: ownership_trace


def smoke(seed):
    """Return finite execution evidence, never model or timing claims."""
    trace, outputs = ownership_trace(seed)
    flattened = [x for value in outputs.values() for x in value]
    return {"source_sha256": hashlib.sha256(
                Path(__file__).read_bytes()).hexdigest(),
            "seed": seed,
            "reductions": {"blocks": 8, "tokens_per_block": 4,
                           "heads": 1, "layers": 1, "dtype": "float64"},
            "shapes": {"physical_keys": [8, 4, 3],
                       "physical_values": [8, 4, 2], "query": [3],
                       "A_output": [2], "B_output": [2]},
            "attention_outputs": outputs,
            "finite_outputs": bool(np.isfinite(flattened).all()),
            "ownership_trace": trace,
            "full_pool_recovered": trace[-1]["free"] == list(range(8)),
            "boundary": "offline CPU gather algebra, not a paged GPU kernel"}


def main():
    """Print actual JSON only when the explicitly requested smoke is run."""
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--smoke", action="store_true", required=True)
    parser.add_argument("--seed", type=int, default=7)
    args = parser.parse_args()
    if args.seed < 0:
        parser.error("seed must be nonnegative")
    print(json.dumps(smoke(args.seed), sort_keys=True, allow_nan=False))


if __name__ == "__main__":
    main()