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

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

# Step 1: Convert uint16 to float and scale to physical height range
height_min = -531.156  # nm
height_max = -455.713  # nm
img_float = img_raw.astype(np.float64)
img_float = img_float / 65535.0 * (height_max - height_min) + height_min
print(f"Height range: [{img_float.min():.3f}, {img_float.max():.3f}] nm")

# Step 2: Fit and subtract 2nd-order polynomial background
rows, cols = img_float.shape
y_coords, x_coords = np.mgrid[0:rows, 0:cols]
y_flat = y_coords.ravel().astype(np.float64)
x_flat = x_coords.ravel().astype(np.float64)
z_flat = img_float.ravel()

# Build design matrix for 2nd order polynomial: 1, x, y, x^2, xy, y^2
A = np.column_stack([
    np.ones_like(x_flat),
    x_flat, y_flat,
    x_flat**2, x_flat*y_flat, y_flat**2
])
coeffs, _, _, _ = np.linalg.lstsq(A, z_flat, rcond=None)
background = (A @ coeffs).reshape(rows, cols)
img_corrected = img_float - background
print(f"Background-corrected range: [{img_corrected.min():.3f}, {img_corrected.max():.3f}] nm")

# Step 3: Apply light Gaussian smoothing (sigma=1.0 px)
img_smoothed = ndimage.gaussian_filter(img_corrected, sigma=1.0)

# Step 4: Convert corrected height map to uint8 for SAM input
img_norm = img_smoothed - img_smoothed.min()
if img_norm.max() > 0:
    img_norm = img_norm / img_norm.max()
img_uint8 = (img_norm * 255).astype(np.uint8)

# Pixel size calculation
pixel_size_nm = 2000.0 / 256.0  # 7.8125 nm/px
pixel_area_nm2 = pixel_size_nm ** 2  # 61.035 nm^2/px

# Step 5: Run SAM instance segmentation
# Start with default parameters, using appropriate area filters
# Expected grain size: ~40-80 nm diameter -> ~5-10 px diameter -> area ~20-80 px
# But with 2000nm / 256px = 7.8125 nm/px, grains of ~100nm diameter = ~13 px diameter = ~130 px area
# Larger grains ~200nm = ~26 px = ~530 px area
# Use relaxed min_area to catch smaller grains
min_area_px = 400
max_area_px = 30000

print("Running SAM with default parameters...")
result = run_sam_analysis(img_uint8, params={
    "sam_parameters": "default",
    "min_area": min_area_px,
    "max_area": max_area_px,
    "pruning_iou_threshold": 0.5,
    "use_clahe": False
})

particles = result.get("particles", [])
total_count = result.get("total_count", len(particles))
print(f"SAM detected {total_count} particles")

# Build labeled mask from SAM particles
labeled = np.zeros(img_uint8.shape[:2], dtype=np.int32)
for i, p in enumerate(particles):
    mask = np.array(p["mask"], dtype=bool)
    labeled[mask] = i + 1

# Check coverage
coverage = np.sum(labeled > 0) / (rows * cols)
print(f"Initial coverage: {coverage:.3f}")

# If coverage is too low or count too low, try sensitive
if coverage < 0.5 or total_count < 15:
    print("Coverage or count too low, retrying with sensitive parameters...")
    result2 = run_sam_analysis(img_uint8, params={
        "sam_parameters": "sensitive",
        "min_area": 300,
        "max_area": max_area_px,
        "pruning_iou_threshold": 0.6,
        "use_clahe": False
    })
    particles2 = result2.get("particles", [])
    total_count2 = result2.get("total_count", len(particles2))
    labeled2 = np.zeros(img_uint8.shape[:2], dtype=np.int32)
    for i, p in enumerate(particles2):
        mask = np.array(p["mask"], dtype=bool)
        labeled2[mask] = i + 1
    coverage2 = np.sum(labeled2 > 0) / (rows * cols)
    print(f"Sensitive: {total_count2} particles, coverage: {coverage2:.3f}")
    if coverage2 > coverage:
        labeled = labeled2
        particles = particles2
        total_count = total_count2
        coverage = coverage2

# Step 6: Morphological closing on each mask
print("Applying morphological closing...")
selem = morphology.disk(3)
unique_labels = np.unique(labeled)
unique_labels = unique_labels[unique_labels > 0]
labeled_closed = np.zeros_like(labeled)
for lbl in unique_labels:
    mask = labeled == lbl
    mask_closed = morphology.binary_closing(mask, selem)
    labeled_closed[mask_closed & (labeled_closed == 0)] = lbl
