'''
BVH, Nov 2025.
Any4D / AnyView / DVS (dynamic view synthesis) gradio app.

Supports many things as input, including vidar datasets, webdataset tarfiles, canonical
trajectories, MP4 videos, YouTube URLs. There is somewhat complex logic stitching together input
video and camera poses as appropriate from these various sources.

Also supports 3D camera visualizations and uncertainty heatmaps.
'''

import os
import sys
sys.path.insert(0, os.getcwd())

# Library imports
import argparse
import copy
import cv2
import gradio as gr
import glob
import lovely_tensors
import numpy as np
import pathlib
import random
import time
import torch
import warnings
from einops import rearrange
from functools import partial
from lovely_numpy import lo
from rich import print

# Internal imports
from custom.eval import infer_utils, gradio_utils
from custom.any4d.a4d_logistics import pack_streams_from_entries, unpack_entries_from_streams
from custom.any4d.a4d_metrics import calculate_metrics_a4d
from custom.any4d.a4d_visuals import create_save_visuals_a4d
from custom.dataset.vidar_dataset import VidarDataset
from torch.utils.data import DataLoader

lovely_tensors.monkey_patch()
np.set_printoptions(precision=3, suppress=True)
torch.set_printoptions(precision=3, sci_mode=False, threshold=1000)
warnings.filterwarnings('ignore', category=FutureWarning)

_TITLE = 'AnyView ({})'

# Camera trajectory presets
# Format: 'Button Label': [setting, azimuth, input_elev, output_elev, radius]
CAMERA_PRESETS = {
    # Exo2Exo: External camera rotating around scene
    'Rotate Left': ['Exo2Exo', -90, 10, 10, 5],
    'Rotate Right': ['Exo2Exo', 90, 10, 10, 5],
    'Rotate Up': ['Exo2Exo', 0, 10, 60, 5],
    
    # Ego2Ego: First-person view rotating in place
    'Turn Left': ['Ego2Ego', -90, 0, 0, 1],
    'Turn Right': ['Ego2Ego', 90, 0, 0, 1],
    
    # Ego2Exo: First-person to external view (moving camera away)
    'Move Back': ['Ego2Exo', 0, 0, 20, 5],
    'Move Back+Left': ['Ego2Exo', -90, 0, 20, 5],
    'Move Back+Right': ['Ego2Exo', 90, 0, 20, 5],
    'Move Back+Up': ['Ego2Exo', 0, 0, 60, 5],
}

_DESCRIPTION = '''
This demo showcases **AnyView** for extreme monocular *dynamic view synthesis* (DVS).

Currently loaded checkpoint: **{checkpoint_path}**.

Reference dataset config: **{dset_cfg}**.
'''

# Random video directory for the "Load Random Video" button
RANDOM_VIDEO_DIR = '/datasets/basile/av_gradio_rnd'

os.environ['GRADIO_TEMP_DIR'] = '/tmp/gradio_anyview'


def parse_tarfile_camera_indices(tarfile_cam_indices):
    '''
    Parse tarfile camera indices string into camera indices.
    :param tarfile_cam_indices (str): Camera indices in format "in,out" e.g. "1,0"
    :return: Tuple of (cam_in, cam_out, cam_indices_list)
             where cam_indices_list is [cam_out, cam_in] in VIDAR order
    '''
    # Default VIDAR convention
    cam_in = 1   # Input/context camera
    cam_out = 0  # Output/target camera
    
    try:
        if tarfile_cam_indices and tarfile_cam_indices.strip():
            parts = tarfile_cam_indices.strip().split(',')
            if len(parts) == 2:
                cam_in = int(parts[0].strip())
                cam_out = int(parts[1].strip())
                print(f'[cyan]Using custom camera indices: in={cam_in}, out={cam_out}')
    except Exception as e:
        print(f'[yellow]Warning: Failed to parse camera indices "{tarfile_cam_indices}": {e}')
        print(f'[yellow]Falling back to default: in=1, out=0')
    
    # Return in VIDAR order: [output, input]
    cam_indices = [cam_out, cam_in]
    return cam_in, cam_out, cam_indices


def extract_frames_from_rgb_dict(rgb_dict, cam_idx, max_frames=None):
    '''
    Extract RGB frames from an rgb_dict for a specific camera.
    Handles cases where frame indices don't start at 0.
    
    :param rgb_dict (dict): Dictionary with (time_idx, cam_idx) -> tensor mappings
    :param cam_idx (int): Camera index to extract frames for
    :param max_frames (int): Maximum number of frames to extract (None = all)
    :return: Tuple of (stacked_frames tensor, time_indices list, error_msg or None)
    '''
    # Get all time indices for the specified camera (frames might not start at 0)
    time_indices = sorted([key[0] for key in rgb_dict.keys() if key[1] == cam_idx])
    
    if len(time_indices) == 0:
        return None, [], f'No RGB frames found for camera {cam_idx}'
    
    # Limit to max_frames if specified
    if max_frames is not None:
        time_indices = time_indices[:max_frames]
    
    # Extract frames in order of their time indices
    frames_list = []
    for t in time_indices:
        frame = rgb_dict[(t, cam_idx)]  # (B, C, H, W) or (1, C, H, W)
        frames_list.append(frame[0])  # Take first batch element -> (C, H, W)
    
    # Stack frames: (T, C, H, W)
    stacked_frames = torch.stack(frames_list, dim=0)
    
    return stacked_frames, time_indices, None


def override_batch_with_video(data_batch, input_rgb, device):
    '''
    Override RGB frames in a vidar data batch with custom video frames.
    :param data_batch (dict): Data batch from vidar dataset
    :param input_rgb (tensor): Input RGB video (T, C, H, W) in [0, 1]
    :param device (str): Device to use
    :return (dict): Updated data batch
    '''
    import torch
    import torch.nn.functional as F
    
    # Ensure input_rgb is on CPU for interpolation
    if input_rgb.device.type != 'cpu':
        input_rgb = input_rgb.cpu()
    
    # Get number of unique timesteps
    timestep_indices = set([key[0] for key in data_batch['anydata']['timestep'].keys()])
    target_T = len(timestep_indices)
    
    # Derive target resolution from anydata rgb tensors
    first_rgb_key = next(k for k in data_batch['anydata']['rgb'] if isinstance(k, tuple))
    target_H = data_batch['anydata']['rgb'][first_rgb_key].shape[-2]
    target_W = data_batch['anydata']['rgb'][first_rgb_key].shape[-1]

    B = data_batch['anydata']['rgb'][(0, 0)].shape[0] if 'rgb' in data_batch['anydata'] else 1
    
    # Resize frames to match target
    if input_rgb.shape[2] != target_H or input_rgb.shape[3] != target_W:
        frames_tensor = F.interpolate(
            input_rgb,
            size=(int(target_H), int(target_W)),
            mode='bilinear',
            align_corners=False
        )
    else:
        frames_tensor = input_rgb
    
    # Trim or repeat frames to match target length
    if frames_tensor.shape[0] > target_T:
        frames_tensor = frames_tensor[:target_T]
    elif frames_tensor.shape[0] < target_T:
        last_frame = frames_tensor[-1:].repeat(target_T - frames_tensor.shape[0], 1, 1, 1)
        frames_tensor = torch.cat([frames_tensor, last_frame], dim=0)
    
    # Replace rgb frames in vidar dict (camera 1 = input view = rgb1)
    for t in range(target_T):
        key = (t, 1)
        if key in data_batch['anydata']['rgb']:
            data_batch['anydata']['rgb'][key] = frames_tensor[t:t+1].repeat(B, 1, 1, 1).to(
                data_batch['anydata']['rgb'][key].device,
                dtype=data_batch['anydata']['rgb'][key].dtype
            )
    
    # Convert timestep values from int to tensor
    if 'timestep' in data_batch['anydata']:
        for key in data_batch['anydata']['timestep'].keys():
            if isinstance(data_batch['anydata']['timestep'][key], int):
                data_batch['anydata']['timestep'][key] = torch.tensor([data_batch['anydata']['timestep'][key]])
    
    return data_batch


def inference(model_bundle, base_data_batch, input_rgb, output_dir, num_samples=1,
              num_frames=41, num_steps=25, frame_rate=12, skip_override=False):
    '''
    Run model inference on input video using a base vidar data batch.
    :param model_bundle (tuple): (model, config, exp_cfg, device)
    :param base_data_batch (dict): Base data batch from vidar dataset
    :param input_rgb (tensor): Input RGB video (T, C, H, W) in [0, 1]
    :param output_dir (str): Output directory for saving visualizations
    :param num_samples (int): Number of samples to generate
    :param num_frames (int): Number of frames
    :param num_steps (int): Number of diffusion steps
    :param frame_rate (int): Frame rate for video visualization
    :param skip_override (bool): If True, skip overriding camera 1 with input video
    :return (list): List of output samples
    '''
    (model, config, exp_cfg, device) = model_bundle
    
    # Verify dimensions
    if input_rgb is not None:
        (Tc, C, Hp, Wp) = input_rgb.shape
        print(f'[cyan]Input shape: {input_rgb.shape}')
    else:
        print(f'[cyan]Using data_batch as-is (YAML/tarfile source)')
    
    # Find the highest val_iter used in existing files to avoid overwriting
    import glob
    import os
    existing_files = glob.glob(os.path.join(output_dir, '*_val_*_i*_*.mp4'))
    max_val_iter = -1
    for fpath in existing_files:
        # Extract val_iter from filename: ...i{val_iter}_...
        fname = os.path.basename(fpath)
        if '_i' in fname:
            try:
                parts = fname.split('_i')[1]
                val_iter_str = parts.split('_')[0]
                val_iter = int(val_iter_str)
                max_val_iter = max(max_val_iter, val_iter)
            except:
                pass
    
    # Start from next available val_iter
    val_iter_start = max_val_iter + 1
    print(f'[cyan]Starting val_iter from {val_iter_start} (found {len(existing_files)} existing files)')
    
    # Temporarily override inference parameters
    original_num_steps = config.val_num_steps
    original_detail = config.val_visuals_detail
    config.val_num_steps = num_steps
    config.val_visuals_detail = 1  # Good detail for gradio
    
    pred_samples = []
    
    with torch.no_grad():
        for sample_idx in range(num_samples):
            current_val_iter = val_iter_start + sample_idx
            print(f'[cyan]Generating sample {sample_idx + 1}/{num_samples} (seed={sample_idx}, val_iter={current_val_iter})...')
            
            # Make a copy of the base data batch
            import copy
            data_batch = copy.deepcopy(base_data_batch)
            
            # Override with input video only if not using YAML/tarfile source
            if not skip_override and input_rgb is not None:
                data_batch = override_batch_with_video(data_batch, input_rgb, device)
            
            # Set FPS from UI parameter (used by a4d_visuals.py for video saving)
            data_batch['fps'] = torch.tensor([frame_rate]).float()
            
            # Run validation step with direct saving to output directory
            val_dict, loss = model.validation_step(
                data_batch=data_batch,
                iteration=current_val_iter,  # Use incremented val_iter to avoid overwriting
                dataloader_key='gradio',
                local_path=output_dir,
                directives={},
                val_iter=sample_idx,
                seed=random.randint(0, 2**31 - 1),  # Completely random seed
            )
            
            pred_samples.append(val_dict)
            print(f'[green]Sample {sample_idx + 1}/{num_samples} complete')
    
    # Restore original parameters
    config.val_num_steps = original_num_steps
    config.val_visuals_detail = original_detail
    
    return pred_samples


