seg

import cv2
import numpy as np
import os
import glob
import sys
import pandas as pd
from skimage.segmentation import watershed
from skimage.feature import peak_local_max
from scipy import ndimage
from dotenv import load_dotenv

# ================= CONFIGURATION DEFAULTS =================
DEFAULT_CONFIG = {
    'MICRONS_PER_PIXEL': 0.25,
    'R_MAX_UM': 6.0,          # (Step 3) Limita expansiune citoplasma (6-8 um)
    'T_RING_UM': 1.0,         # (Step 4) Grosime inel membrana (1-2 um)
    'DAB_THRESHOLD_P90': 50,
    'MIN_CELL_AREA_PX': 30
}

def load_config():
    # ... (Aceeasi logica de incarcare config ca inainte) ...
    if len(sys.argv) < 2:
        print("Usage: python non_nucl_seg.py /path/to/config.env")
        sys.exit(1)
    env_path = sys.argv[1]
    if not os.path.exists(env_path):
        print(f"ERROR: Config file not found: {env_path}")
        sys.exit(1)
    load_dotenv(env_path)
    cfg = {}
    cfg['INPUT_ROOT'] = os.getenv('INPUT_ROOT')
    cfg['OUTPUT_ROOT'] = os.getenv('OUTPUT_ROOT')
    for key, default_val in DEFAULT_CONFIG.items():
        val = os.getenv(key)
        cfg[key] = float(val) if val else default_val
    mpp = cfg['MICRONS_PER_PIXEL']
    cfg['R_MAX_PX'] = int(cfg['R_MAX_UM'] / mpp)
    cfg['T_RING_PX'] = int(max(1, cfg['T_RING_UM'] / mpp))
    return cfg

def get_all_images_recursive(root_dir):
    extensions = ['*.jpg', '*.jpeg', '*.png', '*.tif', '*.tiff']
    image_files = []
    for dirpath, _, filenames in os.walk(root_dir):
        for filename in filenames:
            if any(filename.lower().endswith(ext.replace('*', '')) for ext in extensions):
                image_files.append(os.path.join(dirpath, filename))
    return image_files

def get_dab_channel(img_rgb):
    # Ruifrok-Johnston Deconvolution for DAB
    img_float = img_rgb.astype(float) + 1
    img_float = img_float / 255.0
    od = -np.log(img_float)
    dab_vector = np.array([0.268, 0.570, 0.776])
    dab_channel = np.dot(od, dab_vector)
    dab_channel = np.clip(dab_channel, 0, None)
    if dab_channel.max() > 0:
        dab_channel = (dab_channel / dab_channel.max()) * 255
    return dab_channel.astype(np.uint8)

def separate_nuclei_watershed(binary_mask):
    """
    Simuleaza 'Instance Segmentation': Sparge nucleii lipiti folosind distanta.
    (Step 1 - Replacement for HoVer-Net if not available)
    """
    # 1. Distanta de la centru la margine
    distance = ndimage.distance_transform_edt(binary_mask)

    # 2. Gasim varfurile locale (centrii nucleilor)
    # min_distance=5 asigura ca nu avem nuclei mai aproape de 5px unul de altul
    coords = peak_local_max(distance, footprint=np.ones((3, 3)), labels=binary_mask, min_distance=5)
    mask = np.zeros(distance.shape, dtype=bool)
    mask[tuple(coords.T)] = True
    markers, _ = ndimage.label(mask)

    # 3. Watershed pe inversul distantei pentru a separa nucleii
    labels = watershed(-distance, markers, mask=binary_mask)
    return labels

