import numpy as np
import json
import cv2
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from scipy import ndimage
from skimage import measure, morphology, segmentation, feature
from skimage.filters import threshold_otsu, gaussian
import os

# ========== Step 1: Load image and Tier 1 labels ==========
image_path = '/Users/maxim.ziatdinov/Code/SciLink/test_image_annealing_data/output/image_0000/temp_image_0.npy'
image = np.load(image_path)

# Load Tier 1 labels if available
tier1_labels = None
tier1_label_path = 'analysis_labels.npy'
if os.path.exists(tier1_label_path):
    tier1_labels = np.load(tier1_label_path)
    tier1_count = len(np.unique(tier1_labels)) - 1  # exclude background
    print(f"Loaded Tier 1 labels with {tier1_count} particles")
else:
    tier1_count = 0
    print("No Tier 1 labels found, will detect all particles from scratch")

print(f"Image shape: {image.shape}, dtype: {image.dtype}")
print(f"Intensity range: [{image.min():.2f}, {image.max():.2f}]")

# Normalize image to uint8 for thresholding operations
img_min, img_max = image.min(), image.max()
image_uint8 = ((image - img_min) / (img_max - img_min) * 255).astype(np.uint8)

# ========== Step 2: Adaptive and lower-global thresholding ==========
# Otsu threshold for reference
otsu_thresh = threshold_otsu(image_uint8)
print(f"Otsu threshold on uint8: {otsu_thresh}")

# Adaptive Gaussian thresholding to capture dim particles
# block_size must be odd
block_size = 51
adaptive_offset = 10  # Lower offset captures dimmer particles
adaptive_mask = cv2.adaptiveThreshold(
    image_uint8, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
    cv2.THRESH_BINARY, block_size, -adaptive_offset
)
adaptive_mask = adaptive_mask > 0

# Lower global threshold to catch dim particles missed by Otsu
lower_thresh = max(int(otsu_thresh * 0.65), 70)
print(f"Lower global threshold: {lower_thresh}")
lower_mask = image_uint8 > lower_thresh

# Also keep Otsu-level mask
otsu_mask = image_uint8 > otsu_thresh

# ========== Step 3: Combine masks and subtract Tier 1 regions ==========
combined_mask = adaptive_mask | lower_mask | otsu_mask

# Identify tier1-covered regions
if tier1_labels is not None:
    tier1_binary = tier1_labels > 0
    # Dilate tier1 regions slightly to avoid edge re-detection
    tier1_dilated = ndimage.binary_dilation(tier1_binary, iterations=2)
    # New candidate mask: regions NOT covered by Tier 1
    new_candidate_mask = combined_mask & (~tier1_dilated)
else:
    new_candidate_mask = combined_mask
    tier1_binary = np.zeros_like(combined_mask, dtype=bool)

# ========== Step 4: Morphological cleanup on new candidates ==========
# Opening to remove noise, closing to fill holes
new_candidate_mask = morphology.binary_opening(new_candidate_mask, morphology.disk(2))
new_candidate_mask = morphology.binary_closing(new_candidate_mask, morphology.disk(3))
# Fill holes
new_candidate_mask = ndimage.binary_fill_holes(new_candidate_mask)

# ========== Step 5: Label new components and filter ==========
new_labels_raw = measure.label(new_candidate_mask, connectivity=2)
new_props_raw = measure.regionprops(new_labels_raw)

min_area = 50
max_area = 5000  # generous upper bound
min_circularity = 0.3  # relaxed to catch elongated particles too

filtered_new_mask = np.zeros_like(new_candidate_mask, dtype=bool)
for prop in new_props_raw:
    area = prop.area
    perimeter = prop.perimeter
    if perimeter == 0:
        continue
    circularity = 4 * np.pi * area / (perimeter ** 2)
    if min_area <= area <= max_area and circularity >= min_circularity:
        filtered_new_mask[new_labels_raw == prop.label] = True

