import os
import cv2

from diffgro.utils import make_dir


def save_image(filename, frames):
    for (idx, frame) in enumerate(frames):
        frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
        path = os.path.join(filename, f'{idx}.png')
        cv2.imwrite(path, frame)


def save_video(filename, frames):
    h, w, c = frames[0].shape
    video_writer = cv2.VideoWriter(
        filename, cv2.VideoWriter_fourcc("m", "p", "4", "v"), 30, (w, h)
    )

    for frame in frames:
        frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
        video_writer.write(frame)

    video_writer.release()


def write_annotation(frame, annotation):
    font = cv2.FONT_HERSHEY_SIMPLEX
    bottomLeftCornerOfText = (10, 30)
    fontScale = 1
    fontColor = (0, 0, 255)
    lineType = 2

    cv2.putText(
        frame, annotation, bottomLeftCornerOfText, font, fontScale, fontColor, lineType
    )
    return frame
