"""HandX bimanual text-to-motion — Hugging Face Space demo.

Generates two-hand motion from text (left hand / right hand / interaction) using the
released HandX diffusion checkpoints, and renders a skeleton animation. MANO is not used.
"""
import os
import tempfile

import numpy as np
import torch
import gradio as gr

# Workaround for a gradio_client bug: it crashes parsing JSON schemas where a value is a
# bare bool (e.g. additionalProperties: true). Patch it to degrade gracefully.
import gradio_client.utils as _gcu
_orig_j2p = _gcu._json_schema_to_python_type
def _safe_j2p(schema, defs=None):
    if isinstance(schema, bool):
        return "Any"
    return _orig_j2p(schema, defs)
_gcu._json_schema_to_python_type = _safe_j2p

from einops import rearrange
from omegaconf import OmegaConf
from huggingface_hub import hf_hub_download

# ZeroGPU support: no-op decorator when running outside HF Spaces.
try:
    import spaces
    GPU = spaces.GPU(duration=120)
except Exception:  # noqa: BLE001
    def GPU(fn):
        return fn

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation

from src.diffusion.utils.model_utils import create_model_and_diffusion
from src.diffusion.model.cls_free_sampler import ClassifierFreeSampleWrapper
from src.visualize.skeleton_visualizer import Skeleton_Visualize_Helper

MODEL_REPO = "alexzhang598/HandX-diffusion"
VARIANTS = {"layers12 (best)": "layers12", "layers8": "layers8", "layers4": "layers4"}
MOTION_LENGTH = 60
GUIDANCE = 2.5
REPR_DIMS = {"joint_pos": (42, 3), "joint_pos_w_scalar_rot": (42, 4), "joint_rot": (34, 6)}

_MEAN = np.load(os.path.join(os.path.dirname(__file__), "mean_std", "mean.npy"))
_STD = np.load(os.path.join(os.path.dirname(__file__), "mean_std", "std.npy"))

_cache = {}  # variant -> (model, diffusion, njoints, nfeats)


def _load(variant: str):
    if variant in _cache:
        return _cache[variant]
    cfg = OmegaConf.load(hf_hub_download(MODEL_REPO, f"{variant}/config.yaml"))
    njoints, nfeats = REPR_DIMS[cfg.data.repr]
    model, diffusion = create_model_and_diffusion(cfg.model)
    sd = torch.load(hf_hub_download(MODEL_REPO, f"{variant}/model.pt"), map_location="cpu")["state_dict"]
    model.load_state_dict(sd, strict=False)  # missing keys = frozen T5 (from t5-base)
    model = ClassifierFreeSampleWrapper(model, scale=GUIDANCE).eval()
    _cache[variant] = (model, diffusion, njoints, nfeats)
    return _cache[variant]


@GPU
def generate(variant_label, left_text, right_text, relation_text, seed, guidance):
    variant = VARIANTS[variant_label]
    model, diffusion, njoints, nfeats = _load(variant)

    device = "cuda" if torch.cuda.is_available() else "cpu"
    model = model.to(device)
    model.scale = float(guidance)  # classifier-free guidance strength
    seed = int(seed)
    if seed < 0:
        seed = int(torch.randint(0, 2**31 - 1, (1,)).item())
    torch.manual_seed(seed)

    shape = (1, njoints, nfeats, MOTION_LENGTH)
    model_kwargs = dict(y=dict(
        lengths=torch.as_tensor([MOTION_LENGTH], device=device),
        text=dict(left=[left_text or ""], right=[right_text or ""],
                  two_hands_relation=[relation_text or ""]),
    ))
    with torch.no_grad():
        samples = diffusion.p_sample_loop(model, shape, clip_denoised=False,
                                          model_kwargs=model_kwargs, device=device, progress=False)
    samples = rearrange(samples, "b j f t -> b t (j f)").cpu().numpy()
    samples = samples * _STD.reshape(-1) + _MEAN.reshape(-1)  # inv_transform

    motion = samples[0].reshape(MOTION_LENGTH, njoints, nfeats)[:, :, :3]  # (T, 42, 3)
    left_motion, right_motion = motion[:, :21], motion[:, 21:]

    out_path = os.path.join(tempfile.mkdtemp(), "handx.mp4")
    render(left_motion, right_motion, out_path)
    return out_path, f"🎉 Done · seed={seed} · guidance={guidance:g} · {variant}"