def pick_random_video(random_video_dir=RANDOM_VIDEO_DIR):
    '''
    Pick a random MP4 file from the specified directory.
    :param random_video_dir: Directory to search for MP4 files
    :return: Tuple of (video_path, feedback_message)
    '''
    import glob
    import random
    
    # Find all MP4 files in the directory
    mp4_pattern = os.path.join(random_video_dir, '**/*.mp4')
    mp4_files = glob.glob(mp4_pattern, recursive=True)
    
    if not mp4_files:
        return '', f'❌ No MP4 files found in {random_video_dir}'
    
    # Pick a random video
    random_video = random.choice(mp4_files)
    feedback = f'✅ Selected random video: {os.path.basename(random_video)}'
    
    return random_video, feedback


def pick_random_anyviewbench_tar(anyviewbench_yaml_dir='custom/config/cvpr/quant_v2_HR',
                                  anyviewbench_txt_dir='filenames/quant_v2'):
    '''
    Pick a random tarfile from AnyViewBench and use it for BOTH video and camera paths.
    :param anyviewbench_yaml_dir: Directory containing YAML config files
    :param anyviewbench_txt_dir: Directory containing TXT files with tarfile paths
    :return: Tuple of (tarfile_path, tarfile_path, invert_checkbox_value, feedback_message)
    '''
    import glob
    import random
    
    # Find all YAML files in the directory (to know which TXT files to check)
    yaml_pattern = os.path.join(anyviewbench_yaml_dir, '*.yaml')
    yaml_files = glob.glob(yaml_pattern)
    
    # Also check for .yml extension
    yml_pattern = os.path.join(anyviewbench_yaml_dir, '*.yml')
    yaml_files.extend(glob.glob(yml_pattern))
    
    if not yaml_files:
        return '', '', True, f'❌ No YAML files found in {anyviewbench_yaml_dir}'
    
    # Pick a random YAML file to get the dataset name
    random_yaml = random.choice(yaml_files)
    yaml_basename = os.path.splitext(os.path.basename(random_yaml))[0]
    
    # Find corresponding TXT file
    txt_file = os.path.join(anyviewbench_txt_dir, f'{yaml_basename}.txt')
    
    if not os.path.exists(txt_file):
        return '', '', True, f'❌ No corresponding TXT file found: {txt_file}'
    
    # Read non-empty lines from TXT file
    with open(txt_file, 'r') as f:
        lines = [line.strip() for line in f.readlines() if line.strip()]
    
    if not lines:
        return '', '', True, f'❌ No valid lines found in {txt_file}'
    
    # Pick a random tarfile path and use it for both video and camera
    random_tarfile = random.choice(lines)
    
    # Check the YAML file for invert_pose setting and set checkbox to opposite
    invert_checkbox_value = True  # Default
    try:
        with open(random_yaml, 'r') as f:
            yaml_content = f.read()
            if 'invert_pose: [True]' in yaml_content:
                invert_checkbox_value = False  # Set to opposite
                print(f'[cyan]Found "invert_pose: [True]" in {random_yaml}, setting invert checkbox to False')
            else:
                invert_checkbox_value = True  # Set to opposite (yaml has False or missing)
                print(f'[cyan]No "invert_pose: [True]" in {random_yaml}, setting invert checkbox to True')
    except Exception as e:
        print(f'[yellow]Warning: Could not read YAML file {random_yaml}: {e}')
        invert_checkbox_value = True  # Default on error
    
    feedback = f'✅ Selected AnyViewBench TAR: {yaml_basename} (line {lines.index(random_tarfile)+1}/{len(lines)})'
    
    return random_tarfile, random_tarfile, invert_checkbox_value, feedback


def pick_random_anyviewbench_yaml(anyviewbench_yaml_dir='custom/config/cvpr/quant_v2_HR'):
    '''
    Pick a random YAML file from AnyViewBench and use it for BOTH video and camera paths.
    :param anyviewbench_yaml_dir: Directory containing YAML config files
    :return: Tuple of (yaml_path, yaml_path, feedback_message)
    '''
    import glob
    import random
    
    # Find all YAML files in the directory
    yaml_pattern = os.path.join(anyviewbench_yaml_dir, '*.yaml')
    yaml_files = glob.glob(yaml_pattern)
    
    # Also check for .yml extension
    yml_pattern = os.path.join(anyviewbench_yaml_dir, '*.yml')
    yaml_files.extend(glob.glob(yml_pattern))
    
    if not yaml_files:
        return '', '', f'❌ No YAML files found in {anyviewbench_yaml_dir}'
    
    # Pick a random YAML file and use it for both video and camera
    random_yaml = random.choice(yaml_files)
    yaml_basename = os.path.splitext(os.path.basename(random_yaml))[0]
    
    feedback = f'✅ Selected AnyViewBench YAML: {yaml_basename}'
    
    return random_yaml, random_yaml, feedback


def preview_process_video(which_video, uploaded_video, video_path, num_frames, frame_offset, 
                             frame_stride, center_crop, resolution, frame_rate, tarfile_indices='', base_data_batch=None):
    '''
    Load and preview video from uploaded file or path (for the preview button).
    :param which_video: 'Upload', 'Path', or 'Reference' - which video source to use
    :param uploaded_video: Uploaded video file (from Gradio upload)
    :param video_path: Path to video file, YAML, or tarfile
    :param num_frames: Number of frames to extract
    :param frame_offset: Frame offset
    :param frame_stride: Frame stride
    :param center_crop: Whether to center crop
    :param resolution: Target resolution string
    :param frame_rate: Frame rate for preview video
    :param tarfile_indices: Camera/video indices for tarfile in format "in,out" e.g. "1,0"
    :param base_data_batch: Base data batch (needed for Reference mode)
    :return: (Video path for Gradio display, feedback message)
    '''
    # Determine video source based on which_video selection
    video_source = None
    
    if which_video == 'Reference':
        # Use reference dataset RGB
        # Note: For reference mode, we'll use default cam_in=1 since we don't have access to the config here
        # The config will be properly used in main_run
        cam_in = 1
        input_rgb, metadata, error_msg = gradio_utils.load_rgb_from_reference_dataset(
            base_data_batch, cam_idx=cam_in, extract_frames_fn=extract_frames_from_rgb_dict
        )
        if input_rgb is None:
            return None, f'❌ {error_msg}'
        
    elif which_video == 'Upload':
        if uploaded_video is not None:
            video_source = uploaded_video
            print(f'[cyan]Previewing uploaded video: {video_source}')
        else:
            return None, '❌ No video uploaded'
    else:  # which_video == 'Path'
        if video_path and video_path.strip():
            video_source = video_path.strip()
            print(f'[cyan]Previewing video from path: {video_source}')
        else:
            return None, '❌ No video path provided'
    
    try:
        frame_width = int(resolution.split('x')[0].strip())
        frame_height = int(resolution.split('x')[1].split('(')[0].strip())
        num_frames = int(num_frames)
        
        # VIDAR convention: cam_in=1 (input/context), cam_out=0 (output/target)
        cam_in = 1
        cam_out = 0
        
        # For Reference mode, input_rgb is already loaded, skip video loading
        if which_video != 'Reference':
            # Check if YAML dataset
            if isinstance(video_source, str) and (video_source.endswith('.yaml') or video_source.endswith('.yml')):
                input_rgb, sample, metadata, error_msg = gradio_utils.load_rgb_from_yaml_dataset(
                    video_source, cam_idx=cam_in, extract_frames_fn=extract_frames_from_rgb_dict
                )
                if input_rgb is None:
                    return None, f'❌ {error_msg}'
            
            # Check if tarfile
            elif isinstance(video_source, str) and video_source.endswith('.tar'):
                print(f'[cyan]Loading preview from tarfile: {video_source}')
                # Parse custom camera indices if provided
                cam_in, cam_out, cam_indices = parse_tarfile_camera_indices(tarfile_indices)
                
                print(f'[cyan]Using camera indices: {cam_indices}')
                tarfile_data = infer_utils.load_data_from_tarfile(
                    video_source, num_timesteps=num_frames, cam_indices=cam_indices,
                    remap_camera_indices=True
                )
                
                if tarfile_data is None:
                    return None, '❌ Failed to load tarfile'
                
                # Extract RGB frames for context camera (input/conditioning view)
                rgb_dict = tarfile_data['rgb']
                input_rgb, time_indices, error_msg = extract_frames_from_rgb_dict(
                    rgb_dict, cam_idx=1, max_frames=num_frames)
                
                if input_rgb is None:
                    return None, f'❌ {error_msg} in tarfile'
                
                print(f'[green]Loaded {input_rgb.shape[0]} frames from tarfile (camera {cam_in}, frames {time_indices[0]}-{time_indices[-1]})')
            
            # Load video from S3 or YouTube
            elif isinstance(video_source, str) and (video_source.startswith('s3://') or 'youtube.com' in video_source or 'youtu.be' in video_source):
                input_rgb = infer_utils.load_video_from_source(
                    video_source, target_T=num_frames, target_H=frame_height, target_W=frame_width
                )
            else:
                # Works for both uploaded files and local paths
                input_rgb = infer_utils.create_input_from_video(
                    video_source, num_frames=num_frames, 
                    frame_width=frame_width, frame_height=frame_height,
                    center_crop=center_crop, frame_stride=frame_stride, frame_offset=frame_offset
                )
        
        if input_rgb is None:
            print('[red]Failed to load video for preview')
            return None, '❌ Failed to load video - check path and format'
        
        # Save temporary preview video
        import tempfile
        preview_path = tempfile.mktemp(suffix='.mp4', dir='/tmp/gradio_anyview')
        os.makedirs('/tmp/gradio_anyview', exist_ok=True)
        
        # Convert tensor to numpy array (T, H, W, C)
        preview_frames = rearrange(input_rgb, 'T C H W -> T H W C').numpy()
        preview_frames = np.clip(preview_frames, 0.0, 1.0)
        
        # Save preview video using frame_rate from slider
        infer_utils.save_video(preview_path, preview_frames, fps=int(frame_rate), quality=8)
        print(f'[green]Preview saved: {preview_path}')
        
        feedback = f'✅ Preview loaded successfully ({input_rgb.shape[0]} frames, {frame_width}x{frame_height})'
        return preview_path, feedback
        
    except Exception as e:
        print(f'[red]Error previewing video: {e}')
        import traceback
        print(traceback.format_exc())
        error_msg = f'❌ Error: {str(e)}'
        return None, error_msg


