from __future__ import annotations

import logging
from dataclasses import dataclass
from pathlib import Path
from typing import Optional, Tuple, Dict, Any

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

logger = logging.getLogger(__name__)


# -----------------------------
# Vision modules
# -----------------------------

class PatchEmbedding(nn.Module):
    def __init__(self, img_size: int = 112, patch_size: int = 16, in_channels: int = 3, embed_dim: int = 768):
        super().__init__()
        self.img_size = img_size
        self.patch_size = patch_size
        self.num_patches = (img_size // patch_size) ** 2
        self.projection = nn.Conv2d(in_channels, embed_dim, kernel_size=patch_size, stride=patch_size)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.projection(x)
        x = x.flatten(2).transpose(1, 2)
        return x


class MultiHeadAttention(nn.Module):
    def __init__(self, embed_dim: int, num_heads: int):
        super().__init__()
        assert embed_dim % num_heads == 0
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads
        self.scale = self.head_dim ** -0.5
        self.q_proj = nn.Linear(embed_dim, embed_dim)
        self.k_proj = nn.Linear(embed_dim, embed_dim)
        self.v_proj = nn.Linear(embed_dim, embed_dim)
        self.out_proj = nn.Linear(embed_dim, embed_dim)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, N, C = x.shape
        q = self.q_proj(x).reshape(B, N, self.num_heads, self.head_dim).transpose(1, 2)
        k = self.k_proj(x).reshape(B, N, self.num_heads, self.head_dim).transpose(1, 2)
        v = self.v_proj(x).reshape(B, N, self.num_heads, self.head_dim).transpose(1, 2)
        attn = (q @ k.transpose(-2, -1)) * self.scale
        attn = attn.softmax(dim=-1)
        out = attn @ v
        out = out.transpose(1, 2).reshape(B, N, C)
        return self.out_proj(out)


class MLP(nn.Module):
    def __init__(self, embed_dim: int, mlp_ratio: float = 4.0):
        super().__init__()
        hidden = int(embed_dim * mlp_ratio)
        self.fc1 = nn.Linear(embed_dim, hidden)
        self.fc2 = nn.Linear(hidden, embed_dim)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = F.gelu(self.fc1(x))
        x = self.fc2(x)
        return x


class TransformerBlock(nn.Module):
    def __init__(self, embed_dim: int, num_heads: int, mlp_ratio: float = 4.0):
        super().__init__()
        self.norm1 = nn.LayerNorm(embed_dim)
        self.attn = MultiHeadAttention(embed_dim, num_heads)
        self.norm2 = nn.LayerNorm(embed_dim)
        self.mlp = MLP(embed_dim, mlp_ratio)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = x + self.attn(self.norm1(x))
        x = x + self.mlp(self.norm2(x))
        return x


class LisaViT(nn.Module):
    """Vision Transformer. Returns (cls_token, patch_tokens) so downstream modules
    receive the correct spatial tokens."""

    def __init__(self, img_size: int = 112, patch_size: int = 16, embed_dim: int = 768, depth: int = 12, num_heads: int = 12):
        super().__init__()
        self.patch_embed = PatchEmbedding(img_size, patch_size, 3, embed_dim)
        self.num_patches = self.patch_embed.num_patches
        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
        self.pos_embed = nn.Parameter(torch.zeros(1, self.num_patches + 1, embed_dim))
        self.blocks = nn.ModuleList([TransformerBlock(embed_dim, num_heads) for _ in range(depth)])
        self.norm = nn.LayerNorm(embed_dim)

    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        """Returns (cls_token [B, D], patch_tokens [B, N, D])."""
        B = x.shape[0]
        x = self.patch_embed(x)
        x = torch.cat([self.cls_token.expand(B, -1, -1), x], dim=1)
        x = x + self.pos_embed
        for blk in self.blocks:
            x = blk(x)
        x = self.norm(x)
        # FIX: return cls token AND patch tokens separately
        return x[:, 0], x[:, 1:]  # (cls_token, patch_tokens)


class LisaObjectDetector(nn.Module):
    def __init__(self, embed_dim: int = 768, num_classes: int = 80):
        super().__init__()
        self.feature_conv = nn.Sequential(
            nn.Conv2d(embed_dim, 512, 3, padding=1),
            nn.BatchNorm2d(512),
            nn.ReLU(inplace=True),
            nn.Conv2d(512, 256, 3, padding=1),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True),
        )
        self.cls_head = nn.Conv2d(256, 3 * num_classes, 1)
        self.reg_head = nn.Conv2d(256, 3 * 4, 1)
        self.obj_head = nn.Conv2d(256, 3, 1)

    def forward(self, patch_tokens: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        B, N, C = patch_tokens.shape
        S = int(N ** 0.5)
        x = patch_tokens.transpose(1, 2).reshape(B, C, S, S)
        x = self.feature_conv(x)
        return self.cls_head(x), self.reg_head(x), self.obj_head(x)


# -----------------------------
# Audio analysis modules (perception only — NOT text output)
# -----------------------------

class PositionalEncoding(nn.Module):
    def __init__(self, embed_dim: int, max_len: int = 5000):
        super().__init__()
        pe = torch.zeros(max_len, embed_dim)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, embed_dim, 2).float() * (-torch.log(torch.tensor(10000.0)) / embed_dim))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        self.register_buffer('pe', pe.unsqueeze(0))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return x + self.pe[:, : x.size(1)]


