import gradio as gr
import torch
import torch.nn as nn
import sentencepiece as spm
import math

# Load tokenizers
sp_pseudo = spm.SentencePieceProcessor(model_file="pseudocode_tokenizer.model")  # For decoding pseudocode (target)
sp_code = spm.SentencePieceProcessor(model_file="code_tokenizer.model")      # For encoding C++ (source)

# Load the full saved model (architecture + weights)
model_path = "code2pseudo.pth"
model = torch.load(model_path, map_location=torch.device("cuda" if torch.cuda.is_available() else "cpu"), weights_only=False)
model.eval()

# Function to generate pseudocode
def generate_pseudocode(cpp_code, max_len):
    model.eval()
    src = torch.tensor([sp_code.encode_as_ids(cpp_code)], dtype=torch.long)  # Tokenize C++ code
    tgt = torch.tensor([[2]], dtype=torch.long)  # <bos_id>=2
    generated_tokens = [2]  # Start with <START>
    for _ in range(max_len):
        output = model(src, tgt)
        next_token = output[:, -1, :].argmax(-1).item()
        generated_tokens.append(next_token)
        tgt = torch.cat([tgt, torch.tensor([[next_token]])], dim=1)
        if next_token == 3:  # <END>=3
            break
    return sp_pseudo.decode_ids(generated_tokens)  # Final decoded output

# Gradio interface
demo = gr.Interface(
    fn=generate_pseudocode,
    inputs=[
        gr.Textbox(placeholder="Enter C++ code here", label="C++ Code"),
        gr.Slider(minimum=10, maximum=1000, value=50, step=1, label="Max Tokens")
    ],
    outputs=gr.Textbox(label="Generated Pseudocode"),
    title="C++ to Pseudocode Converter",
    description="Enter C++ code and get its pseudocode equivalent using a transformer model."
)

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