tools/segmentation.py

"""
Segmentation tools for cellpose-sam pipeline with proper smolagents VLM integration.

Key change: Tools now resolve image paths from global context when the provided path
is invalid or empty, preventing LLM path corruption issues.
"""
import base64
import json
import re
from pathlib import Path
from typing import Any, Dict, TYPE_CHECKING
import numpy as np
import cv2
import torch
from PIL import Image
from skimage.measure import regionprops
from cellpose import models
from segment_anything import sam_model_registry, SamPredictor

from smolagents import tool
from smolagents.agents import ActionStep
from langfuse import get_client

from stores import chroma_store
from models.embeddings import get_image_embedding
from utils.image_utils import resize_and_encode_image
from config import settings


langfuse = get_client()


# =============================================================================
# PATH RESOLUTION HELPER
# =============================================================================
def resolve_image_path(provided_path: str, context_type: str = "image") -> str:
    """
    Resolve the actual image path, falling back to global context if needed.
    
    This function handles the case where the LLM corrupts the file path by:
    1. Checking if the provided path exists
    2. If not, retrieving the correct path from global context
    
    Args:
        provided_path: The path provided by the LLM (may be corrupted)
        context_type: Either "image" for input or "output" for segmentation result
    
    Returns:
        The resolved, valid path
        
    Raises:
        FileNotFoundError: If no valid path can be resolved
    """
    # Import here to avoid circular imports
    from agents.agent import get_current_image_path, get_current_output_path
    
    # First, check if the provided path is valid
    if provided_path and Path(provided_path).exists():
        print(f"[Path Resolution] Using provided path: {provided_path}")
        return provided_path
    
    # Path is invalid - try to get from context
    if context_type == "image":
        context_path = get_current_image_path()
    else:
        context_path = get_current_output_path()
    
    if context_path and Path(context_path).exists():
        print(f"[Path Resolution] Provided path invalid, using context: {context_path}")
        print(f"[Path Resolution] (LLM provided: '{provided_path}')")
        return context_path
    
    # Neither worked
    error_msg = f"Could not resolve {context_type} path. Provided: '{provided_path}', Context: '{context_path}'"
    print(f"[Path Resolution] ERROR: {error_msg}")
    raise FileNotFoundError(error_msg)


# --- Global State and Caching ---
_image_cache: Dict[str, tuple[str, str]] = {}
_cellpose_model = None
_sam_predictor = None


def get_cellpose_model():
    """Initialize Cellpose model (singleton)"""
    global _cellpose_model
    if _cellpose_model is None:
        _cellpose_model = models.CellposeModel(gpu=torch.cuda.is_available())
    return _cellpose_model


def get_sam_predictor():
    """Initialize SAM predictor (singleton)"""
    global _sam_predictor
    if _sam_predictor is None:
        device = "cuda" if torch.cuda.is_available() else "cpu"
        sam = sam_model_registry["vit_h"](checkpoint=settings.CELLPOSE_SAM)
        sam.to(device=device)
        _sam_predictor = SamPredictor(sam)
    return _sam_predictor


def _get_cached_image(image_path: str) -> tuple[str, str] | None:
    """Helper to retrieve an image from the cache."""
    if image_path in _image_cache:
        return _image_cache[image_path]
    return None


def _load_and_cache_image(image_path: str) -> tuple[str, str]:
    """Helper to load, encode, and cache an image."""
    image_base64, media_type = resize_and_encode_image(image_path)
    _image_cache[image_path] = (image_base64, media_type)
    return image_base64, media_type


def parse_parameters_from_text(param_text: str) -> dict:
    """Extract parameter values from parameter text string."""
    defaults = {
        'diameter': 25,
        'flow_threshold': 0.6,
        'cellprob_threshold': 0,
        'min_size': 15
    }
    
    params = defaults.copy()
    
    patterns = {
        'diameter': r'diameter[=:]\s*(\d+)',
        'flow_threshold': r'flow_threshold[=:]\s*([\d.]+)',
        'cellprob_threshold': r'cellprob_threshold[=:]\s*([-\d.]+)',
        'min_size': r'min_size[=:]\s*(\d+)'
    }
    
    for param_name, pattern in patterns.items():
        match = re.search(pattern, param_text, re.IGNORECASE)
        if match:
            value = match.group(1)
            if param_name in ['diameter', 'min_size']:
                params[param_name] = int(value)
            else:
                params[param_name] = float(value)
    
    return params


