rootscope/predict.py

"""
Predict cell types on new TIF images using ITERATIVE refinement.

Same as predict_cell_types.py but instead of a fixed 2-pass approach,
iterates predictions until neighbor-celltype features stabilize:

  Round 1: predict with neighbor celltypes = -1 (unknown)
  Round 2: fill neighbor celltypes from round 1, re-predict
  Round 3: fill neighbor celltypes from round 2, re-predict
  ...
  Stop when predictions no longer change OR max_rounds reached.

Usage:
  python predict_cell_types_iterative.py \
    --tif path/to/new_image.tif \
    --model-dir trained_model \
    --um-per-px 1.0 --gpu \
    --max-rounds 10

Output (per TIF, per model):
  predictions/
    {stem}_RandomForest_predictions.csv
    {stem}_RandomForest_overlay.png
    {stem}_XGBoost_predictions.csv
    {stem}_XGBoost_overlay.png
    {stem}_LightGBM_predictions.csv
    {stem}_LightGBM_overlay.png
"""

import argparse
import numpy as np
import pandas as pd
import joblib
from pathlib import Path
from PIL import Image, ImageDraw, ImageFont
from scipy import ndimage
from skimage.io import imread
from skimage.measure import regionprops
from skimage.morphology import binary_closing, binary_opening, remove_small_objects, disk
from skimage.segmentation import find_boundaries
from scipy.ndimage import binary_fill_holes, label as ndlabel
from collections import deque, defaultdict

from .extract_features import (
    ensure_rgb_uint8,
    to_grayscale_float,
    segment_cellpose_sam,
    build_tissue_mask,
    compute_layer_index_edt,
    build_cell_adjacency,
    extract_all_features,
    count_vascular_poles,
    compute_pole_features,
    compute_neighbor_celltypes,
    extract_cnn_embedding_features,
    CELL_CLASSES,
    LABEL_TO_NAME,
)


# Display palette (RGB)
DISPLAY_PALETTE = {
    0: (255, 105, 180),  # root_cap - pink
    1: (0,   0,   255),  # epidermis - blue
    2: (255, 255,   0),  # exodermis - yellow
    3: (0,   200,   0),  # cortex - green
    4: (255, 165,   0),  # endodermis - orange
    5: (128,   0, 128),  # pericycle - purple
    6: (255,   0,   0),  # xylem - red
    7: (255, 255, 255),  # phloem - white
    8: (0,   255, 255),  # stele - cyan
}


def _load_font(size):
    for path in [
        "/usr/share/fonts/liberation/LiberationMono-Regular.ttf",
        "/usr/share/fonts/dejavu/DejaVuSans.ttf",
        "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf",
    ]:
        try:
            return ImageFont.truetype(path, size)
        except (OSError, IOError):
            pass
    return ImageFont.load_default()


MODEL_NAMES = ["RandomForest", "XGBoost", "LightGBM"]


def load_models(model_dir):
    """Load all three models + scalers + encoder.

    Supports two layouts:
      - Per-model scalers: feature_scaler_RandomForest.joblib, etc. (iterative training)
      - Shared scaler: feature_scaler.joblib (old training)
    """
    model_dir = Path(model_dir)
    feature_cols = joblib.load(model_dir / "feature_columns.joblib")
    le = joblib.load(model_dir / "label_encoder.joblib")

    # Check for shared scaler (old layout)
    shared_scaler_path = model_dir / "feature_scaler.joblib"
    shared_scaler = None
    if shared_scaler_path.exists():
        shared_scaler = joblib.load(shared_scaler_path)

    models = {}
    scalers = {}
    for name in MODEL_NAMES:
        model_path = model_dir / f"model_{name}.joblib"
        if not model_path.exists():
            print(f"  WARNING: {model_path} not found, skipping {name}")
            continue
        models[name] = joblib.load(model_path)

        # Per-model scaler (iterative training) or shared scaler
        per_model_scaler_path = model_dir / f"feature_scaler_{name}.joblib"
        if per_model_scaler_path.exists():
            scalers[name] = joblib.load(per_model_scaler_path)
        elif shared_scaler is not None:
            scalers[name] = shared_scaler
        else:
            raise FileNotFoundError(
                f"No scaler found for {name} in {model_dir}")
        print(f"  Loaded {name} from {model_path}")

    print(f"  Classes: {list(le.classes_)}")
    print(f"  Features: {len(feature_cols)}")
    print(f"  Models loaded: {list(models.keys())}")
    return models, scalers, feature_cols, le


