"""
FlexPart — A Unified Framework for Image-to-3D Part Generation via Variable Granularity.

Interactive demo: upload an object image (the background is removed automatically),
click on the object to mark each part you want, and the model generates a 3D mesh
for every marked part, composed into a single colored GLB.

Paper:  https://huggingface.co/papers/2610.31671
Code:   https://github.com/Gucci233/FlexPart
Weights: https://huggingface.co/gucci233/FlexPart
"""

import spaces  # MUST be the first import — before torch / diffusers / transformers

import json
import os
import sys
import tempfile
import time
from typing import List, Tuple, Union

import gradio as gr
import numpy as np
import torch
import torch.nn.functional as F
import trimesh
from accelerate.utils import set_seed
from huggingface_hub import hf_hub_download
from PIL import Image, ImageDraw
from torchvision import transforms
from transformers import BitImageProcessor, Dinov2Model

# This Space runs from its own repo root, so `src.*` imports resolve directly.
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from src.models.autoencoders import TripoSGVAEModel
from src.models.briarmbg import BriaRMBG
from src.models.transformers.flexpart_transformer import FlexPartDiTModel
from src.pipelines.pipeline_flexpart import FlexPartPipeline
from src.schedulers import RectifiedFlowScheduler

DEVICE = "cuda"
DTYPE = torch.float16
FLEXPART_REPO = "gucci233/FlexPart"
TRIPOSG_REPO = "VAST-AI/TripoSG"
RMBG_REPO = "briaai/RMBG-1.4"


# --------------------------------------------------------------------------- #
# Model loading (module scope, eager — ZeroGPU packs weights to disk and       #
# streams them into VRAM on the first @spaces.GPU call).                       #
# --------------------------------------------------------------------------- #
def _load_pipeline() -> FlexPartPipeline:
    """Assemble the FlexPart pipeline from its Hub components."""
    # 1. Transformer: FlexPart repo ships the bare safetensors; its config comes
    #    from the repo's configs/flexpart_inference.json, so build a local
    #    from_pretrained directory that pairs the two.
    ckpt_dir = hf_hub_download(FLEXPART_REPO, "diffusion_pytorch_model.safetensors")
    transformer_dir = os.path.join(os.path.dirname(ckpt_dir), "flexpart_transformer")
    os.makedirs(transformer_dir, exist_ok=True)
    config_src = os.path.join(os.path.dirname(os.path.abspath(__file__)), "configs", "flexpart_inference.json")
    with open(config_src) as f:
        import json as _json
        cfg = _json.load(f)
    with open(os.path.join(transformer_dir, "config.json"), "w") as f:
        _json.dump(cfg, f)
    ckpt_dst = os.path.join(transformer_dir, "diffusion_pytorch_model.safetensors")
    if not os.path.exists(ckpt_dst):
        try:
            os.symlink(ckpt_dir, ckpt_dst)
        except OSError:
            import shutil
            shutil.copy2(ckpt_dir, ckpt_dst)
    transformer = FlexPartDiTModel.from_pretrained(transformer_dir, torch_dtype=DTYPE)

    # 2. VAE / image encoder / feature extractor / scheduler come from TripoSG.
    vae = TripoSGVAEModel.from_pretrained(TRIPOSG_REPO, subfolder="vae", torch_dtype=DTYPE)
    image_encoder = Dinov2Model.from_pretrained(TRIPOSG_REPO, subfolder="image_encoder_dinov2", torch_dtype=DTYPE)
    feature_extractor = BitImageProcessor.from_pretrained(TRIPOSG_REPO, subfolder="feature_extractor_dinov2")
    scheduler = RectifiedFlowScheduler.from_pretrained(TRIPOSG_REPO, subfolder="scheduler")

    # 3. Background removal.
    rmbg = BriaRMBG.from_pretrained(RMBG_REPO)
    rmbg.to(DEVICE).eval()
    global _RMBG_NET
    _RMBG_NET = rmbg

    pipe = FlexPartPipeline(
        vae=vae,
        transformer=transformer,
        scheduler=scheduler,
        image_encoder_dinov2=image_encoder,
        feature_extractor_dinov2=feature_extractor,
    )
    return pipe.to(DEVICE, DTYPE)


