import numpy as np
import json
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from skimage import measure
from scipy.spatial import Delaunay
from scipy import ndimage
import warnings
import os
import glob
warnings.filterwarnings('ignore')

# Step 1: Load the label map and original image
# Determine the correct directory - search for files
base_dir = '/Users/maxim.ziatdinov/Code/SciLink/analysis_session_20260331_173929/results/analysis_nanoparticles_ImageAnalysis_20260331_174128_001/image_0000'
session_dir = '/Users/maxim.ziatdinov/Code/SciLink/analysis_session_20260331_173929'
cwd = os.getcwd()

def find_file(filename, search_roots):
    """Search for a file by name in multiple root directories."""
    for root_dir in search_roots:
        if not os.path.isdir(root_dir):
            continue
        for dirpath, dirnames, filenames in os.walk(root_dir):
            if filename in filenames:
                return os.path.join(dirpath, filename)
    return None

search_roots = [cwd, base_dir, session_dir]

label_map_path = os.path.join(base_dir, 'analysis_labels.npy')
if not os.path.isfile(label_map_path):
    found = find_file('analysis_labels.npy', search_roots)
    if found:
        label_map_path = found
        base_dir = os.path.dirname(found)
    else:
        # List available .npy files for diagnostics and try alternative names
        all_npy = []
        for root_dir in search_roots:
            if os.path.isdir(root_dir):
                for dirpath, dirnames, filenames in os.walk(root_dir):
                    for f in filenames:
                        if f.endswith('.npy'):
                            all_npy.append(os.path.join(dirpath, f))
        # Try to find any label file
        label_candidates = [p for p in all_npy if 'label' in os.path.basename(p).lower()]
        if label_candidates:
            label_map_path = label_candidates[0]
            base_dir = os.path.dirname(label_map_path)
        else:
            raise FileNotFoundError(f"Cannot find analysis_labels.npy. Available .npy files: {all_npy}")

image_path = os.path.join(base_dir, 'temp_image_0.npy')
if not os.path.isfile(image_path):
    found = find_file('temp_image_0.npy', search_roots)
    if found:
        image_path = found
    else:
        # Try alternative image file names
        all_npy = []
        for root_dir in search_roots:
            if os.path.isdir(root_dir):
                for dirpath, dirnames, filenames in os.walk(root_dir):
                    for f in filenames:
                        if f.endswith('.npy'):
                            all_npy.append(os.path.join(dirpath, f))
        image_candidates = [p for p in all_npy if 'image' in os.path.basename(p).lower() or 'temp' in os.path.basename(p).lower()]
        if image_candidates:
            image_path = image_candidates[0]
        else:
            raise FileNotFoundError(f"Cannot find temp_image_0.npy. Available .npy files: {all_npy}")

print(f"Loading label map from: {label_map_path}")
print(f"Loading image from: {image_path}")
print(f"Output base dir: {base_dir}")
print(f"Current working directory: {cwd}")

label_map = np.load(label_map_path)
image = np.load(image_path)

# Normalize image for display
img_display = image.astype(float)
img_display = (img_display - img_display.min()) / (img_display.max() - img_display.min())

# Step 2: Compute regionprops for all labeled regions
props = measure.regionprops(label_map, intensity_image=image)
print(f"Total labeled regions: {len(props)}")

# We need a pixel-to-nm conversion. For ~75 nm diameter gold nanoparticles in a 2048x2048 image,
# we estimate from the median equivalent diameter in pixels.
# First pass: get all equivalent diameters
all_eq_diams_px = [p.equivalent_diameter for p in props]
median_diam_px = np.median(all_eq_diams_px)
print(f"Median equivalent diameter (pixels): {median_diam_px:.1f}")

# Assume the nanoparticles are ~75 nm in diameter (typical for gold NPs in these images)
# This gives us the pixel size
assumed_particle_diameter_nm = 75.0
pixel_size_nm = assumed_particle_diameter_nm / median_diam_px
print(f"Estimated pixel size: {pixel_size_nm:.3f} nm/pixel")

# Step 3: Filter particles
image_h, image_w = label_map.shape

genuine_particles = []
border_particles = []
spurious_particles = []

