# import os
# import shutil
# import zipfile
# from os.path import join, isfile, basename
#
# import cv2
# import numpy as np
import gradio as gr
# 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))]

    zip_path = "asd.zip"
    # 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",
    fn=predict,
    inputs=[gr.inputs.Video(label="Upload Video File"),
            gr.inputs.Number(label="Downsample size")],
    outputs=gr.outputs.File(label="Zip"),
)

demo.launch(debug=True)