import gradio as gr
import numpy as np
from PIL import Image
import cv2
from sklearn.cluster import KMeans
from scipy import ndimage as ndi
import tempfile

# ===============================
# ابزارهای پایه
# ===============================
def pil_to_np_rgb(pil_image: Image.Image) -> np.ndarray:
    return np.array(pil_image.convert("RGB"))

def save_as_bmp(np_img: np.ndarray, name_prefix: str) -> str:
    f = tempfile.NamedTemporaryFile(delete=False, suffix=".bmp", prefix=f"{name_prefix}_")
    Image.fromarray(np_img).save(f.name, format="BMP")
    return f.name

# ===============================
# KMeans: کاهش رنگ‌ها (15..30)
# ===============================
def quantize_colors_kmeans(img_rgb: np.ndarray, n_clusters: int):
    h, w, _ = img_rgb.shape
    pixels = img_rgb.reshape(-1, 3)
    kmeans = KMeans(n_clusters=n_clusters, n_init=10, random_state=42).fit(pixels)
    labels = kmeans.labels_.reshape(h, w)
    palette = kmeans.cluster_centers_.astype(np.uint8)

    reduced = np.zeros_like(img_rgb)
    for i in range(n_clusters):
        reduced[labels == i] = palette[i]
    return labels, palette, reduced

# ===============================
# حذف ریزناحیه‌ها (ادغام با همسایه غالب)
# ===============================
def clean_small_regions(labels: np.ndarray, min_area_px: int, connectivity: int = 8) -> np.ndarray:
    if min_area_px <= 0:
        return labels
    labels = labels.copy()
    kernel = np.ones((3, 3), np.uint8)
    for lbl in np.unique(labels):
        mask = (labels == lbl).astype(np.uint8)
        num, comp_map, stats, _ = cv2.connectedComponentsWithStats(mask, connectivity=connectivity)
        for comp_id in range(1, num):
            area = int(stats[comp_id, cv2.CC_STAT_AREA])
            if area >= min_area_px:
                continue
            comp_mask = (comp_map == comp_id)

            dil = cv2.dilate(comp_mask.astype(np.uint8), kernel, iterations=1).astype(bool)
            border = dil & (~comp_mask)
            neighbor_labels = labels[border]
            neighbor_labels = neighbor_labels[neighbor_labels != lbl]
            if neighbor_labels.size == 0:
                dil2 = cv2.dilate(comp_mask.astype(np.uint8), kernel, iterations=2).astype(bool)
                border2 = dil2 & (~comp_mask)
                neighbor_labels = labels[border2]
                neighbor_labels = neighbor_labels[neighbor_labels != lbl]
                if neighbor_labels.size == 0:
                    continue

            new_lbl = int(np.bincount(neighbor_labels).argmax())
            labels[comp_mask] = new_lbl
    return labels

# ===============================
# نرم‌سازی برچسب‌ها با «اکثریت همسایه‌ها»
# ===============================
def smooth_labels_majority(labels: np.ndarray, radius: int) -> np.ndarray:
    if radius <= 0:
        return labels
    size = 2 * radius + 1
    def mode_func(window):
        w = window.astype(np.int32)
        return np.bincount(w).argmax()
    return ndi.generic_filter(labels, mode_func, size=size, mode="nearest")

# ===============================
# مرز جهانی 1px از روی لیبل‌ها
# ===============================
def build_global_boundary_mask(labels: np.ndarray) -> np.ndarray:
    H, W = labels.shape
    boundary = np.zeros((H, W), dtype=np.uint8)
    h_diff = labels[:, 1:] != labels[:, :-1]
    v_diff = labels[1:, :] != labels[:-1, :]
    boundary[:, :-1] |= h_diff
    boundary[:,  1:] |= h_diff
    boundary[:-1, :] |= v_diff
    boundary[ 1:, :] |= v_diff
    return (boundary.astype(np.uint8) * 255)

# ===============================
# ابزار کانتور + Resample + پایین‌گذر فوریه
# ===============================
def dedupe_consecutive_points(cnt: np.ndarray):
    cnt = cnt.squeeze()
    if cnt.ndim != 2 or len(cnt) < 5:
        return None
    keep = np.ones(len(cnt), dtype=bool)
    keep[1:] = np.any(np.diff(cnt, axis=0) != 0, axis=1)
    cnt = cnt[keep]
    if len(cnt) < 5:
        return None
    return cnt

def resample_by_arclength(points: np.ndarray, step: float = 1.0, closed: bool = True) -> np.ndarray:
    pts = points.astype(np.float64)
    if closed and not np.array_equal(pts[0], pts[-1]):
        pts = np.vstack([pts, pts[0]])
    seg = np.sqrt(((pts[1:] - pts[:-1]) ** 2).sum(axis=1))
    s = np.hstack([[0.0], np.cumsum(seg)])
    total = s[-1]
    if total < step:
        return pts.astype(np.float32)
    n_new = int(np.floor(total / step))
    s_new = np.linspace(0, total, n_new, endpoint=False)
    x = np.interp(s_new, s, pts[:, 0])
    y = np.interp(s_new, s, pts[:, 1])
    out = np.vstack([x, y]).T
    return out.astype(np.float32)

