import numpy as np
import json
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from scipy import ndimage
from skimage import filters, morphology, measure, color
import warnings
warnings.filterwarnings('ignore')

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

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

# Step 1: Gaussian blur (sigma=2) to suppress granular noise
image_blurred = filters.gaussian(image_2d, sigma=2)

# Step 2: Otsu threshold to segment light vs dark phases
threshold_otsu = filters.threshold_otsu(image_blurred)
binary_light = image_blurred >= threshold_otsu  # light phase
binary_dark = image_blurred < threshold_otsu   # dark phase

print(f"Otsu threshold: {threshold_otsu:.4f}")
print(f"Light phase fraction before morphology: {binary_light.mean():.4f}")

# Step 3: Morphological opening (disk r=2) to remove small noise artifacts
selem = morphology.disk(2)
binary_light_clean = morphology.opening(binary_light, selem)
binary_dark_clean = morphology.opening(binary_dark, selem)

# Ensure complementary masks after cleaning (assign unassigned pixels to nearest phase)
# After independent openings, some pixels may be neither. Reassign based on original threshold.
both_zero = ~binary_light_clean & ~binary_dark_clean
binary_light_clean[both_zero] = binary_light[both_zero]
binary_dark_clean[both_zero] = binary_dark[both_zero]

# Step 4: Connected component labeling on each phase
labeled_light = measure.label(binary_light_clean, connectivity=2)
labeled_dark = measure.label(binary_dark_clean, connectivity=2)

# Step 5: Region property extraction
props_light = measure.regionprops(labeled_light, intensity_image=image_2d)
props_dark = measure.regionprops(labeled_dark, intensity_image=image_2d)

# Filter out very small artifacts (< 10 pixels)
min_area_filter = 10
props_light = [p for p in props_light if p.area >= min_area_filter]
props_dark = [p for p in props_dark if p.area >= min_area_filter]

# Phase fractions
total_pixels = image_2d.size
phase_fraction_light = float(binary_light_clean.sum()) / total_pixels
phase_fraction_dark = float(binary_dark_clean.sum()) / total_pixels

# Domain counts
domain_count_light = len(props_light)
domain_count_dark = len(props_dark)

# Domain area distributions
areas_light = [float(p.area) for p in props_light]
areas_dark = [float(p.area) for p in props_dark]

# Equivalent diameter distributions
eq_diameters_light = [float(p.equivalent_diameter) for p in props_light]
eq_diameters_dark = [float(p.equivalent_diameter) for p in props_dark]

# Perimeters for boundary length density
perimeters_light = [float(p.perimeter) for p in props_light]
perimeters_dark = [float(p.perimeter) for p in props_dark]

# Boundary length density (total boundary per unit area)
# Use the interface between phases
boundary = binary_light_clean.astype(np.uint8)
boundary_edges = filters.sobel(boundary.astype(float)) > 0
total_boundary_pixels = boundary_edges.sum()
boundary_length_density = float(total_boundary_pixels) / total_pixels

# Step 6: Local thickness measurement via distance transform
dist_light = ndimage.distance_transform_edt(binary_light_clean)
dist_dark = ndimage.distance_transform_edt(binary_dark_clean)

# Local thickness ~ 2 * distance transform value at skeleton/medial axis
# Use the medial axis (skeleton) to sample representative thickness values
from skimage.morphology import skeletonize

skel_light = skeletonize(binary_light_clean)
skel_dark = skeletonize(binary_dark_clean)

# Local thickness at skeleton points
thickness_at_skel_light = dist_light[skel_light] * 2  # diameter = 2 * radius
thickness_at_skel_dark = dist_dark[skel_dark] * 2

if len(thickness_at_skel_light) > 0:
    mean_local_thickness_light = float(np.mean(thickness_at_skel_light))
    mean_lamellar_spacing_light = float(np.median(thickness_at_skel_light))
else:
    mean_local_thickness_light = 0.0
    mean_lamellar_spacing_light = 0.0

if len(thickness_at_skel_dark) > 0:
    mean_local_thickness_dark = float(np.mean(thickness_at_skel_dark))
    mean_lamellar_spacing_dark = float(np.median(thickness_at_skel_dark))
else:
    mean_local_thickness_dark = 0.0
    mean_lamellar_spacing_dark = 0.0

print(f"Mean local thickness light: {mean_local_thickness_light:.2f} px")
print(f"Mean local thickness dark: {mean_local_thickness_dark:.2f} px")
print(f"Mean lamellar spacing light: {mean_lamellar_spacing_light:.2f} px")
print(f"Mean lamellar spacing dark: {mean_lamellar_spacing_dark:.2f} px")

# Step 7: Spatial analysis of local phase fraction using sliding window (64x64)
window_size = 64
H, W = image_2d.shape
rows = H // window_size
cols = W // window_size
local_phase_fraction_map = np.zeros((rows, cols))

for i in range(rows):
    for j in range(cols):
        y0 = i * window_size
        x0 = j * window_size
        patch = binary_light_clean[y0:y0+window_size, x0:x0+window_size]
        local_phase_fraction_map[i, j] = patch.mean()

# Compute gradient of local phase fraction to detect spatial gradients
phase_fraction_gradient_mag = float(np.std(local_phase_fraction_map))
phase_fraction_range = float(local_phase_fraction_map.max() - local_phase_fraction_map.min())

print(f"Local phase fraction std: {phase_fraction_gradient_mag:.4f}")
print(f"Local phase fraction range: {phase_fraction_range:.4f}")

# Statistics for distributions
def distribution_stats(values):
    if len(values) == 0:
        return {'mean': 0, 'std': 0, 'min': 0, 'max': 0, 'median': 0, 'count': 0}
    return {
        'mean': float(np.mean(values)),
        'std': float(np.std(values)),
        'min': float(np.min(values)),
        'max': float(np.max(values)),
        'median': float(np.median(values)),
        'count': len(values)
    }

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

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

