src/models/model_loader.py

"""
============================================================================
MODEL LOADER - Carga y Predicción con DistilBERT
============================================================================

Este módulo proporciona la clase ModelLoader para:
1. Cargar modelos pre-entrenados de HuggingFace (DistilBERT)
2. Realizar predicciones de análisis de sentimientos
3. Gestionar device (CPU/GPU/MPS) automáticamente
4. Manejar tokenización eficiente

Flujo de uso:
    from src.models.model_loader import ModelLoader
    from src.config import Config
    
    config = Config()
    model = ModelLoader(config)
    
    result = model.predict("This movie is fantastic!")
    print(result['prediction'])  # 'POSITIVE'

ModelLoader
├── Inicialización
│   ├── _detect_device()       # Auto-detecta GPU/CPU
│   ├── _load_tokenizer()      # Carga tokenizer
│   ├── _load_model()          # Carga modelo
│   └── _create_pipeline()     # Crea pipeline HF
│
├── Predicción
│   ├── predict()              # Método principal (auto-detect)
│   ├── predict_single()       # 1 texto
│   ├── predict_batch()        # N textos (eficiente)
│   └── _predict_manual()      # Fallback sin pipeline
│
└── Utilidades
    ├── tokenize()             # Tokenización sin predicción
    ├── get_model_info()       # Info del modelo
    └── __repr__()             # Representación
============================================================================
"""

import torch
import logging
from typing import Union, List, Dict, Any, Optional
from pathlib import Path
from transformers import (
    AutoTokenizer,
    AutoModelForSequenceClassification,
    pipeline
)
import warnings

# Suprimir warnings de transformers (opcional)
warnings.filterwarnings('ignore', category=UserWarning)


# ----------------------------------------------------------------------------
# LOGGING
# ----------------------------------------------------------------------------
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)


# ----------------------------------------------------------------------------
# CLASE PRINCIPAL: ModelLoader
# ----------------------------------------------------------------------------