def fourier_lowpass(points: np.ndarray, min_wavelength_px: int, upsample_factor: float = 2.0) -> np.ndarray:
    pts = points
    N = len(pts)
    if N < 8:
        return pts
    # سیگنال مختلط
    z = pts[:, 0].astype(np.float64) + 1j * pts[:, 1].astype(np.float64)
    Z = np.fft.fft(z)
    k_max = max(1, int(np.floor(N / max(1, min_wavelength_px))))
    Z_lp = np.zeros_like(Z)
    Z_lp[:k_max + 1] = Z[:k_max + 1]
    if k_max > 0:
        Z_lp[-k_max:] = Z[-k_max:]
    z_s = np.fft.ifft(Z_lp)
    xs, ys = np.real(z_s), np.imag(z_s)

    # بازنمونه‌گیری نرم‌تر (upsample)
    idx = np.arange(N)
    N_new = max(N, int(N * upsample_factor))
    idx_new = np.linspace(0, N - 1, N_new)
    x_new = np.interp(idx_new, idx, xs)
    y_new = np.interp(idx_new, idx, ys)
    return np.vstack([x_new, y_new]).T.astype(np.float32)

# ===============================
# رسم مرزهای نرم بدون اورلپ (روی خود ماسک مرز)
# ===============================
def render_boundaries_smooth_no_overlap(labels: np.ndarray,
                                        palette: np.ndarray,
                                        boundary_mask: np.ndarray,
                                        min_wavelength_px: int,
                                        arclen_step: float = 1.0,
                                        upsample_factor: float = 2.0) -> np.ndarray:
    H, W = labels.shape
    canvas = np.zeros((H, W, 3), dtype=np.uint8)
    occupied = np.zeros((H, W), dtype=bool)

    cnts_data = cv2.findContours(boundary_mask, cv2.RETR_LIST, cv2.CHAIN_APPROX_NONE)
    contours = cnts_data[0] if len(cnts_data) == 2 else cnts_data[1]
    kernel = np.ones((3, 3), np.uint8)

    for cnt in contours:
        cnt_xy = dedupe_consecutive_points(cnt)
        if cnt_xy is None or len(cnt_xy) < 12:
            continue

        # بازنمونه‌گیری به گام ثابت طول‌قوس + پایین‌گذر فوریه
        rs = resample_by_arclength(cnt_xy, step=arclen_step, closed=True)
        smooth = fourier_lowpass(rs, min_wavelength_px=min_wavelength_px, upsample_factor=upsample_factor)

        # رستر کردن روی ماسک موقت (خط 1px)
        temp = np.zeros((H, W), dtype=np.uint8)
        smooth_i = np.round(smooth).astype(np.int32)
        # ensure closed
        if not (smooth_i[0] == smooth_i[-1]).all():
            smooth_i = np.vstack([smooth_i, smooth_i[0]])
        cv2.polylines(temp, [smooth_i], isClosed=True, color=255, thickness=1, lineType=cv2.LINE_8)

        # فقط پیکسل‌های مرزی اجازه‌ی رسم دارند (کلاپ روی مرز)
        write_mask = (temp > 0) & (boundary_mask > 0) & (~occupied)
        if not np.any(write_mask):
            continue

        # تعیین رنگ هر نقطه‌ی مرز از اکثریت همسایه‌ها
        ring = cv2.dilate((write_mask).astype(np.uint8), kernel, iterations=1).astype(bool)
        ring = ring & (~write_mask)
        neigh_lbls = labels[ring]
        if neigh_lbls.size == 0:
            # fallback: از خود مرز، یک رشد کوچک
            ring2 = cv2.dilate((write_mask).astype(np.uint8), kernel, iterations=2).astype(bool)
            ring2 = ring2 & (~write_mask)
            neigh_lbls = labels[ring2]
            if neigh_lbls.size == 0:
                continue
        owner_lbl = int(np.bincount(neigh_lbls).argmax())
        color = tuple(int(v) for v in palette[owner_lbl])

        ys, xs = np.where(write_mask)
        canvas[ys, xs] = color
        occupied[ys, xs] = True

    return canvas

# ===============================
# ترکیب نهایی بدون تولید رنگ جدید
# ===============================
def overlay_contours_on_reduced(reduced_rgb: np.ndarray, contours_rgb: np.ndarray) -> np.ndarray:
    final = reduced_rgb.copy()
    mask = np.any(contours_rgb != 0, axis=2)
    final[mask] = contours_rgb[mask]
    return final

