import torch
from transformers import ViTImageProcessor, ViTForImageClassification, AutoImageProcessor, AutoModelForImageClassification
from PIL import Image
import matplotlib.pyplot as plt
import numpy as np
import gradio as gr
import io
import base64
import torch.nn.functional as F
import warnings
import os

# Suprimir warnings
warnings.filterwarnings("ignore")

print("🔍 Iniciando sistema de análisis de lesiones de piel...")

# --- CONFIGURACIÓN DE MODELOS VERIFICADOS ---
# Separamos los modelos en dos categorías para mejor explicación al usuario.
# Los modelos especializados en piel son generalmente más fiables para esta tarea.
MODEL_CONFIGS = {
    "especializados": [
        {
            'name': 'Syaha Skin Cancer',
            'id': 'syaha/skin_cancer_detection_model',
            'type': 'custom',
            'accuracy': 0.82,
            'description': 'CNN entrenado en HAM10000',
            'emoji': '🩺'
        },
        {
            'name': 'VRJBro Skin Detection',
            'id': 'VRJBro/skin-cancer-detection',
            'type': 'custom',
            'accuracy': 0.85,
            'description': 'Detector especializado 2024',
            'emoji': '🎯'
        },
        {
            'name': 'Anwarkh1 Skin Cancer',
            'id': 'Anwarkh1/Skin_Cancer-Image_Classification', 
            'type': 'vit',
            'accuracy': 0.89,
            'description': 'Clasificador multi-clase de lesiones de piel',
            'emoji': '🧠'
        },
        {
            'name': 'Jhoppanne SMOTE',
            'id': 'jhoppanne/SkinCancerClassifier_smote-V0',
            'type': 'custom',
            'accuracy': 0.86,
            'description': 'Modelo ISIC 2024 con SMOTE para desequilibrio de clases',
            'emoji': '⚖️'
        },
    ],
    "generales": [
        {
            'name': 'ViT Base General',
            'id': 'google/vit-base-patch16-224',
            'type': 'vit',
            'accuracy': 0.78, 
            'description': 'ViT base pre-entrenado en ImageNet-1k. Excelente para características visuales generales.',
            'emoji': '📈'
        },
        {
            'name': 'ResNet-50 (Microsoft)',
            'id': 'microsoft/resnet-50',
            'type': 'custom', 
            'accuracy': 0.77, 
            'description': 'Un clásico ResNet-50, robusto y de alto rendimiento en clasificación de imágenes generales.',
            'emoji': '⚙️'
        },
        {
            'name': 'DeiT Base (Facebook)',
            'id': 'facebook/deit-base-patch16-224',
            'type': 'vit',
            'accuracy': 0.79, 
            'description': 'Data-efficient Image Transformer, eficiente y de buen rendimiento general.',
            'emoji': '💡'
        },
        {
            'name': 'MobileNetV2 (Google)',
            'id': 'google/mobilenet_v2_1.0_224',
            'type': 'custom',
            'accuracy': 0.72, 
            'description': 'MobileNetV2, modelo ligero y rápido, ideal para entornos con recursos limitados.',
            'emoji': '📱'
        },
        {
            'name': 'Swin Tiny (Microsoft)',
            'id': 'microsoft/swin-tiny-patch4-window7-224', 
            'type': 'custom', 
            'accuracy': 0.81, 
            'description': 'Swin Transformer (Tiny), potente para visión por computadora.',
            'emoji': '🌀'
        },
        # Modelo de respaldo genérico final (si nada más funciona)
        {
            'name': 'ViT Base General (Fallback)',
            'id': 'google/vit-base-patch16-224-in21k',
            'type': 'vit',
            'accuracy': 0.75,
            'description': 'ViT genérico como respaldo final',
            'emoji': '🔄'
        }
    ]
}

# --- CARGA SEGURA DE MODELOS ---
loaded_models = {}
model_performance = {}

