import numpy as np
import cv2
import json
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from skimage import measure
from scipy.ndimage import gaussian_filter
from scilink.tools.sam import run_sam_analysis

# ── 1. Load image ──────────────────────────────────────────────────────────
path = '/Users/maxim.ziatdinov/Code/SciLink/analysis_session_20260401_122257/results/analysis__easy__vaterite_ImageAnalysis_20260401_122418_001/image_0000/temp_image_0.npy'
image = np.load(path)  # shape (1536, 2048, 3), uint8 RGB
assert image.ndim == 3 and image.shape[2] == 3, f"Unexpected shape {image.shape}"

# ── 2. Scale factor ────────────────────────────────────────────────────────
SCALE_UM_PER_PX = 0.08057  # µm per pixel
SCALE_UM2_PER_PX2 = SCALE_UM_PER_PX ** 2

# ── 3. Preprocessing per plan: green channel + Gaussian blur ───────────────
# Convert to grayscale using green channel (highest particle-background contrast)
green_channel = image[:, :, 1]  # uint8 2D array

# Apply gentle Gaussian blur (sigma=1.0) to reduce noise while preserving edges
green_blurred = gaussian_filter(green_channel.astype(np.float64), sigma=1.0)
green_blurred = np.clip(green_blurred, 0, 255).astype(np.uint8)

# ── 4. SAM segmentation ───────────────────────────────────────────────────
# Pass the preprocessed 2D grayscale image to SAM.
# min_area and max_area are in PIXELS (critical fix from plan).
MIN_AREA_PX = 2500   # minimum particle area in pixels
MAX_AREA_PX = 80000  # maximum particle area in pixels

# Per plan: lower pred_iou_thresh and stability_score_thresh toward 0.8
# to improve sensitivity for darker edge-on particles.
# Try passing these as part of sam_parameters. If the API only accepts
# string values for sam_parameters, fall back to "default".
try:
    result = run_sam_analysis(
        green_blurred,
        params={
            "sam_parameters": {
                "pred_iou_thresh": 0.8,
                "stability_score_thresh": 0.8,
            },
            "min_area": MIN_AREA_PX,
            "max_area": MAX_AREA_PX,
            "pruning_iou_threshold": 0.5,
        }
    )
except Exception:
    # Fallback: use "default" string if dict-based sam_parameters not supported
    result = run_sam_analysis(
        green_blurred,
        params={
            "sam_parameters": "default",
            "min_area": MIN_AREA_PX,
            "max_area": MAX_AREA_PX,
            "pruning_iou_threshold": 0.5,
        }
    )

particles_raw = result["particles"]
print(f"SAM returned {len(particles_raw)} particles after area filtering.")

# ── 5. Build initial label map & extract regionprops ──────────────────────
h, w = image.shape[:2]
labeled_raw = np.zeros((h, w), dtype=np.int32)
for i, p in enumerate(particles_raw):
    mask = np.array(p["mask"], dtype=bool)
    labeled_raw[mask] = i + 1

# Use green channel (unblurred) as intensity image for regionprops
props_raw = measure.regionprops(labeled_raw, intensity_image=green_channel)

# ── 6. Post-processing filters ────────────────────────────────────────────
# Stage 1: Verify area in pixel units (double-check against MIN/MAX)
# Stage 2: Solidity floor >= 0.3
# Stage 3: Circularity floor >= 0.15

def compute_circularity(area, perimeter):
    if perimeter == 0:
        return 0.0
    return 4.0 * np.pi * area / (perimeter ** 2)

valid_labels = []
rejected_labels = []
rejection_reasons = {}

for prop in props_raw:
    label = prop.label
    area_px = prop.area
    perimeter_px = prop.perimeter
    solidity = prop.solidity
    circ = compute_circularity(area_px, perimeter_px)
    
    # Stage 1: area filter in pixel units (safety check)
    if area_px < MIN_AREA_PX or area_px > MAX_AREA_PX:
        rejected_labels.append(label)
        rejection_reasons[label] = f"area_px={area_px} outside [{MIN_AREA_PX},{MAX_AREA_PX}]"
        continue
    
    # Stage 2: solidity floor
    if solidity < 0.3:
        rejected_labels.append(label)
        rejection_reasons[label] = f"solidity={solidity:.4f} < 0.3"
        continue
    
    # Stage 3: circularity floor
    if circ < 0.15:
        rejected_labels.append(label)
        rejection_reasons[label] = f"circularity={circ:.4f} < 0.15"
        continue
    
    valid_labels.append(label)

print(f"After filtering: {len(valid_labels)} valid, {len(rejected_labels)} rejected.")
for lbl, reason in rejection_reasons.items():
    print(f"  Rejected label {lbl}: {reason}")

# ── 7. Build clean label map ──────────────────────────────────────────────
label_map_clean = np.zeros((h, w), dtype=np.int32)
new_label = 0
label_mapping = {}  # old_label -> new_label
for old_label in valid_labels:
    new_label += 1
    label_mapping[old_label] = new_label
    label_map_clean[labeled_raw == old_label] = new_label

