import cv2
import numpy as np
import pandas as pd
import gradio as gr
import matplotlib.pyplot as plt
from datetime import datetime

def preprocess_image(image):
    """Enhance image contrast, apply thresholding, and clean noise."""
    if len(image.shape) == 2:
        image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)

    gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
    
    # Apply Gaussian blur to remove noise
    blurred = cv2.GaussianBlur(gray, (5, 5), 0)
    
    # Otsu's Thresholding (more robust than adaptive for blood cells)
    _, binary = cv2.threshold(blurred, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)

    # Morphological operations to improve segmentation
    kernel = np.ones((3, 3), np.uint8)
    clean_mask = cv2.morphologyEx(binary, cv2.MORPH_CLOSE, kernel, iterations=2)
    clean_mask = cv2.morphologyEx(clean_mask, cv2.MORPH_OPEN, kernel, iterations=1)

    return clean_mask

def detect_blood_cells(image):
    """Detect blood cells and extract features."""
    mask = preprocess_image(image)
    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    features = []
    total_area = 0

    for i, contour in enumerate(contours, 1):
        area = cv2.contourArea(contour)
        perimeter = cv2.arcLength(contour, True)
        circularity = (4 * np.pi * area / (perimeter * perimeter)) if perimeter > 0 else 0

        # Filtering: Only count reasonable-sized circular objects
        if 100 < area < 5000 and circularity > 0.7:
            M = cv2.moments(contour)
            if M["m00"] != 0:
                cx = int(M["m10"] / M["m00"])
                cy = int(M["m01"] / M["m00"])
                features.append({
                    'ID': i, 'Area': area, 'Perimeter': perimeter, 
                    'Circularity': circularity, 'Centroid_X': cx, 'Centroid_Y': cy
                })
                total_area += area

    # Summary Statistics
    avg_cell_size = total_area / len(features) if features else 0
    cell_density = len(features) / (image.shape[0] * image.shape[1])  # Density per pixel

    summary = {
        'Total Cells': len(features),
        'Avg Cell Size': avg_cell_size,
        'Cell Density': cell_density
    }

    return contours, features, mask, summary

def process_image(image):
    if image is None:
        return None, None, None, None, None

    contours, features, mask, summary = detect_blood_cells(image)
    vis_img = image.copy()

    for feature in features:
        contour = contours[feature['ID'] - 1]
        cv2.drawContours(vis_img, [contour], -1, (0, 255, 0), 2)
        cv2.putText(vis_img, str(feature['ID']), (feature['Centroid_X'], feature['Centroid_Y']),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 1)

    df = pd.DataFrame(features)
    return vis_img, mask, df, summary

def analyze(image):
    vis_img, mask, df, summary = process_image(image)

    plt.style.use('dark_background')
    fig, axes = plt.subplots(1, 2, figsize=(12, 5))

    if not df.empty:
        axes[0].hist(df['Area'], bins=20, color='cyan', edgecolor='black')
        axes[0].set_title('Cell Size Distribution')

        axes[1].scatter(df['Area'], df['Circularity'], alpha=0.6, c='magenta')
        axes[1].set_title('Area vs Circularity')

    return vis_img, mask, fig, df, summary

# Gradio Interface
demo = gr.Interface(
    fn=analyze,
    inputs=gr.Image(type="numpy"),
    outputs=[
        gr.Image(label="Processed Image"),
        gr.Image(label="Binary Mask"),
        gr.Plot(label="Analysis Plots"),
        gr.Dataframe(label="Detected Cells Data"),
        gr.JSON(label="Summary Statistics")
    ]
)

demo.launch()