import numpy as np
import json
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from matplotlib.colors import ListedColormap
import cv2
from scipy import ndimage, stats
from skimage import morphology, measure, segmentation, feature, filters
from skimage.morphology import disk, skeletonize
import warnings
warnings.filterwarnings('ignore')

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

# Ensure 2D
if image.ndim == 3:
    image = image[:, :, 0]

# Step 1: Gaussian smoothing
smoothed = filters.gaussian(image, sigma=2)

# Step 2: Otsu threshold to segment dark phase
# Dark phase = below threshold
thresh = filters.threshold_otsu(smoothed)
binary_dark = smoothed < thresh
print(f"Otsu threshold: {thresh:.4f}, dark phase pixels: {binary_dark.sum()}, fraction: {binary_dark.sum()/(image.shape[0]*image.shape[1]):.4f}")

# Step 3: Morphological opening to clean
binary_clean = morphology.binary_opening(binary_dark, disk(4))

# Step 4: Connected component labeling
labeled_initial, n_initial = ndimage.label(binary_clean)
print(f"Initial connected components: {n_initial}")

# Step 5: Remove small artifacts < 500 pixels
props_initial = measure.regionprops(labeled_initial)
total_dark_area = binary_clean.sum()

# Build cleaned label map
labeled_cleaned = np.zeros_like(labeled_initial)
new_label = 0
initial_domains = []
for prop in props_initial:
    if prop.area >= 500:
        new_label += 1
        labeled_cleaned[labeled_initial == prop.label] = new_label
        initial_domains.append({'label': new_label, 'area': prop.area, 'area_fraction': prop.area / total_dark_area})

print(f"Domains after artifact removal: {len(initial_domains)}")
for d in initial_domains:
    print(f"  Domain {d['label']}: area={d['area']}, fraction={d['area_fraction']:.4f}")

# Provenance tracking
provenance = {}  # domain_id -> {'method': str, 'parent_id': int, 'pass_level': int}

# Track constriction detection stats
skeleton_constriction_points_detected = 0
skeleton_constriction_widths = []
bridge_breaking_cuts_count = 0
erosion_fallback_used = False

def find_constriction_points(domain_mask, width_threshold=8):
    """Find constriction points using skeleton + distance transform."""
    # Distance transform
    dist = ndimage.distance_transform_edt(domain_mask)
    
    # Skeletonize
    skel = skeletonize(domain_mask)
    
    # Get skeleton pixel coordinates
    skel_coords = np.argwhere(skel)
    if len(skel_coords) < 10:
        return [], [], skel, dist
    
    # Sample distance transform along skeleton
    skel_widths = dist[skel]
    
    # Find constriction points: where width < threshold
    constriction_mask = skel_widths < width_threshold
    constriction_coords = skel_coords[constriction_mask]
    constriction_widths_vals = skel_widths[constriction_mask]
    
    return constriction_coords, constriction_widths_vals, skel, dist

def cut_at_constrictions(domain_mask, constriction_coords, dist, cut_width=20):
    """Cut the binary mask at constriction points."""
    cut_mask = domain_mask.copy()
    
    if len(constriction_coords) == 0:
        return cut_mask, 0
    
    cuts_made = 0
    # Group nearby constriction points and cut at the narrowest in each group
    # Use simple clustering: if points are within 15 pixels, group them
    if len(constriction_coords) > 0:
        from sklearn.cluster import DBSCAN
        clustering = DBSCAN(eps=15, min_samples=1).fit(constriction_coords)
        labels_clust = clustering.labels_
        
        for cl in set(labels_clust):
            cl_mask = labels_clust == cl
            cl_coords = constriction_coords[cl_mask]
            # Find the narrowest point in this cluster
            cl_dists = dist[cl_coords[:, 0], cl_coords[:, 1]]
            narrowest_idx = np.argmin(cl_dists)
            cy, cx = cl_coords[narrowest_idx]
            local_width = max(cl_dists[narrowest_idx], 1)
            
            # Determine skeleton direction at this point using neighbors
            # Use a small window to estimate direction
            neighborhood = 7
            y_min = max(0, cy - neighborhood)
            y_max = min(domain_mask.shape[0], cy + neighborhood)
            x_min = max(0, cx - neighborhood)
            x_max = min(domain_mask.shape[1], cx + neighborhood)
            
            local_skel = np.zeros_like(domain_mask)
            # Find skeleton in neighborhood
            local_domain = domain_mask[y_min:y_max, x_min:x_max]
            if local_domain.sum() < 5:
                continue
            
            # Estimate perpendicular direction
            # Use image gradient of distance transform for direction
            gy, gx = np.gradient(dist)
            angle = np.arctan2(gy[cy, cx], gx[cy, cx])
            
            # Perpendicular to gradient (along the skeleton)
            perp_angle = angle  # gradient points away from boundary, perpendicular to skeleton
            
            # Cut line perpendicular to skeleton direction
            half_len = int(local_width + 4)
            for t in range(-half_len, half_len + 1):
                py = int(cy + t * np.cos(perp_angle))
                px = int(cx + t * np.sin(perp_angle))
                if 0 <= py < cut_mask.shape[0] and 0 <= px < cut_mask.shape[1]:
                    # Cut a wider swath (3 pixels wide)
                    for dy in range(-1, 2):
                        for dx in range(-1, 2):
                            ny, nx = py + dy, px + dx
                            if 0 <= ny < cut_mask.shape[0] and 0 <= nx < cut_mask.shape[1]:
                                cut_mask[ny, nx] = False
            cuts_made += 1
    
    return cut_mask, cuts_made

