#!/usr/bin/env python3
"""
Scout and Zoom Visualization Tool

This tool visualizes the scout-and-zoom process used in the anonymization pipeline:
1. Shows original image with bounding boxes from classify_pii
2. Shows cropped regions used for precise detection  
3. Shows segmentation masks constrained to their bounding box regions (no pixels outside bboxes)
4. Shows final mapped masks overlaid on the full image

The segmentation visualization ensures that only mask pixels within the corresponding
bounding box are displayed, providing clear debugging of bbox-guided detection results.

Usage:
    python visualize_scout_zoom.py <image_path> <bboxes_json> [--masks mask1.png mask2.png]
    python visualize_scout_zoom.py <image_path> <bboxes_json> --detect  # Use Grounded SAM 2 to generate masks

Example:
    python visualize_scout_zoom.py test_data/image.png '[{"bbox": [100, 200, 150, 100], "description": "text sign"}, ...]'
    python visualize_scout_zoom.py test_data/image.png '[{"bbox": [100, 200, 150, 100], "description": "text sign"}, ...]' --detect
"""

import json
import argparse
import os
import subprocess
import tempfile
import shutil
from pathlib import Path
from typing import List, Dict, Any, Optional
import matplotlib.pyplot as plt
import matplotlib.patches as patches
from PIL import Image, ImageDraw
import numpy as np
import sys

# Add project root to path for imports
sys.path.insert(0, str(Path(__file__).parent.parent))


def load_image(image_path: str) -> Image.Image:
    """Load image from path."""
    return Image.open(image_path)


def detect_with_grounded_sam2(image_path: str, bbox: List[int], description: str, 
                               output_dir: Path) -> Optional[str]:
    """Run Grounded SAM 2 detection on a cropped region.
    
    Args:
        image_path: Path to the full image
        bbox: Bounding box [x, y, w, h] for the region to detect
        description: Text description of what to detect
        output_dir: Directory to save temporary files
        
    Returns:
        Path to the generated mask file, or None if detection failed
    """
    # Get Grounded SAM 2 environment settings
    gsam2_dir = os.getenv("GROUNDED_SAM2_DIR", "models/Grounded-SAM-2")
    gsam2_conda_env = os.getenv("GROUNDED_SAM2_CONDA_ENV", "gsam2")
    
    # Get absolute path to project root
    project_root = Path(__file__).parent.parent.absolute()
    
    # Check if Grounded SAM 2 is available (use absolute path)
    gsam2_path = project_root / gsam2_dir if not Path(gsam2_dir).is_absolute() else Path(gsam2_dir)
    if not gsam2_path.exists():
        print(f"Warning: Grounded SAM 2 not found at {gsam2_path}", file=sys.stderr)
        return None
    
    inference_script = gsam2_path / "inference.py"
    if not inference_script.exists():
        print(f"Warning: inference.py not found at {inference_script}", file=sys.stderr)
        return None
    
    try:
        # Load and crop the image
        img = Image.open(image_path)
        x, y, w, h = bbox
        
        # Clamp coordinates to image bounds
        img_w, img_h = img.size
        x = max(0, min(x, img_w - 1))
        y = max(0, min(y, img_h - 1))
        w = min(w, img_w - x)
        h = min(h, img_h - y)
        
        # Crop the region
        crop = img.crop((x, y, x + w, y + h))
        
        # Save cropped region temporarily
        crop_path = output_dir / f"crop_{hash(description) & 0xffffffff:08x}.png"
        crop.save(crop_path)
        
        # Prepare detection prompt (ensure it ends with a dot)
        det_prompt = description.strip()
        if not det_prompt.endswith('.'):
            det_prompt += ' .'
        
        # Output mask path
        mask_path = output_dir / f"mask_{hash(description) & 0xffffffff:08x}.png"
        
        # Build command to run Grounded SAM 2 inference with absolute paths
        cmd = [
            "conda", "run", "-n", gsam2_conda_env,
            "python", str(inference_script.absolute()),
            "--image", str(crop_path.absolute()),
            "--text", det_prompt,
            "--out", str(mask_path.absolute())
        ]
        
        # Run the inference
        print(f"Running detection for: {description[:50]}...", flush=True)
        result = subprocess.run(
            cmd,
            capture_output=True,
            text=True,
            timeout=60
        )
        
        if result.returncode != 0:
            print(f"Warning: Detection failed for '{description}': {result.stderr}", file=sys.stderr)
            return None
        
        # Check if mask was created
        if not mask_path.exists():
            print(f"Warning: Mask file not created for '{description}'", file=sys.stderr)
            return None
        
        # Map the mask back to full image coordinates
        mapped_mask = Image.new('L', img.size, 0)
        crop_mask = Image.open(mask_path)
        mapped_mask.paste(crop_mask, (x, y))
        
        # Save the mapped mask
        mapped_mask_path = output_dir / f"mapped_{mask_path.name}"
        mapped_mask.save(mapped_mask_path)
        
        print(f"✓ Detection complete for: {description[:50]}...", flush=True)
        return str(mapped_mask_path)
        
    except subprocess.TimeoutExpired:
        print(f"Warning: Detection timed out for '{description}'", file=sys.stderr)
        return None
    except Exception as e:
        print(f"Warning: Detection error for '{description}': {e}", file=sys.stderr)
        return None


