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

# Load image
path = '/Users/maxim.ziatdinov/Code/SciLink/analysis_session_20260331_184040/results/analysis_polycrystalline_grains_ImageAnalysis_20260331_184235_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()}]")

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

# Scale factor: 100 um / 512 pixels
scale_um_per_pixel = 100.0 / 512.0  # 0.1953 um/pixel

# ---- PRIMARY: SAM-based segmentation ----
try:
    from scilink.tools.sam import run_sam_analysis
    print("Running SAM analysis with aggressive parameters...")
    result = run_sam_analysis(image_gray, params={
        "sam_parameters": "default",
        "min_area": 1000,
        "max_area": 150000,
        "pruning_iou_threshold": 0.5
    })
    
    # Build labeled mask
    labeled_sam = np.zeros(image_gray.shape[:2], dtype=np.int32)
    for i, p in enumerate(result["particles"]):
        mask = np.array(p["mask"], dtype=bool)
        labeled_sam[mask] = i + 1
    
    num_sam_grains = len(result["particles"])
    print(f"SAM detected {num_sam_grains} grains")
    
    if num_sam_grains >= 5:
        labeled = labeled_sam
        method_used = f"SAM (default parameters, detected {num_sam_grains} grains)"
    else:
        print("SAM detected fewer than 5 grains, trying sensitive...")
        result2 = run_sam_analysis(image_gray, params={
            "sam_parameters": "sensitive",
            "min_area": 1000,
            "max_area": 150000,
            "pruning_iou_threshold": 0.6
        })
        labeled_sam2 = np.zeros(image_gray.shape[:2], dtype=np.int32)
        for i, p in enumerate(result2["particles"]):
            mask = np.array(p["mask"], dtype=bool)
            labeled_sam2[mask] = i + 1
        num_sam2 = len(result2["particles"])
        print(f"SAM sensitive detected {num_sam2} grains")
        if num_sam2 >= 5:
            labeled = labeled_sam2
            method_used = f"SAM (sensitive parameters, detected {num_sam2} grains)"
        else:
            labeled = None
            method_used = None
except Exception as e:
    print(f"SAM failed: {e}")
    labeled = None
    method_used = None

# ---- FALLBACK: Boundary-skeleton-based segmentation ----
if labeled is None or len(np.unique(labeled)) - 1 < 5:
    print("Falling back to boundary-skeleton-based segmentation...")
    
    # Otsu thresholding to detect grain boundaries
    # Grain boundaries appear as dark lines in the image
    thresh_val, binary_otsu = cv2.threshold(image_gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
    print(f"Otsu threshold value: {thresh_val}")
    
    # In polycrystalline grain images, boundaries are typically darker
    # binary_otsu: 255 = bright (grain interior), 0 = dark (boundaries)
    # We want boundaries as foreground for skeletonization
    boundary_mask = (binary_otsu == 0).astype(np.uint8)
    
    # Clean up the boundary mask
    kernel_small = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3))
    boundary_clean = cv2.morphologyEx(boundary_mask, cv2.MORPH_CLOSE, kernel_small, iterations=2)
    boundary_clean = cv2.morphologyEx(boundary_clean, cv2.MORPH_OPEN, kernel_small, iterations=1)
    
    # Skeletonize the boundary mask
    skeleton = skimage.morphology.skeletonize(boundary_clean.astype(bool))
    
    # Dilate skeleton slightly to ensure it forms closed boundaries
    skeleton_dilated = cv2.dilate(skeleton.astype(np.uint8), kernel_small, iterations=1)
    
    # Invert to get grain interiors
    grain_interior = (skeleton_dilated == 0).astype(np.uint8)
    
    # Distance transform + watershed
    distance = scipy.ndimage.distance_transform_edt(grain_interior)
    
    # Use Sobel gradient as watershed landscape (per advanced guidance)
    sobel_x = cv2.Sobel(image_gray, cv2.CV_64F, 1, 0, ksize=3)
    sobel_y = cv2.Sobel(image_gray, cv2.CV_64F, 0, 1, ksize=3)
    gradient_magnitude = np.sqrt(sobel_x**2 + sobel_y**2)
    
    # Find markers from distance transform peaks
    # Estimate grain radius: if ~15 grains in 512x512, average area ~ 512*512/15 ~ 17476 px
    # radius ~ sqrt(17476/pi) ~ 74 pixels
    estimated_radius = 40  # conservative to avoid missing small grains
    
    coords = skimage.feature.peak_local_max(
        distance, 
        min_distance=estimated_radius,
        threshold_abs=10,
        labels=grain_interior
    )
    
    # Create marker image
    marker_mask = np.zeros(image_gray.shape, dtype=bool)
    marker_mask[tuple(coords.T)] = True
    labeled_markers = scipy.ndimage.label(marker_mask)[0]
    
    print(f"Number of watershed markers: {labeled_markers.max()}")
    
    # Watershed using gradient as landscape
    labeled = skimage.segmentation.watershed(
        gradient_magnitude, 
        labeled_markers, 
        mask=grain_interior
    )
    
    method_used = f"Boundary-skeleton watershed (Otsu threshold={thresh_val:.0f}, {labeled.max()} grains)"
    print(f"Watershed segmentation found {labeled.max()} grains")

