import os

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
os.environ.setdefault("MPLBACKEND", "Agg")

import re
import types
import tempfile

import spaces
import torch
import numpy as np
import gradio as gr
import matplotlib

matplotlib.use("Agg")

from huggingface_hub import hf_hub_download
from transformers import AutoModelForCausalLM, AutoTokenizer

from motionvq.vqvae import HumanVQVAE
from motionvq.motion_utils import recover_from_ric, plot_3d_motion
from motionvq.paramUtil import t2m_kinematic_chain

# ----------------------------------------------------------------------------
# Config -- values verified against the IRG-MotionLLM / Motion-Agent source and
# the actual Stage-3 checkpoint tensor shapes:
#   * base LLM  : google/gemma-2-2b-it (bf16)
#   * added toks : <Motion>, </Motion>, and 512 motion codes  (nb_code = 512)
#   * VQVAE     : HumanML3D 263-dim, codebook 512x512, down_t=2, stride_t=2
#   * Stage-3 checkpoint bundles BOTH the merged full LLM and the VQVAE weights
# ----------------------------------------------------------------------------
BASE_LLM = "google/gemma-2-2b-it"
CKPT_REPO = "Lymann/IRG-MotionLLM-HumanML3D"
CKPT_FILE = "Stage-3/irg_motionllm_unified_rl_stage3.bin"

NB_CODE = 512
JOINTS_NUM = 22
FPS = 20
HF_TOKEN = os.environ.get("HF_TOKEN")


class VQArgs:
    """Minimal args namespace expected by HumanVQVAE / VQVAE_251."""

    dataname = "t2m"
    quantizer = "ema_reset"
    mu = 0.99


VQ_ARGS = VQArgs()

# System prompt for the IRG (interleaved generation-assessment-refinement) task,
# copied verbatim from the repo's TextProcessor (unified_mogen_cot_v3).
IRG_SYSTEM_PROMPT = (
    "You are an assistant who helps users understand or generate 3D human "
    "motion representations."
)
IRG_USER_TEMPLATE = (
    "Given a text outlining a human motion objective, employ a step-by-step "
    "thought process to realize the motion: (1) analyze the text, providing a "
    "clear explanation of the reasoning to identify essential elements; "
    "(2) conduct several rounds of motion generation and self-assessment until "
    "the motion is satisfactory. Wrap all responses in <think> and </think> "
    "tags, and formulate a plan before each step.\nGoal Text: {caption}"
)

# ----------------------------------------------------------------------------
# Load model + VQVAE at module scope (ZeroGPU: .to("cuda") is intercepted)
# ----------------------------------------------------------------------------
print("Loading tokenizer + base LLM ...")
tokenizer = AutoTokenizer.from_pretrained(BASE_LLM, token=HF_TOKEN)

llm = AutoModelForCausalLM.from_pretrained(
    BASE_LLM,
    torch_dtype=torch.bfloat16,
    attn_implementation="sdpa",
    token=HF_TOKEN,
)

# Recreate the exact tokenizer vocabulary the model was trained with:
#   base vocab + <Motion> + </Motion> + <Motion_0> .. <Motion_{NB_CODE-1}>
NB_TEXT_TOKENS = len(tokenizer)
tokenizer.add_tokens(["<Motion>", "</Motion>"])
for i in range(NB_CODE):
    tokenizer.add_tokens([f"<Motion_{i}>"])
llm.resize_token_embeddings(len(tokenizer))

print("Loading Stage-3 checkpoint ...")
ckpt_path = hf_hub_download(CKPT_REPO, CKPT_FILE, token=HF_TOKEN)
state = torch.load(ckpt_path, map_location="cpu")
if isinstance(state, dict) and "state_dict" in state:
    state = state["state_dict"]

# The Stage-3 checkpoint is a full fine-tune (LoRA already merged) and stores
# both the LLM weights (prefixed "llm.") and the VQVAE weights (prefixed "net.").
llm_sd = {}
net_sd = {}
for k, v in state.items():
    if k.startswith("llm."):
        llm_sd[k[len("llm."):]] = v
    elif k.startswith("net."):
        net_sd[k[len("net."):]] = v

missing, unexpected = llm.load_state_dict(llm_sd, strict=False)
print(f"LLM load -> missing={len(missing)} unexpected={len(unexpected)}")
llm = llm.to(torch.bfloat16).eval()

print("Building + loading VQVAE ...")
net = HumanVQVAE(
    VQ_ARGS,
    nb_code=NB_CODE,
    code_dim=512,
    output_emb_width=512,
    down_t=2,
    stride_t=2,
    width=512,
    depth=3,
    dilation_growth_rate=3,
    activation="relu",
    norm=None,
)
vq_missing, vq_unexpected = net.load_state_dict(net_sd, strict=False)
print(f"VQVAE load -> missing={len(vq_missing)} unexpected={len(vq_unexpected)}")
net = net.eval().float()
for p in net.parameters():
    p.requires_grad = False