props_clean = measure.regionprops(label_map_clean, intensity_image=green_channel)

# ── 8. Extract features & detect border-touching particles ────────────────
feature_records = []

for prop in props_clean:
    label = prop.label
    area_px = prop.area
    perimeter_px = prop.perimeter
    eq_diam_px = prop.equivalent_diameter
    solidity = prop.solidity
    eccentricity = prop.eccentricity
    centroid_y, centroid_x = prop.centroid
    major_ax_px = prop.major_axis_length
    minor_ax_px = prop.minor_axis_length
    convex_area_px = prop.convex_area
    circ = compute_circularity(area_px, perimeter_px)
    aspect_ratio = (major_ax_px / minor_ax_px) if minor_ax_px > 0 else float('inf')
    
    # Border-touching detection
    coords = prop.coords  # (row, col)
    rows = coords[:, 0]
    cols = coords[:, 1]
    border_touching = bool(
        np.any(rows == 0) or np.any(rows == h - 1) or
        np.any(cols == 0) or np.any(cols == w - 1)
    )
    
    # Convert to physical units
    area_um2 = area_px * SCALE_UM2_PER_PX2
    perimeter_um = perimeter_px * SCALE_UM_PER_PX
    eq_diam_um = eq_diam_px * SCALE_UM_PER_PX
    major_ax_um = major_ax_px * SCALE_UM_PER_PX
    minor_ax_um = minor_ax_px * SCALE_UM_PER_PX
    
    feature_records.append({
        "label": int(label),
        "area_pixels": int(area_px),
        "area_um2": round(float(area_um2), 4),
        "perimeter_pixels": round(float(perimeter_px), 2),
        "perimeter_um": round(float(perimeter_um), 4),
        "equivalent_diameter_um": round(float(eq_diam_um), 4),
        "circularity": round(float(circ), 4),
        "solidity": round(float(solidity), 4),
        "eccentricity": round(float(eccentricity), 4),
        "centroid_x": round(float(centroid_x), 2),
        "centroid_y": round(float(centroid_y), 2),
        "major_axis_length_um": round(float(major_ax_um), 4),
        "minor_axis_length_um": round(float(minor_ax_um), 4),
        "aspect_ratio": round(float(aspect_ratio), 4),
        "border_touching_flag": border_touching,
        "convex_area_pixels": int(convex_area_px),
    })

num_particles = len(feature_records)
num_border = sum(1 for f in feature_records if f["border_touching_flag"])
num_interior = num_particles - num_border

# Statistics for interior particles only
interior = [f for f in feature_records if not f["border_touching_flag"]]
if interior:
    diams = [f["equivalent_diameter_um"] for f in interior]
    circs = [f["circularity"] for f in interior]
    solidities = [f["solidity"] for f in interior]
    mean_diam = float(np.mean(diams))
    std_diam = float(np.std(diams))
    median_diam = float(np.median(diams))
    mean_circ = float(np.mean(circs))
    mean_solidity = float(np.mean(solidities))
else:
    mean_diam = std_diam = median_diam = mean_circ = mean_solidity = 0.0
    diams = []
    circs = []

all_diams = [f["equivalent_diameter_um"] for f in feature_records]
all_circs = [f["circularity"] for f in feature_records]

# ── 9. Save label map ─────────────────────────────────────────────────────
np.save("analysis_labels.npy", label_map_clean)

# ── 10. Visualization ─────────────────────────────────────────────────────
fig, axes = plt.subplots(2, 3, figsize=(20, 13), dpi=100)

# 10a. Original image
ax = axes[0, 0]
ax.imshow(image)
ax.set_title("Original Image", fontsize=12)
ax.axis('off')

# 10b. Segmentation overlay
ax = axes[0, 1]
overlay = image.copy()
np.random.seed(42)
colors_list = []
for i in range(num_particles):
    c = np.random.randint(50, 255, size=3).tolist()
    colors_list.append(c)

for idx, prop in enumerate(props_clean):
    mask_bool = label_map_clean == prop.label
    color = colors_list[idx]
    for ch in range(3):
        overlay[:, :, ch] = np.where(mask_bool,
            np.clip(overlay[:, :, ch].astype(np.float32) * 0.5 + color[ch] * 0.5, 0, 255).astype(np.uint8),
            overlay[:, :, ch])

