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

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

# Image is [1536, 2048, 3] uint8 RGB
image_h, image_w = image.shape[:2]

# Known field of view: 165 um across the image width
scale_um_per_px = 165.0 / image_w
print(f"Scale: {scale_um_per_px:.6f} um/pixel, image size: {image_w}x{image_h}")

# Step 1: CLAHE preprocessing on each RGB channel to enhance dark/low-contrast particles
clahe = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(8, 8))
image_clahe = np.zeros_like(image)
for c in range(3):
    image_clahe[:, :, c] = clahe.apply(image[:, :, c])

print("CLAHE preprocessing complete")

# Step 2: Convert CLAHE-enhanced image to grayscale for SAM input
# Per domain guidance: pass single most informative channel (2D grayscale)
# Try grayscale from CLAHE-enhanced image
gray_clahe = cv2.cvtColor(image_clahe, cv2.COLOR_RGB2GRAY)

# Additional histogram equalization on grayscale to further boost contrast
gray_enhanced = cv2.equalizeHist(gray_clahe)

print(f"Gray enhanced range: [{gray_enhanced.min()}, {gray_enhanced.max()}]")

# Step 3: Run SAM with increased sensitivity parameters
# Using 'default' first as required, but with tuned area parameters
# Expected particle diameter: ~5-15 um -> ~62-186 px at this scale
# Expected area: ~3000 - ~27000 px^2 for typical particles
# Use min_area=300 to capture small/partial particles, max_area=15000*3=45000

# First attempt with default parameters
print("Running SAM with default parameters...")
result = run_sam_analysis(gray_enhanced, params={
    "sam_parameters": "default",
    "min_area": 300,
    "max_area": 45000,
    "pruning_iou_threshold": 0.5
})

print(f"SAM default detected {result['total_count']} particles")

# If we get few detections, try sensitive mode
if result['total_count'] < 15:
    print("Low count with default, trying sensitive mode...")
    result_sensitive = run_sam_analysis(gray_enhanced, params={
        "sam_parameters": "sensitive",
        "min_area": 300,
        "max_area": 45000,
        "pruning_iou_threshold": 0.5
    })
    print(f"SAM sensitive detected {result_sensitive['total_count']} particles")
    if result_sensitive['total_count'] > result['total_count']:
        result = result_sensitive
        print("Using sensitive results")

# Step 4: Build labeled mask from SAM particles
labeled = np.zeros((image_h, image_w), dtype=np.int32)
for i, p in enumerate(result["particles"]):
    mask = np.array(p["mask"], dtype=bool)
    labeled[mask] = i + 1

num_particles = result['total_count']
print(f"Total particles after SAM: {num_particles}")

# Step 5: Extract region properties with relaxed edge exclusion
# Convert original image to grayscale for intensity measurements
gray_original = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
props = measure.regionprops(labeled, intensity_image=gray_original)

# Filter and extract features
particle_data = []
filtered_label = np.zeros_like(labeled)
new_id = 0

for prop in props:
    area_px = prop.area
    
    # Calculate circularity
    perimeter = prop.perimeter
    if perimeter > 0:
        circularity = 4 * np.pi * area_px / (perimeter ** 2)
    else:
        circularity = 0
    
    # Filter out noise/artifacts with very low circularity (< 0.25)
    if circularity < 0.25:
        continue
    
    # Relaxed edge exclusion: keep particles that have at least 30% of expected area
    # Check if touching border
    minr, minc, maxr, maxc = prop.bbox
    touches_border = (minr <= 1 or minc <= 1 or maxr >= image_h - 2 or maxc >= image_w - 2)
    
    if touches_border:
        # Estimate expected full area from the visible portion
        # Use equivalent diameter to estimate full circle area
        eq_diam = prop.equivalent_diameter
        expected_full_area = np.pi * (eq_diam / 2) ** 2
        # For border particles, we're generous - keep if area > 300
        if area_px < 300:
            continue
    
    # Min/max area filters already applied by SAM, but double-check
    if area_px < 300 or area_px > 45000:
        continue
    
    new_id += 1
    filtered_label[labeled == prop.label] = new_id
    
    # Extract features
    eq_diam_um = prop.equivalent_diameter * scale_um_per_px
    area_um2 = area_px * (scale_um_per_px ** 2)
    perimeter_um = perimeter * scale_um_per_px
    major_um = prop.major_axis_length * scale_um_per_px
    minor_um = prop.minor_axis_length * scale_um_per_px
    aspect_ratio = prop.major_axis_length / prop.minor_axis_length if prop.minor_axis_length > 0 else 1.0
    
    particle_data.append({
        "id": new_id,
        "area_um2": float(round(area_um2, 4)),
        "equivalent_diameter_um": float(round(eq_diam_um, 4)),
        "perimeter_um": float(round(perimeter_um, 4)),
        "circularity": float(round(circularity, 4)),
        "solidity": float(round(prop.solidity, 4)),
        "eccentricity": float(round(prop.eccentricity, 4)),
        "major_axis_length_um": float(round(major_um, 4)),
        "minor_axis_length_um": float(round(minor_um, 4)),
        "orientation": float(round(prop.orientation, 4)),
        "centroid_x_px": float(round(prop.centroid[1], 2)),
        "centroid_y_px": float(round(prop.centroid[0], 2)),
        "aspect_ratio": float(round(aspect_ratio, 4))
    })