# HumanML3D normalization statistics (263-dim), extracted from the same
# T2M-GPT extractor bundle the authors use.
MEAN = np.load(os.path.join(os.path.dirname(__file__), "mean.npy"))
STD = np.load(os.path.join(os.path.dirname(__file__), "std.npy"))

llm.to("cuda")
net.to("cuda")

EOS_ID = tokenizer.eos_token_id
MOTION_ID_START = len(tokenizer) - (NB_CODE + 2)  # first <Motion_i> token id


# ----------------------------------------------------------------------------
# Motion-token extraction helpers (ported from mllm_single_lora.py)
# ----------------------------------------------------------------------------
def _find_seq(tokens_list, target):
    for i in range(len(tokens_list) - len(target) + 1):
        if tokens_list[i:i + len(target)] == target:
            return i
    return -1


def _extract_answer_scores(scores, ids_list):
    """Return the score rows and matching ids inside <answer>..</answer>,
    else fall back to the last <Motion>..</Motion> span, else all."""
    ans_start = tokenizer.encode("<answer>", add_special_tokens=False)
    ans_end = tokenizer.encode("</answer>", add_special_tokens=False)
    s = _find_seq(ids_list, ans_start)
    if s != -1:
        s2 = s + len(ans_start)
        e = _find_seq(ids_list[s2:], ans_end)
        if e != -1:
            e += s2
            return scores[s2:e], ids_list[s2:e]
    return scores, ids_list


def _tokens_from_scores_and_ids(scores, ids_list):
    """Prefer already-decoded motion token ids that fall in the motion range;
    otherwise argmax over the motion-logit slice. Mirrors the repo logic."""
    motion_logits = scores[:, -(NB_CODE + 2):]
    argmax_ids = torch.argmax(motion_logits, dim=-1)
    out = []
    for i, tid in enumerate(ids_list):
        if MOTION_ID_START <= tid < len(tokenizer):
            out.append(tid - MOTION_ID_START)
        else:
            out.append(int(argmax_ids[i].item()))
    motion_tokens = torch.tensor(out, dtype=torch.long)
    if 1 in motion_tokens.tolist():
        motion_tokens = motion_tokens[:motion_tokens.tolist().index(1)]
    if 0 in motion_tokens.tolist():
        motion_tokens = motion_tokens[motion_tokens.tolist().index(0) + 1:]
    motion_tokens = torch.clamp(motion_tokens - 2, min=0)
    return motion_tokens


def _extract_last_motion_span_ids(ids_list):
    st_id = tokenizer.encode("<Motion>", add_special_tokens=False)[0]
    ed_id = tokenizer.encode("</Motion>", add_special_tokens=False)[0]
    starts = [i for i, t in enumerate(ids_list) if t == st_id]
    ends = [i for i, t in enumerate(ids_list) if t == ed_id]
    if not starts or not ends:
        return None
    # last valid <Motion> ... </Motion> pair
    for st in reversed(starts):
        later_ends = [e for e in ends if e > st]
        if later_ends:
            return st, later_ends[0]
    return None


def build_prompt(caption):
    caption = (caption or "").strip()
    user_prompt = IRG_USER_TEMPLATE.format(caption=caption)
    text = IRG_SYSTEM_PROMPT + "\n\n" + "User: " + user_prompt + "\n\n" + " Response:"
    return text


def decode_motion_to_joints(motion_tokens):
    """VQVAE token ids -> 263-dim features -> denormalize -> 22-joint xyz."""
    motion_tokens = motion_tokens.to("cuda").long()
    feats = net.forward_decoder(motion_tokens)          # (1, T, 263)
    feats = feats.detach().cpu().numpy()[0]
    feats = MEAN + feats * STD                          # denormalize
    joints = recover_from_ric(torch.from_numpy(feats).float(), JOINTS_NUM)
    return joints.numpy()                               # (T, 22, 3)


def render_video(joints, title):
    out_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
    plot_3d_motion(out_path, t2m_kinematic_chain, joints, title=title, fps=FPS, radius=4)
    return out_path


# ----------------------------------------------------------------------------
# Inference
# ----------------------------------------------------------------------------
def _duration(caption, reasoning, max_new_tokens, *a, **k):
    base = 55
    if reasoning:
        base = 110
    return min(160, base + int(int(max_new_tokens) / 20))


