import os
import sys
import json
import socket
import tempfile
import pathlib
from dataclasses import dataclass
from typing import List, Dict, Any
import pickle

import cv2
import torch
import numpy as np
import imageio
import gradio as gr
import spaces
import smplx
import pyrender
import trimesh
import trimesh.transformations as tra
import requests
from download_precomputed import download_and_extract_precomputed

# if intermediate_results/8_kid_crossing doesn't exist, download it
if not os.path.exists('intermediate_results/8_kid_crossing'):
    download_and_extract_precomputed()

# Constants and configuration
CHECKPOINT_PATH = './models/checkpoint_150.pt'
NUM_SAMPLES = 50
DISPLAYED_PREDS = 3
FRAME_LIMIT = 30
FPS = 60
DESCRIPTION = "# SkeletonDiffusion Demo"

# Create necessary directories
for dir_name in ['downloads', 'predictions', 'vis', 'intermediate_results', 'outputs', 'assets']:
    os.makedirs(dir_name, exist_ok=True)

# Create a simple loading image if it doesn't exist
LOADING_IMAGE_PATH = os.path.join('assets', 'loading.png')
if not os.path.exists(LOADING_IMAGE_PATH):
    # Create a simple loading image with text
    img = np.zeros((200, 400, 3), dtype=np.uint8)
    img.fill(255)  # White background
    cv2.putText(img, "Processing...", (50, 100), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 0), 2)
    cv2.putText(img, "Please wait", (50, 150), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 0), 2)
    cv2.imwrite(LOADING_IMAGE_PATH, img)

# Available options for displayed predictions
DISPLAYED_PREDS_OPTIONS = [2, 3, 4, 5, 6]

# Fix display and setup issues
os.environ['PYOPENGL_PLATFORM'] = 'egl'
os.system('export IMAGEMAGICK_BINARY=./magick')
os.system('bash ./SkeletonDiffusion_demo/setup_headless.bash')

# Download ImageMagick if not present
MAGICK_PATH = "./magick"
if not os.path.exists(MAGICK_PATH):
    response = requests.get("https://imagemagick.org/archive/binaries/magick", stream=True)
    if response.status_code == 200:
        with open(MAGICK_PATH, "wb") as f:
            for chunk in response.iter_content(chunk_size=8192):
                if chunk:
                    f.write(chunk)
        print(f"Download completed and saved as '{MAGICK_PATH}'")
    else:
        print(f"Download failed with status code: {response.status_code}")
os.system('chmod +x ./magick')

# Setup device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# Local imports
import SkeletonDiffusion_demo.plot_several_meshes as plot_several_meshes
import SkeletonDiffusion_demo.combine_video as combine
from SkeletonDiffusion.src.data.skeleton import create_skeleton
from SkeletonDiffusion.src.eval_prepare_model import get_prediction, load_model_config_exp, prepare_model
from src_joints2smpl_demo.joints2smpl.convert_joints2smpl import process_motion
from SkeletonDiffusion.src.metrics.ranking import get_closest_and_nfurthest_maxapd
@dataclass
class SMPLParams:
    """Data structure to hold SMPL parameters."""
    global_orient: torch.Tensor
    body_pose: torch.Tensor
    betas: torch.Tensor
    transl: torch.Tensor
    joints3d: torch.Tensor

def handle_video_input(video_file: str) -> str:
    """Handle video input from either a local file or YouTube URL.
    
    Args:
        video_file: Path to local video file
        
    Returns:
        str: Path to the video file
    """
    if video_file:
        return video_file
    return None

def correct_vertices(vertices: np.ndarray) -> np.ndarray:
    """Correct SMPL vertices to convert from SMPL to renderer coordinate system.
    
    Applies a rotation about the Y-axis by 180 degrees so that the original +X axis
    (face direction) is transformed to the -Z axis.
    
    Args:
        vertices: SMPL vertices in shape (1, N, 3)
        
    Returns:
        np.ndarray: Corrected vertices in shape (1, N, 3)
    """
    angle = np.radians(180)
    R = tra.rotation_matrix(angle, [1, 0, 0])
    vertices_homo = np.hstack([vertices[0], np.ones((vertices[0].shape[0], 1))])
    vertices_corrected = (R @ vertices_homo.T).T
    return vertices_corrected[:, :3].reshape(1, -1, 3)