def load_model_safe(config):
    """Carga segura de modelos con manejo de errores mejorado"""
    try:
        model_id = config['id']
        model_type = config['type']
        print(f"🔄 Cargando {config['emoji']} {config['name']}...")
        
        try:
            processor = AutoImageProcessor.from_pretrained(model_id)
            model = AutoModelForImageClassification.from_pretrained(model_id)
        except Exception as e_auto:
            if model_type == 'vit':
                try:
                    processor = ViTImageProcessor.from_pretrained(model_id)
                    model = ViTForImageClassification.from_pretrained(model_id)
                except Exception as e_vit:
                    raise e_vit 
            else:
                raise e_auto
        
        model.eval()
        
        # Verificar que el modelo funciona con una entrada dummy
        test_input = processor(Image.new('RGB', (224, 224), color='white'), return_tensors="pt")
        with torch.no_grad():
            test_output = model(**test_input)
            
        print(f"✅ {config['emoji']} {config['name']} cargado exitosamente")
        
        return {
            'processor': processor,
            'model': model,
            'config': config,
            'output_dim': test_output.logits.shape[-1] if hasattr(test_output, 'logits') else len(test_output[0]),
            'category': config.get('category', 'general') # Añadimos la categoría aquí
        }
        
    except Exception as e:
        print(f"❌ {config['emoji']} {config['name']} falló: {e}")
        print(f"    Error detallado: {type(e).__name__}")
        return None

# Cargar modelos
print("\n📦 Cargando modelos...")
# Recorrer ambas categorías de modelos
for category, configs in MODEL_CONFIGS.items():
    for config in configs:
        # Añadir la categoría al diccionario de configuración antes de pasar a load_model_safe
        config['category'] = category 
        model_data = load_model_safe(config)
        if model_data:
            loaded_models[config['name']] = model_data
            model_performance[config['name']] = config.get('accuracy', 0.8)

if not loaded_models:
    print("❌ No se pudo cargar ningún modelo específico. Usando modelos de respaldo...")
    # Modelos de respaldo - más amplios
    fallback_models = [
        'google/vit-base-patch16-224-in21k',
        'microsoft/resnet-50',
        'google/vit-large-patch16-224'
    ]
    
    for fallback_id in fallback_models:
        try:
            print(f"🔄 Intentando modelo de respaldo: {fallback_id}")
            processor = AutoImageProcessor.from_pretrained(fallback_id)
            model = AutoModelForImageClassification.from_pretrained(fallback_id)
            model.eval()
            
            loaded_models[f'Respaldo-{fallback_id.split("/")[-1]}'] = {
                'processor': processor,
                'model': model,
                'config': {
                    'name': f'Respaldo {fallback_id.split("/")[-1]}', 
                    'emoji': '🏥', 
                    'accuracy': 0.75,
                    'type': 'fallback',
                    'category': 'general' # El de respaldo es general
                },
                'category': 'general', # El de respaldo es general
                'type': 'standard'
            }
            print(f"✅ Modelo de respaldo {fallback_id} cargado")
            break
        except Exception as e:
            print(f"❌ Respaldo {fallback_id} falló: {e}")
            continue
    
    if not loaded_models:
        print(f"❌ ERROR CRÍTICO: No se pudo cargar ningún modelo")
        print("💡 Verifica tu conexión a internet y que tengas transformers instalado")
        loaded_models['Modelo Dummy'] = {
            'type': 'dummy',
            'config': {'name': 'Modelo No Disponible', 'emoji': '❌', 'accuracy': 0.0},
            'category': 'dummy'
        }

# Clases de lesiones de piel (HAM10000 dataset)
CLASSES = [
    "Queratosis actínica / Bowen (AKIEC)", 
    "Carcinoma células basales (BCC)",
    "Lesión queratósica benigna (BKL)", 
    "Dermatofibroma (DF)", 
    "Melanoma maligno (MEL)", 
    "Nevus melanocítico (NV)", 
    "Lesión vascular (VASC)"
]

# Sistema de riesgo
RISK_LEVELS = {
    0: {'level': 'Alto', 'color': '#ff6b35', 'urgency': 'Derivación en 48h'},
    1: {'level': 'Crítico', 'color': '#cc0000', 'urgency': 'Derivación inmediata'},
    2: {'level': 'Bajo', 'color': '#44ff44', 'urgency': 'Control rutinario'},
    3: {'level': 'Bajo', 'color': '#44ff44', 'urgency': 'Control rutinario'},
    4: {'level': 'Crítico', 'color': '#990000', 'urgency': 'URGENTE - Oncología'},
    5: {'level': 'Bajo', 'color': '#66ff66', 'urgency': 'Seguimiento 6 meses'},
    6: {'level': 'Moderado', 'color': '#ffaa00', 'urgency': 'Control en 3 meses'}
}

