import gradio as gr
from gradio_client import Client
import matplotlib.pyplot as plt
import numpy as np
import io
from PIL import Image

client = Client("https://duchaba-friendly-text-moderation.hf.space/--replicas/gffry/")

def moderate_text(text, safer_value):
    result = client.predict(
        text,
        safer_value,
        api_name="/censor_me"
    )

    # Assuming 'result' contains moderation categories and their respective counts
    categories = result.get("categories", {"Safe": 1, "Unsafe": 1})  # Example structure
    labels = list(categories.keys())
    sizes = list(categories.values())
    
    # Generate a pie chart
    fig, ax = plt.subplots()
    ax.pie(sizes, labels=labels, autopct='%1.1f%%', startangle=90)
    ax.axis('equal')  # Equal aspect ratio ensures that pie is drawn as a circle.
    ax.set_title('Moderation Result Distribution')

    # Save plot to a bytes buffer
    buf = io.BytesIO()
    plt.savefig(buf, format='png')
    buf.seek(0)
    
    # Convert bytes buffer to PIL Image
    plot_image = Image.open(buf)

    return result, plot_image

# Define the Gradio interface
demo = gr.Interface(
    fn=moderate_text,
    inputs=[
        gr.Textbox(label="Enter Text:", placeholder="Type your text here...", lines=5),
        gr.Slider(minimum=0.005, maximum=0.1, value=0.005, label="Personalize Safer Value: (larger value is less safe)")
    ],
    outputs=[gr.Textbox(label="Moderated Text:", lines=5), gr.Image(type="pil", label="Moderation Pie Chart")],
    title="Friendly Text Moderator",
    description="Enter text to be moderated and adjust the safer value to see how it affects the moderation.",
    theme="compact"
)

# Customize the CSS
custom_css = """
body {
    background-color: #f5f5f5;
    font-family: Arial, sans-serif;
}
.gradio-container {
    max-width: 800px;
    margin: auto;
    padding: 20px;
    background-color: white;
    border: 1px solid #ddd;
    border-radius: 8px;
    box-shadow: 0 2px 10px rgba(0, 0, 0, 0.1);
}
.gr-button {
    background-color: #4CAF50;
    color: white;
}
.gr-button:hover {
    background-color: #45a049;
}
"""

# Add the custom CSS to the Gradio app
demo.css = custom_css

# Launch the app
demo.launch()