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

# Load image
path = '/Users/maxim.ziatdinov/Code/SciLink/test_tier2_output/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:
    img_2d = image[:, :, 0]
else:
    img_2d = image.copy()

# Try SAM first
try:
    from scilink.tools.sam import run_sam_analysis
    
    # Run SAM with default parameters, min_area=1500 to filter small fragments
    result = run_sam_analysis(img_2d, params={
        "sam_parameters": "default",
        "min_area": 1500,
        "max_area": 200000,
        "pruning_iou_threshold": 0.5
    })
    
    print(f"SAM detected {result['total_count']} particles before post-filtering")
    
    # Build labeled mask from SAM particles
    labeled = np.zeros(img_2d.shape[:2], dtype=np.int32)
    for i, p in enumerate(result["particles"]):
        mask = np.array(p["mask"], dtype=bool)
        labeled[mask] = i + 1
    
    use_sam = True
    print("SAM analysis succeeded, applying post-filtering...")
    
except Exception as e:
    print(f"SAM failed: {e}, falling back to watershed")
    use_sam = False

if not use_sam:
    # Fallback: Canny edge detection + watershed
    # Canny edge detection with specified parameters
    edges = feature.canny(img_2d.astype(float), sigma=2, low_threshold=0.05*255, high_threshold=0.15*255)
    
    # Morphological dilation to thicken edges
    dilated_edges = morphology.dilation(edges, morphology.disk(2))
    
    # Invert to get grain interior mask
    grain_mask = ~dilated_edges
    
    # Morphological opening to clean boundary residue (disk r=5)
    grain_mask = morphology.opening(grain_mask, morphology.disk(5))
    
    # Distance transform
    distance = ndimage.distance_transform_edt(grain_mask)
    
    # Peak local max for markers (min_distance=25)
    coords = feature.peak_local_max(distance, min_distance=25, labels=grain_mask)
    
    # Create markers
    marker_mask = np.zeros(distance.shape, dtype=bool)
    marker_mask[tuple(coords.T)] = True
    labeled_markers = ndimage.label(marker_mask)[0]
    
    # Use Sobel gradient as watershed landscape (advanced tip for better boundaries)
    gradient = ndimage.generic_gradient_magnitude(img_2d.astype(float), ndimage.sobel)
    
    # Watershed segmentation
    labeled = segmentation.watershed(gradient, labeled_markers, mask=grain_mask)
    print(f"Watershed detected {labeled.max()} regions before post-filtering")

# Post-filtering
props = measure.regionprops(labeled, intensity_image=img_2d)

regions_removed_area = 0
regions_removed_solidity = 0
regions_removed_circularity = 0
regions_removed_zero_perimeter = 0

valid_labels = []
for prop in props:
    # Filter by area
    if prop.area < 1500:
        regions_removed_area += 1
        continue
    
    # Filter by zero perimeter
    if prop.perimeter == 0:
        regions_removed_zero_perimeter += 1
        continue
    
    # Calculate circularity
    circularity = (4 * np.pi * prop.area) / (prop.perimeter ** 2) if prop.perimeter > 0 else 0
    
    # Filter by solidity
    if prop.solidity < 0.5:
        regions_removed_solidity += 1
        continue
    
    # Filter by circularity
    if circularity < 0.1:
        regions_removed_circularity += 1
        continue
    
    valid_labels.append(prop.label)

print(f"Removed by area: {regions_removed_area}, solidity: {regions_removed_solidity}, circularity: {regions_removed_circularity}, zero perimeter: {regions_removed_zero_perimeter}")
print(f"Valid grains after filtering: {len(valid_labels)}")

# Create filtered labeled image
filtered_labeled = np.zeros_like(labeled)
for i, lbl in enumerate(valid_labels):
    filtered_labeled[labeled == lbl] = i + 1

# Extract properties from filtered regions
filtered_props = measure.regionprops(filtered_labeled, intensity_image=img_2d)

# Extract all features
grain_areas = []
grain_perimeters = []
grain_eq_diameters = []
grain_solidities = []
grain_circularities = []
grain_aspect_ratios = []
grain_major_axes = []
grain_minor_axes = []
grain_orientations = []

for prop in filtered_props:
    grain_areas.append(float(prop.area))
    grain_perimeters.append(float(prop.perimeter))
    grain_eq_diameters.append(float(prop.equivalent_diameter))
    grain_solidities.append(float(prop.solidity))
    
    circ = (4 * np.pi * prop.area) / (prop.perimeter ** 2) if prop.perimeter > 0 else 0
    grain_circularities.append(float(circ))
    
    major = prop.major_axis_length
    minor = prop.minor_axis_length
    grain_major_axes.append(float(major))
    grain_minor_axes.append(float(minor))
    grain_orientations.append(float(prop.orientation))
    
    ar = major / minor if minor > 0 else float('inf')
    grain_aspect_ratios.append(float(ar))