MALIGNANT_INDICES = [0, 1, 4]  # AKIEC, BCC, Melanoma

def predict_with_model(image, model_data):
    """Predicción con un modelo específico - versión mejorada"""
    try:
        config = model_data['config']
        
        # Redimensionar imagen
        image_resized = image.resize((224, 224), Image.LANCZOS)
        
        if model_data.get('type') == 'pipeline': # Esto debería ser poco común con la lista actual
            pipeline = model_data['pipeline']
            results = pipeline(image_resized)
            
            if isinstance(results, list) and len(results) > 0:
                mapped_probs = np.ones(7) / 7 
                confidence = results[0]['score'] if 'score' in results[0] else 0.5
                
                label = results[0].get('label', '').lower()
                if any(word in label for word in ['melanoma', 'mel', 'malignant', 'cancer']):
                    predicted_idx = 4 
                elif any(word in label for word in ['carcinoma', 'bcc', 'basal']):
                    predicted_idx = 1 
                elif any(word in label for word in ['keratosis', 'akiec']):
                    predicted_idx = 0 
                elif any(word in label for word in ['nevus', 'nv', 'benign']):
                    predicted_idx = 5 
                else:
                    predicted_idx = 2 
                    
                mapped_probs[predicted_idx] = confidence
                remaining_sum = (1.0 - confidence)
                if remaining_sum < 0: remaining_sum = 0 
                
                num_other_classes = 6 
                if num_other_classes > 0:
                    remaining_per_class = remaining_sum / num_other_classes
                    for i in range(7):
                        if i != predicted_idx:
                            mapped_probs[i] = remaining_per_class
                
            else:
                mapped_probs = np.ones(7) / 7
                predicted_idx = 5 
                confidence = 0.3
                
        else: # Usar modelo estándar (AutoModel/ViT)
            processor = model_data['processor']
            model = model_data['model']
            
            inputs = processor(image_resized, return_tensors="pt")
            
            with torch.no_grad():
                outputs = model(**inputs)
                
                if hasattr(outputs, 'logits'):
                    logits = outputs.logits
                else:
                    logits = outputs[0] if isinstance(outputs, (tuple, list)) else outputs
                
                probabilities = F.softmax(logits, dim=-1).cpu().numpy()[0]
            
            if len(probabilities) == 7:
                mapped_probs = probabilities
            elif len(probabilities) == 1000: # General ImageNet models
                mapped_probs = np.ones(7) / 7 
                # Ajuste heurístico para modelos generales:
                mapped_probs[5] += 0.1 
                mapped_probs[2] += 0.05 
                mapped_probs = mapped_probs / np.sum(mapped_probs) 
                
            elif len(probabilities) == 2: # Binary classification
                mapped_probs = np.zeros(7)
                if probabilities[1] > 0.5:  # Maligno
                    mapped_probs[4] = probabilities[1] * 0.5  
                    mapped_probs[1] = probabilities[1] * 0.3 
                    mapped_probs[0] = probabilities[1] * 0.2  
                else:  # Benigno
                    mapped_probs[5] = probabilities[0] * 0.6  
                    mapped_probs[2] = probabilities[0] * 0.2  
                    mapped_probs[3] = probabilities[0] * 0.1  
                    mapped_probs[6] = probabilities[0] * 0.1  
                mapped_probs = mapped_probs / np.sum(mapped_probs) 
            else:
                mapped_probs = np.ones(7) / 7
                
            predicted_idx = int(np.argmax(mapped_probs))
            confidence = float(mapped_probs[predicted_idx])
            
        return {
            'model': f"{config['emoji']} {config['name']}",
            'class': CLASSES[predicted_idx],
            'confidence': confidence,
            'probabilities': mapped_probs,
            'is_malignant': predicted_idx in MALIGNANT_INDICES,
            'predicted_idx': predicted_idx,
            'success': True,
            'category': model_data['category'] # Añadir la categoría de vuelta
        }
        
    except Exception as e:
        print(f"❌ Error en {config['name']}: {e}")
        return {
            'model': f"{config.get('name', 'Modelo desconocido')}",
            'success': False,
            'error': str(e),
            'category': model_data.get('category', 'unknown')
        }

