import gradio as gr
import cv2
import numpy as np
import pandas as pd
import os

def process_image(image, txt_file, crop_size, blur_kernel_size, clahe_clip_limit, clahe_tile_grid_size, adaptive_thresh_block_size, adaptive_thresh_c, close_kernel_size_morph, open_kernel_size_morph, min_contour_area, max_contour_area, min_circularity, max_circularity, selection_method, px_per_um, manual_offset_x, manual_offset_y):
    # Convert Gradio image (RGB) to grayscale if it\\'s not already
    if len(image.shape) == 3 and image.shape[2] == 3:
        image_gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
    else:
        image_gray = image # Assume it\\\'s already grayscale

    # Step 1: Crop to Center to Focus on Small Cell
    height, width = image_gray.shape
    center_x, center_y = width // 2, height // 2
    
    y_start = max(0, center_y - crop_size // 2)
    x_start = max(0, center_x - crop_size // 2)
    y_end = min(height, center_y + crop_size // 2)
    x_end = min(width, center_x + crop_size // 2)
    
    img_crop = image_gray[y_start:y_end, x_start:x_end]

    # Step 2: Enhance Contrast (CLAHE)
    clahe = cv2.createCLAHE(clipLimit=clahe_clip_limit, tileGridSize=(clahe_tile_grid_size, clahe_tile_grid_size))
    enhanced_crop = clahe.apply(img_crop)

    # Step 3: Apply Gaussian blur to smooth and highlight contrasts
    # This blurred_crop is used for adaptive thresholding in Colab
    blurred_crop = cv2.GaussianBlur(enhanced_crop, (blur_kernel_size, blur_kernel_size), 0)

    # Step 4: Adaptive Thresholding
    thresh = cv2.adaptiveThreshold(blurred_crop, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
                                   cv2.THRESH_BINARY_INV, adaptive_thresh_block_size, adaptive_thresh_c)

    # Step 5: Morphological Operations (after adaptive thresholding)
    close_kernel = np.ones((close_kernel_size_morph, close_kernel_size_morph), np.uint8)
    thresh = cv2.morphologyEx(thresh, cv2.MORPH_CLOSE, close_kernel)
    open_kernel = np.ones((open_kernel_size_morph, open_kernel_size_morph), np.uint8)
    thresh = cv2.morphologyEx(thresh, cv2.MORPH_OPEN, open_kernel)

    # Step 6: Connected Component Analysis (Crucial for isolating the main cell)
    num_labels, labels, stats, centroids = cv2.connectedComponentsWithStats(thresh, 8, cv2.CV_32S)
    best_component_mask = np.zeros_like(thresh, dtype=np.uint8)
    
    # Center of the cropped image for distance calculation, defined here for scope
    img_crop_center_x, img_crop_center_y = img_crop.shape[1] // 2, img_crop.shape[0] // 2

    if num_labels > 1: # Exclude background
        max_area = 0
        best_label = -1
        
        for i in range(1, num_labels): # Iterate through components, skipping background (label 0)
            area = stats[i, cv2.CC_STAT_AREA]
            # Filter by area (adjust min/max as needed for cell size)
            if 200 < area < 200000: # These values are from Colab, can be exposed as Gradio sliders if needed
                # Calculate distance from centroid to image center
                centroid_x, centroid_y = centroids[i]
                distance = np.sqrt((centroid_x - img_crop_center_x)**2 + (centroid_y - img_crop_center_y)**2)
                
                # Prioritize components closer to the center and with reasonable area
                if area > max_area and distance < crop_size / 4: # Consider components within central quarter
                    max_area = area
                    best_label = i

        if best_label != -1:
            best_component_mask[labels == best_label] = 255
            thresh = best_component_mask # Use the best component as the final thresholded image
        else:
            thresh = np.zeros_like(thresh) # No suitable component found
    else:
        thresh = np.zeros_like(thresh) # No components found (or only background)

    # Step 7: Find Contours on the (now clean) thresholded image
    contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    filtered_contours = []
    if contours:
        for i, cnt in enumerate(contours):
            area = cv2.contourArea(cnt)
            perimeter = cv2.arcLength(cnt, True)
            circularity = 0
            if perimeter > 0:
                circularity = 4 * np.pi * (area / (perimeter * perimeter))
            
            # Apply contour filtering parameters from Colab
            if min_contour_area < area < max_contour_area:
                if min_circularity < circularity < max_circularity:
                    filtered_contours.append(cnt)

    # Step 8: Select Best Contour
    best_contour = None
    if filtered_contours:
        if selection_method == "Largest Area":
            best_contour = max(filtered_contours, key=cv2.contourArea)
        elif selection_method == "Most Circular":
            # Calculate circularity for each contour and find the one closest to 1
            best_contour = min(filtered_contours, key=lambda cnt: abs(1 - (4 * np.pi * cv2.contourArea(cnt) / (cv2.arcLength(cnt, True)**2))))
        elif selection_method == "Closest to Center":
            # img_crop_center_x, img_crop_center_y are already defined above
            def dist_to_center(cnt):
                M = cv2.moments(cnt)
                if M["m00"] == 0: return float("inf")
                cx = int(M["m10"] / M["m00"])
                cy = int(M["m01"] / M["m00"])
                return np.sqrt((cx - img_crop_center_x)**2 + (cy - img_crop_center_y)**2)
            best_contour = min(filtered_contours, key=dist_to_center)

    # Step 9: Final Visualization
    # Create a BGR image from the original (full size) for drawing
    final_image_display = cv2.cvtColor(image_gray, cv2.COLOR_GRAY2BGR)
    
    # Adjust best_contour to full image coordinates for drawing
    best_contour_full_coords = None
    ellipse_full_coords = None
    if best_contour is not None:
        # Shift contour points from cropped coordinates back to full image coordinates
        best_contour_full_coords = best_contour + np.array([x_start, y_start])
        cv2.drawContours(final_image_display, [best_contour_full_coords], -1, (0, 0, 255), 2) # Draw in red
        
        if len(best_contour) >= 5: # Need at least 5 points to fit an ellipse
            ellipse = cv2.fitEllipse(best_contour)
            (ell_center_x, ell_center_y), (major, minor), angle = ellipse
            # Shift ellipse center from cropped coordinates back to full image coordinates
            ellipse_full_coords = ((ell_center_x + x_start, ell_center_y + y_start), (major, minor), angle)
            cv2.ellipse(final_image_display, ellipse_full_coords, (0, 255, 0), 2) # Draw in green

    # Process TXT file if provided
    modified_txt_output = None
    if txt_file is not None:
        df = pd.read_csv(txt_file.name, sep="\\s+", header=None)
        if df.shape[1] < 4:
            raise ValueError("TXT must have at least 4 columns: x_um, y_um, z_um/wavelength, count")
        df.columns = ["x_um", "y_um", "z_um", "count"]
        
        # Adjust y_um for large offset (normalize around 0 to ignore absolute stage position)
        df["y_um"] = df["y_um"] - df["y_um"].mean()

        # Calculate pixel coordinates for data points
        # Use ellipse center for offset if available, otherwise image center
        if ellipse_full_coords is not None:
            ellipse_center_x_full = ellipse_full_coords[0][0]
            ellipse_center_y_full = ellipse_full_coords[0][1]
        else:
            ellipse_center_x_full = width // 2
            ellipse_center_y_full = height // 2

        x_um_mean = df["x_um"].mean() + manual_offset_x
        y_um_mean = df["y_um"].mean() + manual_offset_y
        
        df["x_px"] = (df["x_um"] - x_um_mean) * px_per_um + ellipse_center_x_full
        df["y_px"] = (df["y_um"] - y_um_mean) * px_per_um + ellipse_center_y_full
        
        df["inside"] = df.apply(lambda row: cv2.pointPolygonTest(best_contour_full_coords, (row["x_px"], row["y_px"]), False) >= 0 if best_contour_full_coords is not None else False, axis=1)
        # Set \'count\' to 0 for points outside the contour
        df.loc[df["inside"] == False, "count"] = 0

        for index, row in df.iterrows():
            # Draw points on the full image
            color = (255, 0, 0) if row["inside"] else (0, 165, 255) # Blue for inside, Orange for outside
            cv2.circle(final_image_display, (int(row["x_px"]), int(row["y_px"])), 3, color, -1)

        # Save modified TXT
        modified_txt_path = "modified_streamline.txt"
        df_output = df[["x_um", "y_um", "z_um", "count"]].copy() # Ensure \'count\' column is the modified one
        df_output.to_csv(modified_txt_path, sep="\t", index=False, header=False, float_format="%.6f")
        modified_txt_output = modified_txt_path

    return final_image_display, modified_txt_output

with gr.Blocks() as demo:
    gr.Markdown("## Cell Image Visualization Tool - Step by Step")
    gr.Markdown("Process your cell images step-by-step to detect boundaries and visualize data points.")
    
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### File Upload")
            image_input = gr.Image(type="numpy", label="Cell Image (BMP)")
            txt_input = gr.File(label="Data File (TXT)")

            gr.Markdown("### Step 1: Preprocess Image")
            with gr.Accordion("Adjust Preprocessing Parameters", open=True):
                crop_size = gr.Slider(100, 2000, value=1000, step=100, label="Crop Size (pixels)", info="Size of the square crop around the image center.")
                blur_kernel_size = gr.Slider(3, 15, step=2, value=7, label="Gaussian Blur Kernel Size (odd)", info="Size of the kernel for Gaussian blur. Must be odd.")
                clahe_clip_limit = gr.Slider(1, 100, value=50, step=1, label="CLAHE Clip Limit", info="Threshold for contrast limiting.")
                clahe_tile_grid_size = gr.Slider(2, 16, step=2, value=4, label="CLAHE Tile Grid Size", info="Size of grid for histogram equalization. e.g., 4 means 4x4 grid.")

            gr.Markdown("### Step 2: Adaptive Thresholding")
            with gr.Accordion("Adjust Thresholding Parameters", open=True):
                adaptive_thresh_block_size = gr.Slider(3, 101, step=2, value=41, label="Adaptive Threshold Block Size (odd)", info="Size of a pixel neighborhood that is used to calculate a threshold value. Must be odd.")
                adaptive_thresh_c = gr.Slider(1, 50, value=10, step=1, label="Adaptive Threshold C Value", info="Constant subtracted from the mean or weighted mean.")

            gr.Markdown("### Step 3: Morphological Operations")
            with gr.Accordion("Adjust Morphological Kernels", open=True):
                close_kernel_size_morph = gr.Slider(3, 15, step=2, value=9, label="Closing Kernel Size (odd)", info="Kernel size for morphological closing. Helps fill small holes. Must be odd.")
                open_kernel_size_morph = gr.Slider(3, 15, step=2, value=5, label="Opening Kernel Size (odd)", info="Kernel size for morphological opening. Helps remove small objects. Must be odd.")
            
            gr.Markdown("### Step 4: Find and Filter Contours")
            with gr.Accordion("Adjust Contour Filtering", open=True):
                min_contour_area = gr.Slider(0, 1000, value=100, step=10, label="Min Contour Area (pixels)", info="Minimum area for a contour to be considered.")
                max_contour_area = gr.Slider(1000, 1000000, value=500000, step=1000, label="Max Contour Area (pixels)", info="Maximum area for a contour to be considered.")
                min_circularity = gr.Slider(0.001, 1.0, value=0.010, step=0.001, label="Min Circularity", info="Minimum circularity (4π * Area / Perimeter^2) for a contour. 1.0 is a perfect circle.")
                max_circularity = gr.Slider(1.0, 2.0, value=1.2, step=0.01, label="Max Circularity", info="Maximum circularity for a contour.")

            gr.Markdown("### Step 5: Select Best Contour")
            with gr.Accordion("Choose Selection Method", open=True):
                selection_method = gr.Radio(["Largest Area", "Most Circular", "Closest to Center"], label="Selection Method", value="Closest to Center", info="Method to select the single best contour if multiple are found.")
            
            gr.Markdown("### Step 6: Final Visualization & Data Processing")
            with gr.Accordion("Adjust Visualization and Data Parameters", open=True):
                px_per_um = gr.Slider(0.1, 50, value=21.3, step=0.1, label="Pixels per Micrometer (px/µm)", info="Scale factor to convert micrometers to pixels.")
                manual_offset_x = gr.Slider(-100, 100, value=0, step=1, label="Manual Offset X (µm)", info="Manual adjustment for X-coordinate offset in micrometers.")
                manual_offset_y = gr.Slider(-100, 100, value=0, step=1, label="Manual Offset Y (µm)", info="Manual adjustment for Y-coordinate offset in micrometers.")
            
            run_button = gr.Button("Run Cell Vision Analysis")

        with gr.Column(scale=1):
            gr.Markdown("### Results")
            output_image = gr.Image(label="Processed Image with Detections")
            output_txt = gr.File(label="Download Modified TXT File")

    run_button.click(
        fn=process_image,
        inputs=[
            image_input, txt_input, crop_size, blur_kernel_size, clahe_clip_limit, clahe_tile_grid_size, 
            adaptive_thresh_block_size, adaptive_thresh_c, close_kernel_size_morph, open_kernel_size_morph, 
            min_contour_area, max_contour_area, min_circularity, max_circularity, selection_method, 
            px_per_um, manual_offset_x, manual_offset_y
        ],
        outputs=[output_image, output_txt]
    )

demo.launch()