def extract_videos_from_samples(pred_samples, entry_name='rgb0_pred'):
    '''
    Extract video arrays from prediction samples.
    :param pred_samples (list): List of prediction dictionaries
    :param entry_name (str): Name of entry to extract
    :return (list): List of video arrays (T, H, W, C) in [0, 1]
    '''
    videos = []
    
    for sample in pred_samples:
        if 'y0_pred_entries' in sample and entry_name in sample['y0_pred_entries']:
            # (B, C, T, H, W) -> (T, H, W, C)
            video = sample['y0_pred_entries'][entry_name]
            video = rearrange(video[0], 'C T H W -> T H W C')
            video = video.detach().cpu().numpy()
            video = np.clip(video, 0.0, 1.0)
            videos.append(video)
        elif 'pred_rgb' in sample:
            # Alternative key
            video = sample['pred_rgb']
            if len(video.shape) == 5:  # (B, C, T, H, W)
                video = rearrange(video[0], 'C T H W -> T H W C')
            elif len(video.shape) == 4:  # (T, C, H, W)
                video = rearrange(video, 'T C H W -> T H W C')
            video = video.detach().cpu().numpy()
            video = np.clip(video, 0.0, 1.0)
            videos.append(video)
    
    return videos


def load_camera_data_from_path(camera_path, num_frames, tarfile_cam_indices=''):
    '''
    Helper function to load camera data from tar or yaml file.
    :param camera_path (str): Path to camera file (tar/yaml)
    :param num_frames (int): Number of frames for tarfile loading
    :param tarfile_cam_indices (str): Camera indices for tarfile in format "in,out" e.g. "1,0"
    :return: Tuple of (camera_data dict or None, is_tarfile bool, is_yaml_batch bool, error_msg or None)
    '''
    if not camera_path or not camera_path.strip():
        return None, False, False, None
    
    camera_source = camera_path.strip()
    
    try:
        if camera_source.endswith('.yaml') or camera_source.endswith('.yml'):
            # Load from YAML - return full batch to use as base_data_batch
            print(f'[cyan]Loading camera from YAML: {camera_source}')
            from torch.utils.data import DataLoader
            
            camera_dataset = VidarDataset(dataset_dir=camera_source, phase='test')
            camera_dataloader = DataLoader(camera_dataset, batch_size=1, shuffle=True, num_workers=0, pin_memory=True)
            camera_sample = next(iter(camera_dataloader))
            
            print(f'[green]Loaded full batch from camera YAML (will be used as base_data_batch)')
            return camera_sample, False, True, None  # is_tarfile=False, is_yaml_batch=True
            
        elif camera_source.endswith('.tar'):
            # Load from tarfile - allow custom camera indices
            cam_in_tar, cam_out_tar, cam_indices = parse_tarfile_camera_indices(tarfile_cam_indices)
            
            print(f'[cyan]Loading camera from tarfile: {camera_source}')
            print(f'[cyan]Using camera indices: {cam_indices} (vidar order: cam 0=output, cam 1=input)')
            camera_data = infer_utils.load_data_from_tarfile(
                camera_source, num_timesteps=int(num_frames), cam_indices=cam_indices,
                remap_camera_indices=True
            )
            
            if camera_data:
                return camera_data, True, False, None  # is_tarfile=True, is_yaml_batch=False
            else:
                return None, True, False, 'Failed to load tarfile'
        else:
            return None, False, False, f'Unknown file type: {camera_source}'
            
    except Exception as e:
        import traceback
        error_msg = f'Error loading camera: {e}\n{traceback.format_exc()}'
        print(f'[red]{error_msg}')
        return None, False, False, error_msg


def load_video_data(which_video, video_source, raw_video, 
                   num_frames, frame_width, frame_height, 
                   center_crop, frame_stride, frame_offset, override_resolution=False,
                   tarfile_cam_indices=''):
    '''
    Load RGB video data from various sources.
    Returns RGB data and optional base_data_batch (for YAML datasets).
    
    :param which_video (str): 'Upload', 'Path', or 'Reference'
    :param video_source (str): Path/URL to video
    :param raw_video: Uploaded video file
    :param num_frames (int): Number of frames to extract
    :param frame_width (int): Target width
    :param frame_height (int): Target height
    :param center_crop (bool): Whether to center crop
    :param frame_stride (int): Stride for frame extraction
    :param frame_offset (int): Offset for frame extraction
    :param override_resolution (bool): If True, resize tarfile RGB to target resolution
    :param tarfile_cam_indices (str): Camera indices for tarfile in format "in,out" e.g. "1,0"
    :return: Tuple of (input_rgb, base_data_batch, metadata, skip_override, error_msg)
    '''
    import torch
    
    # VIDAR convention (defaults)
    cam_in = 1   # Input/context camera
    cam_out = 0  # Output/target camera
    
    if which_video == 'Reference':
        # Use reference dataset RGB without overriding
        return None, None, None, True, None
        
    elif which_video == 'Path' and video_source:
        # Check if it's a YAML dataset config file
        if video_source.endswith('.yaml') or video_source.endswith('.yml'):
            print(f'[cyan]Loading RGB from YAML dataset: {video_source}')
            input_rgb, yaml_batch, metadata, error_msg = gradio_utils.load_rgb_from_yaml_dataset(
                video_source, cam_idx=cam_in, extract_frames_fn=None
            )
            if input_rgb is None:
                return None, None, None, False, error_msg
            
            print(f'[green]YAML dataset loaded: camera {cam_out} (output/GT) has {metadata["cam_out_count"]} frames, '
                  f'camera {cam_in} (input/context) has {metadata["cam_in_count"]} frames')
            return input_rgb, yaml_batch, metadata, True, None
            
        # Check if it's a tarfile
        elif video_source.endswith('.tar'):
            print(f'[cyan]Loading RGB from tarfile: {video_source}')
            # Parse custom camera indices if provided
            cam_in, cam_out, cam_indices = parse_tarfile_camera_indices(tarfile_cam_indices)
            
            tarfile_data = infer_utils.load_data_from_tarfile(
                video_source, num_timesteps=num_frames, cam_indices=cam_indices,
                remap_camera_indices=True
            )
            
            if tarfile_data is None:
                return None, None, None, False, 'Failed to load tarfile. Check file format and path.'
            
            # Apply resolution override on raw tarfile data if requested
            if override_resolution:
                print(f'[cyan]Applying resolution override on raw tarfile data: {frame_width}x{frame_height}')
                tarfile_data = gradio_utils.resize_tarfile_rgb(tarfile_data, frame_width, frame_height)
            
            # Extract RGB data only
            rgb_dict = tarfile_data['rgb']
            cam1_keys = sorted([k for k in rgb_dict.keys() if k[1] == 1])
            if len(cam1_keys) > 0:
                # Stack camera 1 frames: (T, C, H, W)
                input_rgb = torch.stack([rgb_dict[k][0] for k in cam1_keys], dim=0)
                print(f'[green]Extracted input_rgb from tarfile camera 1: {input_rgb.shape}')
            else:
                print('[yellow]Warning: No camera 1 frames found in tarfile')
                input_rgb = None
            
            # Verify both cameras are present
            cam0_count = len([k for k in rgb_dict.keys() if k[1] == 0])
            cam1_count = len([k for k in rgb_dict.keys() if k[1] == 1])
            print(f'[green]Tarfile loaded: camera 0 (rgb0/GT) has {cam0_count} frames, camera 1 (rgb1/input) has {cam1_count} frames')
            
            # Return tarfile_data as metadata to be used for RGB update later
            metadata = {'tarfile_data': tarfile_data}
            return input_rgb, None, metadata, True, None
            
        elif video_source.startswith('s3://') or 'youtube.com' in video_source or 'youtu.be' in video_source:
            # Load directly using the new load_video_from_source function
            print(f'[cyan]Downloading/loading video from: {video_source}')
            input_rgb = infer_utils.load_video_from_source(
                video_source, target_T=num_frames, target_H=frame_height, target_W=frame_width
            )
            if input_rgb is None:
                return None, None, None, False, 'Failed to load video from source. Check URL/path and network connection.'
            return input_rgb, None, None, False, None
        else:
            # Use standard video processing for local files
            input_rgb = infer_utils.create_input_from_video(
                video_source, num_frames=num_frames, 
                frame_width=frame_width, frame_height=frame_height,
                center_crop=center_crop, frame_stride=frame_stride, frame_offset=frame_offset
            )
            if input_rgb is None:
                return None, None, None, False, 'Failed to load video. Check file format and path.'
            return input_rgb, None, None, False, None
    
    elif which_video == 'Upload' and raw_video:
        # Use uploaded video file
        print(f'[cyan]Using uploaded video: {raw_video}')
        input_rgb = infer_utils.create_input_from_video(
            raw_video, num_frames=num_frames, 
            frame_width=frame_width, frame_height=frame_height,
            center_crop=center_crop, frame_stride=frame_stride, frame_offset=frame_offset
        )
        if input_rgb is None:
            return None, None, None, False, 'Failed to load video. Check file format and path.'
        return input_rgb, None, None, False, None
    else:
        return None, None, None, False, 'Please provide a video.'


def load_camera_data_full(which_camera, camera_path_or_url, video_source,
                          num_frames, tarfile_cam_indices, 
                          canonical_params, invert_camera_pose):
    '''
    Load camera data from various sources in a unified way.
    
    NOTE: When loading from YAML (is_yaml_batch=True), camera_data_or_batch contains
    a full data batch (base_data_batch), not just camera poses/intrinsics.
    For tarfiles (is_tarfile=True), it contains just {'pose': ..., 'intrinsics': ...}.
    
    :param which_camera (str): 'Canonical', 'Path', or 'Reference'
    :param camera_path_or_url (str): Path to camera file
    :param video_source (str): Video path (used as fallback if camera_path empty)
    :param num_frames (int): Number of frames
    :param tarfile_cam_indices (str): Camera indices for tarfile
    :param canonical_params (dict): Parameters for canonical trajectory
    :param invert_camera_pose (bool): Whether to invert tarfile poses
    :return: Tuple of (camera_data_or_batch, is_tarfile, is_yaml_batch, source_description, error_msg)
             where camera_data_or_batch is:
             - None for 'Canonical' or 'Reference' modes
             - Full data batch if is_yaml_batch=True
             - Dict with 'pose'/'intrinsics' if is_tarfile=True
    '''
    if which_camera == 'Canonical':
        print('[cyan]Using canonical camera parameters')
        return None, False, False, 'canonical', None
        
    elif which_camera == 'Reference':
        print('[cyan]Using reference camera parameters (no modifications)')
        return None, False, False, 'reference', None
        
    elif which_camera == 'Path':
        # Determine camera source
        camera_source = None
        if camera_path_or_url and camera_path_or_url.strip():
            camera_source = camera_path_or_url.strip()
        elif video_source and (video_source.endswith('.tar') or video_source.endswith('.yaml') or video_source.endswith('.yml')):
            # If no camera path provided but video is a tarfile or yaml, use video path for camera
            camera_source = video_source
            print(f'[cyan]Camera path empty, using video path for camera: {camera_source}')
        
        if camera_source:
            print(f'[cyan]Loading camera from: {camera_source}')
            camera_data_or_batch, is_tarfile, is_yaml_batch, error_msg = load_camera_data_from_path(
                camera_source, num_frames, tarfile_cam_indices)
            if error_msg:
                print(f'[yellow]Warning: {error_msg}, falling back to reference')
                return None, False, False, 'reference', error_msg
            
            if is_yaml_batch:
                print(f'[cyan]Loaded full data batch from YAML')
            elif is_tarfile:
                print(f'[cyan]Loaded pose/intrinsics from tarfile')
            
            source_type = 'yaml' if is_yaml_batch else ('tarfile' if is_tarfile else 'unknown')
            return camera_data_or_batch, is_tarfile, is_yaml_batch, source_type, None
        else:
            print('[cyan]No camera path provided and video is not tar/yaml, using reference camera')
            return None, False, False, 'reference', None
    
    return None, False, False, 'unknown', 'Invalid camera mode'