def render(left_motion, right_motion, out_path, fps=30):
    """Render a single bimanual skeleton animation as a browser-playable H.264 mp4."""
    fig = plt.figure(figsize=(7, 7))
    ax = fig.add_subplot(111, projection="3d")
    helper = Skeleton_Visualize_Helper(ax, left_motion, right_motion)
    helper.initialize_ax()
    ax.view_init(elev=15, azim=-70)
    fig.tight_layout()

    frames = left_motion.shape[0]

    def update(frame):
        helper.draw(frame)

    ani = FuncAnimation(fig, update, frames=frames)
    # H.264 + yuv420p + faststart so the video plays inline in browsers (moov atom at front),
    # not just downloadable.
    ani.save(out_path, writer="ffmpeg", fps=fps,
             extra_args=["-vcodec", "libx264", "-pix_fmt", "yuv420p", "-movflags", "+faststart"])
    plt.close(fig)


EXAMPLES = [
    ["layers12 (best)",
     "The four fingers slowly straighten out from a bent position and then begin to bend again.",
     "The fingers mirror the left hand's motion, opening and then closing quickly. As the hand closes, the ring finger's middle joint hyperextends.",
     "The hands move in a synchronized manner while remaining far apart and never interacting.", 42, GUIDANCE],
    ["layers12 (best)",
     "The hand holds a static pose, with the only motion being the index and middle fingers closing together at the start.",
     "The hand is mostly still, but the pinky finger quickly bends twice mid-sequence.",
     "The fingertips of the index and middle fingers of both hands repeatedly make contact with each other throughout the motion.", 42, GUIDANCE],
]

APP_CSS = """
.gradio-container {max-width: 1200px !important; margin: auto;}
#title h1 {margin-bottom: 2px;}
#gen-btn {font-size: 1.05rem;}
#video-out video {border-radius: 12px; box-shadow: 0 4px 20px rgba(0,0,0,0.12); background:#000;}
footer {visibility: hidden; height: 0;}
"""

with gr.Blocks(title="HandX Text-to-Motion", theme=gr.themes.Soft(), css=APP_CSS) as demo:
    gr.Markdown(
        "# 🤲 HandX — Bimanual Text-to-Motion\n"
        "Generate two-hand motion from text. Describe the **left hand**, the **right hand**, "
        "and their **interaction**. "
        "[Paper](https://arxiv.org/abs/2603.28766) · "
        "[Dataset](https://huggingface.co/datasets/alexzhang598/HandX) · "
        "[Models](https://huggingface.co/alexzhang598/HandX-diffusion)",
        elem_id="title",
    )
    with gr.Row(equal_height=False):
        with gr.Column(scale=2):
            variant = gr.Dropdown(list(VARIANTS.keys()), value="layers12 (best)", label="Model")
            left = gr.Textbox(label="Left hand", lines=2,
                              placeholder="e.g. The fingers slowly straighten from a bent position.")
            right = gr.Textbox(label="Right hand", lines=2,
                               placeholder="e.g. The fingers mirror the left hand, opening then closing.")
            relation = gr.Textbox(label="Two-hand interaction", lines=2,
                                  placeholder="e.g. The hands move in a synchronized manner.")
            with gr.Accordion("Advanced settings", open=False):
                with gr.Row():
                    seed = gr.Number(value=42, label="Seed (-1 = random)", precision=0, scale=3)
                    rand_btn = gr.Button("🎲", scale=1, min_width=50)
                guidance = gr.Slider(1.0, 10.0, value=GUIDANCE, step=0.1,
                                     label="Guidance scale (CFG)",
                                     info="Higher = follow text more closely")
            btn = gr.Button("Generate motion", variant="primary", elem_id="gen-btn")
            status = gr.Markdown("")
        with gr.Column(scale=3):
            out = gr.Video(label="Generated motion", autoplay=True, loop=True, elem_id="video-out")

    inputs = [variant, left, right, relation, seed, guidance]
    rand_btn.click(lambda: int(torch.randint(0, 2**31 - 1, (1,)).item()), outputs=seed)
    btn.click(lambda: "⏳ Generating motion, please wait…", outputs=status) \
       .then(generate, inputs, [out, status])
    gr.Examples(EXAMPLES, inputs, cache_examples=False)

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