def detect_all_bboxes(image_path: str, bboxes: List[Dict[str, Any]], 
                      output_dir: Path) -> List[str]:
    """Run Grounded SAM 2 detection on all bounding boxes.
    
    Args:
        image_path: Path to the full image
        bboxes: List of bbox dictionaries with 'bbox' and 'description' keys
        output_dir: Directory to save masks
        
    Returns:
        List of mask file paths (in same order as bboxes)
    """
    masks = []
    temp_dir = output_dir / "detection_temp"
    temp_dir.mkdir(parents=True, exist_ok=True)
    
    for i, instance in enumerate(bboxes):
        bbox = instance.get('bbox', [])
        description = instance.get('description', f'Instance {i+1}')
        
        if len(bbox) != 4:
            print(f"Warning: Invalid bbox for instance {i+1}, skipping", file=sys.stderr)
            masks.append(None)
            continue
        
        mask_path = detect_with_grounded_sam2(image_path, bbox, description, temp_dir)
        masks.append(mask_path)
    
    return masks


def load_image(image_path: str) -> Image.Image:
    """Load image from path."""
    return Image.open(image_path)


def draw_bboxes_on_image(image: Image.Image, bboxes: List[Dict[str, Any]], output_path: str = None) -> Image.Image:
    """Draw bounding boxes on image with labels."""
    # Create a copy to draw on
    img_with_boxes = image.copy()
    draw = ImageDraw.Draw(img_with_boxes)

    # Colors for different bboxes
    colors = ['red', 'blue', 'green', 'orange', 'purple', 'brown', 'pink', 'gray', 'olive', 'cyan']

    for i, instance in enumerate(bboxes):
        bbox = instance.get('bbox', [])
        description = instance.get('description', f'Instance {i+1}')

        if len(bbox) != 4:
            continue

        x, y, w, h = bbox
        color = colors[i % len(colors)]

        # Draw rectangle
        draw.rectangle([x, y, x + w, y + h], outline=color, width=3)

        # Add label
        label = f"{i+1}: {description[:30]}{'...' if len(description) > 30 else ''}"
        # Draw text background
        text_bbox = draw.textbbox((x, y - 25), label, font=None)
        draw.rectangle(text_bbox, fill=color)
        # Draw text
        draw.text((x, y - 25), label, fill='white')

    if output_path:
        img_with_boxes.save(output_path)

    return img_with_boxes


