import numpy as np
import json
import os
import warnings
warnings.filterwarnings('ignore')

import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import matplotlib.colors as mcolors

from scipy import ndimage
from scipy.stats import pearsonr, ks_2samp
from scipy.ndimage import distance_transform_edt, gaussian_filter

import skimage
from skimage import measure, morphology, feature, segmentation, filters, exposure
from skimage.feature import peak_local_max
from skimage.segmentation import watershed

import cv2
import pandas as pd

# ============================================================
# 1. Load image and Tier 1 label map
# ============================================================
image_path = '/Users/maxim.ziatdinov/Code/SciLink/analysis_session_20260401_122257/results/analysis__easy__vaterite_ImageAnalysis_20260401_122418_001/image_0000/temp_image_0.npy'
image = np.load(image_path)
print(f"Image shape: {image.shape}, dtype: {image.dtype}")

# Load Tier 1 label map
label_path = 'analysis_labels.npy'
if os.path.exists(label_path):
    tier1_labels = np.load(label_path)
    print(f"Tier 1 labels shape: {tier1_labels.shape}, dtype: {tier1_labels.dtype}")
else:
    raise FileNotFoundError(f"Tier 1 label map not found at {label_path}")

# ============================================================
# 2. Calibration
# ============================================================
H, W = image.shape[:2]
scale_x = 165.0 / W  # um/px in x direction
scale_y = 124.0 / H  # um/px in y direction
pixel_area_um2 = scale_x * scale_y
scale_avg = (scale_x + scale_y) / 2.0  # average scale for length measurements

print(f"Scale: x={scale_x:.6f} um/px, y={scale_y:.6f} um/px, avg={scale_avg:.6f} um/px")
print(f"Pixel area: {pixel_area_um2:.8f} um^2")

# ============================================================
# 3. Compute regionprops for all Tier 1 labels
# ============================================================
# Convert image to grayscale for intensity measurements
if image.ndim == 3 and image.shape[2] == 3:
    gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY).astype(float)
else:
    gray = image[:,:,0].astype(float)

props = measure.regionprops(tier1_labels, intensity_image=gray)
print(f"Number of Tier 1 labels: {len(props)}")

def compute_local_background(gray_img, mask, label_map, label_id, annular_width=20):
    """Compute local background intensity in an annular ring around the object."""
    # Dilate mask
    obj_mask = (label_map == label_id).astype(np.uint8)
    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (2*annular_width+1, 2*annular_width+1))
    dilated = cv2.dilate(obj_mask, kernel, iterations=1)
    annular = (dilated > 0) & (label_map == 0)  # annular ring in background
    if np.sum(annular) > 0:
        return np.mean(gray_img[annular])
    else:
        return np.mean(gray_img)

# Build feature table
features_list = []
for p in props:
    label_id = p.label
    area_px = p.area
    perimeter_px = p.perimeter
    
    # Convert to physical units
    area_um2 = area_px * pixel_area_um2
    perimeter_um = perimeter_px * scale_avg
    major_um = p.major_axis_length * scale_avg
    minor_um = p.minor_axis_length * scale_avg
    eq_diam_um = 2.0 * np.sqrt(area_um2 / np.pi)
    
    # Circularity
    if perimeter_px > 0:
        circularity = 4.0 * np.pi * area_px / (perimeter_px ** 2)
    else:
        circularity = 0.0
    circularity = min(circularity, 1.0)  # cap at 1.0
    
    aspect_ratio = minor_um / major_um if major_um > 0 else 0.0
    
    # Local background
    local_bg = compute_local_background(gray, tier1_labels == label_id, tier1_labels, label_id)
    mean_int = p.mean_intensity
    contrast_ratio = mean_int / local_bg if local_bg > 0 else 1.0
    
    cy, cx = p.centroid
    
    feat = {
        'label_id': label_id,
        'area_px': area_px,
        'area_um2': area_um2,
        'perimeter_px': perimeter_px,
        'perimeter_um': perimeter_um,
        'major_axis_length_um': major_um,
        'minor_axis_length_um': minor_um,
        'equivalent_diameter_um': eq_diam_um,
        'circularity': circularity,
        'aspect_ratio': aspect_ratio,
        'solidity': p.solidity,
        'eccentricity': p.eccentricity,
        'orientation_deg': np.degrees(p.orientation),
        'centroid_row': cy,
        'centroid_col': cx,
        'mean_intensity': mean_int,
        'local_background_intensity': local_bg,
        'intensity_contrast_ratio': contrast_ratio,
        'morphological_class': '',
        'quality_flag': 'passed',
        'detection_source': 'tier1_original'
    }
    features_list.append(feat)

