app/services/attention_compiler.py

from __future__ import annotations

import math
import time
from typing import Iterable

from app.schemas.visual_lesson import (
    AttentionLessonDraft,
    CompiledBranch,
    CompiledLesson,
    EvidenceClaim,
    EvidenceSource,
    NumericAssertion,
    PanelSpec,
    TensorInput,
    TensorOperation,
    TimelineStep,
    VisualLessonPayload,
    VisualLessonSpec,
)


DEFAULT_TOKENS = ["Research", "Mate", "reads", "papers"]
EMBEDDINGS = [[1.0, 0.0, 1.0], [0.0, 1.0, 1.0], [1.0, 1.0, 0.0], [0.0, 0.0, 1.0]]
WQ = [[1.0, 0.0], [0.0, 1.0], [1.0, 1.0]]
WK = [[1.0, 0.0], [0.0, 1.0], [0.5, 0.5]]
WV = [[1.0, 0.0], [0.0, 1.0], [1.0, -1.0]]
DK = 2
TOLERANCE = 1e-6


class AttentionCompilerError(ValueError):
    pass


def _shape(matrix: list[list[float]]) -> tuple[int, int]:
    if not matrix or not matrix[0]:
        raise AttentionCompilerError("Tensor matrices must not be empty")
    width = len(matrix[0])
    if any(len(row) != width for row in matrix):
        raise AttentionCompilerError("Tensor matrices must be rectangular")
    return len(matrix), width


def _matmul(left: list[list[float]], right: list[list[float]]) -> list[list[float]]:
    left_rows, left_cols = _shape(left)
    right_rows, right_cols = _shape(right)
    if left_cols != right_rows:
        raise AttentionCompilerError(f"Cannot multiply shapes {(left_rows, left_cols)} and {(right_rows, right_cols)}")
    return [
        [sum(left[i][k] * right[k][j] for k in range(left_cols)) for j in range(right_cols)]
        for i in range(left_rows)
    ]


def _matmul_transpose_right(left: list[list[float]], right: list[list[float]]) -> list[list[float]]:
    _, left_cols = _shape(left)
    _, right_cols = _shape(right)
    if left_cols != right_cols:
        raise AttentionCompilerError("Query and key widths must match")
    return [[sum(a * b for a, b in zip(left_row, right_row)) for right_row in right] for left_row in left]


def _softmax_rows(matrix: list[list[float]]) -> list[list[float]]:
    result: list[list[float]] = []
    for row in matrix:
        finite_values = [value for value in row if math.isfinite(value)]
        if not finite_values:
            raise AttentionCompilerError("Softmax row has no finite values")
        row_max = max(finite_values)
        exponentials = [math.exp(value - row_max) if math.isfinite(value) else 0.0 for value in row]
        denominator = sum(exponentials)
        result.append([value / denominator for value in exponentials])
    return result


def _apply_causal_mask(matrix: list[list[float]], enabled: bool) -> list[list[float]]:
    if not enabled:
        return [row[:] for row in matrix]
    return [
        [value if column <= row_index else -1e9 for column, value in enumerate(row)]
        for row_index, row in enumerate(matrix)
    ]


def _round_matrix(matrix: list[list[float]]) -> list[list[float]]:
    return [[round(value, 8) if math.isfinite(value) else value for value in row] for row in matrix]


