import os

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import spaces  # noqa: E402  — must be before torch

import torch  # noqa: E402
import gradio as gr  # noqa: E402
import numpy as np  # noqa: E402
from copy import deepcopy  # noqa: E402
from pathlib import Path  # noqa: E402

from PIL import ImageDraw  # noqa: E402
from huggingface_hub import snapshot_download  # noqa: E402
from transformers import (  # noqa: E402
    AutoModelForZeroShotObjectDetection,
    AutoProcessor,
)
from diffusers.pipelines.flux.pipeline_flux_kontext import (  # noqa: E402
    PREFERRED_KONTEXT_RESOLUTIONS,
)

# ---------------------------------------------------------------------------
# Paths
# ---------------------------------------------------------------------------
HERE = Path(__file__).parent.resolve()
os.sys.path.insert(0, str(HERE))

from pipelines.onehoi import OneHOIPipeline  # noqa: E402
from modules.utils import get_xyxy_from_mask  # noqa: E402

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------
BASE_MODEL = "black-forest-labs/FLUX.1-Kontext-dev"
ONEHOI_REPO = "jiuntian/OneHOI"
# Small open-vocabulary detector used to recover the subject / object layout
# from the source image so that OneHOI's spatial Action Grounding receives real
# boxes (as the reference implementation always does) instead of all-zero
# ("dropped") boxes, which let the model recompose the scene freely and
# destroyed subject identity / appearance.
DETECTOR_REPO = "IDEA-Research/grounding-dino-tiny"

# Exactly the negative prompt used by the reference inference script.
NEGATIVE_PROMPT = (
    "worst quality, low quality, monochrome, zombie, interlocked fingers, "
    "Aissist, cleavage, nsfw, blur,"
)
ZERO_BOX = [0.0, 0.0, 0.0, 0.0]

# ---------------------------------------------------------------------------
# Load the pipeline at module scope (ZeroGPU rule #2)
# ---------------------------------------------------------------------------
print("Downloading OneHOI checkpoint weights …")
ckpt_dir = snapshot_download(
    ONEHOI_REPO,
    repo_type="model",
    allow_patterns=[
        "pytorch_lora_weights.safetensors",
        "grounding_encoder.safetensors",
        "grounding_encoder/config.json",
    ],
)
print(f"OneHOI weights cached at {ckpt_dir}")

print(f"Loading OneHOIPipeline from {BASE_MODEL} …")
pipeline = OneHOIPipeline.from_checkpoint(
    base_model=BASE_MODEL,
    checkpoint_path=str(ckpt_dir),
    torch_dtype=torch.bfloat16,
)
pipeline.to("cuda")

# Speed up: fuse LoRA then unload the raw LoRA weights
pipeline.fuse_lora()
pipeline.unload_lora_weights()
pipeline.vae.to(memory_format=torch.channels_last)

print(f"Loading layout detector {DETECTOR_REPO} …")
det_processor = AutoProcessor.from_pretrained(DETECTOR_REPO)
detector = AutoModelForZeroShotObjectDetection.from_pretrained(DETECTOR_REPO)
detector.eval()
detector.to("cuda")

print("Pipeline ready.")


# ---------------------------------------------------------------------------
# Spatial Action Grounding helpers
# ---------------------------------------------------------------------------
def _normalize_box(box, width, height):
    x1, y1, x2, y2 = (float(v) for v in box)
    box = [x1 / width, y1 / height, x2 / width, y2 / height]
    return [min(1.0, max(0.0, v)) for v in box]


def _ground_one(image, query):
    """Return the best normalized xyxy box for *query* in *image*.

    Falls back to an all-zero box when nothing is found; OneHOI treats a
    zero-area box as "unconstrained" (the modality-dropout case), so the model
    is then free to place that entity.
    """
    query = query.strip().lower().rstrip(".")
    if not query:
        return list(ZERO_BOX)
    inputs = det_processor(images=image, text=query + ".", return_tensors="pt")
    inputs = {k: v.to(detector.device) for k, v in inputs.items()}
    with torch.inference_mode():
        outputs = detector(**inputs)
    post = det_processor.post_process_grounded_object_detection
    common = dict(text_threshold=0.15, target_sizes=[(image.height, image.width)])
    results = None
    for args in ((outputs, inputs["input_ids"]), (outputs,)):
        for thr in ("threshold", "box_threshold"):
            try:
                results = post(*args, **{thr: 0.15}, **common)[0]
                break
            except TypeError:
                continue
        if results is not None:
            break
    if results is None:
        raise RuntimeError("incompatible post_process_grounded_object_detection API")
    scores = results["scores"]
    if len(scores) == 0:
        return list(ZERO_BOX)
    best = int(torch.as_tensor(scores).argmax())
    return _normalize_box(
        torch.as_tensor(results["boxes"])[best].tolist(), image.width, image.height
    )