@spaces.GPU(duration=_duration)
def generate(caption, reasoning=True, max_new_tokens=1300, seed=0,
             progress=gr.Progress(track_tqdm=True)):
    if not caption or not caption.strip():
        raise gr.Error("Please enter a motion description.")

    if seed and int(seed) > 0:
        torch.manual_seed(int(seed))

    do_sample = bool(seed and int(seed) > 0)
    max_new_tokens = int(max_new_tokens)

    prompt = build_prompt(caption)
    enc = tokenizer(prompt, return_tensors="pt").to("cuda")
    input_len = enc.input_ids.shape[1]

    if reasoning:
        max_len = input_len + max_new_tokens
    else:
        # skip the interleaved reasoning; ask for a compact motion answer
        max_len = input_len + 260

    gen_kwargs = dict(
        max_length=max_len,
        do_sample=do_sample,
        return_dict_in_generate=True,
        output_scores=True,
        use_cache=True,
    )
    if do_sample:
        gen_kwargs["temperature"] = 1.0

    with torch.inference_mode():
        outputs = llm.generate(enc.input_ids, attention_mask=enc.attention_mask, **gen_kwargs)

    gen_ids = outputs.sequences[0, input_len:]
    # scores are per generated step (tuple len == n_new); stack into (n_new, vocab)
    scores = torch.stack(outputs.scores)[:, 0, :]

    # truncate at EOS
    gen_ids_list = gen_ids.tolist()
    if EOS_ID in gen_ids_list:
        cut = gen_ids_list.index(EOS_ID) + 1
        gen_ids_list = gen_ids_list[:cut]
        scores = scores[:cut]

    full_text = tokenizer.decode(gen_ids_list, skip_special_tokens=False)

    # Prefer <answer>..</answer>; fall back to last <Motion>..</Motion> span.
    ans_scores, ans_ids = _extract_answer_scores(scores, gen_ids_list)
    if ans_ids is gen_ids_list or len(ans_ids) == len(gen_ids_list):
        span = _extract_last_motion_span_ids(gen_ids_list)
        if span is not None:
            st, ed = span
            ans_scores = scores[st + 1:ed]
            ans_ids = gen_ids_list[st + 1:ed]

    motion_tokens = _tokens_from_scores_and_ids(ans_scores, ans_ids)

    if motion_tokens.numel() == 0:
        raise gr.Error(
            "The model did not produce a valid motion for this prompt. "
            "Try rephrasing, or toggle the reasoning option."
        )

    joints = decode_motion_to_joints(motion_tokens)
    title = caption.strip()
    if len(title) > 60:
        title = title[:57] + "..."
    video = render_video(joints, title)

    # Build a readable reasoning trace (strip the internal [plan]/[tag] markup).
    trace = full_text
    trace = trace.replace("<eos>", "").replace("<pad>", "")
    n_rounds = len(re.findall(r"\[generate\]", trace))
    info = (
        f"Frames: {joints.shape[0]}  |  Motion tokens: {motion_tokens.numel()}"
        f"  |  Internal generate/refine rounds: {max(n_rounds, 1)}"
    )
    return video, info, trace.strip()


# ----------------------------------------------------------------------------
# UI
# ----------------------------------------------------------------------------
DESCRIPTION = """
# 🕺 IRG-MotionLLM — Text-to-3D-Motion

Generate 3D human motion from a text description with
**[IRG-MotionLLM](https://huggingface.co/papers/2512.10730)** — an LLM
(Gemma-2-2B) that *interleaves* motion generation, self-assessment and
refinement, then decodes discrete motion tokens through a VQ-VAE into a
HumanML3D skeleton animation.
"""

EXAMPLES = [
    ["a person walks forward, then turns around and walks back.", True, 1300, 0],
    ["a man is doing cartwheels.", True, 1300, 0],
    ["a person jumps up high with both hands raised.", True, 1300, 0],
    ["someone sits down on a chair and crosses their legs.", True, 1300, 0],
    ["a person raises their right hand and waves.", True, 1300, 0],
]

with gr.Blocks(theme=gr.themes.Citrus()) as demo:
    gr.Markdown(DESCRIPTION)

    with gr.Row():
        with gr.Column(scale=1):
            caption = gr.Textbox(
                label="Motion description",
                placeholder="e.g. a person walks forward, then turns around and walks back.",
                lines=2,
            )
            run_btn = gr.Button("Generate motion", variant="primary")
            with gr.Accordion("Advanced options", open=False):
                reasoning = gr.Checkbox(
                    value=True,
                    label="Interleaved reasoning (generate → assess → refine)",
                    info="Uses the full IRG chain-of-thought. Turn off for a faster, direct generation.",
                )
                max_new_tokens = gr.Slider(
                    400, 2000, value=1300, step=50,
                    label="Max new tokens (reasoning budget)",
                )
                seed = gr.Slider(
                    0, 100000, value=0, step=1,
                    label="Seed (0 = greedy / deterministic, >0 = sampling)",
                )
        with gr.Column(scale=1):
            video_out = gr.Video(label="Generated motion", autoplay=True)
            info_out = gr.Textbox(label="Summary", lines=2)

    with gr.Accordion("Model reasoning trace", open=False):
        trace_out = gr.Textbox(label="Interleaved generation / assessment / refinement", lines=10)

    gr.Examples(
        examples=EXAMPLES,
        inputs=[caption, reasoning, max_new_tokens, seed],
        outputs=[video_out, info_out, trace_out],
        fn=generate,
        cache_examples=True,
        cache_mode="lazy",
    )

    run_btn.click(
        fn=generate,
        inputs=[caption, reasoning, max_new_tokens, seed],
        outputs=[video_out, info_out, trace_out],
    )
    caption.submit(
        fn=generate,
        inputs=[caption, reasoning, max_new_tokens, seed],
        outputs=[video_out, info_out, trace_out],
    )

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