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.lastframe = None
        self.learning_rate = learning_rate
        self.heatmap = np.array([])
        self.backsub = cv2.createBackgroundSubtractorMOG2()
        self.teach_background()
        self.lastsum = self.to_floats(self.lastframe)

    @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(self, frame):
        if self.append_blur:
            frame = cv2.blur(frame,(5,5))
        self.backsub.apply(frame, self.lastframe, self.learning_rate)
        self.lastsum += self.to_floats(self.lastframe)
        self.heatmap = self.gray_to_heat(
            self.float_to_gray(
                self.to_floats(self.lastsum)
            )
        )

    def teach_background(self):
        for frame in self.frames:
            if self.append_blur:
                frame = cv2.blur(frame,(5,5))
            self.lastframe = self.backsub.apply(frame, None, self.learning_rate)