"""FlashReFocus: interactive image refocusing (Hugging Face Space, ZeroGPU).

1. Restore: Stage I (one-step diffusion with detail weight omega and the decoder refiner) restores the all-in-focus
   image once, and Depth Anything V2 estimates its depth once.
2. Click the restored image (or the depth map) to focus; only the 1.3M-parameter renderer runs per edit.
   Modes: one focal plane; two focal planes (first click A, second click B); tilt-shift (a click sets the sharp row).
3. The aperture (blur strength K), the depth curve and an optional user depth map re-render the current edit.
"""
import time

import gradio as gr
import numpy as np
import spaces
from PIL import Image, ImageDraw, ImageFont

import engine

SINGLE, DUAL, TILT = "📍 One focal plane", "👥 Two focal planes", "🏙️ Tilt-shift"
EXAMPLES = ["examples/portrait_city.jpg", "examples/portrait_aquarium.jpg", "examples/flower.jpg",
            "examples/dock.jpg", "examples/sculpture.jpg"]
FONT = "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf"
try:
    ImageFont.truetype(FONT, 12)
except OSError:
    FONT = None


def fit(img, long_side):
    """Long side at most `long_side`, both sides multiples of 16."""
    w, h = img.size
    s = min(1.0, long_side / max(w, h))
    w2, h2 = max(16, int(w * s) // 16 * 16), max(16, int(h * s) // 16 * 16)
    return img.convert("RGB").resize((w2, h2), Image.LANCZOS)


@spaces.GPU(duration=40)
def gpu_restore(rgb, omega):
    t0 = time.perf_counter()
    aif, depth = engine.restore(rgb, omega)
    return aif, depth, time.perf_counter() - t0


@spaces.GPU(duration=10)
def gpu_render(aif, depth, focal_depth, strength):
    t0 = time.perf_counter()
    out = engine.render(aif, depth, focal_depth, strength)
    return out, time.perf_counter() - t0


def depth_image(depth):
    return Image.fromarray((np.clip(1 - depth, 0, 1) * 255).astype(np.uint8))   # near is bright


def scene_depth(state, gamma):
    """The depth the edits start from: the estimate (or the user's map), with the depth curve D^gamma applied."""
    d = state["user_depth"] if state.get("user_depth") is not None else state["depth"]
    return np.power(np.clip(d, 0, 1), float(gamma)).astype(np.float32)


def depth_at(depth, point):
    h, w = depth.shape
    return float(depth[min(max(point[1], 0), h - 1), min(max(point[0], 0), w - 1)])


def edit(state, mode, gamma, band):
    """(depth to render, focal depth D_f) for the current mode, or None when a focal point is still missing."""
    d = scene_depth(state, gamma)
    h, w = d.shape
    if mode == TILT:
        return engine.tilt_shift(h, w, state.get("tilt_row", 0.5), float(band)), 0.5
    pts = state["pts"]
    if not pts:
        return None
    if mode == DUAL and len(pts) == 2:
        da, db = depth_at(d, pts[0]), depth_at(d, pts[1])
        return engine.two_focal_planes(d, da, db), da
    return d, depth_at(d, pts[0])


def restore(img, long_side, omega, gamma):
    if img is None:
        return None, None, None, None, "Upload an image or pick an example first."
    rgb = np.asarray(fit(img, int(long_side)))
    aif, depth, t = gpu_restore(rgb, float(omega))
    state = {"aif": aif, "depth": depth, "user_depth": None, "pts": [], "tilt_row": 0.5}
    msg = (f"Restored {aif.shape[1]}×{aif.shape[0]} with ω = {float(omega):.2f} and estimated its depth in {t:.2f} s. "
           "**Click the restored image or the depth map to focus.**")
    return Image.fromarray(aif), depth_image(scene_depth(state, gamma)), None, state, msg


def render(state, strength, mode, gamma, band, show_marker):
    """(refocused image, depth view, message)."""
    edited = edit(state, mode, gamma, band)
    if edited is None:
        return None, depth_image(scene_depth(state, gamma)), "Click the restored image or the depth map to focus."
    depth, focal_depth = edited
    out, dt = gpu_render(state["aif"], depth, focal_depth, float(strength))
    out = Image.fromarray(out)
    if show_marker:
        out = draw_markers(out, state, mode)
    k = f"aperture K = {float(strength):.2f}" + (" (extrapolated)" if float(strength) > 1 else "")
    if mode == TILT:
        msg = f"Tilt-shift band at {state.get('tilt_row', 0.5):.0%} of the height, {k}: rendered in {dt:.2f} s."
    elif mode == DUAL and len(state["pts"]) == 2:
        msg = f"Two focal planes A {state['pts'][0]} and B {state['pts'][1]}, {k}: rendered in {dt:.2f} s."
    elif mode == DUAL:
        msg = f"Plane A {state['pts'][0]} set (rendered in {dt:.2f} s). Click a second point for plane B."
    else:
        msg = f"Focus {state['pts'][0]} (D_f = {focal_depth:.2f}), {k}: rendered in {dt:.2f} s."
    return out, depth_image(depth), msg


def draw_markers(im, state, mode):
    """Focal points (A / B with two planes), or the sharp row of tilt-shift, on a 3x supersampled overlay."""
    im = im.convert("RGBA"); S = 3
    ov = Image.new("RGBA", (im.width * S, im.height * S), (0, 0, 0, 0)); d = ImageDraw.Draw(ov)
    r = max(9, round(min(im.size) * 0.022)) * S; lw = max(2, round(r / 6))
    if mode == TILT:
        y = state.get("tilt_row", 0.5) * ov.height
        for x in range(0, ov.width, 12 * lw):
            d.line([x, y, x + 6 * lw, y], fill=(255, 255, 255, 220), width=lw)
        return Image.alpha_composite(im, ov.resize(im.size, Image.LANCZOS)).convert("RGB")
    font = ImageFont.truetype(FONT, round(r * 1.25)) if FONT else ImageFont.load_default()
    for i, (x, y) in enumerate(state["pts"]):
        cx, cy = x * S, y * S
        d.ellipse([cx - r - lw, cy - r - lw, cx + r + lw, cy + r + lw], outline=(0, 0, 0, 170), width=lw * 3)
        d.ellipse([cx - r, cy - r, cx + r, cy + r], outline=(255, 255, 255, 255), width=lw)
        dot = lw * 1.3
        d.ellipse([cx - dot, cy - dot, cx + dot, cy + dot], fill=(255, 255, 255, 255))
        if mode == DUAL:
            tx = cx + r + lw * 2 if cx + r * 3 < ov.width else cx - r * 2.2
            d.text((tx, max(0, cy - r * 1.9)), "AB"[i], font=font, fill=(255, 255, 255, 255),
                   stroke_width=max(2, lw), stroke_fill=(0, 0, 0, 200))
    return Image.alpha_composite(im, ov.resize(im.size, Image.LANCZOS)).convert("RGB")


def aif_view(state, mode):
    return draw_markers(Image.fromarray(state["aif"]), state, mode)


def on_click(state, strength, mode, gamma, band, show_marker, evt: gr.SelectData):
    if not state:
        return None, None, None, state, "Restore an image first."
    p = (int(evt.index[0]), int(evt.index[1]))
    if mode == TILT:
        state = dict(state, tilt_row=p[1] / state["aif"].shape[0])
    else:
        state = dict(state, pts=state["pts"] + [p] if mode == DUAL and len(state["pts"]) == 1 else [p])
    out, depth_view, msg = render(state, strength, mode, gamma, band, show_marker)
    return aif_view(state, mode), out, depth_view, state, msg


def on_change(state, strength, mode, gamma, band, show_marker):
    if not state:
        return gr.update(), None, gr.update(), state, ""
    if mode != DUAL and len(state["pts"]) > 1:
        state = dict(state, pts=state["pts"][:1])
    out, depth_view, msg = render(state, strength, mode, gamma, band, show_marker)
    return aif_view(state, mode), out, depth_view, state, msg


def on_user_depth(state, user_depth, strength, mode, gamma, band, show_marker):
    if not state:
        return gr.update(), None, gr.update(), state, "Restore an image first, then upload its depth map."
    h, w = state["aif"].shape[:2]
    state = dict(state, user_depth=None if user_depth is None else engine.depth_from_image(user_depth, h, w))
    out, depth_view, msg = render(state, strength, mode, gamma, band, show_marker)
    note = "Using your depth map. " if user_depth is not None else "Using the estimated depth again. "
    return aif_view(state, mode), out, depth_view, state, note + msg


def on_clear(state, gamma):
    if not state:
        return gr.update(), None, gr.update(), state, ""
    state = dict(state, pts=[], tilt_row=0.5)
    return (Image.fromarray(state["aif"]), None, depth_image(scene_depth(state, gamma)), state,
            "Focal points cleared. Click the restored image or the depth map to focus.")


RECOMMENDED = {"long_side": 1024, "omega": 1.25, "gamma": 1.0, "band": 0.06, "strength": 0.5, "mode": SINGLE,
               "show_marker": False}


def on_reset(state):
    """Recommended settings (the paper's defaults), the estimated depth again; focal points are kept."""
    r = RECOMMENDED
    controls = [r["long_side"], r["omega"], r["gamma"], r["band"], r["strength"], r["mode"], r["show_marker"], None]
    if not state:
        return controls + [gr.update(), None, gr.update(), state, "Settings reset to the recommended values."]
    state = dict(state, user_depth=None, pts=state["pts"][:1])
    out, depth_view, msg = render(state, r["strength"], r["mode"], r["gamma"], r["band"], r["show_marker"])
    note = "Settings reset to the recommended values (ω applies on the next Restore). "
    return controls + [aif_view(state, r["mode"]), out, depth_view, state, note + msg]


LOGO = """<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 64 64" width="64" height="64"><circle cx="32" cy="32" r="30" fill="#FFFFFF" stroke="#2E7D32" stroke-width="3"/><g fill="#48A43F" stroke="#FFFFFF" stroke-width="1.6" stroke-linejoin="round"><path d="M32 4 A28 28 0 0 1 56.2 18 L40.6 27 Z" transform="rotate(0 32 32)"/><path d="M32 4 A28 28 0 0 1 56.2 18 L40.6 27 Z" transform="rotate(60 32 32)"/><path d="M32 4 A28 28 0 0 1 56.2 18 L40.6 27 Z" transform="rotate(120 32 32)"/><path d="M32 4 A28 28 0 0 1 56.2 18 L40.6 27 Z" transform="rotate(180 32 32)"/><path d="M32 4 A28 28 0 0 1 56.2 18 L40.6 27 Z" transform="rotate(240 32 32)"/><path d="M32 4 A28 28 0 0 1 56.2 18 L40.6 27 Z" transform="rotate(300 32 32)"/></g><path d="M35.5 19 L25.5 34 H31.5 L28.5 45 L39 29 H33 Z" fill="#FFD21E" stroke="#363636" stroke-width="1.4" stroke-linejoin="round"/></svg>
"""

HEADER = f"""
<div id="header">
  <div class="title">{LOGO}<div><h1>FlashReFocus</h1><p>Interactive image refocusing in seconds: restore a photo once,
  then click to refocus, open the aperture, or keep two subjects sharp.</p></div></div>
  <div class="authors">Ching-Heng Cheng<sup>1*</sup> · <a href="https://ming053l.github.io/" target="_blank">Chia-Ming Lee</a><sup>2,3*</sup> ·
  <a href="https://www.albany.edu/faculty/mchang2/" target="_blank">Ming-Ching Chang</a><sup>2</sup> ·
  <a href="https://scholar.google.com/citations?user=gMBvzGoAAAAJ&hl=en" target="_blank">Xin Li</a><sup>2</sup> ·
  <a href="https://yulunalexliu.github.io/" target="_blank">Yu-Lun Liu</a><sup>3</sup> ·
  <a href="https://cchsu.info/" target="_blank">Chih-Chung Hsu</a><sup>1,3&dagger;</sup><br>
  <small><sup>1</sup>National Cheng Kung University · <sup>2</sup>State University of New York at Albany ·
  <sup>3</sup>National Yang Ming Chiao Tung University · * equal contribution · &dagger; corresponding author</small></div>
  <div class="links">
    <a href="https://ming053l.github.io/FlashReFocus/" target="_blank">🌐 Project page</a>
    <a href="https://ming053l.github.io/FlashReFocus/assets/FlashReFocus.pdf" target="_blank">📄 Paper</a>
    <a href="https://github.com/ming053l/FlashReFocus" target="_blank">💻 Code</a>
    <a href="https://youtu.be/g032M7moWQ4" target="_blank">🎬 Video</a>
    <a href="https://huggingface.co/datasets/ming0531/3CReal" target="_blank">🗂️ 3CReal dataset</a>
  </div>
  <div class="tip">👉 <b>Try this first:</b> pick the aquarium example → <b>Restore</b> → choose <b>👥 Two focal planes</b>
  → click the face, then the shark. Every click after Restore runs only the 1.3M-parameter renderer.</div>
</div>
"""

CSS = """
#header {margin-bottom: 4px}
#header .title {display: flex; align-items: center; justify-content: center; gap: 16px}
#header h1 {margin: 0; font-size: 2.1em; color: #2E7D32}
#header .title p {margin: 2px 0 0; color: var(--body-text-color-subdued)}
#header .authors {text-align: center; margin-top: 10px; line-height: 1.6}
#header .authors a {color: #2E7D32; text-decoration: none; font-weight: 600}
#header .authors small {color: var(--body-text-color-subdued)}
#header .links {display: flex; flex-wrap: wrap; justify-content: center; gap: 8px; margin: 12px 0}
#header .links a {padding: 5px 14px; border-radius: 999px; border: 1px solid #48A43F; color: #2E7D32;
                  text-decoration: none; font-weight: 600}
#header .links a:hover {background: #48A43F; color: white}
#header .tip {background: #f1f8ef; border-left: 4px solid #48A43F; padding: 10px 14px; border-radius: 8px;
              color: #1f3d1f}
.step h3 {margin: 6px 0 2px; color: #2E7D32}
#status {background: #f1f8ef; border-left: 4px solid #48A43F; padding: 8px 14px; border-radius: 8px; min-height: 2.4em}
#status p {margin: 0; color: #1f3d1f}
"""

with gr.Blocks(title="FlashReFocus", theme=gr.themes.Soft(primary_hue="green"), css=CSS) as demo:
    gr.HTML(HEADER)
    state = gr.State(None)
    with gr.Row():
        with gr.Column():
            gr.Markdown("### ① Upload a photo and restore it once", elem_classes="step")
            inp = gr.Image(label="Input photograph", type="pil", sources=["upload", "clipboard"], height=320)
            gr.Examples(EXAMPLES, inputs=inp, label="Examples (Unsplash portraits and our own photographs)")
            with gr.Accordion("Restoration settings", open=False):
                long_side = gr.Slider(512, 1536, value=1024, step=64, label="Long side (pixels)")
                omega = gr.Slider(0.0, 1.5, value=1.25, step=0.25, label="Detail weight ω",
                                  info="0: pixel LoRA only (most faithful); higher restores more detail.")
            btn = gr.Button("✨ Restore", variant="primary", size="lg")
            gr.Markdown("### ③ Depth map", elem_classes="step")
            depth = gr.Image(label="Depth used for rendering (near is bright): click to focus", type="pil",
                             interactive=False, height=280)
            with gr.Accordion("Depth controls", open=False):
                gamma = gr.Slider(0.4, 2.5, value=1.0, step=0.05, label="Depth curve γ (D → D^γ)",
                                  info="< 1 spreads out the near range, > 1 the far range; blur follows the new spacing.")
                band = gr.Slider(0.02, 0.3, value=0.06, step=0.01, label="Tilt-shift: half-height of the sharp band")
                user_depth = gr.Image(label="Your own depth map (optional; near is bright, any size)", type="numpy",
                                      sources=["upload", "clipboard"], height=150)
        with gr.Column():
            gr.Markdown("### ② Click the restored image to focus", elem_classes="step")
            aif = gr.Image(label="Restored all-in-focus image (click to focus)", type="pil", interactive=False,
                           height=320)
            mode = gr.Radio([SINGLE, DUAL, TILT], value=SINGLE, label="Focus mode",
                            info="Two planes: the first click sets A, the second B. Tilt-shift: click the row to keep sharp.")
            strength = gr.Slider(0.0, 1.5, value=0.5, step=0.05, label="🔆 Aperture (blur strength K)",
                                 info="0: everything sharp; 1: the widest aperture seen in training; above 1 extrapolates.")
            with gr.Row():
                show_marker = gr.Checkbox(value=False, label="Show focal markers")
                clear = gr.Button("Clear focal points", size="sm")
                reset = gr.Button("↺ Recommended settings", size="sm")
            status = gr.Markdown("Upload a photo or pick an example, then press ✨ Restore.", elem_id="status")
            out = gr.Image(label="Refocused result", type="pil", interactive=False, height=430, format="png")
    controls = [state, strength, mode, gamma, band, show_marker]
    outputs = [aif, out, depth, state, status]
    btn.click(restore, [inp, long_side, omega, gamma], [aif, depth, out, state, status])
    aif.select(on_click, controls, outputs)
    depth.select(on_click, controls, outputs)
    for slider in (strength, gamma, band):
        slider.release(on_change, controls, outputs)
    mode.change(on_change, controls, outputs)
    show_marker.change(on_change, controls, outputs)
    user_depth.change(on_user_depth, [state, user_depth, strength, mode, gamma, band, show_marker], outputs)
    clear.click(on_clear, [state, gamma], outputs)
    reset.click(on_reset, [state], [long_side, omega, gamma, band, strength, mode, show_marker, user_depth] + outputs)

if __name__ == "__main__":
    demo.queue(max_size=20).launch(show_api=False)