import numpy as np
import cv2
import json
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from scipy import ndimage
from scipy.interpolate import griddata
from skimage import measure, morphology, feature, segmentation, exposure, filters
from sklearn.mixture import GaussianMixture
import warnings
warnings.filterwarnings('ignore')

# =====================================================================
# LOAD IMAGE
# =====================================================================
path = '/Users/maxim.ziatdinov/Code/SciLink/test_image_annealing_data/output_escalation/image_0000/temp_image_0.npy'
image_raw = np.load(path)
print(f"Image shape: {image_raw.shape}, dtype: {image_raw.dtype}")
print(f"Intensity range: [{image_raw.min()}, {image_raw.max()}]")

if image_raw.ndim == 3:
    image_2d = image_raw[:, :, 0]
else:
    image_2d = image_raw.copy()

img_min, img_max = image_2d.min(), image_2d.max()
image_norm = ((image_2d - img_min) / (img_max - img_min) * 255).astype(np.uint8)

# =====================================================================
# STEP 1: Light Gaussian blur
# =====================================================================
blur_sigma = 0.9
image_blurred = cv2.GaussianBlur(image_2d, (0, 0), blur_sigma)
image_blurred_uint8 = ((image_blurred - image_blurred.min()) / (image_blurred.max() - image_blurred.min()) * 255).astype(np.uint8)

# =====================================================================
# STEP 2: LOCAL BACKGROUND ESTIMATION
# =====================================================================
def sigma_clipped_stats(data, sigma=3.0, iters=3):
    d = data.flatten().astype(float)
    for _ in range(iters):
        if len(d) == 0:
            return 0.0, 1.0
        m, s = np.mean(d), np.std(d)
        if s == 0:
            return m, 0.001
        mask = np.abs(d - m) < sigma * s
        d = d[mask]
    return np.mean(d), max(np.std(d), 0.001)

tile_size = 64
overlap = 32
H, W = image_blurred.shape

tile_centers_y = []
tile_centers_x = []
tile_means = []
tile_stds = []

for y_start in range(0, H, overlap):
    for x_start in range(0, W, overlap):
        y_end = min(y_start + tile_size, H)
        x_end = min(x_start + tile_size, W)
        tile = image_blurred[y_start:y_end, x_start:x_end]
        m, s = sigma_clipped_stats(tile)
        tile_centers_y.append((y_start + y_end) / 2.0)
        tile_centers_x.append((x_start + x_end) / 2.0)
        tile_means.append(m)
        tile_stds.append(s)

tile_centers_y = np.array(tile_centers_y)
tile_centers_x = np.array(tile_centers_x)
tile_means = np.array(tile_means)
tile_stds = np.array(tile_stds)

grid_y, grid_x = np.mgrid[0:H, 0:W]
points = np.column_stack([tile_centers_y, tile_centers_x])

local_bg_mean_map = griddata(points, tile_means, (grid_y, grid_x), method='linear', fill_value=np.median(tile_means))
local_bg_std_map = griddata(points, tile_stds, (grid_y, grid_x), method='linear', fill_value=np.median(tile_stds))
local_bg_std_map = np.maximum(local_bg_std_map, 0.001)

global_bg_mean, global_bg_std = sigma_clipped_stats(image_blurred)
print(f"Global background: mean={global_bg_mean:.2f}, std={global_bg_std:.2f}")

# =====================================================================
# STEP 3: CLAHE preprocessing
# =====================================================================
clahe = cv2.createCLAHE(clipLimit=2.5, tileGridSize=(8, 8))
image_clahe = clahe.apply(image_blurred_uint8)

# =====================================================================
# STEP 4: Border exclusion mask
# =====================================================================
border_px = 2
border_mask = np.ones((H, W), dtype=bool)
border_mask[:border_px, :] = False
border_mask[-border_px:, :] = False
border_mask[:, :border_px] = False
border_mask[:, -border_px:] = False

# =====================================================================
# STEP 5: Primary detection via SAM
# =====================================================================
from scilink.tools.sam import run_sam_analysis

try:
    sam_result = run_sam_analysis(image_2d, params={
        "sam_parameters": "default",
        "min_area": 30,
        "max_area": 8000,
        "pruning_iou_threshold": 0.5
    })
    sam_particles = sam_result.get("particles", [])
    print(f"SAM detected {len(sam_particles)} particles")
except Exception as e:
    print(f"SAM failed: {e}")
    sam_particles = []

samlabeled = np.zeros((H, W), dtype=np.int32)
for i, p in enumerate(sam_particles):
    mask = np.array(p["mask"], dtype=bool)
    samlabeled[mask] = i + 1

# =====================================================================
# STEP 6: Local intensity gating on SAM detections
# =====================================================================
def passes_intensity_gate(mask_bool, image, local_mean_map, local_std_map,
                          global_mean, global_std, snr_thresh=2.0,
                          global_floor_mult=1.5, regional_exception=False):
    ys, xs = np.where(mask_bool)
    if len(ys) == 0:
        return False, 0.0, 0.0
    cy, cx = int(np.mean(ys)), int(np.mean(xs))
    cy = min(max(cy, 0), H-1)
    cx = min(max(cx, 0), W-1)
    mean_int = np.mean(image[mask_bool])
    local_mean = local_mean_map[cy, cx]
    local_std = local_std_map[cy, cx]
    local_snr = (mean_int - local_mean) / max(local_std, 0.001)

    snr_threshold = 1.5 if regional_exception else snr_thresh
    floor_mult = 1.0 if regional_exception else global_floor_mult

    if mean_int < (local_mean + snr_threshold * local_std):
        return False, mean_int, local_snr
    if mean_int < (global_mean + floor_mult * global_std):
        return False, mean_int, local_snr
    return True, mean_int, local_snr

def is_in_left_center_roi(cy, cx):
    return 100 <= cy <= 190 and 50 <= cx <= 160

# We use a dict-based metadata system: each object is a dict with 'mask' and 'meta'
# meta tracks: detection_method, detection_pass, split_strategy_used, parent_label,
# gaussian_fit_separation, gaussian_fit_component_weights, concavity_deficit_fraction,
# erosion_iterations_required