labeled = labeled_closed

# Step 7: Merge adjacent segments with similar heights
print("Merging adjacent segments...")
def compute_grain_props(labeled_map, height_map):
    props = {}
    for lbl in np.unique(labeled_map):
        if lbl == 0:
            continue
        mask = labeled_map == lbl
        area = np.sum(mask)
        if area == 0:
            continue
        cy, cx = ndimage.center_of_mass(mask)
        mean_h = np.mean(height_map[mask])
        props[lbl] = {'area': area, 'cx': cx, 'cy': cy, 'mean_h': mean_h}
    return props

def merge_pass(labeled_map, height_map, dist_thresh=20, height_thresh=2.0):
    props = compute_grain_props(labeled_map, height_map)
    labels_list = list(props.keys())
    merged = False
    merge_map = {l: l for l in labels_list}
    
    def find_root(l):
        while merge_map[l] != l:
            merge_map[l] = merge_map[merge_map[l]]
            l = merge_map[l]
        return l
    
    for i in range(len(labels_list)):
        for j in range(i+1, len(labels_list)):
            li, lj = labels_list[i], labels_list[j]
            ri, rj = find_root(li), find_root(lj)
            if ri == rj:
                continue
            pi, pj = props[li], props[lj]
            dist = np.sqrt((pi['cx'] - pj['cx'])**2 + (pi['cy'] - pj['cy'])**2)
            if dist < dist_thresh and abs(pi['mean_h'] - pj['mean_h']) < height_thresh:
                merge_map[rj] = ri
                merged = True
    
    if merged:
        new_labeled = np.zeros_like(labeled_map)
        remap = {}
        new_id = 1
        for lbl in labels_list:
            root = find_root(lbl)
            if root not in remap:
                remap[root] = new_id
                new_id += 1
            new_labeled[labeled_map == lbl] = remap[root]
        return new_labeled, True
    return labeled_map, False

for iteration in range(10):
    labeled, did_merge = merge_pass(labeled, img_corrected, dist_thresh=20, height_thresh=2.0)
    if not did_merge:
        break
print(f"After merging: {len(np.unique(labeled)) - 1} segments")

# Step 8: Remove small segments after merging
min_area_post = 350
props_post = measure.regionprops(labeled)
labeled_clean = np.zeros_like(labeled)
new_id = 1
for prop in props_post:
    if prop.area >= min_area_post:
        labeled_clean[labeled == prop.label] = new_id
        new_id += 1
labeled = labeled_clean
print(f"After size filter: {len(np.unique(labeled)) - 1} segments")

# Step 9: Validate count - if over-segmented, increase merge threshold
num_grains = len(np.unique(labeled)) - 1
print(f"Grain count: {num_grains}")

if num_grains > 50:
    print("Over-segmented, re-merging with relaxed height threshold...")
    for iteration in range(10):
        labeled, did_merge = merge_pass(labeled, img_corrected, dist_thresh=20, height_thresh=3.0)
        if not did_merge:
            break
    # Re-clean
    props_post = measure.regionprops(labeled)
    labeled_clean = np.zeros_like(labeled)
    new_id = 1
    for prop in props_post:
        if prop.area >= min_area_post:
            labeled_clean[labeled == prop.label] = new_id
            new_id += 1
    labeled = labeled_clean
    num_grains = len(np.unique(labeled)) - 1
    print(f"After re-merge: {num_grains} grains")

# Step 10: Identify edge-touching grains
edge_mask = np.zeros((rows, cols), dtype=bool)
edge_mask[0, :] = True
edge_mask[-1, :] = True
edge_mask[:, 0] = True
edge_mask[:, -1] = True

props_final = measure.regionprops(labeled, intensity_image=img_corrected)

partial_labels = set()
full_labels = set()
for prop in props_final:
    grain_mask = labeled == prop.label
    if np.any(grain_mask & edge_mask):
        partial_labels.add(prop.label)
    else:
        full_labels.add(prop.label)

grain_count_total = len(props_final)
grain_count_full = len(full_labels)
grain_count_partial = len(partial_labels)
print(f"Total: {grain_count_total}, Full: {grain_count_full}, Partial (edge): {grain_count_partial}")

# Step 11 & 12: Extract grain areas and equivalent diameters
grain_data = []
for prop in props_final:
    area_px = prop.area
    area_nm2 = area_px * pixel_area_nm2
    equiv_diam_nm = 2 * np.sqrt(area_nm2 / np.pi)
    is_partial = prop.label in partial_labels
    cy, cx = prop.centroid
    mean_height = prop.mean_intensity
    grain_data.append({
        'label': prop.label,
        'area_px': area_px,
        'area_nm2': area_nm2,
        'equiv_diam_nm': equiv_diam_nm,
        'is_partial': is_partial,
        'centroid_x': cx,
        'centroid_y': cy,
        'mean_height_nm': mean_height
    })

