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

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

# Scale factor: 100 um / 512 pixels
scale = 100.0 / 512.0  # um per pixel
scale_area = scale ** 2  # um^2 per pixel^2

# Step 1: Gaussian blur (sigma=1.5)
blurred = cv2.GaussianBlur(image.astype(np.float64), (0, 0), 1.5)

# Step 2: Sobel gradient magnitude
sobel_x = cv2.Sobel(blurred, cv2.CV_64F, 1, 0, ksize=3)
sobel_y = cv2.Sobel(blurred, cv2.CV_64F, 0, 1, ksize=3)
gradient_mag = np.sqrt(sobel_x**2 + sobel_y**2)

# Normalize gradient to 0-255 for thresholding
grad_norm = ((gradient_mag - gradient_mag.min()) / (gradient_mag.max() - gradient_mag.min()) * 255).astype(np.uint8)

# Step 3: Adaptive thresholding on blurred image for dark boundary detection
blurred_uint8 = np.clip(blurred, 0, 255).astype(np.uint8)
adaptive_thresh = cv2.adaptiveThreshold(blurred_uint8, 255, cv2.ADAPTIVE_THRESH_MEAN_C, 
                                         cv2.THRESH_BINARY_INV, blockSize=51, C=15)

# Step 4: Otsu threshold on gradient magnitude
otsu_val, sobel_binary = cv2.threshold(grad_norm, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
print(f'Otsu threshold on gradient: {otsu_val}')

# Combine Sobel edges with adaptive threshold
boundary_mask = np.logical_or(sobel_binary > 0, adaptive_thresh > 0).astype(np.uint8)

# Step 5: Morphological closing to connect gaps (disk r=2)
selem_close = morphology.disk(2)
boundary_closed = morphology.binary_closing(boundary_mask.astype(bool), selem_close).astype(np.uint8)

# Step 6: Skeletonization
boundary_skeleton = morphology.skeletonize(boundary_closed.astype(bool))

# Step 7: Label connected regions in inverse of skeleton
grain_regions = ~boundary_skeleton
labeled, num_features_raw = ndimage.label(grain_regions)
print(f'Raw labeled regions: {num_features_raw}')

# Step 8: Remove edge-touching grains and small spurious regions
min_area_pixels = 500  # ~19 um^2

# Find edge-touching labels
edge_labels = set()
edge_labels.update(np.unique(labeled[0, :]))   # top
edge_labels.update(np.unique(labeled[-1, :]))  # bottom
edge_labels.update(np.unique(labeled[:, 0]))   # left
edge_labels.update(np.unique(labeled[:, -1]))  # right
edge_labels.discard(0)

# Get region properties to filter by area
props_raw = measure.regionprops(labeled)

# Create filtered label map
filtered_labels = np.zeros_like(labeled)
new_label = 0
label_mapping = {}

for prop in props_raw:
    if prop.label in edge_labels:
        continue
    if prop.area < min_area_pixels:
        continue
    new_label += 1
    label_mapping[prop.label] = new_label
    filtered_labels[labeled == prop.label] = new_label

print(f'Filtered grain count: {new_label}')

# Step 9: Extract region properties
props = measure.regionprops(filtered_labels, intensity_image=image)

grain_count = len(props)
grain_areas_px = []
grain_areas_um2 = []
eq_diameters_um = []
perimeters_um = []
eccentricities = []
solidities = []
aspect_ratios = []
centroids = []
mean_intensities = []

for prop in props:
    area_px = prop.area
    area_um2 = area_px * scale_area
    eq_diam_um = prop.equivalent_diameter * scale
    perim_um = prop.perimeter * scale
    ecc = prop.eccentricity
    sol = prop.solidity
    # Aspect ratio from major/minor axis
    if prop.minor_axis_length > 0:
        ar = prop.major_axis_length / prop.minor_axis_length
    else:
        ar = float('inf')
    centroid = prop.centroid  # (row, col)
    
    grain_areas_px.append(area_px)
    grain_areas_um2.append(round(area_um2, 2))
    eq_diameters_um.append(round(eq_diam_um, 2))
    perimeters_um.append(round(perim_um, 2))
    eccentricities.append(round(ecc, 4))
    solidities.append(round(sol, 4))
    aspect_ratios.append(round(ar, 4))
    centroids.append([round(centroid[0] * scale, 2), round(centroid[1] * scale, 2)])
    mean_intensities.append(round(float(prop.mean_intensity), 2))

# Step 10: Compute distribution statistics
areas_arr = np.array(grain_areas_um2)
diams_arr = np.array(eq_diameters_um)

if grain_count > 0:
    area_stats = {
        'mean': round(float(np.mean(areas_arr)), 2),
        'median': round(float(np.median(areas_arr)), 2),
        'std': round(float(np.std(areas_arr)), 2),
        'min': round(float(np.min(areas_arr)), 2),
        'max': round(float(np.max(areas_arr)), 2)
    }
    diam_stats = {
        'mean': round(float(np.mean(diams_arr)), 2),
        'median': round(float(np.median(diams_arr)), 2),
        'std': round(float(np.std(diams_arr)), 2),
        'min': round(float(np.min(diams_arr)), 2),
        'max': round(float(np.max(diams_arr)), 2)
    }
else:
    area_stats = {'mean': 0, 'median': 0, 'std': 0, 'min': 0, 'max': 0}
    diam_stats = {'mean': 0, 'median': 0, 'std': 0, 'min': 0, 'max': 0}

# Save label map
np.save('analysis_labels.npy', filtered_labels.astype(np.int32))

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

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

# 2. Segmentation overlay
# Create colored overlay
from matplotlib.colors import ListedColormap
import matplotlib.cm as cm

overlay = np.stack([image/255.0]*3, axis=-1)

if grain_count > 0:
    # Generate random colors for each grain
    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 = (filtered_labels == i)
        for c in range(3):
            overlay[:, :, c] = np.where(mask_i, 
                overlay[:, :, c] * 0.5 + colors[i, c] * 0.5, 
                overlay[:, :, c])
    
    # Draw contour boundaries
    contour_img = np.zeros(filtered_labels.shape, dtype=np.uint8)
    for i in range(1, grain_count + 1):
        mask_i = (filtered_labels == i).astype(np.uint8)
        contours, _ = cv2.findContours(mask_i, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        cv2.drawContours(contour_img, contours, -1, 255, 1)
    
    # Overlay contours in white
    contour_mask = contour_img > 0
    overlay[contour_mask] = [1.0, 1.0, 1.0]

axes[0, 1].imshow(np.clip(overlay, 0, 1))
axes[0, 1].set_title(f'Segmentation Overlay ({grain_count} grains)')
axes[0, 1].axis('off')

# 3. Boundary detection result
axes[0, 2].imshow(boundary_skeleton, cmap='gray')
axes[0, 2].set_title('Boundary Skeleton')
axes[0, 2].axis('off')

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

# 5. Grain size distribution histogram
if grain_count > 0:
    axes[1, 1].hist(eq_diameters_um, bins=min(30, max(5, grain_count // 3)), 
                    color='steelblue', edgecolor='black', alpha=0.7)
    axes[1, 1].axvline(diam_stats['mean'], color='red', linestyle='--', 
                       label=f"Mean: {diam_stats['mean']:.1f} µm")
    axes[1, 1].axvline(diam_stats['median'], color='orange', linestyle='--', 
                       label=f"Median: {diam_stats['median']:.1f} µm")
    axes[1, 1].legend()
axes[1, 1].set_xlabel('Equivalent Diameter (µm)')
axes[1, 1].set_ylabel('Count')
axes[1, 1].set_title('Grain Size Distribution')

# 6. Area vs eccentricity scatter
if grain_count > 0:
    axes[1, 2].scatter(grain_areas_um2, eccentricities, c='steelblue', 
                       alpha=0.6, edgecolors='black', s=30)
axes[1, 2].set_xlabel('Area (µm²)')
axes[1, 2].set_ylabel('Eccentricity')
axes[1, 2].set_title('Grain Area vs Eccentricity')

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

# Prepare results
results = {
    'analysis_type': 'Polycrystalline grain boundary segmentation and grain size analysis using Sobel edge detection, adaptive thresholding, morphological processing, and skeletonization',
    'extracted_features': {
        'grain_count': grain_count,
        'grain_areas_um2': grain_areas_um2,
        'equivalent_diameters_um': eq_diameters_um,
        'grain_size_distribution_statistics': {
            'area_stats_um2': area_stats,
            'diameter_stats_um': diam_stats
        },
        'perimeters_um': perimeters_um,
        'eccentricity': eccentricities,
        'solidity': solidities,
        'aspect_ratio': aspect_ratios,
        'centroid_positions_um': centroids,
        'mean_intensities': mean_intensities,
        'scale_factor_um_per_pixel': round(scale, 6),
        'min_area_threshold_pixels': min_area_pixels
    },
    'quality_metrics': {
        'raw_regions_before_filtering': int(num_features_raw),
        'edge_touching_removed': len(edge_labels),
        'otsu_gradient_threshold': float(otsu_val),
        'boundary_pixel_fraction': round(float(np.sum(boundary_skeleton)) / boundary_skeleton.size, 4),
        'mean_grain_area_um2': area_stats['mean'] if grain_count > 0 else 0,
        'mean_equivalent_diameter_um': diam_stats['mean'] if grain_count > 0 else 0,
        'mean_eccentricity': round(float(np.mean(eccentricities)), 4) if grain_count > 0 else 0,
        'mean_solidity': round(float(np.mean(solidities)), 4) if grain_count > 0 else 0
    },
    'summary': f'Detected {grain_count} grains with mean equivalent diameter {diam_stats["mean"]:.1f} µm (std: {diam_stats["std"]:.1f} µm) and mean area {area_stats["mean"]:.1f} µm² using Sobel gradient + adaptive thresholding combined boundary detection followed by skeletonization. Parameters: Gaussian sigma=1.5, adaptive block_size=51 offset=15, closing disk r=2, min_area=500px (~{500*scale_area:.1f} µm²). Scale: {scale:.4f} µm/pixel.',
    'saved_arrays': {
        'analysis_labels.npy': {
            'description': f'Integer label map, {grain_count} grains labeled 1-{grain_count}, background=0. Edge-touching and small regions removed.',
            'shape': list(filtered_labels.shape),
            'dtype': 'int32'
        }
    }
}

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