def erosion_fallback_split(domain_mask, erosion_radius=10):
    """Split domain using heavy erosion followed by dilation back."""
    global erosion_fallback_used
    erosion_fallback_used = True
    
    n_eroded = 0
    labeled_eroded = None
    for radius in [erosion_radius, erosion_radius + 2]:
        eroded = morphology.binary_erosion(domain_mask, disk(radius))
        labeled_eroded, n_eroded = ndimage.label(eroded)
        if n_eroded >= 3:
            break
    
    if n_eroded < 2:
        return None, 0
    
    # Assign each pixel in original domain to nearest eroded component
    # Use distance-based assignment
    result = np.zeros_like(labeled_eroded)
    # For each eroded component, dilate and assign
    dist_maps = []
    for i in range(1, n_eroded + 1):
        comp_mask = labeled_eroded == i
        d = ndimage.distance_transform_edt(~comp_mask)
        dist_maps.append(d)
    
    dist_stack = np.stack(dist_maps, axis=0)
    nearest = np.argmin(dist_stack, axis=0) + 1
    result = np.where(domain_mask, nearest, 0)
    
    return result, n_eroded

# Step 6-8: Process domains that need splitting
current_label_map = labeled_cleaned.copy()
max_label = current_label_map.max()

# Initialize provenance for original domains
for d in initial_domains:
    provenance[d['label']] = {'method': 'original', 'parent_id': 0, 'pass_level': 0}

def process_large_domain(label_map, domain_label, width_threshold, min_sub_area, min_expected_splits, pass_level):
    """Process a large domain with constriction-based splitting."""
    global max_label, skeleton_constriction_points_detected, skeleton_constriction_widths
    global bridge_breaking_cuts_count, provenance
    
    domain_mask = label_map == domain_label
    domain_area = domain_mask.sum()
    
    # Try constriction-point bridge-breaking
    constriction_coords, constriction_widths_vals, skel, dist = find_constriction_points(
        domain_mask, width_threshold=width_threshold
    )
    
    skeleton_constriction_points_detected += len(constriction_coords)
    skeleton_constriction_widths.extend(constriction_widths_vals.tolist() if len(constriction_widths_vals) > 0 else [])
    
    cut_mask, cuts = cut_at_constrictions(domain_mask, constriction_coords, dist)
    bridge_breaking_cuts_count += cuts
    
    # Re-label
    cut_labeled, n_cut = ndimage.label(cut_mask)
    
    # Filter small fragments
    valid_cuts = []
    for i in range(1, n_cut + 1):
        comp_area = (cut_labeled == i).sum()
        if comp_area >= min_sub_area:
            valid_cuts.append(i)
    
    print(f"  Constriction cutting: {len(constriction_coords)} points, {cuts} cuts, {len(valid_cuts)} valid sub-domains")
    
    method_used = 'bridge_cut'
    
    # If not enough sub-domains, try erosion fallback
    if len(valid_cuts) < min_expected_splits:
        print(f"  Constriction cutting insufficient ({len(valid_cuts)} < {min_expected_splits}), trying erosion fallback...")
        eroded_result, n_eroded = erosion_fallback_split(domain_mask)
        if eroded_result is not None and n_eroded >= min_expected_splits:
            # Validate sub-domain sizes
            valid_eroded = []
            for i in range(1, n_eroded + 1):
                comp_area = (eroded_result == i).sum()
                if comp_area >= min_sub_area:
                    valid_eroded.append(i)
            
            if len(valid_eroded) >= min_expected_splits:
                print(f"  Erosion fallback: {len(valid_eroded)} valid sub-domains")
                # Use erosion result
                # Remove old domain from label map
                new_map = label_map.copy()
                new_map[domain_mask] = 0
                
                # Remove old provenance
                if domain_label in provenance:
                    del provenance[domain_label]
                
                for vi in valid_eroded:
                    max_label += 1
                    new_map[eroded_result == vi] = max_label
                    provenance[max_label] = {'method': 'erosion_split', 'parent_id': int(domain_label), 'pass_level': pass_level}
                
                return new_map
            else:
                method_used = 'erosion_split'
    
    if len(valid_cuts) >= 2:  # Accept any meaningful split
        # Use cut result
        new_map = label_map.copy()
        new_map[domain_mask] = 0
        
        if domain_label in provenance:
            del provenance[domain_label]
        
        # Assign uncut pixels (small fragments) to nearest valid component
        # First assign valid cuts
        for vi in valid_cuts:
            max_label += 1
            new_map[cut_labeled == vi] = max_label
            provenance[max_label] = {'method': method_used, 'parent_id': int(domain_label), 'pass_level': pass_level}
        
        # Assign remaining domain pixels to nearest component
        remaining = domain_mask & (new_map == 0)
        if remaining.sum() > 0:
            # Distance-based assignment
            assigned_mask = (new_map > 0) & domain_mask
            if assigned_mask.sum() > 0:
                temp_map = new_map.copy()
                for _ in range(50):
                    dilated = ndimage.maximum_filter(temp_map, size=3)
                    update = remaining & (temp_map == 0) & (dilated > 0)
                    temp_map[update] = dilated[update]
                    remaining = domain_mask & (temp_map == 0)
                    if remaining.sum() == 0:
                        break
                new_map = temp_map
        
        return new_map
    
    # If constriction cutting failed with less than 2, keep original
    print(f"  No effective split achieved for domain {domain_label}")
    return label_map

