"""PatchAlign3D — open-vocabulary (zero-shot) 3D part segmentation from point clouds.

Paper:   https://huggingface.co/papers/2601.02457
Code:    https://github.com/souhail-hadgi/PatchAlign3D
Weights: https://huggingface.co/patchalign3d/patchalign3d-encoder
"""

import os

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

import spaces  # noqa: E402  (must precede torch)

import tempfile  # noqa: E402
import time  # noqa: E402
from pathlib import Path  # noqa: E402

import gradio as gr  # noqa: E402
import numpy as np  # noqa: E402
import plotly.graph_objects as go  # noqa: E402
import torch  # noqa: E402
import trimesh  # noqa: E402
from huggingface_hub import hf_hub_download  # noqa: E402
from transformers import CLIPTextModelWithProjection, CLIPTokenizer  # noqa: E402

import patchalign3d as pa  # noqa: E402

# --------------------------------------------------------------------------------------
# Models — module scope, eager .to("cuda"); ZeroGPU streams them in on the first call
# --------------------------------------------------------------------------------------

CKPT = hf_hub_download("patchalign3d/patchalign3d-encoder", "patchalign3d.pt")
model, proj = pa.load_patchalign3d(CKPT)
model = model.to("cuda")
proj = proj.to("cuda")

tokenizer = CLIPTokenizer.from_pretrained(pa.CLIP_TEXT_REPO, subfolder=pa.CLIP_TOKENIZER_SUBFOLDER)
text_model = (
    CLIPTextModelWithProjection.from_pretrained(
        pa.CLIP_TEXT_REPO, subfolder=pa.CLIP_TEXT_SUBFOLDER, variant="fp16", dtype=torch.float32
    )
    .eval()
    .to("cuda")
)

print(
    f"[init] PatchAlign3D encoder {sum(p.numel() for p in model.parameters()) / 1e6:.1f}M params | "
    f"CLIP ViT-bigG-14 text tower {sum(p.numel() for p in text_model.parameters()) / 1e6:.1f}M params | "
    f"tokenizer pad={tokenizer.pad_token_id} ctx={tokenizer.model_max_length}"
)

MAX_LABELS = 12
MESH_EXTS = {".obj", ".glb", ".gltf", ".stl", ".off", ".ply", ".dae", ".3mf"}

# Distinguishable qualitative palette
PALETTE = [
    "#e6194b", "#3cb44b", "#4363d8", "#f58231", "#911eb4", "#00b8d4",
    "#f032e6", "#a1c800", "#fabed4", "#469990", "#9a6324", "#7f0000",
]

# ShapeNetPart part vocabularies, verbatim from the official eval.py
PRESETS = {
    "— custom —": ("", ""),
    "Airplane": ("body, wing, tail, engine or frame", "airplane"),
    "Bag": ("handle, body", "bag"),
    "Cap": ("crown, brim", "cap"),
    "Car": ("roof, hood, wheel, body", "car"),
    "Chair": ("back, seat, leg, arm", "chair"),
    "Earphone": ("earcup, headband, data wire", "earphone"),
    "Guitar": ("headstock, neck, body", "guitar"),
    "Knife": ("blade, handle", "knife"),
    "Lamp": ("base, lampshade, fixing bracket, pole", "lamp"),
    "Laptop": ("keyboard, screen", "laptop"),
    "Motorbike": ("gas tank, seat, wheel, handles or handlebars, headlight, engine or frame", "motorbike"),
    "Mug": ("handle, cup", "mug"),
    "Pistol": ("barrel, handle or grip, trigger and guard", "pistol"),
    "Rocket": ("body, fin, nose", "rocket"),
    "Skateboard": ("wheel, deck, belt for foot", "skateboard"),
    "Table": ("desktop, leg or support, drawer", "table"),
}


# --------------------------------------------------------------------------------------
# Shape loading
# --------------------------------------------------------------------------------------


