import numpy as np
import json
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import skimage.measure
from scilink.tools.sam import run_sam_analysis

# Load image
image_path = '/Users/maxim.ziatdinov/Code/SciLink/test_image_annealing_data/output_high_threshold/image_0000/temp_image_0.npy'
image = np.load(image_path)

# Verify shape and dtype
print(f"Image shape: {image.shape}, dtype: {image.dtype}")
print(f"Intensity range: [{image.min()}, {image.max()}]")

# Ensure 2D grayscale
if image.ndim == 3:
    image_2d = image[:, :, 0]
else:
    image_2d = image

# Convert to uint8 for SAM (normalize to 0-255)
image_norm = ((image_2d - image_2d.min()) / (image_2d.max() - image_2d.min()) * 255).astype(np.uint8)

# Run SAM instance segmentation with default parameters as specified in the plan
result = run_sam_analysis(image_norm, params={
    "sam_parameters": "default",
    "min_area": 30,
    "max_area": 2000,
    "pruning_iou_threshold": 0.5
})

particles = result["particles"]
total_count = result["total_count"]
print(f"SAM detected {total_count} particles")

# Build labeled mask from SAM particles
labeled = np.zeros(image_2d.shape[:2], dtype=np.int32)
for i, p in enumerate(particles):
    mask = np.array(p["mask"], dtype=bool)
    labeled[mask] = i + 1

# Extract region properties
props = skimage.measure.regionprops(labeled, intensity_image=image_2d)

# Extract per-particle features
particle_areas = []
particle_radii = []
particle_centroids = []

for prop in props:
    area = prop.area
    equiv_radius = np.sqrt(area / np.pi)
    centroid = prop.centroid  # (row, col)
    particle_areas.append(float(area))
    particle_radii.append(float(equiv_radius))
    particle_centroids.append([float(centroid[0]), float(centroid[1])])

particle_areas = np.array(particle_areas)
particle_radii = np.array(particle_radii)

# Compute statistics
particle_count = len(props)
mean_radius = float(np.mean(particle_radii)) if particle_count > 0 else 0.0
std_radius = float(np.std(particle_radii)) if particle_count > 0 else 0.0
mean_area = float(np.mean(particle_areas)) if particle_count > 0 else 0.0
std_area = float(np.std(particle_areas)) if particle_count > 0 else 0.0
median_radius = float(np.median(particle_radii)) if particle_count > 0 else 0.0
min_radius = float(np.min(particle_radii)) if particle_count > 0 else 0.0
max_radius = float(np.max(particle_radii)) if particle_count > 0 else 0.0

# Size distribution histogram
if particle_count > 0:
    n_bins = min(20, max(5, particle_count // 3))
    hist_counts, hist_edges = np.histogram(particle_radii, bins=n_bins)
    hist_counts_list = hist_counts.tolist()
    hist_edges_list = hist_edges.tolist()
else:
    hist_counts_list = []
    hist_edges_list = []

# Save labeled mask
np.save("analysis_labels.npy", labeled)

# Save per-particle data
per_particle_data = np.column_stack([particle_areas, particle_radii]) if particle_count > 0 else np.array([])
if particle_count > 0:
    np.save("particle_properties.npy", per_particle_data)

# Visualization
fig, axes = plt.subplots(1, 3, figsize=(18, 6), dpi=100)

# Subplot 1: Original image
axes[0].imshow(image_2d, cmap='gray')
axes[0].set_title('Original Image')
axes[0].axis('off')

# Subplot 2: Segmentation overlay
axes[1].imshow(image_2d, cmap='gray')
if particle_count > 0:
    # Create colored overlay
    from matplotlib.colors import ListedColormap
    # Generate random colors for each particle
    np.random.seed(42)
    colors = np.random.rand(particle_count, 3)
    overlay = np.zeros((*image_2d.shape, 4), dtype=np.float32)
    for i, prop in enumerate(props):
        mask_i = labeled == prop.label
        overlay[mask_i, 0] = colors[i, 0]
        overlay[mask_i, 1] = colors[i, 1]
        overlay[mask_i, 2] = colors[i, 2]
        overlay[mask_i, 3] = 0.4  # semi-transparent
    axes[1].imshow(overlay)
    # Draw contour boundaries
    for prop in props:
        mask_i = (labeled == prop.label).astype(np.uint8)
        contours_list = skimage.measure.find_contours(mask_i, 0.5)
        for contour in contours_list:
            axes[1].plot(contour[:, 1], contour[:, 0], 'r-', linewidth=0.5)
axes[1].set_title(f'Segmentation Overlay ({particle_count} particles)')
axes[1].axis('off')

# Subplot 3: Size distribution histogram
if particle_count > 0:
    axes[2].hist(particle_radii, bins=n_bins, color='steelblue', edgecolor='black', alpha=0.7)
    axes[2].axvline(mean_radius, color='red', linestyle='--', label=f'Mean={mean_radius:.1f}')
    axes[2].axvline(median_radius, color='green', linestyle='--', label=f'Median={median_radius:.1f}')
    axes[2].legend()
axes[2].set_xlabel('Equivalent Radius (pixels)')
axes[2].set_ylabel('Count')
axes[2].set_title('Particle Size Distribution')

plt.tight_layout()
plt.savefig('analysis_visualization.png', dpi=100, bbox_inches='tight')
plt.close()

# Prepare results
results = {
    "analysis_type": "SAM instance segmentation for circular particle detection and size analysis",
    "extracted_features": {
        "particle_count": particle_count,
        "per_particle_area_pixels": particle_areas.tolist() if particle_count > 0 else [],
        "per_particle_equivalent_radius_pixels": particle_radii.tolist() if particle_count > 0 else [],
        "mean_radius": mean_radius,
        "std_radius": std_radius,
        "median_radius": median_radius,
        "min_radius": min_radius,
        "max_radius": max_radius,
        "mean_area": mean_area,
        "std_area": std_area,
        "particle_centroids": particle_centroids,
        "size_distribution_histogram": {
            "counts": hist_counts_list,
            "bin_edges": hist_edges_list
        }
    },
    "quality_metrics": {
        "total_sam_detections": total_count,
        "particles_after_filtering": particle_count,
        "image_coverage_fraction": float(np.sum(labeled > 0) / labeled.size) if particle_count > 0 else 0.0
    },
    "summary": f"SAM instance segmentation detected {particle_count} circular particles with mean equivalent radius {mean_radius:.2f} +/- {std_radius:.2f} pixels. Parameters: min_area=30, max_area=2000, pruning_iou_threshold=0.5 (plan defaults).",
    "saved_arrays": {
        "analysis_labels.npy": {
            "description": f"Integer label map, {particle_count} particles labeled 1-{particle_count}, background=0",
            "shape": list(labeled.shape),
            "dtype": "int32"
        },
        "particle_properties.npy": {
            "description": "Per-particle properties: column 0 = area (pixels), column 1 = equivalent radius (pixels)",
            "shape": list(per_particle_data.shape) if particle_count > 0 else [0],
            "dtype": "float64"
        }
    }
}

print(f"IMAGE_ANALYSIS_RESULTS_JSON:{json.dumps(results)}")