def render_smpl(vertices: np.ndarray, width: int, height: int) -> np.ndarray:
    """Render SMPL 3D model using PyRender.
    
    Args:
        vertices: SMPL vertices in shape (1, N, 3)
        width: Output image width
        height: Output image height
        
    Returns:
        np.ndarray: Rendered image in BGR format
    """
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    smpl_model = smplx.create("./models/SMPL_NEUTRAL.pkl", model_type="smpl", gender="neutral").to(device)
    
    vertices_corrected = correct_vertices(vertices)
    mesh = trimesh.Trimesh(vertices_corrected[0], smpl_model.faces)
    
    scene = pyrender.Scene(bg_color=[1.0, 1.0, 1.0, 0.9])
    mesh_node = pyrender.Mesh.from_trimesh(mesh)
    scene.add(mesh_node)

    camera = pyrender.OrthographicCamera(xmag=1.0, ymag=1.0)
    camera_pose = np.eye(4)
    camera_pose[:3, 3] = [0, 0, 5.0]  # distance = 5.0
    scene.add(camera, pose=camera_pose)

    renderer = pyrender.OffscreenRenderer(width, height)
    color, _ = renderer.render(scene)
    return cv2.cvtColor(color, cv2.COLOR_RGB2BGR)

def save_intermediate_results(video_name: str, results: Dict[str, Any], displayed_preds: int = DISPLAYED_PREDS):
    """Save intermediate results for a video.
    
    Args:
        video_name: Name of the video file
        results: Dictionary containing intermediate results
        displayed_preds: Number of displayed predictions
    """
    base_name = os.path.splitext(os.path.basename(video_name))[0]
    results_dir = os.path.join('intermediate_results', base_name)
    os.makedirs(results_dir, exist_ok=True)
    
    # Save results with displayed_preds in filename
    results_path = os.path.join(results_dir, f'results_{displayed_preds}.pkl')
    with open(results_path, 'wb') as f:
        pickle.dump(results, f)

def load_intermediate_results(video_name: str, displayed_preds: int = DISPLAYED_PREDS) -> Dict[str, Any]:
    """Load intermediate results for a video.
    
    Args:
        video_name: Name of the video file
        displayed_preds: Number of displayed predictions
        
    Returns:
        Dictionary containing intermediate results or None if not found
    """
    base_name = os.path.splitext(os.path.basename(video_name))[0]
    results_path = os.path.join('intermediate_results', base_name, f'results_{displayed_preds}.pkl')
    
    if os.path.exists(results_path):
        # Set default tensor type to CPU before loading
        torch.set_default_tensor_type(torch.FloatTensor)
        with open(results_path, 'rb') as f:
            # Use map_location to ensure tensors are loaded on CPU
            results = torch.load(f, map_location='cpu')
        # Reset default tensor type
        torch.set_default_tensor_type(torch.cuda.FloatTensor)
        return results
    return None

