"""
Ekalavya Multi-Modal - Vision + Audio + Video + Text
Complete multi-modal AI that can see, hear, read, and understand
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Optional, Tuple, List, Union
import math


class VisionEncoder(nn.Module):
    """
    Vision Encoder - Processes images and video frames
    Converts images to embeddings that can be understood by the language model
    """
    def __init__(self, image_size=224, patch_size=16, embed_dim=1024):
        super().__init__()
        self.image_size = image_size
        self.patch_size = patch_size
        self.num_patches = (image_size // patch_size) ** 2
        
        # Patch embedding
        self.patch_embed = nn.Conv2d(
            in_channels=3,
            out_channels=embed_dim,
            kernel_size=patch_size,
            stride=patch_size
        )
        
        # Position embedding
        self.pos_embed = nn.Parameter(torch.zeros(1, self.num_patches + 1, embed_dim))
        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
        
        # Transformer layers for vision
        self.layers = nn.ModuleList([
            VisionTransformerBlock(embed_dim) for _ in range(12)
        ])
        
        self.norm = nn.LayerNorm(embed_dim)
        
    def forward(self, x):
        """
        Args:
            x: Image tensor [batch, channels, height, width]
        Returns:
            Image embeddings [batch, num_patches+1, embed_dim]
        """
        batch_size = x.shape[0]
        
        # Patch embedding
        x = self.patch_embed(x)  # [B, embed_dim, H/P, W/P]
        x = x.flatten(2).transpose(1, 2)  # [B, num_patches, embed_dim]
        
        # Add CLS token
        cls_tokens = self.cls_token.expand(batch_size, -1, -1)
        x = torch.cat([cls_tokens, x], dim=1)  # [B, num_patches+1, embed_dim]
        
        # Add position embeddings
        x = x + self.pos_embed
        
        # Transformer layers
        for layer in self.layers:
            x = layer(x)
        
        x = self.norm(x)
        return x


class VisionTransformerBlock(nn.Module):
    """Transformer block for vision encoder"""
    def __init__(self, dim, num_heads=16, mlp_ratio=4.0):
        super().__init__()
        self.norm1 = nn.LayerNorm(dim)
        self.attn = nn.MultiheadAttention(dim, num_heads, batch_first=True)
        self.norm2 = nn.LayerNorm(dim)
        self.mlp = nn.Sequential(
            nn.Linear(dim, int(dim * mlp_ratio)),
            nn.GELU(),
            nn.Linear(int(dim * mlp_ratio), dim)
        )
        
    def forward(self, x):
        x = x + self.attn(self.norm1(x), self.norm1(x), self.norm1(x))[0]
        x = x + self.mlp(self.norm2(x))
        return x


class AudioEncoder(nn.Module):
    """
    Audio Encoder - Processes speech and audio
    Converts audio waveforms to embeddings
    """
    def __init__(self, embed_dim=1024, sample_rate=16000):
        super().__init__()
        self.sample_rate = sample_rate
        
        # Mel spectrogram conversion
        self.mel_transform = nn.Sequential(
            nn.Conv1d(1, 64, kernel_size=400, stride=160),
            nn.GELU(),
            nn.Conv1d(64, 128, kernel_size=3, stride=2),
            nn.GELU(),
            nn.Conv1d(128, 256, kernel_size=3, stride=2),
            nn.GELU()
        )
        
        # Project to embedding dimension
        self.proj = nn.Linear(256, embed_dim)
        
        # Transformer layers for audio
        self.layers = nn.ModuleList([
            AudioTransformerBlock(embed_dim) for _ in range(6)
        ])
        
        self.norm = nn.LayerNorm(embed_dim)
        
    def forward(self, x):
        """
        Args:
            x: Audio waveform [batch, samples]
        Returns:
            Audio embeddings [batch, time_steps, embed_dim]
        """
        # Add channel dimension
        x = x.unsqueeze(1)  # [B, 1, samples]
        
        # Convert to features
        x = self.mel_transform(x)  # [B, 256, time_steps]
        x = x.transpose(1, 2)  # [B, time_steps, 256]
        
        # Project to embedding dimension
        x = self.proj(x)  # [B, time_steps, embed_dim]
        
        # Transformer layers
        for layer in self.layers:
            x = layer(x)
        
        x = self.norm(x)
        return x


class AudioTransformerBlock(nn.Module):
    """Transformer block for audio encoder"""
    def __init__(self, dim, num_heads=16, mlp_ratio=4.0):
        super().__init__()
        self.norm1 = nn.LayerNorm(dim)
        self.attn = nn.MultiheadAttention(dim, num_heads, batch_first=True)
        self.norm2 = nn.LayerNorm(dim)
        self.mlp = nn.Sequential(
            nn.Linear(dim, int(dim * mlp_ratio)),
            nn.GELU(),
            nn.Linear(int(dim * mlp_ratio), dim)
        )
        
    def forward(self, x):
        x = x + self.attn(self.norm1(x), self.norm1(x), self.norm1(x))[0]
        x = x + self.mlp(self.norm2(x))
        return x


class MultiModalProjector(nn.Module):
    """
    Projects vision and audio embeddings to language model space
    """
    def __init__(self, vision_dim=1024, audio_dim=1024, llm_dim=2048):
        super().__init__()
        self.vision_proj = nn.Sequential(
            nn.Linear(vision_dim, llm_dim),
            nn.GELU(),
            nn.Linear(llm_dim, llm_dim)
        )
        self.audio_proj = nn.Sequential(
            nn.Linear(audio_dim, llm_dim),
            nn.GELU(),
            nn.Linear(llm_dim, llm_dim)
        )
        
    def forward(self, vision_embeds=None, audio_embeds=None):
        """
        Project multi-modal embeddings to LLM space
        """
        projected = {}
        if vision_embeds is not None:
            projected['vision'] = self.vision_proj(vision_embeds)
        if audio_embeds is not None:
            projected['audio'] = self.audio_proj(audio_embeds)
        return projected


class EkalavyaMultiModal(nn.Module):
    """
    Ekalavya Multi-Modal - Complete multi-modal AI
    
    Can process:
    - Text (language model)
    - Images (vision encoder)
    - Videos (multiple frames)
    - Audio/Voice (audio encoder)
    
    All modalities are fused and understood together
    """
    def __init__(
        self,
        llm_dim=2048,
        vision_dim=1024,
        audio_dim=1024,
        image_size=224,
        patch_size=16
    ):
        super().__init__()
        
        # Vision encoder
        self.vision_encoder = VisionEncoder(image_size, patch_size, vision_dim)
        
        # Audio encoder
        self.audio_encoder = AudioEncoder(audio_dim)
        
        # Multi-modal projector
        self.projector = MultiModalProjector(vision_dim, audio_dim, llm_dim)
        
        # Language model (simplified for demo)
        self.llm_dim = llm_dim
        self.text_embed = nn.Embedding(32000, llm_dim)
        
        # Fusion layers
        self.fusion = nn.ModuleList([
            FusionLayer(llm_dim) for _ in range(8)
        ])
        
        # Output head
        self.output_head = nn.Linear(llm_dim, 32000)
        
    def forward(
        self,
        text_tokens=None,
        image=None,
        audio=None,
        video_frames=None
    ):
        """
        Process multi-modal input
        
        Args:
            text_tokens: Text token IDs [batch, seq_len]
            image: Image tensor [batch, 3, height, width]
            audio: Audio waveform [batch, samples]
            video_frames: Video frames [batch, num_frames, 3, height, width]
        
        Returns:
            Fused multi-modal embeddings
        """
        batch_size = text_tokens.shape[0] if text_tokens is not None else 1
        embeddings = []
        
        # Process text
        if text_tokens is not None:
            text_embeds = self.text_embed(text_tokens)
            embeddings.append(text_embeds)
        
        # Process image
        if image is not None:
            vision_embeds = self.vision_encoder(image)
            vision_projected = self.projector(vision_embeds=vision_embeds)['vision']
            embeddings.append(vision_projected)
        
        # Process video (multiple frames)
        if video_frames is not None:
            num_frames = video_frames.shape[1]
            frame_embeds = []
            for i in range(num_frames):
                frame = video_frames[:, i]  # [B, 3, H, W]
                frame_embed = self.vision_encoder(frame)
                frame_embeds.append(frame_embed)
            # Average frame embeddings
            video_embeds = torch.stack(frame_embeds, dim=1).mean(dim=1)
            video_projected = self.projector(vision_embeds=video_embeds)['vision']
            embeddings.append(video_projected)
        
        # Process audio
        if audio is not None:
            audio_embeds = self.audio_encoder(audio)
            audio_projected = self.projector(audio_embeds=audio_embeds)['audio']
            embeddings.append(audio_projected)
        
        # Concatenate all embeddings
        if embeddings:
            combined = torch.cat(embeddings, dim=1)
            
            # Fusion layers
            for layer in self.fusion:
                combined = layer(combined)
            
            # Output
            logits = self.output_head(combined)
            return logits
        
        return None
    
    def generate_response(self, text_prompt, image=None, audio=None, video=None):
        """
        Generate response from multi-modal input
        """
        # Process inputs
        logits = self.forward(
            text_tokens=text_prompt,
            image=image,
            audio=audio,
            video_frames=video
        )
        
        if logits is not None:
            # Get last token predictions
            next_token_logits = logits[:, -1, :]
            next_token = torch.argmax(next_token_logits, dim=-1)
            return next_token
        
        return None


class FusionLayer(nn.Module):
    """Fusion layer for combining multi-modal embeddings"""
    def __init__(self, dim, num_heads=16):
        super().__init__()
        self.norm = nn.LayerNorm(dim)
        self.attn = nn.MultiheadAttention(dim, num_heads, batch_first=True)
        self.mlp = nn.Sequential(
            nn.Linear(dim, dim * 4),
            nn.GELU(),
            nn.Linear(dim * 4, dim)
        )
        
    def forward(self, x):
        x = x + self.attn(self.norm(x), self.norm(x), self.norm(x))[0]
        x = x + self.mlp(self.norm(x))
        return x


# Utility functions for processing different modalities

def process_image(image_path_or_tensor):
    """
    Process image for the model
    """
    import torchvision.transforms as T
    
    if isinstance(image_path_or_tensor, str):
        from PIL import Image
        image = Image.open(image_path_or_tensor).convert('RGB')
    else:
        image = image_path_or_tensor
    
    # Resize and normalize
    transform = T.Compose([
        T.Resize((224, 224)),
        T.ToTensor(),
        T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    ])
    
    if isinstance(image, torch.Tensor):
        return image
    return transform(image).unsqueeze(0)


def process_audio(audio_path_or_tensor, sample_rate=16000):
    """
    Process audio for the model
    """
    if isinstance(audio_path_or_tensor, str):
        import torchaudio
        waveform, sr = torchaudio.load(audio_path_or_tensor)
        if sr != sample_rate:
            resampler = torchaudio.transforms.Resample(sr, sample_rate)
            waveform = resampler(waveform)
        waveform = waveform.mean(dim=0)  # Convert to mono
    else:
        waveform = audio_path_or_tensor
    
    return waveform.unsqueeze(0)


def process_video(video_path, num_frames=8):
    """
    Extract frames from video
    """
    import cv2
    import numpy as np
    
    cap = cv2.VideoCapture(video_path)
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    
    # Sample frames evenly
    frame_indices = np.linspace(0, total_frames-1, num_frames, dtype=int)
    
    frames = []
    for idx in frame_indices:
        cap.set(cv2.CAP_PROP_POS_FRAMES, idx)
        ret, frame = cap.read()
        if ret:
            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            frames.append(process_image(frame))
    
    cap.release()
    
    return torch.cat(frames, dim=0).unsqueeze(0)  # [1, num_frames, 3, H, W]


if __name__ == '__main__':
    print("="*70)
    print("EKALAVYA MYTHOS MULTI-MODAL")
    print("="*70)
    print("\n🎯 Capabilities:")
    print("  ✅ Image Understanding")
    print("  ✅ Video Analysis")
    print("  ✅ Voice/Audio Processing")
    print("  ✅ Text Generation")
    print("  ✅ Multi-modal Fusion")
    
    # Create model
    model = EkalavyaMultiModal()
    print(f"\n📊 Model Parameters: {sum(p.numel() for p in model.parameters()):,}")
    
    # Test with dummy inputs
    batch_size = 2
    text_tokens = torch.randint(0, 32000, (batch_size, 10))
    image = torch.randn(batch_size, 3, 224, 224)
    audio = torch.randn(batch_size, 16000)  # 1 second of audio
    
    print("\n🧪 Testing multi-modal processing...")
    output = model(text_tokens=text_tokens, image=image, audio=audio)
    print(f"✅ Output shape: {output.shape}")
    
    print("\n" + "="*70)
    print("✅ Multi-modal model ready!")
    print("="*70)