def make_obj(mask, detection_method='SAM', detection_pass='primary', **extra_meta):
    meta = {
        'detection_method': detection_method,
        'detection_pass': detection_pass,
        'split_strategy_used': 'none',
        'parent_label': None,
        'gaussian_fit_separation': None,
        'gaussian_fit_component_weights': None,
        'concavity_deficit_fraction': None,
        'erosion_iterations_required': None,
    }
    meta.update(extra_meta)
    return {'mask': mask, 'meta': meta}

valid_sam_objs = []
for i in range(1, samlabeled.max() + 1):
    mask = samlabeled == i
    if not np.any(mask):
        continue
    ys, xs = np.where(mask)
    if ys.min() < border_px or ys.max() >= H - border_px or xs.min() < border_px or xs.max() >= W - border_px:
        continue
    cy, cx_val = int(np.mean(ys)), int(np.mean(xs))
    regional = is_in_left_center_roi(cy, cx_val)
    passed, mi, snr = passes_intensity_gate(mask, image_blurred, local_bg_mean_map, local_bg_std_map,
                                             global_bg_mean, global_bg_std,
                                             snr_thresh=2.0, regional_exception=regional)
    if passed:
        valid_sam_objs.append(make_obj(mask, 'SAM', 'primary'))

print(f"SAM after intensity gating: {len(valid_sam_objs)} objects")

# =====================================================================
# STEP 7: Supplementary adaptive thresholding
# =====================================================================
def adaptive_threshold_detect(img_uint8, block_size, c_offset, min_area=30, roi=None):
    if roi is not None:
        y1, y2, x1, x2 = roi
        sub = img_uint8[y1:y2, x1:x2]
    else:
        sub = img_uint8

    binary = cv2.adaptiveThreshold(sub, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
                                    cv2.THRESH_BINARY, block_size, c_offset)
    binary = morphology.remove_small_objects(binary.astype(bool), min_size=min_area)
    binary = morphology.remove_small_holes(binary, area_threshold=20)

    labeled_sub = measure.label(binary)
    masks = []
    for region in measure.regionprops(labeled_sub):
        if region.area < min_area:
            continue
        m = labeled_sub == region.label
        if roi is not None:
            full_mask = np.zeros((H, W), dtype=bool)
            full_mask[y1:y2, x1:x2] = m
            masks.append(full_mask)
        else:
            masks.append(m)
    return masks

suppl_masks_1 = adaptive_threshold_detect(image_clahe, 17, -3, min_area=25)
print(f"Supplementary pass 1: {len(suppl_masks_1)} candidates")

suppl_masks_2 = adaptive_threshold_detect(image_clahe, 15, -2, min_area=25)
print(f"Supplementary pass 2: {len(suppl_masks_2)} candidates")

region_rois = [
    (80, 200, 30, 150),
    (180, 250, 0, min(256, W)),
    (30, 80, 80, 200),
]
suppl_masks_3 = []
for roi in region_rois:
    roi_clamped = (min(roi[0], H), min(roi[1], H), min(roi[2], W), min(roi[3], W))
    masks = adaptive_threshold_detect(image_clahe, 13, -1, min_area=20, roi=roi_clamped)
    suppl_masks_3.extend(masks)
print(f"Supplementary pass 3 (targeted): {len(suppl_masks_3)} candidates")

lc_roi = (min(100, H), min(190, H), min(50, W), min(160, W))
suppl_masks_4 = adaptive_threshold_detect(image_clahe, 11, 0, min_area=15, roi=lc_roi)
print(f"Supplementary pass 4 (ultra-sensitive left-center): {len(suppl_masks_4)} candidates")

def gate_supplementary_masks(masks, snr_thresh=2.0, global_floor_mult=1.5, pass_name='suppl'):
    valid = []
    for mask in masks:
        if not np.any(mask):
            continue
        ys, xs = np.where(mask)
        if ys.min() < border_px or ys.max() >= H - border_px or xs.min() < border_px or xs.max() >= W - border_px:
            continue
        cy, cx_val = int(np.mean(ys)), int(np.mean(xs))
        regional = is_in_left_center_roi(cy, cx_val)
        actual_snr = 1.5 if regional else snr_thresh
        actual_floor = 1.0 if regional else global_floor_mult
        passed, mi, snr = passes_intensity_gate(mask, image_blurred, local_bg_mean_map, local_bg_std_map,
                                                 global_bg_mean, global_bg_std,
                                                 snr_thresh=actual_snr, global_floor_mult=actual_floor)
        if passed:
            if regional and snr < 2.0:
                props = measure.regionprops(mask.astype(int))
                if props:
                    area = props[0].area
                    perim = props[0].perimeter
                    circ = 4 * np.pi * area / (perim**2 + 1e-10) if perim > 0 else 0
                    if circ < 0.5:
                        continue
            valid.append(make_obj(mask, 'adaptive_threshold', pass_name))
    return valid

valid_suppl_1 = gate_supplementary_masks(suppl_masks_1, pass_name='pass1')
valid_suppl_2 = gate_supplementary_masks(suppl_masks_2, pass_name='pass2')
valid_suppl_3 = gate_supplementary_masks(suppl_masks_3, pass_name='pass3_targeted')
valid_suppl_4 = gate_supplementary_masks(suppl_masks_4, snr_thresh=1.5, global_floor_mult=1.0, pass_name='pass4_ultrasensitive')

print(f"Valid supplementary: pass1={len(valid_suppl_1)}, pass2={len(valid_suppl_2)}, pass3={len(valid_suppl_3)}, pass4={len(valid_suppl_4)}")

# =====================================================================
# Merge all objects with deduplication
# =====================================================================
def get_centroid(mask):
    ys, xs = np.where(mask)
    return np.mean(ys), np.mean(xs)

def merge_objs_dedup(existing, new_objs, dist_thresh=12):
    added = 0
    for nobj in new_objs:
        nm = nobj['mask']
        cy_n, cx_n = get_centroid(nm)
        is_dup = False
        for eidx, eobj in enumerate(existing):
            em = eobj['mask']
            cy_e, cx_e = get_centroid(em)
            dist = np.sqrt((cy_n - cy_e)**2 + (cx_n - cx_e)**2)
            iou = np.sum(nm & em) / max(np.sum(nm | em), 1)
            if dist < dist_thresh or iou > 0.3:
                is_dup = True
                if np.sum(nm) > np.sum(em):
                    existing[eidx] = nobj
                break
        if not is_dup:
            existing.append(nobj)
            added += 1
    return added

