"""Shared-space positional and spatial encoding for every modality."""

from __future__ import annotations

import math

import torch

from ._source_bound import SourceBoundModule
from .configuration_dendro_omni import DendroOmniConfig
from .source import DendroSourceLayer


MODALITY_TEXT = 0
MODALITY_IMAGE = 1
MODALITY_AUDIO = 2
MODALITY_VIDEO = 3
MODALITY_SENSOR = 4
MODALITY_MEMORY = 5
MODALITY_WORKSPACE = 6
MODALITY_REASONING = 7

MODALITY_NAMES = {
    MODALITY_TEXT: "text",
    MODALITY_IMAGE: "image",
    MODALITY_AUDIO: "audio",
    MODALITY_VIDEO: "video",
    MODALITY_SENSOR: "sensor",
    MODALITY_MEMORY: "memory",
    MODALITY_WORKSPACE: "workspace",
    MODALITY_REASONING: "reasoning",
}

# Logical offsets ensure modality-local coordinates never collide semantically even
# though physical packed sequence positions remain contiguous for cache handling.
MODALITY_POSITION_OFFSETS = {
    MODALITY_TEXT: 0,
    MODALITY_IMAGE: 1_000_000,
    MODALITY_AUDIO: 2_000_000,
    MODALITY_VIDEO: 3_000_000,
    MODALITY_SENSOR: 4_000_000,
    MODALITY_MEMORY: 5_000_000,
    MODALITY_WORKSPACE: 6_000_000,
    MODALITY_REASONING: 7_000_000,
}


class DendroSpatialEncoder(SourceBoundModule):
    """Map sequence, modality and N-D coordinates into the shared hidden space."""

    def __init__(self, config: DendroOmniConfig, source: DendroSourceLayer) -> None:
        super().__init__(source)
        self.config = config

    def _fourier_features(self, coordinates: torch.Tensor, logical_positions: torch.Tensor) -> torch.Tensor:
        dtype = coordinates.dtype
        bands = self.config.spatial_fourier_bands
        frequencies = torch.pow(
            torch.tensor(2.0, device=coordinates.device, dtype=dtype),
            torch.arange(bands, device=coordinates.device, dtype=dtype),
        )
        phase = coordinates.unsqueeze(-1) * frequencies * math.pi
        spatial = torch.cat([coordinates, phase.sin().flatten(-2), phase.cos().flatten(-2)], dim=-1)

        # Logical offsets are encoded continuously rather than through a giant table.
        logical = logical_positions.to(dtype=dtype).unsqueeze(-1) / 1_000_000.0
        logical_phase = logical * frequencies * math.pi
        logical_features = torch.cat([logical, logical_phase.sin(), logical_phase.cos()], dim=-1)
        return torch.cat([spatial, logical_features], dim=-1)

    def forward(
        self,
        hidden: torch.Tensor,
        *,
        modality_ids: torch.Tensor,
        sequence_positions: torch.Tensor,
        logical_positions: torch.Tensor,
        coordinates: torch.Tensor,
        is_prefix: torch.Tensor,
    ) -> torch.Tensor:
        if coordinates.shape[-1] != 4:
            raise ValueError("coordinates must have four axes: temporal/sequence, vertical, horizontal, frequency")
        source = self.source
        hidden_size = self.config.hidden_size
        modality = source.embedding(
            modality_ids,
            "spatial/modality",
            self.config.modality_vocab_size,
            hidden_size,
        )
        role = source.embedding(is_prefix.long(), "spatial/prefix_role", 2, hidden_size)
        features = self._fourier_features(coordinates.to(hidden.dtype), logical_positions)
        spatial = source.project(features, "spatial/fourier", hidden_size, low_bit=False)

        # A bounded sequence code improves recurrence-depth distinction without a
        # max-position table.  It remains valid beyond the training context window.
        seq = sequence_positions.to(hidden.dtype).unsqueeze(-1)
        inv = torch.exp(
            -math.log(self.config.rope_theta)
            * torch.arange(0, hidden_size, 2, device=hidden.device, dtype=hidden.dtype)
            / max(1, hidden_size)
        )
        seq_phase = seq * inv
        seq_code = torch.stack([seq_phase.sin(), seq_phase.cos()], dim=-1).flatten(-2)
        if seq_code.shape[-1] < hidden_size:
            seq_code = torch.nn.functional.pad(seq_code, (0, hidden_size - seq_code.shape[-1]))
        seq_code = seq_code[..., :hidden_size]
        seq_gate = source.gate(hidden, "spatial/sequence_gate", hidden_size)

        return hidden + 0.20 * modality + 0.10 * role + 0.20 * spatial + 0.10 * seq_gate * seq_code


def apply_rotary_position_embedding(
    q: torch.Tensor,
    k: torch.Tensor,
    positions: torch.Tensor,
    *,
    theta: float,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Apply stable RoPE to ``[batch, heads, sequence, head_dim]`` Q and K."""

    head_dim = q.shape[-1]
    if head_dim % 2:
        raise ValueError("RoPE requires an even head dimension")
    inv_freq = torch.exp(
        -math.log(theta)
        * torch.arange(0, head_dim, 2, device=q.device, dtype=torch.float32)
        / head_dim
    )
    phase = positions.to(device=q.device, dtype=torch.float32).unsqueeze(-1) * inv_freq
    cos = phase.cos().to(q.dtype).unsqueeze(1)
    sin = phase.sin().to(q.dtype).unsqueeze(1)

    def rotate(x: torch.Tensor) -> torch.Tensor:
        even, odd = x[..., 0::2], x[..., 1::2]
        rotated_even = even * cos - odd * sin
        rotated_odd = even * sin + odd * cos
        return torch.stack([rotated_even, rotated_odd], dim=-1).flatten(-2)

    return rotate(q), rotate(k)