from typing import Callable

import cv2
import numpy as np
from PIL import Image
from tqdm import tqdm, trange


def compare(frame1, frame2):
    absdiff = cv2.absdiff(frame1, frame2)
    return np.mean(absdiff)


def load_frames(cap: cv2.VideoCapture, frame_interval: int, progress: Callable) -> list:
    frames = []
    num_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))

    i = 0
    res = cap.grab()
    while i < num_frames:
        if res:
            frames.append(cap.retrieve()[1])
        for _ in range(frame_interval):
            res = cap.grab()
        i += frame_interval
        progress(i)

    return frames


def differences(frames: list, progress: Callable) -> list[float]:
    diffs = []

    for index, (frame1, frame2) in enumerate(zip(frames[:-1], frames[1:]), start=1):
        diff = compare(frame1, frame2)
        diffs.append(diff)
        progress(index)

    return diffs


def select_frames(frames: list, diffs: list, threshold: float) -> list[Image.Image]:
    selected_frames = [frames[0]]

    for frame, diff in zip(frames[1:], diffs, strict=True):
        if diff > threshold:
            selected_frames.append(frame)

    # Convert to Images
    return [Image.fromarray(frame[:, :, ::-1]) for frame in selected_frames]