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

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

# Convert int32 to uint8 via min-max normalization
img_min = float(image.min())
img_max = float(image.max())
if img_max > img_min:
    image_norm = ((image.astype(np.float64) - img_min) / (img_max - img_min) * 255).astype(np.uint8)
else:
    image_norm = np.zeros(image.shape, dtype=np.uint8)

print(f"Normalized image: shape={image_norm.shape}, dtype={image_norm.dtype}, range=[{image_norm.min()}, {image_norm.max()}]")

# Run SAM instance segmentation
# Expected particle diameter ~200px -> area ~31000 px²
# min_area=5000 filters small artifacts, max_area=200000 for largest particles
result = run_sam_analysis(image_norm, params={
    "sam_parameters": "default",
    "min_area": 5000,
    "max_area": 200000,
    "pruning_iou_threshold": 0.5
})

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

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

print(f"Labeled mask: unique labels = {len(np.unique(labeled)) - 1}")

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

# Scale factor: 811 nm / 2048 pixels
scale_nm_per_px = 811.0 / 2048.0  # ~0.396 nm/pixel
scale_nm2_per_px2 = scale_nm_per_px ** 2

# Extract features for each particle
particle_areas_px = []
particle_areas_nm2 = []
equivalent_diameters_nm = []
centroids = []
circularities = []
solidities = []

for prop in props:
    area_px = prop.area
    area_nm2 = area_px * scale_nm2_per_px2
    eq_diam_px = prop.equivalent_diameter
    eq_diam_nm = eq_diam_px * scale_nm_per_px
    centroid = prop.centroid  # (row, col)
    solidity = prop.solidity
    
    # Circularity = 4*pi*area / perimeter^2
    perimeter = prop.perimeter
    if perimeter > 0:
        circularity = 4.0 * np.pi * area_px / (perimeter ** 2)
    else:
        circularity = 0.0
    
    particle_areas_px.append(area_px)
    particle_areas_nm2.append(area_nm2)
    equivalent_diameters_nm.append(eq_diam_nm)
    centroids.append(list(centroid))
    circularities.append(circularity)
    solidities.append(solidity)

particle_count = len(props)
particle_areas_nm2 = np.array(particle_areas_nm2)
equivalent_diameters_nm = np.array(equivalent_diameters_nm)
circularities = np.array(circularities)
solidities = np.array(solidities)
centroids_arr = np.array(centroids) if centroids else np.empty((0, 2))

# Compute nearest neighbor distances (in nm)
nn_distances_nm = []
if particle_count > 1:
    centroid_positions = np.array(centroids)  # in pixels
    dist_matrix = distance.cdist(centroid_positions, centroid_positions)
    np.fill_diagonal(dist_matrix, np.inf)
    nn_distances_px = np.min(dist_matrix, axis=1)
    nn_distances_nm = (nn_distances_px * scale_nm_per_px).tolist()
else:
    nn_distances_nm = []

# Size distribution statistics
if particle_count > 0:
    size_mean = float(np.mean(equivalent_diameters_nm))
    size_std = float(np.std(equivalent_diameters_nm))
    size_median = float(np.median(equivalent_diameters_nm))
    area_mean = float(np.mean(particle_areas_nm2))
    area_std = float(np.std(particle_areas_nm2))
    circ_mean = float(np.mean(circularities))
    circ_std = float(np.std(circularities))
    sol_mean = float(np.mean(solidities))
    sol_std = float(np.std(solidities))
    nn_mean = float(np.mean(nn_distances_nm)) if nn_distances_nm else 0.0
    nn_std = float(np.std(nn_distances_nm)) if nn_distances_nm else 0.0
else:
    size_mean = size_std = size_median = 0.0
    area_mean = area_std = 0.0
    circ_mean = circ_std = 0.0
    sol_mean = sol_std = 0.0
    nn_mean = nn_std = 0.0

print(f"Particle count: {particle_count}")
print(f"Mean diameter: {size_mean:.2f} nm, Std: {size_std:.2f} nm")

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

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

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

# 2) Segmentation overlay
ax = axes[0, 1]
ax.imshow(image_norm, cmap='gray')
if particle_count > 0:
    # Create colored overlay
    overlay = np.zeros((*labeled.shape, 4), dtype=np.float32)
    np.random.seed(42)
    colors = np.random.rand(particle_count, 3)
    for i in range(particle_count):
        mask = labeled == (i + 1)
        overlay[mask, :3] = colors[i]
        overlay[mask, 3] = 0.4  # semi-transparent
    ax.imshow(overlay)
    # Draw contours
    for prop in props:
        contour_coords = skimage.measure.find_contours(labeled == prop.label, 0.5)
        for contour in contour_coords:
            ax.plot(contour[:, 1], contour[:, 0], linewidth=0.8, color='lime')
ax.set_title(f'Segmentation Overlay ({particle_count} particles)', fontsize=12)
ax.axis('off')

# 3) Labeled mask
ax = axes[0, 2]
if particle_count > 0:
    cmap = plt.cm.nipy_spectral
    display_labeled = labeled.copy().astype(float)
    display_labeled[display_labeled == 0] = np.nan
    ax.imshow(image_norm, cmap='gray', alpha=0.5)
    ax.imshow(display_labeled, cmap=cmap, alpha=0.6)