def save_prediction_overlay(img_rgb, masks, cell_predictions, out_path,
                             le, alpha=0.5, label_cells=False):
    """
    Overlay predicted cell-type colors on TIF, with a legend.

    label_cells=True also prints an abbreviated cell-type name inside every
    cell. Off by default: on a dense cross-section the per-cell text overlaps
    and hides the image. The color + legend already carry the same information.
    """
    H, W = masks.shape
    gray = to_grayscale_float(img_rgb)
    base = np.stack([gray, gray, gray], axis=-1)

    # Build color overlay from predictions
    color_rgb = np.zeros((H, W, 3), dtype=np.uint8)
    for cid, pred_name in cell_predictions.items():
        label_int = CELL_CLASSES.get(pred_name, -1)
        if label_int >= 0 and label_int in DISPLAY_PALETTE:
            color_rgb[masks == cid] = DISPLAY_PALETTE[label_int]

    has_color = color_rgb.sum(axis=-1) > 0
    mask_f = has_color[..., None].astype(np.float32)
    blended = base * (1.0 - alpha * mask_f) + color_rgb.astype(np.float32) * (alpha * mask_f)

    # Cell boundaries
    boundaries = find_boundaries(masks, mode="outer")
    blended[boundaries] = [255, 255, 255]

    im = Image.fromarray(np.clip(blended, 0, 255).astype(np.uint8))
    draw = ImageDraw.Draw(im)

    # Optional: draw the cell-type name inside each cell (--label-cells)
    if label_cells:
        props = {p.label: p for p in regionprops(masks)}
        med_area = float(np.median([p.area for p in props.values()])) if props else 100.0
        base_fs = max(7, min(14, int(np.sqrt(med_area) * 0.25)))
        font = _load_font(base_fs)

        for cid, pred_name in cell_predictions.items():
            if cid not in props:
                continue
            cy, cx = props[cid].centroid
            x, y = int(round(cx)), int(round(cy))
            # Abbreviate: first 4 chars
            abbr = pred_name[:4]
            for dx, dy in [(-1, 0), (1, 0), (0, -1), (0, 1)]:
                draw.text((x + dx, y + dy), abbr, fill=(0, 0, 0), font=font, anchor="mm")
            draw.text((x, y), abbr, fill=(255, 255, 255), font=font, anchor="mm")

    # Legend
    legend_font = _load_font(13)
    legend_x = W - 140
    legend_y = 10
    present_types = sorted(set(cell_predictions.values()))
    n_legend = len(present_types)
    draw.rectangle(
        [legend_x - 4, legend_y - 4, W - 4, legend_y + n_legend * 20 + 4],
        fill=(0, 0, 0),
    )
    for i, name in enumerate(present_types):
        label_int = CELL_CLASSES.get(name, -1)
        color = DISPLAY_PALETTE.get(label_int, (128, 128, 128))
        ly = legend_y + i * 20
        draw.rectangle([legend_x, ly, legend_x + 14, ly + 14], fill=color,
                       outline=(255, 255, 255))
        draw.text((legend_x + 20, ly), name, fill=(255, 255, 255), font=legend_font)

    im.save(str(out_path))