def create_crop_visualization(image: Image.Image, bboxes: List[Dict[str, Any]], output_path: str = None) -> Image.Image:
    """Create a visualization showing all cropped regions in a grid."""
    if not bboxes:
        return image

    # Calculate grid dimensions
    n_crops = len(bboxes)
    cols = min(3, n_crops)
    rows = (n_crops + cols - 1) // cols

    # Create figure with subplots
    fig, axes = plt.subplots(rows, cols, figsize=(15, 5 * rows))

    # Handle different subplot configurations
    if rows == 1 and cols == 1:
        axes = [[axes]]
    elif rows == 1:
        axes = [axes]
    elif cols == 1:
        axes = [[ax] for ax in axes]

    # Flatten axes for easier iteration
    axes_flat = []
    for row in axes:
        if isinstance(row, (list, np.ndarray)):
            axes_flat.extend(row)
        else:
            axes_flat.append(row)

    for i, instance in enumerate(bboxes):
        if i >= len(axes_flat):
            break

        ax = axes_flat[i]
        bbox = instance.get('bbox', [])
        description = instance.get('description', f'Crop {i+1}')

        if len(bbox) == 4:
            x, y, w, h = bbox

            # Clamp coordinates to image bounds
            img_w, img_h = image.size
            x = max(0, min(x, img_w - 1))
            y = max(0, min(y, img_h - 1))
            w = min(w, img_w - x)
            h = min(h, img_h - y)

            # Crop the region
            crop = image.crop((x, y, x + w, y + h))

            # Display crop
            ax.imshow(crop)
            ax.set_title(f'Crop {i+1}: {description[:40]}{"..." if len(description) > 40 else ""}')
            ax.axis('off')

            # Add bbox info as text
            ax.text(0.02, 0.98, f'Bbox: [{x}, {y}, {w}, {h}]',
                   transform=ax.transAxes, fontsize=8, verticalalignment='top',
                   bbox=dict(boxstyle='round', facecolor='white', alpha=0.8))

    # Hide empty subplots
    for i in range(n_crops, len(axes_flat)):
        axes_flat[i].axis('off')

    plt.tight_layout()

    if output_path:
        plt.savefig(output_path, dpi=150, bbox_inches='tight')

    return fig


def create_mask_overlay_visualization(image: Image.Image, masks: List[str], bboxes: List[Dict[str, Any]] = None,
                                    output_path: str = None) -> Image.Image:
    """Create visualization showing masks overlaid on the image."""
    # Convert PIL image to numpy array for matplotlib
    img_array = np.array(image)

    fig, ax = plt.subplots(1, 1, figsize=(15, 10))
    ax.imshow(img_array)

    # Colors for different masks
    colors = ['red', 'blue', 'green', 'orange', 'purple', 'brown', 'pink', 'gray', 'olive', 'cyan']

    for i, mask_path in enumerate(masks):
        if not Path(mask_path).exists():
            continue

        try:
            # Load mask
            mask_img = Image.open(mask_path).convert('L')
            mask_array = np.array(mask_img) > 127  # Convert to binary mask

            # Create colored overlay
            color = colors[i % len(colors)]
            color_rgb = {
                'red': [255, 0, 0],
                'blue': [0, 0, 255],
                'green': [0, 255, 0],
                'orange': [255, 165, 0],
                'purple': [128, 0, 128],
                'brown': [165, 42, 42],
                'pink': [255, 192, 203],
                'gray': [128, 128, 128],
                'olive': [128, 128, 0],
                'cyan': [0, 255, 255]
            }[color]

            # Create overlay image
            overlay = np.zeros((*mask_array.shape, 4), dtype=np.uint8)
            overlay[mask_array, :3] = color_rgb
            overlay[mask_array, 3] = 128  # 50% transparency

            # Display overlay
            ax.imshow(overlay, alpha=0.5)

        except Exception as e:
            print(f"Error loading mask {mask_path}: {e}")
            continue

    # Add bounding boxes if provided
    if bboxes:
        for i, instance in enumerate(bboxes):
            bbox = instance.get('bbox', [])
            if len(bbox) == 4:
                x, y, w, h = bbox
                rect = patches.Rectangle((x, y), w, h,
                                       linewidth=2,
                                       edgecolor=colors[i % len(colors)],
                                       facecolor='none',
                                       linestyle='--')
                ax.add_patch(rect)

    ax.set_title('Image with Mask Overlays')
    ax.axis('off')

    if output_path:
        plt.savefig(output_path, dpi=150, bbox_inches='tight')

    return fig


