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

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

        # Fill up with first frame
        ret, frame = self.cap.read()
        # Suppose the first frame is background, teach it for a few iterations
        for i in range(bg_teach_iters):
            ret, frame = self.cap.read()
            self.backsub.apply(cv2.blur(frame,(5,5)))
        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(cv2.blur(frame,(5,5)))
        self.lastsum += self.to_floats(self.lastframe)
        self.heatmap = self.gray_to_heat(
            self.float_to_gray(
                self.to_floats(self.lastsum)
            )
        )

if __name__ == '__main__':
    # cap = cv2.VideoCapture(0)
    projects_path = "C:\\UntrackedGit\\opencv_test\\him_projects"
    data = ArkiteData(projects_path, 1)
    for uc in data.use_cases():
        for detection in data.detections(uc):
            for recording in detection.recordings():
                print("New recording: " + str(recording.name))
                with recording.ir() as ir:
                    for frame in ir.frame_sequence():
                        converted = to_8_bit_image(frame, display_min=100, display_max=8000)
                        # converted = cv2.applyColorMap(frame, colormap=cv2.COLORMAP_JET)
                        cv2.imshow("IR", converted)
                        if cv2.waitKey(1) & 0xFF == ord('q'):
                            break
                        # print("Showing frame...")
                        # time.sleep(0.1)

    cv2.destroyAllWindows()

# # 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()