#from pyqtgraph.Qt import QtCore
from threading import Thread
from filevideostream import FileVideoStream
import numpy as np
import cv2
import time
import copy
from pyqtgraph.Qt import QtCore

DEBUG = True

class DepthSlice:

    def __init__(self):
        self.points = None
        self.nom_depth = -1
        self.label = None

    def add_points(self,points):
        if self.points is None:
            self.points = points
        else:
            self.points = np.concatenate((self.points, points))


class FrameProcessor:
    def __init__(self, low=240, high=250):
        self.crop = False
        self.low = low
        self.high = high

    def proc(self, frame):
        edges = cv2.Canny(frame, self.low, self.high)
        kernel = np.ones((3, 3), np.uint8)
        edges_filtered = cv2.morphologyEx(edges, cv2.MORPH_CLOSE, kernel)
        contours, hierarchy = cv2.findContours(edges_filtered, cv2.RETR_LIST, cv2.CHAIN_APPROX_NONE)
        # find the contour with the largest area
        max_area_index = 0
        max_area = -1
        cont_image = frame.copy()
        ds = DepthSlice()
        for ind, cont in enumerate(contours):
            if cv2.arcLength(contours[ind],False) > 40:
                cv2.drawContours(cont_image, contours, ind, (0, 255, 0), 2, cv2.LINE_8, hierarchy, 0)
        ds.add_points(np.vstack(contours).squeeze())
        return cont_image, ds

class ThreadedProcessor(Thread):

    def __init__(self, frames, processor=FrameProcessor()):
        Thread.__init__(self)
        self.frames = frames
        self.processed_frames = copy.deepcopy(frames)
        self.finished = False
        self.three_dr_data = []
        self.processor = processor
        self.processed_counter = 0

    def apply_processor(self, frame):
        return self.processor.proc(frame)

    def run(self):
        self.finished = False
        self.processed_counter = 0
        for i, frame in enumerate(self.frames):
            self.processed_frames[i], ds = self.apply_processor(frame)
            self.three_dr_data.append(ds)
            self.processed_counter += 1
        self.finished = True

class VideoManager(QtCore.QObject):

    frames_loaded = QtCore.Signal(int)
    frames_processed = QtCore.Signal(int)

    def __init__(self, video_path, transform=None, ncpus = 16):
        QtCore.QObject.__init__(self)

        self.video_path = video_path
        self.is_loaded = False
        self.current_frame = 0
        self.frames_in_file = -1
        self.video = None
        self.frames = []
        self.frames_chunks = []
        self.processors = []
        self.video_load_time = 0.0
        self.transform = transform
        self.ncpus = ncpus

    def chunk_video(self):

        total_frames = len(self.frames)
        chunk_size = int(total_frames / self.ncpus)
        idx = 0
        for i in range(self.ncpus):
            if (idx + chunk_size) <= total_frames:
                self.frames_chunks.append(self.frames[idx:(idx + chunk_size)])
            else:
                self.frames_chunks.append(self.frames[idx:])
            idx += chunk_size

    def setup_processors(self):
        self.processors = []
        for i in range(self.ncpus):
            self.processors.append(ThreadedProcessor(self.frames_chunks[i]))

    def start_processors(self):
        for i in range(self.ncpus):
            self.processors[i].start()

    def stop_processors(self):
        for i in range(self.ncpus):
            self.processors[i].stop()

    def done(self):
        nproc = 0
        is_done = True
        for i in range(self.ncpus):
            if not self.processors[i].finished:
                is_done = False
            nproc += self.processors[i].processed_counter
        self.frames_processed.emit(nproc)
        return is_done

    def print_load_stats(self):
        print('Loading Time: ' + str(self.video_load_time))
        print('Loading FPS: ' + str(self.frames_in_file / self.video_load_time))

    def do_process(self):

        self.setup_processors()
        self.start_processors()
        done = False
        tstart = time.time()
        while not done:
            done = self.done()
            time.sleep(0.1)
        if DEBUG:
            print('Proc time: ' + str(time.time() - tstart))

    def get_proc_frames(self):
        output = []
        for i in range(self.ncpus):
            output += self.processors[i].processed_frames
        return output

    def get_3dr_data(self):
        output = []
        for i in range(self.ncpus):
            output += self.processors[i].three_dr_data
        return output


    def load_FVS(self):

        if DEBUG:
            print('Loading video using FVS...')

        self.frames = []
        self.frames_in_file = 0
        self.video = FileVideoStream(self.video_path,
                    transform=self.transform, queue_size=4096).start()
        self.video_load_time = time.time()
        while self.video.more():
            f = self.video.read()
            if f is not None:
                self.frames_loaded.emit(self.frames_in_file)
                self.frames.append(f)
                self.frames_in_file += 1
        self.video.stop()
        self.video_load_time = time.time() - self.video_load_time
        self.print_load_stats()

    def load_CV2(self):

        if DEBUG:
            print('Loading video using CV2...')

        self.frames = []
        self.frames_in_file = 0
        self.video = cv2.VideoCapture(self.video_path)
        self.video_load_time = time.time()
        while True:
            ok, frame = self.video.read()
            if ok:
                if self.transform is not None:
                    frame = self.transform(frame)
                self.frames.append(frame)
                self.frames_in_file += 1
            else:
                break
        self.video.release()
        self.video_load_time = time.time() - self.video_load_time
        self.print_load_stats()

if __name__=="__main__":

    #video_path = "C:\\Users\\proberts\\Videos\\vlc-record-2019-09-30-10h56m40s-DEEPPIV_S001_S102_T001.MOV-.avi"
    video_path = "C:\\Users\\proberts\\Desktop\\151120_SC1ATK62_Batho_10.mov"
    mgr = VideoManager(video_path, ncpus=32)
    mgr.load_FVS()
    mgr.chunk_video()
    for i in range(1):
        mgr.do_process()
        proc_frames = mgr.get_proc_frames()
        for f in proc_frames:
            cv2.imshow('Frames',f)
            cv2.waitKey(33)
        time.sleep(0.5)