import os
import time
from pathlib import Path
import uuid
import tempfile
from typing import Union
import spaces
import atexit
from concurrent.futures import ThreadPoolExecutor

import gradio as gr
import cv2
import torch
import numpy as np

from moge.model import MoGeModel
from moge.utils.vis import colorize_depth
import utils3d

model = MoGeModel.from_pretrained('Ruicheng/moge-vitl').cuda().eval()
thread_pool_executor = ThreadPoolExecutor(max_workers=1)


def delete_later(path: Union[str, os.PathLike], delay: int = 300):
    def _delete():
        try: 
            os.remove(path) 
        except: 
            pass
    def _wait_and_delete():
        time.sleep(delay)
        _delete(path)
    thread_pool_executor.submit(_wait_and_delete)
    atexit.register(_delete)

@spaces.GPU
def run(image: np.ndarray, remove_edge: bool = True):
    run_id = str(uuid.uuid4())

    larger_size = max(image.shape[:2])
    if larger_size > 1024:
        scale = 1024 / larger_size
        image = cv2.resize(image, (0, 0), fx=scale, fy=scale, interpolation=cv2.INTER_AREA)

    image_tensor = torch.tensor(image, dtype=torch.float32, device=torch.device('cuda')).permute(2, 0, 1) / 255
    output = model.infer(image_tensor, resolution_level=9, apply_mask=True)
    points, depth, mask = output['points'].cpu().numpy(), output['depth'].cpu().numpy(), output['mask'].cpu().numpy()

    if remove_edge:
        mask = mask & ~utils3d.numpy.depth_edge(depth, mask=mask, rtol=0.02)
    mask = mask & (depth > 0)

    _, faces, indices = utils3d.numpy.image_mesh(width=image.shape[1], height=image.shape[0], mask=mask)
    faces = utils3d.numpy.triangulate(faces)

    tempdir = Path(tempfile.gettempdir(), 'moge')
    tempdir.mkdir(exist_ok=True)

    output_glb_path = Path(tempdir, f'{run_id}.glb')
    output_glb_path.parent.mkdir(exist_ok=True)
    tempfile.TemporaryFile()
    utils3d.io.write_glb(
        output_glb_path,
        vertices=points.reshape(-1, 3)[indices] * [-1, -1, 1], 
        faces=faces,
        vertex_colors=image.reshape(-1, 3)[indices] / 255,
    )

    output_ply_path = Path(tempdir, f'{run_id}.ply')
    output_ply_path.parent.mkdir(exist_ok=True)
    utils3d.io.write_ply(
        output_ply_path,
        vertices=points.reshape(-1, 3)[indices] * [-1, -1, 1], 
        faces=faces,
        vertex_colors=image.reshape(-1, 3)[indices] / 255,
    )

    colorized_depth = colorize_depth(depth)

    delete_later(output_glb_path, delay=300)
    delete_later(output_ply_path, delay=300)
        
    return colorized_depth, output_glb_path, output_ply_path.as_posix()


DESCRIPTION = """
MoGe turns 2D images into 3D point maps. 

NOTE: 
* If the image is too large (> 1024px), it will be resized accordingly.
* The color in the 3D viewer may look dark due to rendering of 3D viewer. You may download the 3D model as .glb or .ply file to view it in other 3D viewers.
"""

if __name__ == '__main__':
    
    gr.Interface(
        fn=run,
        inputs=[
            gr.Image(type="numpy", image_mode="RGB"),
            gr.Checkbox(True, label="Remove edges"),
        ],
        outputs=[
            gr.Image(type="numpy", label="Depth map (colorized)"),
            gr.Model3D(display_mode="solid", clear_color=[1.0, 1.0, 1.0, 1.0], label="3D Viewer"),
            gr.File(type="filepath", label="Download the model as .ply file"),
        ],
        title="MoGe Live Demo",
        description=DESCRIPTION,
        clear_btn=None,
        allow_flagging="never",
    ).launch(share=False)