import gradio as gr
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from gradio_client import Client
import json

# Initialize the Gradio Client if not already done
client = Client("duchaba/Friendly_Text_Moderation")

def moderate_and_plot(text_input):
    if not text_input:
        raise gr.Error("Please enter some text for moderation.")

    # Use the existing client to predict toxicity levels
    prediction_result = client.predict(
        msg=text_input,
        safer=0.02, # This parameter might be specific to the original model's API
        api_name="/fetch_toxicity_level"
    )
    # Parse the JSON string into a Python dictionary
    dictionary_output = json.loads(prediction_result[-1])

    # Filter out non-toxicity related keys from the dictionary_output for the plot
    toxicity_data = {k: v for k, v in dictionary_output.items() if isinstance(v, (int, float)) and not k.startswith(('is_', 'max_', 'sum_', 'safer_')) and k != 'message'}

    # Convert the dictionary to a pandas DataFrame for easier plotting
    df_toxicity = pd.DataFrame(list(toxicity_data.items()), columns=['Toxicity Type', 'Score'])

    # Sort the values for better visualization (optional)
    df_toxicity = df_toxicity.sort_values(by='Score', ascending=False)

    # Create a bar plot
    plt.figure(figsize=(12, 7))
    sns.barplot(x='Toxicity Type', y='Score', hue='Toxicity Type', data=df_toxicity, palette='viridis', legend=False)
    plt.title('Toxicity Scores for the Sample Message')
    plt.xlabel('Toxicity Type')
    plt.ylabel('Score')
    plt.xticks(rotation=45, ha='right') # Rotate labels to prevent overlap
    plt.tight_layout()

    # Add is_flagged and max_key as text labels on the plot
    is_flagged_status = "Unsafe" if dictionary_output.get('is_flagged', False) else "Safe"
    max_key_value = dictionary_output.get('max_key', 'N/A')

    plt.text(0.02, 0.98, f"Flagged Status: {is_flagged_status}", transform=plt.gca().transAxes, fontsize=12, verticalalignment='top', bbox=dict(boxstyle='round,pad=0.5', fc='yellow', alpha=0.5))
    plt.text(0.02, 0.90, f"Max Toxicity Type: {max_key_value}", transform=plt.gca().transAxes, fontsize=12, verticalalignment='top', bbox=dict(boxstyle='round,pad=0.5', fc='lightblue', alpha=0.5))

    # Prepare detailed toxicity scores for text output
    detailed_scores_output = "Toxicity Details:\n"
    keys_to_display = [
        'harassment', 'harassment_threatening', 'hate', 'hate_threatening',
        'illicit', 'illicit_violent', 'self_harm', 'self_harm_instructions',
        'self_harm_intent', 'sexual', 'sexual_minors', 'violence',
        'violence_graphic', 'harassment/threatening', 'hate/threatening',
        'illicit/violent', 'self-harm/intent', 'self-harm/instructions',
        'self-harm', 'sexual/minors', 'violence/graphic'
    ]
    for key in keys_to_display:
        score = dictionary_output.get(key)
        if score is not None:
            detailed_scores_output += f"- {key.replace('_', ' ').replace('/', '/ ')}: {score:.6f}\n"

    return plt.gcf(), detailed_scores_output

# Define custom CSS for centering text
custom_css = """
.centered-text textarea {
    text-align: center;
}
"""

# Create the Gradio interface
iface = gr.Interface(
    fn=moderate_and_plot,
    inputs=gr.Textbox(lines=5, label="Enter text for moderation"),
    outputs=[
        gr.Plot(label="Toxicity Levels Bar Graph"),
        gr.Textbox(label="Detailed Toxicity Scores", lines=10, elem_classes=["centered-text"])
    ],
    title="Friendly Text Moderation Interface",
    description="Enter text to get a visual breakdown of various toxicity levels.",
    css=custom_css # Moved css argument here
)

# Launch the interface
iface.launch(debug=True)