from __future__ import annotations

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


class LongMemoryState(nn.Module):
    """Chunk-level long-memory reader/writer driven by clean observation latents."""

    def __init__(
        self,
        *,
        latent_channels: int,
        memory_dim: int,
        num_tokens: int,
        num_heads: int,
        video_hidden_size: int,
        action_hidden_size: int,
        memory_noise_std: float = 0.0,
        memory_dropout_prob: float = 0.0,
    ) -> None:
        super().__init__()
        self.memory_dim = int(memory_dim)
        self.num_tokens = int(num_tokens)
        self.memory_noise_std = float(memory_noise_std)
        self.memory_dropout_prob = float(memory_dropout_prob)
        self.init_memory = nn.Parameter(torch.zeros(self.num_tokens, self.memory_dim))
        nn.init.trunc_normal_(self.init_memory, std=0.02)

        self.obs_proj = nn.Linear(int(latent_channels), self.memory_dim)
        self.obs_norm = nn.LayerNorm(self.memory_dim)
        self.mem_norm = nn.LayerNorm(self.memory_dim)
        self.write_attn = nn.MultiheadAttention(
            embed_dim=self.memory_dim,
            num_heads=int(num_heads),
            batch_first=True,
        )
        self.write_ffn = nn.Sequential(
            nn.LayerNorm(self.memory_dim),
            nn.Linear(self.memory_dim, self.memory_dim * 4),
            nn.GELU(),
            nn.Linear(self.memory_dim * 4, self.memory_dim),
        )
        self.write_gate = nn.Linear(self.memory_dim * 2, self.memory_dim)

        self.read_norm = nn.LayerNorm(self.memory_dim)
        self.to_video_bias = nn.Linear(self.memory_dim, int(video_hidden_size), bias=False)
        self.to_action_bias = nn.Linear(self.memory_dim, int(action_hidden_size), bias=False)
        nn.init.zeros_(self.to_video_bias.weight)
        nn.init.zeros_(self.to_action_bias.weight)

    def init_state(self, batch_size: int, *, device: torch.device, dtype: torch.dtype) -> torch.Tensor:
        return self.init_memory.unsqueeze(0).expand(int(batch_size), -1, -1).to(device=device, dtype=dtype)

    def encode_observation_tokens(self, latents: torch.Tensor) -> torch.Tensor:
        if latents.ndim != 5:
            raise ValueError(f"latents must be [B, C, T, H, W], got {tuple(latents.shape)}")
        tokens = latents.permute(0, 2, 3, 4, 1).reshape(int(latents.shape[0]), -1, int(latents.shape[1]))
        tokens = self.obs_proj(tokens)
        return self.obs_norm(tokens)

    def update(self, memory: torch.Tensor, observation_latents: torch.Tensor) -> torch.Tensor:
        obs_tokens = self.encode_observation_tokens(observation_latents)
        norm_memory = self.mem_norm(memory)
        attn_update, _ = self.write_attn(norm_memory, obs_tokens, obs_tokens, need_weights=False)
        proposal = self.write_ffn(attn_update)
        gate = torch.sigmoid(self.write_gate(torch.cat([memory, proposal], dim=-1)))
        updated = memory + gate * proposal
        if self.training and self.memory_noise_std > 0.0:
            updated = updated + torch.randn_like(updated) * float(self.memory_noise_std)
        if self.training and self.memory_dropout_prob > 0.0:
            updated = F.dropout(updated, p=float(self.memory_dropout_prob), training=True)
        return updated

    def summarize(self, memory: torch.Tensor) -> dict[str, torch.Tensor]:
        summary = self.read_norm(memory).mean(dim=1)
        return {
            "summary": summary,
            "video_bias": self.to_video_bias(summary),
            "action_bias": self.to_action_bias(summary),
        }