import torch
import sys
import numpy as np
import tempfile

try:
    torch.ones(1, 1).cuda().float()
    sys.path.insert(0,"/home/user/app/wrapper")
    sys.path.append("/home/user/app/Dynamic3DGaussians")
    import Dynamic3DGaussians.visualize
    Dynamic3DGaussians.visualize.FORCE_LOOP = True
    from Dynamic3DGaussians.visualize import visualize
    import open3d as o3d
    IS_CPU = False
except Exception as e:
    print(f"Error: {e}")
    IS_CPU = True

import gradio as gr
import cv2
from threading import Thread
import os

def video_thread(self):
    if self.counter != -1:
        file_path_video, file_path_pcd, file_path_png = self.callback_thread_params
        if self.videoWriter is None:
            self.counter = 200
            fourcc = cv2.VideoWriter_fourcc(*'mp4v')
            self.videoSize = (640, 360)
            self.videoWriter = cv2.VideoWriter(file_path_video, fourcc, 10, self.videoSize)
            pcd = self.lastPcd
            if not o3d.io.write_point_cloud(file_path_pcd, pcd):
                result = "ERROR #2"
        visualizer = self.delegate
        visualizer.capture_screen_image(file_path_png,do_render=True)
        img = cv2.imread(file_path_png)
        img = cv2.resize(img, self.videoSize)
        self.videoWriter.write(img)
        self.counter -= 1
        if self.counter == 0:
            self.counter = -1
            self.videoWriter.release()
            print("videoWriter released")

def do_visualize(what, file_path_video, file_path_pcd, file_path_png):
    o3d.callback_thread = video_thread
    o3d.callback_thread_params = (file_path_video, file_path_pcd, file_path_png)
    visualize(what, "pretrained")

def run_simulation(what):
    temp_name = next(tempfile._get_candidate_names())
    file_path_video = f"/tmp/{temp_name}.mp4"
    file_path_pcd = f"/tmp/{temp_name}.pcd"
    file_path_png = f"/tmp/{temp_name}.png"
    print("start thread")
    do_visualize(what, file_path_video, file_path_pcd, file_path_png)
    print("done")
    result = "Done"
    #TODO return path to FIFO-video file
    return file_path_video, file_path_pcd, file_path_png, result

def do_nothing():
    return "Not possible on CPU"

iface = None
if IS_CPU:
    iface = gr.Interface(fn=do_nothing, inputs=None, outputs=[gr.Textbox(label="Error", value="Not possible on CPU")], allow_flagging="never")
else:
    iface = gr.Interface(fn=run_simulation, inputs=gr.Dropdown(["basketball", "boxes", "football", "juggle", "softball", "tennis"], value="basketball"), outputs=[gr.Video(format="mp4"), gr.File(), gr.Image(), gr.Textbox(label="debug output")], allow_flagging="never")
iface.launch(server_name="0.0.0.0")