import torch

print(torch.__version__)
print(torch.version.cuda)
print(torch.cuda.is_available())

import os, subprocess, shutil
import uuid, tempfile
from glob import glob

env_list = os.environ['PATH'].split(':')
env_list.append('/usr/local/cuda/bin')
os.environ['PATH'] = ':'.join(env_list)
os.environ['TORCH_CUDA_ARCH_LIST'] = '8.6'

import gradio as gr
from huggingface_hub import snapshot_download

os.makedirs("pretrained", exist_ok=True)
snapshot_download(
    repo_id = "jiawei011/L4GM",
    local_dir = "./pretrained"
)

# Folder containing example images
examples_folder = "data_test"

# Retrieve all file paths in the folder
video_examples = [
    os.path.join(examples_folder, file)
    for file in os.listdir(examples_folder)
    if os.path.isfile(os.path.join(examples_folder, file))
]


def generate(input_video):

    unique_id = str(uuid.uuid4())
    
    workdir = f"results_{unique_id}"
    recon_model = "pretrained/recon.safetensors"
    interp_model = "pretrained/interp.safetensors"
    num_frames = 16
    test_path = input_video
    
    try:
        # Run the inference command
        subprocess.run(
            [
                "python", "infer_3d.py", "big",
                "--workspace", f"{workdir}",
                "--resume", f"{recon_model}",
                "--num_frames", f"1",
                "--test_path", f"{test_path}",
            ],
            check=True
        )

        subprocess.run(
            [
                "python", "infer_4d.py", "big",
                "--workspace", f"{workdir}",
                "--resume", f"{recon_model}",
                "--interpresume", f"{interp_model}",
                "--num_frames", f"{num_frames}",
                "--test_path", f"{test_path}",
            ],
            check=True
        )

        # Get all .mp4 files in the workdir
        output_videos = glob(os.path.join(workdir, "*.mp4"))
        print("Found videos:", output_videos)
    
        # Check if the 5th video exists
        if len(output_videos) < 5:
            raise IndexError("Less than 5 .mp4 files found in the workdir.")
    
        # Get the 5th video
        selected_video = output_videos[4]
        print("Selected video:", selected_video)
    
        # Create a new temporary directory
        temp_dir = tempfile.mkdtemp()
        print("Temporary directory created:", temp_dir)
    
        # Copy the selected video to the temporary directory
        new_video_path = os.path.join(temp_dir, os.path.basename(selected_video))
        shutil.copy(selected_video, new_video_path)
        print(f"Copied {selected_video} to {new_video_path}")
    
        # Delete the workdir folder
        shutil.rmtree(workdir)
        print(f"Deleted workdir: {workdir}")
    
        # Return the new path of the copied video
        return new_video_path
   
    except subprocess.CalledProcessError as e:
        raise gr.Error(f"Error during inference: {str(e)}")

with gr.Blocks() as demo:
    with gr.Column():
        gr.Markdown("# L4GM: Large 4D Gaussian Reconstruction Model")
        gr.HTML("""
            <div style="display:flex;column-gap:4px;">
                <a href="https://github.com/nv-tlabs/L4GM-official/tree/main">
                    <img src='https://img.shields.io/badge/GitHub-Repo-blue'>
                </a> 
                <a href="https://research.nvidia.com/labs/toronto-ai/l4gm/">
                    <img src='https://img.shields.io/badge/Project-Page-green'>
                </a>
    			<a href="https://arxiv.org/abs/2406.10324">
                    <img src='https://img.shields.io/badge/ArXiv-Paper-red'>
                </a>
                <a href="https://huggingface.co/spaces/fffiloni/L4GM-demo?duplicate=true">
    				<img src="https://huggingface.co/datasets/huggingface/badges/resolve/main/duplicate-this-space-sm.svg" alt="Duplicate this Space">
    			</a>
    			<a href="https://huggingface.co/fffiloni">
    				<img src="https://huggingface.co/datasets/huggingface/badges/resolve/main/follow-me-on-HF-sm-dark.svg" alt="Follow me on HF">
    			</a>
            </div>
         """)
        with gr.Row():
            with gr.Column():
                input_video = gr.Video(label="Input Video", interactive=False)
                submit_btn = gr.Button("Submit")
                gr.Examples(
                    examples = video_examples,
                    inputs = [input_video],
                    examples_per_page = 5
                )
            with gr.Column():
    
                output_result_4 = gr.Video(label="Result")

    submit_btn.click(
        fn = generate,
        inputs = [input_video],
        outputs = [
            output_result_4
        ]
    )

demo.queue().launch(show_api=False, show_error=True)