def create_segmentation_visualization(masks: List[str], bboxes: List[Dict[str, Any]], output_path: str = None) -> Image.Image:
    """Create a visualization showing segmentation masks constrained to their bounding boxes.
    
    This shows the actual segmentation results from bbox-guided detection,
    ensuring that only mask pixels within each bounding box are displayed.
    No segmentation extends outside the bounding box boundaries.
    
    Args:
        masks: List of paths to mask images (should be bbox-constrained)
        bboxes: List of bbox dictionaries with 'bbox' and 'description' keys
        output_path: Optional path to save the visualization
        
    Returns:
        PIL Image of the segmentation visualization
    """
    if not masks or not bboxes:
        # Create empty figure if no data
        fig, ax = plt.subplots(1, 1, figsize=(10, 8))
        ax.text(0.5, 0.5, 'No masks or bboxes provided', ha='center', va='center', fontsize=16)
        ax.axis('off')
        if output_path:
            plt.savefig(output_path, dpi=150, bbox_inches='tight')
        return fig

    # Filter to only existing masks and match with bboxes
    valid_pairs = []
    for i, mask_path in enumerate(masks):
        if Path(mask_path).exists() and i < len(bboxes):
            valid_pairs.append((mask_path, bboxes[i]))

    if not valid_pairs:
        # Create empty figure if no valid pairs
        fig, ax = plt.subplots(1, 1, figsize=(10, 8))
        ax.text(0.5, 0.5, 'No valid mask-bbox pairs found', ha='center', va='center', fontsize=16)
        ax.axis('off')
        if output_path:
            plt.savefig(output_path, dpi=150, bbox_inches='tight')
        return fig

    # Calculate grid dimensions
    n_pairs = len(valid_pairs)
    cols = min(3, n_pairs)
    rows = (n_pairs + cols - 1) // cols

    # Create figure with subplots
    fig, axes = plt.subplots(rows, cols, figsize=(15, 5 * rows))

    # Handle different subplot configurations
    if rows == 1 and cols == 1:
        axes = [[axes]]
    elif rows == 1:
        axes = [axes]
    elif cols == 1:
        axes = [[ax] for ax in axes]

    # Flatten axes for easier iteration
    axes_flat = []
    for row in axes:
        if isinstance(row, (list, np.ndarray)):
            axes_flat.extend(row)
        else:
            axes_flat.append(row)

    for i, (mask_path, bbox_dict) in enumerate(valid_pairs):
        if i >= len(axes_flat):
            break

        ax = axes_flat[i]
        bbox = bbox_dict.get('bbox', [])
        description = bbox_dict.get('description', f'Instance {i+1}')

        if len(bbox) != 4:
            ax.text(0.5, 0.5, f'Invalid bbox: {bbox}', ha='center', va='center', fontsize=10)
            ax.axis('off')
            continue

        try:
            # Load mask
            mask_img = Image.open(mask_path).convert('L')
            mask_array = np.array(mask_img)

            # Crop mask to bounding box
            x, y, w, h = bbox
            # Ensure coordinates are within mask bounds
            mask_h, mask_w = mask_array.shape
            x1, y1 = max(0, x), max(0, y)
            x2, y2 = min(mask_w, x + w), min(mask_h, y + h)
            
            if x2 > x1 and y2 > y1:
                cropped_mask = mask_array[y1:y2, x1:x2]
                
                # Display cropped mask
                ax.imshow(cropped_mask, cmap='gray')
                ax.set_title(f'{description[:30]}{"..." if len(description) > 30 else ""}', fontsize=10)
                ax.axis('off')

                # Add mask info
                total_pixels = cropped_mask.size
                if total_pixels > 0:
                    nonzero_pixels = (cropped_mask > 127).sum()
                    coverage = nonzero_pixels / total_pixels * 100
                    ax.text(0.02, 0.98, f'Coverage: {coverage:.1f}%',
                           transform=ax.transAxes, fontsize=8, verticalalignment='top',
                           bbox=dict(boxstyle='round', facecolor='white', alpha=0.8))
            else:
                ax.text(0.5, 0.5, 'Bbox outside mask bounds', ha='center', va='center', fontsize=10)
                ax.axis('off')

        except Exception as e:
            ax.text(0.5, 0.5, f'Error loading mask:\n{Path(mask_path).name}',
                   ha='center', va='center', fontsize=10, transform=ax.transAxes)
            ax.axis('off')
            print(f"Error loading mask {mask_path}: {e}")

    # Hide empty subplots
    for i in range(n_pairs, len(axes_flat)):
        axes_flat[i].axis('off')

    plt.suptitle('Segmentation within Bounding Boxes', fontsize=14)
    plt.tight_layout()

    if output_path:
        plt.savefig(output_path, dpi=150, bbox_inches='tight')

    return fig