def anatomical_postprocess(df, layer_lookup, adjacency, masks, le=None, y_proba=None,
                           rf_proba_lookup=None):
    """
    Enforce known root anatomical ring order using the ADJACENCY GRAPH.

    Root anatomy (outside -> inside):
      root_cap -> epidermis -> exodermis -> cortex -> endodermis -> pericycle -> stele

    Rules (layer + adjacency based):
      - Exodermis = ONE layer only (epidermis_layer + 1)
      - Endodermis = ONE layer only (the layer neighboring pericycle/stele)
      - Pericycle = ONE layer only (just inside endodermis)
      - Cortex fills the multiple layers between exodermis and endodermis
      - Recover phloem and xylem from stele using probability scores
    """
    predictions = dict(zip(df["cell_id"].values, df["predicted_cell_type"].values))

    # Build probability lookup for phloem recovery
    proba_lookup = {}
    if y_proba is not None and le is not None:
        class_names = list(le.classes_)
        for i, cid in enumerate(df["cell_id"].values):
            proba_lookup[int(cid)] = dict(zip(class_names, y_proba[i]))

    # -- Step 1: Identify tissue groups from model predictions --
    epidermis_cells = {cid for cid, pred in predictions.items()
                       if pred == "epidermis"}
    root_cap_cells = {cid for cid, pred in predictions.items()
                      if pred == "root_cap"}
    outer_cells = epidermis_cells | root_cap_cells

    inner_types = ("pericycle", "stele", "xylem", "phloem")
    inner_cells = {cid for cid, pred in predictions.items()
                   if pred in inner_types}

    # -- Step 2: Find exodermis ring -- SINGLE LAYER only --
    epi_layer_counts = defaultdict(int)
    for cid in epidermis_cells:
        lv = layer_lookup.get(cid, -1)
        if lv >= 0:
            epi_layer_counts[lv] += 1
    main_epi_layer = max(epi_layer_counts, key=epi_layer_counts.get) if epi_layer_counts else 0

    epi_neighbor_cells = set()
    for epi_cid in epidermis_cells:
        for nbr in adjacency.get(epi_cid, set()):
            if nbr not in outer_cells and nbr not in inner_cells:
                epi_neighbor_cells.add(nbr)

    epi_nbr_layer_counts = defaultdict(int)
    for cid in epi_neighbor_cells:
        lv = layer_lookup.get(cid, -1)
        if lv >= 0 and lv > main_epi_layer:
            epi_nbr_layer_counts[lv] += 1

    if epi_nbr_layer_counts:
        exo_layer = max(epi_nbr_layer_counts, key=epi_nbr_layer_counts.get)
    else:
        exo_layer = main_epi_layer + 1

    n_exo_promoted = 0
    n_exo_demoted = 0
    for cid, pred in list(predictions.items()):
        lv = layer_lookup.get(cid, -1)
        if lv == exo_layer and cid not in outer_cells and cid not in inner_cells:
            if pred != "exodermis":
                n_exo_promoted += 1
            predictions[cid] = "exodermis"
        elif pred == "exodermis":
            predictions[cid] = "cortex"
            n_exo_demoted += 1

    n_final_exo = sum(1 for p in predictions.values() if p == "exodermis")
    print(f"    Post-processing: exodermis restricted to layer {exo_layer} "
          f"(epi_layer={main_epi_layer}): {n_final_exo} exodermis cells, "
          f"promoted {n_exo_promoted} into exo, demoted {n_exo_demoted} to cortex")

    # -- Step 3: Find endodermis ring -- SINGLE LAYER only --
    cells_neighboring_inner = set()
    for inner_cid in inner_cells:
        for nbr in adjacency.get(inner_cid, set()):
            if nbr not in inner_cells and nbr not in outer_cells:
                cells_neighboring_inner.add(nbr)

    exo_cells_final = {cid for cid, pred in predictions.items()
                       if pred == "exodermis"}
    endo_candidates = cells_neighboring_inner - exo_cells_final

    endo_layer_counts = defaultdict(int)
    for cid in endo_candidates:
        lv = layer_lookup.get(cid, -1)
        if lv >= 0:
            endo_layer_counts[lv] += 1
    endo_layer = max(endo_layer_counts, key=endo_layer_counts.get) if endo_layer_counts else -1

    n_relabeled_endo = 0
    for cid in endo_candidates:
        lv = layer_lookup.get(cid, -1)
        if lv == endo_layer:
            if predictions.get(cid) != "endodermis":
                n_relabeled_endo += 1
            predictions[cid] = "endodermis"

    n_fixed_endo = 0
    for cid, pred in list(predictions.items()):
        if pred == "endodermis":
            lv = layer_lookup.get(cid, -1)
            if lv != endo_layer:
                predictions[cid] = "cortex"
                n_fixed_endo += 1

    n_final_endo = sum(1 for p in predictions.values() if p == "endodermis")
    print(f"    Post-processing: endodermis restricted to layer {endo_layer}: "
          f"{n_final_endo} endodermis cells, relabeled {n_relabeled_endo} into endo, "
          f"demoted {n_fixed_endo} to cortex")

    # -- Step 4: Pericycle -- SINGLE RING only (adjacency-based) --
    endo_cells_final = {cid for cid, pred in predictions.items()
                        if pred == "endodermis"}

    peri_candidates = set()
    for endo_cid in endo_cells_final:
        endo_lv = layer_lookup.get(endo_cid, -1)
        for nbr in adjacency.get(endo_cid, set()):
            nbr_lv = layer_lookup.get(nbr, -1)
            if nbr_lv > endo_lv and nbr not in endo_cells_final:
                nbr_pred = predictions.get(nbr)
                if nbr_pred not in ("epidermis", "root_cap", "exodermis", "cortex",
                                    "endodermis"):
                    peri_candidates.add(nbr)

    peri_layer_counts = defaultdict(int)
    for cid in peri_candidates:
        lv = layer_lookup.get(cid, -1)
        if lv >= 0:
            peri_layer_counts[lv] += 1
    peri_layer = max(peri_layer_counts, key=peri_layer_counts.get) if peri_layer_counts else (endo_layer + 1 if endo_layer >= 0 else -1)

    allowed_pericycle = set()
    n_peri_promoted = 0
    for cid in peri_candidates:
        lv = layer_lookup.get(cid, -1)
        if lv == peri_layer:
            allowed_pericycle.add(cid)
            if predictions.get(cid) != "pericycle":
                n_peri_promoted += 1
            predictions[cid] = "pericycle"

    n_peri_demoted = 0
    for cid, pred in list(predictions.items()):
        if pred == "pericycle" and cid not in allowed_pericycle:
            predictions[cid] = "stele"
            n_peri_demoted += 1

    n_final_peri = sum(1 for p in predictions.values() if p == "pericycle")
    print(f"    Post-processing: pericycle restricted to endodermis-adjacent cells "
          f"in layer {peri_layer} (endo_layer={endo_layer}): {n_final_peri} pericycle "
          f"cells, promoted {n_peri_promoted} into peri, demoted {n_peri_demoted} to stele")

    # -- Step 5: Fix stray epidermis deep inside the root --
    n_fixed_epi = 0
    for cid, pred in list(predictions.items()):
        if pred == "epidermis":
            lv = layer_lookup.get(cid, -1)
            if lv >= main_epi_layer + 3:
                nbrs = adjacency.get(cid, set())
                nbr_preds = [predictions.get(n) for n in nbrs
                             if n in predictions and predictions.get(n) != "epidermis"]
                if nbr_preds:
                    from collections import Counter
                    predictions[cid] = Counter(nbr_preds).most_common(1)[0][0]
                    n_fixed_epi += 1
    if n_fixed_epi > 0:
        print(f"    Post-processing: fixed {n_fixed_epi} stray epidermis cells")

    # -- Step 6: Recover phloem and xylem from inner tissue using probabilities --
    if proba_lookup:
        def _effective_prob(cid, cls_name):
            """Return the best available probability for xylem/phloem."""
            probs = proba_lookup.get(cid, {})
            model_prob = probs.get(cls_name, 0.0)
            if rf_proba_lookup and model_prob < 0.05:
                rf_probs = rf_proba_lookup.get(cid, {})
                rf_prob = rf_probs.get(cls_name, 0.0)
                if rf_prob > model_prob:
                    return rf_prob
            return model_prob

        # Recover phloem from stele (and xylem)
        n_recovered_phloem = 0
        n_recovered_phloem_rf = 0
        for cid, pred in list(predictions.items()):
            if pred in ("stele", "xylem"):
                phloem_prob = _effective_prob(cid, "phloem")
                probs = proba_lookup.get(cid, {})
                current_prob = probs.get(pred, 0.0)
                if phloem_prob >= 0.12 and phloem_prob >= current_prob - 0.25:
                    predictions[cid] = "phloem"
                    n_recovered_phloem += 1
                    if probs.get("phloem", 0.0) < 0.05:
                        n_recovered_phloem_rf += 1
        if n_recovered_phloem > 0:
            msg = f"    Post-processing: recovered {n_recovered_phloem} phloem cells from stele/xylem"
            if n_recovered_phloem_rf > 0:
                msg += f" ({n_recovered_phloem_rf} via RF fallback)"
            print(msg)

        # Recover xylem from stele
        n_recovered_xylem = 0
        n_recovered_xylem_rf = 0
        for cid, pred in list(predictions.items()):
            if pred == "stele":
                xylem_prob = _effective_prob(cid, "xylem")
                probs = proba_lookup.get(cid, {})
                current_prob = probs.get(pred, 0.0)
                if xylem_prob >= 0.08 and xylem_prob >= current_prob - 0.30:
                    predictions[cid] = "xylem"
                    n_recovered_xylem += 1
                    if probs.get("xylem", 0.0) < 0.05:
                        n_recovered_xylem_rf += 1
        if n_recovered_xylem > 0:
            msg = f"    Post-processing: recovered {n_recovered_xylem} xylem cells from stele"
            if n_recovered_xylem_rf > 0:
                msg += f" ({n_recovered_xylem_rf} via RF fallback)"
            print(msg)

        # Xylem strand growing
        n_strand_grown = 0
        for _round in range(3):
            grown_this_round = 0
            current_xylem = {cid for cid, p in predictions.items() if p == "xylem"}
            for cid, pred in list(predictions.items()):
                if pred != "stele":
                    continue
                nbrs = adjacency.get(cid, set())
                n_xylem_nbrs = sum(1 for n in nbrs if n in current_xylem)
                if n_xylem_nbrs >= 2:
                    xylem_prob = _effective_prob(cid, "xylem")
                    if xylem_prob >= 0.05:
                        predictions[cid] = "xylem"
                        grown_this_round += 1
            n_strand_grown += grown_this_round
            if grown_this_round == 0:
                break
        if n_strand_grown > 0:
            print(f"    Post-processing: grew xylem strands by {n_strand_grown} cells")

    # Update dataframe
    df["predicted_cell_type"] = df["cell_id"].map(predictions)
    return df, predictions