def create_probability_chart(predictions, consensus_class):
    """Crear gráfico de barras con probabilidades"""
    try:
        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 6))
        
        # Gráfico 1: Probabilidades por clase (consenso)
        if predictions:
            avg_probs = np.zeros(7)
            valid_predictions = [p for p in predictions if p.get('success', False)]
            
            if len(valid_predictions) > 0:
                for pred in valid_predictions:
                    if isinstance(pred['probabilities'], np.ndarray) and len(pred['probabilities']) == 7 and not np.isnan(pred['probabilities']).any():
                        avg_probs += pred['probabilities']
                    else:
                        print(f"Advertencia: Probabilidades no válidas para {pred['model']}: {pred['probabilities']}")
                avg_probs /= len(valid_predictions)
            else:
                avg_probs = np.ones(7) / 7 
            
            colors = ['#ff6b35' if i in MALIGNANT_INDICES else '#44ff44' for i in range(7)]
            bars = ax1.bar(range(7), avg_probs, color=colors, alpha=0.8)
            
            if consensus_class in CLASSES:
                consensus_idx = CLASSES.index(consensus_class)
                bars[consensus_idx].set_color('#2196F3')
                bars[consensus_idx].set_linewidth(3)
                bars[consensus_idx].set_edgecolor('black')
            
            ax1.set_xlabel('Tipos de Lesión')
            ax1.set_ylabel('Probabilidad Promedio')
            ax1.set_title('📊 Distribución de Probabilidades por Clase')
            ax1.set_xticks(range(7))
            ax1.set_xticklabels([cls.split('(')[1].rstrip(')') for cls in CLASSES], rotation=45)
            ax1.grid(True, alpha=0.3)
            
            for i, bar in enumerate(bars):
                height = bar.get_height()
                ax1.text(bar.get_x() + bar.get_width()/2., height + 0.01,
                                 f'{height:.2%}', ha='center', va='bottom', fontsize=9)
        
        # Gráfico 2: Confianza por modelo
        valid_predictions = [p for p in predictions if p.get('success', False)]
        model_names = [pred['model'].split(' ')[1] if len(pred['model'].split(' ')) > 1 else pred['model'] for pred in valid_predictions]
        confidences = [pred['confidence'] for pred in valid_predictions]
        
        colors_conf = ['#ff6b35' if pred['is_malignant'] else '#44ff44' for pred in valid_predictions]
        bars2 = ax2.bar(range(len(valid_predictions)), confidences, color=colors_conf, alpha=0.8)
        
        ax2.set_xlabel('Modelos')
        ax2.set_ylabel('Confianza')
        ax2.set_title('🎯 Confianza por Modelo')
        ax2.set_xticks(range(len(valid_predictions)))
        ax2.set_xticklabels(model_names, rotation=45)
        ax2.grid(True, alpha=0.3)
        ax2.set_ylim(0, 1)
        
        for i, bar in enumerate(bars2):
            height = bar.get_height()
            ax2.text(bar.get_x() + bar.get_width()/2., height + 0.01,
                            f'{height:.1%}', ha='center', va='bottom', fontsize=9)
        
        plt.tight_layout()
        
        buf = io.BytesIO()
        plt.savefig(buf, format='png', dpi=300, bbox_inches='tight')
        buf.seek(0)
        chart_b64 = base64.b64encode(buf.getvalue()).decode()
        plt.close()
        
        return f'<img src="data:image/png;base64,{chart_b64}" style="width:100%; max-width:800px;">'
        
    except Exception as e:
        print(f"Error creando gráfico: {e}")
        return "<p>❌ Error generando gráfico de probabilidades</p>"