class ModelLoader:
    """
    Clase para cargar y usar modelos de HuggingFace para análisis de sentimientos.

    Esta clase encapsula toda la lógica de:
    - Carga del modelo y tokenizer
    - Detección automática de device
    - Predicciones individuales y en batch
    - Manejo de cache

    Attributes:
        config: Objeto de configuración del proyecto
        model_name (str): Nombre del modelo en HuggingFace Hub
        device (str): Device en uso ('cpu', 'cuda', 'mps')
        model: Modelo de transformers cargado
        tokenizer: Tokenizer del modelo
        pipeline: Pipeline de HuggingFace para predicción rápida

    Example:
        >>> from src.config import Config
        >>> config = Config()
        >>> loader = ModelLoader(config)
        >>> result = loader.predict("Amazing movie!")
        >>> print(result['prediction'])
        'POSITIVE'
    """

    def __init__(self, config):
        """
        Inicializa el ModelLoader con configuración.

        Args:
            config: Objeto Config con todas las configuraciones del proyecto

        Raises:
            RuntimeError: Si falla la carga del modelo
        """
        self.config = config
        self.model_name = config.model.name

        # Configuraciones del modelo
        self.max_length = config.model.max_length
        self.truncation = config.model.truncation
        self.padding = config.model.padding
        self.return_tensors = config.model.return_tensors
        self.batch_size = config.model.batch_size

        # Cache
        self.cache_dir = Path(config.paths.cache_dir) / "models"
        self.cache_dir.mkdir(parents=True, exist_ok=True)

        logger.info(
            f"🚀 Inicializando ModelLoader con modelo: {self.model_name}")

        # Detectar y configurar device
        self.device = self._detect_device()
        logger.info(f"🖥️  Device detectado: {self.device}")

        # Cargar modelo y tokenizer
        self.tokenizer = self._load_tokenizer()
        self.model = self._load_model()

        # Crear pipeline para predicciones rápidas
        self.pipeline = self._create_pipeline()

        logger.info("✅ ModelLoader inicializado exitosamente")

    # ------------------------------------------------------------------------
    # MÉTODOS PRIVADOS: Inicialización
    # ------------------------------------------------------------------------

    def _detect_device(self) -> str:
        """
        Detecta automáticamente el device óptimo disponible.

        Prioridad:
        1. CUDA (GPU NVIDIA)
        2. MPS (GPU Apple Silicon M1/M2)
        3. CPU (fallback)

        Returns:
            str: 'cuda', 'mps', o 'cpu'
        """
        # Si el usuario especificó manualmente, usar ese
        if self.config.model.device != "auto":
            return self.config.model.device

        # Detectar automáticamente
        if torch.cuda.is_available():
            device = "cuda"
            gpu_name = torch.cuda.get_device_name(0)
            logger.info(f"✅ GPU CUDA disponible: {gpu_name}")
        elif hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
            device = "mps"
            logger.info("✅ GPU Apple Silicon (MPS) disponible")
        else:
            device = "cpu"
            logger.warning(
                "⚠️  No se detectó GPU, usando CPU (será más lento)")

        return device

    def _load_tokenizer(self):
        """
        Carga el tokenizer desde HuggingFace.

        Returns:
            Tokenizer de transformers

        Raises:
            RuntimeError: Si falla la carga del tokenizer
        """
        try:
            logger.info(f"📥 Cargando tokenizer: {self.model_name}")

            tokenizer = AutoTokenizer.from_pretrained(
                self.model_name,
                cache_dir=str(self.cache_dir),
                use_fast=self.config.model.use_fast_tokenizer
            )

            logger.info(
                f"✅ Tokenizer cargado. Vocab size: {tokenizer.vocab_size}")
            return tokenizer

        except Exception as e:
            error_msg = f"❌ Error al cargar tokenizer: {e}"
            logger.error(error_msg)
            raise RuntimeError(error_msg)

    def _load_model(self):
        """
        Carga el modelo desde HuggingFace.

        Returns:
            Modelo de transformers

        Raises:
            RuntimeError: Si falla la carga del modelo
        """
        try:
            logger.info(f"📥 Cargando modelo: {self.model_name}")

            model = AutoModelForSequenceClassification.from_pretrained(
                self.model_name,
                cache_dir=str(self.cache_dir)
            )

            # Mover modelo al device
            model = model.to(self.device)

            # Modo evaluación (importante para inference)
            model.eval()

            # Información del modelo
            num_params = sum(p.numel() for p in model.parameters())
            logger.info(f"✅ Modelo cargado. Parámetros: {num_params:,}")

            return model

        except Exception as e:
            error_msg = f"❌ Error al cargar modelo: {e}"
            logger.error(error_msg)
            raise RuntimeError(error_msg)

    def _create_pipeline(self):
        """
        Crea un pipeline de HuggingFace para predicciones rápidas.

        Los pipelines son wrappers convenientes que manejan:
        - Tokenización
        - Inferencia
        - Post-procesamiento

        Returns:
            Pipeline de sentiment-analysis
        """
        try:
            pipe = pipeline(
                "sentiment-analysis",
                model=self.model,
                tokenizer=self.tokenizer,
                device=0 if self.device == "cuda" else -1,  # 0 = GPU, -1 = CPU
                return_all_scores=True,  # Retornar probabilidades de todas las clases
                truncation=True,           # Truncar textos largos
                max_length=512             # Máximo de tokens
            )

            logger.info("✅ Pipeline de predicción creado")
            return pipe

        except Exception as e:
            logger.warning(f"⚠️  No se pudo crear pipeline: {e}")
            return None

    # ------------------------------------------------------------------------
    # MÉTODOS PÚBLICOS: Predicción
    # ------------------------------------------------------------------------

    def predict(
        self,
        texts: Union[str, List[str]],
        return_all_scores: bool = True
    ) -> Union[Dict[str, Any], List[Dict[str, Any]]]:
        """
        Predice el sentimiento de uno o múltiples textos.

        Este es el método principal de predicción. Detecta automáticamente
        si es un texto individual o una lista y llama al método apropiado.

        Args:
            texts: Texto individual (str) o lista de textos
            return_all_scores: Si True, retorna probabilidades de todas las clases

        Returns:
            Dict o List[Dict] con predicciones

            Formato del dict:
            {
                'text': str,                    # Texto original
                'prediction': str,              # 'POSITIVE' o 'NEGATIVE'
                'confidence': float,            # Probabilidad de la clase predicha
                'probabilities': dict,          # {'POSITIVE': 0.92, 'NEGATIVE': 0.08}
                'label_id': int                 # 0 o 1
            }

        Example:
            >>> # Texto individual
            >>> result = model.predict("Great movie!")
            >>> print(result['prediction'])
            'POSITIVE'

            >>> # Lista de textos
            >>> results = model.predict(["Great!", "Terrible!"])
            >>> print([r['prediction'] for r in results])
            ['POSITIVE', 'NEGATIVE']

        Raises:
            ValueError: Si texts es None o vacío
        """
        # Validación de input
        if texts is None:
            raise ValueError("❌ El parámetro 'texts' no puede ser None")

        # Detectar si es individual o batch
        if isinstance(texts, str):
            return self.predict_single(texts, return_all_scores)
        elif isinstance(texts, list):
            return self.predict_batch(texts, return_all_scores)
        else:
            raise ValueError(f"❌ Tipo de input inválido: {type(texts)}")

    def predict_single(
        self,
        text: str,
        return_all_scores: bool = True
    ) -> Dict[str, Any]:
        """
        Predice el sentimiento de un texto individual.

        Args:
            text: Texto a analizar
            return_all_scores: Si True, retorna probabilidades de todas las clases

        Returns:
            Dict con la predicción

        Raises:
            ValueError: Si el texto está vacío
        """
        # Validación
        if not text or not text.strip():
            raise ValueError("❌ El texto no puede estar vacío")

        logger.debug(f"🔍 Prediciendo: '{text[:50]}...'")

        try:
            # Usar pipeline si está disponible (más rápido)
            if self.pipeline is not None:
                result = self.pipeline(text)[0]

                # Parsear resultado del pipeline
                prediction = {
                    'text': text,
                    'prediction': result[0]['label'],
                    'confidence': result[0]['score'],
                    'probabilities': {
                        item['label']: item['score'] for item in result
                    },
                    'label_id': 1 if result[0]['label'] == 'POSITIVE' else 0
                }

            else:
                # Fallback: predicción manual
                prediction = self._predict_manual(text)

            return prediction

        except Exception as e:
            logger.error(f"❌ Error en predicción: {e}")
            raise

    def predict_batch(
        self,
        texts: List[str],
        return_all_scores: bool = True
    ) -> List[Dict[str, Any]]:
        """
        Predice el sentimiento de múltiples textos en batch.

        El procesamiento en batch es más eficiente que procesar
        cada texto individualmente, especialmente con GPU.

        Args:
            texts: Lista de textos a analizar
            return_all_scores: Si True, retorna probabilidades de todas las clases

        Returns:
            Lista de diccionarios con predicciones

        Example:
            >>> texts = ["Great movie!", "Terrible waste of time"]
            >>> results = model.predict_batch(texts)
            >>> for r in results:
            ...     print(f"{r['text']}: {r['prediction']}")
            Great movie!: POSITIVE
            Terrible waste of time: NEGATIVE
        """
        if not texts:
            raise ValueError("❌ La lista de textos no puede estar vacía")

        logger.info(f"🔍 Prediciendo batch de {len(texts)} textos")

        try:
            # Usar pipeline para batch (más eficiente)
            if self.pipeline is not None:
                results = self.pipeline(texts)

                # Parsear resultados
                predictions = []
                for text, result in zip(texts, results):
                    prediction = {
                        'text': text,
                        'prediction': result[0]['label'],
                        'confidence': result[0]['score'],
                        'probabilities': {
                            item['label']: item['score'] for item in result
                        },
                        'label_id': 1 if result[0]['label'] == 'POSITIVE' else 0
                    }
                    predictions.append(prediction)

                return predictions

            else:
                # Fallback: procesar uno por uno
                return [self.predict_single(text) for text in texts]

        except Exception as e:
            logger.error(f"❌ Error en batch prediction: {e}")
            raise

    def _predict_manual(self, text: str) -> Dict[str, Any]:
        """
        Predicción manual sin pipeline (fallback).

        Este método hace la predicción paso a paso:
        1. Tokenizar
        2. Forward pass
        3. Aplicar softmax
        4. Decodificar labels

        Args:
            text: Texto a analizar

        Returns:
            Dict con la predicción
        """
        # 1. Tokenizar
        inputs = self.tokenizer(
            text,
            max_length=self.max_length,
            truncation=self.truncation,
            padding=self.padding,
            return_tensors=self.return_tensors
        )

        # Mover inputs al device
        inputs = {k: v.to(self.device) for k, v in inputs.items()}

        # 2. Forward pass (sin calcular gradientes)
        with torch.no_grad():
            outputs = self.model(**inputs)

        # 3. Obtener logits y aplicar softmax
        logits = outputs.logits
        probabilities = torch.softmax(logits, dim=-1)

        # 4. Obtener predicción
        predicted_class = torch.argmax(probabilities, dim=-1).item()
        confidence = probabilities[0][predicted_class].item()

        # 5. Decodificar label
        label = self.model.config.id2label[predicted_class]

        # 6. Construir resultado
        result = {
            'text': text,
            'prediction': label,
            'confidence': confidence,
            'probabilities': {
                self.model.config.id2label[i]: prob.item()
                for i, prob in enumerate(probabilities[0])
            },
            'label_id': predicted_class
        }

        return result

    # ------------------------------------------------------------------------
    # MÉTODOS PÚBLICOS: Utilidades
    # ------------------------------------------------------------------------

    def tokenize(self, text: str) -> Dict[str, Any]:
        """
        Tokeniza un texto sin hacer predicción.

        Útil para:
        - Debugging
        - Análisis de tokenización
        - Preparación para explicadores (SHAP/LIME)

        Args:
            text: Texto a tokenizar

        Returns:
            Dict con información de tokenización:
            {
                'input_ids': List[int],        # IDs de tokens
                'attention_mask': List[int],   # Máscara de atención
                'tokens': List[str],           # Tokens como strings
                'num_tokens': int              # Número de tokens
            }

        Example:
            >>> result = model.tokenize("Hello world!")
            >>> print(result['tokens'])
            ['[CLS]', 'hello', 'world', '!', '[SEP]']
        """
        # Tokenizar
        encoded = self.tokenizer(
            text,
            max_length=self.max_length,
            truncation=self.truncation,
            padding=self.padding,
            return_tensors=self.return_tensors
        )

        # Decodificar tokens
        tokens = self.tokenizer.convert_ids_to_tokens(
            encoded['input_ids'][0]
        )

        return {
            'input_ids': encoded['input_ids'][0].tolist(),
            'attention_mask': encoded['attention_mask'][0].tolist(),
            'tokens': tokens,
            'num_tokens': len(tokens),
            'text': text
        }

    def get_model_info(self) -> Dict[str, Any]:
        """
        Obtiene información detallada del modelo cargado.

        Returns:
            Dict con información del modelo

        Example:
            >>> info = model.get_model_info()
            >>> print(f"Modelo: {info['model_name']}")
            >>> print(f"Parámetros: {info['num_parameters']:,}")
        """
        num_params = sum(p.numel() for p in self.model.parameters())
        trainable_params = sum(p.numel()
                               for p in self.model.parameters() if p.requires_grad)

        return {
            'model_name': self.model_name,
            'model_type': self.model.config.model_type,
            'num_parameters': num_params,
            'trainable_parameters': trainable_params,
            'num_labels': self.model.config.num_labels,
            'label2id': self.model.config.label2id,
            'id2label': self.model.config.id2label,
            'max_position_embeddings': self.model.config.max_position_embeddings,
            'vocab_size': self.tokenizer.vocab_size,
            'device': self.device,
            'hidden_size': self.model.config.hidden_size,
            'num_attention_heads': self.model.config.num_attention_heads,
            'num_hidden_layers': self.model.config.num_hidden_layers
        }

    def __repr__(self) -> str:
        """Representación legible del ModelLoader"""
        return (
            f"ModelLoader(model={self.model_name}, "
            f"device={self.device})"
        )


