"""Практикум: повторные проекции K/V и KV-кэш в учебном самовнимании на CPU.

Запуск после скачивания файла: python -B -X utf8 kv-cache-cpu.py
Зависимостей кроме стандартной библиотеки Python нет. Сравнивается один слой
и одна голова на фиксированном потоке входов. Счётчик показывает реально
созданные векторы K и V; это не число операций (FLOP), объём памяти, время
или оценка скорости LLM либо GPU. Ограничения: нет пакетной обработки,
нескольких голов и слоёв, позиционных кодировок, выходной проекции и реальных
вычислительных ядер или памяти GPU; токены не генерируются.
"""

from __future__ import annotations

import math
from dataclasses import dataclass
from typing import Sequence

Vector = tuple[float, ...]
Matrix = tuple[Vector, ...]

W_Q: Matrix = ((0.7, -0.2, 0.4), (0.1, 0.8, -0.3), (-0.5, 0.2, 0.9))
W_K: Matrix = ((0.3, 0.9, -0.4), (-0.8, 0.2, 0.5), (0.6, -0.1, 0.7))
W_V: Matrix = ((0.4, -0.7, 0.2), (0.9, 0.3, -0.5), (-0.2, 0.8, 0.6))
INPUTS: tuple[Vector, ...] = (
    (0.8, -0.3, 0.5),
    (-0.4, 0.9, 0.2),
    (0.1, -0.6, 1.1),
    (0.7, 0.4, -0.8),
    (-0.9, 0.2, 0.6),
    (0.3, 1.0, -0.5),
    (1.2, -0.7, 0.1),
)
PROMPT_LENGTH = 3
DECODE_STEPS = 4


@dataclass
class ProjectionCounts:
    """Число фактически созданных K и V векторов, отдельно друг от друга."""

    keys: int = 0
    values: int = 0


def project(matrix: Matrix, vector: Vector) -> Vector:
    return tuple(sum(weight * value for weight, value in zip(row, vector)) for row in matrix)


def dot(left: Vector, right: Vector) -> float:
    return sum(a * b for a, b in zip(left, right))


def softmax_stable(scores: Sequence[float]) -> tuple[float, ...]:
    peak = max(scores)
    exps = tuple(math.exp(score - peak) for score in scores)
    total = sum(exps)
    return tuple(value / total for value in exps)


def attend(query: Vector, keys: Sequence[Vector], values: Sequence[Vector]) -> Vector:
    if not keys or len(keys) != len(values):
        raise ValueError("Для внимания нужны одинаковые непустые списки K и V")
    weights = softmax_stable(tuple(dot(query, key) / math.sqrt(len(query)) for key in keys))
    return tuple(sum(weight * value[axis] for weight, value in zip(weights, values))
                 for axis in range(len(values[0])))


def kv_for(vector: Vector, counts: ProjectionCounts | None = None) -> tuple[Vector, Vector]:
    key = project(W_K, vector)
    value = project(W_V, vector)
    if counts is not None:
        counts.keys += 1
        counts.values += 1
    return key, value


def full_forward(
    inputs: Sequence[Vector], *, causal: bool = True,
    counts: ProjectionCounts | None = None,
) -> list[Vector]:
    """Проецирует весь префикс; causal=False нужен только ошибочному контролю."""
    queries, keys, values = project_qkv(inputs, counts if counts is not None else ProjectionCounts())
    return causal_outputs(queries, keys, values, causal=causal)


def allclose(left: Sequence[float], right: Sequence[float], *, atol: float = 1e-12) -> bool:
    return len(left) == len(right) and all(abs(a - b) <= atol for a, b in zip(left, right))


def run_negative_controls(inputs: Sequence[Vector] = INPUTS) -> tuple[float, float]:
    """Возвращает изменения выхода в двух конкретных ошибочных вариантах."""
    if len(inputs) < 4:
        raise ValueError("Для контролей нужны три позиции входа и одна новая позиция")
    correct = full_forward(inputs[:3], causal=True)
    leaked = full_forward(inputs[:3], causal=False)
    # Query 1 — не последняя позиция prefill длины 3.
    future_delta = max(abs(a - b) for a, b in zip(correct[1], leaked[1]))

    # Ошибочный вариант кэша, в котором оставлена только новая пара K/V.
    keys, values = zip(*(kv_for(item) for item in inputs[:4]))
    query = project(W_Q, inputs[3])
    correct_decode = attend(query, keys, values)
    only_new_decode = attend(query, keys[-1:], values[-1:])
    lost_old_delta = max(abs(a - b) for a, b in zip(correct_decode, only_new_decode))
    return lost_old_delta, future_delta