# Process domains by size
domains_to_process = []
for d in initial_domains:
    if d['area_fraction'] > 0.25:
        domains_to_process.append((d['label'], d['area_fraction'], 8, 2000, 3, 1))
    elif d['area_fraction'] > 0.20:
        domains_to_process.append((d['label'], d['area_fraction'], 10, 2000, 2, 2))
    elif d['area_fraction'] > 0.10:
        domains_to_process.append((d['label'], d['area_fraction'], 6, 1500, 2, 3))

domains_to_process.sort(key=lambda x: -x[1])  # Process largest first

for dlabel, dfrac, wthresh, min_area, min_splits, plevel in domains_to_process:
    # Check if domain still exists in current map
    if (current_label_map == dlabel).sum() < 500:
        continue
    print(f"\nProcessing domain {dlabel} (fraction={dfrac:.4f}, width_thresh={wthresh}, min_splits={min_splits})")
    current_label_map = process_large_domain(current_label_map, dlabel, wthresh, min_area, min_splits, plevel)

# Step 9: Global domain count check
unique_labels = [l for l in np.unique(current_label_map) if l > 0]
print(f"\nDomain count after bridge-breaking: {len(unique_labels)}")

# If count < 12, try watershed on 3 largest remaining domains
if len(unique_labels) < 12:
    print("Count below 12, attempting watershed on largest remaining domains...")
    domain_areas = [(l, (current_label_map == l).sum()) for l in unique_labels]
    domain_areas.sort(key=lambda x: -x[1])
    
    for dlabel, darea in domain_areas[:3]:
        if darea < 5000:  # Don't split small domains
            continue
        domain_mask = current_label_map == dlabel
        
        # Watershed with distance transform
        dist = ndimage.distance_transform_edt(domain_mask)
        
        # Try Sobel gradient as landscape (per advanced guidance)
        gradient = filters.sobel(image)
        gradient_masked = np.where(domain_mask, gradient, 0)
        
        # Find peaks in distance transform for markers
        local_max_coords = feature.peak_local_max(dist, min_distance=30, labels=domain_mask.astype(int))
        
        if len(local_max_coords) < 2:
            continue
        
        markers = np.zeros_like(domain_mask, dtype=int)
        for i, (y, x) in enumerate(local_max_coords):
            markers[y, x] = i + 1
        
        # Use gradient as landscape
        ws = segmentation.watershed(gradient_masked, markers, mask=domain_mask)
        
        n_ws = len(set(np.unique(ws)) - {0})
        if n_ws < 2:
            continue
        
        # Validate: both sub-domains > 2000 pixels and boundary length check
        valid = True
        ws_labels = [l for l in np.unique(ws) if l > 0]
        ws_areas = [(l, (ws == l).sum()) for l in ws_labels]
        
        for wl, wa in ws_areas:
            if wa < 2000:
                valid = False
                break
        
        if valid and len(ws_areas) >= 2:
            # Check boundary length between segments
            # Find boundary pixels
            dilated_ws = ndimage.maximum_filter(ws, size=3)
            eroded_ws = ndimage.minimum_filter(ws, size=3)
            boundary = (dilated_ws != eroded_ws) & domain_mask
            boundary_len = boundary.sum()
            min_ws_area = min(wa for _, wa in ws_areas)
            
            # Estimate perimeter of smaller domain
            smaller_label = min(ws_areas, key=lambda x: x[1])[0]
            smaller_perim = measure.regionprops((ws == smaller_label).astype(int))
            if len(smaller_perim) > 0:
                smaller_perim_val = smaller_perim[0].perimeter
                if boundary_len > 0.15 * smaller_perim_val:
                    print(f"  Watershed split rejected for domain {dlabel}: boundary too long ({boundary_len} > 15% of {smaller_perim_val:.0f})")
                    valid = False
        
        if valid and len(ws_areas) >= 2:
            print(f"  Watershed split domain {dlabel} into {len(ws_areas)} sub-domains")
            # Apply split
            if dlabel in provenance:
                del provenance[dlabel]
            current_label_map[domain_mask] = 0
            for wl, wa in ws_areas:
                max_label += 1
                current_label_map[ws == wl] = max_label
                provenance[max_label] = {'method': 'watershed', 'parent_id': int(dlabel), 'pass_level': 4}