def _resample(pts: np.ndarray, npoints: int, seed: int) -> np.ndarray:
    n = len(pts)
    if n == npoints:
        return pts
    rng = np.random.default_rng(seed)
    return pts[rng.choice(n, size=npoints, replace=n < npoints)]


def load_shape(path: str, npoints: int = pa.DEFAULT_NPOINTS, seed: int = 0):
    """Read a mesh or point cloud; return `npoints` unit-sphere-normalised points + a description."""
    p = Path(path)
    ext = p.suffix.lower()
    src = "point cloud"

    if ext in (".npz", ".npy"):
        if ext == ".npy":
            arr = np.load(p)
        else:
            d = np.load(p, allow_pickle=True)
            key = next((k for k in ("points", "xyz", "pos", "vertices") if k in d), None)
            if key is None:
                raise gr.Error(f"NPZ must contain points/xyz/pos/vertices — found {list(d.keys())}")
            arr = d[key]
        arr = np.asarray(arr, dtype=np.float32)
        pts = arr.reshape(-1, arr.shape[-1])[:, :3]
    elif ext in (".txt", ".pts", ".xyz", ".csv", ".asc"):
        raw = np.loadtxt(p, delimiter="," if ext == ".csv" else None, dtype=np.float32)
        pts = np.atleast_2d(raw)[:, :3]
    elif ext in MESH_EXTS:
        obj = trimesh.load(str(p), process=False)
        if isinstance(obj, trimesh.Scene):
            faced = [g for g in obj.geometry.values() if getattr(g, "faces", None) is not None and len(g.faces)]
            if faced:
                try:
                    obj = obj.to_mesh()
                except Exception:
                    obj = trimesh.util.concatenate(faced)
            else:
                verts = [np.asarray(g.vertices) for g in obj.geometry.values() if hasattr(g, "vertices")]
                if not verts:
                    raise gr.Error("No geometry found in this file.")
                obj = trimesh.PointCloud(np.concatenate(verts, axis=0))
        if getattr(obj, "faces", None) is not None and len(obj.faces) > 0:
            np.random.seed(int(seed) % (2**31))
            pts = np.asarray(trimesh.sample.sample_surface(obj, int(npoints))[0], dtype=np.float32)
            src = f"mesh, {len(obj.faces):,} faces, surface-sampled"
        else:
            pts = np.asarray(obj.vertices, dtype=np.float32)[:, :3]
    else:
        raise gr.Error(
            f"Unsupported file type '{ext}'. Use a mesh (.obj/.glb/.gltf/.stl/.off/.ply) "
            "or a point cloud (.ply/.npz/.txt/.xyz)."
        )

    pts = np.ascontiguousarray(pts[np.isfinite(pts).all(axis=1)], dtype=np.float32)
    if len(pts) < 32:
        raise gr.Error(f"Only {len(pts)} usable points found — need at least 32.")
    raw_n = len(pts)
    pts = _resample(pts, int(npoints), int(seed))
    return pa.pc_normalize(pts.astype(np.float32)), f"{src}, {raw_n:,} pts → {len(pts):,} used"


# --------------------------------------------------------------------------------------
# Plotting
# --------------------------------------------------------------------------------------

_AXIS = dict(showbackground=False, showgrid=False, zeroline=False, showticklabels=False, title="")


def _style(fig: go.Figure, title: str, height: int) -> go.Figure:
    fig.update_layout(
        title=dict(text=title, x=0.02, font=dict(size=12, color="#8a8a8a")),
        scene=dict(xaxis=_AXIS, yaxis=_AXIS, zaxis=_AXIS, aspectmode="data",
                   camera=dict(eye=dict(x=1.6, y=1.2, z=1.0))),
        margin=dict(l=0, r=0, t=28, b=0),
        height=height,
        showlegend=len(fig.data) > 1,
        legend=dict(orientation="h", yanchor="bottom", y=0.0, xanchor="left", x=0.0,
                    font=dict(color="#8a8a8a", size=11), bgcolor="rgba(0,0,0,0)"),
        paper_bgcolor="rgba(0,0,0,0)",
        plot_bgcolor="rgba(0,0,0,0)",
        font=dict(color="#8a8a8a"),
    )
    return fig


