import gradio as gr
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns

def detect_ransomware(file):
    # Read uploaded CSV file
    df = pd.read_csv(file.name)

    # Simulated prediction and actual columns for demo
    df['predicted'] = df.index % 2
    df['actual'] = (df.index + 1) % 2

    # Confusion Matrix
    cm = pd.crosstab(df['actual'], df['predicted'], rownames=['Actual'], colnames=['Predicted'])
    plt.figure(figsize=(4, 4))
    sns.heatmap(cm, annot=True, fmt='d', cmap='Blues')
    confusion_path = "confusion_matrix.png"
    plt.title("Confusion Matrix")
    plt.savefig(confusion_path)
    plt.close()

    # Feature Importance (dummy)
    features = ['duration', 'bytes_sent', 'bytes_received']
    importance = [0.4, 0.35, 0.25]
    plt.figure(figsize=(6, 4))
    sns.barplot(x=importance, y=features)
    plt.title("Feature Importance")
    feature_path = "feature_importance.png"
    plt.savefig(feature_path)
    plt.close()

    # Safe vs Ransomware Pie Chart
    safe_count = (df['predicted'] == 0).sum()
    ransomware_count = (df['predicted'] == 1).sum()
    
    plt.figure(figsize=(4, 4))
    plt.pie(
        [safe_count, ransomware_count],
        labels=['Safe', 'Ransomware'],
        colors=['#4CAF50', '#F44336'],
        autopct='%1.1f%%',
        startangle=90
    )
    plt.title("Prediction Distribution: Safe vs Ransomware")
    class_dist_path = "class_distribution.png"
    plt.savefig(class_dist_path)
    plt.close()

    # Detection Results
    results = {
        "Total Records": len(df),
        "Safe Connections": int(safe_count),
        "Ransomware Connections": int(ransomware_count),
        "Ransomware Percentage": f"{(ransomware_count / len(df)) * 100:.2f}%"
    }

    return confusion_path, feature_path, class_dist_path, results

# Create Gradio interface (Modern syntax)
def create_interface():
    iface = gr.Interface(
        fn=detect_ransomware,
        inputs=gr.File(label="Upload Network Traffic CSV File", file_types=[".csv"]),
        outputs=[
            gr.Image(label="Confusion Matrix", type="filepath"),
            gr.Image(label="Feature Importance", type="filepath"),
            gr.Image(label="Prediction Distribution (Pie Chart)", type="filepath"),
            gr.JSON(label="Detection Results")
        ],
        title="Ransomware Detection System",
        description="""
        Upload a network traffic CSV file to detect potential ransomware activity.
        The system will analyze the data and provide visualizations of the detection results.
        """
    )
    return iface

# Launch the app
interface = create_interface()
interface.launch()
interface.launch(share=True)