def _prepare_features(df_base, feature_cols, scaler):
    """Align feature columns and scale. Returns X_scaled, handling NaN."""
    for col in feature_cols:
        if col not in df_base.columns:
            df_base[col] = 0.0

    X = df_base[feature_cols].values.astype(np.float32)

    # Handle NaN
    nan_mask = np.isnan(X)
    if nan_mask.any():
        col_means = np.nanmean(X, axis=0)
        for j in range(X.shape[1]):
            X[nan_mask[:, j], j] = col_means[j] if not np.isnan(col_means[j]) else 0.0

    X_scaled = scaler.transform(X)
    return X_scaled


def iterative_predict(model, df_base, adjacency, feature_cols, scaler, le,
                      max_rounds=10):
    """
    Iteratively predict cell types until neighbor-celltype features converge.

    Round 1: predict with neighbor celltypes = -1 (unknown)
    Round N: fill neighbor celltypes from round N-1 predictions, re-predict
    Stop when no predictions change or max_rounds reached.

    Returns: (df, pred_names, y_proba) from the final round.
    """
    neighbor_ct_cols = ["radial_inward_neighbor_celltype", "radial_outward_neighbor_celltype",
                        "tangential_cw_neighbor_celltype", "tangential_ccw_neighbor_celltype"]
    has_neighbor_ct = any(c in feature_cols for c in neighbor_ct_cols)

    prev_pred_names = None
    best_pred_names = None
    best_proba = None
    best_df = None
    best_n_changed = float("inf")
    prev_n_changed = float("inf")
    n_stalled = 0  # count rounds where n_changed doesn't decrease

    for round_num in range(1, max_rounds + 1):
        df_round = df_base.copy()

        if has_neighbor_ct:
            if round_num == 1:
                # First round: no neighbor celltype info
                df_round = compute_neighbor_celltypes(
                    df_round, adjacency, cell_type_labels=None
                )
            else:
                # Use previous round's predictions as neighbor celltypes
                prev_labels = {}
                for i, cid in enumerate(df_base["cell_id"].values):
                    prev_labels[int(cid)] = CELL_CLASSES.get(prev_pred_names[i], -1)

                df_round = compute_neighbor_celltypes(
                    df_round, adjacency, cell_type_labels=prev_labels
                )

        X_scaled = _prepare_features(df_round, feature_cols, scaler)
        y_pred = model.predict(X_scaled)
        y_proba = model.predict_proba(X_scaled)
        pred_names = le.inverse_transform(y_pred)

        # Count changes from previous round
        if prev_pred_names is not None:
            n_changed = sum(1 for a, b in zip(prev_pred_names, pred_names) if a != b)
            n_total = len(pred_names)
            print(f"      Round {round_num}: {n_changed}/{n_total} cells changed "
                  f"({100.0 * n_changed / n_total:.2f}%)")

            if n_changed == 0:
                print(f"      Converged at round {round_num}!")
                best_pred_names = pred_names.copy()
                best_proba = y_proba.copy()
                best_df = df_round.copy()
                break

            # Track best round (fewest changes = most stable)
            if n_changed < best_n_changed:
                best_n_changed = n_changed
                best_pred_names = pred_names.copy()
                best_proba = y_proba.copy()
                best_df = df_round.copy()
                n_stalled = 0
            else:
                n_stalled += 1

            # Detect oscillation: if no improvement for 3 rounds, stop
            if n_stalled >= 3:
                print(f"      Oscillation detected (no improvement for 3 rounds), "
                      f"using best round with {best_n_changed} changes.")
                break

            prev_n_changed = n_changed
        else:
            print(f"      Round 1: initial prediction ({len(pred_names)} cells)")
            best_pred_names = pred_names.copy()
            best_proba = y_proba.copy()
            best_df = df_round.copy()

        prev_pred_names = pred_names.copy()

    else:
        print(f"      Reached max rounds ({max_rounds}), "
              f"using best round with {best_n_changed} changes.")

    return best_df, best_pred_names, best_proba


