# Setting global vars 
CSV_PATH = ""
VID_PATH = ""
BATCH_SIZE = 5 # CHANGE THIS TO FIT YOUR GPU, YOU MAY NEED SOME TRIAL AND ERROR
CUDA_PATH = "/usr/local/cuda-11.8/" # CHANGE THIS TO YOUR LOCAL CUDA PATH

import subprocess
import os 
import sys
# set CUDA_HOME
os.environ["CUDA_HOME"] = CUDA_PATH
os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:5000'
subprocess.run(['pip', 'install', '-e', 'GroundingDINO'])
sys.path.append(os.path.join(os.getcwd(), "GroundingDINO"))
sys.path.append(os.path.join(os.getcwd(), "segment_anything"))

import gradio as gr
from dino_sam import sam_dino_vid



def run_sam_dino(input_vid,
                 grounding_caption,
                 box_threshold,
                 text_threshold,
                 fps_processed,
                 scaling_factor,
                 video_options):
    new_input_vid = input_vid.replace(" ", "_")
    os.rename(input_vid, new_input_vid)
    csv_path, vid_path = sam_dino_vid(vid_path=new_input_vid,
                                      text_prompt=grounding_caption,
                                      box_threshold=box_threshold,
                                      text_threshold=text_threshold,
                                      fps_processed=fps_processed,
                                      scaling_factor=scaling_factor,
                                      video_options=video_options,
                                      batch_size=BATCH_SIZE)
    global CSV_PATH
    CSV_PATH = csv_path
    global VID_PATH
    VID_PATH = vid_path
    return vid_path

def vid_download():
    """
    """
    print(CSV_PATH, VID_PATH)
    return [CSV_PATH, VID_PATH]


with gr.Blocks() as demo:
    gr.HTML(
        """
            <h1 align="center" style="font-size:xxx-large">🦍 Primate Detection</h1>
        """
    )
    
    with gr.Row():
        with gr.Column():
            input = gr.Video(label="Input Video", interactive=True)
            grounding_caption = gr.Textbox(label="What do you want to detect? (Muliple species should be separated by commas. i.e. 'baboons, chimpanzees')")
            with gr.Accordion("Advanced Options", open=False):
                box_threshold = gr.Slider(
                    label="Box Threshold",
                    info="Adjust the threshold to change the sensitivity of the model, lower thresholds being more sensitive.",
                    minimum=0.0,
                    maximum=1.0,
                    value=0.25,
                    step=0.01
                )
                text_threshold = gr.Slider(
                    label="Text Threshold", minimum=0.0, maximum=1.0, value=0.25, step=0.01
                )
                fps_processed = gr.Slider(
                    label="Frame Detection Rate",
                    info="Adjust the frame detection rate. I.e. a value of 120 will run detection every 120 frames, a value of 1 will run detection on every frame. Note: the lower the number the slower the processing time.",
                    minimum=1,
                    maximum=120,
                    value=30, 
                    step=1)
                scaling_factor = gr.Slider(
                    label="Downsample Factor",
                    info="Adjust the downsample factor. I.e. a value of 0.5 will downsample the pixels in the video by a factor of 2. Note: the higher the number the slower the processing time.",
                    minimum=0.1,
                    maximum=1.0,
                    value=0.5,
                    step=0.1
                )
                video_options = gr.CheckboxGroup(choices=["Bounding boxes", "Masks"],
                                                 label="Video Output Options",
                                                 info="Select the options to display in the output video. Note: if masks are selected, runtime will increase.",
                                                 value=["Bounding boxes"],
                                                 interactive=True)

            # TODO: Make button visible only after a file has been uploaded
            run_btn = gr.Button(value="Run Detection", visible=True)
        with gr.Column():
            vid = gr.Video(label="Output Video", height=350, interactive=False, visible=True)
            # download_btn = gr.Button(value="Generate Download", visible=True)
            download_file = gr.Files(label="CSV, Video Output", interactive=False)
    
    run_btn.click(fn=run_sam_dino, inputs=[input, grounding_caption, box_threshold, text_threshold, fps_processed, scaling_factor, video_options], outputs=[vid])
    vid.change(fn=vid_download, outputs=download_file)

    gr.Examples(
        [["baboon_15s.mp4", "baboon", 0.25, 0.25, 1, 1, ["Bounding boxes", "Masks"]]],
        inputs = [input, grounding_caption, box_threshold, text_threshold, fps_processed, scaling_factor, video_options],
        outputs = [vid],
        fn=run_sam_dino,
        cache_examples=True,
        label='Example'
      )
    
    gr.DuplicateButton()
    
    gr.Markdown(
        """
        ## Frequently Asked Questions 
        
        ##### How can I run the interface on my own computer? 
        By clicking on the three dots on the top right corner of the interface, you will be able to clone the repository or run it with a Docker image on your local machine. \
        For local machine setup instructions please check the README file. 

        ##### The video is very slow to process, how can I speed it up?
        You can speed up the processing by adjusting the frame detection rate in the advanced options. The lower the number the slower the processing time. Choosing only\
        bounding boxes will make the processing faster. You can also duplicate the space using the Duplicate Button and choose a different GPU which will make the processing faster.

        """
    )

demo.launch(share=False)