df = pd.DataFrame(features_list)

# ============================================================
# 4. Flag problematic labels
# ============================================================
# Merged masks: solidity < 0.65
merged_mask = df['solidity'] < 0.65
merged_labels = df.loc[merged_mask, 'label_id'].tolist()
df.loc[merged_mask, 'quality_flag'] = 'merged'
print(f"Merged mask labels (solidity < 0.65): {merged_labels}")

# Sub-threshold labels: equivalent_diameter < 5 um
sub_thresh = df['equivalent_diameter_um'] < 5.0
df.loc[sub_thresh & (df['quality_flag'] == 'passed'), 'quality_flag'] = 'sub_threshold'

# Flag ambiguous small dark speck near top center
# Top center: row < H/4, col in middle third
top_center = (df['centroid_row'] < H/4) & (df['centroid_col'] > W/3) & (df['centroid_col'] < 2*W/3)
small_dark = (df['equivalent_diameter_um'] < 5.0) | (df['mean_intensity'] < np.percentile(gray[gray > 0], 10))
uncertain = top_center & small_dark & (df['equivalent_diameter_um'] < 7.0)
df.loc[uncertain & (df['quality_flag'] == 'passed'), 'quality_flag'] = 'uncertain_debris'

print(f"Sub-threshold labels: {df.loc[sub_thresh, 'label_id'].tolist()}")
print(f"Uncertain debris: {df.loc[df['quality_flag'] == 'uncertain_debris', 'label_id'].tolist()}")

# ============================================================
# 5. Classify face-on vs edge-on
# ============================================================
valid_for_class = df['quality_flag'].isin(['passed', 'uncertain_debris'])
df.loc[valid_for_class & (df['circularity'] > 0.80), 'morphological_class'] = 'face-on'
df.loc[valid_for_class & (df['circularity'] <= 0.80), 'morphological_class'] = 'edge-on'

face_on = df[(df['morphological_class'] == 'face-on') & (df['quality_flag'] == 'passed')]
edge_on = df[(df['morphological_class'] == 'edge-on') & (df['quality_flag'] == 'passed')]

print(f"\nFace-on particles (circularity > 0.80): N={len(face_on)}")
if len(face_on) > 0:
    print(f"  Diameter: mean={face_on['equivalent_diameter_um'].mean():.2f}, std={face_on['equivalent_diameter_um'].std():.2f}, range=[{face_on['equivalent_diameter_um'].min():.2f}, {face_on['equivalent_diameter_um'].max():.2f}]")

print(f"Edge-on particles (circularity <= 0.80): N={len(edge_on)}")
if len(edge_on) > 0:
    print(f"  Diameter: mean={edge_on['equivalent_diameter_um'].mean():.2f}, std={edge_on['equivalent_diameter_um'].std():.2f}, range=[{edge_on['equivalent_diameter_um'].min():.2f}, {edge_on['equivalent_diameter_um'].max():.2f}]")

# ============================================================
# 6. Pearson correlation (diameter >= 5 um)
# ============================================================
filter_corr = (df['equivalent_diameter_um'] >= 5.0) & (df['quality_flag'].isin(['passed', 'uncertain_debris']))
df_corr = df[filter_corr]

if len(df_corr) >= 3:
    r_val, p_corr = pearsonr(df_corr['equivalent_diameter_um'], df_corr['circularity'])
else:
    r_val, p_corr = 0.0, 1.0
print(f"\nPearson r (diameter vs circularity, >=5um): r={r_val:.4f}, p={p_corr:.4f}")

# ============================================================
# 7. Watershed splitting for merged masks
# ============================================================
corrected_labels = tier1_labels.copy()
next_label = tier1_labels.max() + 1
new_features = []

