import cv2
import numpy as np
import pandas as pd
import gradio as gr
from skimage import measure, morphology
from skimage.segmentation import watershed
import matplotlib.pyplot as plt

def process_image(image):
    """Process uploaded image and extract cell features"""
    # Convert to BGR if image is RGB
    if len(image.shape) == 3:
        image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
    
    # Basic preprocessing
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    clahe = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(8,8))
    enhanced = clahe.apply(gray)
    blurred = cv2.medianBlur(enhanced, 5)
    
    # Segmentation
    thresh = cv2.adaptiveThreshold(
        blurred, 255,
        cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
        cv2.THRESH_BINARY_INV, 21, 4
    )
    
    # Clean small noise
    cleaned = morphology.opening(thresh, morphology.disk(2))
    
    # Watershed segmentation
    sure_bg = cv2.dilate(cleaned, morphology.disk(3), iterations=3)
    dist = cv2.distanceTransform(cleaned, cv2.DIST_L2, 5)
    ret, sure_fg = cv2.threshold(dist, 0.5*dist.max(), 255, 0)
    sure_fg = np.uint8(sure_fg)
    unknown = cv2.subtract(sure_bg, sure_fg)
    
    # Marker labelling
    ret, markers = cv2.connectedComponents(sure_fg)
    markers += 1
    markers[unknown == 255] = 0
    markers = watershed(-dist, markers, mask=cleaned)
    
    # Extract features
    features = []
    props = measure.regionprops(markers, intensity_image=gray)
    
    for i, prop in enumerate(props):
        if prop.area < 50:  # Filter small regions
            continue
        
        features.append({
            'cell_id': i+1,
            'area': prop.area,
            'perimeter': prop.perimeter,
            'circularity': (4 * np.pi * prop.area) / (prop.perimeter**2 + 1e-6),
            'mean_intensity': prop.mean_intensity,
            'centroid_x': prop.centroid[1],
            'centroid_y': prop.centroid[0]
        })
    
    # Create visualization
    vis_img = image.copy()
    contours = measure.find_contours(markers, 0.5)
    
    # Draw contours and cell IDs
    for contour in contours:
        coords = contour.astype(int)
        cv2.drawContours(vis_img, [coords], -1, (0,255,0), 1)
    
    for region in measure.regionprops(markers):
        if region.area >= 50:
            y, x = region.centroid
            cv2.putText(vis_img, str(region.label), 
                       (int(x), int(y)),
                       cv2.FONT_HERSHEY_SIMPLEX, 
                       0.4, (255,0,0), 1)
    
    # Create summary plots
    fig, axes = plt.subplots(1, 2, figsize=(12, 6))
    
    # Cell size distribution
    df = pd.DataFrame(features)
    if not df.empty:
        df['area'].hist(ax=axes[0], bins=20)
        axes[0].set_title('Cell Size Distribution')
        axes[0].set_xlabel('Area')
        axes[0].set_ylabel('Count')
        
        # Circularity vs Intensity
        axes[1].scatter(df['circularity'], df['mean_intensity'])
        axes[1].set_title('Circularity vs Intensity')
        axes[1].set_xlabel('Circularity')
        axes[1].set_ylabel('Mean Intensity')
    else:
        axes[0].text(0.5, 0.5, 'No cells detected', ha='center')
        axes[1].text(0.5, 0.5, 'No cells detected', ha='center')
    
    plt.tight_layout()
    
    return (
        cv2.cvtColor(vis_img, cv2.COLOR_BGR2RGB),
        fig,
        df
    )

# Create Gradio interface
with gr.Blocks(title="Cell Analysis Tool") as demo:
    gr.Markdown("""
    # 🔬 Bioengineering Cell Analysis Tool
    
    Upload microscopy images to analyze cell properties:
    - Automated cell detection
    - Feature extraction (size, shape, intensity)
    - Statistical analysis
    
    **Instructions:**
    1. Upload an image containing cells
    2. Wait for analysis to complete
    3. Review results in three tabs:
       - Detected cells visualization
       - Statistical plots
       - Detailed measurements table
    """)
    
    with gr.Row():
        with gr.Column():
            # Fixed the Image component configuration
            input_image = gr.Image(
                label="Upload Image",
                type="numpy"
            )
            analyze_btn = gr.Button(
                "Analyze Image", 
                variant="primary"
            )
        
        with gr.Column():
            with gr.Tabs():
                with gr.Tab("Detection Results"):
                    output_image = gr.Image(
                        label="Detected Cells"
                    )
                with gr.Tab("Statistics"):
                    output_plot = gr.Plot(
                        label="Statistical Analysis"
                    )
                with gr.Tab("Measurements"):
                    output_table = gr.DataFrame(
                        label="Cell Features"
                    )
    
    analyze_btn.click(
        fn=process_image,
        inputs=input_image,
        outputs=[output_image, output_plot, output_table]
    )

# Launch the demo
if __name__ == "__main__":
    demo.launch()