import torch
import gradio as gr
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer

# Load pretrained model and tokenizer
model_name = "zonghaoyang/DistilRoBERTa-base"
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)

# Define function to analyze input code
def analyze_code(input_code):             
    code_str = " ".join(input_code.split())        
    sentences = [s.strip() for s in code_str.split(".") if s.strip()]   
    variables = []              
    functions = []    
    logic = []       
    for sentence in sentences: 
        if "=" in sentence:           
            variables.append(sentence.split("=")[0].strip())       
        elif "(" in sentence:            
            functions.append(sentence.split("(")[0].strip())       
        else:           
            logic.append(sentence)               
    return {"variables": variables, "functions": functions, "logic": logic}

# Define function to generate prompt from analyzed code  
def generate_prompt(code_analysis):       
    prompt = f"Generate code with the following: \n\n"   
    prompt += f"Variables: {', '.join(code_analysis['variables'])} \n\n"   
    prompt += f"Functions: {', '.join(code_analysis['functions'])} \n\n"   
    prompt += f"Logic: {' '.join(code_analysis['logic'])}"  
    return prompt
       
# Generate code from model and prompt  
def generate_code(prompt):
    input_ids = tokenizer.encode(prompt, return_tensors="pt")
    generated_ids = model.generate(input_ids, max_length=100, num_beams=5, early_stopping=True)
    generated_code = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
    return generated_code

# Suggest improvements to code
def suggest_improvements(code):
    suggestions = ["Use more descriptive variable names", "Add comments to explain complex logic", "Refactor duplicated code into functions"]
    return suggestions

# Main function to integrate the other functions and generate_code
def main_function(input_code):
    code_analysis = analyze_code(input_code)
    prompt = generate_prompt(code_analysis)
    generated_code = generate_code(prompt)
    improvements = suggest_improvements(input_code)
    return generated_code, improvements

# Create Gradio interface
iface = gr.Interface(
    fn=main_function,
    inputs=gr.inputs.Textbox(lines=5, label="Input Code"),
    outputs=[gr.outputs.Textbox(lines=5, label="Generated Code"), gr.outputs.Textbox(lines=5, label="Suggested Improvements")]
)

# Launch Gradio interface
iface.launch()