for p in props:
    eq_diam_nm = p.equivalent_diameter * pixel_size_nm
    solidity = p.solidity
    
    # Check if region touches image border
    minr, minc, maxr, maxc = p.bbox
    touches_border = (minr == 0 or minc == 0 or maxr == image_h or maxc == image_w)
    
    if touches_border:
        border_particles.append(p)
    elif solidity < 0.5 or eq_diam_nm < 20:
        spurious_particles.append(p)
    else:
        genuine_particles.append(p)

num_genuine = len(genuine_particles)
num_border_excluded = len(border_particles)
num_spurious_removed = len(spurious_particles)

print(f"Genuine particles: {num_genuine}")
print(f"Border particles excluded: {num_border_excluded}")
print(f"Spurious detections removed: {num_spurious_removed}")

# Corrected size statistics
genuine_diameters_nm = np.array([p.equivalent_diameter * pixel_size_nm for p in genuine_particles])
corrected_mean_diam = float(np.mean(genuine_diameters_nm)) if len(genuine_diameters_nm) > 0 else 0.0
corrected_std_diam = float(np.std(genuine_diameters_nm)) if len(genuine_diameters_nm) > 0 else 0.0
corrected_cv = float(corrected_std_diam / corrected_mean_diam * 100) if corrected_mean_diam > 0 else 0.0

# Get centroids for genuine particles
centroids_px = np.array([p.centroid for p in genuine_particles])  # (row, col)
centroids_nm = centroids_px * pixel_size_nm