all_objs = list(valid_sam_objs)
added_1 = merge_objs_dedup(all_objs, valid_suppl_1)
added_2 = merge_objs_dedup(all_objs, valid_suppl_2)
added_3 = merge_objs_dedup(all_objs, valid_suppl_3)
added_4 = merge_objs_dedup(all_objs, valid_suppl_4)
print(f"After merging: {len(all_objs)} objects (added: {added_1}+{added_2}+{added_3}+{added_4} from supplementary)")

# =====================================================================
# STEP 8: MULTI-STRATEGY SPLITTING FOR MERGED OBJECTS
# =====================================================================
split_log = []
split_stats = {'attempted': 0, 'succeeded': 0, 'failed': 0,
               'by_strategy': {s: 0 for s in ['Gaussian', 'Concavity', 'Erosion', 'IntensityPeak', 'DistanceTransform']}}

# Target labels tracking: labels 10 (~x:127,y:212) and 15 (~x:213,y:56)
# We identify them by centroid proximity
TARGET_CENTROIDS = [
    {'name': 'target_label10', 'approx_y': 212, 'approx_x': 127, 'match_radius': 25, 'matched_idx': None, 'split_success': False, 'split_strategy': None, 'children': []},
    {'name': 'target_label15', 'approx_y': 56, 'approx_x': 213, 'match_radius': 25, 'matched_idx': None, 'split_success': False, 'split_strategy': None, 'children': []},
]

def compute_circularity(region):
    if region.perimeter == 0:
        return 1.0
    return 4 * np.pi * region.area / (region.perimeter ** 2)

def try_gaussian_split(mask, intensity_img, local_mean_map_arg, min_sep_frac=0.5, min_pix_frac=0.15, cov_type='full'):
    ys, xs = np.where(mask)
    if len(ys) < 20:
        return None, "Too few pixels", {}

    intensities = intensity_img[ys, xs].astype(float)
    local_bg = local_mean_map_arg[ys, xs]
    weights = np.maximum(intensities - local_bg, 0.1)

    coords = np.column_stack([xs.astype(float), ys.astype(float)])

    try:
        gmm = GaussianMixture(n_components=2, covariance_type=cov_type,
                              random_state=42, max_iter=200, n_init=5)
        gmm.fit(coords, sample_weight=weights)
        labels = gmm.predict(coords)
    except Exception as e:
        return None, f"GMM fit failed: {e}", {}

    centers = gmm.means_
    sep = np.sqrt(np.sum((centers[0] - centers[1])**2))
    eq_diam = np.sqrt(4 * len(ys) / np.pi)

    if sep < min_sep_frac * eq_diam:
        return None, f"Centers too close: sep={sep:.1f}, need>{min_sep_frac*eq_diam:.1f}", {}

    n0 = np.sum(labels == 0)
    n1 = np.sum(labels == 1)
    total = len(labels)
    if min(n0, n1) / total < min_pix_frac:
        return None, f"Unbalanced: {n0}/{n1} pixels, need>{min_pix_frac*100:.0f}%", {}

    sub_masks = []
    for lbl in [0, 1]:
        sub = np.zeros_like(mask)
        sub[ys[labels == lbl], xs[labels == lbl]] = True
        sub_masks.append(sub)

    extra = {
        'gaussian_fit_separation': float(sep),
        'gaussian_fit_component_weights': [float(n0/total), float(n1/total)],
    }
    return sub_masks, f"Success: sep={sep:.1f}px, sizes={n0}/{n1}", extra

def try_concavity_split(mask, min_deficit_frac=0.2, min_area_frac=0.15):
    from skimage.morphology import convex_hull_image
    try:
        hull = convex_hull_image(mask)
    except Exception:
        return None, "Convex hull failed", {}

    deficit = hull & ~mask
    deficit_area = np.sum(deficit)
    mask_area = np.sum(mask)
    eq_diam = np.sqrt(4 * mask_area / np.pi)

    deficit_frac = deficit_area / max(mask_area, 1)

    if deficit_area < 5:
        return None, "No significant concavity", {}

    labeled_deficit = measure.label(deficit)
    deficit_props = measure.regionprops(labeled_deficit)

    large_deficits = [d for d in deficit_props if d.equivalent_diameter > min_deficit_frac * eq_diam]
    if not large_deficits:
        return None, f"No deficit large enough (need>{min_deficit_frac*eq_diam:.1f}px diameter)", {}

    props = measure.regionprops(mask.astype(int))
    if not props:
        return None, "No region props", {}

    region = props[0]
    cy_r, cx_r = region.centroid

    largest_deficit = max(large_deficits, key=lambda d: d.area)
    dy, dx = largest_deficit.centroid

    vec = np.array([dx - cx_r, dy - cy_r])
    vec_len = np.linalg.norm(vec)
    if vec_len < 1:
        return None, "Deficit at center", {}

    perp = np.array([-vec[1], vec[0]]) / vec_len

    cut_mask = mask.copy()
    for t in np.linspace(-eq_diam, eq_diam, int(2*eq_diam)+1):
        py_c, px_c = int(dy + t * perp[1]), int(dx + t * perp[0])
        if 0 <= py_c < H and 0 <= px_c < W:
            cut_mask[py_c, px_c] = False

    labeled_cut = measure.label(cut_mask)
    cut_props = measure.regionprops(labeled_cut)

    valid = [p for p in cut_props if p.area > min_area_frac * mask_area]
    if len(valid) < 2:
        return None, f"Cut produced {len(valid)} valid components", {}

    sub_masks = []
    for p in valid[:2]:
        sub = labeled_cut == p.label
        sub_masks.append(sub)

    neck_width = 0
    extra = {
        'concavity_deficit_fraction': float(deficit_frac),
    }
    return sub_masks, f"Success: neck at ({dx:.0f},{dy:.0f}), {len(valid)} components", extra

