"""
Arquitectura Hybrid Fusion Production-Ready para Multimodal AI
Combina CLIP (vision), Sentence-BERT (text), y Whisper (audio) con 
attention mechanism para cross-modal interactions.
"""
import torch
import torch.nn as nn
from transformers import CLIPModel, CLIPProcessor
from sentence_transformers import SentenceTransformer
import whisper

class MultimodalHybridFusion(nn.Module):
    """
    Hybrid fusion architecture con attention mechanism.
    Soporta text, image, y audio inputs con graceful degradation 
    si alguna modalidad falla.
    """
    
    def __init__(self, 
                 clip_model_name: str = "openai/clip-vit-base-patch32",
                 text_model_name: str = "all-MiniLM-L6-v2",
                 audio_model_name: str = "base",
                 fusion_dim: int = 512):
        super().__init__()
        
        # Encoders especializados por modalidad
        self.clip_model = CLIPModel.from_pretrained(clip_model_name)
        self.clip_processor = CLIPProcessor.from_pretrained(clip_model_name)
        self.text_encoder = SentenceTransformer(text_model_name)
        self.audio_encoder = whisper.load_model(audio_model_name)
        
        # Projection layers (alinear dimensiones)
        self.text_proj = nn.Linear(384, fusion_dim)      # MiniLM output: 384
        self.image_proj = nn.Linear(512, fusion_dim)     # CLIP output: 512
        self.audio_proj = nn.Linear(512, fusion_dim)     # Whisper: 512
        
        # Attention mechanism para cross-modal interactions
        self.cross_attention = nn.MultiheadAttention(
            embed_dim=fusion_dim,
            num_heads=8,
            dropout=0.1,
            batch_first=True
        )
        
        # Fusion layer final
        self.fusion_fc = nn.Sequential(
            nn.Linear(fusion_dim, fusion_dim // 2),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(fusion_dim // 2, 128)
        )
    
    def encode_text(self, text_input: str) -> torch.Tensor:
        """Encode texto con Sentence-BERT"""
        try:
            embedding = self.text_encoder.encode(
                text_input, 
                convert_to_tensor=True
            )
            return self.text_proj(embedding)
        except Exception as e:
            print(f"Error encoding text: {e}")
            return None
    
    def encode_image(self, image_input) -> torch.Tensor:
        """Encode imagen con CLIP"""
        try:
            inputs = self.clip_processor(
                images=image_input, 
                return_tensors="pt"
            )
            image_features = self.clip_model.get_image_features(**inputs)
            return self.image_proj(image_features)
        except Exception as e:
            print(f"Error encoding image: {e}")
            return None
    
    def encode_audio(self, audio_path: str) -> torch.Tensor:
        """Encode audio con Whisper"""
        try:
            audio = whisper.load_audio(audio_path)
            audio = whisper.pad_or_trim(audio)
            mel = whisper.log_mel_spectrogram(audio).to(self.audio_encoder.device)
            _, embedding = self.audio_encoder.embed_audio(mel)
            return self.audio_proj(embedding)
        except Exception as e:
            print(f"Error encoding audio: {e}")
            return None
    
    def forward(self, 
                text_input: str = None, 
                image_input = None, 
                audio_input: str = None) -> tuple:
        """
        Forward pass con graceful degradation.
        Si alguna modalidad falla, continúa con las disponibles.
        """
        embeddings = []
        modality_flags = []
        
        # Encode cada modalidad disponible
        if text_input:
            text_emb = self.encode_text(text_input)
            if text_emb is not None:
                embeddings.append(text_emb)
                modality_flags.append("text")
        
        if image_input is not None:
            image_emb = self.encode_image(image_input)
            if image_emb is not None:
                embeddings.append(image_emb)
                modality_flags.append("image")
        
        if audio_input:
            audio_emb = self.encode_audio(audio_input)
            if audio_emb is not None:
                embeddings.append(audio_emb)
                modality_flags.append("audio")
        
        # Validación: al menos 1 modalidad debe estar disponible
        if len(embeddings) == 0:
            raise ValueError("No valid modalities encoded. System degradation failed.")
        
        # Stack embeddings para attention
        stacked_emb = torch.stack(embeddings).unsqueeze(0)  # [1, num_modalities, fusion_dim]
        
        # Cross-modal attention
        attended_emb, _ = self.cross_attention(stacked_emb, stacked_emb, stacked_emb)
        
        # Pooling (mean) sobre modalidades
        pooled_emb = attended_emb.mean(dim=1)  # [1, fusion_dim]
        
        # Fusion final
        output = self.fusion_fc(pooled_emb)
        
        # Return output + metadata sobre qué modalidades se usaron
        return output, modality_flags

# Ejemplo de uso en producción
if __name__ == "__main__":
    model = MultimodalHybridFusion()
    
    # Caso 1: Todas las modalidades disponibles
    output, modalities = model(
        text_input="El producto llegó dañado, ver foto adjunta",
        image_input="path/to/damaged_product.jpg",
        audio_input="path/to/customer_voice.wav"
    )
    print(f"Output shape: {output.shape}, Modalities used: {modalities}")
    # Expected: [1, 128], ['text', 'image', 'audio']
    
    # Caso 2: Solo texto (graceful degradation)
    output, modalities = model(text_input="Consulta sobre mi pedido")
    print(f"Output shape: {output.shape}, Modalities used: {modalities}")
    # Expected: [1, 128], ['text']

"""
Sistema de rollback automático con health checks y feature flags.
Monitorea métricas críticas cada 30s y revierte a modo unimodal 
si detecta degradación.
"""
import time
import requests
from dataclasses import dataclass
from typing import Dict, Any

@dataclass
class HealthMetrics:
    """Métricas críticas para health check"""
    accuracy: float              # 0-1
    latency_p95: float           # seconds
    error_rate: float            # 0-1
    cross_modal_consistency: float  # STS score 0-1

class AutomatedRollbackSystem:
    """
    Sistema de rollback que monitorea métricas y revierte automáticamente
    si detecta degradación.
    """
    
    def __init__(self, prometheus_url: str, feature_flag_api: str, 
                 thresholds: Dict[str, float]):
        self.prometheus_url = prometheus_url
        self.feature_flag_api = feature_flag_api
        self.thresholds = thresholds
        self.consecutive_failures = 0
        self.max_failures = 3  # Rollback después de 3 checks fallidos
    
    def fetch_current_metrics(self) -> HealthMetrics:
        """Fetch métricas actuales desde Prometheus"""
        try:
            # Query Prometheus para métricas
            accuracy_query = 'avg_over_time(model_accuracy[5m])'
            latency_query = 'histogram_quantile(0.95, rate(request_duration_seconds_bucket[5m]))'
            error_query = 'sum(rate(http_requests_total{status=~"5.."}[5m])) / sum(rate(http_requests_total[5m]))'
            consistency_query = 'avg_over_time(cross_modal_sts_score[5m])'
            
            accuracy = self._query_prometheus(accuracy_query)
            latency = self._query_prometheus(latency_query)
            error_rate = self._query_prometheus(error_query)
            consistency = self._query_prometheus(consistency_query)
            
            return HealthMetrics(
                accuracy=accuracy,
                latency_p95=latency,
                error_rate=error_rate,
                cross_modal_consistency=consistency
            )
        except Exception as e:
            print(f"Error fetching metrics: {e}")
            return None
    
    def _query_prometheus(self, query: str) -> float:
        """Query Prometheus API"""
        response = requests.get(
            f"{self.prometheus_url}/api/v1/query",
            params={'query': query}
        )
        result = response.json()['data']['result']
        if result:
            return float(result[0]['value'][1])
        return 0.0
    
    def check_health(self, metrics: HealthMetrics) -> bool:
        """Validar si métricas cumplen thresholds"""
        checks = {
            'accuracy': metrics.accuracy >= self.thresholds['min_accuracy'],
            'latency': metrics.latency_p95