# Step 4: Compute Delaunay triangulation
if num_genuine >= 4:
    # Delaunay expects (x, y) = (col, row)
    points_for_delaunay = centroids_px[:, ::-1]  # convert (row,col) to (col,row) = (x,y)
    tri = Delaunay(points_for_delaunay)
    
    # Step 5: Extract neighbor pairs from Delaunay edges with distance cutoff
    # Build edge set from Delaunay simplices
    edges = set()
    for simplex in tri.simplices:
        for i in range(3):
            for j in range(i+1, 3):
                a, b = simplex[i], simplex[j]
                edges.add((min(a, b), max(a, b)))
    
    # Compute distances for all Delaunay edges
    edge_list = list(edges)
    edge_distances_px = []
    for a, b in edge_list:
        d = np.linalg.norm(centroids_px[a] - centroids_px[b])
        edge_distances_px.append(d)
    edge_distances_px = np.array(edge_distances_px)
    edge_distances_nm = edge_distances_px * pixel_size_nm
    
    # Compute nearest-neighbor distance for each particle
    nn_distances_px = np.full(num_genuine, np.inf)
    for (a, b), d in zip(edge_list, edge_distances_px):
        if d < nn_distances_px[a]:
            nn_distances_px[a] = d
        if d < nn_distances_px[b]:
            nn_distances_px[b] = d
    
    median_nn_px = np.median(nn_distances_px[nn_distances_px < np.inf])
    
    # Distance cutoff for valid neighbors: 1.3 * median NN distance
    distance_cutoff_px = 1.3 * median_nn_px
    
    # Filter edges
    valid_edges = [(a, b) for (a, b), d in zip(edge_list, edge_distances_px) if d < distance_cutoff_px]
    valid_edge_distances_px = np.array([d for d in edge_distances_px if d < distance_cutoff_px])
    valid_edge_distances_nm = valid_edge_distances_px * pixel_size_nm
    
    # Rebuild neighbor list with cutoff
    neighbors = {i: [] for i in range(num_genuine)}
    for a, b in valid_edges:
        da = np.linalg.norm(centroids_px[a] - centroids_px[b])
        neighbors[a].append((b, da))
        neighbors[b].append((a, da))
    
    # Recompute NN distances with valid edges only
    nn_distances_px_filtered = []
    for i in range(num_genuine):
        if neighbors[i]:
            min_d = min(d for _, d in neighbors[i])
            nn_distances_px_filtered.append(min_d)
        else:
            nn_distances_px_filtered.append(np.nan)
    nn_distances_px_filtered = np.array(nn_distances_px_filtered)
    nn_distances_nm_filtered = nn_distances_px_filtered * pixel_size_nm
    
    # Step 6: NN distance statistics
    valid_nn = nn_distances_nm_filtered[~np.isnan(nn_distances_nm_filtered)]
    mean_nn_dist = float(np.mean(valid_nn))
    std_nn_dist = float(np.std(valid_nn))
    nn_to_diam_ratio = float(mean_nn_dist / corrected_mean_diam) if corrected_mean_diam > 0 else 0.0
    
    # Step 7: Compute psi6 for each particle
    psi6_values = np.zeros(num_genuine)
    for i in range(num_genuine):
        if len(neighbors[i]) == 0:
            psi6_values[i] = 0.0
            continue
        n_neigh = len(neighbors[i])
        psi6_complex = 0.0
        for j, _ in neighbors[i]:
            dx = centroids_px[j, 1] - centroids_px[i, 1]  # col difference
            dy = centroids_px[j, 0] - centroids_px[i, 0]  # row difference
            theta = np.arctan2(dy, dx)
            psi6_complex += np.exp(6j * theta)
        psi6_values[i] = abs(psi6_complex / n_neigh)
    
    mean_psi6 = float(np.mean(psi6_values))
    std_psi6 = float(np.std(psi6_values))
    median_psi6 = float(np.median(psi6_values))
    fraction_highly_ordered = float(np.mean(psi6_values > 0.8))
    
    # Step 8: Radial pair correlation function g(r)
    # Compute all pairwise distances
    from scipy.spatial.distance import pdist, squareform
    all_pair_dists_px = pdist(centroids_px)
    all_pair_dists_nm = all_pair_dists_px * pixel_size_nm
    
    # g(r) calculation
    r_max = 500.0  # nm
    dr = 5.0  # nm bin width
    r_bins = np.arange(0, r_max + dr, dr)
    r_centers = 0.5 * (r_bins[:-1] + r_bins[1:])
    
    hist_counts, _ = np.histogram(all_pair_dists_nm, bins=r_bins)
    
    # Normalization: for 2D ideal gas
    # Area of the analysis region
    # Use convex hull area or bounding box of centroids
    from scipy.spatial import ConvexHull
    try:
        hull = ConvexHull(centroids_px)
        area_px2 = hull.volume  # In 2D, volume = area
        area_nm2 = area_px2 * pixel_size_nm**2
    except:
        # Fallback to image area
        area_nm2 = image_h * image_w * pixel_size_nm**2
    
    rho = num_genuine / area_nm2  # number density
    
    # g(r) = hist / (N * rho * 2*pi*r*dr)
    g_r = np.zeros_like(r_centers)
    for k in range(len(r_centers)):
        shell_area = 2.0 * np.pi * r_centers[k] * dr
        expected = num_genuine * rho * shell_area / 2.0  # divide by 2 because pdist gives each pair once
        if expected > 0:
            g_r[k] = hist_counts[k] / expected
    
    # Find first and second peaks in g(r)
    # Skip first few bins (r < 30 nm)
    start_idx = int(30 / dr)
    
    # First peak
    if len(g_r) > start_idx + 5:
        from scipy.signal import find_peaks
        peaks, peak_props = find_peaks(g_r[start_idx:], height=1.0, distance=3)
        if len(peaks) >= 1:
            first_peak_idx = peaks[0] + start_idx
            g_r_first_peak_pos = float(r_centers[first_peak_idx])
            g_r_first_peak_height = float(g_r[first_peak_idx])
        else:
            # Fallback: find max in reasonable range
            search_range = g_r[start_idx:min(start_idx+30, len(g_r))]
            first_peak_idx = np.argmax(search_range) + start_idx
            g_r_first_peak_pos = float(r_centers[first_peak_idx])
            g_r_first_peak_height = float(g_r[first_peak_idx])
        
        # Second peak: search after first peak
        second_peak_start = first_peak_idx + 5
        if len(peaks) >= 2:
            second_peak_idx = peaks[1] + start_idx
            g_r_second_peak_pos = float(r_centers[second_peak_idx])
        elif second_peak_start < len(g_r) - 5:
            search_range2 = g_r[second_peak_start:min(second_peak_start+30, len(g_r))]
            if len(search_range2) > 0:
                second_peak_idx = np.argmax(search_range2) + second_peak_start
                g_r_second_peak_pos = float(r_centers[second_peak_idx])
            else:
                g_r_second_peak_pos = None
        else:
            g_r_second_peak_pos = None
    else:
        g_r_first_peak_pos = None
        g_r_first_peak_height = None
        g_r_second_peak_pos = None
    
    # Step 9: Connectivity graph for islands
    # Use 1.2 * median NN distance as threshold
    island_threshold_px = 1.2 * median_nn_px
    
    # Build adjacency using scipy sparse graph or simple union-find
    parent = list(range(num_genuine))
    
    def find(x):
        while parent[x] != x:
            parent[x] = parent[parent[x]]
            x = parent[x]
        return x
    
    def union(x, y):
        px, py = find(x), find(y)
        if px != py:
            parent[px] = py
    
    for a, b in valid_edges:
        d = np.linalg.norm(centroids_px[a] - centroids_px[b])
        if d < island_threshold_px:
            union(a, b)
    
    # Count islands
    from collections import Counter
    island_labels = [find(i) for i in range(num_genuine)]
    island_counter = Counter(island_labels)
    island_sizes = sorted(island_counter.values(), reverse=True)
    number_of_islands = len(island_sizes)
    max_island_size = island_sizes[0] if island_sizes else 0
    
    # Map each particle to its island ID for coloring
    island_id_map = {}
    for i, label in enumerate(island_labels):
        root = find(label)
        if root not in island_id_map:
            island_id_map[root] = len(island_id_map)
    particle_island_ids = [island_id_map[find(i)] for i in range(num_genuine)]
    
