#!/usr/bin/env python
# -*- coding: utf-8 -*-

import os
import cv2
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from scipy.ndimage import gaussian_laplace
from skimage.feature import peak_local_max

# =========================
# CONFIG
# =========================
image_folder = "images"
output_folder = "FINAL_PIPELINE"
os.makedirs(output_folder, exist_ok=True)

N_ROI = 5
ROI_SIZE = 120   # pixel（固定，不再用pt/mm）

# =========================
# 🔬 SAFE ENHANCEMENT (FIX CLAHE BUG)
# =========================
def enhance(img):

    # ✔ 保证 uint8
    if img.dtype != np.uint8:
        img = cv2.normalize(img, None, 0, 255, cv2.NORM_MINMAX).astype(np.uint8)

    clahe = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(8,8))
    img = clahe.apply(img)

    # float for analysis
    img = img.astype(np.float32)
    img = (img - img.min()) / (img.max() + 1e-6)

    return img


# =========================
# 🔬 AUTO TYPE DETECTION (S2 vs SF9)
# =========================
def detect_type(img):

    img = enhance(img)

    sigmas = np.linspace(1.0, 6.0, 10)

    responses = []
    for s in sigmas:
        responses.append(-gaussian_laplace(img, sigma=s))

    stack = np.stack(responses, axis=-1)

    scale_map = np.argmax(stack, axis=-1)

    dominant_scale = np.median(scale_map)

    # heuristic separation
    if dominant_scale < 3.5:
        return "S2"
    else:
        return "SF9"


# =========================
# 🔬 ROI COUNT (LoG + peak)
# =========================
def count_roi(roi, cell_type):

    img = enhance(roi)

    if cell_type == "S2":
        sigmas = np.linspace(1.0, 2.8, 6)
        min_dist = 3
        thr_percentile = 97
    else:
        sigmas = np.linspace(2.5, 5.5, 6)
        min_dist = 5
        thr_percentile = 96

    responses = []
    for s in sigmas:
        responses.append(-gaussian_laplace(img, sigma=s))

    stack = np.stack(responses, axis=-1)
    resp = np.max(stack, axis=-1)

    thr = np.percentile(resp, thr_percentile)

    coords = peak_local_max(
        resp,
        min_distance=min_dist,
        threshold_abs=thr
    )

    return len(coords), coords


# =========================
# MAIN
# =========================
summary = []
roi_table = []

for fname in os.listdir(image_folder):

    if not fname.lower().endswith((".jpg", ".png", ".tif")):
        continue

    print("\nProcessing:", fname)

    img = cv2.imread(os.path.join(image_folder, fname), 0)

    cell_type = detect_type(img)
    print("Detected:", cell_type)

    h, w = img.shape

    fig, ax = plt.subplots(figsize=(6,6))
    ax.imshow(img, cmap="gray")

    counts = []

    for i in range(N_ROI):

        y = np.random.randint(ROI_SIZE, h-ROI_SIZE)
        x = np.random.randint(ROI_SIZE, w-ROI_SIZE)

        roi = img[y-ROI_SIZE:y+ROI_SIZE, x-ROI_SIZE:x+ROI_SIZE]

        c, coords = count_roi(roi, cell_type)
        counts.append(c)

        roi_table.append({
            "Image": fname,
            "ROI_ID": i,
            "Type": cell_type,
            "Count": c
        })

        # ROI box
        ax.add_patch(plt.Rectangle(
            (x-ROI_SIZE, y-ROI_SIZE),
            ROI_SIZE*2,
            ROI_SIZE*2,
            edgecolor="red",
            fill=False
        ))

        ax.text(x-ROI_SIZE, y-ROI_SIZE, str(c), color="yellow")

        # center visualization
        for cy, cx in coords:
            ax.add_patch(plt.Circle(
                (x-ROI_SIZE+cx, y-ROI_SIZE+cy),
                radius=3,
                color="lime",
                alpha=0.4
            ))

    mean = np.mean(counts)
    std = np.std(counts)

    summary.append({
        "Image": fname,
        "Type": cell_type,
        "Mean": mean,
        "Std": std
    })

    ax.set_title(f"{fname} | {cell_type} | mean={mean:.2f}")
    ax.axis("off")

    plt.savefig(os.path.join(output_folder, fname + "_QC.png"), dpi=300)
    plt.close()

# =========================
# OUTPUT FILES
# =========================
pd.DataFrame(summary).to_csv(
    os.path.join(output_folder, "summary.csv"),
    index=False
)

pd.DataFrame(roi_table).to_csv(
    os.path.join(output_folder, "roi_details.csv"),
    index=False
)

print("\n✅ DONE — STABLE PIPELINE READY")