particle_count = len(particle_data)
print(f"Particles after filtering: {particle_count}")

# Step 6: Compute summary statistics
diameters = [p["equivalent_diameter_um"] for p in particle_data]
circularities = [p["circularity"] for p in particle_data]

if len(diameters) > 0:
    mean_diam = float(np.mean(diameters))
    std_diam = float(np.std(diameters))
    median_diam = float(np.median(diameters))
else:
    mean_diam = std_diam = median_diam = 0.0

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

# Subplot 1: Original image
axes[0, 0].imshow(image)
axes[0, 0].set_title('Original Image')
axes[0, 0].axis('off')

# Subplot 2: Segmentation overlay
overlay = image.copy()
# Generate random colors for each particle
np.random.seed(42)
colors = np.random.randint(50, 255, size=(particle_count + 1, 3), dtype=np.uint8)

# Create colored mask overlay
color_mask = np.zeros_like(image)
for pid in range(1, particle_count + 1):
    mask_region = (filtered_label == pid)
    color_mask[mask_region] = colors[pid]

# Blend
alpha = 0.4
blended = cv2.addWeighted(image, 1.0, color_mask, alpha, 0)

# Draw contours
for pid in range(1, particle_count + 1):
    mask_region = (filtered_label == pid).astype(np.uint8)
    contours, _ = cv2.findContours(mask_region, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    cv2.drawContours(blended, contours, -1, (255, 255, 0), 2)

axes[0, 1].imshow(blended)
axes[0, 1].set_title(f'Segmentation Overlay ({particle_count} particles)')
axes[0, 1].axis('off')

# Subplot 3: CLAHE enhanced image
axes[0, 2].imshow(image_clahe)
axes[0, 2].set_title('CLAHE Enhanced Image')
axes[0, 2].axis('off')

# Subplot 4: Label map
axes[1, 0].imshow(filtered_label, cmap='nipy_spectral', interpolation='nearest')
axes[1, 0].set_title('Label Map')
axes[1, 0].axis('off')

# Subplot 5: Size distribution histogram
if len(diameters) > 0:
    n_bins = min(20, max(5, particle_count // 2))
    axes[1, 1].hist(diameters, bins=n_bins, color='steelblue', edgecolor='black', alpha=0.7)
    axes[1, 1].axvline(mean_diam, color='red', linestyle='--', label=f'Mean={mean_diam:.2f} µm')
    axes[1, 1].axvline(median_diam, color='green', linestyle='--', label=f'Median={median_diam:.2f} µm')
    axes[1, 1].set_xlabel('Equivalent Diameter (µm)')
    axes[1, 1].set_ylabel('Count')
    axes[1, 1].set_title('Size Distribution')
    axes[1, 1].legend()
else:
    axes[1, 1].text(0.5, 0.5, 'No particles detected', ha='center', va='center', transform=axes[1, 1].transAxes)
    axes[1, 1].set_title('Size Distribution')

# Subplot 6: Circularity distribution histogram
if len(circularities) > 0:
    axes[1, 2].hist(circularities, bins=min(20, max(5, particle_count // 2)), color='coral', edgecolor='black', alpha=0.7)
    axes[1, 2].set_xlabel('Circularity')
    axes[1, 2].set_ylabel('Count')
    axes[1, 2].set_title('Circularity Distribution')
    axes[1, 2].axvline(np.mean(circularities), color='red', linestyle='--', label=f'Mean={np.mean(circularities):.3f}')
    axes[1, 2].legend()
else:
    axes[1, 2].text(0.5, 0.5, 'No particles detected', ha='center', va='center', transform=axes[1, 2].transAxes)
    axes[1, 2].set_title('Circularity Distribution')

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

# Step 8: Save label map
np.save('analysis_labels.npy', filtered_label)
print(f"Label map saved: shape={filtered_label.shape}, dtype={filtered_label.dtype}")

# Compute histogram bin data for JSON output
size_hist_counts = []
size_hist_edges = []
circ_hist_counts = []
circ_hist_edges = []

if len(diameters) > 0:
    n_bins = min(20, max(5, particle_count // 2))
    sh_counts, sh_edges = np.histogram(diameters, bins=n_bins)
    size_hist_counts = sh_counts.tolist()
    size_hist_edges = [round(e, 4) for e in sh_edges.tolist()]
    
    ch_counts, ch_edges = np.histogram(circularities, bins=min(20, max(5, particle_count // 2)))
    circ_hist_counts = ch_counts.tolist()
    circ_hist_edges = [round(e, 4) for e in ch_edges.tolist()]

# Step 9: Compile results
results = {
    "analysis_type": "Vaterite particle instance segmentation using CLAHE-enhanced SAM with morphological characterization",
    "extracted_features": {
        "particle_count": particle_count,
        "mean_diameter_um": round(mean_diam, 4),
        "std_diameter_um": round(std_diam, 4),
        "median_diameter_um": round(median_diam, 4),
        "scale_um_per_pixel": round(scale_um_per_px, 6),
        "image_width_px": image_w,
        "image_height_px": image_h,
        "field_of_view_um": 165.0,
        "size_distribution_histogram": {
            "counts": size_hist_counts,
            "bin_edges_um": size_hist_edges
        },
        "circularity_distribution_histogram": {
            "counts": circ_hist_counts,
            "bin_edges": circ_hist_edges
        },
        "per_particle": particle_data[:50]  # Limit to 50 for JSON readability
    },
    "quality_metrics": {
        "mean_circularity": round(float(np.mean(circularities)), 4) if circularities else 0.0,
        "std_circularity": round(float(np.std(circularities)), 4) if circularities else 0.0,
        "mean_solidity": round(float(np.mean([p['solidity'] for p in particle_data])), 4) if particle_data else 0.0,
        "mean_eccentricity": round(float(np.mean([p['eccentricity'] for p in particle_data])), 4) if particle_data else 0.0,
        "mean_aspect_ratio": round(float(np.mean([p['aspect_ratio'] for p in particle_data])), 4) if particle_data else 0.0,
        "min_diameter_um": round(float(min(diameters)), 4) if diameters else 0.0,
        "max_diameter_um": round(float(max(diameters)), 4) if diameters else 0.0,
        "sam_raw_count": result['total_count'],
        "filtered_count": particle_count,
        "preprocessing": "CLAHE (clipLimit=3.0, tileGrid=8x8) + histogram equalization on grayscale"
    },
    "summary": f"Detected {particle_count} vaterite particles (from {result['total_count']} SAM detections) with mean equivalent diameter {mean_diam:.2f} +/- {std_diam:.2f} um (median {median_diam:.2f} um). CLAHE preprocessing enhanced visibility of dark/low-contrast particles. Circularity filter (>=0.25) removed noise artifacts while retaining elongated/tilted particles. Scale: {scale_um_per_px:.4f} um/pixel over 165 um field of view.",
    "saved_arrays": {
        "analysis_labels.npy": {
            "description": f"Integer label map, {particle_count} particles labeled 1-{particle_count}, background=0",
            "shape": list(filtered_label.shape),
            "dtype": str(filtered_label.dtype)
        }
    }
}

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