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

# Load image
path = '/Users/maxim.ziatdinov/Code/SciLink/analysis_session_20260331_212834/results/analysis__easy__AuNP_ImageAnalysis_20260331_212930_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()}]")

# Scale factor
scale_nm_per_px = 810.0 / 2048.0  # 0.3955 nm/px

# Run SAM instance segmentation
result = run_sam_analysis(image, params={
    "sam_parameters": "default",
    "min_area": 3000,
    "max_area": 50000,
    "pruning_iou_threshold": 0.5
})

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

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

# Get region properties
props = measure.regionprops(labeled, intensity_image=image)

# Extract properties for each particle
all_data = []
for prop in props:
    area = prop.area
    perimeter = prop.perimeter
    eq_diam_px = prop.equivalent_diameter
    eq_diam_nm = eq_diam_px * scale_nm_per_px
    if perimeter > 0:
        circularity = 4.0 * np.pi * area / (perimeter ** 2)
    else:
        circularity = 0.0
    centroid = prop.centroid  # (row, col)
    solidity = prop.solidity
    all_data.append({
        "label": prop.label,
        "area_px": area,
        "perimeter_px": perimeter,
        "eq_diam_px": eq_diam_px,
        "eq_diam_nm": eq_diam_nm,
        "circularity": circularity,
        "centroid_row": centroid[0],
        "centroid_col": centroid[1],
        "solidity": solidity
    })

print(f"Total particles with properties: {len(all_data)}")

# Filter out edge-truncated particles (centroid within 80px of border)
rows, cols = image.shape[:2]
border = 80
filtered = [d for d in all_data if
            d["centroid_row"] > border and d["centroid_row"] < (rows - border) and
            d["centroid_col"] > border and d["centroid_col"] < (cols - border)]
print(f"After edge filtering: {len(filtered)} particles")

# Post-filter: remove merged blobs and fragments based on median area
if len(filtered) > 0:
    areas = np.array([d["area_px"] for d in filtered])
    median_area = np.median(areas)
    print(f"Median area: {median_area:.1f} px")
    filtered = [d for d in filtered if
                d["area_px"] <= 2.5 * median_area and
                d["area_px"] >= 0.25 * median_area]
    print(f"After area filtering: {len(filtered)} particles")

# Compute statistics
if len(filtered) > 0:
    diameters = np.array([d["eq_diam_nm"] for d in filtered])
    circularities = np.array([d["circularity"] for d in filtered])
    solidities = np.array([d["solidity"] for d in filtered])
    areas_px = np.array([d["area_px"] for d in filtered])
    
    diam_mean = float(np.mean(diameters))
    diam_std = float(np.std(diameters))
    diam_min = float(np.min(diameters))
    diam_max = float(np.max(diameters))
    diam_median = float(np.median(diameters))
    
    circ_mean = float(np.mean(circularities))
    circ_std = float(np.std(circularities))
    circ_min = float(np.min(circularities))
    circ_max = float(np.max(circularities))
    circ_median = float(np.median(circularities))
    
    sol_mean = float(np.mean(solidities))
    sol_std = float(np.std(solidities))
else:
    diameters = np.array([])
    circularities = np.array([])
    solidities = np.array([])
    areas_px = np.array([])
    diam_mean = diam_std = diam_min = diam_max = diam_median = 0.0
    circ_mean = circ_std = circ_min = circ_max = circ_median = 0.0
    sol_mean = sol_std = 0.0

particle_count = len(filtered)
print(f"Final particle count: {particle_count}")
print(f"Diameter: {diam_mean:.2f} +/- {diam_std:.2f} nm")
print(f"Circularity: {circ_mean:.4f} +/- {circ_std:.4f}")

# Save label map
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')
ax.axis('off')

# 2. Segmentation overlay
ax = axes[0, 1]
# Normalize image for display
img_norm = (image - image.min()) / (image.max() - image.min() + 1e-10)
img_rgb = np.stack([img_norm, img_norm, img_norm], axis=-1)

# Create colored overlay
np.random.seed(42)
colors = np.random.rand(labeled.max() + 1, 3)
colors[0] = [0, 0, 0]  # background

