from enum import Enum
from data.library import ArkiteData, ArkiteDetection, ArkiteRecording
from cv2util import to_8_bit_image
from typing import Iterable
import numpy
from recording_grouper import RecordingGrouper

class VideoType(Enum):
    IR = 0
    DEPTH = 1

class VideoLoader:
    def __init__(self, projects_path: str):
        self.projects_path = projects_path
        # Search folders with depth of 1:
        self.data = ArkiteData(self.projects_path, 1)

    def get_all_recordings(self) -> Iterable[ArkiteRecording]:
        # Loop over projects
        for detection in self.data.detections(self.data.use_cases()):
            for recording in detection.recordings():
                yield recording

    def get_recordings_grouped(self):
        grouper = RecordingGrouper(self.data)
        return grouper.getRecordingGroups()
    
    @staticmethod
    def extract_frames(recording: ArkiteRecording, video_type: VideoType = None) -> Iterable[numpy.ndarray]:
        if video_type == VideoType.IR or video_type == None:
            disp_min = 100
            disp_max = 8000
            with recording.ir() as ir:
                frames = ir.frame_sequence()
                for frame in frames:
                    yield to_8_bit_image(frame, display_min=disp_min, display_max=disp_max)
        elif video_type == VideoType.DEPTH:
            # TODO: check if these values make sense for depth images.
            disp_min = 100
            disp_max = 8000
            with recording.depth() as depth:
                frames = depth.frame_sequence()
                for frame in frames:
                    yield to_8_bit_image(frame, display_min=disp_min, display_max=disp_max)
        else:
            raise NotImplementedError("VideoType \'{}\' not implemented".format(video_type))

    # Yielding in sub method after "with" statement in top method does not work!
    # @staticmethod
    # def frames_to_8_bit(frame_sequence, disp_min=None, disp_max=None):
    #     for frame in frame_sequence:
    #         yield to_8_bit_image(frame, display_min=disp_min, display_max=disp_max)