print(f"New candidate regions after filtering: {measure.label(filtered_new_mask)[1]}")

# ========== Step 6: Watershed splitting for merged particles ==========
distance = ndimage.distance_transform_edt(filtered_new_mask)
# Use Sobel gradient as landscape for watershed (per domain guidance)
smoothed = gaussian(image, sigma=1.5)
sobel_x = ndimage.sobel(smoothed, axis=1)
sobel_y = ndimage.sobel(smoothed, axis=0)
sobel_mag = np.sqrt(sobel_x**2 + sobel_y**2)

# Markers from distance transform peaks
coords = feature.peak_local_max(distance, min_distance=7, labels=filtered_new_mask.astype(int))
marker_mask = np.zeros(distance.shape, dtype=bool)
for c in coords:
    marker_mask[c[0], c[1]] = True
markers_labeled = measure.label(marker_mask)

if markers_labeled.max() > 0:
    new_labels_ws = segmentation.watershed(sobel_mag, markers_labeled, mask=filtered_new_mask)
else:
    new_labels_ws = measure.label(filtered_new_mask)

print(f"New particles after watershed: {new_labels_ws.max()}")

# ========== Step 7: Merge Tier 1 and new labels ==========
if tier1_labels is not None:
    merged_labels = tier1_labels.copy().astype(np.int32)
    max_tier1_label = tier1_labels.max()
else:
    merged_labels = np.zeros(image.shape, dtype=np.int32)
    max_tier1_label = 0

# Add new labels with offset
for new_lbl in range(1, new_labels_ws.max() + 1):
    region_mask = new_labels_ws == new_lbl
    # Check it doesn't heavily overlap with existing labels
    overlap = np.sum(region_mask & (merged_labels > 0))
    region_area = np.sum(region_mask)
    if region_area > 0 and overlap / region_area < 0.3:
        merged_labels[region_mask] = max_tier1_label + new_lbl

total_particles = len(np.unique(merged_labels)) - 1  # exclude 0
print(f"Total particles after merging: {total_particles}")

# Create tier flag array
tier_flag = {}  # label -> 'tier1' or 'new'
for lbl in np.unique(merged_labels):
    if lbl == 0:
        continue
    if lbl <= max_tier1_label:
        tier_flag[int(lbl)] = 'tier1'
    else:
        tier_flag[int(lbl)] = 'new'

# ========== Step 8: Measure all particles ==========
props = measure.regionprops(merged_labels, intensity_image=image)

particle_data = []
for prop in props:
    area = prop.area
    perimeter = prop.perimeter
    circularity = 4 * np.pi * area / (perimeter ** 2) if perimeter > 0 else 0
    equiv_diam = prop.equivalent_diameter
    cy, cx = prop.centroid
    mean_int = prop.mean_intensity
    max_int = prop.max_intensity
    integrated_int = mean_int * area
    lbl = prop.label
    
    particle_data.append({
        'label': int(lbl),
        'centroid_x': float(cx),
        'centroid_y': float(cy),
        'area': int(area),
        'equivalent_diameter': float(equiv_diam),
        'circularity': float(circularity),
        'mean_intensity': float(mean_int),
        'max_intensity': float(max_int),
        'integrated_intensity': float(integrated_int),
        'tier': tier_flag.get(int(lbl), 'unknown')
    })

# ========== Step 9: Bimodality assessment ==========
sizes = np.array([p['equivalent_diameter'] for p in particle_data])
intensities = np.array([p['mean_intensity'] for p in particle_data])

def bimodality_coefficient(data):
    """Compute bimodality coefficient: BC = (skewness^2 + 1) / kurtosis
    BC > 0.555 suggests bimodality."""
    from scipy.stats import skew, kurtosis
    n = len(data)
    if n < 4:
        return 0.0
    g = skew(data)
    k = kurtosis(data, fisher=False)  # excess=False -> Pearson's kurtosis
    if k == 0:
        return 0.0
    bc = (g**2 + 1) / k
    return float(bc)