# Segmentation overlay
ax = axes[0, 1]
overlay = np.stack([image_2d]*3, axis=-1)
# Light phase in blue, dark phase in red
light_mask_rgb = np.zeros_like(overlay)
light_mask_rgb[:,:,2] = binary_light_clean.astype(float) * 0.4  # blue
dark_mask_rgb = np.zeros_like(overlay)
dark_mask_rgb[:,:,0] = binary_dark_clean.astype(float) * 0.4   # red
overlay_combined = np.clip(overlay * 0.6 + light_mask_rgb + dark_mask_rgb, 0, 1)
# Add boundary contours
from skimage import segmentation
boundary_contour = segmentation.find_boundaries(binary_light_clean, mode='thick')
overlay_combined[boundary_contour] = [1, 1, 0]  # yellow boundaries
ax.imshow(overlay_combined)
ax.set_title(f'Phase Segmentation\n(Blue=Light, Red=Dark, Yellow=Boundary)', fontsize=11)
ax.axis('off')

# Thresholded binary
ax = axes[0, 2]
ax.imshow(binary_light_clean, cmap='gray')
ax.set_title(f'Light Phase Mask (Otsu={threshold_otsu:.3f})', fontsize=12)
ax.axis('off')

# Local thickness map (light phase)
ax = axes[1, 0]
thickness_vis = dist_light * 2  # diameter
thickness_vis[~binary_light_clean] = 0
im = ax.imshow(thickness_vis, cmap='hot')
plt.colorbar(im, ax=ax, fraction=0.046)
ax.set_title(f'Local Thickness (Light)\nMean={mean_local_thickness_light:.1f} px', fontsize=11)
ax.axis('off')

# Local thickness map (dark phase)
ax = axes[1, 1]
thickness_vis_dark = dist_dark * 2
thickness_vis_dark[~binary_dark_clean] = 0
im2 = ax.imshow(thickness_vis_dark, cmap='hot')
plt.colorbar(im2, ax=ax, fraction=0.046)
ax.set_title(f'Local Thickness (Dark)\nMean={mean_local_thickness_dark:.1f} px', fontsize=11)
ax.axis('off')

# Local phase fraction map
ax = axes[1, 2]
im3 = ax.imshow(local_phase_fraction_map, cmap='RdBu', vmin=0, vmax=1, 
                extent=[0, W, H, 0], interpolation='nearest')
plt.colorbar(im3, ax=ax, fraction=0.046, label='Light phase fraction')
ax.set_title(f'Local Phase Fraction (64x64 window)\nStd={phase_fraction_gradient_mag:.3f}', fontsize=11)
ax.axis('off')

plt.suptitle('Eutectic Microstructure Analysis: Phase Segmentation & Lamellar Spacing', fontsize=14, fontweight='bold')
plt.tight_layout()
plt.savefig('analysis_visualization.png', dpi=100, bbox_inches='tight')
plt.close()
print("Saved analysis_visualization.png")

# Build results JSON
results = {
    "analysis_type": "Eutectic microstructure phase segmentation and lamellar spacing analysis using Otsu thresholding, morphological cleaning, distance-transform local thickness, and sliding-window spatial analysis",
    "extracted_features": {
        "phase_fraction_light": round(phase_fraction_light, 4),
        "phase_fraction_dark": round(phase_fraction_dark, 4),
        "otsu_threshold": round(float(threshold_otsu), 4),
        "mean_lamellar_spacing_light": round(mean_lamellar_spacing_light, 2),
        "mean_lamellar_spacing_dark": round(mean_lamellar_spacing_dark, 2),
        "mean_local_thickness_light": round(mean_local_thickness_light, 2),
        "mean_local_thickness_dark": round(mean_local_thickness_dark, 2),
        "domain_count_light": domain_count_light,
        "domain_count_dark": domain_count_dark,
        "domain_area_distribution_light": distribution_stats(areas_light),
        "domain_area_distribution_dark": distribution_stats(areas_dark),
        "domain_equivalent_diameter_distribution_light": distribution_stats(eq_diameters_light),
        "domain_equivalent_diameter_distribution_dark": distribution_stats(eq_diameters_dark),
        "boundary_length_density": round(boundary_length_density, 4),
        "local_phase_fraction_map_std": round(phase_fraction_gradient_mag, 4),
        "local_phase_fraction_map_range": round(phase_fraction_range, 4),
        "local_phase_fraction_map_mean": round(float(np.mean(local_phase_fraction_map)), 4)
    },
    "quality_metrics": {
        "image_size": list(image_2d.shape),
        "intensity_range": [round(float(image_2d.min()), 4), round(float(image_2d.max()), 4)],
        "gaussian_sigma": 2,
        "morphological_disk_radius": 2,
        "sliding_window_size": window_size,
        "min_area_filter": min_area_filter,
        "skeleton_points_light": int(skel_light.sum()),
        "skeleton_points_dark": int(skel_dark.sum())
    },
    "summary": f"Eutectic microstructure with {phase_fraction_light:.1%} light phase and {phase_fraction_dark:.1%} dark phase. Light phase mean lamellar thickness is {mean_local_thickness_light:.1f} px and dark phase is {mean_local_thickness_dark:.1f} px. {domain_count_light} light domains and {domain_count_dark} dark domains detected. Local phase fraction std={phase_fraction_gradient_mag:.3f} indicates {'relatively uniform' if phase_fraction_gradient_mag < 0.05 else 'spatially varying'} phase distribution. Parameters: Gaussian sigma=2, morphological disk radius=2 as specified in plan; min_area filter=10 px to remove noise fragments."
}

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