"""Interactive demo for Deep ContourFlow (unsupervised mode).

Upload an image: a circular contour is initialized and then *evolved* by the
training-free DCF algorithm to wrap around the main object, guided only by the
multi-scale features of a frozen VGG16 — no training, no labels.
"""
import os
import tempfile

import cv2
import gradio as gr
import matplotlib

matplotlib.use("Agg")  # headless backend for the Space
import numpy as np
import torch
from torch_contour import CleanContours

from deep_contourflow import UnsupervisedDCF
from deep_contourflow.features import define_contour_init
from deep_contourflow.visualization import plot_contour_evolution

DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
USE_AMP = DEVICE.type == "cuda"
NB_NODES = 200


def segment(image, height, n_epochs, init_size, area_force):
    if image is None:
        raise gr.Error("Please upload an image first.")

    height = int(height)
    img = cv2.resize(image, (height, height), interpolation=cv2.INTER_AREA).astype(np.uint8)
    tensor = (torch.tensor(np.moveaxis(img, -1, 0)[None]) / 255.0).to(DEVICE)

    # Circular contour initialization, resampled to NB_NODES points in [0, 1].
    contour_init, _ = define_contour_init(n=height, shape="circle", size=float(init_size))
    contour_init = CleanContours().interpolate(contour_init, NB_NODES).clip(0, 1)
    contour_init = torch.tensor(contour_init)[None, None].float().to(DEVICE)

    dcf = UnsupervisedDCF(
        model="vgg16",
        n_epochs=int(n_epochs),
        learning_rate=1e-2,
        area_force=float(area_force),
        sigma=5e-1,
        clip=1e-1,
        use_mixed_precision=USE_AMP,
    )
    contours, loss_history, _ = dcf.predict(tensor, contour_init)

    fig = plot_contour_evolution(img, contours, loss_history)
    out_path = os.path.join(tempfile.gettempdir(), "dcf_result.png")
    fig.savefig(out_path, dpi=140, bbox_inches="tight", facecolor="#FAF9F6")
    return out_path


DESCRIPTION = """
# 🪢 Deep ContourFlow — interactive demo

**Training-free** image segmentation: a circle is evolved into the object's
boundary using only the features of a frozen VGG16. No training, no labels.

Upload an image and hit **Submit**. The result shows the contour at several
steps, from the initial circle to its converged shape.

> ⏳ Running on free CPU — a segmentation takes a few seconds to ~1–2 min
> depending on resolution and the number of iterations.
> 📄 [Paper (arXiv:2407.10696)](https://arxiv.org/abs/2407.10696) ·
> 💻 [Code](https://github.com/antoinehabis/Deep-ContourFlow)
"""

EXAMPLE_DIR = os.path.join(os.path.dirname(__file__), "examples")
_examples = [
    [os.path.join(EXAMPLE_DIR, name), 384, 60, 0.5, 1e-3]
    for name in ("lion.jpg", "flower0.jpg", "pineapple.jpg")
    if os.path.exists(os.path.join(EXAMPLE_DIR, name))
]

demo = gr.Interface(
    fn=segment,
    inputs=[
        gr.Image(type="numpy", label="Input image"),
        gr.Slider(192, 512, value=384, step=64, label="Resolution (px)"),
        gr.Slider(10, 120, value=60, step=10, label="Iterations (epochs)"),
        gr.Slider(0.2, 0.9, value=0.5, step=0.05, label="Initial circle size"),
        gr.Slider(0.0, 5e-3, value=1e-3, step=5e-4, label="Area regularization"),
    ],
    outputs=gr.Image(type="filepath", label="Contour evolution"),
    title="Deep ContourFlow",
    description=DESCRIPTION,
    examples=_examples or None,
    cache_examples=False,
    flagging_mode="never",
)

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