# Draw contours
for idx, prop in enumerate(props_clean):
    mask_uint8 = ((label_map_clean == prop.label).astype(np.uint8)) * 255
    contours, _ = cv2.findContours(mask_uint8, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    rec = feature_records[idx]
    if rec["border_touching_flag"]:
        cv2.drawContours(overlay, contours, -1, (255, 0, 0), 2)  # Red for border
    else:
        cv2.drawContours(overlay, contours, -1, (0, 255, 0), 2)  # Green for interior

ax.imshow(overlay)
ax.set_title(f"Segmentation Overlay ({num_particles} particles)\nGreen=interior, Red=border-touching", fontsize=11)
ax.axis('off')

# 10c. Rejected detections overlay (for verification)
ax = axes[0, 2]
reject_overlay = image.copy()
for rlabel in rejected_labels:
    mask_bool = labeled_raw == rlabel
    reject_overlay[:, :, 0] = np.where(mask_bool, 255, reject_overlay[:, :, 0])
    reject_overlay[:, :, 1] = np.where(mask_bool,
        (reject_overlay[:, :, 1].astype(np.float32) * 0.3).astype(np.uint8),
        reject_overlay[:, :, 1])
    reject_overlay[:, :, 2] = np.where(mask_bool,
        (reject_overlay[:, :, 2].astype(np.float32) * 0.3).astype(np.uint8),
        reject_overlay[:, :, 2])
ax.imshow(reject_overlay)
ax.set_title(f"Rejected Detections ({len(rejected_labels)})", fontsize=12)
ax.axis('off')

# 10d. Size distribution histogram (all particles)
ax = axes[1, 0]
if all_diams:
    ax.hist(all_diams, bins=15, color='steelblue', edgecolor='black', alpha=0.8)
ax.set_xlabel('Equivalent Diameter (\u00b5m)', fontsize=11)
ax.set_ylabel('Count', fontsize=11)
ax.set_title(f'Size Distribution (all {num_particles} particles)', fontsize=12)
ax.axvline(np.mean(all_diams) if all_diams else 0, color='red', linestyle='--', label=f'Mean={np.mean(all_diams):.1f} \u00b5m' if all_diams else '')
ax.legend(fontsize=9)

# 10e. Circularity distribution
ax = axes[1, 1]
if all_circs:
    ax.hist(all_circs, bins=15, color='coral', edgecolor='black', alpha=0.8)
ax.set_xlabel('Circularity', fontsize=11)
ax.set_ylabel('Count', fontsize=11)
ax.set_title(f'Circularity Distribution', fontsize=12)
if all_circs:
    ax.axvline(np.mean(all_circs), color='red', linestyle='--', label=f'Mean={np.mean(all_circs):.2f}')
    ax.legend(fontsize=9)

# 10f. Scatter: diameter vs circularity
ax = axes[1, 2]
for f in feature_records:
    c = 'red' if f['border_touching_flag'] else 'steelblue'
    m = 'x' if f['border_touching_flag'] else 'o'
    ax.scatter(f['equivalent_diameter_um'], f['circularity'], c=c, marker=m, s=50, alpha=0.8)
ax.set_xlabel('Equivalent Diameter (\u00b5m)', fontsize=11)
ax.set_ylabel('Circularity', fontsize=11)
ax.set_title('Diameter vs Circularity\nBlue=interior, Red=border', fontsize=12)
ax.grid(True, alpha=0.3)

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

# ── 11. JSON results ──────────────────────────────────────────────────────
results = {
    "analysis_type": "SAM instance segmentation of vaterite particles with green-channel grayscale + Gaussian blur preprocessing and post-processing filters (solidity>=0.3, circularity>=0.15, area 2500-80000 pixels)",
    "extracted_features": {
        "total_particles_detected": num_particles,
        "interior_particles": num_interior,
        "border_touching_particles": num_border,
        "rejected_detections": len(rejected_labels),
        "mean_equivalent_diameter_um": round(mean_diam, 4),
        "std_equivalent_diameter_um": round(std_diam, 4),
        "median_equivalent_diameter_um": round(median_diam, 4),
        "mean_circularity_interior": round(mean_circ, 4),
        "mean_solidity_interior": round(mean_solidity, 4),
        "min_equivalent_diameter_um": round(min(all_diams), 4) if all_diams else 0,
        "max_equivalent_diameter_um": round(max(all_diams), 4) if all_diams else 0,
        "scale_um_per_pixel": SCALE_UM_PER_PX,
        "min_area_pixels": MIN_AREA_PX,
        "max_area_pixels": MAX_AREA_PX,
        "per_particle_features": feature_records,
    },
    "quality_metrics": {
        "sam_raw_count": len(particles_raw),
        "post_filter_count": num_particles,
        "solidity_filter_threshold": 0.3,
        "circularity_filter_threshold": 0.15,
        "num_rejected_by_filters": len(rejected_labels),
        "rejection_details": {str(k): v for k, v in rejection_reasons.items()},
        "preprocessing": "Green channel extraction + Gaussian blur (sigma=1.0)",
    },
    "summary": f"Detected {num_particles} vaterite particles ({num_interior} interior, {num_border} border-touching) after SAM segmentation on green-channel grayscale with Gaussian blur (sigma=1.0) and post-processing. {len(rejected_labels)} detections rejected. Interior particles: mean equivalent diameter = {mean_diam:.2f} \u00b1 {std_diam:.2f} \u00b5m, mean circularity = {mean_circ:.3f}, mean solidity = {mean_solidity:.3f}. Area filters applied strictly in pixel units (2500-80000 px) before \u00b5m conversion.",
    "saved_arrays": {
        "analysis_labels.npy": {
            "description": f"Integer label map, {num_particles} particles labeled 1-{num_particles}, background=0. Border-touching particles included but flagged in feature records.",
            "shape": list(label_map_clean.shape),
            "dtype": str(label_map_clean.dtype),
        }
    }
}

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