from fastapi import FastAPI
from pydantic import BaseModel
import tensorflow as tf
import numpy as np
from tokenizer_utils import load_index_to_word

app = FastAPI()

# Carrega o modelo apenas uma vez
decoder_model = tf.keras.models.load_model("decoder_model.h5")

# Pydantic model para a requisição
class EmbeddingRequest(BaseModel):
    embedding: list[float]

@app.post("/decode")
async def decode(req: EmbeddingRequest):
    try:
        # Prepara tensor com shape (1, embedding_dim)
        input_tensor = np.array([req.embedding], dtype=np.float32)

        # Faz a previsão
        prediction = decoder_model.predict(input_tensor)

        # Pega os índices com maior probabilidade
        token_ids = prediction.argmax(axis=-1)

        # Normaliza para lista de inteiros
        if isinstance(token_ids, np.ndarray):
            token_ids = token_ids[0] if token_ids.ndim > 1 else token_ids
            token_ids = token_ids.tolist()

        if isinstance(token_ids, (np.integer, int)):
            token_ids = [int(token_ids)]

        # Carrega dicionário index → word
        index_to_word = load_index_to_word()

        # Decodifica
        decoded = " ".join(index_to_word.get(i, "") for i in token_ids).strip()

        return {"decoded": decoded}
    
    except Exception as e:
        return {"error": str(e)}