def try_erosion_split(mask, max_iter_frac=3, min_area_frac=0.10):
    eq_diam = np.sqrt(4 * np.sum(mask) / np.pi)
    max_iters = max(int(eq_diam / max_iter_frac), 3)

    selem = morphology.disk(1)
    eroded = mask.copy()

    for i in range(1, max_iters + 1):
        eroded = morphology.binary_erosion(eroded, selem)
        if np.sum(eroded) < 10:
            return None, f"Eroded to nothing at iter {i}", {}

        labeled_eroded = measure.label(eroded)
        n_components = labeled_eroded.max()

        if n_components >= 2:
            markers = measure.label(eroded)
            distance = ndimage.distance_transform_edt(mask)
            grown = segmentation.watershed(-distance, markers, mask=mask)

            grown_props = measure.regionprops(grown)
            valid = [p for p in grown_props if p.area > min_area_frac * np.sum(mask)]

            if len(valid) >= 2:
                sub_masks = []
                for p in valid[:3]:
                    sub = grown == p.label
                    sub_masks.append(sub)
                extra = {
                    'erosion_iterations_required': int(i),
                }
                return sub_masks, f"Success: {i} erosion iters, {len(valid)} components", extra

    return None, f"No split after {max_iters} iterations", {}

def try_intensity_peak_watershed(mask, intensity_img, local_mean_map_arg, local_std_map_arg, min_distance=6):
    masked_intensity = intensity_img.copy().astype(float)
    masked_intensity[~mask] = 0

    ys, xs = np.where(mask)
    cy, cx = int(np.mean(ys)), int(np.mean(xs))
    cy = min(max(cy, 0), H-1)
    cx = min(max(cx, 0), W-1)
    local_thresh = local_mean_map_arg[cy, cx] + 2 * local_std_map_arg[cy, cx]

    coords = feature.peak_local_max(masked_intensity, min_distance=min_distance,
                                     threshold_abs=local_thresh, exclude_border=False)

    if len(coords) < 2:
        return None, f"Only {len(coords)} intensity peaks found", {}

    markers = np.zeros_like(mask, dtype=int)
    for idx_p, (y_p, x_p) in enumerate(coords[:5]):
        markers[y_p, x_p] = idx_p + 1

    ws = segmentation.watershed(-masked_intensity, markers, mask=mask)
    ws_props = measure.regionprops(ws)
    valid = [p for p in ws_props if p.area > 0.1 * np.sum(mask)]

    if len(valid) >= 2:
        sub_masks = []
        for p in valid[:3]:
            sub = ws == p.label
            sub_masks.append(sub)
        return sub_masks, f"Success: {len(coords)} peaks, {len(valid)} components", {}

    return None, f"{len(coords)} peaks but only {len(valid)} valid components", {}

def try_distance_watershed(mask, min_distance=5):
    distance = ndimage.distance_transform_edt(mask)
    coords = feature.peak_local_max(distance, min_distance=min_distance, exclude_border=False)

    if len(coords) < 2:
        return None, f"Only {len(coords)} distance peaks", {}

    markers = np.zeros_like(mask, dtype=int)
    for idx_p, (y_p, x_p) in enumerate(coords[:5]):
        markers[y_p, x_p] = idx_p + 1

    ws = segmentation.watershed(-distance, markers, mask=mask)
    ws_props = measure.regionprops(ws)
    valid = [p for p in ws_props if p.area > 0.1 * np.sum(mask)]

    if len(valid) >= 2:
        sub_masks = []
        for p in valid[:3]:
            sub = ws == p.label
            sub_masks.append(sub)
        return sub_masks, f"Success: {len(coords)} peaks, {len(valid)} components", {}

    return None, f"{len(coords)} peaks but only {len(valid)} valid components", {}

def split_object(mask, intensity_img, local_mean_map_arg, local_std_map_arg, obj_idx, relaxed=False):
    strategies = [
        ('Gaussian', lambda: try_gaussian_split(mask, intensity_img, local_mean_map_arg,
                                                  min_sep_frac=0.3 if relaxed else 0.5,
                                                  min_pix_frac=0.12 if relaxed else 0.15,
                                                  cov_type='tied' if relaxed else 'full')),
        ('Concavity', lambda: try_concavity_split(mask, min_deficit_frac=0.1 if relaxed else 0.2)),
        ('Erosion', lambda: try_erosion_split(mask, max_iter_frac=2 if relaxed else 3)),
        ('IntensityPeak', lambda: try_intensity_peak_watershed(mask, intensity_img, local_mean_map_arg, local_std_map_arg)),
        ('DistanceTransform', lambda: try_distance_watershed(mask)),
    ]

    reasons = {}
    for name, func in strategies:
        try:
            result, reason, extra = func()
            reasons[name] = reason
            if result is not None:
                return result, name, reasons, extra
        except Exception as e:
            reasons[name] = f"Exception: {e}"

    return None, None, reasons, {}

# Identify target objects by centroid proximity before splitting
for tidx, target in enumerate(TARGET_CENTROIDS):
    best_dist = float('inf')
    best_oidx = None
    for oidx, obj in enumerate(all_objs):
        cy, cx = get_centroid(obj['mask'])
        dist = np.sqrt((cy - target['approx_y'])**2 + (cx - target['approx_x'])**2)
        if dist < best_dist:
            best_dist = dist
            best_oidx = oidx
    if best_dist < target['match_radius'] and best_oidx is not None:
        TARGET_CENTROIDS[tidx]['matched_idx'] = best_oidx
        cy, cx = get_centroid(all_objs[best_oidx]['mask'])
        print(f"TARGET TRACKING: {target['name']} matched to object index {best_oidx} at centroid ({cy:.1f},{cx:.1f}), distance={best_dist:.1f}px")
    else:
        print(f"TARGET TRACKING: {target['name']} NOT MATCHED (best dist={best_dist:.1f}px, idx={best_oidx})")

# Splitting loop
split_results = []  # list of (original_idx, new_objs_list)
target_split_attempted = {t['name']: False for t in TARGET_CENTROIDS}

