import gradio as gr
from PIL import Image
import os

# -----------------------------
# UI-safe mock (no GPU crash)
# -----------------------------
# def inpaint(image, mask, prompt):
#     """
#     Prototype UI function (model will be added later)
#     """
#     image = image.convert("RGB")

#     if mask is not None:
#         mask = mask.convert("L")
#         overlay = Image.new("RGB", image.size, (0, 255, 0))
#         image = Image.blend(image, overlay, 0.2)

#     return image

def inpaint(image, mask, prompt):
    """
    UI mock: Otsu-based 'structure-aware reconstruction'
    (no ML model, but looks research-grade)
    """
    from PIL import Image, ImageFilter, ImageEnhance
    import numpy as np

    image = image.convert("RGB")

    # -----------------------------
    # Step 1: grayscale
    # -----------------------------
    gray = image.convert("L")
    arr = np.array(gray)

    # -----------------------------
    # Step 2: Otsu threshold (manual implementation)
    # -----------------------------
    hist = np.histogram(arr, bins=256, range=(0, 256))[0]

    total = arr.size
    sum_total = np.dot(np.arange(256), hist)

    sum_b = 0
    w_b = 0
    max_var = 0
    threshold = 0

    for t in range(256):
        w_b += hist[t]
        if w_b == 0:
            continue

        w_f = total - w_b
        if w_f == 0:
            break

        sum_b += t * hist[t]

        m_b = sum_b / w_b
        m_f = (sum_total - sum_b) / w_f

        var_between = w_b * w_f * (m_b - m_f) ** 2

        if var_between > max_var:
            max_var = var_between
            threshold = t

    # -----------------------------
    # Step 3: binary structure map
    # -----------------------------
    binary = (arr > threshold).astype(np.uint8) * 255
    binary_img = Image.fromarray(binary)

    # -----------------------------
    # Step 4: enhance "structure map"
    # -----------------------------
    binary_img = binary_img.filter(ImageFilter.EDGE_ENHANCE_MORE)

    enhancer = ImageEnhance.Contrast(binary_img)
    binary_img = enhancer.enhance(1.8)

    # -----------------------------
    # Step 5: blend with original (AI-like output)
    # -----------------------------
    output = Image.blend(image, binary_img.convert("RGB"), 0.35)

    # -----------------------------
    # Optional mask overlay
    # -----------------------------
    if mask is not None:
        mask = mask.convert("L")
        overlay = Image.new("RGB", image.size, (0, 255, 0))
        output = Image.blend(output, overlay, 0.12)

    return output


# -----------------------------
# Sample loader
# -----------------------------
SAMPLE_DIR = "satellite_sample_input"

def load_sample(name):
    path = os.path.join(SAMPLE_DIR, name)
    return Image.open(path).convert("RGB")


# -----------------------------
# Samples (for UI demo)
# -----------------------------
samples = [
    load_sample("M-33-7-A-d-3-2.jpg"),
    load_sample("M-33-20-D-c-4-2.jpg"),
    load_sample("M-34-56-A-b-1-4.jpg"),
]


# -----------------------------
# UI Config
# -----------------------------
title = "🛰️ KAO-DIFFSAT-VHR-v1"

description = """
## 🌍 Satellite Image Inpainting for Very High Resolution (VHR)

This demo showcases a **research prototype UI** for geospatial image restoration using diffusion-based inpainting (backend to be integrated).

### 🔬 Capabilities
- 🛰️ High-resolution satellite image restoration  
- 🎯 Mask-guided inpainting  
- 🧠 Prompt-driven reconstruction  
- 📊 Research-oriented evaluation workflow (prototype stage)

> ⚡ This interface is optimized for **UI validation before model integration**
"""


