import numpy as np
from typing import List
import cv2

class BackgroundHeatmap:
    def __init__(self, capture):
        self.heatmap = np.array([])
        self.cap = capture
        self.backsub = cv2.createBackgroundSubtractorMOG2()
        # self.backsub = cv2.bgsegm_BackgroundSubtractorGSOC()
        # self.backsub = cv2.back()

        # Fill up with first frame
        ret, frame = self.cap.read()
        self.lastframe = self.backsub.apply(frame)
        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):
        ret, frame = self.cap.read()
        self.lastframe = self.backsub.apply(frame)
        self.lastsum += self.to_floats(self.lastframe)
        self.heatmap = self.gray_to_heat(
            self.float_to_gray(
                self.to_floats(self.lastsum)
            )
        )

cap = cv2.VideoCapture(0)
diffsum = BackgroundHeatmap(cap)

# Load diffsum up with first
first = True
while(True):
    # Update heatmap
    diffsum.update()

    # Display the resulting frame
    cv2.imshow('Heatmap', diffsum.heatmap)
    cv2.imshow('Backsub', diffsum.lastframe)

    if first:
        cv2.moveWindow("Backsub", 1000, 100)
        first = False

    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

# When everything done, release the capture
diffsum.cap.release()
cv2.destroyAllWindows()