@spaces.GPU(duration=60)
def process_video_gpu(video_file: str, displayed_preds: int = DISPLAYED_PREDS) -> tuple:
    """GPU version of process_video that does the actual processing."""
    import time
    start_time = time.time()
    
    # Load models
    model_start_time = time.time()
    smpl_model = smplx.create("./models/SMPL_NEUTRAL.pkl", model_type="smpl", gender="neutral").to(device)
    nlf_model = torch.jit.load("./models/nlf_l_multi.torchscript").to(device).eval()
    print(f"Time for model loading: {time.time() - model_start_time:.2f}s")
    
    # Handle video input
    input_start_time = time.time()
    input_path = handle_video_input(video_file)
    if not input_path:
        return None, None
    print(f"Time for video input handling: {time.time() - input_start_time:.2f}s")

    # Create output path in outputs directory
    base_name = os.path.splitext(os.path.basename(video_file))[0]
    output_path = os.path.join('outputs', f'{base_name}_smpl_{displayed_preds}.gif')

    # Process frames
    frame_start_time = time.time()
    cap = cv2.VideoCapture(input_path)
    frame_count = 0
    smpl_params_list = []
    nlf_frames = []  # Store NLF detection frames
    
    # Initialize time counters
    total_nlf_time = 0
    total_smpl_time = 0
    total_frame_time = 0
    
    while cap.isOpened() and frame_count < FRAME_LIMIT:
        frame_process_start = time.time()
        ret, frame = cap.read()
        if not ret:
            break

        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        image_tensor = torch.from_numpy(frame_rgb).permute(2, 0, 1).int().to(device)
        
        # SMPL detection
        nlf_start_time = time.time()
        with torch.inference_mode():
            pred = nlf_model.detect_smpl_batched(image_tensor.unsqueeze(0))

        pose_params = pred["pose"][0].cpu().numpy()
        betas = pred["betas"][0].cpu().numpy()
        transl = pred["trans"][0].cpu().numpy()
        joints3d = pred['joints3d'][0].cpu().numpy()
        
        nlf_time = time.time() - nlf_start_time
        total_nlf_time += nlf_time
        print(f"Time for NLF model inference: {nlf_time:.2f}s")
        

        if pose_params.shape[0] == 0:
            print(f"No SMPL detected in frame {frame_count}")
            nlf_frames.append(frame_rgb)
            continue

        if pose_params.shape[0] > 0:
            # SMPL model
            render_start_time = time.time()
            smpl_param = SMPLParams(
                global_orient=torch.tensor(pose_params[:, :3]).to(device),
                body_pose=torch.tensor(pose_params).to(device),
                betas=torch.tensor(betas).to(device),
                transl=torch.tensor(transl).to(device),
                joints3d=torch.tensor(joints3d[:, 0:22, :3]).to(device),
            )
            output_smpl = smpl_model(
                global_orient=torch.tensor(pose_params[:, :3]).to(device),
                body_pose=torch.tensor(pose_params[:, 3:]).to(device),
                betas=torch.tensor(betas).to(device),
                transl=torch.tensor(transl).to(device),
                joints3d=torch.tensor(joints3d[:, 0:66]).to(device),
            )

            smpl_time = time.time() - render_start_time
            total_smpl_time += smpl_time
            print(f"Time for SMPL model: {smpl_time:.2f}s")
            
        smpl_params_list.append(smpl_param)
        nlf_frames.append(frame_rgb)  # Store original frame for NLF visualization
        frame_count += 1
        
        frame_time = time.time() - frame_process_start
        total_frame_time += frame_time
        print(f"Total time for frame {frame_count}: {frame_time:.2f}s")

    cap.release()
    gr.Info("Video-to-motion processing completed!")
    print(f"\nTime statistics for {frame_count} frames:")
    print(f"Average NLF model time per frame: {total_nlf_time/frame_count:.2f}s")
    print(f"Average SMPL model time per frame: {total_smpl_time/frame_count:.2f}s")
    print(f"Average total time per frame: {total_frame_time/frame_count:.2f}s")
    print(f"Total time for all frame processing: {time.time() - frame_start_time:.2f}s")

    # Serialize SMPL parameters
    serial_start_time = time.time()
    smpl_params_serialized = [
        {
            "global_orient": p.global_orient.tolist(),
            "body_pose": p.body_pose.tolist(),
            "betas": p.betas.tolist(),
            "transl": p.transl.tolist(),
            "joints3d": p.joints3d.tolist(),
        }
        for p in smpl_params_list
    ]
    print(f"Time for parameter serialization: {time.time() - serial_start_time:.2f}s")
    
    print(f"Total time: {time.time() - start_time:.2f}s")
    
    # Save SMPL params as JSON
    with open('smpl_params.json', 'w') as f:
        json.dump(smpl_params_serialized, f)
    
    # Save intermediate results
    results = {
        'output_path': output_path,
        'smpl_params_serialized': smpl_params_serialized,
        'nlf_frames': nlf_frames,
        'smpl_params_list': smpl_params_list
    }
    save_intermediate_results(video_file, results, displayed_preds)
        
    return output_path, smpl_params_serialized