for ml in merged_labels:
    mask_ml = (tier1_labels == ml).astype(np.uint8)
    if np.sum(mask_ml) == 0:
        continue
    
    dist = distance_transform_edt(mask_ml)
    
    # Try with min_distance = 37 px (~3 um)
    min_dist_px = int(3.0 / scale_avg)
    coords = peak_local_max(dist, min_distance=min_dist_px, labels=mask_ml)
    
    if len(coords) < 2:
        # Reduce to 2 um and smooth
        min_dist_px2 = int(2.0 / scale_avg)
        dist_smooth = gaussian_filter(dist, sigma=2)
        coords = peak_local_max(dist_smooth, min_distance=min_dist_px2, labels=mask_ml)
    
    if len(coords) < 2:
        # h-minima approach
        h_val = 0.3 * dist.max()
        if h_val > 0:
            dist_smooth = gaussian_filter(dist, sigma=2)
            h_min = morphology.h_maxima(dist_smooth, h=h_val)
            h_labels = ndimage.label(h_min)[0]
            coords_hm = ndimage.center_of_mass(h_min, h_labels, range(1, h_labels.max()+1))
            if len(coords_hm) >= 2:
                coords = np.array(coords_hm).astype(int)
    
    if len(coords) >= 2:
        markers = np.zeros_like(mask_ml, dtype=np.int32)
        for i, c in enumerate(coords):
            markers[int(c[0]), int(c[1])] = i + 1
        
        # Use Sobel gradient for better watershed
        gray_region = gray * mask_ml.astype(float)
        sobel_mag = filters.sobel(gray_region)
        ws = watershed(sobel_mag, markers, mask=mask_ml.astype(bool))
        
        # Evaluate split fragments
        valid_splits = []
        for frag_id in range(1, ws.max() + 1):
            frag_mask = (ws == frag_id)
            frag_props = measure.regionprops(frag_mask.astype(int), intensity_image=gray)
            if len(frag_props) == 0:
                continue
            fp = frag_props[0]
            frag_area_um2 = fp.area * pixel_area_um2
            frag_diam_um = 2.0 * np.sqrt(frag_area_um2 / np.pi)
            if frag_diam_um >= 5.0 and fp.solidity >= 0.80:
                valid_splits.append((frag_id, frag_mask, fp))
        
        if len(valid_splits) >= 2:
            # Remove old label
            corrected_labels[tier1_labels == ml] = 0
            
            for frag_id, frag_mask, fp in valid_splits:
                corrected_labels[frag_mask] = next_label
                
                area_um2_f = fp.area * pixel_area_um2
                perim_um_f = fp.perimeter * scale_avg
                major_f = fp.major_axis_length * scale_avg
                minor_f = fp.minor_axis_length * scale_avg
                eq_diam_f = 2.0 * np.sqrt(area_um2_f / np.pi)
                circ_f = min(4.0 * np.pi * fp.area / (fp.perimeter ** 2), 1.0) if fp.perimeter > 0 else 0
                ar_f = minor_f / major_f if major_f > 0 else 0
                cy_f, cx_f = fp.centroid
                local_bg_f = compute_local_background(gray, frag_mask, corrected_labels, next_label)
                contrast_f = fp.mean_intensity / local_bg_f if local_bg_f > 0 else 1.0
                morph_class_f = 'face-on' if circ_f > 0.80 else 'edge-on'
                
                new_feat = {
                    'label_id': next_label,
                    'area_px': fp.area,
                    'area_um2': area_um2_f,
                    'perimeter_px': fp.perimeter,
                    'perimeter_um': perim_um_f,
                    'major_axis_length_um': major_f,
                    'minor_axis_length_um': minor_f,
                    'equivalent_diameter_um': eq_diam_f,
                    'circularity': circ_f,
                    'aspect_ratio': ar_f,
                    'solidity': fp.solidity,
                    'eccentricity': fp.eccentricity,
                    'orientation_deg': np.degrees(fp.orientation),
                    'centroid_row': cy_f,
                    'centroid_col': cx_f,
                    'mean_intensity': fp.mean_intensity,
                    'local_background_intensity': local_bg_f,
                    'intensity_contrast_ratio': contrast_f,
                    'morphological_class': morph_class_f,
                    'quality_flag': 'passed',
                    'detection_source': 'watershed_split'
                }
                new_features.append(new_feat)
                next_label += 1
                
            # Remove the merged label from df
            print(f"  Label {ml}: split into {len(valid_splits)} fragments")
        else:
            print(f"  Label {ml}: watershed split failed (only {len(valid_splits)} valid fragments)")
    else:
        print(f"  Label {ml}: insufficient peaks for watershed ({len(coords)} found)")

# ============================================================
# 8. Detect missed dark edge-on particles
# ============================================================
# Use darkest channel (blue channel = index 2, or find darkest)
if image.ndim == 3 and image.shape[2] >= 3:
    blue_ch = image[:,:,2].astype(float)
    green_ch = image[:,:,1].astype(float)
    # Use green channel which often shows better contrast for dark particles
    detect_ch = green_ch
else:
    detect_ch = gray