# Re-check unique labels
unique_labels = [l for l in np.unique(current_label_map) if l > 0]
print(f"Domain count after watershed pass: {len(unique_labels)}")

# Step 10: If count > 16, merge smallest adjacent pairs
if len(unique_labels) > 16:
    while len(unique_labels) > 16:
        domain_areas = [(l, (current_label_map == l).sum()) for l in unique_labels]
        domain_areas.sort(key=lambda x: x[1])
        
        merged = False
        for dlabel, darea in domain_areas:
            # Find adjacent domains
            domain_mask = current_label_map == dlabel
            dilated = morphology.binary_dilation(domain_mask, disk(2))
            neighbors = set(np.unique(current_label_map[dilated & ~domain_mask])) - {0}
            
            if not neighbors:
                continue
            
            # Find neighbor with longest shared boundary
            best_neighbor = None
            best_shared = 0
            for n in neighbors:
                n_mask = current_label_map == n
                shared = (dilated & n_mask).sum()
                if shared > best_shared:
                    best_shared = shared
                    best_neighbor = n
            
            if best_neighbor is not None:
                perim_est = measure.regionprops(domain_mask.astype(int))
                if len(perim_est) > 0 and best_shared > 0.3 * perim_est[0].perimeter:
                    current_label_map[current_label_map == dlabel] = best_neighbor
                    if dlabel in provenance:
                        provenance[dlabel]['method'] += '_merged'
                    merged = True
                    break
        
        if not merged:
            break
        unique_labels = [l for l in np.unique(current_label_map) if l > 0]

# Relabel sequentially
final_label_map = np.zeros_like(current_label_map)
new_provenance = {}
for new_id, old_id in enumerate(sorted(unique_labels), 1):
    final_label_map[current_label_map == old_id] = new_id
    if old_id in provenance:
        new_provenance[new_id] = provenance[old_id]
    else:
        new_provenance[new_id] = {'method': 'original', 'parent_id': 0, 'pass_level': 0}

provenance = new_provenance
unique_labels = [l for l in np.unique(final_label_map) if l > 0]
final_domain_count = len(unique_labels)
print(f"\nFinal domain count: {final_domain_count}")

# Step 11: Validate count
in_target_range = 12 <= final_domain_count <= 16
print(f"Within target range (12-16): {in_target_range}")

# Step 12: Extract regionprops
props = measure.regionprops(final_label_map, intensity_image=image)

