import cv2
import logging
import numpy as np


def crop_imgs_box_xywh(imgs: list, box, same_sign: bool = False):
    """
    Crop imgs

    input:
        imgs: list, Each img in imgs has the same Width and High.
        box: The YOLO model outputs bounding box data in the format [x, y, w, h, confidence, class_id].
        same_sign: bool, Set True to skip size check if all img in imgs have the same Width and High.
    output:
        crop_imgs: list;
        (x_min, y_min);
    """
    if not imgs:
        logging.warning("imgs is empty")
        return [], (0, 0)

    if not same_sign and len(imgs) != 1:
        for img in imgs:
            if imgs[0].shape != img.shape:
                raise ValueError(f"Img shape are different: {imgs[0].shape} - {img.shape}")

    high, width = imgs[0].shape[:2]
    x_center, y_center, w, h = box[:4]
    x_min, x_max = max(0, int(round(x_center - w/2))), min(int(round(x_center + w/2)), width-1)
    y_min, y_max = max(0, int(round(y_center - h/2))), min(int(round(y_center + h/2)), high-1)

    crop_imgs = [img[y_min:y_max + 1, x_min:x_max + 1] for img in imgs]

    return crop_imgs, (x_min, y_min)


def crop_imgs_box_xyxy(imgs: list, box, same_sign: bool = False):
    """
    Crop imgs

    input:
        imgs: list, Each img in imgs has the same Width and High.
        box: The YOLO model outputs bounding box data in the format [x1, y1, x2, y2, confidence, class_id].
        same_sign: bool, Set True to skip size check if all img in imgs have the same Width and High.
    output:
        crop_imgs: list
        (x_min, y_min)
    """
    if not imgs:
        logging.warning("imgs is empty")
        return [], (0, 0)

    if not same_sign and len(imgs) != 1:
        for img in imgs:
            if imgs[0].shape != img.shape:
                raise ValueError(f"Img shape are different: {imgs[0].shape} - {img.shape}")

    high, width = imgs[0].shape[:2]
    x1, y1, x2, y2 = box[:4]
    x_min, x_max = max(0, int(round(x1))), min(int(round(x2)), width - 1)
    y_min, y_max = max(0, int(round(y1))), min(int(round(y2)), high - 1)

    crop_imgs = [img[y_min:y_max + 1, x_min:x_max + 1] for img in imgs]

    return crop_imgs, (x_min, y_min)


def crop_imgs_mask(imgs: list, mask: np.ndarray, same_sign: bool = False):
    """
    Crop imgs

    input:
        imgs: list, Each img in imgs has the same Width and High.
        mask: np.ndarray
        same_sign: bool, Set True to skip size check if all img in imgs have the same Width and High.
    output:
        crop_imgs: list
        (x_min, y_min)
    """
    if not imgs:
        logging.warning("imgs is empty")
        return [], (0, 0)

    if not same_sign and len(imgs) != 1:
        for img in imgs:
            if imgs[0].shape != img.shape:
                raise ValueError(f"Img shape are different: {imgs[0].shape} - {img.shape}")

    high, width = imgs[0].shape[:2]
    ys, xs = np.where(mask > 0)

    if xs.size == 0 or ys.size == 0:
        # 没有有效像素
        return None, None, (None, None)

    x_min, x_max = max(0, int(round(xs.min()))), min(int(round(xs.max())), width - 1)
    y_min, y_max = max(0, int(round(ys.min()))), min(int(round(ys.max())), high - 1)

    crop_imgs = [img[y_min:y_max + 1, x_min:x_max + 1] for img in imgs]
    mask_crop = mask[y_min:y_max + 1, x_min:x_max + 1]

    return crop_imgs, mask_crop, (x_min, y_min)


def get_map(K: list, D: list, camera_size: list):
    """
    input:
        K: list, shape (9)
        D: list
        camera_size: list, [w, h]
    output:
        map1: np.ndarray
        map2: np.ndarray
        new_K: list, shape (9)
    """
    h, w = camera_size[::-1]
    K = np.array(K).reshape(3, 3)
    D = np.array(D)
    new_K, _ = cv2.getOptimalNewCameraMatrix(K, D, (w, h), 1, (w, h))
    map1, map2 = cv2.initUndistortRectifyMap(K, D, None, new_K, (w, h), cv2.CV_32FC1)

    return map1, map2, new_K.flatten()


def distortion_correction(
        color_image: np.ndarray,
        depth_image: np.ndarray,
        map1: np.ndarray,
        map2: np.ndarray
):
    """
    畸变矫正

    input:
        color_image: np.ndarray
        depth_image: np.ndarray
        map1: np.ndarray
        map2: np.ndarray
    output:
        undistorted_color: np.ndarray
        undistorted_depth: np.ndarray
    """
    undistorted_color = cv2.remap(color_image, map1, map2, cv2.INTER_LINEAR)
    undistorted_color = undistorted_color.astype(color_image.dtype)

    undistorted_depth = cv2.remap(depth_image, map1, map2, cv2.INTER_NEAREST)
    undistorted_depth = undistorted_depth.astype(depth_image.dtype)

    return undistorted_color, undistorted_depth


