import numpy as np
import json
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from scipy import ndimage
from skimage import filters, morphology, measure, segmentation, feature
import cv2

# 1. Load image
path = '/Users/maxim.ziatdinov/Code/SciLink/test_image_annealing_data/output/image_0000/temp_image_0.npy'
image = np.load(path)
print(f"Image shape: {image.shape}, dtype: {image.dtype}, range: [{image.min()}, {image.max()}]")

# Ensure 2D
assert image.ndim == 2, f"Expected 2D image, got shape {image.shape}"

# Normalize to uint8 for processing
img_norm = ((image - image.min()) / (image.max() - image.min()) * 255).astype(np.uint8)

# 2. Gaussian blur (sigma=1.5) to suppress background noise
blurred = filters.gaussian(img_norm, sigma=1.5, preserve_range=True).astype(np.uint8)

# 3. Otsu thresholding
otsu_thresh = filters.threshold_otsu(blurred)
binary = blurred > otsu_thresh
print(f"Otsu threshold: {otsu_thresh}")

# 4. Morphological opening (disk r=2) to clean small noise artifacts
selem = morphology.disk(2)
binary_clean = morphology.binary_opening(binary, selem)

# 5. Distance transform
distance = ndimage.distance_transform_edt(binary_clean)

# 6. peak_local_max (min_distance=5) to find markers
# Use advanced guidance: for potentially elongated objects, use Sobel gradient as watershed landscape
coords = feature.peak_local_max(distance, min_distance=5, labels=binary_clean.astype(int))
print(f"Number of peaks found: {len(coords)}")

# Create marker image
markers = np.zeros(distance.shape, dtype=int)
for i, (r, c) in enumerate(coords):
    markers[r, c] = i + 1

# Dilate markers slightly so watershed has seeds
markers_labeled = ndimage.label(morphology.dilation(markers > 0, morphology.disk(1)))[0]
# Re-assign unique labels based on original peaks
markers_final = np.zeros_like(markers_labeled)
for i, (r, c) in enumerate(coords):
    label_val = markers_labeled[r, c]
    if label_val > 0:
        markers_final[markers_labeled == label_val] = i + 1

# 7. Watershed to split touching particles
# Use Sobel gradient as watershed landscape (advanced guidance)
sobel = filters.sobel(blurred.astype(float))
labels = segmentation.watershed(sobel, markers_final, mask=binary_clean)

# 8. Region property extraction
props = measure.regionprops(labels, intensity_image=image)

# 9. Filter regions by area (min_area=30, max_area=1500 pixels)
min_area = 30
max_area = 1500

filtered_props = [p for p in props if min_area <= p.area <= max_area]
print(f"Regions before filtering: {len(props)}, after filtering: {len(filtered_props)}")

# Build filtered label map
valid_labels = set(p.label for p in filtered_props)
filtered_labels = np.zeros_like(labels)
new_label = 0
label_mapping = {}
for p in filtered_props:
    new_label += 1
    label_mapping[p.label] = new_label

for old_label, new_label in label_mapping.items():
    filtered_labels[labels == old_label] = new_label

# 10. Extract features
particle_count = len(filtered_props)
particle_areas = [p.area for p in filtered_props]
equivalent_diameters = [p.equivalent_diameter for p in filtered_props]
centroids = [list(p.centroid) for p in filtered_props]

# Circularity = 4*pi*area / perimeter^2
circularities = []
for p in filtered_props:
    if p.perimeter > 0:
        circ = 4 * np.pi * p.area / (p.perimeter ** 2)
    else:
        circ = 0.0
    circularities.append(circ)

mean_radius = np.mean([d / 2.0 for d in equivalent_diameters]) if equivalent_diameters else 0.0

