import sys
sys.path.append('src/') # Add src to path

import gradio as gr

from utils import get_array, save_video, make_predictions
from data.processing import Scaler

examples = [["data/2023-05-05.mp4", '2023-05-05', '12:00:00', '21:30:00'],
            ["data/2023-09-07.mp4", '2023-09-07', '07:00:00', '16:30:00'],
            ["data/2023-10-28.mp4", '2023-10-28', '12:00:00', '21:30:00'],
            ["data/2023-11-09.mp4", '2023-11-09', '01:30:00', '11:00:00'],
            ["data/2023-11-13.mp4", '2023-11-15', '00:00:00', '09:30:00'],
            ["data/2023-11-18.mp4", '2023-11-18', '00:00:00', '09:30:00']]

scaler_dict = {"2023-05-05.mp4": "models/scaler_1.pkl",
               "2023-09-07.mp4": "models/scaler_2.pkl",
               "2023-10-28.mp4": "models/scaler_3.pkl",
               "2023-11-09.mp4": "models/scaler_4.pkl",
               "2023-11-13.mp4": "models/scaler_5.pkl",
               "2023-11-18.mp4": "models/scaler_6.pkl"}

def predict(source, date, start, end):

    print(f"Date: {date}")
    print(f"Start: {start}")
    print(f"End: {end}")

    X, y = get_array(source, scaler_dict)
    print(X.shape, y.shape)

    file_path_input = 'data/input.mp4'
    save_video(X, threshold=0, file_path=file_path_input)
    
    file_path_output = 'data/output.mp4'
    save_video(y, threshold=0, file_path=file_path_output)

    ypred = make_predictions(X)

    file_path_predict = 'data/predict.mp4'
    save_video(ypred, threshold=0.1, file_path=file_path_predict)

    return [file_path_input, file_path_output, file_path_predict]

demo = gr.Interface(
    fn=predict,
    inputs=[
        gr.Video(label="Source", visible=False),
        gr.Text(label='Date'),
        gr.Text(label='Start'),
        gr.Text(label='End')
        ],
    outputs=[
        gr.Video(label='Input'),
        gr.Video(label='Output'),
        gr.Video(label='Prediction')
    ],
    examples=examples
)

if __name__ == "__main__":
    demo.launch()