def apply_camera_to_batch(base_data_batch, which_camera, camera_data, is_tarfile,
                          invert_camera_pose, canonical_params):
    '''
    Apply camera parameters to data batch based on selected mode.
    :param base_data_batch (dict): Base data batch
    :param which_camera (str): 'Canonical', 'Path', or 'Reference'
    :param camera_data (dict): Loaded camera data (or None)
    :param is_tarfile (bool): Whether camera data is from tarfile
    :param invert_camera_pose (bool): Whether to invert poses (tarfile only)
    :param canonical_params (dict): Dict with canonical trajectory parameters
    :return: Modified data batch
    '''
    if base_data_batch is None:
        raise ValueError("base_data_batch is None")
    
    data_batch = copy.deepcopy(base_data_batch)
    
    if which_camera == 'Canonical':
        # Apply canonical trajectory
        data_batch = gradio_utils.apply_canonical_trajectory(
            data_batch,
            azimuth_deg=canonical_params.get('azimuth_deg', 0),
            input_elevation_deg=canonical_params.get('input_elevation_deg', 0),
            output_elevation_deg=canonical_params.get('output_elevation_deg', 0),
            radius=canonical_params['radius'],
            trajectory_type=canonical_params['trajectory_type'],
            which_setting=canonical_params.get('which_setting', 'Exo2Exo')
        )
        
    elif which_camera == 'Path' and camera_data:
        if is_tarfile:
            # Tarfile poses are ABSOLUTE CAM2WORLD by default
            # If invert is checked, they're actually ABSOLUTE WORLD2CAM, so convert to CAM2WORLD
            tarfile_poses = camera_data['pose']
            print(f'[cyan]Loaded {len(tarfile_poses)} poses from tarfile with keys: {sorted(tarfile_poses.keys())[:5]}...')
            
            if invert_camera_pose:
                # Convert ABSOLUTE WORLD2CAM -> ABSOLUTE CAM2WORLD
                temp_data = {'pose': tarfile_poses}
                temp_data = gradio_utils.invert_tarfile_poses(temp_data)
                tarfile_poses = temp_data['pose']
            
            # Now convert ABSOLUTE CAM2WORLD -> RELATIVE WORLD2CAM (vidar format)
            vidar_poses = gradio_utils.convert_absolute_to_vidar_format(tarfile_poses)
            print(f'[cyan]Converted to vidar format with {len(vidar_poses)} poses')
            data_batch['anydata']['pose'] = vidar_poses
        else:
            # YAML already provides poses in vidar RELATIVE WORLD2CAM format
            data_batch['anydata']['pose'] = camera_data['pose']
        
        # Override intrinsics if provided
        if 'intrinsics' in camera_data and len(camera_data['intrinsics']) > 0:
            data_batch['anydata']['intrinsics'] = camera_data['intrinsics']
    
    # For 'Reference' or 'Path' with no data, just return the copied batch as-is
    
    return data_batch


def visualize_cameras_only(base_data_batch, which_camera, invert_camera_pose, 
                          which_setting,
                          canonical_azimuth_deg, canonical_input_elev_deg, canonical_output_elev_deg,
                          canonical_radius, canonical_trajectory,
                          camera_path_or_url, tarfile_cam_indices, num_frames):
    '''
    Visualize camera trajectory without running inference.
    '''
    if base_data_batch is None:
        return '❌ **Error**: Base data batch not loaded', None
    
    # Map UI trajectory names to internal names
    trajectory_map = {
        'None (Direct)': 'Jumpy',
        'Gradual (Linear)': 'Gradual',
        'Gradual (Sqrt)': 'GradualSqrt'
    }
    trajectory_type = trajectory_map.get(canonical_trajectory, 'Jumpy')
    
    try:
        # Prepare canonical params
        canonical_params = {
            'azimuth_deg': canonical_azimuth_deg,
            'input_elevation_deg': canonical_input_elev_deg,
            'output_elevation_deg': canonical_output_elev_deg,
            'radius': canonical_radius,
            'trajectory_type': trajectory_type,
            'which_setting': which_setting
        }
        
        # Load camera data using unified function
        camera_data_or_batch, is_tarfile_cam, is_yaml_batch, source_desc, error_msg = load_camera_data_full(
            which_camera, camera_path_or_url, None,  # video_source=None for visualization
            num_frames, tarfile_cam_indices, 
            canonical_params, invert_camera_pose
        )
        
        if error_msg:
            print(f'[yellow]Warning: {error_msg}')
        
        # If camera is a YAML batch, use it as base_data_batch
        # Otherwise, camera_data_or_batch contains just pose/intrinsics for apply_camera_to_batch
        if is_yaml_batch and camera_data_or_batch:
            base_data_batch = camera_data_or_batch
            print(f'[green]Using camera YAML as base_data_batch')
            camera_poses_only = None  # No separate camera data to apply
        else:
            camera_poses_only = camera_data_or_batch  # Tarfile pose/intrinsics or None
        
        # Apply camera to batch
        data_batch = apply_camera_to_batch(
            base_data_batch, which_camera, camera_poses_only, is_tarfile_cam,
            invert_camera_pose, canonical_params
        )
        
        # Generate trajectory info
        if which_camera == 'Canonical':
            trajectory_info = (f'{canonical_trajectory}, setting={which_setting}, '
                             f'azimuth={canonical_azimuth_deg}°, in_elev={canonical_input_elev_deg}°, '
                             f'out_elev={canonical_output_elev_deg}°, radius={canonical_radius}m')
        elif which_camera == 'Reference':
            trajectory_info = 'Reference dataset (no modifications)'
        else:  # 'Path'
            if source_desc == 'yaml':
                trajectory_info = 'From YAML'
            elif source_desc == 'tarfile':
                trajectory_info = f'From tarfile'
                if invert_camera_pose:
                    trajectory_info += ' (inverted)'
                if tarfile_cam_indices:
                    trajectory_info += f' (indices: {tarfile_cam_indices})'
            else:
                trajectory_info = 'No camera path provided, using reference'
        
        camera_plot, extrinsics_text = gradio_utils.extract_and_visualize_cameras(
            data_batch,
        )
        
        description = f'''📊 **Camera Trajectory Visualization**

**Trajectory**: {trajectory_info}

**Camera extrinsics:**

{extrinsics_text}

*Click "Run AnyView Generation!" to generate the video with these camera trajectories.*
'''
        return description, camera_plot
        
    except Exception as e:
        import traceback
        tb_str = traceback.format_exc()
        print(f'[red]Error visualizing cameras: {e}')
        print(f'[red]{tb_str}')
        
        error_msg = f'''❌ **Error visualizing cameras**

**Error message**: {str(e)}

**Stacktrace**:
```
{tb_str}
```
'''
        return error_msg, None