def plot_raw(points: np.ndarray, title: str, height: int = 300) -> go.Figure:
    fig = go.Figure(
        go.Scatter3d(
            x=points[:, 0], y=points[:, 1], z=points[:, 2], mode="markers",
            marker=dict(size=1.8, color="#9aa0a6"), name="input", hoverinfo="skip",
        )
    )
    return _style(fig, title, height)


def plot_segments(points, pred, names, conf, title: str, height: int = 560) -> go.Figure:
    fig = go.Figure()
    for k, name in enumerate(names):
        m = pred == k
        if not m.any():
            continue
        fig.add_trace(
            go.Scatter3d(
                x=points[m, 0], y=points[m, 1], z=points[m, 2], mode="markers",
                marker=dict(size=2.6, color=PALETTE[k % len(PALETTE)]),
                name=f"{name} · {int(m.sum())}",
                customdata=conf[m],
                hovertemplate=f"<b>{name}</b><br>p=%{{customdata:.2f}}<extra></extra>",
            )
        )
    return _style(fig, title, height)


def export_colored_ply(points: np.ndarray, pred: np.ndarray) -> str:
    rgba = np.zeros((len(points), 4), dtype=np.uint8)
    rgba[:, 3] = 255
    for k in range(int(pred.max()) + 1):
        h = PALETTE[k % len(PALETTE)].lstrip("#")
        rgba[pred == k, :3] = [int(h[i:i + 2], 16) for i in (0, 2, 4)]
    f = tempfile.NamedTemporaryFile(suffix="_patchalign3d.ply", delete=False)
    f.close()
    trimesh.PointCloud(points, colors=rgba).export(f.name)
    return f.name


# --------------------------------------------------------------------------------------
# Handlers
# --------------------------------------------------------------------------------------


def preview_shape(shape_file: str, num_points: int = pa.DEFAULT_NPOINTS, seed: int = 0):
    """Show the uploaded shape as a plain point cloud. CPU only — no GPU needed.

    Args:
        shape_file: Path to a mesh or point-cloud file.
        num_points: Number of points to sample for the preview.
        seed: Sampling seed.

    Returns:
        An interactive 3D scatter plot of the sampled input points.
    """
    if not shape_file:
        return None
    points, info = load_shape(shape_file, int(num_points), int(seed))
    return plot_raw(points, f"Input — {info}")


def _parse_labels(labels_text: str):
    names = [x.strip() for x in (labels_text or "").split(",") if x.strip()]
    if not names:
        raise gr.Error("Enter at least one part name, e.g. `back, seat, leg, arm`.")
    if len(names) > MAX_LABELS:
        raise gr.Error(f"At most {MAX_LABELS} part queries at a time (got {len(names)}).")
    return names


def _estimate_duration(shape_file=None, labels_text="", num_points=pa.DEFAULT_NPOINTS, *args, **kwargs):
    """GPU reservation. Measured worst case is ~1.2 s of compute at the heaviest settings; the rest is
    headroom for loading / surface-sampling a large user-supplied mesh inside the same call."""
    try:
        mb = os.path.getsize(shape_file) / 1e6
    except Exception:
        mb = 0.0
    try:
        n = int(num_points)
    except Exception:
        n = pa.DEFAULT_NPOINTS
    return int(min(90, 5 + 0.6 * mb + 2.0 * n / pa.DEFAULT_NPOINTS))