for idx, obj in enumerate(all_objs):
    mask = obj['mask']
    props = measure.regionprops(mask.astype(int), intensity_image=image_blurred)
    if not props:
        split_results.append((idx, [obj]))
        continue
    region = props[0]
    circ = compute_circularity(region)
    ecc = region.eccentricity

    # Check if this is a target object
    is_target = False
    target_name = None
    for target in TARGET_CENTROIDS:
        if target['matched_idx'] == idx:
            is_target = True
            target_name = target['name']
            break

    needs_split = circ < 0.75 or ecc > 0.65 or is_target

    if needs_split and region.area > 80:
        split_stats['attempted'] += 1
        if is_target:
            target_split_attempted[target_name] = True
        ys, xs = np.where(mask)
        cy, cx_val = np.mean(ys), np.mean(xs)

        result, strategy, reasons, extra = split_object(mask, image_blurred, local_bg_mean_map, local_bg_std_map, idx)

        if result is not None:
            split_stats['succeeded'] += 1
            split_stats['by_strategy'][strategy] += 1
            log_entry = f"Label {idx} at ({cx_val:.0f},{cy:.0f}) circ={circ:.3f} ecc={ecc:.3f}: SPLIT by {strategy}"
            if is_target:
                log_entry += f" [TARGET: {target_name}]"
            split_log.append(log_entry)
            print(log_entry)

            new_objs = []
            for si, sub_mask in enumerate(result):
                child_meta = dict(obj['meta'])
                child_meta['split_strategy_used'] = strategy
                child_meta['parent_label'] = idx
                child_meta.update({k: v for k, v in extra.items() if v is not None})
                new_objs.append({'mask': sub_mask, 'meta': child_meta})
            split_results.append((idx, new_objs))

            if is_target:
                for target in TARGET_CENTROIDS:
                    if target['matched_idx'] == idx:
                        target['split_success'] = True
                        target['split_strategy'] = strategy
                        target['children'] = [get_centroid(sm) for sm in result]
        else:
            # Try relaxed
            result2, strategy2, reasons2, extra2 = split_object(mask, image_blurred, local_bg_mean_map, local_bg_std_map, idx, relaxed=True)
            if result2 is not None:
                split_stats['succeeded'] += 1
                split_stats['by_strategy'][strategy2] += 1
                log_entry = f"Label {idx} at ({cx_val:.0f},{cy:.0f}): SPLIT by {strategy2} (RELAXED)"
                if is_target:
                    log_entry += f" [TARGET: {target_name}]"
                split_log.append(log_entry)
                print(log_entry)

                new_objs = []
                for si, sub_mask in enumerate(result2):
                    child_meta = dict(obj['meta'])
                    child_meta['split_strategy_used'] = strategy2
                    child_meta['parent_label'] = idx
                    child_meta.update({k: v for k, v in extra2.items() if v is not None})
                    new_objs.append({'mask': sub_mask, 'meta': child_meta})
                split_results.append((idx, new_objs))

                if is_target:
                    for target in TARGET_CENTROIDS:
                        if target['matched_idx'] == idx:
                            target['split_success'] = True
                            target['split_strategy'] = strategy2
                            target['children'] = [get_centroid(sm) for sm in result2]
            else:
                split_stats['failed'] = split_stats.get('failed', 0) + 1
                log_entry = f"Label {idx} at ({cx_val:.0f},{cy:.0f}): ALL STRATEGIES FAILED. Reasons: {reasons}"
                if is_target:
                    log_entry += f" [TARGET: {target_name}]"
                split_log.append(log_entry)
                print(log_entry)
                split_results.append((idx, [obj]))
    else:
        split_results.append((idx, [obj]))

# CRITICAL VERIFICATION: re-attempt for targets that were not split
for target in TARGET_CENTROIDS:
    if target['matched_idx'] is not None and not target['split_success']:
        print(f"TARGET VERIFICATION: {target['name']} was NOT split. Re-attempting with extra-relaxed parameters...")
        oidx = target['matched_idx']
        mask = all_objs[oidx]['mask']

        # Extra relaxed: even more lenient parameters
        strategies_extra = [
            ('Gaussian', lambda: try_gaussian_split(mask, image_blurred, local_bg_mean_map,
                                                      min_sep_frac=0.2, min_pix_frac=0.10, cov_type='tied')),
            ('Concavity', lambda: try_concavity_split(mask, min_deficit_frac=0.1, min_area_frac=0.10)),
            ('Erosion', lambda: try_erosion_split(mask, max_iter_frac=2, min_area_frac=0.08)),
            ('IntensityPeak', lambda: try_intensity_peak_watershed(mask, image_blurred, local_bg_mean_map, local_bg_std_map, min_distance=4)),
            ('DistanceTransform', lambda: try_distance_watershed(mask, min_distance=4)),
        ]

        for name, func in strategies_extra:
            try:
                result, reason, extra = func()
                if result is not None:
                    split_stats['succeeded'] += 1
                    split_stats['by_strategy'][name] += 1
                    log_entry = f"TARGET RE-ATTEMPT: {target['name']} (idx {oidx}) SPLIT by {name} (extra-relaxed)"
                    split_log.append(log_entry)
                    print(log_entry)

                    # Replace in split_results
                    new_objs = []
                    for si, sub_mask in enumerate(result):
                        child_meta = dict(all_objs[oidx]['meta'])
                        child_meta['split_strategy_used'] = name
                        child_meta['parent_label'] = oidx
                        child_meta.update({k: v for k, v in extra.items() if v is not None})
                        new_objs.append({'mask': sub_mask, 'meta': child_meta})

                    # Find and replace in split_results
                    for sr_idx, (orig_idx, _) in enumerate(split_results):
                        if orig_idx == oidx:
                            split_results[sr_idx] = (oidx, new_objs)
                            break

                    target['split_success'] = True
                    target['split_strategy'] = name
                    target['children'] = [get_centroid(sm) for sm in result]
                    break
            except Exception as e:
                print(f"  {name} failed: {e}")

        if not target['split_success']:
            print(f"TARGET VERIFICATION: {target['name']} COULD NOT BE SPLIT after all re-attempts.")

