import gradio as gr
import torch
import pickle
import numpy as np
from collections import defaultdict

# 1. Configuración segura
DEFAULT_INPUT_SIZE = 4  # Basado en tu error anterior
DEFAULT_HIDDEN_SIZE = 64  # Tamaño común para modelos pequeños

# 2. Función segura para cargar modelos
def safe_load_model(path):
    try:
        # Intentar cargar con pickle
        with open(path, 'rb') as f:
            data = pickle.load(f)
        
        # Verificar estructura
        if not all(k in data for k in ['model', 'vocab', 'idx_to_word']):
            raise ValueError("Estructura inválida del archivo .pkl")
            
        return data
    
    except Exception as e:
        print(f"Error cargando modelo: {e}")
        # Crear datos dummy seguros
        vocab = defaultdict(lambda: 0, {"<unk>": 0, "hola": 1, "mundo": 2, "adiós": 3})
        return {
            'model': torch.nn.LSTM(
                input_size=DEFAULT_INPUT_SIZE,
                hidden_size=DEFAULT_HIDDEN_SIZE,
                num_layers=2
            ).float(),
            'vocab': vocab,
            'idx_to_word': {v: k for k, v in vocab.items()}
        }

# 3. Cargar modelo
model_data = safe_load_model('model.pkl')
model = model_data['model']
vocab = model_data['vocab']
idx_to_word = model_data['idx_to_word']

print(f"Modelo cargado. Input size: {model.input_size}, Hidden size: {model.hidden_size}")

# 4. Preprocesamiento seguro
def preprocess(text, seq_length=5):
    words = text.lower().split()[-seq_length:]
    embeddings = []
    
    for word in words:
        # Embedding básico con exactamente input_size características
        embedding = [
            len(word),
            sum(ord(c) for c in word),
            len(word) * sum(ord(c) for c in word),
            1 if word.isalpha() else 0
        ][:model.input_size]
        
        # Rellenar si es necesario
        if len(embedding) < model.input_size:
            embedding += [0] * (model.input_size - len(embedding))
        
        embeddings.append(embedding)
    
    # Rellenar secuencia
    while len(embeddings) < seq_length:
        embeddings.append([0] * model.input_size)
    
    return torch.tensor([embeddings], dtype=torch.float32)

# 5. Función de predicción robusta
def predict_next_word(text):
    try:
        inputs = preprocess(text)
        print(f"Input shape: {inputs.shape}")
        
        with torch.no_grad():
            outputs, _ = model(inputs)
            
            if outputs.shape[-1] != len(vocab):
                proj = torch.nn.Linear(outputs.shape[-1], len(vocab)).float()
                outputs = proj(outputs)
            
            probs = torch.softmax(outputs[:, -1, :], dim=1)
            pred_idx = probs.argmax().item()
            
            return f"Predicción: {idx_to_word.get(pred_idx, '<unk>')} ({probs[0][pred_idx].item():.2%})"
    
    except Exception as e:
        return f"Error en predicción: {str(e)}"

# 6. Interfaz simplificada
interface = gr.Interface(
    fn=predict_next_word,
    inputs=gr.Textbox(label="Contexto", placeholder="Escribe 3-5 palabras..."),
    outputs="text",
    examples=[["El cielo es"], ["Mi nombre es"], ["Qué hora es"]],
    title="Predictor de Siguiente Palabra",
    description="Escribe un contexto y predice la siguiente palabra"
)

interface.launch(share=True)