nesa_core.py

"""
Project NESA (Non-Executable Semantic Architecture)
Core Implementation: SSoA Gate with Kinematic Clipping and Topological Anchoring
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Optional, Tuple, Dict, List
from dataclasses import dataclass
import math


@dataclass
class NESAConfig:
    """Configuration for NESA security layer"""

    tau_jerk: float = 0.6  # Local jerk threshold (will be calibrated)
    delta_drift: float = 0.45  # Global drift threshold (will be calibrated)
    executive_head_threshold: float = (
        0.7  # Gradient magnitude threshold for executive heads
    )
    window_size: int = 16  # Window for kinematic calculations
    use_sparse_attention: bool = True  # Use sparse masking for efficiency
    async_monitor: bool = (
        True  # Run monitor in separate stream (when available)
    )

    # Optimization flags
    use_triton: bool = False  # Use Triton kernels (when available)
    use_flash_attention: bool = True  # Use FlashAttention-2 if available
    kv_cache_aware: bool = True  # Optimize for KV-cache scenarios


class KinematicMonitor(nn.Module):
    """
    Calculates local jerk (3rd derivative) of embedding trajectory
    to detect semantic discontinuities indicative of injection attacks.

    Uses fused CUDA operations via torch.compile for efficiency.
    """

    def __init__(self, config: NESAConfig, embedding_dim: int):
        super().__init__()
        self.config = config
        self.embedding_dim = embedding_dim

        # Moving average for incremental drift calculation (O(1) per token)
        self.register_buffer("moving_avg_anchor", torch.zeros(embedding_dim))
        self.register_buffer("anchor_count", torch.tensor(0))

    @torch.compile(mode="reduce-overhead")
    def compute_velocity(self, embeddings: torch.Tensor) -> torch.Tensor:
        """
        Compute velocity (1st derivative) of embedding trajectory.
        embeddings: [batch, seq_len, dim]
        returns: [batch, seq_len-1, dim]
        """
        return embeddings[:, 1:, :] - embeddings[:, :-1, :]

    @torch.compile(mode="reduce-overhead")
    def compute_acceleration(self, velocity: torch.Tensor) -> torch.Tensor:
        """
        Compute acceleration (2nd derivative).
        velocity: [batch, seq_len-1, dim]
        returns: [batch, seq_len-2, dim]
        """
        return velocity[:, 1:, :] - velocity[:, :-1, :]

    @torch.compile(mode="reduce-overhead")
    def compute_jerk(self, acceleration: torch.Tensor) -> torch.Tensor:
        """
        Compute jerk (3rd derivative) - rate of change of acceleration.
        acceleration: [batch, seq_len-2, dim]
        returns: [batch, seq_len-3, dim]
        """
        return acceleration[:, 1:, :] - acceleration[:, :-1, :]

    @torch.compile(mode="reduce-overhead")
    def compute_jerk_magnitude(self, embeddings: torch.Tensor) -> torch.Tensor:
        """
        Fused computation of jerk magnitude from embeddings.
        This is the core kinematic detector for semantic discontinuities.

        embeddings: [batch, seq_len, dim]
        returns: [batch, seq_len-3] (scalar jerk magnitude per token)
        """
        # Sequential derivatives
        v = self.compute_velocity(embeddings)  # [B, L-1, D]
        a = self.compute_acceleration(v)  # [B, L-2, D]
        j = self.compute_jerk(a)  # [B, L-3, D]

        # Magnitude (L2 norm across embedding dimension)
        jerk_magnitude = torch.norm(j, p=2, dim=-1)  # [B, L-3]

        return jerk_magnitude

    def forward(
        self,
        embeddings: torch.Tensor,
        sovereign_anchor: torch.Tensor,
        return_drift: bool = True,
    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
        """
        Main forward pass of Kinematic Monitor.

        Args:
            embeddings: [batch, seq_len, dim] - token embeddings
            sovereign_anchor: [dim] - the Ω₀ vector (system prompt embedding)
            return_drift: whether to compute global drift

        Returns:
            jerk_mask: [batch, seq_len] - binary mask (0 = clip, 1 = allow)
            drift_scores: [batch, seq_len] - cosine distance from anchor (optional)
        """
        batch_size, seq_len, dim = embeddings.shape

        # === LOCAL JERK DETECTION ===
        jerk_magnitude = self.compute_jerk_magnitude(embeddings)  # [B, L-3]

        # Pad to match sequence length (first 3 tokens have no jerk)
        # These are typically system/instruction tokens, so we trust them
        jerk_padded = F.pad(jerk_magnitude, (3, 0), value=0.0)  # [B, L]

        # Create binary mask: 1 if jerk below threshold (safe), 0 if above (clip)
        jerk_mask = (jerk_padded < self.config.tau_jerk).float()

        # === GLOBAL DRIFT DETECTION ===
        drift_scores = None
        if return_drift:
            # Normalize embeddings and anchor for cosine similarity
            embeddings_norm = F.normalize(embeddings, p=2, dim=-1)  # [B, L, D]
            anchor_norm = F.normalize(
                sovereign_anchor.unsqueeze(0), p=2, dim=-1
            ).to(embeddings_norm.dtype)  # [1, D]

            # Cosine similarity to anchor
            cosine_sim = torch.matmul(embeddings_norm, anchor_norm.T).squeeze(
                -1
            )  # [B, L]

            # Drift is 1 - similarity (distance from anchor)
            drift_scores = 1.0 - cosine_sim  # [B, L]

            # Apply drift threshold to mask
            drift_mask = (drift_scores < self.config.delta_drift).float()

            # Combine jerk and drift masks (both must pass)
            jerk_mask = jerk_mask * drift_mask

        return jerk_mask, drift_scores


class SovereignBuffer(nn.Module):
    """
    Stores the Sovereign Root Authority (Ω₀) and manages the
    "Mission Vector" for the current session.

    Persistent buffer that doesn't require gradient computation.
    """

    def __init__(self, embedding_dim: int):
        super().__init__()
        self.embedding_dim = embedding_dim

        # The Sovereign Root - initialized from system prompt
        self.register_buffer("omega_zero", torch.zeros(embedding_dim))
        self.register_buffer("is_initialized", torch.tensor(False))

        # Mission Vector - weighted average of system + first user command
        self.register_buffer("mission_vector", torch.zeros(embedding_dim))
        self.register_buffer(
            "mission_weight", torch.tensor(0.7)
        )  # Weight for system prompt

    def initialize(
        self,
        system_prompt_embedding: torch.Tensor,
        first_command_embedding: Optional[torch.Tensor] = None,
    ):
        """
        Initialize the Sovereign Authority.

        Args:
            system_prompt_embedding: [dim] or [seq_len, dim] - system prompt
            first_command_embedding: [dim] or [seq_len, dim] - first user command (optional)
        """
        # Handle sequence dimension if present
        if system_prompt_embedding.dim() > 1:
            # Average pool across sequence
            system_prompt_embedding = system_prompt_embedding.mean(dim=0)

        # Set Ω₀
        self.omega_zero.copy_(system_prompt_embedding)

        # Create mission vector
        if first_command_embedding is not None:
            if first_command_embedding.dim() > 1:
                first_command_embedding = first_command_embedding.mean(dim=0)

            # Weighted average: mission = α * system + (1-α) * command
            self.mission_vector.copy_(
                self.mission_weight * system_prompt_embedding
                + (1 - self.mission_weight) * first_command_embedding
            )
        else:
            # If no command yet, mission = system
            self.mission_vector.copy_(system_prompt_embedding)

        self.is_initialized.fill_(True)

    def get_anchor(self, use_mission: bool = True) -> torch.Tensor:
        """
        Get the current authority anchor.

        Args:
            use_mission: if True, use mission_vector; else use omega_zero

        Returns:
            anchor: [dim] - the authority vector
        """
        if not self.is_initialized:
            raise RuntimeError(
                "SovereignBuffer not initialized! Call initialize() first."
            )

        return self.mission_vector if use_mission else self.omega_zero

    def update_mission(
        self, new_command_embedding: torch.Tensor, alpha: float = 0.1
    ):
        """
        Incrementally update mission vector with new authorized commands.
        Uses exponential moving average to maintain continuity.

        Args:
            new_command_embedding: [dim] - embedding of new command
            alpha: learning rate for EMA update
        """
        if new_command_embedding.dim() > 1:
            new_command_embedding = new_command_embedding.mean(dim=0)

        # EMA update: mission_new = (1-α) * mission_old + α * new_command
        self.mission_vector.mul_(1 - alpha).add_(
            new_command_embedding, alpha=alpha
        )


class HeadProbe(nn.Module):
    """
    Identifies Executive (instruction-following) vs Perceptual (feature-extracting)
    attention heads using gradient-based attribution.

    This runs once during model setup, not during inference.
    """

    def __init__(self, config: NESAConfig, num_layers: int, num_heads: int):
        super().__init__()
        self.config = config
        self.num_layers = num_layers
        self.num_heads = num_heads

        # Store head classifications: [num_layers, num_heads]
        # 1 = Executive, 0 = Perceptual
        self.register_buffer(
            "head_classification",
            torch.zeros(num_layers, num_heads, dtype=torch.bool),
        )
        self.register_buffer("is_probed", torch.tensor(False))

    def probe_heads(
        self,
        model: nn.Module,
        instruction_dataset: List[Dict[str, torch.Tensor]],
        num_samples: int = 50,
    ):
        """
        Run gradient-based attribution to identify Executive heads.

        Args:
            model: The transformer model to probe
            instruction_dataset: List of {input_ids, labels} for instruction tasks
            num_samples: Number of samples to use for probing
        """
        print(
            "🔍 Probing attention heads for Executive/Perceptual classification..."
        )

        # Store gradient magnitudes per head
        head_gradients = torch.zeros(self.num_layers, self.num_heads)

        model.train()  # Need gradients

        for i, sample in enumerate(instruction_dataset[:num_samples]):
            if i % 10 == 0:
                print(f"  Processing sample {i}/{num_samples}")

            input_ids = sample["input_ids"].unsqueeze(0)  # [1, seq_len]
            labels = sample.get("labels", input_ids)

            # Forward pass
            outputs = model(input_ids, labels=labels)
            loss = outputs.loss

            # Backward pass
            loss.backward()

            # Collect gradients from attention heads
            for layer_idx in range(self.num_layers):
                # Access attention module (model-specific - assumes Llama/Mistral structure)
                try:
                    attn_module = model.model.layers[layer_idx].self_attn

                    # Get gradients from query/key/value projections
                    if (
                        hasattr(attn_module, "q_proj")
                        and attn_module.q_proj.weight.grad is not None
                    ):
                        q_grad = attn_module.q_proj.weight.grad.abs().mean()

                        # Accumulate gradient magnitude per head
                        # Note: this is a simplified version - in practice we'd track per-head
                        head_gradients[layer_idx, :] += q_grad.item()
                except:
                    pass

            # Zero gradients for next iteration
            model.zero_grad()

        # Normalize and threshold
        head_gradients /= num_samples
        threshold = (
            head_gradients.mean()
            + self.config.executive_head_threshold * head_gradients.std()
        )

        # Classify heads: high gradient = Executive (instruction-following)
        self.head_classification = head_gradients > threshold

        num_executive = self.head_classification.sum().item()
        total_heads = self.num_layers * self.num_heads

        print(
            f" Probing complete: {num_executive}/{total_heads} heads classified as Executive"
        )
        print(
            f"   Distribution by layer: {self.head_classification.sum(dim=1).tolist()}"
        )

        self.is_probed.fill_(True)
        model.eval()

    def get_executive_mask(self, layer_idx: int) -> torch.Tensor:
        """
        Get binary mask for executive heads at a specific layer.

        Args:
            layer_idx: Layer index

        Returns:
            mask: [num_heads] - 1 for executive heads, 0 for perceptual
        """
        if not self.is_probed:
            raise RuntimeError("Heads not probed! Call probe_heads() first.")

        return self.head_classification[layer_idx].float()


class NESAAttentionWrapper(nn.Module):
    """
    Wraps the standard attention mechanism with NESA security.
    Applies M_s mask to clip non-executive packets from reaching executive heads.

    Uses torch.nn.functional.scaled_dot_product_attention with custom attn_mask.
    """

    def __init__(
        self,
        original_attention: nn.Module,
        config: NESAConfig,
        layer_idx: int,
        num_heads: int,
        head_dim: int,
    ):
        super().__init__()
        self.original_attention = original_attention
        self.config = config
        self.layer_idx = layer_idx
        self.num_heads = num_heads
        self.head_dim = head_dim

    def apply_sovereign_mask(
        self,
        attn_weights: torch.Tensor,
        semantic_packet_mask: torch.Tensor,
        executive_head_mask: torch.Tensor,
    ) -> torch.Tensor:
        """
        Apply M_s mask: clip unsigned packets from executive heads.

        Args:
            attn_weights: [batch, num_heads, seq_len, seq_len] - attention scores
            semantic_packet_mask: [batch, seq_len] - 1 for authorized, 0 for clipped
            executive_head_mask: [num_heads] - 1 for executive, 0 for perceptual

        Returns:
            masked_weights: [batch, num_heads, seq_len, seq_len]
        """
        batch_size, _, seq_len, _ = attn_weights.shape

        # Expand masks to match attention dimensions
        # semantic_packet_mask: [B, 1, 1, L] - broadcasts across heads and query positions
        packet_mask = semantic_packet_mask.unsqueeze(1).unsqueeze(
            2
        )  # [B, 1, 1, L]

        # executive_head_mask: [1, H, 1, 1] - broadcasts across batch and sequence
        exec_mask = executive_head_mask.view(1, -1, 1, 1)  # [1, H, 1, 1]

        # M_s logic:
        # - For Executive heads: mask = packet_mask (only authorized packets)
        # - For Perceptual heads: mask = all ones (no filtering)

        # Create the sovereign mask
        sovereign_mask = torch.where(
            exec_mask.bool(),
            packet_mask,  # Executive: use packet authorization
            torch.ones_like(packet_mask),  # Perceptual: allow all
        )

        # Apply mask (set unauthorized attention to -inf for softmax)
        # Note: we're masking the VALUE dimension (keys being attended TO)
        masked_weights = attn_weights.masked_fill(
            sovereign_mask == 0, float("-inf")
        )

        return masked_weights

    def forward(
        self,
        hidden_states: torch.Tensor,
        semantic_packet_mask: torch.Tensor,
        executive_head_mask: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        **kwargs,
    ) -> torch.Tensor:
        """
        Forward pass with NESA security.

        Args:
            hidden_states: [batch, seq_len, hidden_dim]
            semantic_packet_mask: [batch, seq_len] - output from KinematicMonitor
            executive_head_mask: [num_heads] - from HeadProbe
            attention_mask: Optional standard attention mask

        Returns:
            output: [batch, seq_len, hidden_dim]
        """
        # Get Q, K, V from original attention module
        # Note: This assumes Llama/Mistral architecture - may need adjustment
        batch_size, seq_len, _ = hidden_states.shape

        # Project to Q, K, V
        q = self.original_attention.q_proj(hidden_states)
        k = self.original_attention.k_proj(hidden_states)
        v = self.original_attention.v_proj(hidden_states)

        # Reshape for multi-head attention
        q = q.view(
            batch_size, seq_len, self.num_heads, self.head_dim
        ).transpose(1, 2)
        k = k.view(
            batch_size, seq_len, self.num_heads, self.head_dim
        ).transpose(1, 2)
        v = v.view(
            batch_size, seq_len, self.num_heads, self.head_dim
        ).transpose(1, 2)

        # Compute attention scores
        attn_weights = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(
            self.head_dim
        )

        # Apply sovereign mask
        attn_weights = self.apply_sovereign_mask(
            attn_weights, semantic_packet_mask, executive_head_mask
        )

        # Apply standard attention mask if provided
        if attention_mask is not None:
            attn_weights = attn_weights + attention_mask

        # Softmax and apply to values
        attn_probs = F.softmax(attn_weights, dim=-1)
        attn_output = torch.matmul(attn_probs, v)

        # Reshape back
        attn_output = attn_output.transpose(1, 2).contiguous()
        attn_output = attn_output.view(batch_size, seq_len, -1)

        # Final projection
        output = self.original_attention.o_proj(attn_output)

        return output


def create_nesa_mask_for_sdpa(
    semantic_packet_mask: torch.Tensor,
    executive_head_mask: torch.Tensor,
    seq_len: int,
    num_heads: int,
    device: torch.device,
) -> torch.Tensor:
    """
    Create attention mask compatible with scaled_dot_product_attention.

    Args:
        semantic_packet_mask: [batch, seq_len]
        executive_head_mask: [num_heads]
        seq_len: sequence length
        num_heads: number of attention heads
        device: torch device

    Returns:
        mask: [batch, num_heads, seq_len, seq_len] - additive mask for SDPA
    """
    batch_size = semantic_packet_mask.shape[0]

    # Expand masks
    packet_mask = semantic_packet_mask.unsqueeze(1).unsqueeze(
        2
    )  # [B, 1, 1, L]
    exec_mask = executive_head_mask.view(1, -1, 1, 1)  # [1, H, 1, 1]

    # Create sovereign mask
    sovereign_mask = torch.where(
        exec_mask.bool(), packet_mask, torch.ones_like(packet_mask)
    )

    # Expand to full attention shape
    sovereign_mask = sovereign_mask.expand(
        batch_size, num_heads, seq_len, seq_len
    )

    # Convert to additive mask (0 for allowed, -inf for blocked)
    additive_mask = torch.zeros_like(sovereign_mask, dtype=torch.float32)
    additive_mask = additive_mask.masked_fill(
        sovereign_mask == 0, float("-inf")
    )

    return additive_mask


if __name__ == "__main__":
    print(" NESA Core Modules Loaded")
    print("=" * 60)

    # Quick smoke test
    config = NESAConfig()
    embedding_dim = 4096

    # Test Kinematic Monitor
    monitor = KinematicMonitor(config, embedding_dim)
    test_embeddings = torch.randn(2, 20, embedding_dim)
    test_anchor = torch.randn(embedding_dim)

    jerk_mask, drift = monitor(test_embeddings, test_anchor)
    print(f" Kinematic Monitor: jerk_mask shape = {jerk_mask.shape}")
    print(f"   Detected {(jerk_mask == 0).sum().item()} clipped tokens")

    # Test Sovereign Buffer
    buffer = SovereignBuffer(embedding_dim)
    system_emb = torch.randn(embedding_dim)
    buffer.initialize(system_emb)
    anchor = buffer.get_anchor()
    print(f" Sovereign Buffer: initialized with anchor shape = {anchor.shape}")

    # Test Head Probe
    probe = HeadProbe(config, num_layers=32, num_heads=32)
    print(
        f" Head Probe: ready for {probe.num_layers} layers x {probe.num_heads} heads"
    )

    print("=" * 60)
    print(" All core modules initialized successfully!")