def kde_bimodality_check(data, label=''):
    """Check bimodality using KDE peak counting."""
    from scipy.signal import find_peaks
    from scipy.stats import gaussian_kde
    if len(data) < 3:
        return 0, False
    try:
        kde = gaussian_kde(data, bw_method='silverman')
        x_grid = np.linspace(data.min() - 0.1 * np.ptp(data), data.max() + 0.1 * np.ptp(data), 200)
        density = kde(x_grid)
        peaks, _ = find_peaks(density)
        return len(peaks), len(peaks) >= 2
    except Exception:
        return 0, False

size_bc = bimodality_coefficient(sizes) if len(sizes) >= 4 else 0.0
intensity_bc = bimodality_coefficient(intensities) if len(intensities) >= 4 else 0.0

size_n_peaks, size_bimodal = kde_bimodality_check(sizes, 'size')
intensity_n_peaks, intensity_bimodal = kde_bimodality_check(intensities, 'intensity')

print(f"Size BC: {size_bc:.3f}, KDE peaks: {size_n_peaks}")
print(f"Intensity BC: {intensity_bc:.3f}, KDE peaks: {intensity_n_peaks}")

# ========== Step 10: Nearest-neighbor distances and spatial statistics ==========
centroids = np.array([[p['centroid_x'], p['centroid_y']] for p in particle_data])

from scipy.spatial import distance_matrix, KDTree

if len(centroids) >= 2:
    tree = KDTree(centroids)
    nn_dists, nn_idx = tree.query(centroids, k=2)  # k=2 because nearest is self
    nn_distances = nn_dists[:, 1]  # second nearest = actual nearest neighbor
    mean_nn_dist = float(np.mean(nn_distances))
    
    # Clark-Evans R statistic
    # R = mean_observed_nn / mean_expected_nn
    # Expected mean NN for CSR in 2D: 0.5 / sqrt(density)
    area_total = image.shape[0] * image.shape[1]
    density = len(centroids) / area_total
    expected_nn = 0.5 / np.sqrt(density) if density > 0 else 0
    clark_evans_R = mean_nn_dist / expected_nn if expected_nn > 0 else 0
    # R < 1 => clustered, R = 1 => random, R > 1 => dispersed
    
    # Significance test (z-score)
    se = 0.26136 / np.sqrt(len(centroids) * density) if density > 0 else 1
    z_clark_evans = (mean_nn_dist - expected_nn) / se if se > 0 else 0
else:
    nn_distances = np.array([])
    mean_nn_dist = 0.0
    clark_evans_R = 0.0
    z_clark_evans = 0.0

print(f"Mean NN distance: {mean_nn_dist:.2f}")
print(f"Clark-Evans R: {clark_evans_R:.3f}, z: {z_clark_evans:.3f}")

# Spatial intensity gradient
if len(centroids) >= 3:
    from sklearn.linear_model import LinearRegression
    X_spatial = centroids.copy()
    y_int = intensities.copy()
    reg = LinearRegression().fit(X_spatial, y_int)
    gradient_x, gradient_y = reg.coef_
    gradient_mag = float(np.sqrt(gradient_x**2 + gradient_y**2))
    gradient_dir = float(np.degrees(np.arctan2(gradient_y, gradient_x)))
else:
    gradient_x, gradient_y = 0.0, 0.0
    gradient_mag = 0.0
    gradient_dir = 0.0

# ========== Step 11: 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 with Tier 1 (green) and new (cyan)
ax = axes[0, 1]
overlay = np.stack([(image - img_min) / (img_max - img_min)] * 3, axis=-1)
overlay = np.clip(overlay, 0, 1).copy()

for p in particle_data:
    lbl = p['label']
    mask = merged_labels == lbl
    contours_list = measure.find_contours(mask.astype(float), 0.5)
    if p['tier'] == 'tier1':
        color = [0, 1, 0]  # green
        overlay[mask] = overlay[mask] * 0.6 + np.array([0, 0.4, 0])
    else:
        color = [0, 1, 1]  # cyan
        overlay[mask] = overlay[mask] * 0.6 + np.array([0, 0.2, 0.2])
    for contour in contours_list:
        ax.plot(contour[:, 1], contour[:, 0], color=color, linewidth=1.2)

