import os
import torch
import subprocess
import spaces
import gradio as gr
from PIL import Image
import numpy as np
from moviepy.editor import VideoFileClip

os.environ['CUDA_HOME'] = "/usr/local/lib/python3.10/site-packages/nvidia/cuda_runtime/include"

os.environ['PATH'] = '/usr/local/lib/python3.10/site-packages/nvidia/cuda_nvrtc/lib:' + os.environ.get('PATH', '')
os.environ['LD_LIBRARY_PATH'] = '/usr/local/lib/python3.10/site-packages/nvidia/cuda_nvrtc/lib:' + os.environ.get('LD_LIBRARY_PATH', '')

os.environ['PATH'] = '/usr/local/lib/python3.10/site-packages/nvidia/cuda_runtime/include:' + os.environ.get('PATH', '')
os.environ['LD_LIBRARY_PATH'] = '/usr/local/lib/python3.10/site-packages/nvidia/cuda_runtime/include:' + os.environ.get('LD_LIBRARY_PATH', '')

print(os.environ['PATH'])
print(os.environ['LD_LIBRARY_PATH'])

subprocess.run(["git", "clone", "https://github.com/sniklaus/3d-ken-burns.git"])
subprocess.run(["mv", "./network-disparity.pytorch", "./3d-ken-burns/models/disparity-estimation.pytorch"])
subprocess.run(["mv", "./network-refinement.pytorch", "./3d-ken-burns/models/disparity-refinement.pytorch"])
subprocess.run(["mv", "./network-inpainting.pytorch", "./3d-ken-burns/models/pointcloud-inpainting.pytorch"])

subprocess.run(["cp", "./3d-ken-burns/autozoom.py", "./3d-ken-burns/autozoom-original.py"])

def update_parameters(fltShift=100, fltZoom=1.25, steps1=0.0, steps2=1.0, frames=240, fps=30):
  with open ('./3d-ken-burns/autozoom-original.py','r') as azo:
     script = azo.read()
     script = script.replace("'fltShift': 100.0",f"'fltShift': {fltShift}")
     script = script.replace("'fltZoom': 1.25",f"'fltZoom': {fltZoom}")
     script = script.replace("'fltSteps': numpy.linspace(0.0, 1.0, 75).tolist(),",f"'fltSteps': numpy.linspace({steps1}, {steps2}, {frames}).tolist(),")
     script = script.replace(", fps=25)",f", fps={fps})")
     with open ('./3d-ken-burns/autozoom.py','w') as az:
         az.write(script)
         az.truncate()
         az.seek(0)

@spaces.GPU
def generate_video(image):
    subprocess.run(["nvidia-smi"])
    subprocess.run(["ls","-al","/tmp"])
    image = Image.fromarray(image)
    image.save("/tmp/img.png")
    os.chdir("./3d-ken-burns")
    subprocess.run(["python", "autozoom.py", "--in", "/tmp/img.png", "--out", "/tmp/autozoom.mp4"]) 
    os.chdir("..")
    subprocess.run(["ls","-al","/tmp"])
    return "/tmp/autozoom.mp4"

def get_video(image,fltShift=100, fltZoom=1.25, steps1=0.0, steps2=1.0, frames=240, fps=30):
    update_parameters(fltShift, fltZoom, steps1, steps2, frames, fps)
    video_path = generate_video(image)
    clip = VideoFileClip("/tmp/autozoom.mp4")
    duration = clip.duration / 2
    clip = clip.subclip(0, duration)
    clip.write_videofile("/tmp/final.mp4")
    return "/tmp/final.mp4"

with gr.Blocks() as demo:
    gr.Markdown("""# 3D Zoom Video Generator!""")
    image = gr.Image()
    fltShift = gr.Slider(0, 1000, value=100, step=1, label="Flat Shift")
    fltZoom = gr.Slider(0.0, 5.0, value=1.25, step=0.1,label="Flat Zoom")
    steps1 = gr.Slider(0.0, 5.0, value=0.0, step=0.1, label="Steps1")
    steps2 = gr.Slider(0.0, 5.0, value=1.0, step=0.1, label="Steps2")
    frames = gr.Slider(50, 10000, value=240, step=1, label="Frames")
    fps = gr.Slider(10, 60, value=30, step=1, label="Frames per Second")
    submit_image = gr.Button(value="Generate Video")
    video = gr.Video()
    submit_image.click(fn=get_video,inputs=[image,fltShift,fltZoom,steps1,steps2,frames,fps],outputs=video)
    
demo.launch(show_error=True)