def create_heatmap(predictions):
    """Crear mapa de calor de probabilidades por modelo"""
    try:
        valid_predictions = [p for p in predictions if p.get('success', False)]
        
        if not valid_predictions:
            return "<p>No hay datos suficientes para el mapa de calor</p>"
        
        prob_matrix_list = []
        model_names_for_heatmap = []
        for pred in valid_predictions:
            if isinstance(pred['probabilities'], np.ndarray) and len(pred['probabilities']) == 7 and not np.isnan(pred['probabilities']).any():
                prob_matrix_list.append(pred['probabilities'])
                model_names_for_heatmap.append(pred['model'])
            else:
                print(f"Advertencia: Probabilidades no válidas para heatmap de {pred['model']}: {pred['probabilities']}")
        
        if not prob_matrix_list:
            return "<p>No hay datos válidos para el mapa de calor después de filtrar.</p>"

        prob_matrix = np.array(prob_matrix_list)
        
        fig, ax = plt.subplots(figsize=(10, len(model_names_for_heatmap) * 0.8)) 
        
        im = ax.imshow(prob_matrix, cmap='RdYlGn_r', aspect='auto', vmin=0, vmax=1)
        
        ax.set_xticks(np.arange(7))
        ax.set_yticks(np.arange(len(model_names_for_heatmap)))
        ax.set_xticklabels([cls.split('(')[1].rstrip(')') for cls in CLASSES])
        ax.set_yticklabels(model_names_for_heatmap)
        
        plt.setp(ax.get_xticklabels(), rotation=45, ha="right", rotation_mode="anchor")
        
        for i in range(len(model_names_for_heatmap)):
            for j in range(7):
                text = ax.text(j, i, f'{prob_matrix[i, j]:.2f}',
                                 ha="center", va="center", color="white" if prob_matrix[i, j] > 0.5 else "black",
                                 fontsize=8)
        
        ax.set_title("Mapa de Calor: Probabilidades por Modelo y Clase")
        fig.tight_layout()
        
        cbar = plt.colorbar(im, ax=ax)
        cbar.set_label('Probabilidad', rotation=270, labelpad=15)
        
        buf = io.BytesIO()
        plt.savefig(buf, format='png', dpi=300, bbox_inches='tight')
        buf.seek(0)
        heatmap_b64 = base64.b64encode(buf.getvalue()).decode()
        plt.close()
        
        return f'<img src="data:image/png;base64,{heatmap_b64}" style="width:100%; max-width:800px;">'
        
    except Exception as e:
        print(f"Error creando mapa de calor: {e}")
        return "<p>❌ Error generando mapa de calor</p>"

