import gradio as gr
import clip
import torch
import numpy as np
import models.vqvae as vqvae
import models.t2m_trans as trans
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from mpl_toolkits.mplot3d import Axes3D
import sys
import os
from matplotlib.animation import FuncAnimation
from mpl_toolkits.mplot3d import Axes3D
from mpl_toolkits.mplot3d.art3d import Poly3DCollection

import warnings
warnings.filterwarnings('ignore')

device = "cuda" if torch.cuda.is_available() else "cpu"

vqvae_model = vqvae.VQVAE().to(device)
transformer_model = trans.Text2Motion_Transformer().to(device)
vqvae_model.load_state_dict(torch.load("output/VQVAE_imp_resnet_100k_hml3d/net_last.pth", map_location=device))
transformer_model.load_state_dict(torch.load("output/net_best_fid.pth", map_location=device))
vqvae_model.eval()
transformer_model.eval()

mean = torch.from_numpy(np.load('output/Mean.npy', allow_pickle=True)).to(device)
std = torch.from_numpy(np.load('output/Std.npy', allow_pickle=True)).to(device)

def generate_motion(text, vqvae_model, transformer_model):
    clip_text = [text]
    text_encoded = clip.tokenize(clip_text, truncate=True).to(device)
    with torch.no_grad():
        motion_indices = transformer_model.sample(text_encoded, False)
        pred_pose = vqvae_model.forward_decoder(motion_indices)
        pred_xyz = recover_from_ric((pred_pose * std + mean).float(), 22)
    return pred_xyz.cpu().numpy().reshape(-1, 22, 3)

def create_animation(joints, title="3D Motion", save_path="static/animation.gif"):
    fig = plt.figure(figsize=(10, 10))
    ax = fig.add_subplot(111, projection='3d')
    data = np.array(joints).T
    line, = ax.plot(data[0, 0:1], data[1, 0:1], data[2, 0:1])

    def update(num, data, line):
        line.set_data(data[:2, :num])
        line.set_3d_properties(data[2, :num])
        return line,

    ani = FuncAnimation(fig, update, frames=len(joints), fargs=(data, line), interval=50, blit=True)
    ani.save(save_path, writer=PillowWriter(fps=20))
    plt.close(fig)
    return save_path
    
examples = [
    "Person doing yoga",
    "A person is dancing ballet",
]

def infer(text):
    motion_data = generate_motion(text, vqvae_model, transformer_model)
    gif_path = create_animation(motion_data)
    return gif_path

with gr.Blocks(css=".container { max-width: 800px; margin: auto; }") as demo:
    with gr.Column():
        gr.Markdown("## 3D Motion Generation on " + ("GPU" if device == "cuda" else "CPU"))
        text_input = gr.Textbox(label="Describe the action", placeholder="Enter text description for the action here...")
        output_image = gr.Image(label="Generated Motion Animation")
        submit_button = gr.Button("Generate Motion")
        
        submit_button.click(
            fn=infer,
            inputs=[text_input],
            outputs=[output_image]
        )

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