def create_comprehensive_visualization(image_path: str, bboxes: List[Dict[str, Any]],
                                     masks: List[str] = None, output_dir: str = None) -> Dict[str, str]:
    """Create comprehensive scout-and-zoom visualization."""
    image = load_image(image_path)
    output_files = {}

    if output_dir:
        output_dir = Path(output_dir)
        output_dir.mkdir(parents=True, exist_ok=True)
    else:
        # Use same directory as input image
        output_dir = Path(image_path).parent

    # 1. Original image with bounding boxes
    bbox_output = output_dir / f"{Path(image_path).stem}_bboxes.png"
    draw_bboxes_on_image(image, bboxes, str(bbox_output))
    output_files['bboxes'] = str(bbox_output)

    # 2. Crop visualization
    crop_output = output_dir / f"{Path(image_path).stem}_crops.png"
    create_crop_visualization(image, bboxes, str(crop_output))
    plt.close()  # Close the figure
    output_files['crops'] = str(crop_output)

    # 3. Segmentation masks (individual masks cropped to bboxes)
    if masks:
        segmentation_output = output_dir / f"{Path(image_path).stem}_segmentation.png"
        create_segmentation_visualization(masks, bboxes, str(segmentation_output))
        plt.close()  # Close the figure
        output_files['segmentation'] = str(segmentation_output)

    # 4. Mask overlay (if masks provided)
    if masks:
        mask_output = output_dir / f"{Path(image_path).stem}_masks.png"
        create_mask_overlay_visualization(image, masks, bboxes, str(mask_output))
        plt.close()  # Close the figure
        output_files['masks'] = str(mask_output)

    return output_files


def main():
    parser = argparse.ArgumentParser(description='Visualize scout-and-zoom process')
    parser.add_argument('image_path', help='Path to input image')
    parser.add_argument('bboxes_json', help='JSON string or file path containing bounding boxes')
    parser.add_argument('--masks', nargs='*', help='Paths to mask files to overlay')
    parser.add_argument('--detect', action='store_true', 
                       help='Use Grounded SAM 2 to detect and segment each bbox region')
    parser.add_argument('--output-dir', help='Output directory for visualizations')
    parser.add_argument('--show', action='store_true', help='Display plots instead of saving')

    args = parser.parse_args()

    # Load bounding boxes
    try:
        if args.bboxes_json.startswith('['):  # JSON string
            bboxes = json.loads(args.bboxes_json)
        else:  # File path
            with open(args.bboxes_json, 'r') as f:
                data = json.load(f)
                # Check if it's a full pipeline JSON with 'instances' key
                if isinstance(data, dict) and 'instances' in data:
                    bboxes = data['instances']
                else:
                    bboxes = data
    except Exception as e:
        print(f"Error loading bboxes: {e}")
        return

    # Determine output directory
    if args.output_dir:
        output_dir = Path(args.output_dir)
    else:
        # Use same directory as input image
        output_dir = Path(args.image_path).parent
    
    output_dir.mkdir(parents=True, exist_ok=True)

    # Run detection if requested
    masks = args.masks
    if args.detect:
        print("Running Grounded SAM 2 detection on all bounding boxes...")
        detected_masks = detect_all_bboxes(args.image_path, bboxes, output_dir)
        # Filter out None values
        masks = [m for m in detected_masks if m is not None]
        if not masks:
            print("Warning: No masks were generated from detection")
    
    # Create visualizations
    try:
        output_files = create_comprehensive_visualization(
            args.image_path,
            bboxes,
            masks,
            str(output_dir)
        )

        print("\nVisualization files created:")
        for key, path in output_files.items():
            print(f"  {key}: {path}")

        if args.show:
            plt.show()

    except Exception as e:
        print(f"Error creating visualizations: {e}")
        import traceback
        traceback.print_exc()


if __name__ == '__main__':
    main()