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

# Create a configuration object
config = RobertaConfig.from_pretrained('roberta-base')

# Create the Roberta model
model = RobertaModel.from_pretrained('roberta-base', config=config)

# 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):             
	# Format code into strings and sentences for NLP     
	code_str = " ".join(input_code.split())        
	sentences = [s.strip() for s in code_str.split(".") if s.strip()]   
	#Extract relevant info and intent from code        
	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 info and intent in dictionary    
	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):
	generated_code = model.generate(prompt, max_length=100, num_beams=5, early_stopping=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

# Define Gradio interface
interface = gr.Interface(fn=generate_code, inputs=["textbox"], outputs=["textbox"])

# Have a conversation about the code
input_code = """x = 10
y = 5
def add(a, b):
    return a + b
result = add(x, y)"""
code_analysis = analyze_code(input_code)
prompt = generate_prompt(code_analysis)
reply = f"{prompt}\n\n{generate_code(prompt)}\n\nSuggested improvements: {', '.join(suggest_improvements(input_code))}"
print(reply)

while True:
    change = input("Would you like to make any changes to the code? (Y/N) ")
    if change == "Y":
        new_code = input("Enter the updated code: ")
        code_analysis = analyze_code(new_code)
        prompt = generate_prompt(code_analysis)
        reply = f"{prompt}\n\n{generate_code(prompt)}\n\nSuggested improvements: {', '.join(suggest_improvements(new_code))}"
        print(reply)
    elif change == "N":
        print("OK, conversation ended.")
        break