# Threshold to find dark regions
bg_mean = np.mean(detect_ch[corrected_labels == 0]) if np.sum(corrected_labels == 0) > 0 else np.mean(detect_ch)
bg_std = np.std(detect_ch[corrected_labels == 0]) if np.sum(corrected_labels == 0) > 0 else np.std(detect_ch)

dark_thresh = bg_mean - 1.5 * bg_std
dark_mask = (detect_ch < dark_thresh).astype(np.uint8)

# Remove regions already in label map
existing_mask = (corrected_labels > 0).astype(np.uint8)
# Dilate existing mask slightly to avoid edge detections
kernel_small = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
existing_dilated = cv2.dilate(existing_mask, kernel_small, iterations=2)
new_dark = dark_mask & (~existing_dilated.astype(bool)).astype(np.uint8)

# Morphological cleanup
new_dark = cv2.morphologyEx(new_dark, cv2.MORPH_OPEN, cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3,3)))
new_dark = cv2.morphologyEx(new_dark, cv2.MORPH_CLOSE, cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5,5)))

new_dark_labels = ndimage.label(new_dark)[0]
new_dark_props = measure.regionprops(new_dark_labels, intensity_image=gray)

rejected_candidates = []  # For visualization
accepted_new = 0

for ndp in new_dark_props:
    area_um2_n = ndp.area * pixel_area_um2
    eq_diam_n = 2.0 * np.sqrt(area_um2_n / np.pi)
    cy_n, cx_n = ndp.centroid
    
    # Upper-left quadrant check
    in_upper_left = (cy_n < H/4) and (cx_n < W/4)
    
    min_diam = 8.0 if in_upper_left else 5.0
    
    # Apply filters
    if eq_diam_n < min_diam:
        rejected_candidates.append((cy_n, cx_n, 'too_small'))
        continue
    if ndp.solidity < 0.80:
        rejected_candidates.append((cy_n, cx_n, 'low_solidity'))
        continue
    if ndp.eccentricity < 0.6:
        rejected_candidates.append((cy_n, cx_n, 'low_eccentricity'))
        continue
    
    # Intensity contrast check
    local_bg_n = compute_local_background(gray, new_dark_labels == ndp.label, corrected_labels, 0)
    if local_bg_n > 0:
        if ndp.mean_intensity >= 0.85 * local_bg_n:
            rejected_candidates.append((cy_n, cx_n, 'low_contrast'))
            continue
    
    # Passed all filters - add to corrected labels
    frag_mask_n = (new_dark_labels == ndp.label)
    corrected_labels[frag_mask_n] = next_label
    
    perim_n = ndp.perimeter
    circ_n = min(4.0 * np.pi * ndp.area / (perim_n ** 2), 1.0) if perim_n > 0 else 0
    major_n = ndp.major_axis_length * scale_avg
    minor_n = ndp.minor_axis_length * scale_avg
    ar_n = minor_n / major_n if major_n > 0 else 0
    morph_n = 'face-on' if circ_n > 0.80 else 'edge-on'
    contrast_n = ndp.mean_intensity / local_bg_n if local_bg_n > 0 else 1.0
    
    new_feat_n = {
        'label_id': next_label,
        'area_px': ndp.area,
        'area_um2': area_um2_n,
        'perimeter_px': perim_n,
        'perimeter_um': perim_n * scale_avg,
        'major_axis_length_um': major_n,
        'minor_axis_length_um': minor_n,
        'equivalent_diameter_um': eq_diam_n,
        'circularity': circ_n,
        'aspect_ratio': ar_n,
        'solidity': ndp.solidity,
        'eccentricity': ndp.eccentricity,
        'orientation_deg': np.degrees(ndp.orientation),
        'centroid_row': cy_n,
        'centroid_col': cx_n,
        'mean_intensity': ndp.mean_intensity,
        'local_background_intensity': local_bg_n,
        'intensity_contrast_ratio': contrast_n,
        'morphological_class': morph_n,
        'quality_flag': 'passed',
        'detection_source': 'missed_detection'
    }
    new_features.append(new_feat_n)
    accepted_new += 1
    next_label += 1

print(f"\nMissed edge-on detection: {accepted_new} new particles added, {len(rejected_candidates)} candidates rejected")

# Combine feature tables
if new_features:
    df_new = pd.DataFrame(new_features)
    # Remove merged labels from original df
    df_final = pd.concat([df[~df['label_id'].isin(merged_labels)], df_new], ignore_index=True)
else:
    df_final = df[~df['label_id'].isin(merged_labels)].copy()

