# build upon InstantSplat https://huggingface.co/spaces/kairunwen/InstantSplat/blob/main/app.py
import os, subprocess, shlex, sys, gc
import numpy as np
import shutil
import argparse
import gradio as gr
import uuid
import glob
import re
import torch
import spaces

subprocess.run(shlex.split("pip install wheel/diff_gaussian_rasterization-0.0.0-cp310-cp310-linux_x86_64.whl"))
subprocess.run(shlex.split("pip install wheel/simple_knn-0.0.0-cp310-cp310-linux_x86_64.whl"))
subprocess.run(shlex.split("pip install wheel/curope-0.0.0-cp310-cp310-linux_x86_64.whl"))
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
os.sys.path.append(os.path.abspath(os.path.join(BASE_DIR, "dynamic_predictor")))
os.sys.path.append(os.path.abspath(os.path.join(BASE_DIR)))
GRADIO_CACHE_FOLDER = './gradio_cache_folder'

from dynamic_predictor.launch import main as dynamic_predictor_main
from utils_das3r.rearrange import main as rearrange_main
from train_gui import main as train_main
from render import main as render_main


def natural_sort(l): 
    convert = lambda text: int(text) if text.isdigit() else text.lower()
    alphanum_key = lambda key: [convert(c) for c in re.split('([0-9]+)', key.split('/')[-1])]
    return sorted(l, key=alphanum_key)

def cmd(command):
    print(command)
    os.system(command)

@spaces.GPU(duration=150)
def process(inputfiles, input_path='demo'):
    if inputfiles:
        frames = natural_sort(inputfiles)
    else:
        frames = natural_sort(glob.glob('./assets/example/' + input_path + '/*'))
    if len(frames) > 20:
        stride = int(np.ceil(len(frames) / 20))
        frames = frames[::stride]
    
    # Create a temporary directory to store the selected frames
    temp_dir = os.path.join(GRADIO_CACHE_FOLDER, str(uuid.uuid4()))
    os.makedirs(temp_dir, exist_ok=True)
    
    # Copy the selected frames to the temporary directory
    for i, frame in enumerate(frames):
        shutil.copy(frame, f"{temp_dir}/{i:04d}.{frame.split('.')[-1]}")

    imgs_path = temp_dir
    output_path = f'./results/{input_path}/output'
    rearranged_path = f'{output_path}_rearranged'
    # cmd(f"python dynamic_predictor/launch.py --mode=eval_pose_custom \
    #     --pretrained=Kai422kx/das3r \
    #     --dir_path={imgs_path} \
    #     --output_dir={output_path} \
    #     --use_pred_mask ")
    dynamic_predictor_main(pretrained='Kai422kx/das3r', dir_path=imgs_path, output_dir=output_path, use_pred_mask=True, n_iter=150)
    rearrange_main(output_dir=output_path, rearranged_path = rearranged_path)
    train_main(s = rearranged_path, m = rearranged_path, iter = 2000)
    render_main(s = rearranged_path, m = rearranged_path, iter = 2000, get_video = True)



    output_video_path = f"{rearranged_path}/rendered.mp4"
    output_ply_path = f"{rearranged_path}/point_cloud/iteration_2000/point_cloud.ply"
    return  output_video_path, output_ply_path, output_ply_path



_TITLE = '''DAS3R'''
_DESCRIPTION = '''
<div style="display: flex; justify-content: center; align-items: center;">
    <div style="width: 100%; text-align: center; font-size: 30px;">
        <strong>DAS3R: Dynamics-Aware Gaussian Splatting for Static Scene Reconstruction</strong>
    </div>
</div> 
<p></p>


<div align="center">
    <a style="display:inline-block" href="https://arxiv.org/abs/2412.19584"><img src="https://img.shields.io/badge/ArXiv-2412.19584-b31b1b.svg?logo=arXiv" alt='arxiv'></a>
    <a style="display:inline-block" href="https://kai422.github.io/DAS3R/"><img src='https://img.shields.io/badge/Project-Website-blue.svg'></a>
    <a style="display:inline-block" href="https://github.com/kai422/DAS3R"><img src='https://img.shields.io/badge/GitHub-%23121011.svg?logo=github&logoColor=white'></a>
</div>
<p></p>


* Official demo of [DAS3R: Dynamics-Aware Gaussian Splatting for Static Scene Reconstruction](https://kai422.github.io/DAS3R/).
* You can explore the sample results by clicking the sequence names at the bottom of the page.
* Due to GPU memory and time limitations, processing is restricted to 20 frames and 2000 GS training iterations. Uniform sampling is applied if input frames exceed 20. 
* This Gradio demo is built upon InstantSplat, which can be found at [https://huggingface.co/spaces/kairunwen/InstantSplat](https://huggingface.co/spaces/kairunwen/InstantSplat).

'''

block = gr.Blocks().queue()
with block:
    with gr.Row():
        with gr.Column(scale=1):
            # gr.Markdown('# ' + _TITLE)
            gr.Markdown(_DESCRIPTION)
    
    with gr.Row(variant='panel'):
        with gr.Tab("Input"):
            inputfiles = gr.File(file_count="multiple", label="images")
            input_path = gr.Textbox(visible=False, label="example_path")
            button_gen = gr.Button("RUN")

    with gr.Row(variant='panel'):
        with gr.Tab("Output"):
            with gr.Column(scale=2):
                with gr.Group():
                    output_model = gr.Model3D(
                        label="3D Dense Model under Gaussian Splats Formats, need more time to visualize",
                        interactive=False,
                        camera_position=[0.5, 0.5, 1],  # 稍微偏移一点，以便更好地查看模型
                    )
                    gr.Markdown(
                        """
                        <div class="model-description">
                           &nbsp;&nbsp;Use the left mouse button to rotate, the scroll wheel to zoom, and the right mouse button to move.
                        </div>
                        """
                    )    
                output_file = gr.File(label="ply")
            with gr.Column(scale=1):
                output_video = gr.Video(label="video")
                
    button_gen.click(process, inputs=[inputfiles], outputs=[output_video, output_file, output_model])
    
    gr.Examples(
        examples=[
            "davis-dog",
            # "sintel-market_2",
        ],
        inputs=[input_path],
        outputs=[output_video, output_file, output_model],
        fn=lambda x: process(inputfiles=None, input_path=x),
        cache_examples=True,
        label='Examples'
    )
block.launch(server_name="0.0.0.0", share=False)