import numpy as np
import os
import gradio as gr
import numpy as np
import tempfile
import shutil
import hashlib
import subprocess
import time

TMP_DIR = os.path.join(tempfile.gettempdir(), 'gradio')
os.makedirs(TMP_DIR, exist_ok=True)

gs_train_iter = 1000
pose_lr = '1x'


def move_uploaded_files(files, temp_dir):
    saved_files = []
    for file in files:
        target_file_path = os.path.join(temp_dir, os.path.basename(file.name))
        saved_files.append(target_file_path)
        if os.path.exists(file.name):
            shutil.move(file.name, target_file_path)
    return saved_files


def process_images(files):
    n_view = len(files)
    
    files_str = ','.join([f.name for f in files])
    hashdir_name = hashlib.md5(files_str.encode('utf8')).hexdigest()
    hashdir_path = os.path.join(TMP_DIR, hashdir_name)
    hashdir_path_img = os.path.join(hashdir_path, 'images')
    os.makedirs(hashdir_path_img, exist_ok=True)
    saved_files = move_uploaded_files(files, hashdir_path_img)
    saved_files_str = ' '.join(saved_files)

    stdout, stderr, = [], []
    
    # ----- (1) Dust3r_coarse_geometric_initialization -----
    command = [
        "python", "-W", "ignore",
        "./coarse_init_infer.py",
        "--img_base_path", hashdir_path,
        "--n_views", str(n_view),
        "--focal_avg",
    ]
    start_time = time.time()
    result = subprocess.run(command, capture_output=True, text=True)
    end_time = time.time()
    stdout.append(f"Time taken for (1) Dust3r_coarse_geometric_initialization: {end_time-start_time} sec.")
    stdout.append(result.stdout)
    stderr.append(result.stderr)
    # ----- (2) Train: jointly optimize pose -----
    scene = 'scene'
    subdir = f'{n_view}_views_{gs_train_iter}Iter_{pose_lr}PoseLR'
    model_path = os.path.join(hashdir_path, 'infer', scene, subdir)
    
    # python -W ignore ./train_joint.py \
    #         -s ${SOURCE_PATH} \
    #         -m ${MODEL_PATH}  \
    #         --n_views ${N_VIEW}  \
    #         --scene ${SCENE} \
    #         --iter ${gs_train_iter} \
    #         --optim_pose \
    
    command = [
        "python", "-W", "ignore",
        "./train_joint.py",
        "-s", hashdir_path,
        "-m", model_path,
        "--n_views", str(n_view),
        "--scene", scene,
        "--iter", str(gs_train_iter),
        "--optim_pose",
    ]

    start_time = time.time()
    result = subprocess.run(command, capture_output=True, text=True)
    end_time = time.time()
    stdout.append(f"Time taken for (2) Train: jointly optimize pose: {end_time-start_time} sec.")
    stdout.append(result.stdout)
    stderr.append(result.stderr)

    # ----- (3) Render interpolated pose & output video -----
    # python -W ignore ./render_by_interp.py \
    #         -s ${SOURCE_PATH} \
    #         -m ${MODEL_PATH}  \
    #         --n_views ${N_VIEW}  \
    #         --scene ${SCENE} \
    #         --iter ${gs_train_iter} \
    #         --eval \
    #         --get_video
    
    command = [
        "python", "-W", "ignore",
        "./render_by_interp.py",
        "-s", hashdir_path,
        "-m", model_path,
        "--n_views", str(n_view),
        "--scene", scene,
        "--iter", str(gs_train_iter),
        "--eval",
        "--get_video",
    ]
    
    start_time = time.time()
    result = subprocess.run(command, capture_output=True, text=True)
    end_time = time.time()
    stdout.append(f"Time taken for (3) Render interpolated pose & output video: {end_time-start_time} sec.")
    stdout.append(result.stdout)
    stderr.append(result.stderr)

    stdout = '\n'.join(stdout)
    stderr = '\n'.join(stderr)
    
    result_tree = subprocess.run(["tree", hashdir_path], capture_output=True, text=True)
    ply_file = os.path.join(model_path, 'point_cloud', f'iteration_{gs_train_iter}', 'point_cloud.ply')
    mp4_file = os.path.join(model_path, f'{scene}_{n_view}_view.mp4')
    
    return mp4_file, ply_file, result_tree.stdout, stdout, stderr

def run():
    with gr.Blocks() as demo:
        gr.Markdown("Image and Video Processor")
        # Gradio interface
        inputs = gr.File(file_count="multiple", file_types=["image"], label="Upload Images")
        btn = gr.Button(value="Run")
        # outputs = [gr.Image(type="numpy", label="Processed Images", interactive=False), gr.Video(label="Output Video")]    
        # outputs = [gr.Textbox(label="Log"), gr.Image(type="filepath", label="Processed Images", interactive=False)]    
        outputs = [
            gr.Video(label="Output Video"),
            gr.Model3D(label="3D Model"),
            gr.Textbox(label="Output dirs"), 
            gr.Textbox(label="Log (stdout)"), 
            gr.Textbox(label="Log (stderr)"),
        ]    
        btn.click(fn=process_images, inputs=inputs, outputs=outputs)

    demo.launch(server_name="0.0.0.0", server_port=7860)


if __name__ == "__main__":
    run()