def _layout_preview(image, boxes, labels):
    """Draw the grounded subject / object boxes on top of the source image."""
    preview = image.convert("RGB").copy()
    draw = ImageDraw.Draw(preview)
    colors = ["#ff2d55", "#00c853"]
    w, h = preview.size
    for box, label, color in zip(boxes, labels, colors):
        if sum(box) <= 0:
            continue
        xy = [box[0] * w, box[1] * h, box[2] * w, box[3] * h]
        draw.rectangle(xy, outline=color, width=max(2, int(min(w, h) * 0.006)))
        draw.text((xy[0] + 4, max(0, xy[1] - 12)), label, fill=color)
    return preview


def _kontext_size(image):
    """Pick the FLUX-Kontext resolution matching the source aspect ratio.

    Without this the pipeline always renders a 1024x1024 target while the
    conditioning image is resized to its own aspect ratio, which stretches the
    subject and hurts appearance preservation for non-square inputs.
    """
    aspect_ratio = image.width / image.height
    _, width, height = min(
        (abs(aspect_ratio - w / h), w, h) for w, h in PREFERRED_KONTEXT_RESOLUTIONS
    )
    return int(height), int(width)


# ---------------------------------------------------------------------------
# Inference
# ---------------------------------------------------------------------------
@spaces.GPU(duration=240)
def edit_hoi(
    source_image,
    prompt,
    subject_label,
    object_label,
    action_label,
    guidance_scale=3.5,
    true_cfg=3.5,
    num_steps=28,
    seed=2026,
    progress=gr.Progress(track_tqdm=True),
):
    """Run OneHOI HOI editing on *source_image*.

    Parameters
    ----------
    source_image : PIL.Image or None
        The input image to edit.  If None the model generates from scratch.
    prompt : str
        Natural-language description of the desired edit.
    subject_label : str
        Label for the subject (almost always "person").
    object_label : str
        Label for the object (e.g. "skateboard", "dog", "broccoli").
    action_label : str
        The target interaction / action (e.g. "ride", "hug", "eat").

    Returns
    -------
    (PIL.Image, PIL.Image)
        The edited image and a preview of the grounded subject / object layout.
    """
    if not prompt or not prompt.strip():
        prompt = f"the person is now {action_label} the {object_label}."
    if not subject_label.strip():
        subject_label = "person"

    generator = torch.Generator(device="cpu").manual_seed(int(seed))

    box_labels = [subject_label, object_label]
    hoi_labels = [action_label]

    kwargs = dict(
        prompt=prompt,
        box_labels=box_labels,
        hoi_labels=hoi_labels,
        generator=generator,
        negative_prompt=NEGATIVE_PROMPT,
        true_cfg_scale=float(true_cfg),
        guidance_scale=float(guidance_scale),
        num_inference_steps=int(num_steps),
        hoi_seq_len=512,
    )

    # Spatial Action Grounding: recover the subject / object boxes from the
    # source image (normalized xyxy, [subject, object] order — the same
    # convention as the reference implementation) so the interaction is edited
    # in place instead of the scene being recomposed from scratch.
    boxes = [list(ZERO_BOX), list(ZERO_BOX)]
    preview = None
    if source_image is not None:
        kwargs["image"] = source_image
        height, width = _kontext_size(source_image)
        kwargs["height"] = height
        kwargs["width"] = width
        try:
            boxes = [_ground_one(source_image, label) for label in box_labels]
        except Exception as exc:  # pragma: no cover - never fail the edit
            print(f"[grounding] detection failed, falling back to layout-free: {exc}")
            boxes = [list(ZERO_BOX), list(ZERO_BOX)]
        print(f"[grounding] {box_labels} -> {boxes} (target {width}x{height})")
        preview = _layout_preview(source_image, boxes, box_labels)

    kwargs["boxes"] = boxes

    result = pipeline(**kwargs)
    out = result.images[0]
    return out, preview if preview is not None else out