@spaces.GPU(duration=200)
def generate_motion_video_gpu(smpl_params_json: List[Dict[str, Any]], video_file: str, displayed_preds: int = DISPLAYED_PREDS) -> str:
    """GPU version of generate_motion_video that does the actual processing."""
    import time
    start_time = time.time()
    
    # Setup device
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    
    # Create output path in outputs directory with displayed_preds in filename
    base_name = os.path.splitext(os.path.basename(video_file))[0]
    output_path = os.path.join('outputs', f'{base_name}_motion_{displayed_preds}.gif')
    
    # Load input video frames
    input_frames = []
    cap = cv2.VideoCapture(video_file)
    frame_count = 0
    while cap.isOpened() and frame_count < FRAME_LIMIT:
        ret, frame = cap.read()
        if not ret:
            break
        input_frames.append(frame)
        frame_count += 1
    cap.release()
    
    # Ensure we have exactly FRAME_LIMIT frames
    if len(input_frames) < FRAME_LIMIT:
        # Pad with the last frame if needed
        last_frame = input_frames[-1] if input_frames else np.zeros((480, 640, 3), dtype=np.uint8)
        input_frames.extend([last_frame] * (FRAME_LIMIT - len(input_frames)))
    elif len(input_frames) > FRAME_LIMIT:
        # Trim to FRAME_LIMIT frames
        input_frames = input_frames[:FRAME_LIMIT]
    
    # Deserialize JSON back into SMPLParams objects
    smpl_params_list = [
        SMPLParams(
            global_orient=torch.tensor(p["global_orient"]).to(device),
            body_pose=torch.tensor(p["body_pose"]).to(device),
            betas=torch.tensor(p["betas"]).to(device),
            transl=torch.tensor(p["transl"]).to(device),
            joints3d=torch.tensor(p["joints3d"]).to(device),
        )
        for p in smpl_params_json
    ]
    print(f"Time for deserialization: {time.time() - start_time:.2f}s")
    
    # Calculate number of frames for 0.5 seconds
    # frames_for_half_second = int(FPS * 0.5)
    frames_for_half_second = 30
    # Collect last 0.5 seconds of frames
    obs = torch.stack([p.joints3d[0] for p in smpl_params_list[-frames_for_half_second:]])
    
    # Pad with first frame if needed
    if len(obs) < frames_for_half_second:
        padding = torch.stack([obs[0]] * (frames_for_half_second - len(obs)))
        obs = torch.cat([padding, obs], dim=0)
    
    # Load model and prepare data
    model_start_time = time.time()
    config, exp_folder = load_model_config_exp(CHECKPOINT_PATH)
    config['checkpoint_path'] = CHECKPOINT_PATH
    skeleton = create_skeleton(**config)   

    model, device, *_ = prepare_model(config, skeleton, **config)
    
    # Convert from mm to m and prepare input
    obs = obs / 1000.0
    obs = obs.reshape(1, frames_for_half_second, 22, 3).to(device)
    obs = torch.stack([-obs[..., 2], obs[..., 0], -obs[..., 1]], dim=-1)
    obs = obs - obs[..., 0:1, :]
    obs_in = skeleton.tranform_to_input_space(obs).to(device)
    
    # Get predictions
    pred = get_prediction(obs_in, model, num_samples=NUM_SAMPLES, **config)
    pred = torch.cat((torch.zeros(1, NUM_SAMPLES, pred.shape[2], 1, 3).to(device), pred), dim=3)
    obs_in = torch.cat((torch.zeros(1, frames_for_half_second, 1, 3).to(device), obs_in), dim=2)
    
    print(f"Time for model inference: {time.time() - model_start_time:.2f}s")

    # Convert predictions to SMPL parameters
    smpl_start_time = time.time()
    
    print(f"Time for SMPL conversion: {time.time() - smpl_start_time:.2f}s")
    
    # Prepare data for visualization
    obs_in_50 = obs_in.unsqueeze(1).repeat(1, NUM_SAMPLES, 1, 1, 1)
    obs_and_pred = torch.cat((obs_in_50, pred), dim=2)
    
    # Save predictions
    pred_np = obs_and_pred.cpu().numpy()
    os.makedirs('predictions', exist_ok=True)
    np.save('predictions/joints3d.npy', pred_np)
    print(f"Joints3D data saved to predictions/joints3d.npy")
    
    # Calculate metrics and select best samples
    metric_start_time = time.time()
    from SkeletonDiffusion.src.metrics.body_realism import limb_stretching_normed_rmse
    
    limbstretching = limb_stretching_normed_rmse(
        pred[..., 1:, :], 
        target=obs[0, ..., 1:, :].unsqueeze(0), 
        limbseq=skeleton.get_limbseq(), 
        reduction='persample', 
        obs_as_target=True
    )
    
    # Sort samples by limb stretching and take the half with smallest values
    limbstretching_sorted, indices = torch.sort(limbstretching.squeeze(1), dim=-1, descending=False)
    half_size = len(indices[0]) // 2
    best_half_indices = indices[0, :15]
    
    # Get predictions for best half
    y_pred = pred.squeeze(0)[best_half_indices]
    y_gt = y_pred[0].unsqueeze(0)
    
    # Use get_closest_and_nfurthest_maxapd to select diverse samples
    _, _, top_indices = get_closest_and_nfurthest_maxapd(y_pred, y_gt, nsamples=displayed_preds)
    
    print(f"Selected {len(best_half_indices)} samples with smallest limb stretching")
    print(f"Selected {len(top_indices)} diverse samples from best half")

    # Generate visualization
    vis_start_time = time.time()
    with torch.no_grad():
        # Create video-specific obj directory
        obj_dir = os.path.join('outputs', f'{base_name}_obj/')
        os.makedirs(obj_dir, exist_ok=True)
        process_motion("smpl_params.json", "predictions/joints3d.npy", device=device, sorted_idx=top_indices, output_dir=obj_dir)
    
    print(f"Checking obj_dir contents: {obj_dir}")
    if os.path.exists(obj_dir):
        print("Contents of obj_dir:")
        for root, dirs, files in os.walk(obj_dir):
            print(f"Directory: {root}")
            print(f"Files: {files}")
            print(f"Subdirectories: {dirs}")
    
    plot_several_meshes.main(obj_dir, displayed_preds)
    
    # Show completion message before combining video
    gr.Info("Step 3/3: Combine GIFs")
    # Save the motion video to the output path with input frames as picture-in-picture
    output_path = combine.combine_video(obj_dir, output_path, input_frames, displayed_preds)
    
    return output_path