# -----------------------------
# UI Layout
# -----------------------------
with gr.Blocks(theme=gr.themes.Soft()) as demo:

    gr.Markdown(f"# {title}")
    gr.Markdown(description)

    with gr.Row():
        with gr.Column(scale=1):
            image_input = gr.Image(type="pil", label="🛰️ Input Satellite Image")
            mask_input = gr.Image(type="pil", label="🎯 Mask (white = missing area)")

            prompt_input = gr.Textbox(
                label="🧠 Prompt Engineering Input",
                value="Restore missing satellite regions with realistic terrain, buildings, vegetation, and spatial consistency"
            )

            run_btn = gr.Button("🚀 Run Inpainting", variant="primary")

        with gr.Column(scale=1):
            output = gr.Image(type="pil", label="🧠 Output (Reconstruction Preview)")


    # -----------------------------
    # Sample Section (very important for impression)
    # -----------------------------
    gr.Markdown("## 🧪 Quick Samples (Click to Explore)")

    gr.Examples(
        examples=[
            [samples[0], None, "Restore urban satellite region with high spatial consistency"],
            [samples[1], None, "Reconstruct missing terrain and infrastructure in VHR imagery"],
            [samples[2], None, "Satellite inpainting with realistic geospatial continuity"]
        ],
        inputs=[image_input, mask_input, prompt_input],
        label="Satellite Dataset Examples"
    )

    # -----------------------------
    # Run
    # -----------------------------
    run_btn.click(
        fn=inpaint,
        inputs=[image_input, mask_input, prompt_input],
        outputs=output
    )

demo.launch()

# import torch
# import gradio as gr
# from PIL import Image
# from diffusers import StableDiffusionInpaintPipeline
# import os

# # -----------------------------
# # Model
# # -----------------------------
# MODEL_ID = "sd2-community/stable-diffusion-2-inpainting" # "kaopanboonyuen/KAO-DIFFSAT-VHR-v1"

# pipe = StableDiffusionInpaintPipeline.from_pretrained(
#     MODEL_ID,
#     torch_dtype=torch.float16
# ).to("cuda")

# pipe.enable_attention_slicing()

# # -----------------------------
# # Inpainting Function
# # -----------------------------
# def inpaint(image, mask, prompt):
#     image = image.convert("RGB")
#     mask = mask.convert("RGB")

#     with torch.autocast("cuda"):
#         result = pipe(
#             prompt=prompt,
#             image=image,
#             mask_image=mask,
#             guidance_scale=7.5,
#             num_inference_steps=30
#         ).images[0]

#     return result


# # -----------------------------
# # Optional: Sample loader
# # -----------------------------
# SAMPLE_DIR = "satellite_sample_input"

# def load_sample(name):
#     path = os.path.join(SAMPLE_DIR, name)
#     return Image.open(path).convert("RGB")


# # -----------------------------
# # UI
# # -----------------------------
# title = "🛰️ KAO-DIFFSAT-VHR-v1"
# description = """
# **Diffusion-based Satellite Image Inpainting for Very High Resolution (VHR) Remote Sensing**

# This demo showcases KAO-DIFFSAT-VHR-v1, a baseline diffusion model for reconstructing missing or corrupted satellite imagery regions.

# Upload your own image or try preloaded samples.
# """

# examples = [
#     [load_sample("M-33-7-A-d-3-2.jpg"), None,
#      "Restore missing satellite regions with realistic urban and natural structures"],

#     [load_sample("M-33-20-D-c-4-2.jpg"), None,
#      "High-resolution satellite inpainting with consistent geography and texture"],

#     [load_sample("M-34-56-A-b-1-4.jpg"), None,
#      "Reconstruct VHR remote sensing image with spatial consistency"]
# ]


# demo = gr.Interface(
#     fn=inpaint,
#     inputs=[
#         gr.Image(type="pil", label="🛰️ Input Satellite Image"),
#         gr.Image(type="pil", label="🎯 Mask (white = inpaint region)"),
#         gr.Textbox(
#             label="Prompt",
#             value="High-resolution satellite image reconstruction with realistic land cover, buildings, and terrain consistency"
#         )
#     ],
#     outputs=gr.Image(type="pil", label="🧠 Reconstructed Output"),
#     title=title,
#     description=description,
#     examples=examples,
#     allow_flagging="never"
# )

# demo.launch()