# ---------------------------------------------------------------------------
# UI
# ---------------------------------------------------------------------------
custom_css = """
#title-block { text-align: center; margin-bottom: 20px; }
#title-block h1 { font-size: 2em; margin-bottom: 0; }
#title-block p { opacity: 0.7; margin-top: 0.5em; }
"""

with gr.Blocks(theme=gr.themes.Citrus(), css=custom_css) as demo:
    gr.HTML(
        """
        <div id="title-block">
            <h1>🎯 OneHOI — Human-Object Interaction Editing</h1>
            <p>Edit human-object interactions in images with text guidance. Based on
            <a href="https://arxiv.org/abs/2604.14062">OneHOI (CVPR 2026)</a> &amp;
            <a href="https://arxiv.org/abs/2503.09130">InteractEdit</a> by Jiun Tian Hoe et al.</p>
        </div>
        """
    )

    with gr.Row():
        with gr.Column(scale=1):
            source_image = gr.Image(
                label="Source Image",
                type="pil",
                height=400,
            )
            prompt = gr.Textbox(
                label="Edit Prompt",
                placeholder="e.g. 'the person is now riding the skateboard.'",
                value="the person is now riding the skateboard.",
            )
            with gr.Row():
                subject_label = gr.Textbox(
                    label="Subject",
                    value="person",
                    scale=1,
                )
                object_label = gr.Textbox(
                    label="Object",
                    value="skateboard",
                    scale=1,
                )
                action_label = gr.Textbox(
                    label="Action / Interaction",
                    value="ride",
                    scale=1,
                )

            with gr.Accordion("Advanced Settings", open=False):
                guidance_scale = gr.Slider(
                    label="Guidance Scale",
                    minimum=1.0,
                    maximum=10.0,
                    step=0.5,
                    value=3.5,
                )
                true_cfg = gr.Slider(
                    label="True CFG Scale",
                    minimum=1.0,
                    maximum=10.0,
                    step=0.5,
                    value=3.5,
                )
                num_steps = gr.Slider(
                    label="Inference Steps",
                    minimum=4,
                    maximum=50,
                    step=1,
                    value=28,
                )
                seed = gr.Number(
                    label="Seed",
                    value=2026,
                    precision=0,
                )

            run_btn = gr.Button("Edit Interaction", variant="primary")

        with gr.Column(scale=1):
            output_image = gr.Image(
                label="Edited Result",
                type="pil",
                height=400,
            )
            layout_preview = gr.Image(
                label="Action Grounding (detected subject / object layout)",
                type="pil",
                height=240,
            )

    gr.Examples(
        examples=[
            [
                "examples/HICO_train2015_00004579.png",
                "the person is now riding the skateboard.",
                "person",
                "skateboard",
                "ride",
            ],
            [
                "examples/HICO_train2015_00004579.png",
                "the person is now sitting on the skateboard.",
                "person",
                "skateboard",
                "sit on",
            ],
            [
                "examples/HICO_train2015_00014446.png",
                "the person is now eating the broccoli.",
                "person",
                "broccoli",
                "eat",
            ],
            [
                "examples/HICO_train2015_00014446.png",
                "the person is now cutting the broccoli.",
                "person",
                "broccoli",
                "cut",
            ],
            [
                "examples/HICO_train2015_00015541.png",
                "the person is now jumping with the skateboard.",
                "person",
                "skateboard",
                "jump",
            ],
            [
                "examples/HICO_train2015_00015541.png",
                "the person is now sitting on the skateboard.",
                "person",
                "skateboard",
                "sit on",
            ],
        ],
        inputs=[
            source_image,
            prompt,
            subject_label,
            object_label,
            action_label,
        ],
        fn=edit_hoi,
        outputs=[output_image, layout_preview],
        cache_examples=True,
        cache_mode="lazy",
    )

    run_btn.click(
        fn=edit_hoi,
        inputs=[
            source_image,
            prompt,
            subject_label,
            object_label,
            action_label,
            guidance_scale,
            true_cfg,
            num_steps,
            seed,
        ],
        outputs=[output_image, layout_preview],
    )

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