# Re-classify for the final table
valid_final = df_final['quality_flag'].isin(['passed', 'uncertain_debris'])
df_final.loc[valid_final & (df_final['circularity'] > 0.80) & (df_final['morphological_class'] == ''), 'morphological_class'] = 'face-on'
df_final.loc[valid_final & (df_final['circularity'] <= 0.80) & (df_final['morphological_class'] == ''), 'morphological_class'] = 'edge-on'

total_corrected = len(df_final)
print(f"\nTotal corrected particle count: {total_corrected}")

# ============================================================
# 9. Oblate spheroid aspect ratio estimation
# ============================================================
# Use high-quality Tier 1 particles
hq_mask = (df['solidity'] > 0.85) & (df['equivalent_diameter_um'] >= 5.0) & (~df['label_id'].isin(merged_labels))
df_hq = df[hq_mask].copy()
print(f"\nHigh-quality particles for oblate fit: N={len(df_hq)}")

observed_ar = df_hq['aspect_ratio'].values

def simulate_projected_aspect_ratios(c_over_a, n_samples=10000):
    """Simulate projected aspect ratios for randomly oriented oblate spheroids."""
    # theta sampled from P(theta) ~ sin(theta), 0 to pi/2
    # For oblate spheroid with semi-axes a, a, c (c < a)
    u = np.random.uniform(0, 1, n_samples)
    theta = np.arccos(u)  # This gives P(theta) ~ sin(theta)
    
    c = c_over_a  # normalized, a=1
    a = 1.0
    
    # Projected semi-axes for oblate spheroid viewed at angle theta from pole
    # Major projected axis = a
    # Minor projected axis = sqrt(a^2 * sin^2(theta) + c^2 * cos^2(theta)) ... wait
    # For oblate spheroid (a=b>c), viewed at inclination theta from the symmetry axis (c):
    # The projected ellipse has semi-axes:
    #   a_proj = a (always, the equatorial radius perpendicular to viewing direction)
    #   b_proj = sqrt(c^2 * sin^2(theta) + a^2 * cos^2(theta))... let me think carefully
    
    # Actually for oblate spheroid with semi-axes (a, a, c), c < a:
    # Viewed at angle theta from the c-axis (symmetry axis):
    # One projected axis is always a (perpendicular to both c-axis and line of sight)
    # Other projected axis is sqrt(a^2 * sin^2(theta) + c^2 * cos^2(theta))... no
    # 
    # When theta=0 (face-on, looking along c-axis): circle of radius a, aspect ratio = 1
    # When theta=pi/2 (edge-on): ellipse with axes a and c, aspect ratio = c/a
    #
    # The projection: semi-axes are a and sqrt(c^2*sin^2(theta) + a^2*cos^2(theta))... 
    # Hmm, let me reconsider.
    # 
    # For an oblate spheroid x^2/a^2 + y^2/a^2 + z^2/c^2 = 1, viewed along direction
    # making angle theta with z-axis in the xz plane:
    # The projected ellipse has semi-axes:
    #   semi_1 = a  (along y-axis, always)
    #   semi_2 = a*c / sqrt(c^2*cos^2(theta) + a^2*sin^2(theta))
    # 
    # Check: theta=0: semi_2 = a*c/c = a -> circle, correct
    # theta=pi/2: semi_2 = a*c/a = c -> axes a,c, correct
    
    semi_1 = a * np.ones(n_samples)
    semi_2 = a * c / np.sqrt(c**2 * np.cos(theta)**2 + a**2 * np.sin(theta)**2)
    
    # Aspect ratio = minor/major
    projected_ar = np.minimum(semi_1, semi_2) / np.maximum(semi_1, semi_2)
    
    return projected_ar

best_ca = 0.5
best_ks = 1.0
best_p = 0.0

if len(observed_ar) >= 5:
    ca_range = np.arange(0.10, 0.96, 0.01)
    ks_stats = []
    
    for ca in ca_range:
        sim_ar = simulate_projected_aspect_ratios(ca, n_samples=10000)
        ks_stat, p_val = ks_2samp(observed_ar, sim_ar)
        ks_stats.append((ca, ks_stat, p_val))
    
    ks_stats = sorted(ks_stats, key=lambda x: x[1])
    best_ca, best_ks, best_p = ks_stats[0]
    
    print(f"\nBest-fit oblate c/a = {best_ca:.2f}, KS = {best_ks:.4f}, p = {best_p:.4f}")
    
    # Flag suspect results
    oblate_flag = ''
    if best_ca < 0.2 or best_ca > 0.8:
        oblate_flag = 'suspect_value'
        print(f"  WARNING: c/a = {best_ca:.2f} is outside physically reasonable range [0.2, 0.8]")
    if best_ks > 0.20:
        oblate_flag = 'poor_fit'
        print(f"  WARNING: KS statistic {best_ks:.4f} > 0.20, poor fit")
    if best_p < 0.05:
        oblate_flag += '_low_pvalue'
        print(f"  WARNING: p-value {best_p:.4f} < 0.05")