def analizar_lesion(img):
    """Función principal para analizar la lesión"""
    try:
        if img is None:
            return "<h3>⚠️ Por favor, carga una imagen</h3>"
        
        if not loaded_models or all(m.get('type') == 'dummy' for m in loaded_models.values()):
            return "<h3>❌ Error del Sistema</h3><p>No hay modelos disponibles. Por favor, recarga la aplicación.</p>"
        
        if img.mode != 'RGB':
            img = img.convert('RGB')
        
        predictions = []
        
        for model_name, model_data in loaded_models.items():
            if model_data.get('type') != 'dummy':
                pred = predict_with_model(img, model_data)
                if pred.get('success', False):
                    predictions.append(pred)
        
        if not predictions:
            return "<h3>❌ Error</h3><p>No se pudieron obtener predicciones de ningún modelo.</p>"
        
        # Análisis de consenso
        class_votes = {}
        confidence_sum = {}
        
        for pred in predictions:
            class_name = pred['class']
            confidence = pred['confidence']
            
            if class_name not in class_votes:
                class_votes[class_name] = 0
                confidence_sum[class_name] = 0
            
            class_votes[class_name] += 1
            confidence_sum[class_name] += confidence
        
        consensus_class = max(class_votes.keys(), key=lambda x: class_votes[x])
        avg_confidence = confidence_sum[consensus_class] / class_votes[consensus_class]
        
        consensus_idx = CLASSES.index(consensus_class)
        is_malignant = consensus_idx in MALIGNANT_INDICES
        risk_info = RISK_LEVELS[consensus_idx]
        
        probability_chart = create_probability_chart(predictions, consensus_class)
        heatmap = create_heatmap(predictions)
        
        html_report = f"""
        <div style="font-family: Arial, sans-serif; max-width: 1200px; margin: 0 auto;">
            <h2 style="color: #2c3e50; text-align: center;">🏥 Análisis Completo de Lesión Cutánea</h2>
            
            <div style="background: linear-gradient(135deg, #667eea 0%, #764ba2 100%); color: white; padding: 20px; border-radius: 10px; margin: 20px 0;">
                <h3 style="margin: 0; text-align: center;">📋 Resultado de Consenso</h3>
                <p style="font-size: 18px; text-align: center; margin: 10px 0;"><strong>{consensus_class}</strong></p>
                <p style="text-align: center; margin: 5px 0;">Confianza Promedio: <strong>{avg_confidence:.1%}</strong></p>
                <p style="text-align: center; margin: 5px 0;">Consenso: <strong>{class_votes[consensus_class]}/{len(predictions)} modelos</strong></p>
            </div>
            
            <div style="background: {risk_info['color']}; color: white; padding: 15px; border-radius: 8px; margin: 15px 0;">
                <h4 style="margin: 0;">⚠️ Nivel de Riesgo: {risk_info['level']}</h4>
                <p style="margin: 5px 0;"><strong>{risk_info['urgency']}</strong></p>
                <p style="margin: 5px 0;">Tipo: {'🔴 Potencialmente maligna' if is_malignant else '🟢 Probablemente benigna'}</p>
            </div>
            
            <div style="background: #e3f2fd; padding: 15px; border-radius: 8px; margin: 15px 0;">
                <h4 style="color: #1976d2;">🤖 Resultados Individuales por Modelo</h4>
                <p style="font-size: 0.9em; color: #555;">
                    A continuación se detallan las predicciones de cada modelo. Es importante destacar que los <strong>modelos entrenados específicamente en lesiones de piel (Categoría: Especializados) suelen ser más fiables</strong> para este tipo de análisis que los modelos generales.
                </p>
        """
        
        # RESULTADOS INDIVIDUALES DETALLADOS - Separados por categoría
        
        # Especializados
        html_report += """
        <h5 style="color: #007bff; border-bottom: 1px solid #007bff; padding-bottom: 5px; margin-top: 20px;">
            Modelos Especializados en Lesiones de Piel
        </h5>
        """
        specialized_models_found = False
        for i, pred in enumerate(predictions):
            if pred['success'] and pred['category'] == 'especializados':
                specialized_models_found = True
                model_risk = RISK_LEVELS[pred['predicted_idx']]
                malignant_status = "🔴 Maligna" if pred['is_malignant'] else "🟢 Benigna"
                
                html_report += f"""
                <div style="margin: 15px 0; padding: 15px; background: white; border-radius: 8px; border-left: 5px solid {'#ff6b35' if pred['is_malignant'] else '#44ff44'}; box-shadow: 0 2px 4px rgba(0,0,0,0.1);">
                    <div style="display: flex; justify-content: space-between; align-items: center; margin-bottom: 10px;">
                        <h5 style="margin: 0; color: #333;">{pred['model']}</h5>
                        <span style="background: {model_risk['color']}; color: white; padding: 4px 8px; border-radius: 4px; font-size: 12px;">{model_risk['level']}</span>
                    </div>
                    
                    <div style="display: grid; grid-template-columns: 1fr 1fr 1fr; gap: 10px; font-size: 14px;">
                        <div><strong>Diagnóstico:</strong><br>{pred['class']}</div>
                        <div><strong>Confianza:</strong><br>{pred['confidence']:.1%}</div>
                        <div><strong>Clasificación:</strong><br>{malignant_status}</div>
                    </div>
                    
                    <div style="margin-top: 10px;">
                        <strong>Top 3 Probabilidades:</strong><br>
                        <div style="font-size: 12px; color: #666;">
                """
                
                top_indices = np.argsort(pred['probabilities'])[-3:][::-1]
                for idx in top_indices:
                    prob = pred['probabilities'][idx]
                    if prob > 0.01:
                        html_report += f"• {CLASSES[idx].split('(')[1].rstrip(')')}: {prob:.1%}<br>"
                
                html_report += f"""
                        </div>
                        <div style="margin-top: 8px; font-size: 12px; color: #888;">
                            <strong>Recomendación:</strong> {model_risk['urgency']}
                        </div>
                    </div>
                </div>
                """
        if not specialized_models_found:
            html_report += "<p style='color: #888;'>No se cargaron modelos especializados o fallaron al predecir.</p>"

        # Generales
        html_report += """
        <h5 style="color: #6c757d; border-bottom: 1px solid #6c757d; padding-bottom: 5px; margin-top: 20px;">
            Modelos Generales de Visión
        </h5>
        <p style="font-size: 0.85em; color: #777;">
            Estos modelos son pre-entrenados en grandes datasets de imágenes generales (como ImageNet). Aunque no están optimizados específicamente para lesiones cutáneas, contribuyen al consenso general con su capacidad para reconocer patrones visuales. Sus predicciones son un complemento útil, pero pueden ser menos precisas que las de los modelos especializados.
        </p>
        """
        general_models_found = False
        for i, pred in enumerate(predictions):
            if pred['success'] and pred['category'] == 'generales':
                general_models_found = True
                model_risk = RISK_LEVELS[pred['predicted_idx']]
                malignant_status = "🔴 Maligna" if pred['is_malignant'] else "🟢 Benigna"
                
                html_report += f"""
                <div style="margin: 15px 0; padding: 15px; background: white; border-radius: 8px; border-left: 5px solid {'#ff6b35' if pred['is_malignant'] else '#44ff44'}; box-shadow: 0 2px 4px rgba(0,0,0,0.1);">
                    <div style="display: flex; justify-content: space-between; align-items: center; margin-bottom: 10px;">
                        <h5 style="margin: 0; color: #333;">{pred['model']}</h5>
                        <span style="background: {model_risk['color']}; color: white; padding: 4px 8px; border-radius: 4px; font-size: 12px;">{model_risk['level']}</span>
                    </div>
                    
                    <div style="display: grid; grid-template-columns: 1fr 1fr 1fr; gap: 10px; font-size: 14px;">
                        <div><strong>Diagnóstico:</strong><br>{pred['class']}</div>
                        <div><strong>Confianza:</strong><br>{pred['confidence']:.1%}</div>
                        <div><strong>Clasificación:</strong><br>{malignant_status}</div>
                    </div>
                    
                    <div style="margin-top: 10px;">
                        <strong>Top 3 Probabilidades:</strong><br>
                        <div style="font-size: 12px; color: #666;">
                """
                
                top_indices = np.argsort(pred['probabilities'])[-3:][::-1]
                for idx in top_indices:
                    prob = pred['probabilities'][idx]
                    if prob > 0.01:
                        html_report += f"• {CLASSES[idx].split('(')[1].rstrip(')')}: {prob:.1%}<br>"
                
                html_report += f"""
                        </div>
                        <div style="margin-top: 8px; font-size: 12px; color: #888;">
                            <strong>Recomendación:</strong> {model_risk['urgency']}
                        </div>
                    </div>
                </div>
                """
        if not general_models_found:
            html_report += "<p style='color: #888;'>No se cargaron modelos generales o fallaron al predecir.</p>"
            
        html_report += f"""
            </div>
            
            <div style="background: #f8f9fa; padding: 15px; border-radius: 8px; margin: 15px 0;">
                <h4 style="color: #495057;">📊 Análisis Estadístico</h4>
                <div style="display: grid; grid-template-columns: repeat(auto-fit, minmax(200px, 1fr)); gap: 15px;">
                    <div>
                        <strong>Modelos Activos:</strong> {len([p for p in predictions if p['success']])}/{len(predictions)}<br>
                        <strong>Acuerdo Total:</strong> {class_votes[consensus_class]}/{len([p for p in predictions if p['success']])}<br>
                        <strong>Confianza Máxima:</strong> {max([p['confidence'] for p in predictions if p['success']]):.1%}
                    </div>
                    <div>
                        <strong>Diagnósticos Malignos:</strong> {len([p for p in predictions if p.get('success') and p.get('is_malignant')])}<br>
                        <strong>Diagnósticos Benignos:</strong> {len([p for p in predictions if p.get('success') and not p.get('is_malignant')])}<br>
                        <strong>Consenso Maligno:</strong> {'Sí' if is_malignant else 'No'}
                    </div>
                </div>
            </div>
            
            <div style="background: #ffffff; padding: 15px; border-radius: 8px; margin: 15px 0; border: 1px solid #ddd;">
                <h4 style="color: #333;">📈 Gráficos de Análisis</h4>
                {probability_chart}
            </div>
            
            <div style="background: #ffffff; padding: 15px; border-radius: 8px; margin: 15px 0; border: 1px solid #ddd;">
                <h4 style="color: #333;">🔥 Mapa de Calor de Probabilidades</h4>
                {heatmap}
            </div>
            
            <div style="background: #fff3e0; padding: 15px; border-radius: 8px; margin: 15px 0; border: 1px solid #ff9800;">
                <h4 style="color: #f57c00;">⚠️ Advertencia Médica</h4>
                <p style="margin: 5px 0;">Este análisis es solo una herramienta de apoyo diagnóstico basada en IA.</p>
                <p style="margin: 5px 0;"><strong>Siempre consulte con un dermatólogo profesional para un diagnóstico definitivo.</strong></p>
                <p style="margin: 5px 0;">No utilice esta información como único criterio para decisiones médicas.</p>
                <p style="margin: 5px 0;"><em>Los resultados individuales de cada modelo se muestran para transparencia y análisis comparativo.</em></p>
            </div>
        </div>
        """
        
        return html_report
        
    except Exception as e:
        return f"<h3>❌ Error en el análisis</h3><p>Error técnico: {str(e)}</p><p>Por favor, intente con otra imagen.</p>"