def process_pipeline():
    cfg = load_config()
    print(f"Running pipeline with Instance Separation logic...")

    image_files = get_all_images_recursive(cfg['INPUT_ROOT'])
    all_cells_data = []

    for filepath in image_files:
        try:
            filename = os.path.basename(filepath)
            rel_path = os.path.relpath(os.path.dirname(filepath), cfg['INPUT_ROOT'])
            current_output_dir = os.path.join(cfg['OUTPUT_ROOT'], rel_path)
            if not os.path.exists(current_output_dir): os.makedirs(current_output_dir)

            img = cv2.imread(filepath)
            if img is None: continue
            h, w = img.shape[:2]
            img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

            # --- PASUL 1: NUCLEUS DETECTION & INSTANCE SPLITTING ---
            # Daca aveai HoVer-Net, aici incarcai direct masca de instante.
            # Acum simulam asta detectand rosu si spargand obiectele.
            hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
            # Thresholds ajustate pentru nuclei rosii/maro
            mask1 = cv2.inRange(hsv, np.array([0, 40, 40]), np.array([15, 255, 255]))
            mask2 = cv2.inRange(hsv, np.array([160, 40, 40]), np.array([180, 255, 255]))
            nuclei_binary = mask1 + mask2
            nuclei_binary = ndimage.binary_fill_holes(nuclei_binary) # Umplem gaurile

            # CRITIC: Transformam masca binara in LABEL MAP (Instante separate)
            # Fiecare nucleu primeste un ID unic (1, 2, 3...)
            nuclei_labels = separate_nuclei_watershed(nuclei_binary)
            num_nuclei = nuclei_labels.max()

            if num_nuclei < 1: continue

            # --- PASUL 3: GEODESIC EXPANSION / VORONOI (Cytoplasm) ---
            # Vrem sa impartim tot spatiul dintre nuclei, dar cu limita R_MAX.

            # Calculam distanta de la ORICE pixel la cel mai apropiat nucleu
            # distance_transform_edt pe 'nuclei_labels > 0' inversat
            dist_from_nuclei, indices = ndimage.distance_transform_edt(nuclei_labels == 0, return_indices=True)

            # 'indices' ne spune pentru fiecare pixel de fundal, coordonatele celui mai apropiat nucleu.
            # Putem reconstrui 'Voronoi Map' complet folosind aceste indici:
            voronoi_map = nuclei_labels[indices[0], indices[1]]

            # Acum aplicam limita R_MAX (Step 3: limitare expansiune)
            # Pseudo-celula = Voronoi Region AND (Distance < R_MAX)
            cytoplasm_mask = np.where(dist_from_nuclei <= cfg['R_MAX_PX'], voronoi_map, 0)

            # --- PASII 4, 5, 6: MEMBRANE RING & CLIPPING ---

            final_vis_overlay = img.copy()
            final_mask_membranes = np.zeros((h, w), dtype=np.uint8)

            dab_image = get_dab_channel(img_rgb)

            # Iteram prin fiecare celula (INSTANTA) pentru a calcula inelul specific ei
            # Aceasta este cheia evitarii suprapunerilor

            # Optimizare: Putem face operatiile global, dar logic e mai clar per instanta
            # Pentru viteza, facem dilatarea pe fiecare 'island'

            count_valid = 0

            # Folosim 'regionprops' din skimage ar fi mai rapid, dar ramanem pe numpy/cv2 logic
            for cell_id in range(1, num_nuclei + 1):
                # 1. Extragem masca pseudo-celulei curente (Nucleu + Citoplasma)
                # Aceasta masca este DEJA taiata de vecini datorita Voronoi-ului de la Pasul 3
                current_cell_mask = (cytoplasm_mask == cell_id).astype(np.uint8)

                if np.sum(current_cell_mask) < cfg['MIN_CELL_AREA_PX']:
                    continue

                # 2. (Pasul 4) Dilatam celula curenta
                kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (2*cfg['T_RING_PX']+1, 2*cfg['T_RING_PX']+1))
                outer = cv2.dilate(current_cell_mask, kernel, iterations=1)

                # 3. (Pasul 5) Ring = Outer \ Inner
                ring = cv2.subtract(outer, current_cell_mask)

                # 4. (Pasul 6) CLIPPING STRICT
                # "Ring final = ring intersectat cu regiunea Voronoi a celulei"
                # Asta asigura ca inelul nu sare gardul la vecin.
                # Avem 'voronoi_map' calculat global la pasul 3.

                # Masca regiunii Voronoi pure a acestei celule (fara limita de distanta)
                voronoi_region = (voronoi_map == cell_id).astype(np.uint8)

                # Intersectie
                final_ring = cv2.bitwise_and(ring, voronoi_region)

                # Adaugam la masca globala pentru vizualizare
                final_mask_membranes = cv2.bitwise_or(final_mask_membranes, final_ring)

                # --- PASUL 7: DAB SCORE ---
                ring_pixels = dab_image[final_ring > 0]

                status = 'Negative'
                dab_p90 = 0
                if len(ring_pixels) > 0:
                    dab_p90 = np.percentile(ring_pixels, 90)
                    if dab_p90 > cfg['DAB_THRESHOLD_P90']:
                        status = 'Positive'

                all_cells_data.append({
                    'Image': filename,
                    'Cell_ID': cell_id,
                    'DAB_P90': round(dab_p90, 2),
                    'Status': status
                })
                count_valid += 1

                # Desenam conturul pe imagine (Galben=Pozitiv, Cyan=Negativ)
                color = (0, 255, 255) if status == 'Positive' else (255, 255, 0)
                contours, _ = cv2.findContours(final_ring, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
                cv2.drawContours(final_vis_overlay, contours, -1, color, 1)

            # --- SALVARE ---
            # Salvam imaginea cu contururi subtiri (care arata separarea)
            save_path = os.path.join(current_output_dir, f"analyzed_{filename}")
            cv2.imwrite(save_path, final_vis_overlay)
            print(f"Processed {filename}: {count_valid} cells.")

        except Exception as e:
            print(f"ERROR processing {filepath}: {e}")

    # Save Excel
    if all_cells_data:
        df = pd.DataFrame(all_cells_data)
        excel_path = os.path.join(cfg['OUTPUT_ROOT'], 'results.csv')
        df.to_csv(excel_path, index=False)
        print("Done.")

if __name__ == "__main__":
    process_pipeline()