def run_lab(
    inputs: Sequence[Vector] = INPUTS,
    *,
    prompt_length: int = PROMPT_LENGTH,
    decode_steps: int = DECODE_STEPS,
) -> dict[str, object]:
    if prompt_length < 1 or decode_steps < 1 or prompt_length + decode_steps != len(inputs):
        raise ValueError("Нужны положительные длины prefill и decode, сумма которых равна числу входных векторов")

    base_counts = ProjectionCounts()
    cache_counts = ProjectionCounts()

    # Оба варианта ровно один раз выполняют полный prefill: все prompt Q/K/V,
    # causal outputs для каждой prompt позиции. Считаем реально созданные K/V.
    base_q, base_k, base_v = project_qkv(inputs[:prompt_length], base_counts)
    base_prompt_outputs = causal_outputs(base_q, base_k, base_v)
    cache_q, cached_keys, cached_values = project_qkv(inputs[:prompt_length], cache_counts)
    cache_prompt_outputs = causal_outputs(cache_q, cached_keys, cached_values)
    if any(not allclose(a, b) for a, b in zip(base_prompt_outputs, cache_prompt_outputs)):
        raise AssertionError("Выходы prefill при полном пересчёте и с KV-кэшем не совпали")

    prefill = {
        "length": prompt_length,
        "base_k": base_counts.keys,
        "base_v": base_counts.values,
        "cache_k": cache_counts.keys,
        "cache_v": cache_counts.values,
        "query_key_pairs": prompt_length * (prompt_length + 1) // 2,
    }
    rows = []

    for position in range(prompt_length, prompt_length + decode_steps):
        prefix = inputs[: position + 1]
        before_base = ProjectionCounts(base_counts.keys, base_counts.values)
        base_outputs = full_forward(prefix, counts=base_counts)
        base_new_k = base_counts.keys - before_base.keys
        base_new_v = base_counts.values - before_base.values

        before_cache = ProjectionCounts(cache_counts.keys, cache_counts.values)
        cached_output = cache_decode_step(inputs[position], cached_keys, cached_values, cache_counts)
        cache_new_k = cache_counts.keys - before_cache.keys
        cache_new_v = cache_counts.values - before_cache.values
        if not allclose(base_outputs[-1], cached_output):
            raise AssertionError(f"Выходы полного пересчёта и KV-кэша не совпали для позиции {position}")
        rows.append({
            "position": position,
            "prefix_length": position + 1,
            "base_new_k": base_new_k,
            "base_new_v": base_new_v,
            "cache_new_k": cache_new_k,
            "cache_new_v": cache_new_v,
            "readable_keys": len(cached_keys),
        })

    return {"prefill": prefill, "decode": rows, "base_total": base_counts, "cache_total": cache_counts}


def project_qkv(
    inputs: Sequence[Vector], counts: ProjectionCounts,
) -> tuple[list[Vector], list[Vector], list[Vector]]:
    queries = [project(W_Q, item) for item in inputs]
    pairs = [kv_for(item, counts) for item in inputs]
    return queries, [pair[0] for pair in pairs], [pair[1] for pair in pairs]


def causal_outputs(
    queries: Sequence[Vector], keys: Sequence[Vector], values: Sequence[Vector],
    *, causal: bool = True,
) -> list[Vector]:
    outputs = []
    for position, query in enumerate(queries):
        end = position + 1 if causal else len(keys)
        outputs.append(attend(query, keys[:end], values[:end]))
    return outputs


def cache_decode_step(
    item: Vector,
    keys: list[Vector],
    values: list[Vector],
    counts: ProjectionCounts,
) -> Vector:
    """Создаёт одну новую пару K/V и читает весь сохранённый causal cache."""
    key, value = kv_for(item, counts)
    keys.append(key)
    values.append(value)
    return attend(project(W_Q, item), keys, values)


def main() -> None:
    report = run_lab()
    prefill = report["prefill"]
    rows = report["decode"]
    old_delta, future_delta = run_negative_controls()
    if old_delta <= 1e-9 or future_delta <= 1e-9:
        raise AssertionError("Контроль ошибок не обнаружил потерю старых K/V или доступ к будущему")
    print("Причинное самовнимание: один слой, одна голова, фиксированные входные векторы.")
    print("Счётчики новых K и V показывают число созданных векторов каждого вида.")
    print("Позиции в таблице нумеруются с нуля.")
    print(f"Prefill: позиций — {prefill['length']}; пересчёт K={prefill['base_k']}, V={prefill['base_v']}; кэш K={prefill['cache_k']}, V={prefill['cache_v']}; разрешённых пар Q/K в каждом варианте — {prefill['query_key_pairs']}.")
    print("Decode: позиция | длина префикса | полный пересчёт: новые K/V | кэш: новые K/V | ключей в кэше для нового запроса")
    for row in rows:
        print(f"{row['position']:>4} | {row['prefix_length']:>14} | K={row['base_new_k']}, V={row['base_new_v']} | K={row['cache_new_k']}, V={row['cache_new_v']} | {row['readable_keys']}")
    base = report["base_total"]
    cache = report["cache_total"]
    print(f"Итог с учётом prefill: пересчёт K={base.keys}, V={base.values}; KV-кэш K={cache.keys}, V={cache.values}.")
    print("Выходы внимания совпали при prefill и на всех шагах decode: ПРОЙДЕНО (допуск 1e-12).")
    print(f"Отрицательные проверки: кэш только с новой парой K/V меняет выход на {old_delta:.6g}; доступ к будущей позиции — на {future_delta:.6g}.")
    print("Это учебный подсчёт векторов, а не оценка числа операций (FLOP), памяти, времени или скорости LLM.")


if __name__ == "__main__":
    main()