from spacy.lang.en import English
import pandas as pd
import gradio as gr
from transformers import pipeline
from gradio.themes.utils.colors import red, green

# Initialize the NLP pipeline
nlp = English()
nlp.add_pipe("sentencizer")

# Initialize the text classification pipeline
detector = pipeline(task='text-classification', model='SJTU-CL/RoBERTa-large-ArguGPT-sent')

# Define color map for highlighted text
color_map = {
    '0%': green.c400,
    '10%': green.c300,
    '20%': green.c200,
    '30%': green.c100,
    '40%': green.c50,
    '50%': red.c50,
    '60%': red.c100,
    '70%': red.c200,
    '80%': red.c300,
    '90%': red.c400,
    '100%': red.c500
}

def predict_doc(doc):
    sents = [s.text for s in nlp(doc).sents]
    data = {'sentence': [], 'label': [], 'score': []}
    res = []
    for sent in sents:
        prob = predict_one_sent(sent)

        data['sentence'].append(sent)
        data['score'].append(round(prob, 4))
        if prob <= 0.5:
            data['label'].append('Human')
        else:
            data['label'].append('Machine')

        if prob < 0.1: label = '0%'
        elif prob < 0.2: label = '10%'
        elif prob < 0.3: label = '20%'
        elif prob < 0.4: label = '30%'
        elif prob < 0.5: label = '40%'
        elif prob < 0.6: label = '50%'
        elif prob < 0.7: label = '60%'
        elif prob < 0.8: label = '70%'
        elif prob < 0.9: label = '80%'
        elif prob < 1: label = '90%'
        else: label = '100%'
        res.append((sent, label))

    df = pd.DataFrame(data)
    df.to_csv('result.csv')
    overall_score = df.score.mean()
    overall_label = 'Human' if overall_score <= 0.5 else 'Machine'
    sum_str = f'The essay is probably written by {overall_label}. The probability of being generated by AI is {overall_score:.2f}'

    return sum_str, res, df, 'result.csv'

def predict_one_sent(sent):
    res = detector(sent)[0]
    org_label, prob = res['label'], res['score']
    if org_label == 'LABEL_0': prob = 1 - prob
    return prob

# Custom CSS for modern look
custom_css = """
.gradio-container {
    font-family: 'Arial', sans-serif;
}
.gradio-header {
    background-color: #4CAF50;
    color: white;
    padding: 10px;
    text-align: center;
}
.gradio-button {
    background-color: #4CAF50;
    color: white;
    border: none;
    padding: 10px 20px;
    text-align: center;
    text-decoration: none;
    display: inline-block;
    font-size: 16px;
    margin: 4px 2px;
    cursor: pointer;
    border-radius: 5px;
}
.gradio-button:hover {
    background-color: #45a049;
}
"""

with gr.Blocks(css=custom_css) as demo:
    gr.Markdown("## AI vs Human Essay Detector")
    gr.Markdown("This tool helps you determine whether an essay is written by a human or generated by AI.")
    
    with gr.Row():
        with gr.Column():
            text_in = gr.Textbox(
                lines=5, 
                label='Essay Input', 
                info='Please enter the essay in the textbox',
                placeholder="Paste your essay here..."
            )
            btn = gr.Button('Predict who writes this essay!', variant="primary")

        sent_res = gr.HighlightedText(label='Labeled Result', color_map=color_map)
    
    with gr.Row():
        summary = gr.Text(label='Result Summary')
        csv_f = gr.File(label='CSV File Storing Data with All Sentences')

    tab = gr.Dataframe(label='Table with Probability Score', row_count=100)
    btn.click(predict_doc, inputs=[text_in], outputs=[summary, sent_res, tab, csv_f], api_name='predict_doc')

demo.launch()