import subprocess
import re
import pandas as pd
import plotly.express as px
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
from prettytable import PrettyTable
import streamlit as st

st.title('Code Generation on the CoNaLa Dataset')

class CodeGenerator:
    def __init__(self):
        self.tokenizer = AutoTokenizer.from_pretrained("AhmedSSoliman/MarianCG-CoNaLa-Large")
        self.model = AutoModelForSeq2SeqLM.from_pretrained("AhmedSSoliman/MarianCG-CoNaLa-Large")
        
    def generate_code(self, nl_input):
        input_ids = self.tokenizer.encode(nl_input, return_tensors="pt")
        output_ids = self.model.generate(input_ids)
        output_code = self.tokenizer.decode(output_ids[0], skip_special_tokens=True)
        return output_code

   

  
  
    def check_code(self, code):
        with open("temp.py", "w") as f:
            f.write(code)
        result = subprocess.run(["flake8", "temp.py"], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        output = result.stdout.decode()
        error = result.stderr.decode()

        return self._process_output(output, error)

    def check_code_list(self, code_list):
        output = ""
        error = ""
        for code in code_list:
            with open("temp.py", "w") as f:
                f.write(code)
            result = subprocess.run(["flake8", "--count", "temp.py"], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
            output += result.stdout.decode()
            error += result.stderr.decode()

        return self._process_output(output, error)

    def _process_output(self, output, error):
        if output:
            output_counts = self._get_error_counts(output)
            self.show_variables_in_table(output_counts, output)
            self.visualize_all_errors(output_counts)
            self.visualize_error_types(output_counts)
            
            return self._format_error_counts(output_counts)
        else:
            error_counts = self._get_error_counts(error)
            self.show_variables_in_table(output_counts, output)
            self.visualize_all_errors(error_counts)
            self.visualize_error_types(error_counts)

            return self._format_error_counts(error_counts)

    def _get_error_counts(self, output):
        error_counts = {}
        error_messages = re.findall(r"temp.py:(\d+):\d+: (\w\d+)", output)
        for message in error_messages:
            error_type = message[1]
            if error_type in error_counts:
                error_counts[error_type] += 1
            else:
                error_counts[error_type] = 1
        return error_counts

    def _format_error_counts(self, error_counts):
        error_message = "\n".join([f"{error_type}: {count}" for error_type, count in error_counts.items()])
        return error_message

    def visualize_all_errors(self, error_counts):
        for error_type, count in error_counts.items():
            print(f"{error_type}: {count}\n")


    def visualize_error_types(self, error_counts):
        df = pd.DataFrame({'Error Type': list(error_counts.keys()), 'Count': list(error_counts.values())})
        fig = px.bar(df, x='Count', y='Error Type', title='Error Occurrences in The Generated Code')
        fig.update_layout(
            title={
                'text': "Error Occurrences in The Generated Code",
                'x': 0.5,
                'y': 0.96,
                'xanchor': 'center',
                'yanchor': 'top'
            },
            xaxis_title="Error Counts",
            yaxis_title="Error Codes"
        )
        fig.show()

    def show_variables_in_table(self, output_counts, output):
        table = PrettyTable()
        table.field_names = ["Error Code", "Message"]
        table.add_row([output_counts, output])
        #table.add_row(["Error", error])
        print(table)

    def display_variables(self, output, error):
        output_df = pd.DataFrame({"Output": [output]})
        error_df = pd.DataFrame({"Error": [error]})
        display(pd.concat([output_df, error_df], axis=1))



code_generator = CodeGenerator()


# Streamlit app
def main():
    st.title('Code Generator and Error Checker')
    nl_input = st.text_area('Enter natural language input for code generation')
    if st.button('Generate Code'):
        # Generate code
        output_code = code_generator.generate_code(nl_input)
        st.subheader('Generated Code')
        st.code(output_code, language='python')

        # Check code for errors
        st.subheader('Error Check')
        error_message = code_generator.check_code(output_code)
        st.write('Error Counts:')
        st.write(error_message)


if __name__ == '__main__':
    
    main()