@spaces.GPU(duration=_estimate_duration)
def segment(
    shape_file: str,
    labels_text: str = "back, seat, leg, arm",
    num_points: int = pa.DEFAULT_NPOINTS,
    num_group: int = pa.DEFAULT_NUM_GROUP,
    group_size: int = pa.DEFAULT_GROUP_SIZE,
    text_setting: str = "part_only",
    category: str = "",
    assign: str = "nearest",
    tau: float = pa.DEFAULT_TAU,
    seed: int = 0,
):
    """Zero-shot 3D part segmentation of a shape, driven by free-form text part names.

    Args:
        shape_file: Path to a mesh (.obj/.glb/.gltf/.stl/.off/.ply) or point cloud (.ply/.npz/.txt/.xyz).
        labels_text: Comma-separated part names to look for, e.g. "back, seat, leg, arm".
        num_points: Points sampled from the shape (2048 matches the training setting).
        num_group: Number of patches (furthest-point-sampled centres) the encoder uses.
        group_size: Points per patch (k-NN neighbourhood size).
        text_setting: Prompt ensemble — "part_only", "part_plus_cat" or "ensemble".
        category: Object category used by the "part_plus_cat" / "ensemble" prompts, e.g. "chair".
        assign: Patch-to-point assignment — "nearest" patch centre, or patch "membership" voting.
        tau: CLIP temperature used to turn cosine similarities into probabilities.
        seed: Seed for point / surface sampling.

    Returns:
        An interactive 3D plot of the segmented shape, the share of points per part,
        a colour-coded .ply download, and a short run summary.
    """
    if not shape_file:
        raise gr.Error("Upload a 3D shape first, or pick one of the examples below.")
    names = _parse_labels(labels_text)

    num_points = int(num_points)
    num_group = max(1, min(int(num_group), num_points))
    group_size = max(1, min(int(group_size), num_points))

    t0 = time.perf_counter()
    points, info = load_shape(shape_file, num_points, int(seed))
    t_load = time.perf_counter() - t0

    t1 = time.perf_counter()
    pred, probs = pa.segment_point_cloud(
        points, names, model, proj, text_model, tokenizer, "cuda",
        category=category or "", text_setting=text_setting, assign=assign,
        tau=float(tau), num_group=num_group, group_size=group_size,
    )
    t_gpu = time.perf_counter() - t1

    conf = probs[np.arange(len(pred)), pred]
    shares = {name: float((pred == k).mean()) for k, name in enumerate(names)}
    fig = plot_segments(points, pred, names, conf, "Predicted parts — drag to rotate, scroll to zoom")
    ply = export_colored_ply(points, pred)
    summary = (
        f"**{len(points):,} points → {num_group} patches → {len(names)} text queries**  \n"
        f"{info} · prompts `{text_setting}`"
        + (f" · category `{category}`" if category and text_setting != "part_only" else "")
        + f"  \nload {t_load:.2f}s · inference **{t_gpu:.2f}s** · mean confidence {conf.mean():.2f}"
    )
    return fig, shares, ply, summary


def apply_preset(preset: str, labels_text: str, category: str):
    if preset in PRESETS and preset != "— custom —":
        return PRESETS[preset]
    return labels_text, category


# --------------------------------------------------------------------------------------
# UI
# --------------------------------------------------------------------------------------

CSS = """
#col-container { max-width: 1240px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

EXAMPLES = [
    ["examples/chair.ply", "back, seat, leg"],
    ["examples/airplane.ply", "body, wing, tail"],
    ["examples/guitar.ply", "headstock, neck, body"],
    ["examples/table.ply", "desktop, leg or support, drawer"],
    ["examples/lamp.ply", "base, lampshade, pole"],
    ["examples/mug.ply", "handle, cup"],
    ["examples/bunny.obj", "ear, head, torso, foot"],
]

with gr.Blocks(title="PatchAlign3D") as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            """
# PatchAlign3D · zero-shot 3D part segmentation

Name the parts you want **in words** and see them highlighted on the 3D shape. One forward pass of a
point-cloud encoder whose *patch* features are aligned to CLIP text space — no test-time multi-view rendering.