@tool
def get_segmentation_parameters(image_path: str = "", agent: Any = None) -> str:
    """
    Finds the best cellpose-sam segmentation parameters for an image using vector similarity.
    The image will be visible to the VLM for visual analysis.
    
    NOTE: If image_path is empty or invalid, the tool will automatically use the 
    current image from the system context.
    
    Args:
        image_path (str): Path to the image file (optional - uses context if invalid).
        agent (Any, optional): The agent instance, passed automatically by smol-agents.
    
    Returns:
        str: JSON string containing recommended parameters and analysis context
    """
    print(f"\n--- TOOL CALLED: get_segmentation_parameters for '{image_path}' ---")

    # Resolve the actual image path
    try:
        actual_path = resolve_image_path(image_path, context_type="image")
    except FileNotFoundError as e:
        return json.dumps({"error": str(e)})
    
    image_path = actual_path  # Use resolved path from here on

    try:
        # Load and cache image (for internal use)
        image_base64, media_type = _get_cached_image(image_path) or _load_and_cache_image(image_path)
        
    except Exception as e:
        print(f"Warning: Could not read/resize image: {e}")
        return json.dumps({"error": f"Could not read image: {e}"})

    try:
        # Get similar parameters from ChromaDB
        client = chroma_store.get_client()
        collection = client.get_collection(name='cellpose-sam_parameters_by_image_similarity')
        query_embedding = get_image_embedding(image_path)
        
        results = collection.query(query_embeddings=[query_embedding], n_results=1)

        if not (results['metadatas'] and results['metadatas'][0]):
            return json.dumps({"error": "No similar images found in the database."})

        matched_parameters = results['metadatas'][0][0].get('parameter_text', 'N/A')
        matched_image = results['metadatas'][0][0].get('image_name', 'N/A')
        distance = results['distances'][0][0]

        print(f"Most similar: {matched_image} (distance: {distance:.3f})")
        print(f"Recommended: {matched_parameters}")

        # Parse parameters
        params = parse_parameters_from_text(matched_parameters)

        # Analyze image
        image = np.array(Image.open(image_path).convert("RGB"))
        image_shape = image.shape
        stats = {
            'size': (image_shape[0] * image_shape[1]),
            'mean_intensity': float(np.mean(image)),
            'stdev_intensity': float(np.std(image)),
            'min_intensity': int(np.min(image)),
            'max_intensity': int(np.max(image)),
        }

        # Log to Langfuse WITH image (for observability)
        try:
            langfuse.update_current_trace(
                input={
                    "image_path": image_path,
                    "query_image": {
                        "type": "image_url",
                        "image_url": {
                            "url": f"data:{media_type};base64,{image_base64}"
                        }
                    },
                    "image_stats": stats
                },
                metadata={
                    "matched_image": matched_image,
                    "similarity_distance": float(distance),
                    "matched_parameters": matched_parameters,
                    "parsed_parameters": params
                }
            )
        except Exception as log_error:
            print(f"Warning: Could not log to Langfuse: {log_error}")

        # Determine confidence level
        if distance < 0.2:
            confidence = "high"
            confidence_note = "Very similar image found. Parameters should work well as-is."
        elif distance < 0.4:
            confidence = "medium"
            confidence_note = "Similar image found. Parameters are a good starting point but may need minor adjustments."
        else:
            confidence = "low"
            confidence_note = "No very similar images found. Parameters may need significant adjustment based on visual inspection."

        # Return WITHOUT base64 (image already attached to ActionStep)
        response = {
            "status": "success",
            "image_path": image_path,
            "recommended_parameters": params,
            "matched_image": matched_image,
            "similarity_distance": float(distance),
            "confidence": confidence,
            "image_stats": stats,
            "raw_parameter_text": matched_parameters,
            "visual_guidance": "IMAGE NOW VISIBLE: The input image is now attached to this step. "
                              "Please visually inspect the image to assess cell morphology, density, "
                              "and boundaries before deciding whether to adjust the recommended parameters.",
            "recommendation": f"{confidence_note}\n\nRecommended parameters:\n"
                             f"- diameter: {params['diameter']}\n"
                             f"- flow_threshold: {params['flow_threshold']}\n"
                             f"- cellprob_threshold: {params['cellprob_threshold']}\n"
                             f"- min_size: {params['min_size']}\n\n"
                             f"Image stats: {image_shape[0]}x{image_shape[1]} pixels, "
                             f"mean intensity {stats['mean_intensity']:.1f}\n\n"
                             f"To run segmentation, use: run_cellpose_sam(image_path='', "
                             f"diameter={params['diameter']}, flow_threshold={params['flow_threshold']}, "
                             f"cellprob_threshold={params['cellprob_threshold']}, min_size={params['min_size']})"
        }

        return json.dumps(response, indent=2)

    except Exception as e:
        return json.dumps({"error": str(e)})


