"""Deterministic evaluation-time interventions for predicate memory.

These interventions are not training features. They are causal probes for asking
whether the learned predicate-memory path is example-specific, span-anchored, and
predicate-gated, rather than just extra residual capacity.
"""

from __future__ import annotations

from typing import Literal

import torch

PredicateMemoryIntervention = Literal[
    "none",
    "typed_graph_memory",
    "zero_memory",
    "memory_zero",
    "uniform_gate",
    "untyped_same_topology",
    "span_roll",
    "span_corrupted",
    "span_reverse",
    "batch_shuffle",
    "batch_shuffled_graph",
    "random_same_degree",
]

PREDICATE_MEMORY_INTERVENTIONS: tuple[str, ...] = (
    "none",
    "typed_graph_memory",
    "zero_memory",
    "memory_zero",
    "uniform_gate",
    "untyped_same_topology",
    "span_roll",
    "span_corrupted",
    "span_reverse",
    "batch_shuffle",
    "batch_shuffled_graph",
    "random_same_degree",
)


def _valid_lengths(attention_mask: torch.Tensor | None, batch: int, seq: int) -> list[int]:
    if attention_mask is None:
        return [seq] * batch
    return [int(x) for x in attention_mask.to(torch.long).sum(dim=1).tolist()]


def _roll_valid(x: torch.Tensor, attention_mask: torch.Tensor | None, *, shifts: int = 1) -> torch.Tensor:
    out = x.clone()
    b, s = x.shape[:2]
    for i, n in enumerate(_valid_lengths(attention_mask, b, s)):
        if n > 1:
            out[i, :n] = torch.roll(x[i, :n], shifts=shifts, dims=0)
    return out


def _reverse_valid(x: torch.Tensor, attention_mask: torch.Tensor | None) -> torch.Tensor:
    out = x.clone()
    b, s = x.shape[:2]
    for i, n in enumerate(_valid_lengths(attention_mask, b, s)):
        if n > 1:
            out[i, :n] = torch.flip(x[i, :n], dims=(0,))
    return out


def _random_permute_valid(x: torch.Tensor, attention_mask: torch.Tensor | None) -> torch.Tensor:
    """Deterministically permute valid memory slots while preserving the budget."""

    out = x.clone()
    b, s = x.shape[:2]
    for i, n in enumerate(_valid_lengths(attention_mask, b, s)):
        if n > 1:
            generator = torch.Generator(device="cpu")
            generator.manual_seed(1729 + i * 1_000_003 + n * 97)
            perm = torch.randperm(n, generator=generator).to(x.device)
            out[i, :n] = x[i, perm]
    return out


def apply_predicate_memory_intervention(
    *,
    candidate_key: torch.Tensor,
    candidate_value: torch.Tensor,
    predicate_gate: torch.Tensor,
    attention_mask: torch.Tensor | None,
    intervention: str = "none",
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """Return intervened ``(key, value, gate)`` tensors.

    ``zero_memory`` makes the predicate-memory update exactly zero by zeroing the
    values read by attention. The other interventions preserve tensor shapes and
    most marginal statistics while corrupting a specific part of the learned
    predicate-memory computation.
    """

    if intervention in {"none", "typed_graph_memory"}:
        return candidate_key, candidate_value, predicate_gate
    if intervention not in PREDICATE_MEMORY_INTERVENTIONS:
        raise ValueError(f"unknown predicate memory intervention {intervention!r}")

    if intervention in {"zero_memory", "memory_zero"}:
        return candidate_key, torch.zeros_like(candidate_value), predicate_gate

    if intervention in {"uniform_gate", "untyped_same_topology"}:
        gate = torch.ones_like(predicate_gate)
        if attention_mask is not None:
            gate = gate * attention_mask.to(gate.dtype).unsqueeze(-1)
        return candidate_key, candidate_value, gate

    if intervention in {"span_roll", "span_corrupted"}:
        return (
            _roll_valid(candidate_key, attention_mask),
            _roll_valid(candidate_value, attention_mask),
            _roll_valid(predicate_gate, attention_mask),
        )

    if intervention == "span_reverse":
        return (
            _reverse_valid(candidate_key, attention_mask),
            _reverse_valid(candidate_value, attention_mask),
            _reverse_valid(predicate_gate, attention_mask),
        )

    if intervention in {"batch_shuffle", "batch_shuffled_graph"}:
        if candidate_key.shape[0] > 1:
            return (
                torch.roll(candidate_key, shifts=1, dims=0),
                torch.roll(candidate_value, shifts=1, dims=0),
                torch.roll(predicate_gate, shifts=1, dims=0),
            )
        return (
            _roll_valid(candidate_key, attention_mask),
            _roll_valid(candidate_value, attention_mask),
            _roll_valid(predicate_gate, attention_mask),
        )

    if intervention == "random_same_degree":
        return (
            _random_permute_valid(candidate_key, attention_mask),
            _random_permute_valid(candidate_value, attention_mask),
            _random_permute_valid(predicate_gate, attention_mask),
        )

    raise AssertionError(f"unhandled intervention {intervention!r}")


__all__ = [
    "PREDICATE_MEMORY_INTERVENTIONS",
    "PredicateMemoryIntervention",
    "apply_predicate_memory_intervention",
]