# Configuración de Gradio
def create_interface():
    # Calcular el número total de modelos posibles
    total_possible_models = sum(len(configs) for configs in MODEL_CONFIGS.values())

    with gr.Blocks(theme=gr.themes.Soft(), title="Análisis de Lesiones Cutáneas") as demo:
        gr.Markdown("""
        # 🏥 Sistema de Análisis de Lesiones Cutáneas
        
        **Herramienta de apoyo diagnóstico basada en IA**
        
        Carga una imagen dermatoscópica para obtener una evaluación automatizada.
        """)
        
        with gr.Row():
            with gr.Column(scale=1):
                input_img = gr.Image(
                    type="pil", 
                    label="📷 Imagen Dermatoscópica",
                    height=400
                )
                analyze_btn = gr.Button(
                    "🚀 Analizar Lesión", 
                    variant="primary",
                    size="lg"
                )
                
                gr.Markdown("""
                ### 📝 Instrucciones:
                1. Carga una imagen clara de la lesión.
                2. La imagen debe estar bien iluminada.
                3. Enfócate en la lesión cutánea.
                4. Formatos soportados: JPG, PNG.
                """)
            
            with gr.Column(scale=2):
                output_html = gr.HTML(label="📊 Resultado del Análisis")
        
        analyze_btn.click(
            fn=analizar_lesion,
            inputs=input_img,
            outputs=output_html
        )
        
        gr.Markdown(f"""
        ---
        **Estado del Sistema:**
        - ✅ Modelos cargados: {len(loaded_models)} de {total_possible_models} configurados.
        - 🎯 Precisión promedio estimada: {np.mean(list(model_performance.values())):.1%}
        - ⚠️ **Este sistema es solo para apoyo diagnóstico. Consulte siempre a un profesional médico.**
        """)
    
    return demo

if __name__ == "__main__":
    print(f"\n🚀 Sistema listo!")
    # Calcular el número total de modelos posibles
    total_possible_models = sum(len(configs) for configs in MODEL_CONFIGS.values())
    print(f"📊 Modelos cargados: {len(loaded_models)} de {total_possible_models} configurados.")
    print(f"🎯 Estado: {'✅ Operativo' if loaded_models else '❌ Sin modelos'}")
    
    demo = create_interface()
    demo.launch(share=True, server_name="0.0.0.0", server_port=7860)