# Size distribution statistics
if particle_areas:
    area_arr = np.array(particle_areas)
    diam_arr = np.array(equivalent_diameters)
    size_stats = {
        "mean_area": float(np.mean(area_arr)),
        "std_area": float(np.std(area_arr)),
        "median_area": float(np.median(area_arr)),
        "min_area": float(np.min(area_arr)),
        "max_area": float(np.max(area_arr)),
        "mean_diameter": float(np.mean(diam_arr)),
        "std_diameter": float(np.std(diam_arr)),
        "median_diameter": float(np.median(diam_arr)),
        "mean_circularity": float(np.mean(circularities)),
        "std_circularity": float(np.std(circularities))
    }
else:
    size_stats = {}

# 11. Save label map
np.save("analysis_labels.npy", filtered_labels.astype(np.int32))

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

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

# Segmentation overlay
axes[1].imshow(image, cmap='gray')
# Create colored overlay from labels
if particle_count > 0:
    # Create random colormap for labels
    np.random.seed(42)
    colors = np.random.rand(particle_count + 1, 4)
    colors[:, 3] = 0.4  # semi-transparent
    colors[0] = [0, 0, 0, 0]  # background transparent
    
    overlay = np.zeros((*filtered_labels.shape, 4))
    for lbl in range(1, particle_count + 1):
        mask = filtered_labels == lbl
        overlay[mask] = colors[lbl]
    
    axes[1].imshow(overlay)
    
    # Draw contour boundaries
    for p in filtered_props:
        lbl_val = label_mapping[p.label]
        mask = (filtered_labels == lbl_val).astype(np.uint8)
        contours_cv, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
        for cnt in contours_cv:
            cnt_squeezed = cnt.squeeze()
            if cnt_squeezed.ndim == 2 and len(cnt_squeezed) > 2:
                axes[1].plot(cnt_squeezed[:, 0], cnt_squeezed[:, 1], 'r-', linewidth=0.8)

axes[1].set_title(f'Segmentation Overlay ({particle_count} particles)')
axes[1].axis('off')

# Distance transform with markers
axes[2].imshow(distance, cmap='hot')
if len(coords) > 0:
    axes[2].plot(coords[:, 1], coords[:, 0], 'c.', markersize=3)
axes[2].set_title('Distance Transform + Markers')
axes[2].axis('off')

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

# 13. Print results as JSON
results = {
    "analysis_type": "Particle segmentation using Otsu thresholding + watershed splitting with Sobel gradient landscape",
    "extracted_features": {
        "particle_count": particle_count,
        "particle_areas": particle_areas,
        "equivalent_diameters": [round(d, 2) for d in equivalent_diameters],
        "mean_radius": round(mean_radius, 2),
        "centroids": [[round(c[0], 2), round(c[1], 2)] for c in centroids],
        "circularity": [round(c, 4) for c in circularities],
        "size_distribution_statistics": {k: round(v, 2) for k, v in size_stats.items()}
    },
    "quality_metrics": {
        "otsu_threshold": float(otsu_thresh),
        "num_peaks_before_filtering": int(len(coords)),
        "num_regions_before_area_filter": len(props),
        "num_regions_after_area_filter": particle_count,
        "min_area_filter": min_area,
        "max_area_filter": max_area
    },
    "summary": f"Detected {particle_count} bright circular particles using Gaussian blur (sigma=1.5) + Otsu thresholding (threshold={otsu_thresh:.1f}) + morphological opening (disk r=2) + watershed with Sobel gradient landscape (advanced guidance) and peak_local_max markers (min_distance=5). Area filter: [{min_area}, {max_area}] pixels. Mean equivalent diameter: {round(np.mean(equivalent_diameters), 1) if equivalent_diameters else 0} px, mean circularity: {round(np.mean(circularities), 3) if circularities else 0}.",
    "saved_arrays": {
        "analysis_labels.npy": {
            "description": f"Integer label map, {particle_count} particles labeled 1-{particle_count}, background=0",
            "shape": list(filtered_labels.shape),
            "dtype": "int32"
        }
    }
}

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