def load_image(tif_path):
    """Read a TIFF and return it as an RGB uint8 array (z-stacks max-projected)."""
    img_raw = imread(str(tif_path))
    return ensure_rgb_uint8(img_raw, stack_mode="max")


# ═══════════════════════════════════════════════════════════════
# PIPELINE STAGES
#
# predict_single_tif() below runs these in order. They are also exposed
# separately so a caller that pays for GPU time by the second (a web app on
# ZeroGPU, say) can wrap only stage_segment and stage_embed — the two steps
# that actually touch the GPU — and run the rest on CPU.
# ═══════════════════════════════════════════════════════════════

def stage_segment(img_rgb, gpu=True, cellpose_model=None):
    """GPU stage 1 — Cellpose-SAM. Returns the int32 label mask."""
    return segment_cellpose_sam(img_rgb, use_gpu=gpu, model=cellpose_model)


def stage_features(masks, img_rgb, um_per_px=1.0, verbose=True):
    """CPU stage — tissue mask, layer index, debris removal, handcrafted
    features.

    Debris removal edits ``masks`` in place, so the (possibly modified) mask is
    returned alongside the feature table.

    Returns ``(masks, df_base, layer_lookup, adjacency, n_layers)``, or
    ``(masks, None, ...)`` if every cell was debris.
    """
    if verbose:
        print("    Computing layer index...")
    tissue = build_tissue_mask(masks)
    layer_lookup, n_layers, adjacency = compute_layer_index_edt(masks, tissue)
    if verbose:
        print(f"    {n_layers} layers")

    # Filter out debris cells outside the main tissue body
    n_cells = int(masks.max())
    debris_cells = set()
    for cid in range(1, n_cells + 1):
        cell_px = masks == cid
        n_total = cell_px.sum()
        if n_total == 0:
            continue
        n_in_tissue = (cell_px & tissue).sum()
        if n_in_tissue / n_total < 0.5:
            debris_cells.add(cid)
    if debris_cells:
        if verbose:
            print(f"    Removed {len(debris_cells)} debris cells outside tissue")
        for cid in debris_cells:
            masks[masks == cid] = 0
            layer_lookup.pop(cid, None)
        n_cells = int(masks.max())

    if n_cells == 0:
        return masks, None, layer_lookup, adjacency, n_layers

    if verbose:
        print("    Extracting features...")
    df_base = extract_all_features(masks, img_rgb, um_per_px, layer_lookup,
                                   adjacency, tissue_mask=tissue)
    df_base["n_layers_total"] = n_layers
    return masks, df_base, layer_lookup, adjacency, n_layers