overlay = img_rgb.copy()
for d in filtered:
    lbl = d["label"]
    mask = labeled == lbl
    color = colors[lbl]
    for c in range(3):
        overlay[:, :, c] = np.where(mask, 0.5 * img_rgb[:, :, c] + 0.5 * color[c], overlay[:, :, c])
    # Draw contour
    contour_mask = mask.astype(np.uint8)
    contours_cv, _ = cv2.findContours(contour_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
    for cnt in contours_cv:
        cv2.drawContours(overlay, [cnt], -1, color.tolist(), 2)

ax.imshow(overlay)
ax.set_title(f'Segmentation Overlay ({particle_count} particles)')
ax.axis('off')

# 3. Annotated image with particle outlines and IDs
ax = axes[0, 2]
ax.imshow(image, cmap='gray')
for d in filtered:
    lbl = d["label"]
    mask = (labeled == lbl).astype(np.uint8)
    contours_cv, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
    for cnt in contours_cv:
        pts = cnt.squeeze()
        if pts.ndim == 2:
            ax.plot(pts[:, 0], pts[:, 1], '-', color='lime', linewidth=0.8)
    ax.text(d["centroid_col"], d["centroid_row"], f"{d['eq_diam_nm']:.0f}",
            fontsize=5, color='yellow', ha='center', va='center')
ax.set_title('Detected Particles (diameter in nm)')
ax.axis('off')

# 4. Diameter histogram
ax = axes[1, 0]
if len(diameters) > 0:
    ax.hist(diameters, bins=20, color='steelblue', edgecolor='black', alpha=0.8)
    ax.axvline(diam_mean, color='red', linestyle='--', label=f'Mean={diam_mean:.1f} nm')
    ax.axvline(diam_median, color='orange', linestyle='--', label=f'Median={diam_median:.1f} nm')
    ax.legend(fontsize=9)
ax.set_xlabel('Equivalent Diameter (nm)')
ax.set_ylabel('Count')
ax.set_title('Diameter Distribution')

# 5. Circularity histogram
ax = axes[1, 1]
if len(circularities) > 0:
    ax.hist(circularities, bins=20, color='coral', edgecolor='black', alpha=0.8)
    ax.axvline(circ_mean, color='red', linestyle='--', label=f'Mean={circ_mean:.3f}')
    ax.legend(fontsize=9)
ax.set_xlabel('Circularity')
ax.set_ylabel('Count')
ax.set_title('Circularity Distribution')

# 6. Diameter vs Circularity scatter
ax = axes[1, 2]
if len(diameters) > 0 and len(circularities) > 0:
    sc = ax.scatter(diameters, circularities, c=solidities, cmap='viridis', 
                    alpha=0.7, edgecolors='black', linewidth=0.5, s=30)
    plt.colorbar(sc, ax=ax, label='Solidity')
ax.set_xlabel('Equivalent Diameter (nm)')
ax.set_ylabel('Circularity')
ax.set_title('Diameter vs Circularity')

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

# Build per-particle data for JSON
per_particle = []
for d in filtered:
    per_particle.append({
        "label": int(d["label"]),
        "area_px": int(d["area_px"]),
        "eq_diam_nm": round(d["eq_diam_nm"], 2),
        "circularity": round(d["circularity"], 4),
        "solidity": round(d["solidity"], 4),
        "centroid_row": round(d["centroid_row"], 1),
        "centroid_col": round(d["centroid_col"], 1)
    })

results = {
    "analysis_type": "Nanoparticle size and shape analysis using SAM instance segmentation",
    "extracted_features": {
        "particle_count": particle_count,
        "diameter_mean_nm": round(diam_mean, 2),
        "diameter_std_nm": round(diam_std, 2),
        "diameter_min_nm": round(diam_min, 2),
        "diameter_max_nm": round(diam_max, 2),
        "diameter_median_nm": round(diam_median, 2),
        "circularity_mean": round(circ_mean, 4),
        "circularity_std": round(circ_std, 4),
        "circularity_min": round(circ_min, 4),
        "circularity_max": round(circ_max, 4),
        "circularity_median": round(circ_median, 4),
        "solidity_mean": round(sol_mean, 4),
        "solidity_std": round(sol_std, 4),
        "scale_nm_per_px": round(scale_nm_per_px, 4),
        "per_particle_data": per_particle
    },
    "quality_metrics": {
        "sam_total_detections": result["total_count"],
        "particles_after_edge_filter": len([d for d in all_data if
            d["centroid_row"] > border and d["centroid_row"] < (rows - border) and
            d["centroid_col"] > border and d["centroid_col"] < (cols - border)]),
        "particles_after_area_filter": particle_count,
        "edge_border_px": border,
        "min_area_sam": 3000,
        "max_area_sam": 50000,
        "pruning_iou_threshold": 0.5
    },
    "summary": f"Detected {particle_count} nanoparticles using SAM instance segmentation. Mean equivalent diameter: {diam_mean:.1f} +/- {diam_std:.1f} nm. Mean circularity: {circ_mean:.3f} +/- {circ_std:.3f}. Scale: {scale_nm_per_px:.4f} nm/px (810 nm / 2048 px). Edge particles (centroid within 80px of border) and outlier areas (>2.5x or <0.25x median area) were filtered out.",
    "saved_arrays": {
        "analysis_labels.npy": {
            "description": f"Integer label map from SAM segmentation. {result['total_count']} particles labeled 1-{result['total_count']}, background=0. Note: post-filtering reduces to {particle_count} valid particles.",
            "shape": list(labeled.shape),
            "dtype": str(labeled.dtype)
        }
    }
}

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