ax.imshow(np.clip(overlay, 0, 1))
ax.set_title(f'Segmentation Overlay (Total: {total_particles})\nGreen=Tier1, Cyan=New')
ax.axis('off')

# (c) Scatter: intensity vs size colored by spatial position
ax = axes[0, 2]
if len(particle_data) > 0:
    xs = [p['centroid_x'] for p in particle_data]
    ys = [p['centroid_y'] for p in particle_data]
    spatial_pos = np.sqrt(np.array(xs)**2 + np.array(ys)**2)
    sc = ax.scatter(sizes, intensities, c=spatial_pos, cmap='viridis', s=60, edgecolors='k', linewidth=0.5)
    plt.colorbar(sc, ax=ax, label='Distance from origin')
ax.set_xlabel('Equivalent Diameter (px)')
ax.set_ylabel('Mean Intensity')
ax.set_title('Intensity vs Size')

# (d) Histogram of sizes with bimodality
ax = axes[1, 0]
if len(sizes) > 0:
    ax.hist(sizes, bins=max(5, len(sizes)//3), color='steelblue', edgecolor='black', alpha=0.7)
    ax.axvline(np.mean(sizes), color='red', linestyle='--', label=f'Mean={np.mean(sizes):.1f}')
    ax.set_title(f'Size Distribution\nBC={size_bc:.3f}, KDE peaks={size_n_peaks}')
ax.set_xlabel('Equivalent Diameter (px)')
ax.set_ylabel('Count')
ax.legend()

# (e) Histogram of intensities
ax = axes[1, 1]
if len(intensities) > 0:
    ax.hist(intensities, bins=max(5, len(intensities)//3), color='coral', edgecolor='black', alpha=0.7)
    ax.axvline(np.mean(intensities), color='red', linestyle='--', label=f'Mean={np.mean(intensities):.1f}')
    ax.set_title(f'Intensity Distribution\nBC={intensity_bc:.3f}, KDE peaks={intensity_n_peaks}')
ax.set_xlabel('Mean Intensity')
ax.set_ylabel('Count')
ax.legend()

# (f) NN distance distribution with CSR reference
ax = axes[1, 2]
if len(nn_distances) > 0:
    ax.hist(nn_distances, bins=max(5, len(nn_distances)//3), color='mediumseagreen', edgecolor='black', alpha=0.7, label='Observed NN dist')
    # CSR expected NN distance
    if density > 0:
        expected_nn_val = 0.5 / np.sqrt(density)
        ax.axvline(expected_nn_val, color='red', linestyle='--', linewidth=2, label=f'CSR expected={expected_nn_val:.1f}')
    ax.axvline(mean_nn_dist, color='blue', linestyle='--', linewidth=2, label=f'Mean obs={mean_nn_dist:.1f}')
    ax.set_title(f'NN Distances\nClark-Evans R={clark_evans_R:.3f}')
ax.set_xlabel('Nearest Neighbor Distance (px)')
ax.set_ylabel('Count')
ax.legend()

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

# ========== Save outputs ==========
np.save('analysis_labels.npy', merged_labels.astype(np.int32))
np.save('nn_distances.npy', nn_distances.astype(np.float32) if len(nn_distances) > 0 else np.array([], dtype=np.float32))

# ========== JSON results ==========
per_particle_centroids = {str(p['label']): [p['centroid_x'], p['centroid_y']] for p in particle_data}
per_particle_diameters = {str(p['label']): round(p['equivalent_diameter'], 2) for p in particle_data}
per_particle_mean_int = {str(p['label']): round(p['mean_intensity'], 2) for p in particle_data}
per_particle_max_int = {str(p['label']): round(p['max_intensity'], 2) for p in particle_data}
per_particle_circ = {str(p['label']): round(p['circularity'], 3) for p in particle_data}
tier_flags = {str(p['label']): p['tier'] for p in particle_data}

spatial_pattern = 'clustered' if clark_evans_R < 0.9 else ('dispersed' if clark_evans_R > 1.1 else 'random')

results = {
    "analysis_type": "Multi-level particle detection with adaptive thresholding, watershed splitting, morphometric characterization, bimodality assessment, and spatial statistics",
    "extracted_features": {
        "total_particle_count": total_particles,
        "tier1_particle_count": tier1_count,
        "new_particle_count": total_particles - tier1_count,
        "per_particle_centroid_xy": per_particle_centroids,
        "per_particle_equivalent_diameter": per_particle_diameters,
        "per_particle_mean_intensity": per_particle_mean_int,
        "per_particle_max_intensity": per_particle_max_int,
        "per_particle_circularity": per_particle_circ,
        "tier1_vs_new_detection_label_flag": tier_flags,
        "size_distribution_bimodality_coefficient": round(size_bc, 4),
        "size_kde_peaks": size_n_peaks,
        "size_bimodal": bool(size_bimodal),
        "intensity_distribution_bimodality_coefficient": round(intensity_bc, 4),
        "intensity_kde_peaks": intensity_n_peaks,
        "intensity_bimodal": bool(intensity_bimodal),
        "mean_nearest_neighbor_distance": round(mean_nn_dist, 2),
        "nearest_neighbor_distances": [round(float(d), 2) for d in nn_distances] if len(nn_distances) > 0 else [],
        "clark_evans_R_statistic": round(float(clark_evans_R), 4),
        "clark_evans_z_score": round(float(z_clark_evans), 4),
        "spatial_pattern": spatial_pattern,
        "spatial_intensity_gradient_direction_and_magnitude": {
            "direction_degrees": round(gradient_dir, 2),
            "magnitude": round(gradient_mag, 4)
        },
        "mean_equivalent_diameter": round(float(np.mean(sizes)), 2) if len(sizes) > 0 else 0,
        "std_equivalent_diameter": round(float(np.std(sizes)), 2) if len(sizes) > 0 else 0,
        "mean_intensity_all": round(float(np.mean(intensities)), 2) if len(intensities) > 0 else 0,
        "std_intensity_all": round(float(np.std(intensities)), 2) if len(intensities) > 0 else 0
    },
    "quality_metrics": {
        "otsu_threshold": int(otsu_thresh),
        "lower_threshold_used": lower_thresh,
        "adaptive_block_size": block_size,
        "adaptive_offset": adaptive_offset,
        "image_intensity_range": [float(img_min), float(img_max)],
        "mean_circularity": round(float(np.mean([p['circularity'] for p in particle_data])), 3) if particle_data else 0,
        "watershed_min_distance": 7
    },
    "summary": f"Detected {total_particles} total particles ({tier1_count} from Tier 1, {total_particles - tier1_count} newly detected). Size BC={size_bc:.3f} ({size_n_peaks} KDE peaks), Intensity BC={intensity_bc:.3f} ({intensity_n_peaks} KDE peaks). Spatial pattern: {spatial_pattern} (Clark-Evans R={clark_evans_R:.3f}, z={z_clark_evans:.3f}). Adaptive thresholding with block_size=51 and offset=10 combined with lower global threshold at {lower_thresh} used to capture dim particles.",
    "saved_arrays": {
        "analysis_labels.npy": {
            "description": f"Integer label map with {total_particles} particles labeled 1-{merged_labels.max()}, background=0. Labels 1-{max_tier1_label} are Tier 1, rest are newly detected.",
            "shape": list(merged_labels.shape),
            "dtype": "int32"
        },
        "nn_distances.npy": {
            "description": "Nearest-neighbor distances for each particle centroid",
            "shape": [len(nn_distances)],
            "dtype": "float32"
        }
    }
}

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