else:
    # Not enough particles for meaningful analysis
    mean_nn_dist = 0.0
    std_nn_dist = 0.0
    nn_to_diam_ratio = 0.0
    mean_psi6 = 0.0
    std_psi6 = 0.0
    median_psi6 = 0.0
    fraction_highly_ordered = 0.0
    g_r_first_peak_pos = None
    g_r_first_peak_height = None
    g_r_second_peak_pos = None
    number_of_islands = 0
    island_sizes = []
    max_island_size = 0
    psi6_values = np.array([])
    valid_edges = []
    valid_nn = np.array([])
    r_centers = np.array([])
    g_r = np.array([])

# Step 10: Generate multi-panel figure
fig, axes = plt.subplots(2, 3, figsize=(18, 12), dpi=100)

# (a) Original image with centroids colored by psi6 and Delaunay edges
ax = axes[0, 0]
ax.imshow(img_display, cmap='gray')
if num_genuine >= 4:
    # Draw valid Delaunay edges
    for a, b in valid_edges:
        ax.plot([centroids_px[a, 1], centroids_px[b, 1]], 
                [centroids_px[a, 0], centroids_px[b, 0]], 
                'w-', alpha=0.3, linewidth=0.5)
    # Scatter centroids colored by psi6
    sc = ax.scatter(centroids_px[:, 1], centroids_px[:, 0], 
                    c=psi6_values, cmap='RdYlGn', vmin=0, vmax=1, 
                    s=30, edgecolors='white', linewidths=0.5)
    plt.colorbar(sc, ax=ax, label='$\\psi_6$', shrink=0.7)
ax.set_title(f'Particles colored by $\\psi_6$ (n={num_genuine})')
ax.set_xlim(0, image_w)
ax.set_ylim(image_h, 0)
ax.axis('off')

# (b) Nearest-neighbor distance histogram
ax = axes[0, 1]
if len(valid_nn) > 0:
    ax.hist(valid_nn, bins=20, color='steelblue', edgecolor='black', alpha=0.8)
    ax.axvline(mean_nn_dist, color='red', linestyle='--', linewidth=2, label=f'Mean={mean_nn_dist:.1f} nm')
    ax.set_xlabel('Nearest-neighbor distance (nm)')
    ax.set_ylabel('Count')
    ax.legend()
ax.set_title('NN Distance Distribution')

# (c) g(r) plot
ax = axes[0, 2]
if len(r_centers) > 0:
    ax.plot(r_centers, g_r, 'b-', linewidth=1.5)
    ax.axhline(1.0, color='gray', linestyle='--', alpha=0.5)
    if g_r_first_peak_pos is not None:
        ax.axvline(g_r_first_peak_pos, color='red', linestyle=':', alpha=0.7, label=f'1st peak={g_r_first_peak_pos:.0f} nm')
    if g_r_second_peak_pos is not None:
        ax.axvline(g_r_second_peak_pos, color='orange', linestyle=':', alpha=0.7, label=f'2nd peak={g_r_second_peak_pos:.0f} nm')
    ax.set_xlabel('r (nm)')
    ax.set_ylabel('g(r)')
    ax.legend(fontsize=8)
    ax.set_xlim(0, 400)