def process_video(video_file: str, displayed_preds: int = DISPLAYED_PREDS) -> tuple:
    """Process input video to extract SMPL parameters and generate visualization."""
    # Check if we have pre-computed results
    pre_computed = load_intermediate_results(video_file, displayed_preds)
    if pre_computed is not None:
        print("Using pre-computed results")
        return pre_computed['output_path'], pre_computed['smpl_params_serialized']
    
    # If no pre-computed results, use GPU processing
    return process_video_gpu(video_file, displayed_preds)

def generate_motion_video(smpl_params_json: List[Dict[str, Any]], video_file: str = None, displayed_preds: int = DISPLAYED_PREDS) -> str:
    """Generate a motion video from SMPL parameters."""
    if smpl_params_json is None:
        raise ValueError("No SMPL parameters provided. Please process a video first.")
    
    # Check if we have pre-computed results
    if video_file:
        pre_computed = load_intermediate_results(video_file, displayed_preds)
        if pre_computed is not None and 'motion_video_path' in pre_computed:
            print("Using pre-computed motion video")
            return pre_computed['motion_video_path']
    
    # If no pre-computed results, use GPU processing
    return generate_motion_video_gpu(smpl_params_json, video_file, displayed_preds)

def video_to_gif(video_path, gif_path, frame_limit=30):
    """Convert video to GIF with specified frame limit and smooth looping."""
    frames = []
    # lower resolution
    reader = imageio.get_reader(video_path)
    for i, frame in enumerate(reader):
        if i >= frame_limit:
            break
        frames.append(frame)
    # Set FPS to 15 for smooth looping of 30 frames (2 seconds per loop)
    imageio.mimsave(gif_path, frames, fps=15, loop=0)  # loop=0 means infinite loop
    return gif_path