pipe = _load_pipeline()

# Pre-annotated example point sets (pixel coordinates in each 2048x2048 source
# image). 1: robot dog (8 parts), 2: creature (6 parts), 3: authors' own
# example_points.npy (7 parts, gray electronic device).
EXAMPLE_POINTS = {
    "1": [[819, 472], [571, 748], [737, 404], [901, 404], [942, 1065], [707, 1588], [1259, 1536], [1414, 870]],
    "2": [[666, 512], [594, 389], [973, 942], [799, 1536], [1280, 1485], [1476, 727]],
    "3": [[1024.0, 1000.0], [430.0, 1260.0], [1610.0, 1250.0], [825.0, 1550.0], [1220.0, 1550.0], [815.0, 530.0], [1285.0, 530.0]],
}

RGB_PALETTE = [
    (82, 170, 220), (215, 91, 78), (45, 136, 117), (247, 172, 83), (124, 121, 121),
    (127, 171, 209), (243, 152, 101), (145, 204, 192), (150, 59, 121), (181, 206, 78),
    (189, 119, 149), (199, 193, 222), (200, 151, 54), (236, 110, 102), (238, 182, 212),
]


# --------------------------------------------------------------------------- #
# Mesh helpers (ported from the reference app / data_utils).                  #
# --------------------------------------------------------------------------- #
def filter_mesh(mesh: trimesh.Trimesh) -> trimesh.Trimesh:
    """Keep the largest connected submeshes (>20% of the max face count)."""
    submeshes = mesh.split(only_watertight=False)
    if len(submeshes) == 1:
        return submeshes[0]
    face_counts = [len(m.faces) for m in submeshes]
    max_faces = max(face_counts)
    filtered = [m for m in submeshes if len(m.faces) > max_faces * 0.2]
    if len(filtered) == 1:
        return filtered[0]
    merged = trimesh.util.concatenate(filtered)
    merged.merge_vertices()
    _ = merged.vertex_normals
    return merged


def normalize_mesh(mesh: Union[trimesh.Trimesh, trimesh.Scene], scale: float = 2.0):
    """Center the composition on its bounding-box centroid and fit it to `scale`."""
    bbox = mesh.bounding_box
    translation = -bbox.centroid
    scale = scale / bbox.primitive.extents.max()
    mesh.apply_translation(translation)
    mesh.apply_scale(scale)
    return mesh


def get_colored_mesh_composition(meshes: List[trimesh.Trimesh]) -> trimesh.Scene:
    """Compose per-part meshes into one Scene with a fixed part-color palette."""
    colored_scene = trimesh.Scene()
    for idx, mesh in enumerate(meshes):
        color = np.array(RGB_PALETTE[idx % len(RGB_PALETTE)])
        mesh.visual = trimesh.visual.ColorVisuals(mesh=mesh, vertex_colors=color)
        colored_scene.add_geometry(mesh)
    return colored_scene


# --------------------------------------------------------------------------- #
# Background removal (reference pipeline, 1:1).                               #
# --------------------------------------------------------------------------- #
def remove_background(pil_image: Image.Image) -> Image.Image:
    """Run BriaRMBG-1.4 and return the image with a soft alpha channel."""
    pil_image = pil_image.convert("RGB")
    w, h = pil_image.size

    im_np = np.array(pil_image)
    im_tensor = torch.tensor(im_np, dtype=torch.float32).permute(2, 0, 1).unsqueeze(0).to(DEVICE)
    im_tensor = torch.div(im_tensor, 255.0)
    im_tensor = transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))(im_tensor)
    im_tensor = F.interpolate(im_tensor, size=(1024, 1024), mode="bilinear", align_corners=False)

    with torch.no_grad():
        result = _RMBG_NET(im_tensor)

    mask_tensor = result
    while isinstance(mask_tensor, (list, tuple)):
        mask_tensor = mask_tensor[0]
    if mask_tensor.dim() == 2:
        mask_tensor = mask_tensor.unsqueeze(0).unsqueeze(0)
    elif mask_tensor.dim() == 3:
        mask_tensor = mask_tensor.unsqueeze(0)
    mask_tensor = F.interpolate(mask_tensor, size=(h, w), mode="bilinear", align_corners=False)

    ma = mask_tensor.squeeze().clamp(0, 1).cpu().numpy()
    pil_image.putalpha(Image.fromarray((ma * 255).astype(np.uint8)))
    return pil_image


