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 (independent)
# -----------------------------

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):
    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) -> torch.Tensor:
        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)
        return x[:, 0]


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 modules (independent)
# -----------------------------

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 LisaSpeechRecognizer(nn.Module):
    def __init__(self, vocab_size: int = 32, embed_dim: int = 512, num_layers: int = 6):
        super().__init__()
        self.encoder = AudioEncoder(80, embed_dim, num_layers)
        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):
        enc = self.encoder(mel_features)
        return self.ctc_projection(enc), self.lm_head(enc), enc


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:
        return self.conv_layers(mel_features.transpose(1, 2))


# -----------------------------
# Fusion (independent)
# -----------------------------

class FusionAttention(nn.Module):
    """Matches checkpoint keys under prefix 'model.' (q_proj/k_proj/v_proj/out_proj)."""

    def __init__(self, hidden_dim: int = 1024):
        super().__init__()
        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)


# -----------------------------
# 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:
    """Self-contained runtime with modules matching checkpoint keys."""
    vision_vit: LisaViT
    object_detector: LisaObjectDetector
    speech_recognizer: LisaSpeechRecognizer
    sound_classifier: LisaSoundClassifier
    emotion_detector: LisaEmotionDetector
    vad_model: LisaVAD
    fusion_model: FusionAttention
    
    def __post_init__(self):
        """Initialize runtime state tracking."""
        self._load_stats = {}
        self._detected_dims = {}
    
    def get_model_info(self) -> Dict[str, Any]:
        """Get comprehensive model information."""
        info = {
            'architecture': 'LISA Vision+Audio Runtime',
            'detected_dimensions': getattr(self, '_detected_dims', {}),
            'load_statistics': getattr(self, '_load_stats', {}),
            'components': {}
        }
        
        # Get parameter counts for each component
        components = {
            'vision_vit': self.vision_vit,
            'object_detector': self.object_detector,
            '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,
        }
        
        total_params = 0
        for name, module in components.items():
            param_count = sum(p.numel() for p in module.parameters())
            trainable_count = sum(p.numel() for p in module.parameters() if p.requires_grad)
            info['components'][name] = {
                'total_params': param_count,
                'trainable_params': trainable_count,
                'device': str(next(module.parameters()).device),
                'dtype': str(next(module.parameters()).dtype)
            }
            total_params += param_count
        
        info['total_parameters'] = total_params
        return info
    
    def set_eval_mode(self):
        """Set all models to evaluation mode."""
        for module in [self.vision_vit, self.object_detector, self.speech_recognizer, 
                      self.sound_classifier, self.emotion_detector, self.vad_model, self.fusion_model]:
            module.eval()
    
    def set_train_mode(self):
        """Set all models to training mode."""
        for module in [self.vision_vit, self.object_detector, self.speech_recognizer, 
                      self.sound_classifier, self.emotion_detector, self.vad_model, self.fusion_model]:
            module.train()
    
    def to(self, device: torch.device):
        """Move all models to specified 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)
        return self
    
    def get_vision_features(self, image: torch.Tensor) -> torch.Tensor:
        """Extract vision features from image."""
        with torch.no_grad():
            try:
                result = self.vision_vit(image)
                if isinstance(result, tuple) and len(result) >= 2:
                    cls_token, patch_tokens = result[:2]
                    return cls_token
                else:
                    # Single output or different format
                    return result if not isinstance(result, tuple) else result[0]
            except Exception as e:
                logger.warning(f"Vision feature extraction error: {e}")
                # Return dummy features
                batch_size = image.shape[0]
                return torch.zeros(batch_size, 768)  # Standard ViT feature size
    
    def get_audio_features(self, mel_features: torch.Tensor, model_type: str = 'speech') -> torch.Tensor:
        """Extract audio features from mel spectrogram."""
        with torch.no_grad():
            try:
                if model_type == 'speech':
                    result = self.speech_recognizer(mel_features)
                    if isinstance(result, tuple) and len(result) >= 3:
                        _, _, features = result
                        return features.mean(dim=1) if features.ndim > 2 else features  # Global pool
                    elif isinstance(result, tuple) and len(result) >= 1:
                        features = result[0]
                        return features.mean(dim=1) if features.ndim > 2 else features
                    else:
                        return result.mean(dim=1) if result.ndim > 2 else result
                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}")
                # Return dummy features
                batch_size = mel_features.shape[0]
                return torch.zeros(batch_size, 512)  # Standard audio feature size
               
    
    def detect_objects(self, image: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """Perform object detection on image."""
        with torch.no_grad():
            try:
                vit_result = self.vision_vit(image)
                if isinstance(vit_result, tuple) and len(vit_result) >= 2:
                    cls_token, patch_tokens = vit_result[:2]
                else:
                    # Use the single result as both cls and patch tokens
                    cls_token = vit_result if not isinstance(vit_result, tuple) else vit_result[0]
                    patch_tokens = cls_token
                
                result = self.object_detector(patch_tokens)
                if isinstance(result, tuple):
                    if len(result) >= 3:
                        return result[:3]  # Return first 3 elements
                    elif len(result) == 2:
                        cls_out, reg_out = result
                        # Create appropriate objectness score
                        batch_size, num_patches = cls_out.shape[:2]
                        obj_out = torch.ones(batch_size, num_patches, 1, device=cls_out.device)
                        return cls_out, reg_out, obj_out
                    elif len(result) == 1:
                        output = result[0]
                        batch_size, num_patches = output.shape[:2]
                        # Create dummy outputs
                        cls_out = output
                        reg_out = torch.zeros(batch_size, num_patches, 4, device=output.device)
                        obj_out = torch.ones(batch_size, num_patches, 1, device=output.device)
                        return cls_out, reg_out, obj_out
                    else:
                        raise ValueError(f"Unexpected object detector output length: {len(result)}")
                else:
                    # Single tensor output
                    batch_size, num_patches = result.shape[:2]
                    cls_out = result
                    reg_out = torch.zeros(batch_size, num_patches, 4, device=result.device)
                    obj_out = torch.ones(batch_size, num_patches, 1, device=result.device)
                    return cls_out, reg_out, obj_out
            except Exception as e:
                logger.warning(f"Object detection model error: {e}")
                # Return dummy outputs
                batch_size = image.shape[0]
                num_patches = (image.shape[-1] // 16) ** 2  # Assuming 16x16 patches
                cls_out = torch.zeros(batch_size, num_patches, 80)  # 80 classes
                reg_out = torch.zeros(batch_size, num_patches, 4)
                obj_out = torch.zeros(batch_size, num_patches, 1)
                return cls_out, reg_out, obj_out
    
    def transcribe_speech(self, mel_features: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        """Transcribe speech from mel features."""
        with torch.no_grad():
            try:
                result = self.speech_recognizer(mel_features)
                if isinstance(result, tuple):
                    if len(result) >= 3:
                        ctc_out, lm_out, _ = result
                        return ctc_out, lm_out
                    elif len(result) == 2:
                        ctc_out, lm_out = result
                        return ctc_out, lm_out
                    elif len(result) == 1:
                        output = result[0]
                        return output, output  # Use same output for both
                    else:
                        raise ValueError(f"Unexpected number of outputs: {len(result)}")
                else:
                    # Single tensor output
                    return result, result
            except Exception as e:
                logger.warning(f"Speech recognition model error: {e}")
                # Return dummy outputs with correct shapes
                batch_size = mel_features.shape[0]
                dummy_output = torch.zeros(batch_size, 10)  # Dummy sequence output
                return dummy_output, dummy_output
    
    def classify_sound(self, mel_features: torch.Tensor) -> torch.Tensor:
        """Classify sound from mel features."""
        with torch.no_grad():
            return self.sound_classifier(mel_features)
    
    def detect_emotion(self, mel_features: torch.Tensor) -> torch.Tensor:
        """Detect emotion from audio features."""
        with torch.no_grad():
            return self.emotion_detector(mel_features)
    
    def detect_voice_activity(self, mel_features: torch.Tensor) -> torch.Tensor:
        """Detect voice activity."""
        with torch.no_grad():
            return self.vad_model(mel_features)
    
    def multimodal_fusion(self, vision_features: torch.Tensor, audio_features: torch.Tensor) -> torch.Tensor:
        """Perform multimodal fusion of vision and audio features."""
        # Simple fusion by concatenation then projection
        combined = torch.cat([vision_features, audio_features], dim=-1)
        # This is a simplified version - real fusion would use attention
        return self.fusion_model.q_proj(combined)  # Use one of the projections as example


def infer_dimensions_from_ckpt(ckpt: Dict[str, torch.Tensor]) -> Tuple[int, int, int]:
    """Enhanced dimension inference from checkpoint."""
    # Vision ViT dimension
    vit_norm = ckpt.get('vision.vit.norm.weight')
    vit_dim = int(vit_norm.numel()) if vit_norm is not None else 768
    
    # Audio dimension (try speech recognizer first, then sound classifier)
    audio_dim = 512  # default
    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 dimension
    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:
    """Enhanced runtime builder with automatic dimension detection and robust loading."""
    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)
    
    # Log detected dimensions
    import logging
    logger = logging.getLogger(__name__)
    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
    )

    # Construct modules with detected dimensions
    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).to(device)

    # Enhanced weight loading with better error handling
    def load_prefix_safe(module: nn.Module, prefix: str):
        """Load weights for a module prefix with enhanced compatibility."""
        sd = module.state_dict()
        loadable = {}
        skipped = []
        
        for k, v in ckpt.items():
            if not k.startswith(prefix):
                continue
            subk = k[len(prefix):]
            
            # Apply naming transformations for compatibility
            original_subk = subk
            
            # Handle different MLP naming conventions
            if 'mlp.0.' in subk:
                subk = subk.replace('mlp.0.', 'mlp.fc1.')
            elif 'mlp.3.' in subk:
                subk = subk.replace('mlp.3.', 'mlp.fc2.')
            
            # Handle attention naming differences
            if 'attn.qkv.' in subk:
                # Some models might use separate q, k, v projections
                continue  # Skip for now, could be handled with more complex logic
            
            # Check if parameter exists and has compatible shape
            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}: parameter not found (mapped to {subk})")
        
        # Load compatible parameters
        missing_keys, unexpected_keys = module.load_state_dict(loadable, strict=False)
        
        loaded_count = len(loadable)
        total_saved = len([k for k in ckpt.keys() if k.startswith(prefix)])
        
        logger.info(f"Loaded {loaded_count}/{total_saved} tensors for {prefix}")
        if skipped:
            logger.warning(f"Skipped {len(skipped)} incompatible tensors for {prefix}")
            for skip_msg in skipped[:3]:  # Show first 3 issues
                logger.debug(f"  {skip_msg}")
        
        return loaded_count, len(skipped)

    # Load weights for each component
    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.')
    
    # Report loading statistics
    total_loaded = sum(stat[0] for stat in load_stats.values())
    total_skipped = sum(stat[1] for stat 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,
    )
    
    # Add loading statistics to runtime
    runtime._load_stats = load_stats
    runtime._detected_dims = {'vision': vit_dim, 'audio': audio_dim, 'fusion': fusion_dim}
    
    return runtime


