modalities.py

"""Real multimodal patchification and packing into one Dendro token space."""

from __future__ import annotations

from dataclasses import dataclass
from typing import Any

import torch
from torch.nn import functional as F

from ._source_bound import SourceBoundModule
from .configuration_dendro_omni import DendroOmniConfig
from .source import DendroSourceLayer
from .spatial import (
    MODALITY_AUDIO,
    MODALITY_IMAGE,
    MODALITY_POSITION_OFFSETS,
    MODALITY_SENSOR,
    MODALITY_TEXT,
    MODALITY_VIDEO,
    DendroSpatialEncoder,
)


@dataclass(slots=True)
class ModalitySegment:
    name: str
    modality_id: int
    start: int
    end: int
    shape: tuple[int, ...]

    @property
    def length(self) -> int:
        return self.end - self.start


@dataclass(slots=True)
class DendroModalityLayout:
    modality_ids: torch.Tensor
    sequence_positions: torch.Tensor
    logical_positions: torch.Tensor
    coordinates: torch.Tensor
    is_prefix: torch.Tensor
    attention_mask: torch.Tensor
    segments: tuple[ModalitySegment, ...]
    text_start: int
    text_length: int

    def to_dict(self) -> dict[str, Any]:
        return {
            "segments": [
                {
                    "name": segment.name,
                    "modality_id": segment.modality_id,
                    "start": segment.start,
                    "end": segment.end,
                    "length": segment.length,
                    "shape": segment.shape,
                }
                for segment in self.segments
            ],
            "text_start": self.text_start,
            "text_length": self.text_length,
            "total_length": int(self.modality_ids.shape[-1]),
            "prefix_length": int(self.is_prefix[0].sum().item()) if self.is_prefix.numel() else 0,
        }


@dataclass(slots=True)
class DendroPackedInput:
    hidden_states: torch.Tensor
    layout: DendroModalityLayout
    aligned_labels: torch.Tensor | None = None