def concat_gifs_side_by_side(gif1_path, gif2_path, output_path):
    """Pad both GIFs to the same (max) height, center them vertically, then concatenate side by side."""
    gif1 = imageio.mimread(gif1_path)
    gif2 = imageio.mimread(gif2_path)
    
    # Ensure both GIFs have the same number of frames
    if len(gif1) != len(gif2):
        print(f"Warning: GIFs have different frame counts ({len(gif1)} vs {len(gif2)}). Adjusting to match.")
        # Use the shorter length
        min_frames = min(len(gif1), len(gif2))
        gif1 = gif1[:min_frames]
        gif2 = gif2[:min_frames]
    
    frames = []
    for f1, f2 in zip(gif1, gif2):
        # Convert both frames to RGB if needed (handle RGBA with alpha channel)
        if f1.shape[2] == 4:
            f1 = f1[..., :3]
        if f2.shape[2] == 4:
            f2 = f2[..., :3]
        h1, w1, c1 = f1.shape
        h2, w2, c2 = f2.shape
        max_h = max(h1, h2)
        # Pad f1 to max_h, vertically centered
        pad_top1 = (max_h - h1) // 2
        pad_bot1 = max_h - h1 - pad_top1
        f1_pad = np.pad(f1, ((pad_top1, pad_bot1), (0, 0), (0, 0)), mode='constant', constant_values=255)
        # Pad f2 to max_h, vertically centered
        pad_top2 = (max_h - h2) // 2
        pad_bot2 = max_h - h2 - pad_top2
        f2_pad = np.pad(f2, ((pad_top2, pad_bot2), (0, 0), (0, 0)), mode='constant', constant_values=255)
        # Concatenate horizontally
        frame = np.concatenate([f1_pad, f2_pad], axis=1)
        frames.append(frame)
    imageio.mimsave(output_path, frames, fps=15, loop=0)
    return output_path