# ----------------------------------------------------------------------------
# EJEMPLO DE USO
# ----------------------------------------------------------------------------

if __name__ == "__main__":
    """
    Ejemplo de uso del ModelLoader.
    Ejecutar: python -m src.models.model_loader
    """

    from ..config import setup_project

    print("\n" + "="*60)
    print("🧪 EJEMPLO DE USO: ModelLoader")
    print("="*60 + "\n")

    # 1. Setup del proyecto
    config = setup_project()

    # 2. Inicializar ModelLoader
    loader = ModelLoader(config)

    # 3. Información del modelo
    print("\n📊 INFORMACIÓN DEL MODELO:")
    info = loader.get_model_info()
    for key, value in info.items():
        print(f"  • {key}: {value}")

    # 4. Ejemplos de predicción individual
    print("\n🔍 PREDICCIONES INDIVIDUALES:")

    examples = [
        "This movie was absolutely fantastic! I loved every minute of it.",
        "Terrible waste of time. I want my money back.",
        "It was okay, nothing special but not bad either."
    ]

    for text in examples:
        result = loader.predict(text)
        print(f"\n📝 Texto: {text}")
        print(f"   Predicción: {result['prediction']}")
        print(f"   Confianza: {result['confidence']:.2%}")
        print(f"   Probabilidades: {result['probabilities']}")

    # 5. Predicción en batch
    print("\n🔍 PREDICCIÓN EN BATCH:")
    results = loader.predict_batch(examples)
    for r in results:
        print(
            f"  {r['prediction']:8} ({r['confidence']:.0%}) - {r['text'][:50]}...")

    # 6. Tokenización
    print("\n🔤 TOKENIZACIÓN:")
    tokenized = loader.tokenize(examples[0])
    print(f"  Tokens: {tokenized['tokens'][:10]}...")
    print(f"  Número de tokens: {tokenized['num_tokens']}")

    print("\n" + "="*60)
    print("✅ Ejemplo completado")
    print("="*60 + "\n")