# ===============================
# پایپلاین اصلی Gradio
# ===============================
def process_pipeline(pil_image: Image.Image,
                     n_clusters: int,
                     min_area_px: int,
                     label_smooth_r: int,
                     min_wavelength_px: int,
                     arc_step: float,
                     upsample_factor: float):
    original = pil_to_np_rgb(pil_image)

    # 1) کوانتیزه
    labels_raw, palette, _ = quantize_colors_kmeans(original, n_clusters=n_clusters)

    # 2) حذف ریزناحیه‌ها
    labels = clean_small_regions(labels_raw, min_area_px=min_area_px, connectivity=8)

    # 3) نرم‌سازی برچسب‌ها با اکثریت همسایه‌ها (کاهش زیگزاگ پیکسلی)
    labels = smooth_labels_majority(labels, radius=label_smooth_r)

    # 4) یک پاسِ دیگر حذف ریزناحیه‌ها بعد از نرم‌سازی اختیاری
    labels = clean_small_regions(labels, min_area_px=min_area_px, connectivity=8)

    # 5) بازسازی تصویر کاهش‌یافته از روی لیبل‌های تمیز
    reduced = np.zeros_like(original)
    for i in range(len(palette)):
        reduced[labels == i] = palette[i]

    # 6) مرز 1px سراسری
    boundary_mask = build_global_boundary_mask(labels)

    # 7) مرزهای نرم، یکپارچه، بدون اورلپ، فقط روی مرز
    contours_img = render_boundaries_smooth_no_overlap(
        labels=labels,
        palette=palette,
        boundary_mask=boundary_mask,
        min_wavelength_px=min_wavelength_px,
        arclen_step=arc_step,
        upsample_factor=upsample_factor
    )

    # 8) خروجی نهایی (نواحی رنگ + مرز 1px)
    final_img = overlay_contours_on_reduced(reduced, contours_img)

    # BMP برای دانلود
    p_ori = save_as_bmp(original, "original")
    p_red = save_as_bmp(reduced,  "reduced_clean")
    p_cnt = save_as_bmp(contours_img, "contours_smooth_no_overlap")
    p_fin = save_as_bmp(final_img,"final")

    return (
        Image.fromarray(original), Image.fromarray(reduced),
        Image.fromarray(contours_img), Image.fromarray(final_img),
        p_ori, p_red, p_cnt, p_fin
    )

# ===============================
# UI با اسکرول + دانلود BMP
# ===============================
CSS = """
.scroll-pane { max-height: 720px; overflow: auto; }
.scroll-pane img, .scroll-pane canvas { width: auto !important; max-width: none !important; }
"""

with gr.Blocks(css=CSS) as demo:
    gr.Markdown("## 🎨 مرزهای نرمِ پیوسته (Fourier + Arc-length) بدون اورلپ + حذف ریزناحیه‌ها (BMP)")

    with gr.Row():
        inp = gr.Image(label="آپلود تصویر", type="pil")

    with gr.Row():
        n_clusters = gr.Slider(15, 30, value=20, step=1, label="تعداد کلاسترهای رنگ (15–30)")
        min_area = gr.Slider(0, 20000, value=1500, step=100, label="حداقل مساحت ناحیه (پیکسل) – حذف ریزناحیه‌ها")
        label_smooth_r = gr.Slider(0, 5, value=2, step=1, label="نرم‌سازی برچسب‌ها (شعاع اکثریت همسایه)")

    with gr.Row():
        min_wave = gr.Slider(10, 400, value=140, step=5, label="حداقل طول‌موج مرز (پیکسل) – صاف‌سازی فوریه")
        arc_step = gr.Slider(0.5, 3.0, value=1.0, step=0.5, label="گام بازنمونه‌گیری طول‌قوس (پیکسل)")
        upsample = gr.Slider(1.0, 4.0, value=2.0, step=0.5, label="Upsample منحنی پس از فیلتر")

    btn = gr.Button("🚀 پردازش")

    with gr.Row():
        out_ori = gr.Image(label="۱) تصویر اصلی", elem_classes=["scroll-pane"], height=600)
        out_red = gr.Image(label="۲) تصویر کاهش‌یافته پس از پاک‌سازی", elem_classes=["scroll-pane"], height=600)
    with gr.Row():
        out_cnt = gr.Image(label="۳) مرزهای نرمِ پیوسته (۱px، بدون اورلپ)", elem_classes=["scroll-pane"], height=600)
        out_fin = gr.Image(label="۴) خروجی نهایی (نواحی رنگ‌شده + مرز ۱px)", elem_classes=["scroll-pane"], height=600)

    with gr.Accordion("دانلود فایل‌های BMP", open=True):
        with gr.Row():
            file_ori = gr.File(label="دانلود: تصویر اصلی (BMP)")
            file_red = gr.File(label="دانلود: تصویر کاهش‌یافته پاک‌سازی‌شده (BMP)")
        with gr.Row():
            file_cnt = gr.File(label="دانلود: مرزهای نرمِ پیوسته (BMP)")
            file_fin = gr.File(label="دانلود: خروجی نهایی (BMP)")

    btn.click(
        process_pipeline,
        inputs=[inp, n_clusters, min_area, label_smooth_r, min_wave, arc_step, upsample],
        outputs=[out_ori, out_red, out_cnt, out_fin, file_ori, file_red, file_cnt, file_fin]
    )

demo.launch()