def create_gradio_interface():
    """Create and configure the Gradio interface."""
    with gr.Blocks(css="style.css") as demo:
        gr.Markdown(DESCRIPTION)
        
        # Add user instructions
        gr.Markdown("""
        Demo for the CVPR2025 paper "Nonisotropic Gaussian Diffusion for Realistic 3D Human Motion Prediction", available [here](https://ceveloper.github.io/publications/skeletondiffusion/). Codebase released on [GitHub](https://github.com/Ceveloper/SkeletonDiffusion/tree/main). 
        
        SkeletonDiffusion takes as input a sequence of 3D body joints coordinates, which not everyone has at disposal. In this demo, we use a publicly available model, Neural Localizer Fields ([NLF](https://istvansarandi.com/nlf/)) to extract 3D poses from a given input video. We feed the extracted poses to SkeletonDiffusion to generate corresponding future motions. Note that the poses extracted from the video are noisy and imperfect, but SkeletonDiffusion has been trained only with precise sensor data obtained in laboratory settings. 
        Despite never having seen noisy data and various real-world actions (ballet, basketball, etc.), SkeletonDiffusion can handle most cases reasonably!
        
        ### Instructions
        1. Upload a video or select from examples
        2. Choose whether to use precomputed results (if available)
        3. Select the number of motion predictions to display (2-6)
        4. Click "Run Skeleton Diffusion" to start
        
        **Note:**
        - SkeletonDiffusion requires less than half a second for a forward pass, but extracting the poses from RGB and rendering the output are time consuming 
        - Only the first 30 frames of the input video will be used
        - The first 0.5 seconds of motion will be used to predict future motion
        - Processing time depends on video length and selected number of predictions
        - Precomputed results will be much faster if available
        """)

        with gr.Tabs():
            with gr.Tab("Video Processing"):
                with gr.Row():
                    input_video = gr.Video(label="Input Video", height=600)
                with gr.Row():
                    gr.Examples(
                        examples=sorted(pathlib.Path("downloads").glob("*.mp4")),
                        inputs=input_video,
                        cache_examples=False,
                    )
                with gr.Row():
                    use_precomputed = gr.Checkbox(
                        label="Use precomputed results if available",
                        value=True,
                        info="If checked, will use existing results instead of processing the video again"
                    )
                    displayed_preds = gr.Dropdown(
                        choices=DISPLAYED_PREDS_OPTIONS,
                        value=DISPLAYED_PREDS,
                        label="Number of displayed predictions",
                        info="Select how many motion predictions to display (2-6)"
                    )
                with gr.Row():
                    process_btn = gr.Button("Run Skeleton Diffusion")
                
                # Two-column output: left=input, right=output
                with gr.Row():
                    with gr.Column():
                        input_video_display = gr.Image(label="Input Video (Preview GIF)", height=600)
                    with gr.Column():
                        output_video = gr.Image(label="Generated Motion", height=600)
                
                # Download buttons
                with gr.Row():
                    with gr.Column():
                        download_motion_btn = gr.Button("Download Motion Video")
                        download_motion_btn.click(
                            fn=lambda video_file, displayed_preds: os.path.join('outputs', f'{os.path.splitext(os.path.basename(video_file))[0]}_motion_{displayed_preds}.gif') if video_file else None,
                            inputs=[input_video, displayed_preds],
                            outputs=[gr.File(label="Download Motion Video")]
                        )
                    with gr.Column():
                        download_data_btn = gr.Button("Download Motion Data (SMPL + Joints)")
                        download_data_btn.click(
                            fn=lambda video_file, displayed_preds: os.path.join('intermediate_results', os.path.splitext(os.path.basename(video_file))[0], f'results_{displayed_preds}.pkl') if video_file else None,
                            inputs=[input_video, displayed_preds],
                            outputs=[gr.File(label="Download Motion Data")]
                        )

        def process_video_with_notification(video_file, use_precomputed, displayed_preds):
            # Step 1: Show input video as GIF immediately with exactly 30 frames
            gr.Info("Converting input video to preview GIF...")
            gif_path = os.path.join(tempfile.gettempdir(), f"input_preview_{os.path.splitext(os.path.basename(video_file))[0]}.gif")
            video_to_gif(video_file, gif_path, frame_limit=30)  # Explicitly set to 30 frames
            yield gif_path, LOADING_IMAGE_PATH

            base_name = os.path.splitext(os.path.basename(video_file))[0]
            output_path = os.path.join('outputs', f'{base_name}_motion_{displayed_preds}.gif')
            obs_gif_path = os.path.join('outputs', f'{base_name}_obj', 'shadow_gif', 'obs_obj_tranp.gif')
            concat_gif_path = os.path.join(tempfile.gettempdir(), f"concat_{base_name}.gif")

            # If using precomputed and obs_obj_tranp.gif exists, show concatenated GIF immediately
            if use_precomputed and os.path.exists(obs_gif_path):
                concat_gifs_side_by_side(gif_path, obs_gif_path, concat_gif_path)
                # Replace the preview with the concatenated GIF
                yield concat_gif_path, LOADING_IMAGE_PATH
                # ... continue with rest of workflow
                # (The rest of the workflow remains unchanged)
                if os.path.exists(output_path):
                    gr.Info("Found precomputed video.\nUsing existing results...")
                    yield concat_gif_path, output_path
                    return
                # ... existing code ...

            # If not precomputed, after obs_obj_tranp.gif is generated, show concatenated GIF
            # Continue with the rest of the workflow as before
            # Case 1: Not using precomputed results
            if not use_precomputed:
                gr.Info("Starting video processing with GPU...\nThis may take a few minutes.")
                gr.Info("Step 1/3: Extracting SMPL parameters from video...")
                _, smpl_params = process_video_gpu(video_file, displayed_preds)
                # After SMPL extraction, check if obs_obj_tranp.gif exists and show concatenated GIF
                if os.path.exists(obs_gif_path):
                    concat_gifs_side_by_side(gif_path, obs_gif_path, concat_gif_path)
                    yield concat_gif_path, LOADING_IMAGE_PATH
                gr.Info("Step 2/3: Generating motion predictions...")
                motion_gif = generate_motion_video_gpu(smpl_params, video_file, displayed_preds)
                yield concat_gif_path, motion_gif
                return

            # If no precomputed video, check if we have enough GIFs to generate one
            shadow_gif_dir = os.path.join('outputs', f'{base_name}_obj', 'shadow_gif')
            if os.path.exists(shadow_gif_dir):
                existing_gifs = [f for f in os.listdir(shadow_gif_dir) if f.endswith('_tranp.gif') and not f.startswith('obs')]
                if len(existing_gifs) >= displayed_preds:
                    gr.Info(f"Found {len(existing_gifs)} existing GIFs.\nGenerating video from existing predictions...")
                    # Load input video frames for picture-in-picture
                    input_frames = []
                    cap = cv2.VideoCapture(video_file)
                    frame_count = 0
                    while cap.isOpened() and frame_count < FRAME_LIMIT:
                        ret, frame = cap.read()
                        if not ret:
                            break
                        input_frames.append(frame)
                        frame_count += 1
                    cap.release()
                    # Ensure we have exactly FRAME_LIMIT frames
                    if len(input_frames) < FRAME_LIMIT:
                        # Pad with the last frame if needed
                        last_frame = input_frames[-1] if input_frames else np.zeros((480, 640, 3), dtype=np.uint8)
                        input_frames.extend([last_frame] * (FRAME_LIMIT - len(input_frames)))
                    elif len(input_frames) > FRAME_LIMIT:
                        # Trim to FRAME_LIMIT frames
                        input_frames = input_frames[:FRAME_LIMIT]
                    # If obs_obj_tranp.gif exists, show concatenated GIF
                    if os.path.exists(obs_gif_path):
                        concat_gifs_side_by_side(gif_path, obs_gif_path, concat_gif_path)
                        gif_path = concat_gif_path
                        yield concat_gif_path, LOADING_IMAGE_PATH
                    
                    # Generate video from existing GIFs
                    gr.Info("Combining predictions into final video...")
                    motion_gif = combine.combine_video(
                        os.path.join('outputs', f'{base_name}_obj'),
                        output_path,
                        input_frames,
                        displayed_preds
                    )
                    gr.Info("Processing complete!")
                    yield gif_path, motion_gif
                    return

            # If we don't have enough GIFs, proceed with full processing
            gr.Info("No precomputed results found.\nStarting full video processing...")
            gr.Info("Step 1/3: Extracting SMPL parameters from video...")
            _, smpl_params = process_video_gpu(video_file, displayed_preds)

            gr.Info("Step 2/3: Generating motion predictions...")
            motion_gif = generate_motion_video_gpu(smpl_params, video_file, displayed_preds)
                        # After SMPL extraction, check if obs_obj_tranp.gif exists and show concatenated GIF
            if os.path.exists(obs_gif_path):
                concat_gifs_side_by_side(gif_path, obs_gif_path, concat_gif_path)
                yield concat_gif_path, LOADING_IMAGE_PATH

            yield concat_gif_path, motion_gif
            return

        process_btn.click(
            fn=process_video_with_notification,
            inputs=[input_video, use_precomputed, displayed_preds],
            outputs=[input_video_display, output_video]
        )
    
    return demo

if __name__ == "__main__":
    # Get local IP address
    s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    s.connect(("8.8.8.8", 80))
    print(s.getsockname()[0])
    s.close()
    
    # Create and launch interface
    demo = create_gradio_interface()
    print(demo.get_api_info())
    demo.launch(server_name="0.0.0.0", share=True)