# ---- Post-processing: filter and analyze ----
print(f"Method used: {method_used}")
print(f"Total labels: {labeled.max()}")

# Get region properties
props = skimage.measure.regionprops(labeled, intensity_image=image_gray)

# Filter out very small fragments
min_area_filter = 1000  # pixels
valid_props = [p for p in props if p.area >= min_area_filter]
print(f"Grains after area filter (>={min_area_filter} px): {len(valid_props)}")

# Identify edge-touching grains
H, W = image_gray.shape
edge_labels = set()
for p in valid_props:
    coords = p.coords
    if (np.any(coords[:, 0] == 0) or np.any(coords[:, 0] == H - 1) or
        np.any(coords[:, 1] == 0) or np.any(coords[:, 1] == W - 1)):
        edge_labels.add(p.label)

interior_props = [p for p in valid_props if p.label not in edge_labels]
print(f"Interior grains (not touching edges): {len(interior_props)}")

# ---- Neighbor count via dilated label overlap ----
def compute_neighbors(labeled_map, valid_labels):
    """Compute neighbor count for each grain via dilation overlap."""
    neighbors = {}
    dilation_kernel = skimage.morphology.disk(3)
    for lbl in valid_labels:
        mask = (labeled_map == lbl)
        dilated = skimage.morphology.binary_dilation(mask, dilation_kernel)
        boundary = dilated & ~mask
        neighbor_labels = set(labeled_map[boundary]) - {0, lbl}
        # Only count valid labels as neighbors
        neighbor_labels = neighbor_labels & set(valid_labels)
        neighbors[lbl] = len(neighbor_labels)
    return neighbors

valid_labels = [p.label for p in valid_props]
neighbor_counts = compute_neighbors(labeled, valid_labels)

# ---- Extract features per grain ----
grain_data = []
for p in valid_props:
    area_px = p.area
    area_um2 = area_px * (scale_um_per_pixel ** 2)
    eq_diam_px = p.equivalent_diameter
    eq_diam_um = eq_diam_px * scale_um_per_pixel
    perimeter_px = p.perimeter
    perimeter_um = perimeter_px * scale_um_per_pixel
    
    # Circularity: 4*pi*area/perimeter^2
    if perimeter_px > 0:
        circularity = 4 * np.pi * area_px / (perimeter_px ** 2)
    else:
        circularity = 0.0
    
    # Aspect ratio from major/minor axis
    if p.minor_axis_length > 0:
        aspect_ratio = p.major_axis_length / p.minor_axis_length
    else:
        aspect_ratio = 1.0
    
    solidity = p.solidity
    centroid = p.centroid  # (row, col)
    is_edge = p.label in edge_labels
    n_neighbors = neighbor_counts.get(p.label, 0)
    
    grain_data.append({
        "label": int(p.label),
        "area_pixels": int(area_px),
        "area_um2": round(area_um2, 2),
        "equivalent_diameter_um": round(eq_diam_um, 2),
        "perimeter_um": round(perimeter_um, 2),
        "circularity": round(circularity, 4),
        "aspect_ratio": round(aspect_ratio, 4),
        "solidity": round(solidity, 4),
        "centroid_row": round(centroid[0], 2),
        "centroid_col": round(centroid[1], 2),
        "is_edge_touching": is_edge,
        "neighbor_count": n_neighbors
    })