# Flatten split_results into a single list of objects
post_split_objs = []
for orig_idx, new_objs in split_results:
    post_split_objs.extend(new_objs)

print(f"After splitting: {len(post_split_objs)} objects")
print(f"Split stats: {split_stats}")

# Re-apply local intensity gate to all split products
gated_objs = []
for obj in post_split_objs:
    mask = obj['mask']
    ys, xs = np.where(mask)
    if len(ys) == 0:
        continue
    cy, cx_val = int(np.mean(ys)), int(np.mean(xs))
    regional = is_in_left_center_roi(cy, cx_val)
    passed, mi, snr = passes_intensity_gate(mask, image_blurred, local_bg_mean_map, local_bg_std_map,
                                             global_bg_mean, global_bg_std,
                                             snr_thresh=1.5 if regional else 2.0,
                                             global_floor_mult=1.0 if regional else 1.5)
    if passed:
        gated_objs.append(obj)

print(f"After re-gating split products: {len(gated_objs)} objects")

# =====================================================================
# STEP 9: Combined circularity-intensity filter
# =====================================================================
filtered_objs = []
for obj in gated_objs:
    mask = obj['mask']
    props = measure.regionprops(mask.astype(int), intensity_image=image_blurred)
    if not props:
        continue
    region = props[0]
    circ = compute_circularity(region)
    mi = region.mean_intensity

    if circ < 0.6 and mi < 100:
        continue
    filtered_objs.append(obj)

print(f"After circularity-intensity filter: {len(filtered_objs)} objects")

# =====================================================================
# STEP 11: Fragment filtering
# =====================================================================
areas_list = [np.sum(obj['mask']) for obj in filtered_objs]
if areas_list:
    median_area = np.median(areas_list)
    min_area_thresh = median_area / 4
    expected_single_area = median_area

    defrag_objs = []
    for obj in filtered_objs:
        mask = obj['mask']
        a = np.sum(mask)
        ys, xs = np.where(mask)
        cy, cx_val = int(np.mean(ys)), int(np.mean(xs))
        cy = min(max(cy, 0), H-1)
        cx_val = min(max(cx_val, 0), W-1)
        mi = np.mean(image_blurred[mask])
        local_mean = local_bg_mean_map[cy, cx_val]
        local_std = local_bg_std_map[cy, cx_val]
        snr = (mi - local_mean) / max(local_std, 0.001)

        if a < min_area_thresh:
            if a > 0.5 * expected_single_area:
                defrag_objs.append(obj)
            elif snr > 3.0:
                defrag_objs.append(obj)
            else:
                continue
        else:
            defrag_objs.append(obj)
    filtered_objs = defrag_objs
    print(f"After fragment filter: {len(filtered_objs)} objects")

# =====================================================================
# STEP 12: Secondary intensity filter
# =====================================================================
final_filtered_objs = []

for obj in filtered_objs:
    mask = obj['mask']
    ys, xs = np.where(mask)
    if len(ys) == 0:
        continue
    cy, cx_val = int(np.mean(ys)), int(np.mean(xs))
    cy = min(max(cy, 0), H-1)
    cx_val = min(max(cx_val, 0), W-1)
    mi = np.mean(image_blurred[mask])
    local_mean = local_bg_mean_map[cy, cx_val]
    local_std = local_bg_std_map[cy, cx_val]
    snr = (mi - local_mean) / max(local_std, 0.001)
    regional = is_in_left_center_roi(cy, cx_val)

    if regional:
        if snr >= 1.5:
            obj['meta']['low_confidence'] = snr < 2.0
            obj['meta']['regional_exception'] = True
            final_filtered_objs.append(obj)
        else:
            continue
    else:
        if mi < (local_mean + 1.5 * local_std) and snr < 2.0:
            continue
        obj['meta']['low_confidence'] = (2.0 <= snr < 2.5)
        obj['meta']['regional_exception'] = False
        final_filtered_objs.append(obj)

print(f"After secondary intensity filter: {len(final_filtered_objs)} objects")

# =====================================================================
# STEP 13: Final edge exclusion
# =====================================================================
edge_filtered_objs = []
for obj in final_filtered_objs:
    mask = obj['mask']
    ys, xs = np.where(mask)
    if ys.min() < border_px or ys.max() >= H - border_px or xs.min() < border_px or xs.max() >= W - border_px:
        continue
    edge_filtered_objs.append(obj)

final_filtered_objs = edge_filtered_objs
print(f"After edge filter: {len(final_filtered_objs)} objects")

# =====================================================================
# STEP 14: Verification pass - local peak detection
# =====================================================================
local_max_coords = feature.peak_local_max(image_blurred, min_distance=8,
                                           threshold_abs=global_bg_mean + 2*global_bg_std,
                                           exclude_border=border_px)
print(f"Local maxima found: {len(local_max_coords)}")

new_peaks = []
for (py, px) in local_max_coords:
    already_covered = False
    for obj in final_filtered_objs:
        if obj['mask'][py, px]:
            already_covered = True
            break
    if not already_covered:
        local_mean = local_bg_mean_map[py, px]
        local_std = local_bg_std_map[py, px]
        regional = is_in_left_center_roi(py, px)
        snr_thresh = 1.5 if regional else 2.5

        y_lo, y_hi = max(0, py-5), min(H, py+6)
        x_lo, x_hi = max(0, px-5), min(W, px+6)
        patch = image_blurred[y_lo:y_hi, x_lo:x_hi]
        mi = np.mean(patch)
        snr = (mi - local_mean) / max(local_std, 0.001)

        int_thresh_mult = 1.5 if regional else 2.0
        if snr > snr_thresh and mi > (local_mean + int_thresh_mult * local_std):
            new_peaks.append((py, px, snr, mi))

print(f"New peaks to add: {len(new_peaks)}")