grain_count = len(grain_areas)
total_image_area = img_2d.shape[0] * img_2d.shape[1]
total_segmented_area = sum(grain_areas)
area_fraction = total_segmented_area / total_image_area

# Histogram of grain sizes (equivalent diameters)
if grain_eq_diameters:
    hist_values, hist_edges = np.histogram(grain_eq_diameters, bins=10)
    hist_dict = {f"{hist_edges[i]:.1f}-{hist_edges[i+1]:.1f}": int(hist_values[i]) for i in range(len(hist_values))}
else:
    hist_dict = {}

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

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

# Segmentation overlay
axes[1].imshow(img_2d, cmap='gray')
if grain_count > 0:
    # Create colored overlay
    overlay = np.zeros((*img_2d.shape, 4), dtype=float)
    cmap = plt.cm.get_cmap('tab20', grain_count + 1)
    
    for i, prop in enumerate(filtered_props):
        color = cmap(i % 20)
        mask = filtered_labeled == prop.label
        overlay[mask, 0] = color[0]
        overlay[mask, 1] = color[1]
        overlay[mask, 2] = color[2]
        overlay[mask, 3] = 0.4
        
        # Draw contour
        contour_mask = mask.astype(np.uint8)
        contours_cv, _ = cv2.findContours(contour_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
        for cnt in contours_cv:
            pts = cnt.squeeze()
            if pts.ndim == 2:
                axes[1].plot(pts[:, 0], pts[:, 1], color=color[:3], linewidth=1.5)
    
    axes[1].imshow(overlay)
axes[1].set_title(f'Segmentation Overlay ({grain_count} grains)')
axes[1].axis('off')

# Grain size distribution
if grain_eq_diameters:
    axes[2].hist(grain_eq_diameters, bins=10, color='steelblue', edgecolor='black')
    axes[2].set_xlabel('Equivalent Diameter (pixels)')
    axes[2].set_ylabel('Count')
axes[2].set_title('Grain Size Distribution')

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

# Prepare results
extracted_features = {
    "grain_count_after_filtering": grain_count,
    "individual_grain_areas": grain_areas,
    "individual_grain_perimeters": grain_perimeters,
    "individual_grain_equivalent_diameters": grain_eq_diameters,
    "individual_grain_solidity": grain_solidities,
    "individual_grain_circularity": grain_circularities,
    "individual_grain_aspect_ratio": grain_aspect_ratios,
    "individual_grain_major_axis_length": grain_major_axes,
    "individual_grain_minor_axis_length": grain_minor_axes,
    "individual_grain_orientation": grain_orientations,
    "mean_grain_area": float(np.mean(grain_areas)) if grain_areas else 0,
    "median_grain_area": float(np.median(grain_areas)) if grain_areas else 0,
    "std_grain_area": float(np.std(grain_areas)) if grain_areas else 0,
    "min_grain_area": float(np.min(grain_areas)) if grain_areas else 0,
    "max_grain_area": float(np.max(grain_areas)) if grain_areas else 0,
    "mean_equivalent_diameter": float(np.mean(grain_eq_diameters)) if grain_eq_diameters else 0,
    "grain_size_distribution_histogram": hist_dict,
    "total_segmented_area_fraction": float(area_fraction),
    "regions_removed_by_area_filter": regions_removed_area,
    "regions_removed_by_solidity_filter": regions_removed_solidity,
    "regions_removed_by_circularity_filter": regions_removed_circularity,
    "regions_removed_by_zero_perimeter_filter": regions_removed_zero_perimeter
}

quality_metrics = {
    "total_regions_before_filtering": grain_count + regions_removed_area + regions_removed_solidity + regions_removed_circularity + regions_removed_zero_perimeter,
    "total_regions_after_filtering": grain_count,
    "fraction_regions_kept": grain_count / max(1, grain_count + regions_removed_area + regions_removed_solidity + regions_removed_circularity + regions_removed_zero_perimeter),
    "segmented_area_fraction": float(area_fraction),
    "method_used": "SAM" if use_sam else "Canny+Watershed"
}

method_str = "SAM instance segmentation" if use_sam else "Canny edge detection + Sobel gradient watershed"
summary = f"Detected {grain_count} grains using {method_str} with multi-criteria post-filtering (area>=1500, solidity>=0.5, circularity>=0.1); mean equivalent diameter = {extracted_features['mean_equivalent_diameter']:.1f} pixels, area fraction = {area_fraction:.3f}."

results = {
    "analysis_type": f"Grain segmentation and size analysis using {method_str} with aggressive post-filtering",
    "extracted_features": extracted_features,
    "quality_metrics": quality_metrics,
    "summary": summary
}

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