def compute_convex_hull_perimeter(prop):
    """Compute the convex hull perimeter from the convex hull coordinates."""
    try:
        # Get the binary mask of the region
        region_mask = (final_label_map == prop.label).astype(np.uint8)
        # Find contours
        contours, _ = cv2.findContours(region_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
        if len(contours) == 0:
            return prop.perimeter
        # Get the largest contour
        largest_contour = max(contours, key=cv2.contourArea)
        # Compute convex hull of the contour
        hull = cv2.convexHull(largest_contour)
        # Compute perimeter of the convex hull
        convex_perim = cv2.arcLength(hull, closed=True)
        if convex_perim > 0:
            return float(convex_perim)
        else:
            return prop.perimeter
    except Exception:
        return prop.perimeter

per_domain_data = []
for prop in props:
    # Circularity
    if prop.perimeter > 0:
        circularity = 4 * np.pi * prop.area / (prop.perimeter ** 2)
    else:
        circularity = 0
    
    # Smoothed circularity using convex hull perimeter (FIX: proper convex hull perimeter)
    convex_perim = compute_convex_hull_perimeter(prop)
    
    if convex_perim > 0:
        circularity_smoothed = 4 * np.pi * prop.area / (convex_perim ** 2)
    else:
        circularity_smoothed = 0
    
    aspect_ratio = prop.major_axis_length / max(prop.minor_axis_length, 1)
    
    domain_info = {
        'label': prop.label,
        'area': int(prop.area),
        'area_fraction': prop.area / total_dark_area,
        'centroid': [float(prop.centroid[0]), float(prop.centroid[1])],
        'major_axis_length': float(prop.major_axis_length),
        'minor_axis_length': float(prop.minor_axis_length),
        'aspect_ratio': float(aspect_ratio),
        'perimeter': float(prop.perimeter),
        'perimeter_smoothed_convex_hull': float(convex_perim),
        'circularity': float(circularity),
        'circularity_smoothed': float(circularity_smoothed),
        'eccentricity': float(prop.eccentricity),
        'solidity': float(prop.solidity),
        'equivalent_diameter': float(prop.equivalent_diameter),
        'orientation': float(prop.orientation),
        'bounding_box': [int(x) for x in prop.bbox],
        'convex_area': int(prop.convex_area),
        'euler_number': int(prop.euler_number),
        'provenance': provenance.get(prop.label, {'method': 'original', 'parent_id': 0, 'pass_level': 0})
    }
    per_domain_data.append(domain_info)

print(f"Extracted properties for {len(per_domain_data)} domains")
for d in per_domain_data:
    print(f"  Domain {d['label']}: area={d['area']}, circularity={d['circularity']:.3f}, solidity={d['solidity']:.3f}, provenance={d['provenance']['method']}")

# Step 14: Size distribution analysis
areas = np.array([d['area'] for d in per_domain_data])
area_mean = float(np.mean(areas))
area_std = float(np.std(areas))
area_skewness = float(stats.skew(areas)) if len(areas) > 2 else 0.0
area_kurtosis = float(stats.kurtosis(areas)) if len(areas) > 2 else 0.0

# KDE for peak detection
try:
    from scipy.stats import gaussian_kde
    if len(areas) >= 3:
        kde = gaussian_kde(areas, bw_method='silverman')
        x_kde = np.linspace(areas.min() * 0.5, areas.max() * 1.5, 500)
        y_kde = kde(x_kde)
        # Find peaks
        from scipy.signal import find_peaks
        peaks, _ = find_peaks(y_kde, prominence=0.00001)
        kde_peak_positions = x_kde[peaks].tolist()
        kde_peak_count = len(peaks)
    else:
        kde_peak_positions = [float(np.mean(areas))]
        kde_peak_count = 1
        x_kde = np.linspace(areas.min(), areas.max(), 100)
        y_kde = np.ones(100)
except:
    kde_peak_positions = [float(np.mean(areas))]
    kde_peak_count = 1
    x_kde = np.linspace(areas.min(), areas.max(), 100)
    y_kde = np.ones(100)

# Hartigan's dip test (simplified)
try:
    sorted_areas = np.sort(areas)
    n = len(sorted_areas)
    if n >= 4:
        ecdf = np.arange(1, n + 1) / n
        # Simplified dip: max deviation from uniform
        uniform = np.linspace(0, 1, n)
        dip = float(np.max(np.abs(ecdf - uniform)))
    else:
        dip = 0.0
except:
    dip = 0.0

# Step 15: Correlations
circularities = np.array([d['circularity'] for d in per_domain_data])
aspect_ratios = np.array([d['aspect_ratio'] for d in per_domain_data])
solidities = np.array([d['solidity'] for d in per_domain_data])

if len(areas) >= 3:
    circ_area_corr = float(np.corrcoef(areas, circularities)[0, 1]) if np.std(circularities) > 0 else 0.0
    ar_area_corr = float(np.corrcoef(areas, aspect_ratios)[0, 1]) if np.std(aspect_ratios) > 0 else 0.0
    sol_area_corr = float(np.corrcoef(areas, solidities)[0, 1]) if np.std(solidities) > 0 else 0.0
else:
    circ_area_corr = ar_area_corr = sol_area_corr = 0.0

# Step 16: Local phase fraction heat map
window_size = 64
stride = 8
h, w = image.shape[:2]

ny_steps = (h - window_size) // stride + 1
nx_steps = (w - window_size) // stride + 1

phase_fraction_map = np.zeros((ny_steps, nx_steps))
for iy in range(ny_steps):
    for ix in range(nx_steps):
        y0 = iy * stride
        x0 = ix * stride
        window = binary_clean[y0:y0+window_size, x0:x0+window_size]
        phase_fraction_map[iy, ix] = window.sum() / (window_size * window_size)

# Apply bilinear interpolation to smooth the 2D map (as specified in plan step 16)
from scipy.ndimage import zoom
smooth_factor = 2  # Upsample by factor of 2 using bilinear interpolation
phase_fraction_map_smooth = zoom(phase_fraction_map, smooth_factor, order=1)  # order=1 = bilinear

# Gradient of phase fraction
gy, gx = np.gradient(phase_fraction_map)
gradient_mag = np.sqrt(gy**2 + gx**2)
max_grad_idx = np.unravel_index(np.argmax(gradient_mag), gradient_mag.shape)
max_grad_mag = float(gradient_mag[max_grad_idx])
max_grad_dir = float(np.arctan2(gy[max_grad_idx], gx[max_grad_idx]) * 180 / np.pi)

# Step 17: Nearest-neighbor centroid distances
centroids = np.array([d['centroid'] for d in per_domain_data])
if len(centroids) >= 2:
    from scipy.spatial.distance import cdist
    dist_matrix = cdist(centroids, centroids)
    np.fill_diagonal(dist_matrix, np.inf)
    nn_distances = np.min(dist_matrix, axis=1)
    nn_mean = float(np.mean(nn_distances))
    nn_std = float(np.std(nn_distances))
    nn_cv = nn_std / nn_mean if nn_mean > 0 else 0.0
else:
    nn_distances = np.array([0.0])
    nn_mean = 0.0
    nn_std = 0.0
    nn_cv = 0.0

# Identify central vs peripheral domains
image_center = np.array([h/2, w/2])
centroid_dists_to_center = np.array([np.linalg.norm(np.array(d['centroid']) - image_center) for d in per_domain_data])
median_dist = np.median(centroid_dists_to_center)

central_indices = centroid_dists_to_center < median_dist
peripheral_indices = ~central_indices

central_nn = nn_distances[central_indices] if central_indices.sum() > 0 else np.array([0.0])
peripheral_nn = nn_distances[peripheral_indices] if peripheral_indices.sum() > 0 else np.array([0.0])

central_spacing = float(np.mean(central_nn)) if len(central_nn) > 0 else 0.0
peripheral_spacing = float(np.mean(peripheral_nn)) if len(peripheral_nn) > 0 else 0.0

# NN distance modality
if len(nn_distances) >= 4:
    nn_sorted = np.sort(nn_distances)
    nn_ecdf = np.arange(1, len(nn_sorted) + 1) / len(nn_sorted)
    nn_uniform = np.linspace(0, 1, len(nn_sorted))
    nn_dip = float(np.max(np.abs(nn_ecdf - nn_uniform)))
    nn_modality = 'bimodal' if nn_dip > 0.15 else 'unimodal'
else:
    nn_dip = 0.0
    nn_modality = 'insufficient_data'

# Step 18: Visualization - Composite with all required plots including scatter plots
fig, axes = plt.subplots(3, 3, figsize=(20, 18))

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

# Row 0, Col 1: Segmentation overlay with labeled domains
overlay = np.stack([image]*3, axis=-1) if image.ndim == 2 else image.copy()
overlay = (overlay * 255).astype(np.uint8) if overlay.max() <= 1.0 else overlay.astype(np.uint8)

# Color each domain
np.random.seed(42)
colors = np.random.randint(50, 255, size=(final_domain_count + 1, 3))

overlay_colored = overlay.copy().astype(np.float32)
for d in per_domain_data:
    mask = final_label_map == d['label']
    color = colors[d['label']]
    for c in range(3):
        overlay_colored[:,:,c] = np.where(mask, 
            overlay_colored[:,:,c] * 0.5 + color[c] * 0.5, 
            overlay_colored[:,:,c])

# Draw contours
overlay_uint8 = np.clip(overlay_colored, 0, 255).astype(np.uint8)
for d in per_domain_data:
    mask = (final_label_map == d['label']).astype(np.uint8)
    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    cv2.drawContours(overlay_uint8, contours, -1, (255, 255, 0), 1)

axes[0, 1].imshow(overlay_uint8)
axes[0, 1].set_title(f'Segmented Domains ({final_domain_count})')
axes[0, 1].axis('off')

# Row 0, Col 2: Provenance-colored map
prov_colors = {'original': [0, 0, 1], 'bridge_cut': [1, 0, 0], 'erosion_split': [0, 1, 0], 'watershed': [1, 1, 0]}
prov_map = np.zeros((*final_label_map.shape, 3))
for d in per_domain_data:
    mask = final_label_map == d['label']
    method = d['provenance']['method']
    color = prov_colors.get(method, [0.5, 0.5, 0.5])
    for c in range(3):
        prov_map[:,:,c] = np.where(mask, color[c], prov_map[:,:,c])

axes[0, 2].imshow(prov_map)
axes[0, 2].set_title('Domain Provenance\n(blue=orig, red=bridge, green=erosion, yellow=ws)')
axes[0, 2].axis('off')

# Row 1, Col 0: Phase fraction heat map
im_pf = axes[1, 0].imshow(phase_fraction_map_smooth, cmap='hot', interpolation='bilinear')
axes[1, 0].set_title('Local Phase Fraction')
plt.colorbar(im_pf, ax=axes[1, 0], fraction=0.046)

# Row 1, Col 1: Size distribution with KDE
axes[1, 1].hist(areas, bins=max(5, final_domain_count//2), alpha=0.7, color='steelblue', edgecolor='black')
if len(areas) >= 3:
    ax2 = axes[1, 1].twinx()
    ax2.plot(x_kde, y_kde, 'r-', linewidth=2, label='KDE')
    ax2.set_ylabel('Density', color='red')
axes[1, 1].set_xlabel('Domain Area (pixels)')
axes[1, 1].set_ylabel('Count')
axes[1, 1].set_title(f'Size Distribution (mean={area_mean:.0f}, std={area_std:.0f})')

# Row 1, Col 2: Nearest-neighbor histogram
axes[1, 2].hist(nn_distances, bins=max(5, final_domain_count//2), alpha=0.7, color='seagreen', edgecolor='black')
axes[1, 2].axvline(nn_mean, color='red', linestyle='--', label=f'Mean={nn_mean:.1f}')
axes[1, 2].set_xlabel('Nearest Neighbor Distance (pixels)')
axes[1, 2].set_ylabel('Count')
axes[1, 2].set_title(f'NN Distance (CV={nn_cv:.2f})')
axes[1, 2].legend()

# Row 2, Col 0: Scatter plot - Circularity vs Area (FIX: added per plan Step 15)
axes[2, 0].scatter(areas, circularities, c='steelblue', edgecolors='black', s=60, alpha=0.8)
axes[2, 0].set_xlabel('Domain Area (pixels)')
axes[2, 0].set_ylabel('Circularity')
axes[2, 0].set_title(f'Circularity vs Area (r={circ_area_corr:.3f})')
axes[2, 0].grid(True, alpha=0.3)

# Row 2, Col 1: Scatter plot - Aspect Ratio vs Area (FIX: added per plan Step 15)
axes[2, 1].scatter(areas, aspect_ratios, c='darkorange', edgecolors='black', s=60, alpha=0.8)
axes[2, 1].set_xlabel('Domain Area (pixels)')
axes[2, 1].set_ylabel('Aspect Ratio')
axes[2, 1].set_title(f'Aspect Ratio vs Area (r={ar_area_corr:.3f})')
axes[2, 1].grid(True, alpha=0.3)

# Row 2, Col 2: Scatter plot - Solidity vs Area (FIX: added per plan Step 15)
axes[2, 2].scatter(areas, solidities, c='mediumpurple', edgecolors='black', s=60, alpha=0.8)
axes[2, 2].set_xlabel('Domain Area (pixels)')
axes[2, 2].set_ylabel('Solidity')
axes[2, 2].set_title(f'Solidity vs Area (r={sol_area_corr:.3f})')
axes[2, 2].grid(True, alpha=0.3)

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

# Step 19: Compile results
# FIX: Include local_phase_fraction_heatmap_2d_array in extracted features
extracted_features = {
    'domain_count': final_domain_count,
    'per_domain_area_pixels': [d['area'] for d in per_domain_data],
    'per_domain_area_fraction_of_dark_phase': [round(d['area_fraction'], 4) for d in per_domain_data],
    'per_domain_centroid_xy': [d['centroid'] for d in per_domain_data],
    'per_domain_major_axis_length': [round(d['major_axis_length'], 2) for d in per_domain_data],
    'per_domain_minor_axis_length': [round(d['minor_axis_length'], 2) for d in per_domain_data],
    'per_domain_aspect_ratio': [round(d['aspect_ratio'], 3) for d in per_domain_data],
    'per_domain_perimeter': [round(d['perimeter'], 2) for d in per_domain_data],
    'per_domain_perimeter_smoothed_convex_hull': [round(d['perimeter_smoothed_convex_hull'], 2) for d in per_domain_data],
    'per_domain_circularity': [round(d['circularity'], 4) for d in per_domain_data],
    'per_domain_circularity_smoothed': [round(d['circularity_smoothed'], 4) for d in per_domain_data],
    'per_domain_eccentricity': [round(d['eccentricity'], 4) for d in per_domain_data],
    'per_domain_solidity': [round(d['solidity'], 4) for d in per_domain_data],
    'per_domain_equivalent_diameter': [round(d['equivalent_diameter'], 2) for d in per_domain_data],
    'per_domain_orientation': [round(d['orientation'], 4) for d in per_domain_data],
    'per_domain_bounding_box': [d['bounding_box'] for d in per_domain_data],
    'per_domain_convex_area': [d['convex_area'] for d in per_domain_data],
    'per_domain_euler_number': [d['euler_number'] for d in per_domain_data],
    'per_domain_provenance_split_method': [d['provenance']['method'] for d in per_domain_data],
    'per_domain_provenance_parent_id': [d['provenance']['parent_id'] for d in per_domain_data],
    'per_domain_provenance_pass_level': [d['provenance']['pass_level'] for d in per_domain_data],
    'global_dark_phase_fraction': round(float(total_dark_area / (h * w)), 4),
    'domain_size_distribution_mean': round(area_mean, 2),
    'domain_size_distribution_std': round(area_std, 2),
    'domain_size_distribution_skewness': round(area_skewness, 4),
    'domain_size_distribution_kurtosis': round(area_kurtosis, 4),
    'domain_size_kde_peak_positions': [round(p, 1) for p in kde_peak_positions],
    'domain_size_kde_peak_count': kde_peak_count,
    'domain_size_multimodality_dip_statistic': round(dip, 4),
    'circularity_vs_area_correlation': round(circ_area_corr, 4),
    'aspect_ratio_vs_area_correlation': round(ar_area_corr, 4),
    'solidity_vs_area_correlation': round(sol_area_corr, 4),
    'local_phase_fraction_heatmap_2d_array': phase_fraction_map.tolist(),
    'local_phase_fraction_gradient_magnitude_max': round(max_grad_mag, 4),
    'local_phase_fraction_gradient_direction_at_max': round(max_grad_dir, 2),
    'nearest_neighbor_centroid_distances': [round(float(d), 2) for d in nn_distances],
    'nearest_neighbor_distance_mean': round(nn_mean, 2),
    'nearest_neighbor_distance_std': round(nn_std, 2),
    'nearest_neighbor_distance_cv': round(nn_cv, 4),
    'nearest_neighbor_distance_distribution_modality': nn_modality,
    'central_cluster_inter_domain_spacing': round(central_spacing, 2),
    'peripheral_inter_domain_spacing': round(peripheral_spacing, 2),
    'skeleton_constriction_points_detected': skeleton_constriction_points_detected,
    'skeleton_constriction_widths_at_cuts': [round(float(w_val), 2) for w_val in skeleton_constriction_widths[:20]],
    'bridge_breaking_cuts_applied_count': bridge_breaking_cuts_count,
    'erosion_fallback_used_boolean': erosion_fallback_used,
    'final_domain_count_within_target_range_boolean': in_target_range
}

quality_metrics = {
    'domain_count': final_domain_count,
    'target_range': '12-16',
    'in_target_range': in_target_range,
    'total_dark_phase_fraction': round(float(total_dark_area / (h * w)), 4),
    'mean_circularity': round(float(np.mean(circularities)), 4),
    'mean_solidity': round(float(np.mean(solidities)), 4),
    'nn_distance_cv': round(nn_cv, 4),
    'size_distribution_skewness': round(area_skewness, 4),
    'initial_component_count': len(initial_domains),
    'constriction_points_found': skeleton_constriction_points_detected,
    'cuts_applied': bridge_breaking_cuts_count,
    'erosion_fallback': erosion_fallback_used
}

results = {
    'analysis_type': 'Dark-phase domain morphometric analysis with skeleton-based constriction-point bridge-breaking segmentation and spatial phase fraction gradient mapping',
    'extracted_features': extracted_features,
    'quality_metrics': quality_metrics,
    'summary': f'Segmented {final_domain_count} dark-phase domains (global fraction {total_dark_area/(h*w):.3f}) using skeleton constriction-point bridge-breaking ({bridge_breaking_cuts_count} cuts, {skeleton_constriction_points_detected} constriction points detected, erosion fallback={erosion_fallback_used}). Mean domain area {area_mean:.0f}px (std {area_std:.0f}), mean circularity {np.mean(circularities):.3f}, NN distance mean {nn_mean:.1f}px (CV={nn_cv:.3f}). Central cluster spacing {central_spacing:.1f}px vs peripheral {peripheral_spacing:.1f}px. Phase fraction gradient max {max_grad_mag:.4f} at {max_grad_dir:.1f} degrees. Parameters: Gaussian sigma=2, Otsu threshold={thresh:.4f}, morphological opening disk(4), constriction width thresholds 6-10px.'
}

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