import os
import shutil
import zipfile
from os.path import join, isfile, basename

import cv2
import numpy as np
import gradio as gr
from gradio.components import Video, Number, File
import torch

from resnet50 import resnet18
from sampling_util import furthest_neighbours
from video_reader import video_reader

model = resnet18(
    output_dim=0,
    nmb_prototypes=0,
    eval_mode=True,
    hidden_mlp=0,
    normalize=False)
model.load_state_dict(torch.load("model.pt"))
model.eval()
avg_pool = torch.nn.AdaptiveAvgPool2d((1, 1))


def predict(input_file, downsample_size):
    downsample_size = int(downsample_size)
    base_directory = os.getcwd()
    selected_directory = os.path.join(base_directory, "selected_images")
    if os.path.isdir(selected_directory):
        shutil.rmtree(selected_directory)
    os.mkdir(selected_directory)

    file_name = (input_file.split('/')[-1]).split('.')[-1]
    zip_path = os.path.join(selected_directory, file_name + ".zip")

    mean = np.asarray([0.3156024, 0.33569682, 0.34337464], dtype=np.float32)
    std = np.asarray([0.16568947, 0.17827448, 0.18925823], dtype=np.float32)

    img_vecs = []
    with torch.no_grad():
        for fp_i, file_path in enumerate([input_file]):
            for i, in_img in enumerate(video_reader(file_path,
                                                    targetFPS=9,
                                                    targetWidth=100,
                                                    to_rgb=True)):
                in_img = (in_img.astype(np.float32) / 255.)
                in_img = (in_img - mean) / std
                in_img = np.expand_dims(in_img, 0)
                in_img = np.transpose(in_img, (0, 3, 1, 2))
                in_img = torch.from_numpy(in_img).float()
                encoded = avg_pool(model(in_img))[0, :, 0, 0].cpu().numpy()
                img_vecs += [encoded]
    img_vecs = np.asarray(img_vecs)
    print("images encoded")
    rv_indices, _ = furthest_neighbours(
        x=img_vecs,
        downsample_size=downsample_size,
        seed=0)
    indices = np.zeros((img_vecs.shape[0],))
    indices[np.asarray(rv_indices)] = 1
    print("images selected")

    global_ctr = 0
    for fp_i, file_path in enumerate([input_file]):
        for i, img in enumerate(video_reader(file_path,
                                             targetFPS=9,
                                             targetWidth=None,
                                             to_rgb=False)):
            if indices[global_ctr] == 1:
                cv2.imwrite(join(selected_directory, str(global_ctr) + ".jpg"), img)
            global_ctr += 1
    print("selected images extracted")

    all_selected_imgs_path = [join(selected_directory, f) for f in os.listdir(selected_directory) if
                              isfile(join(selected_directory, f))]

    zipf = zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED)
    for i, f in enumerate(all_selected_imgs_path):
        zipf.write(f, basename(f))
    zipf.close()
    print("selected images zipped")

    return zip_path


demo = gr.Interface(
    enable_queue=True,
    title="Frame selection by visual difference",
    description="",
    fn=predict,
    inputs=[Video(label="Upload Video File"),
            Number(label="Downsample size")],
    outputs=File(label="Zip"),
)

demo.launch(enable_queue=True)