def stage_embed(masks, img_rgb, df_base, gpu=True, cnn_weights=None,
                dinov2_model=None, verbose=True):
    """GPU stage 2 — fine-tuned DINOv2 per-cell embeddings, merged into
    ``df_base``. Returns the merged table (unchanged if embeddings are
    unavailable)."""
    if verbose:
        print("    Extracting CNN embeddings...")
    cnn_df = extract_cnn_embedding_features(masks, img_rgb, use_gpu=gpu,
                                             weights_path=cnn_weights,
                                             model=dinov2_model)
    if cnn_df is not None:
        df_base = df_base.merge(cnn_df, on="cell_id", how="left")
        emb_cols = [c for c in df_base.columns if c.startswith("cnn_emb_")]
        df_base[emb_cols] = df_base[emb_cols].fillna(0.0)
    return df_base


def stage_classify(df_base, masks, img_rgb, layer_lookup, adjacency,
                   models_dict, scalers, feature_cols, le,
                   out_dir="predictions", stem="image", source_name=None,
                   um_per_px=1.0, max_rounds=10, label_cells=False):
    """CPU stage — iterative prediction with each model plus the weighted
    ensemble, anatomical post-processing, and per-model CSV + overlay PNG.

    Returns the concatenated per-cell table across all models."""
    out_dir = Path(out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)
    if source_name is None:
        source_name = f"{stem}.tif"

    all_dfs = []
    model_probas = {}  # model_name -> (cell_ids, y_proba)
    for model_name, model in models_dict.items():
        print(f"\n    --- {model_name} (iterative, max {max_rounds} rounds) ---")

        model_scaler = scalers[model_name]
        df, pred_names, y_proba = iterative_predict(
            model, df_base, adjacency, feature_cols, model_scaler, le,
            max_rounds=max_rounds,
        )

        pred_conf = y_proba.max(axis=1)
        df["predicted_cell_type"] = pred_names
        df["prediction_confidence"] = np.round(pred_conf, 4)
        df["source_file"] = source_name
        df["um_per_px"] = um_per_px
        df["model"] = model_name

        # Report (before post-processing)
        print(f"    Model predictions:")
        for ct, cnt in pd.Series(pred_names).value_counts().items():
            print(f"      {ct}: {cnt}")

        # Anatomical post-processing
        rf_fallback = None
        if model_name != "RandomForest" and "RandomForest" in model_probas:
            rf_cids, rf_proba_arr = model_probas["RandomForest"]
            rf_class_names = list(le.classes_)
            rf_fallback = {}
            for i, cid in enumerate(rf_cids):
                rf_fallback[int(cid)] = dict(zip(rf_class_names, rf_proba_arr[i]))

        print(f"    Applying anatomical post-processing...")
        df, cell_predictions = anatomical_postprocess(
            df, layer_lookup, adjacency, masks, le=le, y_proba=y_proba,
            rf_proba_lookup=rf_fallback
        )

        # Report (after post-processing)
        final_preds = df["predicted_cell_type"].values
        print(f"    Final predictions:")
        for ct, cnt in pd.Series(final_preds).value_counts().items():
            print(f"      {ct}: {cnt}")

        # Count vascular poles
        final_labels = {}
        for cid, pred in cell_predictions.items():
            final_labels[cid] = CELL_CLASSES.get(pred, -1)
        pole_info = count_vascular_poles(final_labels, adjacency, masks)
        n_ph = pole_info["phloem"]["n_poles"]
        n_xy = pole_info["xylem"]["n_poles"]
        print(f"    Vascular poles: phloem={n_ph}, xylem={n_xy}")

        df["n_phloem_poles"] = n_ph
        df["n_xylem_poles"] = n_xy

        # Store probabilities for ensemble
        model_probas[model_name] = (df["cell_id"].values.copy(), y_proba.copy())

        # Save per-model CSV
        csv_path = out_dir / f"{stem}_{model_name}_predictions.csv"
        df.to_csv(csv_path, index=False)
        print(f"    Saved: {csv_path}")

        # Save per-model overlay PNG
        overlay_path = out_dir / f"{stem}_{model_name}_overlay.png"
        save_prediction_overlay(img_rgb, masks, cell_predictions,
                                overlay_path, le, label_cells=label_cells)
        print(f"    Saved: {overlay_path}")

        all_dfs.append(df)

    # -- Ensemble voting (average probabilities across models) --
    if len(model_probas) >= 2:
        print(f"\n    --- Ensemble ({len(model_probas)} models) ---")

        ref_model = list(model_probas.keys())[0]
        ref_cell_ids = model_probas[ref_model][0]
        n_classes = model_probas[ref_model][1].shape[1]

        minority_classes = {"xylem", "phloem"}
        minority_idx = set()
        for cls_name in minority_classes:
            try:
                minority_idx.add(list(le.classes_).index(cls_name))
            except ValueError:
                pass
        rf_trust_threshold = 0.20

        avg_proba = np.zeros((len(ref_cell_ids), n_classes), dtype=np.float64)
        n_models = len(model_probas)
        rf_proba = model_probas.get("RandomForest", (None, None))[1]

        if rf_proba is not None and minority_idx:
            rf_top_idx = np.argmax(rf_proba, axis=1)
            rf_top_conf = rf_proba[np.arange(len(rf_proba)), rf_top_idx]
            rf_minority_mask = np.array([
                (idx in minority_idx and conf >= rf_trust_threshold)
                for idx, conf in zip(rf_top_idx, rf_top_conf)
            ])
            n_rf_trusted = rf_minority_mask.sum()
            print(f"    RF-trust rule: {n_rf_trusted} cells where RF predicts minority class")

            for mname, (cids, proba) in model_probas.items():
                weight = np.ones(len(ref_cell_ids), dtype=np.float64)
                if mname == "RandomForest":
                    weight[rf_minority_mask] = 2.0
                avg_proba += proba * weight[:, np.newaxis]
            total_weight = np.full(len(ref_cell_ids), float(n_models), dtype=np.float64)
            total_weight[rf_minority_mask] += 1.0
            avg_proba /= total_weight[:, np.newaxis]
        else:
            for mname, (cids, proba) in model_probas.items():
                avg_proba += proba
            avg_proba /= n_models

        ens_pred_idx = np.argmax(avg_proba, axis=1)
        ens_pred_names = le.inverse_transform(ens_pred_idx)
        ens_conf = avg_proba.max(axis=1)

        df_ens = df_base.copy()
        df_ens["predicted_cell_type"] = ens_pred_names
        df_ens["prediction_confidence"] = np.round(ens_conf, 4)
        df_ens["source_file"] = source_name
        df_ens["um_per_px"] = um_per_px
        df_ens["model"] = "Ensemble"

        print(f"    Ensemble predictions:")
        for ct, cnt in pd.Series(ens_pred_names).value_counts().items():
            print(f"      {ct}: {cnt}")

        print(f"    Applying anatomical post-processing...")
        df_ens, ens_cell_predictions = anatomical_postprocess(
            df_ens, layer_lookup, adjacency, masks, le=le, y_proba=avg_proba
        )

        final_ens_preds = df_ens["predicted_cell_type"].values
        print(f"    Final predictions:")
        for ct, cnt in pd.Series(final_ens_preds).value_counts().items():
            print(f"      {ct}: {cnt}")

        ens_final_labels = {}
        for cid, pred in ens_cell_predictions.items():
            ens_final_labels[cid] = CELL_CLASSES.get(pred, -1)
        ens_pole_info = count_vascular_poles(ens_final_labels, adjacency, masks)
        n_ph_ens = ens_pole_info["phloem"]["n_poles"]
        n_xy_ens = ens_pole_info["xylem"]["n_poles"]
        print(f"    Vascular poles: phloem={n_ph_ens}, xylem={n_xy_ens}")

        df_ens["n_phloem_poles"] = n_ph_ens
        df_ens["n_xylem_poles"] = n_xy_ens

        ens_csv_path = out_dir / f"{stem}_Ensemble_predictions.csv"
        df_ens.to_csv(ens_csv_path, index=False)
        print(f"    Saved: {ens_csv_path}")

        ens_overlay_path = out_dir / f"{stem}_Ensemble_overlay.png"
        save_prediction_overlay(img_rgb, masks, ens_cell_predictions,
                                ens_overlay_path, le, label_cells=label_cells)
        print(f"    Saved: {ens_overlay_path}")

        all_dfs.append(df_ens)

    return pd.concat(all_dfs, ignore_index=True)


