import torch
from safetensors.torch import load_file
from architecture import ValdamVision2_5
import json

class ValdamVisionLocalPredictor:
    def __init__(self, safetensors_weights_path="model.safetensors", config_path="config.json", vocab_path="vocab.json", device="cpu"):
        self.device = torch.device(device if torch.cuda.is_available() else "cpu")
        
        # Cargar configuraciones y vocabulario
        with open(config_path, "r", encoding="utf-8") as f:
            self.config = json.load(f)
        with open(vocab_path, "r", encoding="utf-8") as f:
            self.vocab_stoi = json.load(f)
        self.vocab_itos = {int(v): k for k, v in self.vocab_stoi.items() if k is not None}

        # Instanciar el modelo con los parámetros configurados
        arch = self.config["architecture"]
        self.model = ValdamVision2_5(
            embed_size=arch["embed_size"],
            vocab_size=len(self.vocab_stoi),
            num_heads=arch.get("num_heads", 8),
            num_layers=arch.get("num_layers", 6)
        ).to(self.device)

        # Cargar pesos estructurados usando safetensors
        weights = load_file(safetensors_weights_path, device=device)
        self.model.load_state_dict(weights)
        self.model.eval()
        print("🚀 Valdam Vision 2.5 cargado de forma ultra-segura con Safetensors.")

    def predict(self, image_tensor, max_len=30):
        # Realiza inferencia local autoregresiva
        with torch.no_grad():
            memory = self.model.encoder(image_tensor.to(self.device))
            input_tokens = torch.tensor([[self.vocab_stoi["<SOS>"]]], device=self.device)

            result = []
            for _ in range(max_len):
                outputs = self.model.decoder(memory, input_tokens)
                next_token = outputs[0, -1, :].argmax().item()
                token_str = self.vocab_itos.get(next_token, "<UNK>")
                
                if token_str == "<EOS>":
                    break
                if token_str not in ["<SOS>", "<PAD>"]:
                    result.append(token_str)
                
                input_tokens = torch.cat([input_tokens, torch.tensor([[next_token]], device=self.device)], dim=1)
            
            return " ".join(result)