ax.set_title('Pair Correlation Function g(r)')

# (d) psi6 distribution
ax = axes[1, 0]
if len(psi6_values) > 0:
    ax.hist(psi6_values, bins=20, range=(0, 1), color='green', edgecolor='black', alpha=0.8)
    ax.axvline(mean_psi6, color='red', linestyle='--', linewidth=2, label=f'Mean={mean_psi6:.3f}')
    ax.set_xlabel('$\\psi_6$')
    ax.set_ylabel('Count')
    ax.legend()
ax.set_title('$\\psi_6$ Distribution')

# (e) Corrected size distribution
ax = axes[1, 1]
if len(genuine_diameters_nm) > 0:
    ax.hist(genuine_diameters_nm, bins=20, color='coral', edgecolor='black', alpha=0.8)
    ax.axvline(corrected_mean_diam, color='red', linestyle='--', linewidth=2, label=f'Mean={corrected_mean_diam:.1f}\u00b1{corrected_std_diam:.1f} nm')
    ax.set_xlabel('Equivalent Diameter (nm)')
    ax.set_ylabel('Count')
    ax.legend(fontsize=8)
ax.set_title('Corrected Size Distribution')

# (f) Segmentation overlay with island coloring
ax = axes[1, 2]
ax.imshow(img_display, cmap='gray')
if num_genuine >= 4:
    # Color by island
    island_colors = np.array(particle_island_ids)
    sc2 = ax.scatter(centroids_px[:, 1], centroids_px[:, 0], 
                     c=island_colors, cmap='tab20', s=40, 
                     edgecolors='white', linewidths=0.5)
    # Draw edges within islands
    for a, b in valid_edges:
        d = np.linalg.norm(centroids_px[a] - centroids_px[b])
        if d < island_threshold_px:
            ax.plot([centroids_px[a, 1], centroids_px[b, 1]], 
                    [centroids_px[a, 0], centroids_px[b, 0]], 
                    'cyan', alpha=0.4, linewidth=0.5)
ax.set_title(f'Close-packed Islands (n={number_of_islands})')
ax.set_xlim(0, image_w)
ax.set_ylim(image_h, 0)
ax.axis('off')

plt.tight_layout()

# Determine all output directories to save to
output_dirs = set()
output_dirs.add(cwd)
output_dirs.add(base_dir)
# Also try the parent of base_dir in case system expects it there
output_dirs.add(os.path.dirname(base_dir))

for out_dir in output_dirs:
    if not os.path.isdir(out_dir):
        try:
            os.makedirs(out_dir, exist_ok=True)
        except:
            continue
    try:
        out_path_1 = os.path.join(out_dir, 'image_0000_analysis.png')
        out_path_2 = os.path.join(out_dir, 'analysis_visualization.png')
        plt.savefig(out_path_1, dpi=100, bbox_inches='tight')
        plt.savefig(out_path_2, dpi=100, bbox_inches='tight')
        print(f"Saved visualization to: {out_path_1}")
        print(f"Saved visualization to: {out_path_2}")
    except Exception as e:
        print(f"Warning: Could not save to {out_dir}: {e}")

plt.close()

# Save arrays to base_dir
np.save(os.path.join(base_dir, 'analysis_labels.npy'), label_map)  # Keep the original label map
np.save(os.path.join(base_dir, 'filtered_centroids_px.npy'), centroids_px)
np.save(os.path.join(base_dir, 'psi6_values.npy'), psi6_values)
np.save(os.path.join(base_dir, 'genuine_diameters_nm.npy'), genuine_diameters_nm)

# Also save arrays to cwd
if cwd != base_dir:
    try:
        np.save(os.path.join(cwd, 'analysis_labels.npy'), label_map)
        np.save(os.path.join(cwd, 'filtered_centroids_px.npy'), centroids_px)
        np.save(os.path.join(cwd, 'psi6_values.npy'), psi6_values)
        np.save(os.path.join(cwd, 'genuine_diameters_nm.npy'), genuine_diameters_nm)
    except Exception as e:
        print(f"Warning: Could not save arrays to cwd: {e}")