def predict_single_tif(tif_path, models_dict, scalers, feature_cols, le,
                        um_per_px=1.0, gpu=True, out_dir="predictions",
                        max_rounds=10, cnn_weights=None, label_cells=False,
                        cellpose_model=None, dinov2_model=None):
    """
    Full pipeline: segment -> features -> iterative predict -> overlay.
    """
    out_dir = Path(out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)
    stem = Path(tif_path).stem.replace(".aivia", "")

    print(f"\n  Processing: {tif_path}")

    # Load TIF
    img_rgb = load_image(tif_path)
    print(f"    Image shape: {img_rgb.shape}")

    # Segment
    print("    Segmenting (Cellpose-SAM)...")
    masks = stage_segment(img_rgb, gpu=gpu, cellpose_model=cellpose_model)
    n_cells = int(masks.max())
    print(f"    {n_cells} cells found")

    if n_cells == 0:
        print("    No cells found, skipping.")
        return None

    # BFS layer index, debris removal, handcrafted features
    masks, df_base, layer_lookup, adjacency, n_layers = stage_features(
        masks, img_rgb, um_per_px=um_per_px)
    if df_base is None:
        print("    No cells left after debris removal, skipping.")
        return None

    # CNN embeddings
    df_base = stage_embed(masks, img_rgb, df_base, gpu=gpu,
                          cnn_weights=cnn_weights, dinov2_model=dinov2_model)

    return stage_classify(
        df_base, masks, img_rgb, layer_lookup, adjacency,
        models_dict, scalers, feature_cols, le,
        out_dir=out_dir, stem=stem, source_name=Path(tif_path).name,
        um_per_px=um_per_px, max_rounds=max_rounds, label_cells=label_cells,
    )