class DendroOmniInputProjector(SourceBoundModule):
    """Parameterless modality adapters backed entirely by ``DendroSourceLayer``."""

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

    @staticmethod
    def _batch_size(*values: torch.Tensor | None) -> int:
        batches = [int(value.shape[0]) for value in values if value is not None]
        if not batches:
            raise ValueError("At least one text, image, audio, video, or sensor input is required")
        if any(batch != batches[0] for batch in batches):
            raise ValueError(f"All modalities must share a batch size, got {batches}")
        return batches[0]

    @staticmethod
    def _normalize_coords(index: torch.Tensor, maximum: int) -> torch.Tensor:
        if maximum <= 1:
            return torch.zeros_like(index, dtype=torch.float32)
        return index.float() / float(maximum - 1) * 2.0 - 1.0

    def _token_features(self, input_ids: torch.Tensor, token_hidden: torch.Tensor) -> torch.Tensor:
        source = self.source
        offset = self.config.byte_offset
        byte = input_ids - offset
        atom_ids = torch.zeros_like(input_ids)
        atom_ids = torch.where((byte >= ord("0")) & (byte <= ord("9")), 1, atom_ids)
        atom_ids = torch.where(
            ((byte >= ord("A")) & (byte <= ord("Z"))) | ((byte >= ord("a")) & (byte <= ord("z"))),
            2,
            atom_ids,
        )
        atom_ids = torch.where((byte == 9) | (byte == 10) | (byte == 13) | (byte == 32), 3, atom_ids)
        atom_ids = torch.where((byte >= 128) & (byte <= 255), 4, atom_ids)
        atom_ids = torch.where(input_ids < offset, 5, atom_ids)
        atoms = source.embedding(atom_ids, "token/atoms", 6, self.config.hidden_size)

        # Atom -> bond -> molecule composition is deliberately pointwise here. Any
        # cross-token neighborhood operation belongs inside cache-aware attention;
        # otherwise a one-token decode chunk would not match full-sequence training.
        bond_input = token_hidden * torch.tanh(atoms)
        bonds = source.project(bond_input, "token/bonds", self.config.hidden_size, low_bit=False)
        molecule_input = F.silu(bonds) + 0.5 * token_hidden + 0.25 * atoms
        molecules = source.project(molecule_input, "token/molecules", self.config.hidden_size, low_bit=False)
        gate = source.gate(torch.cat([token_hidden, atoms], dim=-1), "token/compose_gate", self.config.hidden_size)
        return token_hidden + 0.15 * atoms + 0.10 * gate * bonds + 0.10 * (1.0 - gate) * molecules

    def _text(
        self,
        input_ids: torch.Tensor | None,
        inputs_embeds: torch.Tensor | None,
        attention_mask: torch.Tensor | None,
        *,
        position_start: int,
        prefix_mask: torch.Tensor | None,
    ) -> tuple[torch.Tensor, dict[str, torch.Tensor], tuple[int, ...]] | None:
        if input_ids is None and inputs_embeds is None:
            return None
        if input_ids is not None and inputs_embeds is not None:
            raise ValueError("Pass input_ids or inputs_embeds, not both")
        if inputs_embeds is None:
            assert input_ids is not None
            hidden = self.source.embedding(
                input_ids,
                "token",
                self.config.vocab_size,
                self.config.hidden_size,
            )
            hidden = self._token_features(input_ids, hidden)
        else:
            if inputs_embeds.shape[-1] != self.config.hidden_size:
                raise ValueError("inputs_embeds last dimension must equal hidden_size")
            hidden = inputs_embeds
        batch, length = hidden.shape[:2]
        device = hidden.device
        positions = torch.arange(position_start, position_start + length, device=device).expand(batch, -1)
        logical = positions + MODALITY_POSITION_OFFSETS[MODALITY_TEXT]
        coords = torch.zeros(batch, length, 4, device=device, dtype=hidden.dtype)
        # Absolute coordinates must be invariant to chunking.  Normalizing by the
        # current input length made a token receive different spatial features during
        # full-sequence training and cached one-token decoding.
        denominator = max(1, self.config.max_position_embeddings - 1)
        coords[..., 0] = (positions.to(hidden.dtype) / denominator * 2.0 - 1.0).clamp(-1.0, 1.0)
        modality = torch.full((batch, length), MODALITY_TEXT, device=device, dtype=torch.long)
        is_prefix = (
            prefix_mask.to(device=device, dtype=torch.bool)
            if prefix_mask is not None
            else torch.zeros(batch, length, device=device, dtype=torch.bool)
        )
        mask = (
            attention_mask.to(device=device, dtype=torch.bool)
            if attention_mask is not None
            else torch.ones(batch, length, device=device, dtype=torch.bool)
        )
        return hidden, {
            "modality": modality,
            "positions": positions,
            "logical": logical,
            "coords": coords,
            "prefix": is_prefix,
            "mask": mask,
        }, (length,)

    def _image(self, pixel_values: torch.Tensor | None) -> tuple[torch.Tensor, dict[str, torch.Tensor], tuple[int, ...]] | None:
        if pixel_values is None:
            return None
        if pixel_values.ndim != 4:
            raise ValueError("pixel_values must be [batch, channels, height, width]")
        batch, channels, height, width = pixel_values.shape
        if channels != self.config.image_channels:
            raise ValueError(f"Expected {self.config.image_channels} image channels, got {channels}")
        patch = self.config.image_patch_size
        pad_h, pad_w = (-height) % patch, (-width) % patch
        values = F.pad(pixel_values, (0, pad_w, 0, pad_h))
        grid_h, grid_w = values.shape[-2] // patch, values.shape[-1] // patch
        patches = F.unfold(values, kernel_size=patch, stride=patch).transpose(1, 2)
        hidden = self.source.project(patches, "modality/image_patch", self.config.hidden_size)
        length = hidden.shape[1]
        y = torch.arange(grid_h, device=hidden.device).repeat_interleave(grid_w)
        x = torch.arange(grid_w, device=hidden.device).repeat(grid_h)
        coords = torch.zeros(batch, length, 4, device=hidden.device, dtype=hidden.dtype)
        coords[..., 1] = self._normalize_coords(y, grid_h)
        coords[..., 2] = self._normalize_coords(x, grid_w)
        local = torch.arange(length, device=hidden.device).expand(batch, -1)
        return hidden, self._metadata(local, coords, MODALITY_IMAGE, prefix=True), (grid_h, grid_w)

    def _audio(self, audio_values: torch.Tensor | None) -> tuple[torch.Tensor, dict[str, torch.Tensor], tuple[int, ...]] | None:
        if audio_values is None:
            return None
        if audio_values.ndim == 3:
            audio_values = audio_values.mean(dim=1)
        if audio_values.ndim != 2:
            raise ValueError("audio_values must be [batch, samples] or [batch, channels, samples]")
        patch, stride = self.config.audio_patch_size, self.config.audio_patch_stride
        if audio_values.shape[-1] < patch:
            audio_values = F.pad(audio_values, (0, patch - audio_values.shape[-1]))
        remainder = (audio_values.shape[-1] - patch) % stride
        if remainder:
            audio_values = F.pad(audio_values, (0, stride - remainder))
        windows = audio_values.unfold(-1, patch, stride)
        hidden = self.source.project(windows, "modality/audio_patch", self.config.hidden_size)
        length = hidden.shape[1]
        local = torch.arange(length, device=hidden.device).expand(hidden.shape[0], -1)
        coords = torch.zeros(hidden.shape[0], length, 4, device=hidden.device, dtype=hidden.dtype)
        coords[..., 0] = self._normalize_coords(torch.arange(length, device=hidden.device), length)
        # Frequency-energy coordinate gives raw wave patches a useful second axis.
        spectrum = torch.fft.rfft(windows.float(), dim=-1).abs().mean(dim=-1)
        spectrum = spectrum / spectrum.amax(dim=-1, keepdim=True).clamp_min(1e-8)
        coords[..., 3] = spectrum.to(hidden.dtype) * 2.0 - 1.0
        return hidden, self._metadata(local, coords, MODALITY_AUDIO, prefix=True), (length, patch)

    def _video(self, video_values: torch.Tensor | None) -> tuple[torch.Tensor, dict[str, torch.Tensor], tuple[int, ...]] | None:
        if video_values is None:
            return None
        if video_values.ndim != 5:
            raise ValueError("video_values must be [batch, frames, channels, height, width]")
        batch, frames, channels, height, width = video_values.shape
        if channels != self.config.image_channels:
            raise ValueError(f"Expected {self.config.image_channels} video channels, got {channels}")
        tube, patch = self.config.video_tubelet_size, self.config.video_patch_size
        pad_t, pad_h, pad_w = (-frames) % tube, (-height) % patch, (-width) % patch
        # F.pad follows reverse dimension order for [B,T,C,H,W].
        values = F.pad(video_values, (0, pad_w, 0, pad_h, 0, 0, 0, pad_t))
        tg, hg, wg = values.shape[1] // tube, values.shape[3] // patch, values.shape[4] // patch
        blocks = values.reshape(batch, tg, tube, channels, hg, patch, wg, patch)
        blocks = blocks.permute(0, 1, 4, 6, 2, 3, 5, 7).reshape(batch, tg * hg * wg, -1)
        hidden = self.source.project(blocks, "modality/video_tubelet", self.config.hidden_size)
        t = torch.arange(tg, device=hidden.device).repeat_interleave(hg * wg)
        y = torch.arange(hg, device=hidden.device).repeat_interleave(wg).repeat(tg)
        x = torch.arange(wg, device=hidden.device).repeat(hg * tg)
        coords = torch.zeros(batch, hidden.shape[1], 4, device=hidden.device, dtype=hidden.dtype)
        coords[..., 0] = self._normalize_coords(t, tg)
        coords[..., 1] = self._normalize_coords(y, hg)
        coords[..., 2] = self._normalize_coords(x, wg)
        local = torch.arange(hidden.shape[1], device=hidden.device).expand(batch, -1)
        return hidden, self._metadata(local, coords, MODALITY_VIDEO, prefix=True), (tg, hg, wg)

    def _sensor(self, sensor_values: torch.Tensor | None) -> tuple[torch.Tensor, dict[str, torch.Tensor], tuple[int, ...]] | None:
        if sensor_values is None:
            return None
        if sensor_values.ndim == 2:
            sensor_values = sensor_values.unsqueeze(1)
        if sensor_values.ndim != 3:
            raise ValueError("sensor_values must be [batch, steps, features] or [batch, features]")
        target = self.config.sensor_feature_size
        if sensor_values.shape[-1] < target:
            sensor_values = F.pad(sensor_values, (0, target - sensor_values.shape[-1]))
        elif sensor_values.shape[-1] > target:
            sensor_values = sensor_values[..., :target]
        hidden = self.source.project(sensor_values, "modality/sensor", self.config.hidden_size)
        length = hidden.shape[1]
        local = torch.arange(length, device=hidden.device).expand(hidden.shape[0], -1)
        coords = torch.zeros(hidden.shape[0], length, 4, device=hidden.device, dtype=hidden.dtype)
        coords[..., 0] = self._normalize_coords(torch.arange(length, device=hidden.device), length)
        coords[..., 3] = sensor_values.float().std(dim=-1).to(hidden.dtype).clamp(max=1.0) * 2.0 - 1.0
        return hidden, self._metadata(local, coords, MODALITY_SENSOR, prefix=True), (length, target)

    @staticmethod
    def _metadata(
        local: torch.Tensor,
        coords: torch.Tensor,
        modality_id: int,
        *,
        prefix: bool,
    ) -> dict[str, torch.Tensor]:
        batch, length = local.shape
        device = local.device
        return {
            "modality": torch.full((batch, length), modality_id, device=device, dtype=torch.long),
            "positions": local,
            "logical": local + MODALITY_POSITION_OFFSETS[modality_id],
            "coords": coords,
            "prefix": torch.full((batch, length), prefix, device=device, dtype=torch.bool),
            "mask": torch.ones(batch, length, device=device, dtype=torch.bool),
        }

    def forward(
        self,
        *,
        input_ids: torch.Tensor | None = None,
        inputs_embeds: torch.Tensor | None = None,
        attention_mask: torch.Tensor | None = None,
        pixel_values: torch.Tensor | None = None,
        audio_values: torch.Tensor | None = None,
        video_values: torch.Tensor | None = None,
        sensor_values: torch.Tensor | None = None,
        prefix_mask: torch.Tensor | None = None,
        labels: torch.Tensor | None = None,
        position_start: int = 0,
    ) -> DendroPackedInput:
        self._batch_size(input_ids, inputs_embeds, pixel_values, audio_values, video_values, sensor_values)
        # Non-text modalities form a bidirectional perceptual prefix. Text remains
        # last so causal decoding can append tokens without repacking old inputs.
        parts = [
            ("image", MODALITY_IMAGE, self._image(pixel_values)),
            ("video", MODALITY_VIDEO, self._video(video_values)),
            ("audio", MODALITY_AUDIO, self._audio(audio_values)),
            ("sensor", MODALITY_SENSOR, self._sensor(sensor_values)),
            (
                "text",
                MODALITY_TEXT,
                self._text(
                    input_ids,
                    inputs_embeds,
                    attention_mask,
                    position_start=position_start,
                    prefix_mask=prefix_mask,
                ),
            ),
        ]
        hidden_parts: list[torch.Tensor] = []
        metadata: dict[str, list[torch.Tensor]] = {
            "modality": [],
            "positions": [],
            "logical": [],
            "coords": [],
            "prefix": [],
            "mask": [],
        }
        segments: list[ModalitySegment] = []
        cursor = 0
        text_start, text_length = 0, 0
        for name, modality_id, result in parts:
            if result is None:
                continue
            hidden, info, original_shape = result
            length = int(hidden.shape[1])
            # Physical sequence positions are contiguous across the packed sequence.
            physical = torch.arange(cursor + position_start, cursor + position_start + length, device=hidden.device)
            info["positions"] = physical.expand(hidden.shape[0], -1)
            # Logical modality offsets are applied to the same absolute physical
            # positions so a token keeps identical coordinates when a multimodal
            # prefix is processed in one call or reused through KV cache.
            info["logical"] = info["positions"] + MODALITY_POSITION_OFFSETS[modality_id]
            if modality_id == MODALITY_TEXT:
                denominator = max(1, self.config.max_position_embeddings - 1)
                info["coords"][..., 0] = (
                    info["positions"].to(hidden.dtype) / denominator * 2.0 - 1.0
                ).clamp(-1.0, 1.0)
            hidden_parts.append(hidden)
            for key in metadata:
                metadata[key].append(info[key])
            segments.append(ModalitySegment(name, modality_id, cursor, cursor + length, original_shape))
            if modality_id == MODALITY_TEXT:
                text_start, text_length = cursor, length
            cursor += length
        if not hidden_parts:
            raise RuntimeError("No modality generated tokens")

        hidden = torch.cat(hidden_parts, dim=1)
        combined = {key: torch.cat(values, dim=1) for key, values in metadata.items()}
        hidden = self.spatial_encoder(
            hidden,
            modality_ids=combined["modality"],
            sequence_positions=combined["positions"],
            logical_positions=combined["logical"],
            coordinates=combined["coords"],
            is_prefix=combined["prefix"],
        )
        layout = DendroModalityLayout(
            modality_ids=combined["modality"],
            sequence_positions=combined["positions"],
            logical_positions=combined["logical"],
            coordinates=combined["coords"],
            is_prefix=combined["prefix"],
            attention_mask=combined["mask"],
            segments=tuple(segments),
            text_start=text_start,
            text_length=text_length,
        )

        aligned_labels = labels
        if labels is not None:
            if labels.shape[0] != hidden.shape[0]:
                raise ValueError("labels batch size does not match inputs")
            if labels.shape[1] == text_length and text_start > 0:
                prefix_labels = torch.full(
                    (labels.shape[0], text_start),
                    -100,
                    device=labels.device,
                    dtype=labels.dtype,
                )
                aligned_labels = torch.cat([prefix_labels, labels], dim=1)
            elif labels.shape[1] != hidden.shape[1]:
                raise ValueError(
                    f"labels length {labels.shape[1]} must equal text length {text_length} "
                    f"or packed length {hidden.shape[1]}"
                )
        return DendroPackedInput(hidden_states=hidden, layout=layout, aligned_labels=aligned_labels)