import gradio as gr
import cv2
import numpy as np
import pandas as pd
from collections import Counter
from ultralytics import YOLO
import plotly.express as px
import plotly.graph_objects as go

# Load the model
model = YOLO("best.pt")

def create_size_distribution_plot(df):
    """Create a box plot of cell sizes for each class."""
    fig = px.box(df, x="class_name", y="area", title="Cell Size Distribution by Type")
    fig.update_layout(
        xaxis_title="Cell Type",
        yaxis_title="Area (pixels²)",
        template="plotly_white"
    )
    return fig

def create_density_heatmap(df, image_shape):
    """Create a heatmap showing cell density."""
    heatmap = np.zeros(image_shape[:2])
    for _, row in df.iterrows():
        center_x = int((row['x_min'] + row['x_max']) / 2)
        center_y = int((row['y_min'] + row['y_max']) / 2)
        heatmap[max(0, center_y-20):min(image_shape[0], center_y+20),
                max(0, center_x-20):min(image_shape[1], center_x+20)] += 1
    
    fig = go.Figure(data=go.Heatmap(z=heatmap))
    fig.update_layout(title="Cell Density Heatmap")
    return fig

def process_image(image, conf_threshold=0.25):
    """Detect cells in the image, extract attributes, and return results."""
    if image is None:
        return None, "No image uploaded", None, None, None
    
    # Convert image to RGB
    image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    
    # Perform detection
    results = model.predict(source=image_rgb, imgsz=640, conf=conf_threshold)
    
    # Get annotated image
    annotated_img = results[0].plot()
    
    # Extract detection data
    detections = results[0].boxes.data if results[0].boxes is not None else []
    
    if len(detections) > 0:
        # Count detections
        class_names = [model.names[int(cls)] for cls in detections[:, 5]]
        count = Counter(class_names)
        detection_str = '\n'.join([f"{name}: {count[name]} cells detected" for name in count])
        
        # Create detailed DataFrame
        df = pd.DataFrame(detections.numpy(), columns=["x_min", "y_min", "x_max", "y_max", "confidence", "class"])
        df["class_name"] = df["class"].apply(lambda x: model.names[int(x)])
        df["width"] = df["x_max"] - df["x_min"]
        df["height"] = df["y_max"] - df["y_min"]
        df["area"] = df["width"] * df["height"]
        
        # Generate summary statistics
        summary = df.groupby("class_name").agg({
            'area': ['count', 'mean', 'std', 'min', 'max'],
            'confidence': 'mean'
        }).round(2)
        summary.columns = ['Count', 'Mean Area', 'Std Dev', 'Min Area', 'Max Area', 'Avg Confidence']
        summary = summary.reset_index()
        
        # Create visualizations
        size_dist_plot = create_size_distribution_plot(df)
        density_plot = create_density_heatmap(df, image.shape)
        
        return (
            annotated_img,
            detection_str,
            summary,
            size_dist_plot,
            density_plot
        )
    else:
        return (
            annotated_img,
            "No cells detected",
            pd.DataFrame(),
            None,
            None
        )

# Create Gradio interface with improved layout
with gr.Blocks(theme=gr.themes.Soft()) as app:
    gr.Markdown("""
    # Bioengineering Image Analysis Tool
    Upload microscopy images to detect and analyze cells using YOLOv10.
    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            input_image = gr.Image(type="numpy", label="Upload Image")
            conf_slider = gr.Slider(
                minimum=0.1,
                maximum=1.0,
                value=0.25,
                step=0.05,
                label="Confidence Threshold",
                info="Adjust detection sensitivity"
            )
            analyze_btn = gr.Button("Analyze Image", variant="primary")
        
        with gr.Column(scale=1):
            output_image = gr.Image(type="numpy", label="Detected Cells")
            detection_text = gr.Textbox(label="Detection Summary", lines=3)
    
    with gr.Row():
        with gr.Column(scale=1):
            stats_df = gr.Dataframe(
                label="Cell Statistics",
                headers=['Cell Type', 'Count', 'Mean Area', 'Std Dev', 'Min Area', 'Max Area', 'Avg Confidence']
            )
    
    with gr.Row():
        with gr.Column(scale=1):
            size_plot = gr.Plot(label="Cell Size Distribution")
        with gr.Column(scale=1):
            density_plot = gr.Plot(label="Cell Density Heatmap")
    
    # Handle button click
    analyze_btn.click(
        process_image,
        inputs=[input_image, conf_slider],
        outputs=[output_image, detection_text, stats_df, size_plot, density_plot]
    )
    
    gr.Markdown("""
    ### Instructions:
    1. Upload a microscopy image containing cells
    2. Adjust the confidence threshold if needed (higher values = stricter detection)
    3. Click 'Analyze Image' to process
    4. View results in the various panels:
        - Annotated image shows detected cells
        - Summary provides cell counts
        - Statistics table shows detailed measurements
        - Plots visualize size distribution and spatial density
    """)

# Launch the app
app.launch()