def main():
    parser = argparse.ArgumentParser(
        description="Predict cell types on new TIF images (iterative refinement)"
    )
    parser.add_argument("--tif", default=None,
                        help="Single TIF path")
    parser.add_argument("--tif-dir", default=None,
                        help="Directory of TIF files (batch mode)")
    parser.add_argument("--model-dir", required=True,
                        help="Directory with trained model artifacts")
    parser.add_argument("--um-per-px", type=float, default=1.0,
                        help="Microns per pixel (default 1.0)")
    parser.add_argument("--gpu", action="store_true",
                        help="Use GPU for Cellpose")
    parser.add_argument("--out-dir", default="predictions",
                        help="Output directory")
    parser.add_argument("--max-rounds", type=int, default=10,
                        help="Max iterative prediction rounds (default 10)")
    parser.add_argument("--cnn-weights", default=None,
                        help="Path to fine-tuned DINOv2 backbone.pt (must "
                             "match what the models were trained on). "
                             "If omitted, uses pretrained DINOv2.")
    args = parser.parse_args()

    if not args.tif and not args.tif_dir:
        parser.error("Provide --tif for single image or --tif-dir for batch")

    print("=" * 60)
    print("CELL TYPE PREDICTION (ITERATIVE)")
    print("=" * 60)

    print("\n[1] Loading models...")
    models_dict, scalers, feature_cols, le = load_models(args.model_dir)

    if not models_dict:
        print("  ERROR: No models found. Run train_classifier.py first.")
        return

    # Collect TIF paths
    if args.tif:
        tif_paths = [Path(args.tif)]
    else:
        tif_dir = Path(args.tif_dir)
        tif_paths = sorted(tif_dir.glob("*.tif"))
        print(f"\n  Found {len(tif_paths)} TIF files in {tif_dir}")

    print(f"\n[2] Predicting cell types with {len(models_dict)} models "
          f"({', '.join(models_dict.keys())}) — iterative, max {args.max_rounds} rounds...")
    all_tables = []
    for tif_path in tif_paths:
        try:
            df = predict_single_tif(
                str(tif_path), models_dict, scalers, feature_cols, le,
                um_per_px=args.um_per_px, gpu=args.gpu, out_dir=args.out_dir,
                max_rounds=args.max_rounds,
                cnn_weights=args.cnn_weights,
                label_cells=getattr(args, "label_cells", False),
            )
            if df is not None:
                all_tables.append(df)
        except Exception as e:
            print(f"    FAILED: {e}")
            import traceback
            traceback.print_exc()

    # Combined CSV (all models x all images)
    if all_tables:
        combined = pd.concat(all_tables, ignore_index=True)
        combined_path = Path(args.out_dir) / "all_predictions.csv"
        combined.to_csv(combined_path, index=False)
        print(f"\n  Combined predictions: {combined_path} ({len(combined)} rows)")

    print(f"\n{'=' * 60}")
    print(f"DONE. Results in {args.out_dir}/")
    print(f"{'=' * 60}")


if __name__ == "__main__":
    main()