class AudioEncoder(nn.Module):
    def __init__(self, input_dim: int = 80, embed_dim: int = 512, num_layers: int = 6, num_heads: int = 8):
        super().__init__()
        self.input_proj = nn.Linear(input_dim, embed_dim)
        self.pos_encoding = PositionalEncoding(embed_dim)
        layer = nn.TransformerEncoderLayer(d_model=embed_dim, nhead=num_heads, dim_feedforward=embed_dim * 4, batch_first=True)
        self.transformer = nn.TransformerEncoder(layer, num_layers=num_layers)
        self.output_norm = nn.LayerNorm(embed_dim)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.input_proj(x)
        x = self.pos_encoding(x)
        x = self.transformer(x)
        return self.output_norm(x)


class LisaSoundClassifier(nn.Module):
    def __init__(self, num_classes: int = 50, embed_dim: int = 512):
        super().__init__()
        self.encoder = AudioEncoder(80, embed_dim, num_layers=4)
        self.classifier = nn.Sequential(
            nn.AdaptiveAvgPool1d(1),
            nn.Flatten(),
            nn.Linear(embed_dim, embed_dim // 2),
            nn.ReLU(),
            nn.Dropout(0.1),
            nn.Linear(embed_dim // 2, num_classes),
        )

    def forward(self, mel_features: torch.Tensor) -> torch.Tensor:
        x = self.encoder(mel_features)
        x = x.transpose(1, 2)
        return self.classifier(x)


class LisaEmotionDetector(nn.Module):
    def __init__(self, embed_dim: int = 512):
        super().__init__()
        self.encoder = AudioEncoder(80, embed_dim, num_layers=4)
        self.emotion_head = nn.Sequential(
            nn.AdaptiveAvgPool1d(1),
            nn.Flatten(),
            nn.Linear(embed_dim, embed_dim // 2),
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(embed_dim // 2, 7),
        )

    def forward(self, mel_features: torch.Tensor) -> torch.Tensor:
        x = self.encoder(mel_features)
        x = x.transpose(1, 2)
        return self.emotion_head(x)


class LisaVAD(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv_layers = nn.Sequential(
            nn.Conv1d(80, 128, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.Conv1d(128, 128, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.Conv1d(128, 64, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.AdaptiveAvgPool1d(1),
            nn.Flatten(),
            nn.Linear(64, 32),
            nn.ReLU(),
            nn.Linear(32, 1),
            nn.Sigmoid(),
        )

    def forward(self, mel_features: torch.Tensor) -> torch.Tensor:
        # mel_features: [B, T, 80]
        return self.conv_layers(mel_features.transpose(1, 2))


# -----------------------------
# Speech-like audio encoder (kept for checkpoint compatibility)
# Maps input mel → encoded features.  NO text/CTC output.
# -----------------------------

class LisaSpeechRecognizer(nn.Module):
    """Audio encoder retained for checkpoint weight compatibility.
    The ASR/CTC heads are removed. Only the encoder is used as a
    perceptual audio feature extractor for fusion."""

    def __init__(self, vocab_size: int = 32, embed_dim: int = 512, num_layers: int = 6):
        super().__init__()
        # Keep encoder — trained weights are valuable
        self.encoder = AudioEncoder(80, embed_dim, num_layers)
        # Keep projection heads in skeleton only so saved weights load cleanly;
        # they are not used during inference.
        self.ctc_projection = nn.Linear(embed_dim, vocab_size + 1)
        self.lm_head = nn.Sequential(
            nn.Linear(embed_dim, embed_dim * 2),
            nn.ReLU(),
            nn.Dropout(0.1),
            nn.Linear(embed_dim * 2, vocab_size),
        )

    def forward(self, mel_features: torch.Tensor) -> torch.Tensor:
        """Returns encoded features [B, T, embed_dim]. No CTC/LM output."""
        return self.encoder(mel_features)


# -----------------------------
# MelDecoder: vision → mel spectrogram (TTS-style)
# This is the core "speak from what you see" module.
# Architecture: cross-attention decoder conditioned on vision CLS token.
# Input:  vision_context [B, D_vision] (from ViT CLS token)
# Output: mel_spectrogram [B, T_out, n_mels=80]
# -----------------------------

class MelDecoderLayer(nn.Module):
    """Single decoder layer: self-attn on mel queries, cross-attn on vision context."""

    def __init__(self, mel_dim: int, context_dim: int, num_heads: int = 8):
        super().__init__()
        self.norm1 = nn.LayerNorm(mel_dim)
        self.self_attn = nn.MultiheadAttention(mel_dim, num_heads, batch_first=True)

        self.norm2 = nn.LayerNorm(mel_dim)
        # Cross-attention: queries from mel, keys/values from vision context
        self.cross_attn = nn.MultiheadAttention(
            mel_dim, num_heads,
            kdim=context_dim, vdim=context_dim,
            batch_first=True,
        )

        self.norm3 = nn.LayerNorm(mel_dim)
        self.ffn = nn.Sequential(
            nn.Linear(mel_dim, mel_dim * 4),
            nn.GELU(),
            nn.Linear(mel_dim * 4, mel_dim),
        )

    def forward(
        self,
        mel_query: torch.Tensor,       # [B, T_out, mel_dim]
        vision_context: torch.Tensor,  # [B, N_ctx, context_dim]
        causal_mask: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        # Self-attention (causal for autoregressive mel generation)
        x = mel_query
        sa_out, _ = self.self_attn(
            self.norm1(x), self.norm1(x), self.norm1(x),
            attn_mask=causal_mask,
        )
        x = x + sa_out

        # Cross-attention: attend to vision context
        ca_out, _ = self.cross_attn(self.norm2(x), vision_context, vision_context)
        x = x + ca_out

        # Feed-forward
        x = x + self.ffn(self.norm3(x))
        return x


class MelDecoder(nn.Module):
    """Vision-conditioned mel spectrogram decoder.
    Converts vision features into speech-like mel output without any text.

    Pipeline:
        vision CLS [B, D_v] → expand to [B, 1, D_v] (context)
        learned start frame → autoregressive cross-attention → mel [B, T_out, 80]
        Linear → raw mel values (no softmax — continuous output)
    """

    def __init__(
        self,
        vision_dim: int = 768,
        mel_dim: int = 512,
        n_mels: int = 80,
        num_layers: int = 4,
        num_heads: int = 8,
        max_mel_len: int = 200,
    ):
        super().__init__()
        self.mel_dim = mel_dim
        self.n_mels = n_mels
        self.max_mel_len = max_mel_len

        # Project vision CLS to mel_dim space (context for cross-attn)
        self.vision_proj = nn.Linear(vision_dim, mel_dim)

        # Learned start-of-speech frame embedding
        self.start_frame = nn.Parameter(torch.zeros(1, 1, mel_dim))

        # Positional encoding for mel sequence
        self.pos_encoding = PositionalEncoding(mel_dim, max_len=max_mel_len + 1)

        # Input projection: mel bins → mel_dim
        self.mel_input_proj = nn.Linear(n_mels, mel_dim)

        # Decoder layers
        self.layers = nn.ModuleList([
            MelDecoderLayer(mel_dim, mel_dim, num_heads)
            for _ in range(num_layers)
        ])

        self.norm = nn.LayerNorm(mel_dim)

        # Output projection: mel_dim → n_mels (continuous mel values)
        self.mel_output_proj = nn.Linear(mel_dim, n_mels)

        # Stop token predictor (sigmoid: 0=continue, 1=stop)
        self.stop_proj = nn.Linear(mel_dim, 1)

    def _causal_mask(self, T: int, device: torch.device) -> torch.Tensor:
        """Upper-triangular causal mask for autoregressive decoding."""
        mask = torch.triu(torch.ones(T, T, device=device), diagonal=1).bool()
        return mask

    def forward(self, vision_cls: torch.Tensor, target_mel: Optional[torch.Tensor] = None) -> Dict[str, torch.Tensor]:
        """
        Args:
            vision_cls: [B, D_vision] — CLS token from ViT
            target_mel:  [B, T, n_mels] — ground truth mel for teacher-forced training (None at inference)

        Returns dict:
            'mel':  [B, T_out, n_mels]   — predicted mel spectrogram
            'stop': [B, T_out, 1]        — stop token logits
        """
        B = vision_cls.shape[0]
        device = vision_cls.device

        # Project vision to decoder space → context [B, 1, mel_dim]
        ctx = self.vision_proj(vision_cls).unsqueeze(1)  # [B, 1, mel_dim]

        if target_mel is not None:
            # --- Teacher-forced training path ---
            T = target_mel.shape[1]
            # Project mel input and prepend start frame
            mel_emb = self.mel_input_proj(target_mel)                # [B, T, mel_dim]
            sos = self.start_frame.expand(B, -1, -1)                 # [B, 1, mel_dim]
            queries = torch.cat([sos, mel_emb[:, :-1, :]], dim=1)   # [B, T, mel_dim] (shifted right)
        else:
            # --- Inference: start from SOS ---
            T = self.max_mel_len
            queries = self.start_frame.expand(B, -1, -1)            # [B, 1, mel_dim]

        queries = self.pos_encoding(queries)
        causal_mask = self._causal_mask(queries.shape[1], device)

        x = queries
        for layer in self.layers:
            x = layer(x, ctx, causal_mask)
        x = self.norm(x)

        mel_out = self.mel_output_proj(x)   # [B, T_out, n_mels]
        stop_out = self.stop_proj(x)        # [B, T_out, 1]

        return {'mel': mel_out, 'stop': stop_out}


# -----------------------------
# FusionAttention: cross-modal attention (vision queries, audio keys/values)
# Matches checkpoint keys: model.q_proj / k_proj / v_proj / out_proj
# -----------------------------

class FusionAttention(nn.Module):
    """True cross-attention fusion: vision queries attend to audio (optional) context.

    Checkpoint compatibility: Q/K/V/out projections live at 'model.*' prefix,
    so the state-dict keys match the saved weights.

    Also contains trainable projectors to align vision (768) and audio (512)
    into a common fusion_dim (1024) space — these replace the throwaway
    per-call Linear layers that were previously used in api.py.
    """

    def __init__(self, hidden_dim: int = 1024, vision_dim: int = 768, audio_dim: int = 512, num_heads: int = 8):
        super().__init__()
        assert hidden_dim % num_heads == 0
        self.hidden_dim = hidden_dim
        self.num_heads = num_heads
        self.head_dim = hidden_dim // num_heads
        self.scale = self.head_dim ** -0.5

        # Trained dimension alignment projectors (not throwaway)
        self.vision_proj = nn.Linear(vision_dim, hidden_dim)
        self.audio_proj = nn.Linear(audio_dim, hidden_dim)

        # Checkpoint-compatible attention weights (keys: model.q_proj, etc.)
        self.q_proj = nn.Linear(hidden_dim, hidden_dim)
        self.k_proj = nn.Linear(hidden_dim, hidden_dim)
        self.v_proj = nn.Linear(hidden_dim, hidden_dim)
        self.out_proj = nn.Linear(hidden_dim, hidden_dim)

        self.norm_q = nn.LayerNorm(hidden_dim)
        self.norm_k = nn.LayerNorm(hidden_dim)

    def forward(
        self,
        vision_feat: torch.Tensor,                    # [B, D_v] or [B, N, D_v]
        audio_feat: Optional[torch.Tensor] = None,   # [B, D_a] or [B, T, D_a] — None → self-fusion
    ) -> torch.Tensor:
        """Returns fused features [B, hidden_dim]."""
        # Ensure sequence dim
        if vision_feat.dim() == 2:
            vision_feat = vision_feat.unsqueeze(1)   # [B, 1, D_v]
        if audio_feat is not None and audio_feat.dim() == 2:
            audio_feat = audio_feat.unsqueeze(1)     # [B, 1, D_a]

        # Project to common space
        Q = self.vision_proj(vision_feat)            # [B, N_q, hidden]
        if audio_feat is not None:
            K = self.audio_proj(audio_feat)          # [B, N_k, hidden]
            V = K
        else:
            # Self-fusion when no audio available
            K = Q
            V = Q

        Q = self.norm_q(self.q_proj(Q))
        K = self.norm_k(self.k_proj(K))
        V = self.v_proj(V)

        B, N_q, _ = Q.shape
        B, N_k, _ = K.shape

        # Multi-head cross attention
        Q = Q.reshape(B, N_q, self.num_heads, self.head_dim).transpose(1, 2)
        K = K.reshape(B, N_k, self.num_heads, self.head_dim).transpose(1, 2)
        V = V.reshape(B, N_k, self.num_heads, self.head_dim).transpose(1, 2)

        attn = (Q @ K.transpose(-2, -1)) * self.scale
        attn = attn.softmax(dim=-1)
        out = (attn @ V).transpose(1, 2).reshape(B, N_q, self.hidden_dim)
        out = self.out_proj(out)

        # Global pool to [B, hidden_dim]
        return out.mean(dim=1)


# -----------------------------
# Runtime assembly
# -----------------------------

@dataclass
class RuntimeConfig:
    vision_embed_dim: int = 768
    audio_embed_dim: int = 512
    fusion_hidden_dim: int = 1024
    img_size: int = 112
    vit_depth: int = 12
    vit_heads: int = 12


@dataclass
class LisaAIRuntime:
    """Vision-to-Audio runtime. Sees the world, speaks about it — no text output."""
    vision_vit: LisaViT
    object_detector: LisaObjectDetector
    speech_recognizer: LisaSpeechRecognizer   # audio *encoder* (perceptual), retained for checkpoint compat
    sound_classifier: LisaSoundClassifier
    emotion_detector: LisaEmotionDetector
    vad_model: LisaVAD
    fusion_model: FusionAttention
    mel_decoder: MelDecoder                   # NEW: vision→mel speech output module

    def __post_init__(self):
        self._load_stats = {}
        self._detected_dims = {}

    # ---- mode helpers ----

    def _all_modules(self):
        return [
            self.vision_vit, self.object_detector, self.speech_recognizer,
            self.sound_classifier, self.emotion_detector, self.vad_model,
            self.fusion_model, self.mel_decoder,
        ]

    def set_eval_mode(self):
        for m in self._all_modules():
            m.eval()

    def set_train_mode(self):
        for m in self._all_modules():
            m.train()

    def to(self, device: torch.device):
        self.vision_vit = self.vision_vit.to(device)
        self.object_detector = self.object_detector.to(device)
        self.speech_recognizer = self.speech_recognizer.to(device)
        self.sound_classifier = self.sound_classifier.to(device)
        self.emotion_detector = self.emotion_detector.to(device)
        self.vad_model = self.vad_model.to(device)
        self.fusion_model = self.fusion_model.to(device)
        self.mel_decoder = self.mel_decoder.to(device)
        return self

    def get_model_info(self) -> Dict[str, Any]:
        info = {
            'architecture': 'LISA Vision→Audio Runtime (no text)',
            'detected_dimensions': getattr(self, '_detected_dims', {}),
            'load_statistics': getattr(self, '_load_stats', {}),
            'components': {},
        }
        components = {
            'vision_vit': self.vision_vit,
            'object_detector': self.object_detector,
            'audio_encoder (speech_recognizer)': self.speech_recognizer,
            'sound_classifier': self.sound_classifier,
            'emotion_detector': self.emotion_detector,
            'vad_model': self.vad_model,
            'fusion_model': self.fusion_model,
            'mel_decoder': self.mel_decoder,
        }
        total_params = 0
        for name, module in components.items():
            params = sum(p.numel() for p in module.parameters())
            info['components'][name] = {
                'total_params': params,
                'trainable_params': sum(p.numel() for p in module.parameters() if p.requires_grad),
                'device': str(next(module.parameters()).device),
                'dtype': str(next(module.parameters()).dtype),
            }
            total_params += params
        info['total_parameters'] = total_params
        return info

    # ---- perception ----

    def get_vision_features(self, image: torch.Tensor) -> torch.Tensor:
        """Returns CLS token [B, D_v]."""
        with torch.no_grad():
            try:
                cls_token, _ = self.vision_vit(image)
                return cls_token
            except Exception as e:
                logger.warning(f"Vision feature extraction error: {e}")
                B = image.shape[0]
                return torch.zeros(B, self._detected_dims.get('vision', 768), device=image.device)

    def get_vision_patch_tokens(self, image: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        """Returns (cls_token [B,D], patch_tokens [B,N,D])."""
        with torch.no_grad():
            try:
                return self.vision_vit(image)
            except Exception as e:
                logger.warning(f"Vision patch extraction error: {e}")
                B = image.shape[0]
                D = self._detected_dims.get('vision', 768)
                N = self.vision_vit.num_patches
                dev = image.device
                return torch.zeros(B, D, device=dev), torch.zeros(B, N, D, device=dev)

    def get_audio_features(self, mel_features: torch.Tensor, model_type: str = 'speech') -> torch.Tensor:
        """Extract perceptual audio features. model_type: 'speech'|'sound'|'emotion'."""
        with torch.no_grad():
            try:
                if model_type == 'speech':
                    # encoder only — no CTC text output
                    enc = self.speech_recognizer(mel_features)   # [B, T, D]
                    return enc.mean(dim=1)                        # [B, D]
                elif model_type == 'sound':
                    return self.sound_classifier(mel_features)
                elif model_type == 'emotion':
                    return self.emotion_detector(mel_features)
                else:
                    raise ValueError(f"Unknown audio model type: {model_type}")
            except Exception as e:
                logger.warning(f"Audio feature extraction error: {e}")
                B = mel_features.shape[0]
                return torch.zeros(B, self._detected_dims.get('audio', 512), device=mel_features.device)

    def detect_objects(self, image: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """Object detection using real patch tokens from ViT."""
        with torch.no_grad():
            try:
                # FIX: always use patch_tokens, not cls_token
                _, patch_tokens = self.vision_vit(image)
                return self.object_detector(patch_tokens)
            except Exception as e:
                logger.warning(f"Object detection error: {e}")
                B = image.shape[0]
                N = (image.shape[-1] // 16) ** 2
                dev = image.device
                return (
                    torch.zeros(B, N, 80, device=dev),
                    torch.zeros(B, N, 4, device=dev),
                    torch.zeros(B, N, 1, device=dev),
                )

    def detect_emotion(self, mel_features: torch.Tensor) -> torch.Tensor:
        with torch.no_grad():
            return self.emotion_detector(mel_features)

    def classify_sound(self, mel_features: torch.Tensor) -> torch.Tensor:
        with torch.no_grad():
            return self.sound_classifier(mel_features)

    def detect_voice_activity(self, mel_features: torch.Tensor) -> torch.Tensor:
        with torch.no_grad():
            return self.vad_model(mel_features)

    # ---- fusion ----

    def fuse(
        self,
        vision_feat: torch.Tensor,
        audio_feat: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        """Cross-modal fusion. Returns [B, fusion_dim]."""
        return self.fusion_model(vision_feat, audio_feat)

    # Legacy name kept for backward compat
    def multimodal_fusion(self, vision_features: torch.Tensor, audio_features: torch.Tensor) -> torch.Tensor:
        return self.fuse(vision_features, audio_features)

    # ---- speech output (vision → mel → waveform) ----

    def synthesize_speech(
        self,
        image: torch.Tensor,
        audio_context: Optional[torch.Tensor] = None,
    ) -> Dict[str, torch.Tensor]:
        """Core vision-to-speech output pipeline.

        Args:
            image: [B, 3, H, W] — input image(s) / video frame(s)
            audio_context: optional [B, T, 80] mel — for emotional conditioning via fusion

        Returns dict:
            'mel':          [B, T_out, 80]  — predicted mel spectrogram
            'stop':         [B, T_out, 1]   — stop probabilities
            'vision_cls':   [B, D_v]        — vision embedding used
            'fused':        [B, fusion_dim] — fused vision+audio context
        """
        with torch.no_grad():
            # 1. Perceive the scene
            vision_cls, _ = self.vision_vit(image)

            # 2. Optionally fuse with audio context (e.g. emotion from mic)
            if audio_context is not None:
                audio_feat = self.get_audio_features(audio_context, 'speech')
                fused = self.fuse(vision_cls, audio_feat)
            else:
                fused = self.fuse(vision_cls, None)

            # 3. Decode vision into mel spectrogram (speak what you see)
            mel_out = self.mel_decoder(vision_cls)

        return {
            'mel': mel_out['mel'],
            'stop': mel_out['stop'],
            'vision_cls': vision_cls,
            'fused': fused,
        }


# -----------------------------
# Checkpoint utilities
# -----------------------------

def infer_dimensions_from_ckpt(ckpt: Dict[str, torch.Tensor]) -> Tuple[int, int, int]:
    """Infer model dimensions from checkpoint keys."""
    vit_norm = ckpt.get('vision.vit.norm.weight')
    vit_dim = int(vit_norm.numel()) if vit_norm is not None else 768

    audio_dim = 512
    for key in ['audio.speech_recognizer.encoder.input_proj.weight',
                'audio.sound_classifier.encoder.input_proj.weight']:
        if key in ckpt:
            audio_dim = int(ckpt[key].shape[0])
            break

    fusion = ckpt.get('model.out_proj.weight')
    fusion_dim = int(fusion.shape[0]) if (fusion is not None and fusion.ndim == 2) else 1024

    return vit_dim, audio_dim, fusion_dim


def build_runtime_from_checkpoint(model_dir: str, device: Optional[torch.device] = None) -> LisaAIRuntime:
    """Build LisaAIRuntime from a checkpoint directory."""
    device = device or torch.device('cpu')
    ckpt_path = Path(model_dir) / 'model.safetensors'

    if not ckpt_path.exists():
        raise FileNotFoundError(f"Checkpoint not found: {ckpt_path}")

    ckpt = load_file(str(ckpt_path))
    vit_dim, audio_dim, fusion_dim = infer_dimensions_from_ckpt(ckpt)
    logger.info(f"Detected model dimensions: vision={vit_dim}, audio={audio_dim}, fusion={fusion_dim}")

    cfg = RuntimeConfig(
        vision_embed_dim=vit_dim,
        audio_embed_dim=audio_dim,
        fusion_hidden_dim=fusion_dim,
    )

    # Build modules
    vit = LisaViT(
        img_size=cfg.img_size,
        embed_dim=cfg.vision_embed_dim,
        depth=cfg.vit_depth,
        num_heads=max(1, cfg.vision_embed_dim // 64),
    ).to(device)

    det = LisaObjectDetector(embed_dim=cfg.vision_embed_dim).to(device)
    asr = LisaSpeechRecognizer(vocab_size=32, embed_dim=cfg.audio_embed_dim, num_layers=6).to(device)
    snd = LisaSoundClassifier(num_classes=50, embed_dim=cfg.audio_embed_dim).to(device)
    emo = LisaEmotionDetector(embed_dim=cfg.audio_embed_dim).to(device)
    vad = LisaVAD().to(device)
    fus = FusionAttention(
        hidden_dim=cfg.fusion_hidden_dim,
        vision_dim=cfg.vision_embed_dim,
        audio_dim=cfg.audio_embed_dim,
    ).to(device)
    mel_dec = MelDecoder(
        vision_dim=cfg.vision_embed_dim,
        mel_dim=cfg.audio_embed_dim,
        n_mels=80,
        num_layers=4,
    ).to(device)

    # Weight loading with prefix matching
    def load_prefix_safe(module: nn.Module, prefix: str):
        sd = module.state_dict()
        loadable = {}
        skipped = []

        for k, v in ckpt.items():
            if not k.startswith(prefix):
                continue
            subk = k[len(prefix):]

            # Naming transformations
            if 'mlp.0.' in subk:
                subk = subk.replace('mlp.0.', 'mlp.fc1.')
            elif 'mlp.3.' in subk:
                subk = subk.replace('mlp.3.', 'mlp.fc2.')
            if 'attn.qkv.' in subk:
                continue  # fused QKV not supported here

            if subk in sd:
                if sd[subk].shape == v.shape:
                    loadable[subk] = v
                else:
                    skipped.append(f"{k}: shape mismatch {sd[subk].shape} vs {v.shape}")
            else:
                skipped.append(f"{k}: not found")

        missing, unexpected = module.load_state_dict(loadable, strict=False)
        loaded = len(loadable)
        total = len([k for k in ckpt if k.startswith(prefix)])
        logger.info(f"[{prefix}] loaded {loaded}/{total}")
        if skipped:
            logger.warning(f"[{prefix}] skipped {len(skipped)}")
        return loaded, len(skipped)

    load_stats = {}
    load_stats['vision.vit']               = load_prefix_safe(vit,     'vision.vit.')
    load_stats['vision.object_detector']   = load_prefix_safe(det,     'vision.object_detector.')
    load_stats['audio.speech_recognizer']  = load_prefix_safe(asr,     'audio.speech_recognizer.')
    load_stats['audio.sound_classifier']   = load_prefix_safe(snd,     'audio.sound_classifier.')
    load_stats['audio.emotion_detector']   = load_prefix_safe(emo,     'audio.emotion_detector.')
    load_stats['audio.vad_model']          = load_prefix_safe(vad,     'audio.vad_model.')
    load_stats['model']                    = load_prefix_safe(fus,     'model.')
    # mel_decoder has no pretrained weights — it starts fresh and is fine-tuned

    total_loaded  = sum(s[0] for s in load_stats.values())
    total_skipped = sum(s[1] for s in load_stats.values())
    logger.info(f"Model loading complete: {total_loaded} loaded, {total_skipped} skipped")

    runtime = LisaAIRuntime(
        vision_vit=vit,
        object_detector=det,
        speech_recognizer=asr,
        sound_classifier=snd,
        emotion_detector=emo,
        vad_model=vad,
        fusion_model=fus,
        mel_decoder=mel_dec,
    )
    runtime._load_stats = load_stats
    runtime._detected_dims = {'vision': vit_dim, 'audio': audio_dim, 'fusion': fusion_dim}

    return runtime