areas_nm2 = np.array([g['area_nm2'] for g in grain_data])
diameters_nm = np.array([g['equiv_diam_nm'] for g in grain_data])

# Step 13: Statistics
if len(areas_nm2) > 0:
    mean_area = float(np.mean(areas_nm2))
    median_area = float(np.median(areas_nm2))
    std_area = float(np.std(areas_nm2))
    min_area = float(np.min(areas_nm2))
    max_area = float(np.max(areas_nm2))
    mean_diam = float(np.mean(diameters_nm))
    median_diam = float(np.median(diameters_nm))
    std_diam = float(np.std(diameters_nm))
    min_diam = float(np.min(diameters_nm))
    max_diam = float(np.max(diameters_nm))
else:
    mean_area = median_area = std_area = min_area = max_area = 0.0
    mean_diam = median_diam = std_diam = min_diam = max_diam = 0.0

# Full grains only stats
full_areas = np.array([g['area_nm2'] for g in grain_data if not g['is_partial']])
full_diams = np.array([g['equiv_diam_nm'] for g in grain_data if not g['is_partial']])
if len(full_areas) > 0:
    mean_area_full = float(np.mean(full_areas))
    median_area_full = float(np.median(full_areas))
    std_area_full = float(np.std(full_areas))
else:
    mean_area_full = median_area_full = std_area_full = 0.0

# Coverage
coverage_final = float(np.sum(labeled > 0) / (rows * cols))
print(f"Final coverage: {coverage_final:.3f}")

# Create grain boundary mask
grain_boundary = np.zeros((rows, cols), dtype=bool)
for prop in props_final:
    mask = labeled == prop.label
    eroded = morphology.binary_erosion(mask, morphology.disk(1))
    boundary = mask & ~eroded
    grain_boundary |= boundary

# Step 14: Generate histogram
fig, axes = plt.subplots(2, 3, figsize=(18, 12), dpi=100)

# Original image
ax = axes[0, 0]
im = ax.imshow(img_float, cmap='afmhot')
ax.set_title('Original AFM Height Map')
plt.colorbar(im, ax=ax, label='Height (nm)')

# Background-corrected
ax = axes[0, 1]
im = ax.imshow(img_corrected, cmap='afmhot')
ax.set_title('Background-Corrected Height')
plt.colorbar(im, ax=ax, label='Height (nm)')

# Segmentation overlay
ax = axes[0, 2]
ax.imshow(img_corrected, cmap='gray', alpha=0.7)
# Create colored overlay
from matplotlib.colors import ListedColormap
np.random.seed(42)
colors = np.random.rand(grain_count_total + 1, 4)
colors[:, 3] = 0.4  # semi-transparent
colors[0] = [0, 0, 0, 0]  # background transparent
for prop in props_final:
    mask = labeled == prop.label
    color = colors[prop.label % (grain_count_total + 1)]
    overlay = np.zeros((rows, cols, 4))
    overlay[mask] = color
    ax.imshow(overlay)
# Draw boundaries
boundary_overlay = np.zeros((rows, cols, 4))
boundary_overlay[grain_boundary] = [1, 0, 0, 0.8]
ax.imshow(boundary_overlay)
ax.set_title(f'Segmentation Overlay ({grain_count_total} grains)')

# Labeled map
ax = axes[1, 0]
im = ax.imshow(labeled, cmap='nipy_spectral', interpolation='nearest')
ax.set_title('Labeled Grain Map')
plt.colorbar(im, ax=ax, label='Grain ID')

# Histogram of equivalent diameters
ax = axes[1, 1]
if len(diameters_nm) > 0:
    bin_width = 20  # nm
    bins = np.arange(min_diam - bin_width/2, max_diam + bin_width, bin_width)
    if len(bins) < 3:
        bins = 10
    ax.hist(diameters_nm, bins=bins, color='steelblue', edgecolor='black', alpha=0.7)
    ax.axvline(mean_diam, color='red', linestyle='--', label=f'Mean: {mean_diam:.1f} nm')
    ax.axvline(median_diam, color='green', linestyle='--', label=f'Median: {median_diam:.1f} nm')
    ax.legend()
ax.set_xlabel('Equivalent Diameter (nm)')
ax.set_ylabel('Count')
ax.set_title('Grain Size Distribution')