else:
    print("Insufficient high-quality particles for oblate fit")
    oblate_flag = 'insufficient_data'

# Generate simulated distribution at best fit for plotting
sim_ar_best = simulate_projected_aspect_ratios(best_ca, n_samples=10000)
# Also compute simulated circularity
# Circularity of an ellipse = pi*a*b / ((pi/2)*(3(a+b) - sqrt((3a+b)(a+3b))))**... 
# Actually for an ellipse, circularity = 4*pi*area/perimeter^2
# Area = pi*a*b, Perimeter ~ pi*(3(a+b) - sqrt((3a+b)*(a+3b))) (Ramanujan)
def ellipse_circularity(aspect_ratio):
    """Compute circularity of an ellipse given aspect ratio (minor/major)."""
    a = 1.0  # semi-major
    b = aspect_ratio  # semi-minor
    area = np.pi * a * b
    # Ramanujan approximation for ellipse perimeter
    h = ((a - b) / (a + b)) ** 2
    perimeter = np.pi * (a + b) * (1 + 3 * h / (10 + np.sqrt(4 - 3 * h)))
    circ = 4 * np.pi * area / (perimeter ** 2)
    return np.minimum(circ, 1.0)

observed_circ_hq = df_hq['circularity'].values
sim_circ_best = ellipse_circularity(sim_ar_best)

# ============================================================
# 10. Generate summary figure
# ============================================================
fig, axes = plt.subplots(2, 2, figsize=(16, 12))

# Panel (a): Labeled image with overlays
ax = axes[0, 0]
ax.imshow(image)

# Overlay face-on (green), edge-on (red), new detections (yellow)
overlay = np.zeros((H, W, 4), dtype=float)