def main_run(model_bundle, base_data_batch, output_path, action,
            which_video='Upload', raw_video=None, video_path_or_url='',
            which_camera='Canonical', invert_camera_pose=False,
            which_setting='Exo2Exo',
            canonical_azimuth_deg=0, canonical_input_elev_deg=0, canonical_output_elev_deg=0,
            canonical_radius=10.0, canonical_trajectory='None (Direct)',
            camera_path_or_url='', tarfile_cam_indices='',
            num_frames=41, frame_offset=0, frame_stride=1, frame_rate=10,
            center_crop=True, resolution='512 x 288', override_resolution=False,
            num_samples=1, num_steps=25):
    '''
    Main function to run inference on uploaded video or video from path/URL.
    Uses VIDAR convention: camera 0 = output/target, camera 1 = input/context
    '''
    # Map UI trajectory names to internal names
    trajectory_map = {
        'None (Direct)': 'Jumpy',
        'Gradual (Linear)': 'Gradual',
        'Gradual (Sqrt)': 'GradualSqrt'
    }
    trajectory_type = trajectory_map.get(canonical_trajectory, 'Jumpy')
    
    import torch
    
    # VIDAR convention (hardcoded)
    cam_in = 1   # Input/context camera
    cam_out = 0  # Output/target camera
    
    print('[cyan]═══════════════════════════════════════')
    print(f'[cyan]Action: {action}')
    print(f'[cyan]Video mode: {which_video}')
    print(f'[cyan]Camera mode: {which_camera}')
    print(f'[cyan]Setting: {which_setting}')
    if which_camera == 'Canonical':
        print(f'[cyan]Canonical: {canonical_trajectory}, azimuth={canonical_azimuth_deg}°, '
              f'in_elev={canonical_input_elev_deg}°, out_elev={canonical_output_elev_deg}°, radius={canonical_radius}m')
    elif which_camera == 'Path':
        print(f'[cyan]Invert camera pose: {invert_camera_pose}')
        if tarfile_cam_indices:
            print(f'[cyan]Tarfile camera indices: {tarfile_cam_indices}')
    
    # Parse resolution parameters
    frame_width = int(resolution.split('x')[0].strip())
    frame_height = int(resolution.split('x')[1].split('(')[0].strip())
    num_frames = int(num_frames)  # Convert from slider value
    
    # Determine video source
    video_source = video_path_or_url.strip() if (which_video == 'Path' and video_path_or_url) else raw_video
    if which_video == 'Upload' and not raw_video:
        error_msg = '❌ **Error**: Please upload a video file.'
        print(f'[red]{error_msg}')
        return (error_msg, None, None, None, None, None, None, None)
    elif which_video == 'Path' and not video_source:
        error_msg = '❌ **Error**: Please provide a video path/URL.'
        print(f'[red]{error_msg}')
        return (error_msg, None, None, None, None, None, None, None)
    
    # ═══════════════════════════════════════════════════════════
    # PHASE 1: Load Camera Data (separately from video/RGB)
    # ═══════════════════════════════════════════════════════════
    canonical_params = {
        'azimuth_deg': canonical_azimuth_deg,
        'input_elevation_deg': canonical_input_elev_deg,
        'output_elevation_deg': canonical_output_elev_deg,
        'radius': canonical_radius,
        'trajectory_type': trajectory_type,
        'which_setting': which_setting
    }
    
    camera_data_or_batch, is_tarfile_cam, is_yaml_batch_cam, camera_source_desc, cam_error = load_camera_data_full(
        which_camera, camera_path_or_url, video_source,
        num_frames, tarfile_cam_indices, 
        canonical_params, invert_camera_pose
    )
    
    if cam_error:
        print(f'[yellow]Warning: {cam_error}')
    
    # If camera is a YAML batch, use it as base_data_batch
    # Otherwise, camera_data_or_batch contains just pose/intrinsics for apply_camera_to_batch
    if is_yaml_batch_cam and camera_data_or_batch:
        base_data_batch = camera_data_or_batch
        print(f'[green]Using camera YAML as base_data_batch')
        camera_poses_only = None  # No separate camera data to apply (already in base_data_batch)
    else:
        camera_poses_only = camera_data_or_batch  # Tarfile pose/intrinsics or None
    
    # ═══════════════════════════════════════════════════════════
    # PHASE 2: Load Video/RGB Data (separately from cameras)
    # ═══════════════════════════════════════════════════════════
    print('[yellow]Processing video input...')
    
    try:
        if which_video == 'Reference':
            # Use reference dataset RGB without overriding
            input_rgb, metadata, error_msg = gradio_utils.load_rgb_from_reference_dataset(
                base_data_batch, cam_idx=cam_in, extract_frames_fn=None
            )
            if input_rgb is None:
                error_msg = f'❌ **Error**: {error_msg}'
                print(f'[red]{error_msg}')
                return (error_msg, None, None, None, None, None, None, None)
            
            skip_override = True
            print(f'[green]Reference dataset: camera {cam_out} (output/GT) has {metadata["cam_out_count"]} frames, '
                  f'camera {cam_in} (input/context) has {metadata["cam_in_count"]} frames')
        else:
            # Load video using unified helper function
            input_rgb, video_batch, video_metadata, skip_override, error_msg = load_video_data(
                which_video, video_source, raw_video,
                num_frames, frame_width, frame_height,
                center_crop, frame_stride, frame_offset, override_resolution,
                tarfile_cam_indices
            )
            
            if error_msg:
                error_msg = f'❌ **Error**: {error_msg}'
                print(f'[red]{error_msg}')
                return (error_msg, None, None, None, None, None, None, None)
            
            # If video loaded a YAML batch and we don't have camera YAML, use video's batch
            if video_batch and not is_yaml_batch_cam:
                base_data_batch = video_batch
                print(f'[green]Using video YAML batch as base_data_batch')
            
            # If video is tarfile, update base_data_batch with RGB
            if video_metadata and 'tarfile_data' in video_metadata:
                import copy
                tarfile_data = video_metadata['tarfile_data']
                base_data_batch = copy.deepcopy(base_data_batch)
                base_data_batch['anydata']['rgb'] = tarfile_data['rgb']
                print(f'[green]Updated base_data_batch with tarfile RGB data')
        
        if input_rgb is not None:
            print(f'[green]Processed video shape: {input_rgb.shape}')
        
    except Exception as e:
        import traceback
        tb_str = traceback.format_exc()
        print(f'[red]Error processing video: {e}')
        print(f'[red]{tb_str}')
        
        error_msg = f'''❌ **Error processing video**

**Error message**: {str(e)}

**Stacktrace**:
```
{tb_str}
```
'''
        return (error_msg, None, None, None, None, None, None, None)
    
    # Note: Resolution override for tarfiles is now handled in load_video_data()
    # on the raw tarfile data before it becomes a vidar batch.
    # For YAML/reference sources, resolution override is not supported
    # (as requested by user - it's harder to change data once vidar batch is constructed).
    
    # ═══════════════════════════════════════════════════════════
    # PHASE 3: Apply Camera Parameters to Batch
    # ═══════════════════════════════════════════════════════════
    # Note: canonical_params already defined in PHASE 1
    base_data_batch = apply_camera_to_batch(
        base_data_batch, which_camera, camera_poses_only, is_tarfile_cam,
        invert_camera_pose, canonical_params
    )
    
    # NOTE(bvh): image_size was a legacy Cosmos key, now derived from anydata rgb tensors.
    # Resolution override for inference should resize the rgb tensors directly instead.
    
    # Log what was applied
    if which_camera == 'Canonical':
        print(f'[green]Applied canonical trajectory: {canonical_trajectory}, setting={which_setting}, '
              f'azimuth={canonical_azimuth_deg}°, in_elev={canonical_input_elev_deg}°, '
              f'out_elev={canonical_output_elev_deg}°, radius={canonical_radius}m')
    elif which_camera == 'Path':
        if camera_source_desc == 'yaml':
            print(f'[green]Applied camera from YAML (full data batch)')
            if invert_camera_pose:
                print('[yellow]Inversion checkbox ignored for YAML camera (invert only applies to tarfile)')
        elif camera_source_desc == 'tarfile':
            print(f'[green]Applied camera from tarfile (pose/intrinsics only)')
            if invert_camera_pose:
                print('[green]Applied camera pose inversion')
        elif camera_source_desc == 'reference':
            print(f'[green]Using reference camera (no modifications)')
    
    # Run inference
    try:
        pred_samples = inference(
            model_bundle, base_data_batch, input_rgb, output_path,
            num_samples=num_samples, num_frames=num_frames,
            num_steps=num_steps, frame_rate=frame_rate,
            skip_override=skip_override
        )
        
        print(f'[green]Generated {len(pred_samples)} samples')
        
        # Calculate uncertainty if multiple samples (only for rgb0)
        uncertainty_videos = {}
        diversity_metrics = {}
        if len(pred_samples) > 1:
            print('[cyan]Calculating uncertainty heatmaps...')
            entry_names = ['rgb0']  # Only calculate diversity for rgb0 (output)
            (model, config, exp_cfg, device) = model_bundle
            uncertainty_videos, diversity_metrics = infer_utils.calculate_uncertainty_heatmaps(
                pred_samples, entry_names, a4d_vae=model.vae)
            
            # Print diversity metrics
            for entry_name, metrics in diversity_metrics.items():
                print(f'[green]{entry_name}: mean diversity = {metrics["mean_diversity"]:.4f}')
        
    except Exception as e:
        import traceback
        tb_str = traceback.format_exc()
        print(f'[red]Error during inference: {e}')
        print(f'[red]{tb_str}')
        
        error_msg = f'''❌ **Error during model inference**

**Error message**: {str(e)}

**Stacktrace**:
```
{tb_str}
```

'''
        return (error_msg, None, None, None, None, None, None, None)
    
    # Extract and save videos
    try:
        # Convert input to displayable format
        input_rgb_np = rearrange(input_rgb, 'T C H W -> T H W C').numpy()
        input_rgb_np = np.clip(input_rgb_np, 0.0, 1.0)
        
        # Get file prefix
        fn_prefix = time.strftime('%Y%m%d-%H%M%S')
        fn_idx = 1
        
        # Create camera angle suffix for canonical trajectory
        if which_camera == 'Canonical':
            angle_suffix = f'_az{canonical_azimuth_deg:+.0f}'
        else:
            angle_suffix = ''
        
        # Find unique filename
        while True:
            in_vid_fp = os.path.join(output_path, f'{fn_prefix}{fn_idx:02d}{angle_suffix}-input.mp4')
            if not os.path.exists(in_vid_fp):
                break
            fn_idx += 1
        
        # Use frame_rate directly from slider
        vis_fps = int(frame_rate)
        
        # Save input video
        infer_utils.save_video(in_vid_fp, input_rgb_np, fps=vis_fps, quality=9)
        print(f'[green]Saved input video: {in_vid_fp}')
        
        # Initialize uncertainty videos (will be computed if multiple samples)
        uncertainty_videos = {}
        
        # Save multi-sample gallery if multiple samples generated
        if len(pred_samples) > 1:
            (model, config, exp_cfg, device) = model_bundle
            
            # Calculate uncertainty videos from multiple samples
            uncertainty_videos, diversity_metrics = infer_utils.calculate_uncertainty_heatmaps(
                pred_samples, entry_names=['rgb0'], a4d_vae=model.vae)
            
            gallery_video = infer_utils.create_multi_sample_gallery(
                pred_samples, model.vae, entry_names=['rgb0', 'rgb1'], 
                uncertainty_videos=uncertainty_videos)
            
            if gallery_video is not None:
                gallery_fp = os.path.join(output_path, f'{fn_prefix}{fn_idx:02d}{angle_suffix}-gallery.mp4')
                infer_utils.save_video(gallery_fp, gallery_video, fps=vis_fps, quality=9)
                print(f'[green]Saved multi-sample gallery: {gallery_fp}')
                if diversity_metrics:
                    for entry_name, metrics in diversity_metrics.items():
                        print(f'[cyan]  {entry_name} mean diversity: {metrics["mean_diversity"]:.4f}')
        
        # Extract and save output videos, and find gallery visualizations
        output_videos = []
        gallery_videos = []
        output_resolution = None  # To capture output (rgb0) resolution
        
        # Get VAE from model for decoding latents
        (model, config, exp_cfg, device) = model_bundle
        a4d_vae = model.vae
        
        from custom.eval.visuals import cached_decode_latent_video_auto
        
        for (s, pred_sample) in enumerate(pred_samples):
            # Try to extract video from various possible keys
            video_extracted = False
            
            if 'y0_pred_entries' in pred_sample:
                # Try to find RGB keys (decoded cache first, then latent)
                for key in ['rgb0_cdec', 'rgb1_cdec', 'rgb_cdec', 'rgb0', 'rgb1', 'rgb', 'rgb0_pred', 'rgb1_pred', 'rgb_pred']:
                    if key in pred_sample['y0_pred_entries']:
                        print(f'[cyan]Found key: {key}')
                        pred_tensor = pred_sample['y0_pred_entries'][key]
                        print(f'[cyan]DEBUG: pred_tensor shape: {pred_tensor.shape}')
                        
                        # Check if it's already decoded (3 channels) or needs decoding (16 channels)
                        num_channels = pred_tensor.shape[1] if len(pred_tensor.shape) == 5 else pred_tensor.shape[0]
                        print(f'[cyan]DEBUG: num_channels: {num_channels}')
                        
                        if num_channels > 4:
                            # VAE latent - needs decoding
                            print(f'[cyan]Decoding VAE latents with {num_channels} channels...')
                            # Remove _cdec suffix if present for decoding
                            decode_key = key.replace('_cdec', '')
                            decoded = cached_decode_latent_video_auto(
                                pred_sample['y0_pred_entries'], decode_key, a4d_vae)
                            # decoded is (B, 3, T, H, W) in [0, 1]
                            print(f'[cyan]DEBUG: decoded shape: {decoded.shape}')
                            output_rgb = rearrange(decoded[0].float(), 'C T H W -> T H W C').cpu().numpy()
                        else:
                            # Already RGB - just rearrange
                            print(f'[cyan]Using pre-decoded RGB with {num_channels} channels')
                            if len(pred_tensor.shape) == 5:
                                output_rgb = rearrange(pred_tensor[0], 'C T H W -> T H W C')
                            else:
                                output_rgb = rearrange(pred_tensor, 'C T H W -> T H W C')
                            output_rgb = output_rgb.detach().cpu().numpy()
                        
                        output_rgb = np.clip(output_rgb, 0.0, 1.0)
                        video_extracted = True
                        print(f'[green]Extracted video from key: {key} with shape {output_rgb.shape}')
                        
                        # Capture output resolution from first sample
                        if s == 0 and output_resolution is None:
                            # output_rgb shape is (T, H, W, C)
                            output_resolution = (output_rgb.shape[2], output_rgb.shape[1])  # (W, H)
                        
                        break
            
            if not video_extracted:
                print(f'[yellow]Warning: Could not extract video from sample {s}')
                print(f'Available keys: {list(pred_sample.keys())}')
                if 'y0_pred_entries' in pred_sample:
                    print(f'y0_pred_entries keys: {list(pred_sample["y0_pred_entries"].keys())}')
                continue
            
            # Save output video only
            out_vid_fp = os.path.join(output_path, f'{fn_prefix}{fn_idx:02d}{angle_suffix}-output-{s+1}.mp4')
            infer_utils.save_video(out_vid_fp, output_rgb, fps=vis_fps, quality=9)
            
            output_videos.append(gr.Video(value=out_vid_fp, format='mp4', 
                                         label=f'Output (Sample {s+1})', visible=True))
            
            print(f'[green]Saved output {s+1}: {out_vid_fp}')
        
        # Pad output list to 4 samples
        while len(output_videos) < 4:
            output_videos.append(gr.Video(visible=False))
        
        # Create custom Input | Pred concatenation videos for each sample (no GT)
        # If multiple samples, also include uncertainty heatmap
        cat_videos = []
        for (s, pred_sample) in enumerate(pred_samples):
            concat_video = infer_utils.create_input_pred_concatenation(pred_sample, a4d_vae)
            
            if concat_video is not None:
                # Add uncertainty heatmap if available (only for first sample to show once)
                if s == 0 and len(uncertainty_videos) > 0 and 'rgb0' in uncertainty_videos:
                    uncertainty_video = uncertainty_videos['rgb0']  # (T, H, W, 3)
                    # Concatenate: Input | Pred | Uncertainty (with padding for height mismatches)
                    concat_video = infer_utils.concatenate_videos_horizontal_with_padding(
                        [concat_video, uncertainty_video], verbose=False)
                    print(f'[green]Added uncertainty heatmap to concatenation')
                
                # Save concatenation video
                concat_fp = os.path.join(output_path, f'{fn_prefix}{fn_idx:02d}{angle_suffix}-concat-{s+1}.mp4')
                infer_utils.save_video(concat_fp, concat_video, fps=vis_fps, quality=9)
                cat_videos.append(concat_fp)
                
                if s == 0 and len(uncertainty_videos) > 0:
                    print(f'[green]Saved Input|Pred|Uncertainty concatenation {s+1}: {concat_fp}')
                else:
                    print(f'[green]Saved Input|Pred concatenation {s+1}: {concat_fp}')
        
        # Save camera poses (input + output, absolute cam2world) to .npy file
        poses_fp = os.path.join(output_path, f'{fn_prefix}{fn_idx:02d}{angle_suffix}-poses.npy')
        gradio_utils.save_camera_poses_to_npy(base_data_batch, poses_fp)
        
        # Use first concatenation video for display, or None if none created
        cat_video = cat_videos[0] if len(cat_videos) > 0 else None
        
        # Look for gal.mp4 visualization files for THIS run (use timestamp to filter)
        # Pattern: dvs2_val_s0_gradio--kubric5d_d11_r0_i0_c0,12_gal.mp4
        gal_pattern = os.path.join(output_path, f'*_val_*_gradio*_i*_*_gal.mp4')
        gal_files = sorted(glob.glob(gal_pattern), key=os.path.getmtime, reverse=True)
        gal_video = gal_files[0] if gal_files and os.path.getmtime(gal_files[0]) > time.time() - 300 else None
        
        if cat_video:
            print(f'[green]Created Input|Pred concatenation: {cat_video}')
        else:
            print(f'[yellow]No concatenation video created')
        
        if gal_video:
            print(f'[green]Found gal visualization: {gal_video}')
        else:
            print(f'[yellow]No gal file found matching: {gal_pattern}')
        
        # Create camera visualization and extract extrinsics
        camera_plot, extrinsics_text = gradio_utils.extract_and_visualize_cameras(
            base_data_batch,
        )
        
        # Build description with uncertainty metrics if available
        uncertainty_text = ""
        if len(uncertainty_videos) > 0 and len(diversity_metrics) > 0:
            uncertainty_text = "\n\n**Uncertainty metrics:**\n"
            for entry_name, metrics in diversity_metrics.items():
                uncertainty_text += f"- {entry_name}: mean diversity = {metrics['mean_diversity']:.4f}\n"
        
        # Build resolution info with both input (rgb1) and output (rgb0)
        # input_rgb shape is (T, C, H, W)
        input_res = f"{input_rgb.shape[3]} x {input_rgb.shape[2]}" if input_rgb is not None else "N/A"
        output_res = f"{output_resolution[0]} x {output_resolution[1]}" if output_resolution is not None else "N/A"
        
        description = f'''✅ **Done!** Generated {len(pred_samples)} sample(s).

**Debug info:**
- Input frames: {num_frames}
- Input resolution (rgb1): {input_res}
- Output resolution (rgb0): {output_res}
- Diffusion steps: {num_steps}
{uncertainty_text}
**Camera extrinsics:\n**
{extrinsics_text}

**Output files saved to:** `{output_path}`
'''
        
        return (description, *output_videos, cat_video, gal_video, camera_plot)
        
    except Exception as e:
        import traceback
        tb_str = traceback.format_exc()
        print(f'[red]Error saving results: {e}')
        print(f'[red]{tb_str}')
        
        error_msg = f'''❌ **Error saving results**

**Error message**: {str(e)}

**Stacktrace**:
```
{tb_str}
```
'''
        return (error_msg, None, None, None, None, None, None, None)