@tool
def run_cellpose_sam(
    image_path: str = "",
    diameter: int = None,
    flow_threshold: float = None,
    cellprob_threshold: float = None,
    min_size: int = None,
    output_path: str = None,
    use_recommended_params: bool = True,
    agent: Any = None
) -> str:
    """
    Runs cellpose-sam segmentation pipeline on an image with specified parameters.
    Returns results WITHOUT base64 images to prevent GPU memory issues.
    
    NOTE: If image_path is empty or invalid, the tool will automatically use the 
    current image from the system context.
    
    Args:
        image_path (str): Path to the image file (optional - uses context if invalid)
        diameter (int): Expected diameter of cells in pixels
        flow_threshold (float): Flow error threshold (range: 0-1)
        cellprob_threshold (float): Cell probability threshold (range: -6 to 6)
        min_size (int): Minimum cell size in pixels
        output_path (str): Optional path to save the overlay image
        use_recommended_params (bool): If True and params not provided, get recommendations
        agent (Any, optional): The agent instance
    
    Returns:
        str: JSON string with segmentation results (paths and stats, NO base64)
    """
    print(f"\n--- TOOL CALLED: run_cellpose_sam for '{image_path}' ---")
    
    # Resolve the actual image path
    try:
        actual_path = resolve_image_path(image_path, context_type="image")
    except FileNotFoundError as e:
        return json.dumps({"error": str(e)})
    
    image_path = actual_path  # Use resolved path from here on
    
    try:
        # Load and cache input image
        input_image_base64, input_media_type = _get_cached_image(image_path) or _load_and_cache_image(image_path)
    except Exception as e:
        return json.dumps({"error": f"Could not read input image: {e}"})
    
    # Auto-fetch recommended parameters if needed
    if use_recommended_params and all(p is None for p in [diameter, flow_threshold, cellprob_threshold, min_size]):
        print("No parameters provided. Fetching recommended parameters...")
        param_response = get_segmentation_parameters(image_path, agent=agent)
        
        try:
            param_data = json.loads(param_response)
            if param_data.get("status") == "success":
                rec_params = param_data["recommended_parameters"]
                diameter = diameter or rec_params.get('diameter', 25)
                flow_threshold = flow_threshold or rec_params.get('flow_threshold', 0.6)
                cellprob_threshold = cellprob_threshold or rec_params.get('cellprob_threshold', 0)
                min_size = min_size or rec_params.get('min_size', 15)
            else:
                diameter, flow_threshold, cellprob_threshold, min_size = 25, 0.6, 0, 15
        except json.JSONDecodeError:
            diameter, flow_threshold, cellprob_threshold, min_size = 25, 0.6, 0, 15
    else:
        diameter = diameter if diameter is not None else 25
        flow_threshold = flow_threshold if flow_threshold is not None else 0.6
        cellprob_threshold = cellprob_threshold if cellprob_threshold is not None else 0
        min_size = min_size if min_size is not None else 15
    
    print(f"Final parameters: diameter={diameter}, flow_threshold={flow_threshold}, "
          f"cellprob_threshold={cellprob_threshold}, min_size={min_size}")
    
    try:
        # Read image
        img = cv2.imread(image_path)
        if img is None:
            return json.dumps({"error": f"Could not read image at {image_path}"})
        
        img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        cellpose_model = get_cellpose_model()
        sam_predictor = get_sam_predictor()
        
        # Run Cellpose
        print("Running Cellpose...")
        masks_cellpose, flows, styles = cellpose_model.eval(
            img_rgb,
            diameter=diameter,
            flow_threshold=flow_threshold,
            cellprob_threshold=cellprob_threshold,
            min_size=min_size
        )
        
        if masks_cellpose.max() == 0:
            return json.dumps({
                "status": "no_cells_detected",
                "message": "No cells detected. Try adjusting parameters.",
                "parameters": {
                    "diameter": diameter,
                    "flow_threshold": flow_threshold,
                    "cellprob_threshold": cellprob_threshold,
                    "min_size": min_size
                }
            })
        
        print(f"Cellpose detected {masks_cellpose.max()} regions")
        
        # SAM refinement
        sam_predictor.set_image(img_rgb)
        props = regionprops(masks_cellpose)
        boxes = np.array([prop.bbox for prop in props])
        boxes = boxes[:, [1,0,3,2]]
        
        print(f"Refining {len(boxes)} masks with SAM...")
        
        combined_masks = np.zeros(img_rgb.shape[:2], dtype=np.uint16)
        colored_overlay = img_rgb.copy().astype(np.float32)
        
        for i, box in enumerate(boxes):
            masks, scores, _ = sam_predictor.predict(box=box, multimask_output=True)
            best_mask = masks[np.argmax(scores)]
            combined_masks[best_mask] = i + 1
            color = np.random.randint(0, 255, 3)
            colored_overlay[best_mask] = colored_overlay[best_mask] * 0.6 + color * 0.4
        
        # Generate output path
        if output_path is None:
            base_name = image_path.rsplit('.', 1)[0]
            output_path = f"{base_name}_cellpose_sam_overlay.png"
        
        # Save output
        cv2.imwrite(output_path, cv2.cvtColor(colored_overlay.astype(np.uint8), cv2.COLOR_RGB2BGR))
        
        # Store output path in context for later tools
        from agents.agent import set_current_output_path
        set_current_output_path(output_path)
        
        # Load and cache output image
        output_image_base64, output_media_type = _load_and_cache_image(output_path)
        
        # Log to Langfuse WITH both images
        try:
            langfuse.update_current_trace(
                input={
                    "image_path": image_path,
                    "input_image": {
                        "type": "image_url",
                        "image_url": {"url": f"data:{input_media_type};base64,{input_image_base64}"}
                    }
                },
                output={
                    "cell_count": int(masks_cellpose.max()),
                    "output_image": {
                        "type": "image_url",
                        "image_url": {"url": f"data:{output_media_type};base64,{output_image_base64}"}
                    },
                    "output_path": output_path
                },
                metadata={
                    "parameters": {
                        "diameter": diameter,
                        "flow_threshold": flow_threshold,
                        "cellprob_threshold": cellprob_threshold,
                        "min_size": min_size
                    }
                }
            )
        except Exception as log_error:
            print(f"Warning: Could not log output to Langfuse: {log_error}")
        
        # Return WITHOUT base64
        result = {
            "status": "success",
            "cell_count": int(masks_cellpose.max()),
            "output_path": output_path,
            "input_path": image_path,
            "parameters": {
                "diameter": diameter,
                "flow_threshold": flow_threshold,
                "cellprob_threshold": cellprob_threshold,
                "min_size": min_size
            },
            "summary": f"Detected {masks_cellpose.max()} cells. Output saved to: {output_path}",
            "next_step": "Call refine_cellpose_sam_segmentation to visually analyze the segmentation quality and decide if parameter adjustments are needed."
        }
        
        return json.dumps(result, indent=2)
        
    except Exception as e:
        return json.dumps({"error": f"Error during segmentation: {e}"})