# --------------------------------------------------------------------------- #
# Annotation helpers.                                                         #
# --------------------------------------------------------------------------- #
def draw_annotations(image: np.ndarray, points: List[List[float]]) -> np.ndarray:
    """Draw numbered markers on a copy of the image so users see their clicks."""
    pil = Image.fromarray(image).convert("RGB")
    draw = ImageDraw.Draw(pil)
    scale = max(pil.size) / 1024.0
    radius = max(int(9 * scale), 6)
    for i, (x, y) in enumerate(points):
        x, y = int(round(x)), int(round(y))
        draw.ellipse([x - radius, y - radius, x + radius, y + radius], fill=(255, 60, 60), outline=(255, 255, 255), width=max(radius // 3, 2))
        draw.text((x + radius + 2, y - radius), str(i + 1), fill=(255, 255, 255))
    return np.array(pil)


def parse_points_json(points_json: str, width: int, height: int) -> List[Tuple[float, float]]:
    """Parse the JSON list of [x, y] pixels into normalized (x/w, y/h) points."""
    if not points_json:
        return []
    try:
        raw = json.loads(points_json)
    except json.JSONDecodeError:
        return []
    out = []
    for p in raw:
        if len(p) < 2:
            continue
        x, y = float(p[0]), float(p[1])
        if 0 <= x < width and 0 <= y < height:
            out.append((x / width, y / height))
    return out


# --------------------------------------------------------------------------- #
# Event handlers.                                                             #
# --------------------------------------------------------------------------- #
def on_image_change(image: np.ndarray):
    """Reset annotations whenever a new image is loaded."""
    if image is None:
        return "[]", None
    return "[]", draw_annotations(image, [])


def on_image_select(image: np.ndarray, points_json: str, evt: gr.SelectData):
    """Append the clicked point (original-image pixel coords) to the annotation list."""
    if image is None:
        raise gr.Error("Upload an image first.")
    x, y = evt.index
    try:
        points = json.loads(points_json) if points_json else []
    except json.JSONDecodeError:
        points = []
    points.append([float(x), float(y)])
    return json.dumps(points), draw_annotations(image, points)


def on_clear(points_json: str, image: np.ndarray):
    """Clear all annotations."""
    if image is None:
        return "[]", None
    return "[]", draw_annotations(image, [])


def on_undo(points_json: str, image: np.ndarray):
    """Remove the most recent annotation."""
    if image is None:
        raise gr.Error("Upload an image first.")
    try:
        points = json.loads(points_json) if points_json else []
    except json.JSONDecodeError:
        points = []
    points = points[:-1] if points else []
    return json.dumps(points), draw_annotations(image, points)


def _estimate_duration(image, points_json, seed=1000, guidance_scale=7.0, steps=50, *args, **kwargs):
    """Scale reserved GPU time with the number of parts and diffusion steps.

    Measured baseline: 7 parts x 50 steps took ~84 s end-to-end on ZeroGPU,
    so reserve ~1.8x that and scale linearly with parts/steps, capped to stay
    within visitor tier limits.
    """
    try:
        num_parts = len(json.loads(points_json)) if points_json else 0
    except (json.JSONDecodeError, TypeError):
        num_parts = 0
    try:
        steps = int(steps)
    except (TypeError, ValueError):
        steps = 50
    return min(590, 30 + int(num_parts * 14 + steps * 0.5))


@spaces.GPU(duration=_estimate_duration)
def generate(image: np.ndarray, points_json: str, seed: float = 1000, guidance_scale: float = 7.0, steps: int = 50, progress=gr.Progress(track_tqdm=True)) -> Tuple[str, List[str], int]:
    """Generate a 3D mesh per annotated part.

    Args:
        image: the input object image (background already removed).
        points_json: JSON list of [x, y] click positions, one per part.
        seed: random seed for reproducibility.
        guidance_scale: classifier-free guidance strength.
        steps: number of diffusion denoising steps.

    Returns:
        (composed_glb_path, per_part_glb_paths, num_parts)
    """
    t0 = time.perf_counter()
    if image is None:
        raise gr.Error("Please upload an image first!")
    image_input_np = image
    if image_input_np.ndim == 3 and image_input_np.shape[2] == 3:
        # Background removal was skipped — run it now so the DINOv2 encoder
        # sees a white-backed, alpha-trimmed subject, as the reference does.
        image_input_np = np.array(remove_background(Image.fromarray(image_input_np)))
    h, w = image_input_np.shape[:2]

    points_list = parse_points_json(points_json, w, h)
    num_parts = len(points_list)
    if num_parts == 0:
        raise gr.Error("No annotations detected! Please click on the image to add points (one per part).")
    if num_parts > pipe.transformer.config.max_num_parts:
        raise gr.Error(f"The number of annotated parts ({num_parts}) exceeds the model capacity ({pipe.transformer.config.max_num_parts}).")

    # Point-only prompts (the reference path for this demo): each point gets a
    # full-frame placeholder box flagged invalid, so only the point is active.
    processed_boxes = [(0.0, 0.0, 1.0, 1.0)] * num_parts
    valid_boxes = [False] * num_parts
    prompts = {
        "masks_prompt": None,
        "boxes_prompt": processed_boxes,
        "points_prompt": points_list,
        "valid_masks": [False] * num_parts,
        "valid_boxes": valid_boxes,
        "valid_points": [True] * num_parts,
    }

    pil_image = Image.fromarray(image_input_np)
    if pil_image.mode == "RGBA":
        white_bg = Image.new("RGB", pil_image.size, (255, 255, 255))
        white_bg.paste(pil_image, mask=pil_image.split()[3])
        final_image = white_bg
    else:
        final_image = pil_image.convert("RGB")

    set_seed(int(seed))
    outputs = pipe(
        image=[final_image] * num_parts,
        attention_kwargs={"num_parts": num_parts},
        num_tokens=1024,
        generator=torch.Generator(device=DEVICE).manual_seed(int(seed)),
        num_inference_steps=int(steps),
        guidance_scale=guidance_scale,
        use_flash_decoder=False,
        **prompts,
    ).meshes

    if any(m is None or len(m.faces) == 0 for m in outputs):
        raise gr.Error("Mesh extraction returned an empty part. Try another seed or prompt.")

    clean_outputs = [filter_mesh(m) for m in outputs]
    merged_mesh = get_colored_mesh_composition(clean_outputs)
    merged_mesh = normalize_mesh(merged_mesh)

    composed_path = os.path.join(tempfile.mkdtemp(), "flexpart_composed.glb")
    merged_mesh.export(composed_path)

    part_paths = []
    for i, mesh in enumerate(clean_outputs):
        m = normalize_mesh(mesh)
        p = os.path.join(os.path.dirname(composed_path), f"flexpart_part_{i}.glb")
        m.export(p)
        part_paths.append(p)

    print(f"[FlexPart] generated {num_parts} parts in {time.perf_counter() - t0:.1f}s")
    return composed_path, part_paths, num_parts


# --------------------------------------------------------------------------- #
# UI                                                                          #
# --------------------------------------------------------------------------- #
css = """
.gradio-container {max-width: 1400px !important;}
.annot-preview img {image-rendering: auto;}
"""

with gr.Blocks(theme=gr.themes.Citrus(), css=css, title="FlexPart — Image-to-3D Part Generation") as demo:
    gr.Markdown(
        """
        # 🧩 FlexPart — Image-to-3D Part Generation
        **A Unified Framework for Image-to-3D Part Generation via Variable Granularity**

        Upload an object image, click once on each part you want (each click = one part), and hit **Generate**.
        The model produces one 3D mesh per clicked part, composed into a single colored GLB.
        """
    )

    with gr.Row():
        with gr.Column(scale=1):
            input_image = gr.Image(
                label="Input image (click a part to mark it)",
                type="numpy",
                interactive=True,
                height=420,
                sources=["upload", "clipboard"],
            )
            annotated_preview = gr.Image(
                label="Your annotations (preview)",
                type="numpy",
                interactive=False,
                height=420,
                elem_classes=["annot-preview"],
            )
            points_state = gr.Textbox(
                value="[]",
                label="Part points (JSON, one [x, y] per click — editable)",
                lines=2,
                interactive=True,
            )

            with gr.Row():
                undo_btn = gr.Button("↩ Undo last")
                clear_btn = gr.Button("🗑 Clear all")

            with gr.Accordion("Advanced settings", open=False):
                seed = gr.Number(value=1000, label="Seed", precision=0)
                guidance = gr.Slider(1.0, 20.0, value=7.0, step=0.1, label="Guidance scale")
                steps = gr.Slider(10, 100, value=50, step=1, label="Diffusion steps")

            gen_btn = gr.Button("🚀 Generate 3D Mesh", variant="primary")

        with gr.Column(scale=1):
            part_count = gr.Number(value=0, label="Parts generated", precision=0, interactive=False)
            composed_3d = gr.Model3D(
                label="Composed 3D parts (colored per part)",
                clear_color=[1.0, 1.0, 1.0, 1.0],
                camera_position=(180, 90, 4),
            )
            part_files = gr.Files(label="Individual part GLBs (download)")

    gr.Markdown(
        """
        ### How it works
        1. **Upload** an object image (or pick an example below). The background is removed with BriaRMBG-1.4.
        2. **Click** on the image — every click marks one part you want the model to generate (up to 32).
        3. **Generate** — FlexPart runs part-level 3D diffusion (TripoSG VAE + DINOv2 image conditioning) and
        extracts one mesh per part via hierarchical octree marching cubes, then composes them with per-part colors.

        Model: [gucci233/FlexPart](https://huggingface.co/gucci233/FlexPart) ·
        Code: [Gucci233/FlexPart](https://github.com/Gucci233/FlexPart) ·
        Paper: [2610.31671](https://huggingface.co/papers/2610.31671)
        """
    )

    examples = gr.Examples(
        examples=[
            ["assets/1.png", json.dumps(EXAMPLE_POINTS["1"])],
            ["assets/2.png", json.dumps(EXAMPLE_POINTS["2"])],
            ["assets/3.png", json.dumps(EXAMPLE_POINTS["3"])],
        ],
        inputs=[input_image, points_state],
        fn=generate,
        outputs=[composed_3d, part_files, part_count],
        cache_examples=True,
        cache_mode="lazy",
        label="Examples (pre-annotated — click to run)",
    )

    # Wiring
    input_image.select(on_image_select, inputs=[input_image, points_state], outputs=[points_state, annotated_preview])
    input_image.change(on_image_change, inputs=[input_image], outputs=[points_state, annotated_preview])
    undo_btn.click(on_undo, inputs=[points_state, input_image], outputs=[points_state, annotated_preview])
    clear_btn.click(on_clear, inputs=[points_state, input_image], outputs=[points_state, annotated_preview])
    gen_btn.click(
        generate,
        inputs=[input_image, points_state, seed, guidance, steps],
        outputs=[composed_3d, part_files, part_count],
    )

demo.queue().launch(mcp_server=True)