def run_test_inference(model_bundle, base_data_batch, output_path, test_stub):
    '''
    Run test inference on a stub video to verify pipeline.
    :param model_bundle: Model bundle (model, config, exp_cfg, device)
    :param base_data_batch: Base data batch from vidar dataset
    '''
    if test_stub == 0:
        print('[yellow]Test stub disabled (test_stub=0)')
        return
    
    print('[bold cyan]═══════════════════════════════════════')
    print('[bold cyan]Running test inference on first Kubric5D sample...')
    print('[bold cyan]═══════════════════════════════════════')
    
    # Extract RGB frames from the base_data_batch (first Kubric5D sample)
    rgb_dict = base_data_batch['anydata']['rgb']
    input_rgb, time_indices, error_msg = extract_frames_from_rgb_dict(rgb_dict, cam_idx=0)
    
    if input_rgb is None:
        print(f'[red]Error: {error_msg} in base_data_batch')
        return
    
    num_frames, C, frame_height, frame_width = input_rgb.shape
    print(f'[cyan]Using frames {time_indices[0]}-{time_indices[-1]} from base_data_batch')
    
    print(f'[cyan]Using Kubric5D sample with shape ({num_frames}x{frame_height}x{frame_width})...')
    
    try:
        import signal
        import tempfile
        
        # Save input_rgb to a temporary video file
        temp_video_path = os.path.join(output_path, 'test_stub_input.mp4')
        
        # Convert tensor to numpy and save as video
        # input_rgb is (T, C, H, W), need to convert to (T, H, W, C) for saving
        input_rgb_np = rearrange(input_rgb, 'T C H W -> T H W C').cpu().numpy()
        input_rgb_np = np.clip(input_rgb_np, 0.0, 1.0)
        
        print(f'[cyan]Saving test input video to: {temp_video_path}')
        import infer_utils
        infer_utils.save_video(temp_video_path, input_rgb_np, fps=20, quality=9)  # Use default frame rate
        print(f'[green]Saved test input video')
        
        # Set a timeout for test inference to avoid hanging
        def timeout_handler(signum, frame):
            raise TimeoutError("Test inference timed out")
        
        signal.signal(signal.SIGALRM, timeout_handler)
        signal.alarm(120)  # 120 second timeout (increased for main_run)
        
        # Run main_run to test the full pipeline
        result = main_run(
            model_bundle=model_bundle,
            base_data_batch=base_data_batch,
            output_path=output_path,
            action='generate',
            which_video='Upload',
            raw_video=temp_video_path,
            which_camera='Canonical',
            num_frames=num_frames,
            frame_rate=12,
            num_steps=25,
            num_samples=1,
        )
        
        signal.alarm(0)  # Cancel timeout
        
        print('[bold green]Test inference with main_run completed successfully!')
        print(f'[green]Result: {result}')
        
    except TimeoutError as e:
        print(f'[yellow]Test inference timed out - skipping')
        print('[yellow]Gradio app will still work for uploads')
    
    except Exception as e:
        print(f'[yellow]Test inference failed: {e}')
        print('[yellow]Gradio app will still work for uploads')
        import traceback
        print(traceback.format_exc())