for _, row in df_final.iterrows():
    lid = int(row['label_id'])
    mask_i = (corrected_labels == lid)
    if not np.any(mask_i):
        continue
    
    if row['detection_source'] == 'missed_detection':
        color = [1.0, 1.0, 0.0, 0.3]  # yellow
    elif row['morphological_class'] == 'face-on':
        color = [0.0, 1.0, 0.0, 0.25]  # green
    elif row['morphological_class'] == 'edge-on':
        color = [1.0, 0.0, 0.0, 0.25]  # red
    else:
        color = [0.5, 0.5, 0.5, 0.2]  # gray for unclassified
    
    overlay[mask_i] = color
    
    # Draw contour
    contour_mask = mask_i.astype(np.uint8)
    contours_cv, _ = cv2.findContours(contour_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    for cnt in contours_cv:
        pts = cnt.squeeze()
        if pts.ndim == 2:
            ax.plot(pts[:, 0], pts[:, 1], color=color[:3], linewidth=0.8, alpha=0.8)

ax.imshow(overlay)

# Mark rejected candidates as gray crosses
for rc_y, rc_x, rc_reason in rejected_candidates:
    ax.plot(rc_x, rc_y, 'x', color='gray', markersize=4, alpha=0.5)

ax.set_title(f'Segmentation Overlay (N={total_corrected})', fontsize=10)
ax.set_xlabel('Pixels')
ax.set_ylabel('Pixels')

# Legend
from matplotlib.patches import Patch
legend_elements = [
    Patch(facecolor='green', alpha=0.4, label=f'Face-on (N={len(df_final[df_final["morphological_class"]=="face-on"])})'),
    Patch(facecolor='red', alpha=0.4, label=f'Edge-on (N={len(df_final[df_final["morphological_class"]=="edge-on"])})'),
    Patch(facecolor='yellow', alpha=0.4, label=f'New detections (N={accepted_new})'),
]
ax.legend(handles=legend_elements, loc='lower right', fontsize=7)

# Panel (b): Diameter-circularity scatter
ax = axes[0, 1]
df_plot = df_final.copy()
df_above = df_plot[df_plot['equivalent_diameter_um'] >= 5.0]
df_below = df_plot[df_plot['equivalent_diameter_um'] < 5.0]

# Sub-threshold as gray
if len(df_below) > 0:
    ax.scatter(df_below['equivalent_diameter_um'], df_below['circularity'], 
              facecolors='none', edgecolors='gray', s=30, label='< 5 µm', alpha=0.5)

# Color by class
for cls, color, marker in [('face-on', 'green', 'o'), ('edge-on', 'red', 's')]:
    sub = df_above[df_above['morphological_class'] == cls]
    if len(sub) > 0:
        ax.scatter(sub['equivalent_diameter_um'], sub['circularity'], 
                  c=color, s=40, label=f'{cls} (N={len(sub)})', alpha=0.7, marker=marker)

# Regression line
if len(df_above) >= 3:
    from numpy.polynomial import polynomial as P
    x_fit = df_above['equivalent_diameter_um'].values
    y_fit = df_above['circularity'].values
    coeffs = np.polyfit(x_fit, y_fit, 1)
    x_line = np.linspace(x_fit.min(), x_fit.max(), 100)
    ax.plot(x_line, np.polyval(coeffs, x_line), 'k--', alpha=0.5, label=f'r={r_val:.3f}')

ax.axhline(y=0.80, color='blue', linestyle=':', alpha=0.3, label='Face-on/Edge-on threshold')
ax.set_xlabel('Equivalent Diameter (µm)', fontsize=9)
ax.set_ylabel('Circularity', fontsize=9)
ax.set_title(f'Diameter vs Circularity (Pearson r={r_val:.3f})', fontsize=10)
ax.legend(fontsize=7)
ax.set_ylim(0, 1.1)

# Panel (c): Size distribution histogram
ax = axes[1, 0]
df_passed = df_final[df_final['quality_flag'] == 'passed']
df_fo = df_passed[df_passed['morphological_class'] == 'face-on']['equivalent_diameter_um']
df_eo = df_passed[df_passed['morphological_class'] == 'edge-on']['equivalent_diameter_um']

bins = np.arange(0, max(df_passed['equivalent_diameter_um'].max() + 2, 25), 1.5)
if len(df_fo) > 0:
    ax.hist(df_fo, bins=bins, alpha=0.6, color='green', label=f'Face-on (N={len(df_fo)})', edgecolor='darkgreen')
if len(df_eo) > 0:
    ax.hist(df_eo, bins=bins, alpha=0.6, color='red', label=f'Edge-on (N={len(df_eo)})', edgecolor='darkred')
ax.set_xlabel('Equivalent Diameter (µm)', fontsize=9)
ax.set_ylabel('Count', fontsize=9)
ax.set_title('Corrected Size Distribution', fontsize=10)
ax.legend(fontsize=8)

# Panel (d): Observed vs modeled circularity
ax = axes[1, 1]
bins_circ = np.linspace(0.3, 1.05, 25)
if len(observed_circ_hq) > 0:
    ax.hist(observed_circ_hq, bins=bins_circ, density=True, alpha=0.6, color='steelblue', 
            label='Observed', edgecolor='navy')
ax.hist(sim_circ_best, bins=bins_circ, density=True, alpha=0.4, color='orange', 
        label=f'Model (c/a={best_ca:.2f})', edgecolor='darkorange')
ax.set_xlabel('Circularity', fontsize=9)
ax.set_ylabel('Density', fontsize=9)
ax.set_title(f'Circularity: Observed vs Oblate Model\nc/a={best_ca:.2f}, KS={best_ks:.3f}, p={best_p:.3f}', fontsize=10)
ax.legend(fontsize=8)

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

# ============================================================
# 11. Save outputs
# ============================================================
# Add population-level metrics to df
df_final['pearson_r_diameter_circularity'] = r_val
df_final['best_fit_oblate_c_over_a'] = best_ca
df_final['oblate_fit_KS_statistic'] = best_ks
df_final['oblate_fit_p_value'] = best_p

# Save corrected label map
np.save('corrected_labels.npy', corrected_labels.astype(np.int32))

# Save feature table
df_final.to_csv('corrected_feature_table.csv', index=False)
print("Saved corrected_feature_table.csv")

# Also update analysis_labels.npy
np.save('analysis_labels.npy', corrected_labels.astype(np.int32))

# ============================================================
# Print JSON results
# ============================================================
# Compute summary statistics
df_valid = df_final[df_final['quality_flag'] == 'passed']
fo_stats = {}
eo_stats = {}

df_fo_valid = df_valid[df_valid['morphological_class'] == 'face-on']
df_eo_valid = df_valid[df_valid['morphological_class'] == 'edge-on']

if len(df_fo_valid) > 0:
    fo_stats = {
        'count': int(len(df_fo_valid)),
        'mean_diameter_um': round(float(df_fo_valid['equivalent_diameter_um'].mean()), 2),
        'std_diameter_um': round(float(df_fo_valid['equivalent_diameter_um'].std()), 2),
        'min_diameter_um': round(float(df_fo_valid['equivalent_diameter_um'].min()), 2),
        'max_diameter_um': round(float(df_fo_valid['equivalent_diameter_um'].max()), 2),
        'mean_circularity': round(float(df_fo_valid['circularity'].mean()), 3),
    }

if len(df_eo_valid) > 0:
    eo_stats = {
        'count': int(len(df_eo_valid)),
        'mean_diameter_um': round(float(df_eo_valid['equivalent_diameter_um'].mean()), 2),
        'std_diameter_um': round(float(df_eo_valid['equivalent_diameter_um'].std()), 2),
        'min_diameter_um': round(float(df_eo_valid['equivalent_diameter_um'].min()), 2),
        'max_diameter_um': round(float(df_eo_valid['equivalent_diameter_um'].max()), 2),
        'mean_circularity': round(float(df_eo_valid['circularity'].mean()), 3),
    }

results = {
    "analysis_type": "Tier 2 morphometric analysis: face-on/edge-on classification, diameter-circularity correlation, watershed splitting of merged masks, missed particle detection, oblate spheroid aspect ratio estimation",
    "extracted_features": {
        "total_tier1_particles": int(len(props)),
        "total_corrected_particles": int(total_corrected),
        "merged_labels_identified": merged_labels,
        "watershed_split_fragments": int(len([f for f in new_features if f['detection_source'] == 'watershed_split'])),
        "missed_detections_added": accepted_new,
        "rejected_candidates": len(rejected_candidates),
        "face_on_statistics": fo_stats,
        "edge_on_statistics": eo_stats,
        "pearson_r_diameter_circularity": round(float(r_val), 4),
        "pearson_p_value": round(float(p_corr), 4),
        "best_fit_oblate_c_over_a": round(float(best_ca), 3),
        "oblate_fit_KS_statistic": round(float(best_ks), 4),
        "oblate_fit_p_value": round(float(best_p), 4),
        "oblate_fit_flag": oblate_flag if oblate_flag else 'acceptable',
        "scale_x_um_per_px": round(scale_x, 6),
        "scale_y_um_per_px": round(scale_y, 6),
        "pixel_area_um2": round(pixel_area_um2, 8),
        "n_particles_for_oblate_fit": int(len(df_hq)),
    },
    "quality_metrics": {
        "n_quality_passed": int(len(df_final[df_final['quality_flag'] == 'passed'])),
        "n_merged": int(len(df_final[df_final['quality_flag'] == 'merged'])),
        "n_sub_threshold": int(len(df_final[df_final['quality_flag'] == 'sub_threshold'])),
        "n_uncertain_debris": int(len(df_final[df_final['quality_flag'] == 'uncertain_debris'])),
        "n_artifact": int(len(df_final[df_final['quality_flag'] == 'artifact'])),
        "total_in_feature_table": int(len(df_final)),
    },
    "summary": f"Tier 2 analysis of {len(props)} Tier 1 particles. Identified {len(merged_labels)} merged masks (labels {merged_labels}). After watershed splitting and missed detection, corrected count is {total_corrected}. Face-on: N={fo_stats.get('count',0)}, mean diam={fo_stats.get('mean_diameter_um','N/A')} um. Edge-on: N={eo_stats.get('count',0)}, mean diam={eo_stats.get('mean_diameter_um','N/A')} um. Pearson r(diam,circ)={r_val:.3f}. Best-fit oblate c/a={best_ca:.2f} (KS={best_ks:.3f}, p={best_p:.3f}). Solidity threshold 0.65 used for merged mask detection. Circularity threshold 0.80 for face-on/edge-on classification. Minimum 5 um diameter for correlation and oblate fitting. Dark particle detection used green channel intensity thresholding with strict quality filters.",
    "saved_arrays": {
        "analysis_labels.npy": {
            "description": f"Corrected integer label map with {total_corrected} particles (including watershed splits and missed detections), background=0",
            "shape": list(corrected_labels.shape),
            "dtype": "int32"
        },
        "corrected_labels.npy": {
            "description": f"Same corrected integer label map, {total_corrected} particles labeled, background=0",
            "shape": list(corrected_labels.shape),
            "dtype": "int32"
        }
    }
}

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