import math

import torch
import torch.nn as nn
import torch.nn.functional as F


class GeometrySpatialMemory(nn.Module):
    """Encode a rendered static 3D scene into compact conditioning tokens.

    The input must be VAE latents of a geometry-rendered video produced from a
    persistent 3D representation (for example a TSDF-fused point cloud). This
    module intentionally does not infer geometry from ordinary context tokens.
    """

    def __init__(
        self,
        dim: int,
        latent_channels: int = 16,
        patch_size=(1, 2, 2),
        grid_size: int = 8,
        temporal_bins: int = 4,
        num_tokens: int = 64,
    ):
        super().__init__()
        self.dim = int(dim)
        self.latent_channels = int(latent_channels)
        self.patch_size = tuple(int(x) for x in patch_size)
        self.grid_size = int(grid_size)
        self.temporal_bins = int(temporal_bins)
        self.num_tokens = int(num_tokens)
        self.geometry_patch_embedding = nn.Conv3d(
            self.latent_channels,
            self.dim,
            kernel_size=self.patch_size,
            stride=self.patch_size,
        )
        source_tokens = self.temporal_bins * self.grid_size * self.grid_size
        self.register_buffer(
            "geometry_config",
            torch.tensor(
                [1, self.grid_size, self.temporal_bins],
                dtype=torch.int64,
            ),
            persistent=True,
        )
        self.geometry_to_tokens = nn.Parameter(torch.empty(source_tokens, self.num_tokens))
        self.norm = nn.LayerNorm(self.dim)
        nn.init.normal_(self.geometry_to_tokens, std=0.02)

    def initialize_from_dit_patch_embedding(self, patch_embedding: nn.Conv3d) -> None:
        """Initialize the geometry encoder from the pretrained video tokenizer."""
        if not isinstance(patch_embedding, nn.Conv3d):
            return
        if patch_embedding.weight.shape != self.geometry_patch_embedding.weight.shape:
            return
        with torch.no_grad():
            self.geometry_patch_embedding.weight.copy_(patch_embedding.weight)
            if (
                self.geometry_patch_embedding.bias is not None
                and patch_embedding.bias is not None
                and self.geometry_patch_embedding.bias.shape == patch_embedding.bias.shape
            ):
                self.geometry_patch_embedding.bias.copy_(patch_embedding.bias)

    def forward(self, geometry_latents: torch.Tensor) -> torch.Tensor:
        if geometry_latents is None or geometry_latents.ndim != 5:
            raise ValueError(
                "GeometrySpatialMemory expects geometry VAE latents shaped (B, C, F, H, W)."
            )
        if int(geometry_latents.shape[1]) != self.latent_channels:
            raise ValueError(
                f"GeometrySpatialMemory channel mismatch: input={geometry_latents.shape[1]} "
                f"expected={self.latent_channels}"
            )
        features = self.geometry_patch_embedding(geometry_latents)
        features = F.adaptive_avg_pool3d(
            features,
            (self.temporal_bins, self.grid_size, self.grid_size),
        )
        features = features.flatten(2).transpose(1, 2)
        mix = torch.softmax(self.geometry_to_tokens, dim=0)
        memory = torch.einsum("bsd,sm->bmd", features, mix)
        return self.norm(memory / math.sqrt(max(self.temporal_bins, 1)))