def run_demo(device='cuda', gpu_id=0,
            port=7860, share=False,
            exp_cfg='custom/experiment/basile/eval_cvpr_qual.py',
            ckpt_path=None,
            output_dir='gradio_out_dbg/',
            test_stub=0,
            dset_cfg='custom/config/cvpr/quant_v2_HR/kubric5d.yaml',
            default_resolution='384 x 256',
            anyviewbench_yaml_dir='custom/config/cvpr/quant_v2_HR',
            anyviewbench_txt_dir='filenames/quant_v2'):
    '''
    Launch Gradio demo interface.
    :param dset_cfg: Path to vidar dataset yaml config for loading sample metadata
    '''
    # Setup device
    if 'cuda' in device:
        device = f'cuda:{gpu_id}'
        torch.cuda.set_device(gpu_id)
    
    # Initialize torch.distributed for single GPU (required by some model code)
    if not torch.distributed.is_initialized():
        import random as rand_mod
        port_dist = rand_mod.randint(29500, 29600)
        torch.distributed.init_process_group(
            backend='gloo',
            init_method=f'tcp://127.0.0.1:{port_dist}',
            world_size=1,
            rank=0
        )
    
    print('[bold cyan]═══════════════════════════════════════')
    print('[bold cyan]  Any4D / AnyView Gradio Demo')
    print('[bold cyan]═══════════════════════════════════════')
    print(f'[cyan]Device: {device}')
    print(f'[cyan]Experiment config: {exp_cfg}')
    print(f'[cyan]Output directory: {output_dir}')
    
    # Create output directory
    os.makedirs(output_dir, exist_ok=True)
    
    # Load model
    print('[cyan]Loading model...')
    model, config, exp_config = infer_utils.load_model_bundle(
        exp_cfg, ckpt_path=ckpt_path, device=device)
    
    # Initialize validation metrics list (required by validation_step)
    model.all_val_metrics = []
    
    experiment_name = exp_config['job']['name']
    # Get the actual checkpoint path used
    actual_ckpt_path = ckpt_path if ckpt_path else config.model_manager_config.dit_path
    print(f'[green]Model loaded: {experiment_name}')
    print(f'[green]Checkpoint: {actual_ckpt_path}')
    
    model_bundle = (model, config, exp_config, device)
    
    # Load dataset with dataloader for proper batching (like infer.py)
    print(f'[cyan]Loading dataset from {dset_cfg}...')
    print(f'[cyan]Using VIDAR convention: cam 0=output, cam 1=input')
    
    try:
        from torch.utils.data import DataLoader
        
        dataset = VidarDataset(
            dataset_dir=dset_cfg,
            phase='test',  # Use test phase
        )
        print(f'[green]Dataset loaded: {len(dataset)} samples')
        
        # Create dataloader with batch_size=1 to get properly batched template
        dataloader = DataLoader(
            dataset,
            batch_size=1,
            shuffle=True,  # Shuffle to get random samples
            num_workers=0,
            pin_memory=True,
        )
        
        # Get first batch as template (properly collated by DataLoader)
        base_data_batch = next(iter(dataloader))
        print(f'[green]Using sample 0 as template for data batch structure (via DataLoader)')
        
    except Exception as e:
        print(f'[red]Error loading dataset: {e}')
        print('[yellow]Falling back to manual data batch construction (may fail)')
        base_data_batch = None
    
    # Run test inference if requested
    if test_stub > 0 and base_data_batch is not None:
        run_test_inference(model_bundle, base_data_batch, output_dir, test_stub)
    
    # Create Gradio interface
    print('[cyan]Building Gradio interface...')

    frames_options = ['13', '29', '41', '81', '121']
    default_frames = '41'
    
    # resolution_options = ['1024 x 576', '768 x 512', '576 x 384', '512 x 288', '384 x 256']
    resolution_options = ['576 x 384', '576 x 368', '576 x 352', '576 x 320',
                          '384 x 256', '384 x 240', '384 x 224', '384 x 208']
    # default_resolution is now passed as a function parameter
    
    demo = gr.Blocks(title=_TITLE.format(default_resolution.split(' x ')[0]))
    
    with demo:
        gr.Markdown('# ' + _TITLE.format(default_resolution.split(' x ')[0]))
        gr.Markdown(_DESCRIPTION.format(checkpoint_path=actual_ckpt_path, dset_cfg=dset_cfg))
        
        with gr.Row():
            with gr.Column(scale=8, variant='panel'):
                
                gr.Markdown('*Video clip processing options:*')
                num_frames_sld = gr.Slider(
                    9, 121, value=41, step=4,
                    label='Number of frames')
                frame_offset_sld = gr.Slider(
                    0, 100, value=0, step=1,
                    label='Frame offset (start later)')
                frame_stride_sld = gr.Slider(
                    1, 10, value=2, step=1,
                    label='Frame stride (temporally subsample)')
                frame_rate_sld = gr.Slider(
                    6, 40, value=20, step=2,
                    label='Frame rate (after subsampling)')
                
                resolution_rad = gr.Radio(
                    resolution_options,
                    value=default_resolution,
                    label='Select resolution:')
                center_crop_chk = gr.Checkbox(
                    True, label='Center crop to correct aspect ratio')
                override_resolution_chk = gr.Checkbox(
                    True, label='[Tarfile] Override resolution with selected')
                
                # Tarfile indices selector (applies to both video and camera when using tarfiles)
                tarfile_indices_txt = gr.Textbox(
                    label='[Tarfile] Camera indices (format: in,out) (for video + pose):',
                    placeholder='1,0',
                    value='',
                    lines=1)

                gr.Markdown('*Model inference options:*')
                samples_sld = gr.Slider(
                    1, 4, value=1, step=1,
                    label='Number of samples to generate')
                steps_sld = gr.Slider(
                    5, 100, value=25, step=5,
                    label='Number of diffusion inference timesteps')
                # guidance_sld = gr.Slider(
                #     1.0, 10.0, value=3.0, step=0.5,
                #     label='Guidance scale (CFG)')
                
                desc_output = gr.Markdown('Results will appear on the right.', height=480)
            
            with gr.Column(scale=10, variant='panel'):
                
                gr.Markdown('*Input and conditioning data:*')
                
                video_block = gr.Video(
                    sources=['upload', 'webcam'],
                    include_audio=False,
                    label='Upload video file')
                video_path_txt = gr.Textbox(
                    label='Enter video MP4 or tarfile or yaml path (S3 or DGX), or YouTube URL:',
                    placeholder=('DGX MP4 path: /path/to/video.mp4\n'
                                 'DGX tarfile path: /path/to/camera.tar\n'
                                 'DGX yaml path: /path/to/dataset.yaml\n'
                                 'S3 MP4 path: s3://bucket/path/video.mp4\n'
                                 'S3 tarfile path: s3://bucket/path/camera.tar\n'
                                 'YouTube URL: https://youtube.com/watch?v=...'),
                    lines=3)
                
                which_video = gr.Radio(
                    ['Upload', 'Path', 'Reference'],
                    value='Upload',
                    label='Select which video to use:')
                
                with gr.Row():
                    preview_btn = gr.Button('Preview Input Video From Path', variant='secondary', size='sm')
                    random_video_btn = gr.Button('Load Random In-The-Wild Video', variant='secondary', size='sm')
                    random_anyviewbench_tar_btn = gr.Button('Load Random AnyViewBench (TAR)', variant='secondary', size='sm')
                    random_anyviewbench_yaml_btn = gr.Button('Load Random AnyViewBench (YAML)', variant='secondary', size='sm')
                
                preview_feedback = gr.Markdown('', visible=True)
                
                gr.Markdown('**Camera Parameters:**')
                
                with gr.Tabs() as camera_tabs:
                    with gr.TabItem('Canonical', id='canonical_tab'):
                        
                        gr.Markdown('*Presets:*')
                        with gr.Row():
                            preset_exo_rotleft_btn = gr.Button('Rotate Left', variant='secondary', size='sm')
                            preset_exo_rotright_btn = gr.Button('Rotate Right', variant='secondary', size='sm')
                            preset_exo_rotup_btn = gr.Button('Rotate Up', variant='secondary', size='sm')
                            preset_ego_turnleft_btn = gr.Button('Turn Left', variant='secondary', size='sm')
                            preset_ego_turnright_btn = gr.Button('Turn Right', variant='secondary', size='sm')
                            preset_ego2exo_moveback_btn = gr.Button('Move Back', variant='secondary', size='sm')
                            preset_ego2exo_movebackleft_btn = gr.Button('Move Back+Left', variant='secondary', size='sm')
                            preset_ego2exo_movebackright_btn = gr.Button('Move Back+Right', variant='secondary', size='sm')
                            preset_ego2exo_movebackup_btn = gr.Button('Move Back+Up', variant='secondary', size='sm')

                        gr.Markdown('*Synthetic camera trajectory based on spherical angles:*')
                        which_setting = gr.Radio(
                            ['Exo2Exo', 'Ego2Ego', 'Ego2Exo'],
                            value='Exo2Exo',
                            label='Select which setting to use:')
                        canonical_azimuth_deg = gr.Slider(
                            -180, 180, value=30, step=5,
                            label='Azimuth angle (degrees, negative=left, positive=right):')
                        canonical_input_elev_deg = gr.Slider(
                            -60, 60, value=10, step=5,
                            label='Input elevation (degrees from horizon, 0=horizontal):')
                        canonical_output_elev_deg = gr.Slider(
                            -60, 60, value=10, step=5,
                            label='Output elevation (degrees from horizon):')
                        canonical_radius = gr.Slider(
                            0, 20, value=5, step=1,
                            label='Radial distance (meters) from scene center:')
                        canonical_trajectory = gr.Radio(
                            ['None (Direct)', 'Gradual (Linear)', 'Gradual (Sqrt)'],
                            value='Gradual (Sqrt)',
                            label='Trajectory interpolation type:')
                        
                        with gr.Row():
                            vis_cam_btn_canonical = gr.Button('Visualize Camera Poses', variant='secondary')
                            run_btn_canonical = gr.Button('Run AnyView Generation!', variant='primary')
                    
                    with gr.TabItem('Path', id='path_tab'):
                        gr.Markdown('*Load camera trajectory from file (tar/yaml):*')
                        camera_path_txt = gr.Textbox(
                            label='Path to extract camera from (tar/yaml):',
                            placeholder=('DGX tarfile path: /path/to/camera.tar\n'
                                         'DGX yaml path: /path/to/dataset.yaml\n'
                                         'S3 tarfile path: s3://bucket/path/camera.tar'),
                            lines=4)
                        
                        invert_camera_pose = gr.Checkbox(
                            True, 
                            label='[Tarfile] Invert camera extrinsics (check if poses seem weird)')
                        
                        with gr.Row():
                            vis_cam_btn_path = gr.Button('Visualize Camera Poses', variant='secondary')
                            run_btn_path = gr.Button('Run AnyView Generation!', variant='primary')
                    
                    with gr.TabItem('Reference', id='reference_tab'):
                        gr.Markdown('This will use camera parameters from the loaded vidar dataset config (no modifications), mostly for debugging.')
                        
                        with gr.Row():
                            vis_cam_btn_reference = gr.Button('Visualize Camera Poses', variant='secondary')
                            run_btn_reference = gr.Button('Run AnyView Generation!', variant='primary')
            
            with gr.Column(scale=12, variant='panel'):

                gen1_output = gr.Video(
                    format='mp4',
                    label='Generated video from new viewpoint (Sample 1)',
                    visible=True)
                gen2_output = gr.Video(
                    format='mp4',
                    label='Generated video from new viewpoint (Sample 2)',
                    visible=False)
                gen3_output = gr.Video(
                    format='mp4',
                    label='Generated video from new viewpoint (Sample 3)',
                    visible=False)
                gen4_output = gr.Video(
                    format='mp4',
                    label='Generated video from new viewpoint (Sample 4)',
                    visible=False)

                cat1_output = gr.Video(
                    format='mp4',
                    label='Concatenated video (Input | Pred [| Uncertainty if >1 sample]) (Sample 1)',
                    visible=True)
                gal1_output = gr.Video(
                    format='mp4',
                    label='Full gallery visualization (Sample 1)',
                    visible=True)
                
                camera_plot_output = gr.Plot(
                    label='Camera Trajectory 3D Visualization (XYZ = RGB)',
                    visible=True)
        
        # Set up preview button click handler with auto-switch to Path if path textbox has content
        def preview_with_auto_switch(which_video, uploaded_video, video_path, num_frames, frame_offset, 
                                     frame_stride, center_crop, resolution, frame_rate, tarfile_indices):
            # Auto-switch to 'Path' if video_path textbox has content
            if video_path and video_path.strip() and which_video != 'Reference':
                which_video = 'Path'
            
            # Call the actual preview function
            video_result, feedback = preview_process_video(
                which_video, uploaded_video, video_path, num_frames, frame_offset,
                frame_stride, center_crop, resolution, frame_rate, tarfile_indices, base_data_batch
            )
            
            return gr.update(value=which_video), video_result, feedback
        
        preview_btn.click(
            fn=preview_with_auto_switch,
            inputs=[which_video, video_block, video_path_txt, num_frames_sld, frame_offset_sld, 
                   frame_stride_sld, center_crop_chk, resolution_rad, frame_rate_sld, tarfile_indices_txt],
            outputs=[which_video, video_block, preview_feedback]
        )
        
        # Set up random video button handler
        # First pick random video, then preview it
        def random_video_handler():
            # Return "Path" to switch radio, and pick random video
            video_path, feedback = pick_random_video()
            return gr.update(value='Path'), video_path, feedback
        
        random_video_btn.click(
            fn=random_video_handler,
            inputs=[],
            outputs=[which_video, video_path_txt, preview_feedback]
        ).then(
            fn=partial(preview_process_video, base_data_batch=base_data_batch),
            inputs=[which_video, video_block, video_path_txt, num_frames_sld, frame_offset_sld, 
                   frame_stride_sld, center_crop_chk, resolution_rad, frame_rate_sld, tarfile_indices_txt],
            outputs=[video_block, preview_feedback]
        )
        
        # Set up random AnyViewBench TAR button handler
        def random_anyviewbench_tar_handler():
            # Return "Path" to switch video radio, pick random tarfile for both video and camera, and switch to Path camera tab
            tarfile_path, tarfile_path_dup, invert_checkbox_val, feedback = pick_random_anyviewbench_tar(anyviewbench_yaml_dir, anyviewbench_txt_dir)
            return gr.update(value='Path'), tarfile_path, tarfile_path_dup, gr.update(selected='path_tab'), invert_checkbox_val, feedback
        
        random_anyviewbench_tar_btn.click(
            fn=random_anyviewbench_tar_handler,
            inputs=[],
            outputs=[which_video, video_path_txt, camera_path_txt, camera_tabs, invert_camera_pose, preview_feedback]
        ).then(
            fn=partial(preview_process_video, base_data_batch=base_data_batch),
            inputs=[which_video, video_block, video_path_txt, num_frames_sld, frame_offset_sld, 
                   frame_stride_sld, center_crop_chk, resolution_rad, frame_rate_sld, tarfile_indices_txt],
            outputs=[video_block, preview_feedback]
        )
        
        # Set up random AnyViewBench YAML button handler
        def random_anyviewbench_yaml_handler():
            # Return "Path" to switch video radio, pick random YAML for both video and camera, and switch to Path camera tab
            yaml_path, yaml_path_dup, feedback = pick_random_anyviewbench_yaml(anyviewbench_yaml_dir)
            return gr.update(value='Path'), yaml_path, yaml_path_dup, gr.update(selected='path_tab'), feedback
        
        random_anyviewbench_yaml_btn.click(
            fn=random_anyviewbench_yaml_handler,
            inputs=[],
            outputs=[which_video, video_path_txt, camera_path_txt, camera_tabs, preview_feedback]
        ).then(
            fn=partial(preview_process_video, base_data_batch=base_data_batch),
            inputs=[which_video, video_block, video_path_txt, num_frames_sld, frame_offset_sld, 
                   frame_stride_sld, center_crop_chk, resolution_rad, frame_rate_sld, tarfile_indices_txt],
            outputs=[video_block, preview_feedback]
        )
        
        # Set up visualize camera trajectory buttons - one for each tab
        # Canonical tab
        vis_cam_btn_canonical.click(
            fn=partial(visualize_cameras_only, base_data_batch),
            inputs=[gr.State('Canonical'), invert_camera_pose, which_setting,
                   canonical_azimuth_deg, canonical_input_elev_deg, canonical_output_elev_deg,
                   canonical_radius, canonical_trajectory,
                   camera_path_txt, tarfile_indices_txt, num_frames_sld],
            outputs=[desc_output, camera_plot_output]
        )
        
        # Path tab
        vis_cam_btn_path.click(
            fn=partial(visualize_cameras_only, base_data_batch),
            inputs=[gr.State('Path'), invert_camera_pose, which_setting,
                   canonical_azimuth_deg, canonical_input_elev_deg, canonical_output_elev_deg,
                   canonical_radius, canonical_trajectory,
                   camera_path_txt, tarfile_indices_txt, num_frames_sld],
            outputs=[desc_output, camera_plot_output]
        )
        
        # Reference tab
        vis_cam_btn_reference.click(
            fn=partial(visualize_cameras_only, base_data_batch),
            inputs=[gr.State('Reference'), invert_camera_pose, which_setting,
                   canonical_azimuth_deg, canonical_input_elev_deg, canonical_output_elev_deg,
                   canonical_radius, canonical_trajectory,
                   camera_path_txt, tarfile_indices_txt, num_frames_sld],
            outputs=[desc_output, camera_plot_output]
        )
        
        # Set up preset button click handlers using CAMERA_PRESETS dict
        preset_button_map = {
            'Rotate Left': preset_exo_rotleft_btn,
            'Rotate Right': preset_exo_rotright_btn,
            'Rotate Up': preset_exo_rotup_btn,
            'Turn Left': preset_ego_turnleft_btn,
            'Turn Right': preset_ego_turnright_btn,
            'Move Back': preset_ego2exo_moveback_btn,
            'Move Back+Left': preset_ego2exo_movebackleft_btn,
            'Move Back+Right': preset_ego2exo_movebackright_btn,
            'Move Back+Up': preset_ego2exo_movebackup_btn,
        }
        
        for preset_name, button in preset_button_map.items():
            setting, azimuth, input_elev, output_elev, radius = CAMERA_PRESETS[preset_name]
            button.click(
                fn=partial(visualize_cameras_only, base_data_batch),
                inputs=[gr.State('Canonical'), invert_camera_pose, gr.State(setting),
                       gr.State(azimuth), gr.State(input_elev), gr.State(output_elev),
                       gr.State(radius), canonical_trajectory,
                       camera_path_txt, tarfile_indices_txt, num_frames_sld],
                outputs=[desc_output, camera_plot_output]
            ).then(
                lambda p=CAMERA_PRESETS[preset_name]: p,
                outputs=[which_setting, canonical_azimuth_deg, canonical_input_elev_deg, 
                        canonical_output_elev_deg, canonical_radius]
            )
        
        # Set up run button click handlers - one for each tab
        my_outputs = [desc_output, 
                      gen1_output, gen2_output, gen3_output, gen4_output,
                      cat1_output, gal1_output, camera_plot_output]
        
        # Canonical tab
        my_inputs_canonical = [which_video, video_block, video_path_txt,
                     gr.State('Canonical'), invert_camera_pose, which_setting,
                     canonical_azimuth_deg, canonical_input_elev_deg, canonical_output_elev_deg,
                     canonical_radius, canonical_trajectory,
                     camera_path_txt, tarfile_indices_txt,
                     num_frames_sld, frame_offset_sld, frame_stride_sld, frame_rate_sld,
                     center_crop_chk, resolution_rad, override_resolution_chk,
                     samples_sld, steps_sld]
        
        run_btn_canonical.click(
            fn=partial(main_run, model_bundle, base_data_batch, output_dir, 'run'),
            inputs=my_inputs_canonical,
            outputs=my_outputs
        )
        
        # Path tab
        my_inputs_path = [which_video, video_block, video_path_txt,
                     gr.State('Path'), invert_camera_pose, which_setting,
                     canonical_azimuth_deg, canonical_input_elev_deg, canonical_output_elev_deg,
                     canonical_radius, canonical_trajectory,
                     camera_path_txt, tarfile_indices_txt,
                     num_frames_sld, frame_offset_sld, frame_stride_sld, frame_rate_sld,
                     center_crop_chk, resolution_rad, override_resolution_chk,
                     samples_sld, steps_sld]
        
        run_btn_path.click(
            fn=partial(main_run, model_bundle, base_data_batch, output_dir, 'run'),
            inputs=my_inputs_path,
            outputs=my_outputs
        )
        
        # Reference tab
        my_inputs_reference = [which_video, video_block, video_path_txt,
                     gr.State('Reference'), invert_camera_pose, which_setting,
                     canonical_azimuth_deg, canonical_input_elev_deg, canonical_output_elev_deg,
                     canonical_radius, canonical_trajectory,
                     camera_path_txt, tarfile_indices_txt,
                     num_frames_sld, frame_offset_sld, frame_stride_sld, frame_rate_sld,
                     center_crop_chk, resolution_rad, override_resolution_chk,
                     samples_sld, steps_sld]
        
        run_btn_reference.click(
            fn=partial(main_run, model_bundle, base_data_batch, output_dir, 'run'),
            inputs=my_inputs_reference,
            outputs=my_outputs
        )
    
    print('[bold green]Launching Gradio interface...')
    demo.queue(max_size=20)
    demo.launch(
        share=share,
        server_port=port,
        server_name='0.0.0.0',
        allowed_paths=[output_dir, RANDOM_VIDEO_DIR]
    )


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Any4D Gradio Demo')
    
    parser.add_argument('--device', type=str, default='cuda',
                       help='Device to use')
    parser.add_argument('--gpu_id', type=int, default=0,
                       help='GPU ID to use')
    parser.add_argument('--port', type=int, default=7860,
                       help='Port to run Gradio server on')
    parser.add_argument('--share', type=int, default=1,
                       help='Create shareable Gradio link')
    parser.add_argument('--exp_cfg', type=str, required=True,
                       help='Path to experiment config file')
    parser.add_argument('--ckpt_path', type=str, default=None,
                       help='Optional checkpoint path override')
    parser.add_argument('--output_dir', type=str, default='gradio_out_dbg/',
                       help='Output directory for results')
    parser.add_argument('--test_stub', type=int, default=0,
                       help='Test stub video ID (0=disabled, 1+=various test videos)')
    parser.add_argument('--dset_cfg', type=str, 
                       default='custom/config/cvpr/quant_v2_HR/kubric5d.yaml',
                       help='Path to vidar dataset yaml config for loading sample metadata')
    parser.add_argument('--anyviewbench_yaml_dir', type=str, 
                       default='custom/config/cvpr/quant_v2_HR',
                       help='Directory containing AnyViewBench YAML config files')
    parser.add_argument('--anyviewbench_txt_dir', type=str, 
                       default='filenames/quant_v2',
                       help='Directory containing AnyViewBench TXT files with tarfile paths')
    parser.add_argument('--default_resolution', type=str, default='384 x 256',
                       help='Default resolution for video processing (e.g., "384 x 256")')
    
    args = parser.parse_args()
    
    run_demo(
        device=args.device,
        gpu_id=args.gpu_id,
        port=args.port,
        share=bool(args.share),
        exp_cfg=args.exp_cfg,
        ckpt_path=args.ckpt_path,
        output_dir=args.output_dir,
        test_stub=args.test_stub,
        dset_cfg=args.dset_cfg,
        default_resolution=args.default_resolution,
        anyviewbench_yaml_dir=args.anyviewbench_yaml_dir,
        anyviewbench_txt_dir=args.anyviewbench_txt_dir,
    )