import numpy as np
import cv2
import time
from typing import List, Iterable
from video_loader import VideoLoader
import numpy
from data.library import ArkiteData
from cv2util import to_8_bit_image

class BackgroundHeatmap:
    def __init__(self, frames, append_blur=True, learning_rate=0.0003):
        self.append_blur = append_blur
        self.frames = frames
        self.max_history_len = 2
        self.lastframes = []
        self.learning_rate = learning_rate
        self.heatmap = np.array([])
        self.backsub = cv2.createBackgroundSubtractorMOG2()
        self.teach_background()
        self.reset_sum()

    @staticmethod
    def to_grayscale(frame):
        return cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)

    @staticmethod
    def to_floats(frame):
        return cv2.normalize(frame, None, alpha=0, beta=1, norm_type=cv2.NORM_MINMAX, dtype=cv2.CV_32F)

    @staticmethod
    def float_to_gray(frame):
        return cv2.normalize(frame, None, alpha=0, beta=255, norm_type=cv2.NORM_MINMAX, dtype=cv2.CV_8UC1)

    @staticmethod
    def gray_to_heat(frame):
        return cv2.applyColorMap(frame, cv2.COLORMAP_JET)

    def update_heatmap(self):
        self.heatmap = self.gray_to_heat(
            self.float_to_gray(
                self.to_floats(self.lastsum)
            )
        )

    def update(self, frame):
        if self.append_blur:
            frame = cv2.blur(frame,(5,5))
        self.add_frame_history(self.backsub.apply(frame, None, self.learning_rate))
        self.lastsum += self.to_floats(self.lastframe)
        self.update_heatmap()

    def teach_background(self, frames=None):
        if frames is None:
            frames = self.frames
        for frame in frames:
            if self.append_blur:
                frame = cv2.blur(frame,(5,5))
            # TODO: is it logical here to use the same learning rate?
            self.add_frame_history(self.backsub.apply(frame, None, self.learning_rate))

    def add_frame_history(self, frame):
        if len(self.lastframes) == self.max_history_len:
            self.lastframes.pop(0)
        self.lastframes.append(frame)

    def reset_sum(self):
        self.lastsum = self.to_floats(self.lastframe)

    def reset(self):
        self.reset_sum()
        self.update_heatmap()

    @property
    def lastframe(self):
        return self.lastframes[-1]

    @property
    def bgf_diff(self):
        # make binary map
        first = cv2.threshold(self.lastframes[1], 60, 255, cv2.THRESH_BINARY)[1]
        second = cv2.threshold(self.lastframes[0], 60, 255, cv2.THRESH_BINARY)[1]
        return cv2.subtract(first, second)

class BackgroundHeatmapFromGroup(BackgroundHeatmap):
    def __init__(self, group, append_blur=True, learning_rate=0.0003):
        self.group = group
        assert len(self.group[1]) is not 0
        # Init with first recording
        frames = VideoLoader.extract_frames(self.group[1][0])
        super().__init__(frames, append_blur=append_blur, learning_rate=learning_rate)
        self.teachRestOfVideosInGroup()
    
    def teachRestOfVideosInGroup(self):
        # Slice away first recording since it was passed to super constructor already
        for recording in self.group[1][1:]:
            frames = VideoLoader.extract_frames(recording)
            self.teach_background(frames=frames)