import zipfile
import cv2
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
import os
from os.path import isfile, join


class ScenePreprocessor:
    """
    Preprocesses raw videos into frames so that further downstream extraction and
    preprocessing can occur
    Example usage:
    video_preprocessor = VideoPreprocessor(
        video_folder="drive/My Drive/cs231n-project/datasets/emotiw/train-tiny.zip",
        label_file="drive/My Drive/cs231n-project/datasets/emotiw/Train_labels.txt",
        output_folder="train-tiny-local",
        output_file="drive/My Drive/cs231n-project/datasets/emotiw/train-tiny-local.zip"
    )
    video_preprocessor.preprocess()
    """

    def __init__(self, video_folder, output_folder, output_file=None, label_file=None, height=320,
                 width=480, sample_every=10, max_workers=32):
        """
        @param video_folder   The folder where the list of videos are stored. If
                              this ends with .zip, this should be a single zip
                              file containing the videos. Paths can either by a local
                              folder or a GDrive mounted path.
        @param label_file     The file containing the space-delimited video name to label mapping. If None,
                              we assume that we are in 'test' mode and that there are no categories
        @param output_folder  The local output path where the preprocessed files will be stored for
                              further preprocessing can be done
        @param output_file    If not none, the output_folder will be zipped up and stored at this location
        @param height         Height of the extracted video frames
        @param width          Width of the extracted video frames
        @param sample_every   The frames to skip.
        @param max_workers    The number of workers to use to parallelize work.
        """
        self.video_folder = video_folder
        self.label_file = label_file
        self.output_folder = output_folder
        self.output_file = output_file
        print(
            f"Video Preprocessor created with video_folder = {video_folder} , label_file = {label_file} , output_folder = {output_folder}, output_file = {output_file}")

        self.height = height
        self.width = width
        self.sample_every = sample_every
        self.max_workers = max_workers
        print(f"Frames will be created with height = {height} , width = {width} , sample_every = {sample_every}")

    def preprocess(self):
        if self.video_folder.endswith(".zip"):
            # Unzips files to a temp directory
            tmp_output_folder = self.output_folder.rstrip('/') + "_tmp"
            print(f"Unzipping files to temp dir {tmp_output_folder}...")
            Path(f"{tmp_output_folder}").mkdir(parents=True, exist_ok=True)
            with zipfile.ZipFile(self.video_folder, 'r') as zip_ref:
                zip_ref.extractall(tmp_output_folder)
            print("Finished unzipping files")
        else:
            tmp_output_folder = self.video_folder
            print("Skipping unzipping files as input is a folder")

        if self.label_file is not None:
            # Create the category subfolders in the output folder
            # Path Structure:
            #   output/
            #     1/
            #     2/
            video_to_label = {}
            unique_labels = set()
            with open(self.label_file, "r") as f:
                i = 0
                for line in f:
                    if i == 0:
                        i += 1
                        continue
                    line_arr = line.split(" ")
                    video_to_label[line_arr[0] + ".mp4"] = line_arr[1].strip()
                    if line_arr[1].strip() not in unique_labels:
                        unique_labels.add(line_arr[1].strip())
                        Path(f"{self.output_folder}/{line_arr[1].strip()}").mkdir(parents=True, exist_ok=True)
                    i += 1
        else:
            video_to_label = None
            Path(f"{self.output_folder}").mkdir(parents=True, exist_ok=True)

        # Process each video by extracting the frames in a multi-threaded fashion
        futures = []
        with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
            videos = next(os.walk(tmp_output_folder))[2]
            videos = sorted(videos)

            print(f"Found {len(videos)} videos")
            video_num = 1

            for video_name in videos:
                future = executor.submit(self.process_video, tmp_output_folder, video_name, video_num, len(videos),
                                         video_to_label)
                futures.append(future)
                video_num += 1

        cv2.destroyAllWindows()

        print("***** Submitted all tasks *****")
        for future in futures:
            future.result()
        print("***** Completed *****")

        if self.output_file is not None:
            print(f"Starting to zip files to {self.output_file}")

            def zipdir(path, ziph):
                for root, dirs, files in os.walk(path):
                    folder = root[len(path):]
                    for file in files:
                        ziph.write(join(root, file), join(folder, file))

            zipf = zipfile.ZipFile(self.output_file, 'w', zipfile.ZIP_DEFLATED)
            zipdir(self.output_folder, zipf)
            zipf.close()
            print(f"Done zipping files to {self.output_file}")

        print("Done!")

    def process_video(self, tmp_output_folder, video_name, video_num, total_videos, video_to_label):
        """
        Processes a video by extracting the frames
        Writes out each frame, where each frame is an image resized to the the specified dimensions
        """
        vidcap = cv2.VideoCapture(join(tmp_output_folder, video_name))
        if video_to_label is not None:
            label = video_to_label[video_name]
            print(f"Processing video {video_num}/{total_videos} with name {video_name} and class {label} \n")
        else:
            label = None
            print(f"Processing video {video_num}/{total_videos} with name {video_name} (test-mode) \n")

        input_length = int(vidcap.get(cv2.CAP_PROP_FRAME_COUNT))
        frame_width = int(vidcap.get(cv2.CAP_PROP_FRAME_WIDTH))
        frame_height = int(vidcap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        fps = int(vidcap.get(cv2.CAP_PROP_FPS))

        success, image = vidcap.read()
        count = 0
        frame = 0
        while success:
            if count % self.sample_every == 0:
                height, width = image.shape[:2]
                image = cv2.resize(image, (self.width, self.height), interpolation=cv2.INTER_CUBIC)
                if video_to_label is not None:
                    cv2.imwrite(f"{self.output_folder}/{label}/frame_{video_name}_{frame}.jpg", image)
                else:
                    cv2.imwrite(f"{self.output_folder}/frame_{video_name}_{frame}.jpg", image)

                frame += 1
            success, image = vidcap.read()
            count += 1
        video_num += 1
        vidcap.release()