for py, px, snr, mi in new_peaks:
    mask = np.zeros((H, W), dtype=bool)
    yy, xx = np.ogrid[0:H, 0:W]
    r2 = (yy - py)**2 + (xx - px)**2
    mask = r2 <= 36
    mask = mask & border_mask

    too_close = False
    for eobj in final_filtered_objs:
        eys, exs = np.where(eobj['mask'])
        if len(eys) > 0:
            ecy, ecx = np.mean(eys), np.mean(exs)
            if np.sqrt((py - ecy)**2 + (px - ecx)**2) < 10:
                too_close = True
                break

    if not too_close and np.sum(mask) > 15:
        # Check circularity-intensity filter
        props_check = measure.regionprops(mask.astype(int), intensity_image=image_blurred)
        if props_check:
            circ_check = compute_circularity(props_check[0])
            mi_check = props_check[0].mean_intensity
            if circ_check < 0.6 and mi_check < 100:
                continue

        regional = is_in_left_center_roi(py, px)
        new_obj = make_obj(mask, 'peak_detection', 'verification')
        new_obj['meta']['low_confidence'] = snr < 2.5
        new_obj['meta']['regional_exception'] = regional
        final_filtered_objs.append(new_obj)

print(f"Final object count: {len(final_filtered_objs)}")

# =====================================================================
# Build final label map and extract properties
# =====================================================================
label_map = np.zeros((H, W), dtype=np.int32)
for i, obj in enumerate(final_filtered_objs):
    label_map[obj['mask']] = i + 1

props = measure.regionprops(label_map, intensity_image=image_blurred)

results_objects = []
for i, region in enumerate(props):
    cy, cx = region.centroid
    cy_int, cx_int = int(cy), int(cx)
    cy_int = min(max(cy_int, 0), H-1)
    cx_int = min(max(cx_int, 0), W-1)

    local_mean = local_bg_mean_map[cy_int, cx_int]
    local_std = local_bg_std_map[cy_int, cx_int]
    mi = region.mean_intensity
    snr = (mi - local_mean) / max(local_std, 0.001)
    circ = compute_circularity(region)

    mask_bool = label_map == region.label
    orig_intensities = image_2d[mask_bool]

    # Get metadata from the corresponding object
    # region.label is 1-indexed, objects are 0-indexed
    obj_meta = final_filtered_objs[region.label - 1]['meta'] if region.label - 1 < len(final_filtered_objs) else {}

    obj_info = {
        'label': int(region.label),
        'centroid_y': float(cy),
        'centroid_x': float(cx),
        'area': int(region.area),
        'bbox': [int(x) for x in region.bbox],
        'mean_intensity_original': float(np.mean(orig_intensities)),
        'max_intensity_original': float(np.max(orig_intensities)),
        'min_intensity_original': float(np.min(orig_intensities)),
        'circularity': float(circ),
        'eccentricity': float(region.eccentricity),
        'solidity': float(region.solidity),
        'equivalent_diameter': float(region.equivalent_diameter),
        'perimeter': float(region.perimeter),
        'major_axis_length': float(region.major_axis_length),
        'minor_axis_length': float(region.minor_axis_length),
        'orientation': float(region.orientation),
        'local_background_mean': float(local_mean),
        'local_background_std': float(local_std),
        'local_snr': float(snr),
        'low_confidence': bool(obj_meta.get('low_confidence', False)),
        'regional_exception': bool(obj_meta.get('regional_exception', False)),
        'detection_method': str(obj_meta.get('detection_method', 'unknown')),
        'detection_pass': str(obj_meta.get('detection_pass', 'unknown')),
        'split_strategy_used': str(obj_meta.get('split_strategy_used', 'none')),
        'parent_label': obj_meta.get('parent_label', None),
        'gaussian_fit_separation': obj_meta.get('gaussian_fit_separation', None),
        'gaussian_fit_component_weights': obj_meta.get('gaussian_fit_component_weights', None),
        'concavity_deficit_fraction': obj_meta.get('concavity_deficit_fraction', None),
        'erosion_iterations_required': obj_meta.get('erosion_iterations_required', None),
    }

    results_objects.append(obj_info)

object_count = len(results_objects)
print(f"\nFinal validated object count: {object_count}")

# =====================================================================
# STEP 15: FINAL VALIDATION
# =====================================================================
print("\n=== FINAL VALIDATION ===")
print(f"Total objects: {object_count}")

# Confirm intensity gate compliance
gate_violations = 0
for obj in results_objects:
    mi = obj['mean_intensity_original']
    lm = obj['local_background_mean']
    ls = obj['local_background_std']
    snr = obj['local_snr']
    regional = obj['regional_exception']
    snr_thresh = 1.5 if regional else 2.0
    if snr < snr_thresh:
        gate_violations += 1
        print(f"  GATE VIOLATION: label {obj['label']} snr={snr:.2f} < {snr_thresh} (regional={regional})")
print(f"Intensity gate violations: {gate_violations}")

# Report target label status
for target in TARGET_CENTROIDS:
    status = "SPLIT" if target['split_success'] else "NOT SPLIT"
    print(f"Target {target['name']}: {status}")
    if target['split_success']:
        print(f"  Strategy: {target['split_strategy']}")
        print(f"  Children centroids: {target['children']}")
    else:
        print(f"  Matched index: {target['matched_idx']}")

# Report split stats
print(f"Split attempts: {split_stats['attempted']}, succeeded: {split_stats['succeeded']}, failed: {split_stats.get('failed', 0)}")
print(f"Split by strategy: {split_stats['by_strategy']}")

# Report regional exceptions
regional_objs = [o for o in results_objects if o['regional_exception']]
print(f"Objects recovered in left-center ROI with relaxed SNR: {len(regional_objs)}")
for ro in regional_objs:
    print(f"  Label {ro['label']} at ({ro['centroid_y']:.1f},{ro['centroid_x']:.1f}) snr={ro['local_snr']:.2f}")

print("=== END VALIDATION ===")

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

axes[0].imshow(image_2d, cmap='gray')
axes[0].set_title(f'Original Image\n{image_2d.shape}')
axes[0].axis('off')

