# generators.py
# Vojtěch Golda
import os
import json
import cv2 as cv
import numpy as np

dataset_cache = {}
caching = False





def dataset_generator(source):
    
    if source not in dataset_cache:
        if caching:
            dataset_cache[source] = []
        for filename in os.listdir(source):
            if not filename.endswith(".json"):
                continue
             
            with open(os.path.join(source, filename), 'r') as file:
                #load data and image
                file = json.load(file)
                img_location = os.path.join(source, file["imagePath"])
                image = cv.imread(img_location)
                if image is None:
                     raise ValueError(f"{img_location} not found or corrupted. Terminating.")

                if "imageData" in file:
                    file.pop("imageData")
                if caching:
                    dataset_cache[source].append((image, file))
            yield image, file
    else:
        for file in dataset_cache[source]:
            yield file


def video_frame_generator(input_file):
    if type(input_file) is not int and not os.path.isfile(input_file):
        raise ValueError(f"{input_file} not found. Terminating.")
    video = cv.VideoCapture(input_file)
    while video.isOpened():
        ret, frame = video.read()
        if not ret:
            break
        yield frame, ""

def batch_generator(data, to_tensor, batchsize, none_terminator=False, nth=None):
    data_stack = []
    frame_stack = []
    frames = []
    for i, d in enumerate(data):
        if nth is None:
            data_stack.append(d)
        else:
            frames.append(d)
            if i % nth == 0:
                data_stack.append(d)
                frame_stack.append(frames)
                frames = []

        
        if len(data_stack) >= batchsize:
            if nth is None:
                yield to_tensor(data_stack)
            else:
                yield to_tensor(data_stack), frame_stack
            data_stack = []
            frame_stack = []


    if len(data_stack) > 0:
        if nth is None:
            yield to_tensor(data_stack)
        else:
            yield to_tensor(data_stack), frame_stack

    if none_terminator:
        yield None
