import cv2
import numpy as np
import json


rotation_map = {
    0: None,
    90: cv2.ROTATE_90_COUNTERCLOCKWISE,
    180: cv2.ROTATE_180,
    270: cv2.ROTATE_90_CLOCKWISE
}

class NumpyEncoder(json.JSONEncoder):
    def default(self, obj):
        if isinstance(obj, np.ndarray):
            return obj.tolist()
        return json.JSONEncoder.default(self, obj)

def draw_bboxes(image, bounding_boxes, boxes_id, scores):
    image_with_boxes = image.copy()

    for bbox, bbox_id, score in zip(bounding_boxes, boxes_id, scores):
        x1, y1, x2, y2 = bbox
        cv2.rectangle(image_with_boxes, (x1, y1), (x2, y2), (128, 128, 0), 2)

        label = f'#{bbox_id}: {score:.2f}'

        (label_width, label_height), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
        label_x = x1
        label_y = y1 - 5 if y1 > 20 else y1 + 20

        # Draw a filled rectangle as the background for the label
        cv2.rectangle(image_with_boxes, (x1, label_y - label_height - 5),
                      (x1 + label_width, label_y + 5), (128, 128, 0), cv2.FILLED)
        cv2.putText(image_with_boxes, label, (label_x, label_y),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1)

    return image_with_boxes


def pad_image(image: np.ndarray, aspect_ratio: float) -> np.ndarray:
    # Get the current aspect ratio of the image
    image_height, image_width = image.shape[:2]
    current_aspect_ratio = image_width / image_height

    left_pad = 0
    top_pad = 0
    # Determine whether to pad horizontally or vertically
    if current_aspect_ratio < aspect_ratio:
        # Pad horizontally
        target_width = int(aspect_ratio * image_height)
        pad_width = target_width - image_width
        left_pad = pad_width // 2
        right_pad = pad_width - left_pad

        padded_image = np.pad(image,
                              pad_width=((0, 0), (left_pad, right_pad), (0, 0)),
                              mode='constant')
    else:
        # Pad vertically
        target_height = int(image_width / aspect_ratio)
        pad_height = target_height - image_height
        top_pad = pad_height // 2
        bottom_pad = pad_height - top_pad

        padded_image = np.pad(image,
                              pad_width=((top_pad, bottom_pad), (0, 0), (0, 0)),
                              mode='constant')
    return padded_image, (left_pad, top_pad)


class VideoReader(object):
    def __init__(self, file_name, rotate=0):
        self.file_name = file_name
        self.rotate = rotation_map[rotate]
        try:  # OpenCV needs int to read from webcam
            self.file_name = int(file_name)
        except ValueError:
            pass

    def __iter__(self):
        self.cap = cv2.VideoCapture(self.file_name)
        if not self.cap.isOpened():
            raise IOError('Video {} cannot be opened'.format(self.file_name))
        return self

    def __next__(self):
        was_read, img = self.cap.read()
        if not was_read:
            raise StopIteration
        if self.rotate is not None:
            img = cv2.rotate(img, self.rotate)
        return cv2.cvtColor(img, cv2.COLOR_BGR2RGB)