@tool
def refine_cellpose_sam_segmentation(
    original_image_path: str = "",
    segmentation_output_path: str = "",
    current_parameters: dict = None,
    agent: Any = None,
) -> str:
    """
    Provides both original and segmented images to the VLM for visual quality assessment.
    The VLM will be able to see both images and provide informed analysis.
    
    Use this tool after run_cellpose_sam to check segmentation quality. The tool attaches
    both images to the current step so you can visually compare them.
    
    NOTE: If paths are empty or invalid, the tool will automatically use paths from context.
    
    Before calling, consider using search_knowledge_graph or hybrid_search to refresh
    your understanding of how cellpose parameters affect segmentation.
    
    Common issues and fixes:
    - Under-segmentation (cells merged): decrease flow_threshold or diameter
    - Over-segmentation (cells fragmented): increase flow_threshold or min_size
    - Too few cells: decrease cellprob_threshold or flow_threshold
    - Too many false positives: increase cellprob_threshold or min_size
    
    Args:
        original_image_path: Path to the original input image (optional - uses context)
        segmentation_output_path: Path to the segmented overlay image (optional - uses context)
        current_parameters: Dict with current diameter, flow_threshold, cellprob_threshold, min_size
        agent: The agent instance (passed automatically)
    
    Returns:
        str: JSON with guidance for VLM analysis (NO base64 images)
    """
    print(f"\n--- TOOL CALLED: refine_cellpose_sam_segmentation ---")
    print(f"Original image (provided): {original_image_path}")
    print(f"Segmented image (provided): {segmentation_output_path}")
    print(f"Current parameters: {current_parameters}")
    
    # Resolve paths from context if needed
    try:
        actual_original = resolve_image_path(original_image_path, context_type="image")
    except FileNotFoundError as e:
        return json.dumps({"error": f"Could not resolve original image: {e}"})
    
    try:
        actual_segmented = resolve_image_path(segmentation_output_path, context_type="output")
    except FileNotFoundError as e:
        return json.dumps({"error": f"Could not resolve segmented image: {e}"})
    
    original_image_path = actual_original
    segmentation_output_path = actual_segmented
    
    print(f"Resolved original: {original_image_path}")
    print(f"Resolved segmented: {segmentation_output_path}")
    
    try:
        # Load both images (for cache)
        original_b64, original_type = _get_cached_image(original_image_path) or _load_and_cache_image(original_image_path)
        segmented_b64, segmented_type = _get_cached_image(segmentation_output_path) or _load_and_cache_image(segmentation_output_path)
        
        # CRITICAL: Attach BOTH images to ActionStep so VLM can see them
        if agent is not None and hasattr(agent, 'memory') and hasattr(agent.memory, 'steps'):
            current_steps = [s for s in agent.memory.steps if isinstance(s, ActionStep)]
            if current_steps:
                current_step = current_steps[-1]
                
                # Load both as PIL Images
                original_img = Image.open(original_image_path).convert("RGB")
                segmented_img = Image.open(segmentation_output_path).convert("RGB")
                
                # CRITICAL: Use .copy() for both images
                current_step.observations_images = [original_img.copy(), segmented_img.copy()]
                print(f"✓ Attached both images to ActionStep for VLM comparison")
        
        # Get image dimensions for context
        original_img_array = np.array(Image.open(original_image_path).convert("RGB"))
        img_size = original_img_array.shape[0] * original_img_array.shape[1]
        
        # Log to Langfuse WITH both images
        try:
            langfuse.update_current_trace(
                input={
                    "tool": "refine_cellpose_sam_segmentation",
                    "original_image": {
                        "type": "image_url",
                        "image_url": {"url": f"data:{original_type};base64,{original_b64}"}
                    },
                    "segmented_image": {
                        "type": "image_url",
                        "image_url": {"url": f"data:{segmented_type};base64,{segmented_b64}"}
                    },
                    "current_parameters": current_parameters
                },
                metadata={
                    "original_path": original_image_path,
                    "segmented_path": segmentation_output_path
                }
            )
        except Exception as log_error:
            print(f"Warning: Could not log to Langfuse: {log_error}")
        
        # Return analysis guidance WITHOUT base64
        analysis = {
            "status": "ready_for_visual_analysis",
            "images_attached": "BOTH IMAGES NOW VISIBLE: The first image is the original input, "
                              "the second is the segmented overlay. Compare them visually to assess quality.",
            "image_paths": {
                "original": original_image_path,
                "segmented": segmentation_output_path
            },
            "current_parameters": current_parameters,
            "image_info": {
                "dimensions": f"{original_img_array.shape[1]}x{original_img_array.shape[0]}",
                "total_pixels": img_size
            },
            "visual_analysis_checklist": [
                "1. Do the colored masks accurately cover entire cells without extending beyond boundaries?",
                "2. Are neighboring cells properly separated, or are they merged together?",
                "3. Are there many small false positive detections (noise)?",
                "4. Are any large, obvious cells being missed completely?",
                "5. Overall quality assessment: excellent, good, needs_refinement, or poor?"
            ],
            "parameter_adjustment_guide": {
                "under_segmentation": {
                    "symptoms": "Masks don't reach cell edges, cells appear merged",
                    "solution": "Decrease flow_threshold by 0.1-0.2 OR decrease diameter by 10-20%"
                },
                "over_segmentation": {
                    "symptoms": "Masks extend past boundaries, cells fragmented into pieces",
                    "solution": "Increase flow_threshold by 0.1-0.2 OR increase min_size to 2-3x current value"
                },
                "too_few_cells": {
                    "symptoms": "Obvious cells in image are not being detected",
                    "solution": "Decrease cellprob_threshold by 1-2 OR decrease flow_threshold by 0.1-0.2"
                },
                "too_many_false_positives": {
                    "symptoms": "Many tiny spurious detections, background noise detected as cells",
                    "solution": "Increase cellprob_threshold by 1-2 OR increase min_size to 2-3x current value"
                }
            },
            "next_steps": {
                "if_good": "If segmentation looks accurate, inform the user of success and provide the output_path.",
                "if_needs_refinement": "Based on your visual analysis, adjust the appropriate parameters and call run_cellpose_sam again with the new values.",
                "important": "You can only call refine_cellpose_sam_segmentation AT MOST 2 TIMES total. If this is your second call, you must make a final decision."
            }
        }
        
        return json.dumps(analysis, indent=2)
        
    except Exception as e:
        error_result = {
            "status": "error",
            "error": str(e),
            "message": "Could not load images for refinement. Check that both file paths are valid."
        }
        return json.dumps(error_result, indent=2)