else:
    ax.imshow(image_norm, cmap='gray')
ax.set_title('Label Map', fontsize=12)
ax.axis('off')

# 4) Size distribution histogram
ax = axes[1, 0]
if particle_count > 0:
    ax.hist(equivalent_diameters_nm, bins=max(10, particle_count // 3), color='steelblue', edgecolor='black', alpha=0.7)
    ax.axvline(size_mean, color='red', linestyle='--', label=f'Mean={size_mean:.1f} nm')
    ax.axvline(size_median, color='orange', linestyle='--', label=f'Median={size_median:.1f} nm')
    ax.legend(fontsize=9)
ax.set_xlabel('Equivalent Diameter (nm)', fontsize=11)
ax.set_ylabel('Count', fontsize=11)
ax.set_title('Size Distribution', fontsize=12)

# 5) Circularity vs Solidity scatter
ax = axes[1, 1]
if particle_count > 0:
    sc = ax.scatter(circularities, solidities, c=equivalent_diameters_nm, cmap='viridis', alpha=0.7, edgecolors='black', linewidths=0.5, s=30)
    plt.colorbar(sc, ax=ax, label='Diameter (nm)')
ax.set_xlabel('Circularity', fontsize=11)
ax.set_ylabel('Solidity', fontsize=11)
ax.set_title('Shape Analysis', fontsize=12)
ax.set_xlim(0, 1.1)
ax.set_ylim(0, 1.1)

# 6) Nearest neighbor distances
ax = axes[1, 2]
if len(nn_distances_nm) > 0:
    ax.hist(nn_distances_nm, bins=max(10, particle_count // 3), color='coral', edgecolor='black', alpha=0.7)
    ax.axvline(nn_mean, color='red', linestyle='--', label=f'Mean={nn_mean:.1f} nm')
    ax.legend(fontsize=9)
ax.set_xlabel('Nearest Neighbor Distance (nm)', fontsize=11)
ax.set_ylabel('Count', fontsize=11)
ax.set_title('Nearest Neighbor Distances', fontsize=12)

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

# Prepare results JSON
results = {
    "analysis_type": "SAM instance segmentation of nanoparticles with size and shape analysis",
    "extracted_features": {
        "particle_count": particle_count,
        "scale_nm_per_pixel": round(scale_nm_per_px, 4),
        "particle_areas_nm2_mean": round(area_mean, 2),
        "particle_areas_nm2_std": round(area_std, 2),
        "equivalent_diameter_nm_mean": round(size_mean, 2),
        "equivalent_diameter_nm_std": round(size_std, 2),
        "equivalent_diameter_nm_median": round(size_median, 2),
        "equivalent_diameter_nm_min": round(float(np.min(equivalent_diameters_nm)), 2) if particle_count > 0 else 0.0,
        "equivalent_diameter_nm_max": round(float(np.max(equivalent_diameters_nm)), 2) if particle_count > 0 else 0.0,
        "circularity_mean": round(circ_mean, 4),
        "circularity_std": round(circ_std, 4),
        "solidity_mean": round(sol_mean, 4),
        "solidity_std": round(sol_std, 4),
        "nearest_neighbor_distance_nm_mean": round(nn_mean, 2),
        "nearest_neighbor_distance_nm_std": round(nn_std, 2),
        "individual_diameters_nm": [round(d, 2) for d in equivalent_diameters_nm.tolist()] if particle_count > 0 else [],
        "individual_areas_nm2": [round(a, 2) for a in particle_areas_nm2.tolist()] if particle_count > 0 else [],
        "individual_circularities": [round(c, 4) for c in circularities.tolist()] if particle_count > 0 else [],
        "individual_solidities": [round(s, 4) for s in solidities.tolist()] if particle_count > 0 else [],
        "centroid_positions_pixels": [[round(c[0], 1), round(c[1], 1)] for c in centroids],
        "nearest_neighbor_distances_nm": [round(d, 2) for d in nn_distances_nm]
    },
    "quality_metrics": {
        "sam_total_detected": result.get("total_count", 0),
        "particles_after_filtering": particle_count,
        "image_intensity_range": [int(image.min()), int(image.max())],
        "min_area_filter_px": 5000,
        "max_area_filter_px": 200000,
        "pruning_iou_threshold": 0.5
    },
    "summary": f"SAM instance segmentation detected {particle_count} nanoparticles. Mean equivalent diameter: {size_mean:.1f} +/- {size_std:.1f} nm (median: {size_median:.1f} nm). Mean circularity: {circ_mean:.3f}, mean solidity: {sol_mean:.3f}. Scale: {scale_nm_per_px:.4f} nm/pixel (811 nm / 2048 px). Parameters: min_area=5000, max_area=200000, pruning_iou=0.5.",
    "saved_arrays": {
        "analysis_labels.npy": {
            "description": f"Integer label map from SAM segmentation, {particle_count} particles labeled 1-{particle_count}, background=0",
            "shape": list(labeled.shape),
            "dtype": str(labeled.dtype)
        }
    }
}

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