# ---- Compute statistics (interior grains only for unbiased stats) ----
stats_props = interior_props if len(interior_props) >= 3 else valid_props
stats_label = "interior" if len(interior_props) >= 3 else "all_valid"

areas_um2 = np.array([p.area * (scale_um_per_pixel ** 2) for p in stats_props])
diameters_um = np.array([p.equivalent_diameter * scale_um_per_pixel for p in stats_props])

mean_area = float(np.mean(areas_um2))
std_area = float(np.std(areas_um2))
median_area = float(np.median(areas_um2))
min_area = float(np.min(areas_um2))
max_area = float(np.max(areas_um2))
cv_area = float(std_area / mean_area) if mean_area > 0 else 0.0

mean_diam = float(np.mean(diameters_um))
std_diam = float(np.std(diameters_um))

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

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

# 2. Segmentation overlay
ax = axes[0, 1]
ax.imshow(image_gray, cmap='gray')
# Create colored overlay
np.random.seed(42)
colors = np.random.rand(labeled.max() + 1, 4)
colors[:, 3] = 0.4  # semi-transparent
colors[0] = [0, 0, 0, 0]  # background transparent

for p in valid_props:
    mask = labeled == p.label
    color = colors[p.label]
    overlay = np.zeros((*image_gray.shape, 4))
    overlay[mask] = color
    ax.imshow(overlay)
    # Draw contour
    contour_mask = mask.astype(np.uint8)
    contours, _ = cv2.findContours(contour_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    for c in contours:
        c_squeezed = c.squeeze()
        if c_squeezed.ndim == 2 and len(c_squeezed) > 2:
            ax.plot(c_squeezed[:, 0], c_squeezed[:, 1], 'r-', linewidth=1)

ax.set_title(f'Segmentation Overlay ({len(valid_props)} grains)', fontsize=12)
ax.axis('off')

# 3. Label map
ax = axes[0, 2]
cmap = plt.cm.get_cmap('tab20', labeled.max() + 1)
ax.imshow(labeled, cmap=cmap, interpolation='nearest')
ax.set_title('Label Map', fontsize=12)
ax.axis('off')

# 4. Interior vs edge grains
ax = axes[1, 0]
ax.imshow(image_gray, cmap='gray')
for p in valid_props:
    mask = labeled == p.label
    if p.label in edge_labels:
        color = [1, 0, 0, 0.3]  # red for edge
    else:
        color = [0, 1, 0, 0.3]  # green for interior
    overlay = np.zeros((*image_gray.shape, 4))
    overlay[mask] = color
    ax.imshow(overlay)
ax.set_title(f'Interior (green, {len(interior_props)}) vs Edge (red, {len(edge_labels)})', fontsize=11)
ax.axis('off')

# 5. Grain size distribution histogram
ax = axes[1, 1]
all_areas_um2 = np.array([g["area_um2"] for g in grain_data if not g["is_edge_touching"]])
if len(all_areas_um2) < 3:
    all_areas_um2 = np.array([g["area_um2"] for g in grain_data])
ax.hist(all_areas_um2, bins=min(15, max(5, len(all_areas_um2)//2)), color='steelblue', edgecolor='black', alpha=0.7)
ax.axvline(mean_area, color='red', linestyle='--', label=f'Mean={mean_area:.1f}')
ax.axvline(median_area, color='orange', linestyle='--', label=f'Median={median_area:.1f}')
ax.set_xlabel('Grain Area (µm²)', fontsize=11)
ax.set_ylabel('Count', fontsize=11)
ax.set_title('Grain Size Distribution (Interior)', fontsize=12)
ax.legend(fontsize=9)

# 6. Equivalent diameter distribution
ax = axes[1, 2]
int_diams = np.array([g["equivalent_diameter_um"] for g in grain_data if not g["is_edge_touching"]])
if len(int_diams) < 3:
    int_diams = np.array([g["equivalent_diameter_um"] for g in grain_data])
ax.hist(int_diams, bins=min(15, max(5, len(int_diams)//2)), color='coral', edgecolor='black', alpha=0.7)
ax.axvline(mean_diam, color='red', linestyle='--', label=f'Mean={mean_diam:.1f}')
ax.set_xlabel('Equivalent Diameter (µm)', fontsize=11)
ax.set_ylabel('Count', fontsize=11)
ax.set_title('Grain Diameter Distribution (Interior)', fontsize=12)
ax.legend(fontsize=9)

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

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

# ---- Compute histogram data for JSON ----
hist_counts, hist_edges = np.histogram(all_areas_um2, bins=min(15, max(5, len(all_areas_um2)//2)))
hist_data = {
    "bin_edges_um2": [round(float(e), 2) for e in hist_edges],
    "counts": [int(c) for c in hist_counts]
}

# ---- Build results JSON ----
results = {
    "analysis_type": "Polycrystalline grain segmentation and morphological analysis",
    "method_used": method_used,
    "scale_um_per_pixel": round(scale_um_per_pixel, 4),
    "extracted_features": {
        "grain_count_total": len(valid_props),
        "grain_count_interior": len(interior_props),
        "grain_count_edge": len(edge_labels),
        "statistics_computed_from": stats_label,
        "mean_grain_area_um2": round(mean_area, 2),
        "std_grain_area_um2": round(std_area, 2),
        "median_grain_area_um2": round(median_area, 2),
        "min_grain_area_um2": round(min_area, 2),
        "max_grain_area_um2": round(max_area, 2),
        "coefficient_of_variation_area": round(cv_area, 4),
        "mean_equivalent_diameter_um": round(mean_diam, 2),
        "std_equivalent_diameter_um": round(std_diam, 2),
        "grain_size_distribution_histogram": hist_data,
        "per_grain_data": grain_data
    },
    "quality_metrics": {
        "segmentation_method": method_used,
        "total_labeled_area_fraction": round(float(np.sum(labeled > 0)) / (H * W), 4),
        "mean_circularity": round(float(np.mean([g["circularity"] for g in grain_data])), 4),
        "mean_solidity": round(float(np.mean([g["solidity"] for g in grain_data])), 4),
        "mean_aspect_ratio": round(float(np.mean([g["aspect_ratio"] for g in grain_data])), 4),
        "mean_neighbor_count": round(float(np.mean([g["neighbor_count"] for g in grain_data])), 2)
    },
    "summary": f"Segmented {len(valid_props)} total grains ({len(interior_props)} interior, {len(edge_labels)} edge-touching) using {method_used}. Interior grain mean area = {mean_area:.1f} µm² (CV = {cv_area:.2f}), mean equivalent diameter = {mean_diam:.1f} µm. Scale: {scale_um_per_pixel:.4f} µm/pixel (100 µm / 512 px).",
    "saved_arrays": {
        "analysis_labels.npy": {
            "description": f"Integer label map, {len(valid_props)} grains labeled 1-{labeled.max()}, background=0",
            "shape": list(labeled.shape),
            "dtype": "int32"
        }
    }
}

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