import gradio as gr
import torch
import sentencepiece as spm

# ---------------------- Model & SentencePiece Loading ----------------------
@torch.no_grad()
def load_model():
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = torch.jit.load("best_model_scripted.pt", map_location=device).eval()
    return model, device

def load_sp_model():
    sp = spm.SentencePieceProcessor()
    sp.load("spm.model")
    return sp

# Cache models globally
model, device = load_model()
sp = load_sp_model()

# ---------------------- Prediction Function ----------------------
@torch.no_grad()
def predict_next_words(text, max_predictions=3):
    if not text.strip():
        return []

    token_ids = sp.encode(text.strip(), out_type=int)
    if len(token_ids) == 0:
        return []

    input_seq = torch.tensor(token_ids, dtype=torch.long).unsqueeze(0).to(device)
    logits = model(input_seq)
    probabilities = torch.softmax(logits, dim=-1)
    top_indices = torch.topk(probabilities, k=max_predictions, dim=-1).indices.squeeze(0).tolist()

    predicted_pieces = [sp.id_to_piece(idx).lstrip("▁") for idx in top_indices]
    return predicted_pieces

# ---------------------- Gradio App Functions ----------------------
def submit_and_predict(text):
    suggestions = predict_next_words(text)
    suggestions += [""] * (3 - len(suggestions))  # Ensure 3 buttons always
    return suggestions

def append_suggestion(text, suggestion):
    if suggestion:
        text += suggestion + " "
    return text

# ---------------------- Gradio Interface ----------------------
with gr.Blocks(title="Next Word Predictor") as app:
    gr.Markdown("# Next Word Prediction")
    gr.Markdown("Enter text and click 'Submit' to get word suggestions.")

    text_input = gr.Textbox(label="Your Text", placeholder="Type here...", lines=3)
    submit_btn = gr.Button("Submit", variant="primary")

    with gr.Row():
        suggestion_buttons = [gr.Button(visible=False) for _ in range(3)]

    submit_btn.click(
        fn=submit_and_predict,
        inputs=text_input,
        outputs=suggestion_buttons,
    )

    for btn in suggestion_buttons:
        btn.click(
            fn=append_suggestion,
            inputs=[text_input, btn],
            outputs=text_input
        )

if __name__ == "__main__":
    app.launch()