class AttentionCompiler:
    def operation_graph(self) -> list[TensorOperation]:
        return [
            TensorOperation(operation_id="q-projection", op="matmul", inputs=["embeddings", "Wq"], output="Q"),
            TensorOperation(operation_id="k-projection", op="matmul", inputs=["embeddings", "Wk"], output="K"),
            TensorOperation(operation_id="v-projection", op="matmul", inputs=["embeddings", "Wv"], output="V"),
            TensorOperation(operation_id="attention-logits", op="matmul_transpose_right", inputs=["Q", "K"], output="logits"),
            TensorOperation(operation_id="scale", op="divide_sqrt_dimension", inputs=["logits"], output="scaledLogits", dimension=DK),
            TensorOperation(operation_id="mask", op="apply_causal_mask", inputs=["scaledLogits"], output="maskedLogits"),
            TensorOperation(operation_id="normalize", op="softmax_rows", inputs=["maskedLogits"], output="attentionWeights"),
            TensorOperation(operation_id="aggregate", op="matmul", inputs=["attentionWeights", "V"], output="output"),
        ]

    def compile_operations(
        self,
        operations: list[TensorOperation],
        *,
        causal: bool,
        input_tensors: dict[str, list[list[float]]] | None = None,
    ) -> dict[str, list[list[float]]]:
        expected = [
            ("matmul", ("embeddings", "Wq"), "Q"),
            ("matmul", ("embeddings", "Wk"), "K"),
            ("matmul", ("embeddings", "Wv"), "V"),
            ("matmul_transpose_right", ("Q", "K"), "logits"),
            ("divide_sqrt_dimension", ("logits",), "scaledLogits"),
            ("apply_causal_mask", ("scaledLogits",), "maskedLogits"),
            ("softmax_rows", ("maskedLogits",), "attentionWeights"),
            ("matmul", ("attentionWeights", "V"), "output"),
        ]
        actual = [(op.op, tuple(op.inputs), op.output) for op in operations]
        if actual != expected:
            raise AttentionCompilerError("Operation graph does not match the approved attention operation graph")

        defaults: dict[str, list[list[float]]] = {
            "embeddings": [row[:] for row in EMBEDDINGS],
            "Wq": [row[:] for row in WQ],
            "Wk": [row[:] for row in WK],
            "Wv": [row[:] for row in WV],
        }
        source_inputs = input_tensors or defaults
        required_shapes = {"embeddings": (4, 3), "Wq": (3, 2), "Wk": (3, 2), "Wv": (3, 2)}
        if set(source_inputs) != set(required_shapes):
            raise AttentionCompilerError("Attention inputs must contain only embeddings, Wq, Wk, and Wv")
        for tensor_id, expected_shape in required_shapes.items():
            if _shape(source_inputs[tensor_id]) != expected_shape:
                raise AttentionCompilerError(f"{tensor_id} must have shape {expected_shape}")
            if not all(math.isfinite(value) for row in source_inputs[tensor_id] for value in row):
                raise AttentionCompilerError(f"{tensor_id} contains non-finite values")
        tensors = {tensor_id: [row[:] for row in values] for tensor_id, values in source_inputs.items()}
        for operation in operations:
            if operation.op == "matmul":
                tensors[operation.output] = _matmul(tensors[operation.inputs[0]], tensors[operation.inputs[1]])
            elif operation.op == "matmul_transpose_right":
                tensors[operation.output] = _matmul_transpose_right(tensors[operation.inputs[0]], tensors[operation.inputs[1]])
            elif operation.op == "divide_sqrt_dimension":
                divisor = math.sqrt(operation.dimension)
                tensors[operation.output] = [[value / divisor for value in row] for row in tensors[operation.inputs[0]]]
            elif operation.op == "apply_causal_mask":
                tensors[operation.output] = _apply_causal_mask(tensors[operation.inputs[0]], causal)
            elif operation.op == "softmax_rows":
                tensors[operation.output] = _softmax_rows(tensors[operation.inputs[0]])
            else:
                raise AttentionCompilerError(f"Unsupported operation: {operation.op}")
        return {key: _round_matrix(value) for key, value in tensors.items()}

    def _assertions(self, tensors: dict[str, list[list[float]]], *, causal: bool) -> list[NumericAssertion]:
        results: list[NumericAssertion] = []
        expected_shapes = {"Q": (4, 2), "K": (4, 2), "V": (4, 2), "attentionWeights": (4, 4), "output": (4, 2)}
        for tensor_id, expected in expected_shapes.items():
            passed = _shape(tensors[tensor_id]) == expected
            results.append(NumericAssertion(
                assertion_id=f"shape-{tensor_id}", kind="shape", tensor_id=tensor_id,
                passed=passed, detail=f"expected {expected}, got {_shape(tensors[tensor_id])}",
            ))
        finite = all(math.isfinite(value) for matrix in tensors.values() for row in matrix for value in row)
        results.append(NumericAssertion(assertion_id="finite", kind="finite", tensor_id="all", passed=finite))
        row_sum_ok = all(abs(sum(row) - 1.0) <= TOLERANCE for row in tensors["attentionWeights"])
        results.append(NumericAssertion(assertion_id="attention-row-sum", kind="row_sum", tensor_id="attentionWeights", passed=row_sum_ok))
        masked_ok = True
        if causal:
            masked_ok = all(
                abs(value) <= TOLERANCE
                for row_index, row in enumerate(tensors["attentionWeights"])
                for value in row[row_index + 1 :]
            )
        results.append(NumericAssertion(assertion_id="causal-mask", kind="masked_zero", tensor_id="attentionWeights", passed=masked_ok))
        expected_output = _round_matrix(_matmul(tensors["attentionWeights"], tensors["V"]))
        aggregate_ok = all(
            abs(a - b) <= TOLERANCE
            for actual_row, expected_row in zip(tensors["output"], expected_output)
            for a, b in zip(actual_row, expected_row)
        )
        results.append(NumericAssertion(assertion_id="weighted-output", kind="matmul_close", tensor_id="output", passed=aggregate_ok))
        if not all(result.passed for result in results):
            failed = ", ".join(result.assertion_id for result in results if not result.passed)
            raise AttentionCompilerError(f"Attention assertions failed: {failed}")
        return results

    @staticmethod
    def _safe_tokens(tokens: Iterable[str]) -> list[str]:
        candidate = [str(token).strip() for token in tokens]
        if len(candidate) != 4 or any(not token or len(token) > 18 or any(ch in token for ch in "<>\n\r") for token in candidate):
            return DEFAULT_TOKENS[:]
        return candidate

    @staticmethod
    def _panels(claim_ids: list[str]) -> list[PanelSpec]:
        definitions = [
            ("tokens", "token_strip", "Tokens", ["tokens"]),
            ("projections", "tensor_matrix", "Queries, keys, and values", ["Q", "K", "V"]),
            ("scores", "tensor_matrix", "Scaled query-key scores", ["logits", "scaledLogits"]),
            ("heatmap", "attention_heatmap", "Attention weights", ["attentionWeights"]),
            ("connections", "weighted_attention", "Weighted connections", ["attentionWeights"]),
            ("aggregation", "vector_aggregation", "Weighted value aggregation", ["V", "output"]),
            ("equation", "equation", "Scaled dot-product attention", ["Q", "K", "V", "attentionWeights", "output"]),
            ("evidence", "evidence", "Evidence and assumptions", []),
        ]
        return [
            PanelSpec(panel_id=panel_id, panel_type=panel_type, title=title, bindings=bindings, claim_ids=claim_ids, order=index)
            for index, (panel_id, panel_type, title, bindings) in enumerate(definitions)
        ]

    @staticmethod
    def _timeline(draft: AttentionLessonDraft, claim_ids: list[str]) -> list[TimelineStep]:
        defaults = [
            "Start with four tokens and small illustrative embeddings.",
            "Project each embedding into query, key, and value vectors.",
            "Compare every query with every key using dot products.",
            "Scale logits by the square root of the key dimension.",
            "Apply the optional causal mask before normalization.",
            "Normalize each score row into attention weights.",
            "Use the weights to combine value vectors into outputs.",
        ]
        explanations = draft.teaching_steps if len(draft.teaching_steps) == 7 else defaults
        definitions = [
            ("inputs", "Inputs", [], ["tokens"]),
            ("project", "Project Q, K, and V", ["q-projection", "k-projection", "v-projection"], ["tokens", "projections"]),
            ("compare", "Compare queries and keys", ["attention-logits"], ["scores", "connections"]),
            ("scale", "Scale the logits", ["scale"], ["scores", "equation"]),
            ("mask", "Apply a causal mask", ["mask"], ["scores", "heatmap"]),
            ("softmax", "Normalize with softmax", ["normalize"], ["heatmap", "connections", "equation"]),
            ("aggregate", "Aggregate the values", ["aggregate"], ["aggregation", "equation"]),
        ]
        return [
            TimelineStep(
                step_id=step_id, label=label, explanation=explanations[index], operation_ids=operation_ids,
                active_panel_ids=active_panels, claim_ids=claim_ids,
            )
            for index, (step_id, label, operation_ids, active_panels) in enumerate(definitions)
        ]

    def compile_spec(self, spec: VisualLessonSpec) -> CompiledLesson:
        input_tensors = {item.tensor_id: item.values for item in spec.inputs}
        branches: list[CompiledBranch] = []
        for branch_id, causal in (("unmasked", False), ("causal", True)):
            tensors = self.compile_operations(spec.operations, causal=causal, input_tensors=input_tensors)
            assertions = self._assertions(tensors, causal=causal)
            branches.append(CompiledBranch(branch_id=branch_id, tensors=tensors, assertions=assertions))
        return CompiledLesson(branches=branches, assertions_passed=True)

    def build_payload(
        self,
        *,
        project_id: str,
        prompt: str,
        draft: AttentionLessonDraft,
        evidence_sources: list[EvidenceSource],
        evidence_claims: list[EvidenceClaim],
        warnings: list[str],
    ) -> VisualLessonPayload:
        claim_ids = [claim.claim_id for claim in evidence_claims]
        tokens = self._safe_tokens(draft.tokens)
        operations = self.operation_graph()
        spec = VisualLessonSpec(
            project_id=project_id,
            prompt=prompt,
            title=draft.title or "Inside scaled dot-product attention",
            interpretation=draft.interpretation,
            requested_variant=draft.requested_variant,
            lesson_variant=draft.lesson_variant,
            variant_notice=draft.variant_notice,
            tokens=tokens,
            assumptions=draft.assumptions,
            evidence_sources=evidence_sources,
            evidence_claims=evidence_claims,
            inputs=[
                TensorInput(tensor_id="embeddings", label="Illustrative embeddings", values=EMBEDDINGS),
                TensorInput(tensor_id="Wq", label="Illustrative query projection", values=WQ),
                TensorInput(tensor_id="Wk", label="Illustrative key projection", values=WK),
                TensorInput(tensor_id="Wv", label="Illustrative value projection", values=WV),
            ],
            operations=operations,
            assertions=[
                NumericAssertion(assertion_id="attention-row-sum", kind="row_sum", tensor_id="attentionWeights"),
                NumericAssertion(assertion_id="causal-mask", kind="masked_zero", tensor_id="attentionWeights"),
                NumericAssertion(assertion_id="weighted-output", kind="matmul_close", tensor_id="output"),
            ],
            panels=self._panels(claim_ids),
            timeline=self._timeline(draft, claim_ids),
            created_at=time.time(),
        )
        return VisualLessonPayload(spec=spec, compiled=self.compile_spec(spec), warnings=warnings)