import numpy as np
import cv2
import json
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from scipy import ndimage
from skimage import measure, morphology, segmentation, feature
from skimage.morphology import disk
from scipy.spatial import cKDTree

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

# Ensure 2D grayscale
if image.ndim == 3:
    gray = image[:, :, 0]
else:
    gray = image.copy()

gray_float = gray.astype(np.float64)

# 2. Pipeline implementation per plan

# Step 1: Gaussian blur (sigma=1.5) to suppress minor noise
from scipy.ndimage import gaussian_filter
blurred = gaussian_filter(gray_float, sigma=1.5)

# Step 2: Sobel gradient magnitude to enhance grain boundaries
sobel_x = ndimage.sobel(blurred, axis=1)
sobel_y = ndimage.sobel(blurred, axis=0)
gradient_magnitude = np.hypot(sobel_x, sobel_y)

# Step 3: Otsu threshold on gradient image to create binary boundary mask
# Convert gradient to uint8 for Otsu
grad_normalized = ((gradient_magnitude - gradient_magnitude.min()) / 
                   (gradient_magnitude.max() - gradient_magnitude.min()) * 255).astype(np.uint8)
otsu_thresh, boundary_mask = cv2.threshold(grad_normalized, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
print(f"Otsu threshold on gradient image: {otsu_thresh}")
boundary_binary = boundary_mask > 0

# Step 4: Morphological dilation (disk r=1) to close small gaps in boundary network
selem_dilate = disk(1)
boundary_dilated = morphology.binary_dilation(boundary_binary, selem_dilate)

# Step 5: Invert mask to get grain interiors
grain_interiors = ~boundary_dilated

# Step 6: Morphological erosion (disk r=2) to separate grains at thin contact points
selem_erode = disk(2)
grain_eroded = morphology.binary_erosion(grain_interiors, selem_erode)

# Step 7: Connected component labeling
labeled_seeds, num_seeds = ndimage.label(grain_eroded)
print(f"Number of seed regions before filtering: {num_seeds}")

# Step 8: Remove small spurious regions (area < 500 pixels)
min_area = 500
props_seeds = measure.regionprops(labeled_seeds)
for prop in props_seeds:
    if prop.area < min_area:
        labeled_seeds[labeled_seeds == prop.label] = 0

# Relabel after removal
labeled_seeds_clean = measure.label(labeled_seeds > 0)
num_seeds_clean = labeled_seeds_clean.max()
print(f"Number of seed regions after filtering (area >= {min_area}): {num_seeds_clean}")

# Step 9: Watershed expansion from labeled seeds on original gradient image
# Per advanced guidance: use Sobel gradient as watershed landscape
watershed_labels = segmentation.watershed(
    gradient_magnitude, 
    markers=labeled_seeds_clean, 
    mask=None,  # expand into all space
    compactness=0.01
)

# Actually, we should constrain watershed to grain interiors (not boundary)
# Let's use the non-dilated interior as mask for better results
# But plan says "recover full grain areas" so we use a generous mask
# Use the original grain_interiors (before erosion) as mask, slightly dilated
mask_for_watershed = ~boundary_binary  # use original boundary without dilation for better coverage
watershed_labels = segmentation.watershed(
    gradient_magnitude,
    markers=labeled_seeds_clean,
    mask=mask_for_watershed,
    compactness=0.01
)

# Relabel to ensure contiguous labels
final_labels = measure.label(watershed_labels > 0, connectivity=1)
# But we want to keep watershed label identity, let's use watershed_labels directly
# Just ensure background is 0
final_labels = watershed_labels.copy()
final_labels[final_labels < 0] = 0

# Count grains
grain_count = final_labels.max()
print(f"Final grain count: {grain_count}")

# 3. Extract region properties
props = measure.regionprops(final_labels, intensity_image=gray)

grain_areas = []
grain_perimeters = []
grain_centroids = []
grain_equiv_diameters = []
grain_major_axis = []
grain_minor_axis = []
grain_aspect_ratios = []
grain_solidity = []
grain_eccentricity = []
grain_mean_intensities = []
grain_orientations = []

for prop in props:
    grain_areas.append(float(prop.area))
    grain_perimeters.append(float(prop.perimeter))
    grain_centroids.append([float(prop.centroid[0]), float(prop.centroid[1])])
    grain_equiv_diameters.append(float(prop.equivalent_diameter))
    grain_major_axis.append(float(prop.major_axis_length))
    grain_minor_axis.append(float(prop.minor_axis_length))
    if prop.minor_axis_length > 0:
        grain_aspect_ratios.append(float(prop.major_axis_length / prop.minor_axis_length))
    else:
        grain_aspect_ratios.append(float('inf'))
    grain_solidity.append(float(prop.solidity))
    grain_eccentricity.append(float(prop.eccentricity))
    grain_mean_intensities.append(float(prop.mean_intensity))
    grain_orientations.append(float(prop.orientation))

# Grain size distribution statistics
areas_arr = np.array(grain_areas)
size_stats = {
    "mean_area": float(np.mean(areas_arr)) if len(areas_arr) > 0 else 0,
    "std_area": float(np.std(areas_arr)) if len(areas_arr) > 0 else 0,
    "median_area": float(np.median(areas_arr)) if len(areas_arr) > 0 else 0,
    "min_area": float(np.min(areas_arr)) if len(areas_arr) > 0 else 0,
    "max_area": float(np.max(areas_arr)) if len(areas_arr) > 0 else 0,
    "mean_equiv_diameter": float(np.mean(grain_equiv_diameters)) if len(grain_equiv_diameters) > 0 else 0,
    "std_equiv_diameter": float(np.std(grain_equiv_diameters)) if len(grain_equiv_diameters) > 0 else 0,
}

# Nearest neighbor distances
centroid_array = np.array(grain_centroids)
if len(centroid_array) >= 2:
    tree = cKDTree(centroid_array)
    dists, _ = tree.query(centroid_array, k=2)  # k=2 because closest is itself
    nn_distances = dists[:, 1].tolist()
    nn_distances_float = [float(d) for d in nn_distances]
    mean_nn_dist = float(np.mean(nn_distances))
    std_nn_dist = float(np.std(nn_distances))
else:
    nn_distances_float = []
    mean_nn_dist = 0.0
    std_nn_dist = 0.0

# Spatial gradient of grain size
# Compute local average grain size as function of position
if len(centroid_array) >= 3:
    # Fit a simple linear model: area = a*x + b*y + c
    from sklearn.linear_model import LinearRegression
    X_pos = centroid_array  # [row, col]
    y_area = np.array(grain_areas)
    reg = LinearRegression().fit(X_pos, y_area)
    spatial_gradient = {
        "gradient_row": float(reg.coef_[0]),
        "gradient_col": float(reg.coef_[1]),
        "gradient_magnitude": float(np.sqrt(reg.coef_[0]**2 + reg.coef_[1]**2)),
        "r_squared": float(reg.score(X_pos, y_area)),
        "intercept": float(reg.intercept_)
    }
else:
    spatial_gradient = {"gradient_row": 0.0, "gradient_col": 0.0, "gradient_magnitude": 0.0, "r_squared": 0.0, "intercept": 0.0}

# 4. Visualization
fig, axes = plt.subplots(2, 3, figsize=(18, 12), dpi=100)

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

# Segmentation overlay
from matplotlib.colors import ListedColormap
# Create colored overlay
overlay = np.zeros((*gray.shape, 4), dtype=np.float32)
np.random.seed(42)
colors = np.random.rand(grain_count + 1, 3)
colors[0] = [0, 0, 0]  # background

for i in range(1, grain_count + 1):
    mask_i = final_labels == i
    overlay[mask_i, 0] = colors[i, 0]
    overlay[mask_i, 1] = colors[i, 1]
    overlay[mask_i, 2] = colors[i, 2]
    overlay[mask_i, 3] = 0.4  # semi-transparent

axes[0, 1].imshow(gray, cmap='gray')
axes[0, 1].imshow(overlay)
# Add contour boundaries
for i in range(1, grain_count + 1):
    mask_i = (final_labels == i).astype(np.uint8)
    contours_cv, _ = cv2.findContours(mask_i, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    for cnt in contours_cv:
        cnt_squeezed = cnt.squeeze()
        if cnt_squeezed.ndim == 2 and len(cnt_squeezed) > 2:
            axes[0, 1].plot(cnt_squeezed[:, 0], cnt_squeezed[:, 1], 'r-', linewidth=0.5)
axes[0, 1].set_title(f'Segmentation Overlay ({grain_count} grains)')
axes[0, 1].axis('off')

# Gradient magnitude
axes[0, 2].imshow(gradient_magnitude, cmap='hot')
axes[0, 2].set_title('Sobel Gradient Magnitude')
axes[0, 2].axis('off')

# Label map
label_display = np.zeros((*gray.shape, 3), dtype=np.float32)
for i in range(1, grain_count + 1):
    label_display[final_labels == i] = colors[i]
axes[1, 0].imshow(label_display)
axes[1, 0].set_title('Grain Label Map')
axes[1, 0].axis('off')

# Grain area histogram
if len(grain_areas) > 0:
    axes[1, 1].hist(grain_areas, bins=min(30, max(5, grain_count // 2)), color='steelblue', edgecolor='black')
    axes[1, 1].set_xlabel('Grain Area (pixels)')
    axes[1, 1].set_ylabel('Count')
    axes[1, 1].set_title('Grain Size Distribution')
else:
    axes[1, 1].text(0.5, 0.5, 'No grains detected', ha='center', va='center', transform=axes[1, 1].transAxes)
    axes[1, 1].set_title('Grain Size Distribution')

# Centroid scatter colored by area
if len(centroid_array) > 0:
    sc = axes[1, 2].scatter(centroid_array[:, 1], centroid_array[:, 0], 
                             c=grain_areas, cmap='viridis', s=20, edgecolors='white', linewidth=0.3)
    axes[1, 2].set_xlim(0, gray.shape[1])
    axes[1, 2].set_ylim(gray.shape[0], 0)
    axes[1, 2].set_title('Centroids (colored by area)')
    axes[1, 2].set_aspect('equal')
    plt.colorbar(sc, ax=axes[1, 2], label='Area (px)')
else:
    axes[1, 2].text(0.5, 0.5, 'No grains', ha='center', va='center', transform=axes[1, 2].transAxes)
    axes[1, 2].set_title('Centroids')

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

# 5. Save label map
np.save('analysis_labels.npy', final_labels.astype(np.int32))
print("Saved analysis_labels.npy")

# 6. Print results as JSON
results = {
    "analysis_type": "Grain boundary segmentation and morphology analysis using Sobel gradient, Otsu thresholding, morphological operations, and watershed expansion",
    "extracted_features": {
        "grain_count": int(grain_count),
        "grain_areas": grain_areas,
        "grain_perimeters": grain_perimeters,
        "grain_centroids": grain_centroids,
        "grain_equivalent_diameters": grain_equiv_diameters,
        "grain_major_axis_lengths": grain_major_axis,
        "grain_minor_axis_lengths": grain_minor_axis,
        "grain_aspect_ratios": grain_aspect_ratios,
        "grain_solidity": grain_solidity,
        "grain_eccentricity": grain_eccentricity,
        "grain_mean_intensities": grain_mean_intensities,
        "grain_orientations": grain_orientations,
        "grain_size_distribution_statistics": size_stats,
        "nearest_neighbor_distances": nn_distances_float,
        "mean_nearest_neighbor_distance": mean_nn_dist,
        "std_nearest_neighbor_distance": std_nn_dist,
        "spatial_gradient_of_grain_size": spatial_gradient
    },
    "quality_metrics": {
        "otsu_threshold_on_gradient": float(otsu_thresh),
        "num_seeds_before_filtering": int(num_seeds),
        "num_seeds_after_filtering": int(num_seeds_clean),
        "final_grain_count": int(grain_count),
        "min_area_threshold": int(min_area),
        "gaussian_sigma": 1.5,
        "erosion_disk_radius": 2,
        "dilation_disk_radius": 1,
        "image_coverage_fraction": float(np.sum(final_labels > 0) / final_labels.size)
    },
    "summary": f"Detected {grain_count} grains via Sobel gradient + Otsu boundary detection + watershed expansion. Mean grain area: {size_stats['mean_area']:.1f} px (std: {size_stats['std_area']:.1f}), mean equivalent diameter: {size_stats['mean_equiv_diameter']:.1f} px. Gaussian sigma=1.5, erosion disk r=2, dilation disk r=1, min area filter=500 px. Spatial gradient of grain size has magnitude {spatial_gradient['gradient_magnitude']:.2f} px^2/px (R^2={spatial_gradient['r_squared']:.3f}).",
    "saved_arrays": {
        "analysis_labels.npy": {
            "description": f"Integer label map, {grain_count} grains labeled 1-{grain_count}, background=0",
            "shape": list(final_labels.shape),
            "dtype": "int32"
        }
    }
}

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