import cv2
import numpy as np
import gradio as gr

# ================= FIXED CONFIG =================
RUST_SCORE_THRESH = 0.45

# GrabCut rectangle (tower prior)
GC_LEFT   = 0.20
GC_RIGHT  = 0.80
GC_TOP    = 0.02
GC_BOTTOM = 0.98

# Sky HSV bounds
SKY_LOWER = (80, 20, 40)
SKY_UPPER = (140, 255, 255)

GROUND_FRAC = 0.15
GRABCUT_ITERS = 10
TOWER_MORPH_K = 5
RUST_MORPH_K = 7

# Rust scoring weights
RUST_HUE_CENTER = 15
HUE_WIDTH = 0.20
W_A, W_H, W_S, W_T, W_L = 0.38, 0.22, 0.15, 0.20, 0.05
# ================================================


def remove_sky_mask(img):
    hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
    sky = cv2.inRange(hsv, np.array(SKY_LOWER), np.array(SKY_UPPER))
    return (sky == 0).astype(np.uint8)


def remove_ground_mask(img):
    h, w = img.shape[:2]
    m = np.ones((h, w), np.uint8)
    cut = int(h * GROUND_FRAC)
    m[h - cut:h, :] = 0
    return m


def grabcut_tower_mask(img, non_sky, non_ground):
    h, w = img.shape[:2]

    mask = np.full((h, w), cv2.GC_PR_BGD, np.uint8)
    mask[non_sky == 0] = cv2.GC_BGD
    mask[non_ground == 0] = cv2.GC_BGD

    x1, x2 = int(w * GC_LEFT), int(w * GC_RIGHT)
    y1, y2 = int(h * GC_TOP), int(h * GC_BOTTOM)

    area = mask[y1:y2, x1:x2]
    area[non_ground[y1:y2, x1:x2] == 1] = cv2.GC_PR_FGD
    mask[y1:y2, x1:x2] = area

    bgdModel = np.zeros((1, 65), np.float64)
    fgdModel = np.zeros((1, 65), np.float64)

    cv2.grabCut(img, mask, None, bgdModel, fgdModel, GRABCUT_ITERS, cv2.GC_INIT_WITH_MASK)

    fg = ((mask == cv2.GC_FGD) | (mask == cv2.GC_PR_FGD)).astype(np.uint8)

    k = np.ones((TOWER_MORPH_K, TOWER_MORPH_K), np.uint8)
    fg = cv2.morphologyEx(fg, cv2.MORPH_CLOSE, k, iterations=2)
    fg = cv2.morphologyEx(fg, cv2.MORPH_OPEN, k, iterations=1)

    return fg


def rust_score_map(img):
    lab = cv2.cvtColor(img, cv2.COLOR_BGR2LAB).astype(np.float32)
    L, A = lab[:, :, 0], lab[:, :, 1]

    A_n = (A - A.min()) / (A.max() - A.min() + 1e-6)
    L_n = (L - L.min()) / (L.max() - L.min() + 1e-6)

    hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV).astype(np.float32)
    H, S = hsv[:, :, 0] / 179.0, hsv[:, :, 1] / 255.0

    rust_center = RUST_HUE_CENTER / 179.0
    hue_score = 1.0 - np.clip(np.abs(H - rust_center) / HUE_WIDTH, 0, 1)

    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    tex = np.abs(cv2.Laplacian(gray, cv2.CV_32F))
    tex_n = (tex - tex.min()) / (tex.max() - tex.min() + 1e-6)

    score = (
        W_A * A_n +
        W_H * hue_score +
        W_S * S +
        W_T * tex_n +
        W_L * (1.0 - np.abs(L_n - 0.55))
    )

    return np.clip(score, 0, 1)


def postprocess(m):
    k = np.ones((RUST_MORPH_K, RUST_MORPH_K), np.uint8)
    m = cv2.morphologyEx(m, cv2.MORPH_CLOSE, k, iterations=2)
    m = cv2.morphologyEx(m, cv2.MORPH_OPEN, k, iterations=1)
    return m


def run(img_rgb):
    if img_rgb is None:
        return None, None, None, None, "Upload an image."

    img = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2BGR)

    non_sky = remove_sky_mask(img)
    non_ground = remove_ground_mask(img)
    tower = grabcut_tower_mask(img, non_sky, non_ground)

    roi = tower * non_sky * non_ground
    score = rust_score_map(img)

    rust = ((score >= RUST_SCORE_THRESH).astype(np.uint8)) * roi
    rust = postprocess(rust)

    overlay = img.copy()
    overlay[rust == 1] = [0, 0, 255]

    rust_px = int(rust.sum())
    roi_px = int(roi.sum())
    pct = 100 * rust_px / max(1, roi_px)

    stats = (
        f"Rust threshold: {RUST_SCORE_THRESH}\n"
        f"Rust pixels: {rust_px}\n"
        f"ROI pixels: {roi_px}\n"
        f"Rust on tower ROI: {pct:.2f}%"
    )

    return (
        cv2.cvtColor(tower * 255, cv2.COLOR_GRAY2RGB),
        cv2.cvtColor(rust * 255, cv2.COLOR_GRAY2RGB),
        cv2.cvtColor(overlay, cv2.COLOR_BGR2RGB),
        cv2.applyColorMap((score * 255).astype(np.uint8), cv2.COLORMAP_TURBO),
        stats
    )


with gr.Blocks(title="Tower Rust Detection") as demo:
    gr.Markdown("## Tower Rust Detection \nUpload an image and click **Run**.")

    inp = gr.Image(type="numpy", label="Input Image")
    run_btn = gr.Button("Run", variant="primary")

    with gr.Row():
        out_tower = gr.Image(label="Tower Mask")
        out_rust = gr.Image(label="Rust Mask")
    with gr.Row():
        out_overlay = gr.Image(label="Overlay")
        out_heat = gr.Image(label="Rust Score Heatmap")

    stats = gr.Textbox(label="Stats", lines=5)

    run_btn.click(run, inputs=inp, outputs=[out_tower, out_rust, out_overlay, out_heat, stats])

if __name__ == "__main__":
    demo.launch()