[Paper](https://huggingface.co/papers/2601.02457) · [Code](https://github.com/souhail-hadgi/PatchAlign3D)
· [Weights](https://huggingface.co/patchalign3d/patchalign3d-encoder)
· [Project page](https://souhail-hadgi.github.io/patchalign3dsite)
"""
        )

        with gr.Row():
            with gr.Column(scale=2):
                shape_file = gr.File(
                    label="3D shape — mesh or point cloud",
                    file_types=[".obj", ".glb", ".gltf", ".stl", ".off", ".ply",
                                ".npz", ".npy", ".txt", ".xyz", ".pts"],
                    type="filepath",
                )
                preview = gr.Plot(label="Input")
                preset = gr.Dropdown(
                    label="Part-vocabulary preset (fills the box below)",
                    choices=list(PRESETS.keys()), value="— custom —",
                )
                labels_text = gr.Textbox(
                    label="Part queries (comma-separated)",
                    value="back, seat, leg, arm",
                    placeholder="back, seat, leg, arm",
                    lines=2,
                )
                run = gr.Button("Segment", variant="primary")

            with gr.Column(scale=3):
                plot = gr.Plot(label="Segmentation")
                summary = gr.Markdown()
                with gr.Row():
                    shares = gr.Label(label="Share of points per part", num_top_classes=MAX_LABELS)
                    ply_out = gr.File(label="Colour-coded point cloud (.ply)")

        with gr.Accordion("Advanced settings", open=False):
            with gr.Row():
                num_points = gr.Slider(512, 8192, value=pa.DEFAULT_NPOINTS, step=512, label="Points sampled")
                num_group = gr.Slider(32, 512, value=pa.DEFAULT_NUM_GROUP, step=32, label="Patches (FPS centres)")
                group_size = gr.Slider(8, 64, value=pa.DEFAULT_GROUP_SIZE, step=8, label="Points per patch")
            with gr.Row():
                text_setting = gr.Radio(
                    ["part_only", "part_plus_cat", "ensemble"], value="part_only",
                    label="Prompt ensemble",
                    info="`part_plus_cat` / `ensemble` also use the object category",
                )
                category = gr.Textbox(label="Object category", value="", placeholder="chair")
            with gr.Row():
                assign = gr.Radio(["nearest", "membership"], value="nearest", label="Patch → point assignment")
                tau = gr.Slider(0.01, 1.0, value=pa.DEFAULT_TAU, step=0.01, label="CLIP temperature τ")
                seed = gr.Number(label="Sampling seed", value=0, precision=0)

        inputs = [shape_file, labels_text, num_points, num_group, group_size,
                  text_setting, category, assign, tau, seed]
        outputs = [plot, shares, ply_out, summary]

        gr.Examples(
            examples=EXAMPLES,
            inputs=[shape_file, labels_text],
            outputs=outputs,
            fn=segment,
            cache_examples=True,
            cache_mode="lazy",
            label="Examples · ShapeNetPart test shapes and the Stanford Bunny mesh",
        )

        gr.Markdown(
            """
### How it works
Points are centred and scaled to the unit sphere and the Y/Z axes are swapped to match the training
convention (exactly as in the official `infer.py`). Furthest-point sampling picks patch centres, a k-NN
neighbourhood around each becomes a patch token, and a 12-layer point transformer produces one feature
per patch. A learned linear head projects those into the CLIP `ViT-bigG-14 (laion2b_s39b_b160k)` text
space, where they are matched against the prompt ensemble `{"<part>", "a <part>", "<part> part"}`.
Each point takes the label of its nearest patch centre.

Every query is *forced* to win somewhere, so asking for a part the shape does not have will still colour
something — that is expected for open-vocabulary matching. Shapes close to the ShapeNetPart categories
work best; the Bunny is there to show that arbitrary meshes go through the same path.
"""
        )

    shape_file.change(preview_shape, inputs=[shape_file, num_points, seed], outputs=preview,
                      api_name="preview")
    preset.change(apply_preset, inputs=[preset, labels_text, category], outputs=[labels_text, category],
                  api_name=False)
    run.click(segment, inputs=inputs, outputs=outputs, api_name="segment")
    labels_text.submit(segment, inputs=inputs, outputs=outputs, api_name=False)

if __name__ == "__main__":
    demo.launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True)