# Compile results
results = {
    "analysis_type": "Spatial ordering analysis of nanoparticle assembly: nearest-neighbor distances, hexagonal bond-orientational order (psi6), pair correlation function g(r), and close-packed island identification from Tier 1 segmentation label map",
    "extracted_features": {
        "corrected_mean_diameter_nm": round(corrected_mean_diam, 2),
        "corrected_std_diameter_nm": round(corrected_std_diam, 2),
        "corrected_cv_percent": round(corrected_cv, 2),
        "num_genuine_particles": int(num_genuine),
        "num_border_particles_excluded": int(num_border_excluded),
        "num_spurious_detections_removed": int(num_spurious_removed),
        "mean_nearest_neighbor_distance_nm": round(mean_nn_dist, 2),
        "std_nearest_neighbor_distance_nm": round(std_nn_dist, 2),
        "nn_distance_to_diameter_ratio": round(nn_to_diam_ratio, 3),
        "mean_psi6": round(mean_psi6, 4),
        "std_psi6": round(std_psi6, 4),
        "median_psi6": round(median_psi6, 4),
        "fraction_highly_ordered_psi6_above_0.8": round(fraction_highly_ordered, 4),
        "g_r_first_peak_position_nm": round(g_r_first_peak_pos, 1) if g_r_first_peak_pos is not None else None,
        "g_r_first_peak_height": round(g_r_first_peak_height, 2) if g_r_first_peak_height is not None else None,
        "g_r_second_peak_position_nm": round(g_r_second_peak_pos, 1) if g_r_second_peak_pos is not None else None,
        "number_of_islands": int(number_of_islands),
        "island_sizes_particles": island_sizes,
        "max_island_size": int(max_island_size),
        "pixel_size_nm": round(pixel_size_nm, 4),
        "assumed_particle_diameter_nm": assumed_particle_diameter_nm
    },
    "quality_metrics": {
        "total_labeled_regions": len(props),
        "particles_after_filtering": int(num_genuine),
        "delaunay_edge_distance_cutoff_nm": round(distance_cutoff_px * pixel_size_nm, 1) if num_genuine >= 4 else None,
        "island_connectivity_threshold_nm": round(island_threshold_px * pixel_size_nm, 1) if num_genuine >= 4 else None,
        "median_nn_distance_px": round(float(median_nn_px), 1) if num_genuine >= 4 else None
    },
    "summary": f"Spatial ordering analysis of {num_genuine} genuine nanoparticles (from {len(props)} total detections, {num_border_excluded} border-excluded, {num_spurious_removed} spurious removed). Mean diameter {corrected_mean_diam:.1f}\u00b1{corrected_std_diam:.1f} nm (CV={corrected_cv:.1f}%). Mean NN distance {mean_nn_dist:.1f}\u00b1{std_nn_dist:.1f} nm (ratio to diameter: {nn_to_diam_ratio:.2f}). Mean psi6={mean_psi6:.3f} (fraction highly ordered: {fraction_highly_ordered:.1%}). Found {number_of_islands} close-packed islands with max size {max_island_size} particles. Pixel size assumed from 75 nm nominal particle diameter.",
    "saved_arrays": {
        "analysis_labels.npy": {
            "description": f"Integer label map from Tier 1 segmentation, {len(props)} regions labeled 1-{len(props)}, background=0",
            "shape": list(label_map.shape),
            "dtype": str(label_map.dtype)
        },
        "filtered_centroids_px.npy": {
            "description": f"Centroids (row, col) in pixels for {num_genuine} genuine, non-border particles",
            "shape": list(centroids_px.shape),
            "dtype": str(centroids_px.dtype)
        },
        "psi6_values.npy": {
            "description": f"Hexagonal bond-orientational order parameter psi6 for each of {num_genuine} genuine particles",
            "shape": list(psi6_values.shape),
            "dtype": str(psi6_values.dtype)
        },
        "genuine_diameters_nm.npy": {
            "description": f"Equivalent diameters in nm for {num_genuine} genuine particles",
            "shape": list(genuine_diameters_nm.shape),
            "dtype": str(genuine_diameters_nm.dtype)
        }
    }
}

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