axes[1].imshow(image_2d, cmap='gray')
colors = plt.cm.tab20(np.linspace(0, 1, max(20, object_count + 1)))
for i, region in enumerate(props):
    mask_bool = label_map == region.label
    color = colors[i % len(colors)]
    overlay = np.zeros((H, W, 4))
    overlay[mask_bool] = [color[0], color[1], color[2], 0.35]
    axes[1].imshow(overlay)

    contours_cv, _ = cv2.findContours(mask_bool.astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    for c in contours_cv:
        c = c.squeeze()
        if c.ndim == 2:
            axes[1].plot(c[:, 0], c[:, 1], color=color[:3], linewidth=1)

    cy, cx = region.centroid
    axes[1].text(cx, cy, str(region.label), fontsize=6, color='white', ha='center', va='center',
                fontweight='bold')

axes[1].set_title(f'Segmentation Overlay\n{object_count} objects')
axes[1].axis('off')

im3 = axes[2].imshow(local_bg_mean_map, cmap='viridis')
axes[2].set_title('Local Background Mean Map')
axes[2].axis('off')
plt.colorbar(im3, ax=axes[2], fraction=0.046)

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

# =====================================================================
# SAVE ARRAYS
# =====================================================================
np.save('analysis_labels.npy', label_map)
print("Saved analysis_labels.npy")

# =====================================================================
# OUTPUT JSON
# =====================================================================
centroids = [[obj['centroid_y'], obj['centroid_x']] for obj in results_objects]
areas = [obj['area'] for obj in results_objects]
bboxes = [obj['bbox'] for obj in results_objects]

target_validation = {}
for target in TARGET_CENTROIDS:
    target_validation[target['name']] = {
        'split_success': target['split_success'],
        'split_strategy': target['split_strategy'],
        'children_centroids': [[float(c[0]), float(c[1])] for c in target['children']] if target['children'] else [],
        'matched_index': target['matched_idx'],
    }

results = {
    "analysis_type": "Hybrid multi-pass object detection with local background estimation and multi-strategy splitting for overlapping/merged objects",
    "extracted_features": {
        "object_count": object_count,
        "centroids": centroids,
        "areas": areas,
        "bounding_boxes": bboxes,
        "mean_intensities_original": [obj['mean_intensity_original'] for obj in results_objects],
        "max_intensities_original": [obj['max_intensity_original'] for obj in results_objects],
        "min_intensities_original": [obj['min_intensity_original'] for obj in results_objects],
        "circularities": [obj['circularity'] for obj in results_objects],
        "eccentricities": [obj['eccentricity'] for obj in results_objects],
        "solidities": [obj['solidity'] for obj in results_objects],
        "equivalent_diameters": [obj['equivalent_diameter'] for obj in results_objects],
        "perimeters": [obj['perimeter'] for obj in results_objects],
        "major_axis_lengths": [obj['major_axis_length'] for obj in results_objects],
        "minor_axis_lengths": [obj['minor_axis_length'] for obj in results_objects],
        "orientations": [obj['orientation'] for obj in results_objects],
        "local_background_means": [obj['local_background_mean'] for obj in results_objects],
        "local_background_stds": [obj['local_background_std'] for obj in results_objects],
        "local_snrs": [obj['local_snr'] for obj in results_objects],
        "global_background_mean": float(global_bg_mean),
        "global_background_std": float(global_bg_std),
        "detection_methods": [obj['detection_method'] for obj in results_objects],
        "detection_passes": [obj['detection_pass'] for obj in results_objects],
        "split_strategies_used": [obj['split_strategy_used'] for obj in results_objects],
        "parent_labels": [obj['parent_label'] for obj in results_objects],
        "gaussian_fit_separations": [obj['gaussian_fit_separation'] for obj in results_objects],
        "gaussian_fit_component_weights": [obj['gaussian_fit_component_weights'] for obj in results_objects],
        "concavity_deficit_fractions": [obj['concavity_deficit_fraction'] for obj in results_objects],
        "erosion_iterations_required": [obj['erosion_iterations_required'] for obj in results_objects],
        "low_confidence_flags": [obj['low_confidence'] for obj in results_objects],
        "regional_exception_flags": [obj['regional_exception'] for obj in results_objects],
        "per_object_details": results_objects
    },
    "quality_metrics": {
        "total_objects_detected": object_count,
        "sam_initial_detections": len(sam_particles),
        "sam_after_gating": len(valid_sam_objs),
        "supplementary_added": added_1 + added_2 + added_3 + added_4,
        "split_attempts": split_stats['attempted'],
        "split_successes": split_stats['succeeded'],
        "split_failures": split_stats.get('failed', 0),
        "split_by_strategy": split_stats['by_strategy'],
        "target_label_validation": target_validation,
        "mean_local_snr": float(np.mean([obj['local_snr'] for obj in results_objects])) if results_objects else 0,
        "min_local_snr": float(np.min([obj['local_snr'] for obj in results_objects])) if results_objects else 0,
        "objects_with_low_confidence": sum(1 for obj in results_objects if obj['low_confidence']),
        "objects_with_regional_exception": sum(1 for obj in results_objects if obj['regional_exception']),
        "median_area": float(np.median(areas)) if areas else 0,
        "mean_circularity": float(np.mean([obj['circularity'] for obj in results_objects])) if results_objects else 0,
        "intensity_gate_violations": gate_violations,
        "split_log": split_log,
    },
    "summary": f"Detected {object_count} objects using hybrid SAM + adaptive thresholding with local background estimation. Split stats: {split_stats['succeeded']}/{split_stats['attempted']} objects split, {split_stats.get('failed',0)} failed. Strategies: {split_stats['by_strategy']}. Target label 10 split={TARGET_CENTROIDS[0]['split_success']} (strategy={TARGET_CENTROIDS[0]['split_strategy']}). Target label 15 split={TARGET_CENTROIDS[1]['split_success']} (strategy={TARGET_CENTROIDS[1]['split_strategy']}). Global background: mean={global_bg_mean:.1f}, std={global_bg_std:.1f}. Blur sigma=0.9, CLAHE clip_limit=2.5. Objects gated against local background maps with SNR thresholds (2.0 general, 1.5 left-center ROI). Regional exceptions: {sum(1 for obj in results_objects if obj['regional_exception'])} objects.",
    "saved_arrays": {
        "analysis_labels.npy": {
            "description": f"Integer label map with {object_count} objects labeled 1-{object_count}, background=0",
            "shape": list(label_map.shape),
            "dtype": str(label_map.dtype)
        }
    }
}

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