import numpy as np
import json
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from scipy.spatial.distance import pdist, squareform
from scipy.signal import find_peaks
import skimage.measure
from scilink.tools.sam import run_sam_analysis
import cv2

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

# 2. Prepare image for SAM
# Particles are dark on light background - invert so particles are bright
image_inverted = image.max() - image

# Normalize to 0-255 uint8 for SAM
img_norm = ((image_inverted - image_inverted.min()) / (image_inverted.max() - image_inverted.min()) * 255).astype(np.uint8)

print("Running SAM instance segmentation...")
result = run_sam_analysis(img_norm, params={
    "sam_parameters": "default",
    "min_area": 5000,
    "max_area": 80000,
    "pruning_iou_threshold": 0.5
})

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

# 3. Build labeled mask from SAM particles
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

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

# Extract centroids, areas, equivalent diameters
centroids = []
areas = []
equiv_diameters = []
for prop in props:
    centroids.append(prop.centroid)  # (row, col)
    areas.append(prop.area)
    equiv_diameters.append(prop.equivalent_diameter)

centroids = np.array(centroids)  # shape (N, 2)
areas = np.array(areas)
equiv_diameters = np.array(equiv_diameters)

particle_count = len(centroids)
print(f"Extracted {particle_count} particles with properties")

if particle_count < 2:
    print("Too few particles for graph analysis")
    # Save minimal results
    np.save("analysis_labels.npy", labeled)
    results = {
        "analysis_type": "Nanoparticle instance segmentation and connectivity graph analysis using SAM",
        "extracted_features": {
            "particle_count": int(particle_count),
        },
        "quality_metrics": {},
        "summary": f"Only {particle_count} particle(s) detected, insufficient for graph analysis.",
        "saved_arrays": {
            "analysis_labels.npy": {
                "description": f"Integer label map, {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)}")
else:
    # 5. Compute pairwise center-to-center distances
    # centroids are in (row, col) = (y, x) format
    dist_matrix = squareform(pdist(centroids))
    
    # 6. Determine connectivity threshold
    # For each particle, find nearest neighbor distance
    nn_distances = []
    for i in range(particle_count):
        dists = dist_matrix[i].copy()
        dists[i] = np.inf  # exclude self
        nn_distances.append(dists.min())
    nn_distances = np.array(nn_distances)
    
    median_diameter = np.median(equiv_diameters)
    mean_nn_dist = np.mean(nn_distances)
    median_nn_dist = np.median(nn_distances)
    
    print(f"Median equivalent diameter: {median_diameter:.1f} px")
    print(f"Mean nearest-neighbor distance: {mean_nn_dist:.1f} px")
    print(f"Median nearest-neighbor distance: {median_nn_dist:.1f} px")
    
    # Connectivity threshold: touching particles have center-to-center distance ~ 1 diameter
    # Use 1.5x median diameter as threshold (valley between touching and distant)
    # Also consider the nearest-neighbor distribution
    
    # Try to find bimodal valley in pairwise distance distribution
    all_pairwise = pdist(centroids)
    
    # Build histogram of pairwise distances
    hist_range = (0, np.percentile(all_pairwise, 50))  # focus on closer pairs
    hist_counts, hist_edges = np.histogram(all_pairwise, bins=100, range=hist_range)
    hist_centers = (hist_edges[:-1] + hist_edges[1:]) / 2
    
    # Smooth histogram to find valley
    from scipy.ndimage import gaussian_filter1d
    smoothed = gaussian_filter1d(hist_counts.astype(float), sigma=3)
    
    # Find peaks in smoothed histogram
    peaks, peak_props = find_peaks(smoothed, height=smoothed.max() * 0.1, distance=5)
    
    # Default threshold: 1.5x median diameter
    connectivity_threshold = 1.5 * median_diameter
    threshold_method = "1.5x median diameter"
    
    # If we find two clear peaks, use the valley between them
    if len(peaks) >= 2:
        # Find valley between first two peaks
        p1, p2 = peaks[0], peaks[1]
        valley_region = smoothed[p1:p2+1]
        valley_idx = p1 + np.argmin(valley_region)
        valley_threshold = hist_centers[valley_idx]
        
        # Only use valley if it makes physical sense (between 1x and 2.5x median diameter)
        if 0.8 * median_diameter < valley_threshold < 2.5 * median_diameter:
            connectivity_threshold = valley_threshold
            threshold_method = f"valley between histogram peaks at {hist_centers[peaks[0]]:.1f} and {hist_centers[peaks[1]]:.1f} px"
    
    print(f"Connectivity threshold: {connectivity_threshold:.1f} px (method: {threshold_method})")
    
    # 7. Build adjacency graph
    adjacency_matrix = (dist_matrix < connectivity_threshold).astype(int)
    np.fill_diagonal(adjacency_matrix, 0)  # no self-loops
    
    # Graph statistics
    degrees = adjacency_matrix.sum(axis=1)
    num_edges = adjacency_matrix.sum() // 2  # undirected
    
    # Degree distribution
    unique_degrees, degree_counts = np.unique(degrees, return_counts=True)
    degree_distribution = {int(d): int(c) for d, c in zip(unique_degrees, degree_counts)}
    
    # Connected components using BFS
    visited = set()
    components = []
    for node in range(particle_count):
        if node not in visited:
            # BFS
            component = []
            queue = [node]
            while queue:
                n = queue.pop(0)
                if n not in visited:
                    visited.add(n)
                    component.append(n)
                    neighbors = np.where(adjacency_matrix[n] > 0)[0]
                    for nb in neighbors:
                        if nb not in visited:
                            queue.append(nb)
            components.append(component)
    
    num_connected_components = len(components)
    component_sizes = [len(c) for c in components]
    
    print(f"Number of edges: {num_edges}")
    print(f"Number of connected components: {num_connected_components}")
    print(f"Degree distribution: {degree_distribution}")
    print(f"Mean degree: {degrees.mean():.2f}")
    
    # 8. Save arrays
    np.save("analysis_labels.npy", labeled)
    np.save("centroids.npy", centroids)
    np.save("adjacency_matrix.npy", adjacency_matrix)
    np.save("pairwise_distances.npy", dist_matrix)
    
    # 9. Visualization
    fig, axes = plt.subplots(2, 3, figsize=(18, 12), dpi=100)
    
    # (a) Original image
    ax = axes[0, 0]
    ax.imshow(image, cmap='gray')
    ax.set_title('Original Image')
    ax.axis('off')
    
    # (b) Segmentation overlay
    ax = axes[0, 1]
    ax.imshow(image, cmap='gray')
    # Create colored overlay
    from matplotlib.colors import ListedColormap
    import matplotlib.cm as cm
    
    # Draw colored semi-transparent masks
    overlay = np.zeros((*image.shape[:2], 4), dtype=float)
    cmap = plt.cm.tab20
    for i in range(particle_count):
        mask = labeled == (i + 1)
        color = cmap(i % 20)
        overlay[mask, 0] = color[0]
        overlay[mask, 1] = color[1]
        overlay[mask, 2] = color[2]
        overlay[mask, 3] = 0.35
    ax.imshow(overlay)
    
    # Draw contours
    for i in range(particle_count):
        mask = (labeled == (i + 1)).astype(np.uint8)
        contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        for cnt in contours:
            cnt_squeezed = cnt.squeeze()
            if cnt_squeezed.ndim == 2 and len(cnt_squeezed) > 2:
                ax.plot(cnt_squeezed[:, 0], cnt_squeezed[:, 1], '-', 
                       color=cmap(i % 20), linewidth=0.5)
    
    ax.set_title(f'Segmentation Overlay ({particle_count} particles)')
    ax.axis('off')
    
    # (c) Graph overlay on original image
    ax = axes[0, 2]
    ax.imshow(image, cmap='gray')
    
    # Draw edges
    for i in range(particle_count):
        for j in range(i+1, particle_count):
            if adjacency_matrix[i, j] > 0:
                y1, x1 = centroids[i]
                y2, x2 = centroids[j]
                ax.plot([x1, x2], [y1, y2], 'c-', linewidth=0.8, alpha=0.7)
    
    # Draw nodes colored by degree
    scatter = ax.scatter(centroids[:, 1], centroids[:, 0], c=degrees, 
                        cmap='hot', s=20, edgecolors='white', linewidths=0.3, 
                        zorder=5)
    plt.colorbar(scatter, ax=ax, label='Degree', shrink=0.7)
    ax.set_title(f'Connectivity Graph ({num_edges} edges)')
    ax.axis('off')
    
    # (d) Nearest-neighbor distance histogram
    ax = axes[1, 0]
    ax.hist(nn_distances, bins=30, color='steelblue', edgecolor='black', alpha=0.7)
    ax.axvline(connectivity_threshold, color='red', linestyle='--', linewidth=2, 
              label=f'Threshold = {connectivity_threshold:.1f} px')
    ax.axvline(median_diameter, color='orange', linestyle=':', linewidth=2, 
              label=f'Median diameter = {median_diameter:.1f} px')
    ax.set_xlabel('Nearest-neighbor distance (px)')
    ax.set_ylabel('Count')
    ax.set_title('Nearest-Neighbor Distance Distribution')
    ax.legend(fontsize=8)
    
    # (e) Degree distribution
    ax = axes[1, 1]
    ax.bar(unique_degrees, degree_counts, color='coral', edgecolor='black')
    ax.set_xlabel('Degree (number of neighbors)')
    ax.set_ylabel('Count')
    ax.set_title(f'Degree Distribution (mean={degrees.mean():.1f})')
    ax.set_xticks(unique_degrees)
    
    # (f) Particle size distribution
    ax = axes[1, 2]
    ax.hist(equiv_diameters, bins=30, color='mediumseagreen', edgecolor='black', alpha=0.7)
    ax.set_xlabel('Equivalent diameter (px)')
    ax.set_ylabel('Count')
    ax.set_title(f'Particle Size Distribution')
    ax.axvline(median_diameter, color='red', linestyle='--', 
              label=f'Median = {median_diameter:.1f} px')
    ax.legend(fontsize=8)
    
    plt.tight_layout()
    plt.savefig('analysis_visualization.png', dpi=100, bbox_inches='tight')
    plt.close()
    print("Visualization saved.")
    
    # 10. Prepare JSON results
    # Convert centroid positions to list format
    centroid_list = centroids.tolist()
    
    # Particle areas and diameters
    particle_size_stats = {
        "mean_area_px": float(np.mean(areas)),
        "std_area_px": float(np.std(areas)),
        "median_area_px": float(np.median(areas)),
        "mean_equiv_diameter_px": float(np.mean(equiv_diameters)),
        "std_equiv_diameter_px": float(np.std(equiv_diameters)),
        "median_equiv_diameter_px": float(np.median(equiv_diameters)),
        "min_equiv_diameter_px": float(np.min(equiv_diameters)),
        "max_equiv_diameter_px": float(np.max(equiv_diameters))
    }
    
    results = {
        "analysis_type": "Nanoparticle instance segmentation (SAM) and connectivity graph analysis",
        "extracted_features": {
            "particle_count": int(particle_count),
            "particle_size_statistics": particle_size_stats,
            "connectivity_threshold_px": float(connectivity_threshold),
            "connectivity_threshold_method": threshold_method,
            "number_of_edges": int(num_edges),
            "number_of_connected_components": int(num_connected_components),
            "component_sizes": [int(s) for s in sorted(component_sizes, reverse=True)],
            "degree_distribution": degree_distribution,
            "mean_degree": float(degrees.mean()),
            "max_degree": int(degrees.max()),
            "mean_nearest_neighbor_distance_px": float(mean_nn_dist),
            "median_nearest_neighbor_distance_px": float(median_nn_dist),
            "std_nearest_neighbor_distance_px": float(np.std(nn_distances)),
            "mean_pairwise_distance_px": float(np.mean(all_pairwise)),
        },
        "quality_metrics": {
            "sam_total_detected": int(total_count),
            "particles_after_filtering": int(particle_count),
            "image_shape": list(image.shape),
            "intensity_range": [float(image.min()), float(image.max())]
        },
        "summary": f"Detected {particle_count} nanoparticles via SAM instance segmentation (image inverted for bright-on-dark). Built connectivity graph with threshold {connectivity_threshold:.1f} px ({threshold_method}): {num_edges} edges, {num_connected_components} connected components, mean degree {degrees.mean():.1f}. Median particle equivalent diameter: {median_diameter:.1f} px, mean nearest-neighbor distance: {mean_nn_dist:.1f} px.",
        "saved_arrays": {
            "analysis_labels.npy": {
                "description": f"Integer label map, {particle_count} particles labeled 1-{particle_count}, background=0",
                "shape": list(labeled.shape),
                "dtype": str(labeled.dtype)
            },
            "centroids.npy": {
                "description": f"Centroid coordinates (row, col) for {particle_count} particles",
                "shape": list(centroids.shape),
                "dtype": str(centroids.dtype)
            },
            "adjacency_matrix.npy": {
                "description": f"Binary adjacency matrix ({particle_count}x{particle_count}) for connectivity graph, threshold={connectivity_threshold:.1f} px",
                "shape": list(adjacency_matrix.shape),
                "dtype": str(adjacency_matrix.dtype)
            },
            "pairwise_distances.npy": {
                "description": f"Pairwise center-to-center distance matrix ({particle_count}x{particle_count}) in pixels",
                "shape": list(dist_matrix.shape),
                "dtype": str(dist_matrix.dtype)
            }
        }
    }
    
    print(f"IMAGE_ANALYSIS_RESULTS_JSON:{json.dumps(results)}")