# Area histogram
ax = axes[1, 2]
if len(areas_nm2) > 0:
    ax.hist(areas_nm2, bins=15, color='coral', edgecolor='black', alpha=0.7)
    ax.axvline(mean_area, color='red', linestyle='--', label=f'Mean: {mean_area:.0f} nm²')
    ax.legend()
ax.set_xlabel('Grain Area (nm²)')
ax.set_ylabel('Count')
ax.set_title('Grain Area Distribution')

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)
print("Saved analysis_labels.npy")

# Save grain boundary mask
np.save('grain_boundary_mask.npy', grain_boundary.astype(np.uint8))
print("Saved grain_boundary_mask.npy")

# Build histogram data for JSON
if len(diameters_nm) > 0:
    bin_width = 20
    bins = np.arange(min_diam - bin_width/2, max_diam + bin_width, bin_width)
    hist_counts, hist_edges = np.histogram(diameters_nm, bins=bins)
    histogram_data = {
        'bin_edges_nm': hist_edges.tolist(),
        'counts': hist_counts.tolist(),
        'bin_width_nm': bin_width
    }
else:
    histogram_data = {'bin_edges_nm': [], 'counts': [], 'bin_width_nm': 20}

centroid_positions = [[g['centroid_x'], g['centroid_y']] for g in grain_data]
mean_heights = [g['mean_height_nm'] for g in grain_data]

results = {
    "analysis_type": "AFM Au grain instance segmentation using SAM with polynomial background correction, Gaussian smoothing, morphological post-processing, and adjacent segment merging",
    "extracted_features": {
        "grain_count_total": grain_count_total,
        "grain_count_full": grain_count_full,
        "grain_count_partial_edge": grain_count_partial,
        "individual_grain_areas_nm2": [round(a, 1) for a in areas_nm2.tolist()] if len(areas_nm2) > 0 else [],
        "individual_grain_equivalent_diameters_nm": [round(d, 1) for d in diameters_nm.tolist()] if len(diameters_nm) > 0 else [],
        "mean_grain_area_nm2": round(mean_area, 1),
        "median_grain_area_nm2": round(median_area, 1),
        "std_grain_area_nm2": round(std_area, 1),
        "min_grain_area_nm2": round(min_area, 1),
        "max_grain_area_nm2": round(max_area, 1),
        "mean_equivalent_diameter_nm": round(mean_diam, 1),
        "median_equivalent_diameter_nm": round(median_diam, 1),
        "std_equivalent_diameter_nm": round(std_diam, 1),
        "min_equivalent_diameter_nm": round(min_diam, 1),
        "max_equivalent_diameter_nm": round(max_diam, 1),
        "grain_size_histogram": histogram_data,
        "segmentation_coverage_fraction": round(coverage_final, 4),
        "grain_centroid_positions": [[round(c[0], 1), round(c[1], 1)] for c in centroid_positions],
        "grain_mean_heights_nm": [round(h, 3) for h in mean_heights],
        "mean_grain_area_full_grains_nm2": round(mean_area_full, 1),
        "median_grain_area_full_grains_nm2": round(median_area_full, 1),
        "std_grain_area_full_grains_nm2": round(std_area_full, 1),
        "pixel_size_nm": round(pixel_size_nm, 4),
        "pixel_area_nm2": round(pixel_area_nm2, 3)
    },
    "quality_metrics": {
        "segmentation_coverage_fraction": round(coverage_final, 4),
        "grain_count_in_expected_range": 25 <= grain_count_total <= 50,
        "coverage_above_threshold": coverage_final >= 0.75,
        "background_correction": "2nd-order polynomial",
        "smoothing_sigma_px": 1.0
    },
    "summary": f"Detected {grain_count_total} Au grains ({grain_count_full} full, {grain_count_partial} edge-touching) with mean equivalent diameter {mean_diam:.1f} nm and coverage {coverage_final:.1%}. SAM instance segmentation with polynomial background correction and segment merging was applied. Pixel size: {pixel_size_nm:.4f} nm/px.",
    "saved_arrays": {
        "analysis_labels.npy": {
            "description": f"Integer label map, {grain_count_total} grains labeled 1-{grain_count_total}, background=0",
            "shape": list(labeled.shape),
            "dtype": str(labeled.dtype)
        },
        "grain_boundary_mask.npy": {
            "description": "Binary grain boundary mask (1=boundary, 0=interior/background)",
            "shape": list(grain_boundary.shape),
            "dtype": "uint8"
        }
    }
}

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