import gradio as gr
import numpy as np
from scipy.ndimage import uniform_filter, median_filter, gaussian_filter1d
from skimage import exposure
import matplotlib.pyplot as plt
import io
from PIL import Image

# ---------- util ----------
def to_gray2d(arr):
    arr = np.array(arr, dtype=np.float32)
    if arr.ndim == 3 and arr.shape[2] in (3,4):
        arr = arr[..., :3].mean(axis=2)
    elif arr.ndim > 2:
        arr = np.squeeze(arr)
    if arr.max() > 1.0:
        arr = arr / 255.0
    return arr.astype(np.float32)

# ---------- pipeline ----------
def splitnormalize_balanced(img2d):
    p_lo, p_hi = np.percentile(img2d, [2, 98])
    x = np.clip((img2d - p_lo) / (p_hi - p_lo + 1e-6), 0, 1)
    x = exposure.equalize_adapthist(x, clip_limit=0.003, nbins=256)
    return x.astype(np.float32)

def removebias(img2d, regionwidth=10, fraction=0.75):
    kz = 2*regionwidth + 1
    kx = 2*regionwidth + 1
    bg = uniform_filter(img2d, size=(kz, kx))
    out = img2d - fraction * bg
    out = (out - out.min()) / (out.max() - out.min() + 1e-6)
    return out.astype(np.float32)

def estimate_medline_intensity(img2d):
    H, W = img2d.shape
    p1, p99 = np.percentile(img2d, [1, 99])
    norm = np.clip((img2d - p1) / (p99 - p1 + 1e-6), 0, 1)
    z_low, z_high = int(0.30*H), int(0.70*H)
    return (np.argmax(norm[z_low:z_high, :], axis=0) + z_low).astype(np.float32)

def detect_rpe_simple(img2d, medline):
    H, W = img2d.shape
    sm = gaussian_filter1d(img2d, sigma=4, axis=0)
    grad = np.gradient(gaussian_filter1d(sm, sigma=1, axis=0), axis=0)
    z0, z1 = int(0.30*H), int(0.85*H)
    return (np.argmax(-grad[z0:z1, :], axis=0) + z0).astype(np.float32)

def linesweeter(y):
    from scipy.signal import savgol_filter
    return savgol_filter(y, 9, 2).astype(np.float32)

def overlay(img2d, curve, title=""):
    H, W = img2d.shape
    x = np.arange(W)
    fig, ax = plt.subplots(figsize=(8,4))
    ax.imshow(img2d, cmap="gray")
    ax.plot(x, curve, 'r-', lw=2)
    ax.set_title(title); ax.axis('off')
    buf = io.BytesIO(); plt.savefig(buf, format="png", bbox_inches="tight"); buf.seek(0)
    return Image.open(buf)

# ---------- UI ----------
with gr.Blocks(title="OCT Step-by-Step Visual Lab") as demo:
    gr.Markdown("## 🧠 OCT Step-by-Step Visual Lab — Compare cada etapa lado a lado")

    img_state = gr.State()  # guarda a imagem 2D corrente (float [0,1])

    with gr.Tab("1) Carregar Imagem"):
        img_input = gr.Image(label="Imagem OCT", type="numpy")
        def store(img):
            return to_gray2d(img) if img is not None else None
        img_input.change(store, inputs=img_input, outputs=img_state)

    with gr.Tab("2) Normalização"):
        btn_norm = gr.Button("Aplicar splitnormalize (balanced)")
        before_norm = gr.Image(label="Antes")
        after_norm  = gr.Image(label="Depois")
        def do_norm(img2d):
            img2d = to_gray2d(img2d)
            out = splitnormalize_balanced(img2d)
            return img2d, out, out  # antes, depois, novo estado
        btn_norm.click(do_norm, inputs=img_state, outputs=[before_norm, after_norm, img_state])

    with gr.Tab("3) Remove Bias"):
        btn_bias = gr.Button("Aplicar removebias")
        before_b = gr.Image(label="Antes")
        after_b  = gr.Image(label="Depois")
        def do_bias(img2d):
            img2d = to_gray2d(img2d)
            out = removebias(img2d)
            return img2d, out, out
        btn_bias.click(do_bias, inputs=img_state, outputs=[before_b, after_b, img_state])

    with gr.Tab("4) Mediana 5×9"):
        btn_med = gr.Button("Aplicar filtro mediano (5x9)")
        before_m = gr.Image(label="Antes")
        after_m  = gr.Image(label="Depois")
        def do_med(img2d):
            img2d = to_gray2d(img2d)
            out = median_filter(img2d, size=(5,9)).astype(np.float32)
            return img2d, out, out
        btn_med.click(do_med, inputs=img_state, outputs=[before_m, after_m, img_state])

    with gr.Tab("5) Estimar Medline (IS/OS)"):
        btn_medline = gr.Button("Calcular Medline")
        out_medline = gr.Image(label="Visualização")
        def show_medline(img2d):
            img2d = to_gray2d(img2d)
            med = estimate_medline_intensity(img2d)
            return overlay(img2d, med, "Medline (IS/OS)")
        btn_medline.click(show_medline, inputs=img_state, outputs=out_medline)

    with gr.Tab("6) Detectar RPE"):
        btn_rpe = gr.Button("Detectar RPE simples")
        out_rpe = gr.Image(label="Visualização")
        def step_rpe(img2d):
            img2d = to_gray2d(img2d)
            med = estimate_medline_intensity(img2d)
            rpe = detect_rpe_simple(img2d, med)
            return overlay(img2d, rpe, "RPE simples")
        btn_rpe.click(step_rpe, inputs=img_state, outputs=out_rpe)

    with gr.Tab("7) Suavização da RPE"):
        btn_smooth = gr.Button("Suavizar RPE (linesweeter)")
        out_smooth = gr.Image(label="Visualização")
        def step_smooth(img2d):
            img2d = to_gray2d(img2d)
            med = estimate_medline_intensity(img2d)
            rpe = detect_rpe_simple(img2d, med)
            rpe_s = linesweeter(rpe)
            return overlay(img2d, rpe_s, "RPE suavizada (linesweeter)")
        btn_smooth.click(step_smooth, inputs=img_state, outputs=out_smooth)

demo.launch()