import cv2
import os
import numpy as np
import pyzjr.Z as Z
from PIL import Image
import torch

def cpu():
    return torch.device('cpu')

def gpu(i=0,cuda=True):
    """
    如果cuda为True并且gpu可用,则返回gpu(i),否则返回cpu().
    :param i: 索引i,表示使用第几块gpu
    :param cuda: 布尔值,表示是否使用cuda,没有可以使用cpu,设置为False
    """
    if cuda and torch.cuda.is_available() and torch.cuda.device_count() >= i + 1:
        return torch.device(f'cuda:{i}')
    return torch.device('cpu')

def allgpu():
    """
    Those who can use this function must be very rich.
    """
    devices = []
    for i in range(torch.cuda.device_count()):
        devices.append(i)
    return devices if devices else [torch.device('cpu')]

def num_workers(batch_size):
    """
    确定用于数据加载的并行工作进程的数量
    """
    num_workers = min([os.cpu_count(), batch_size if batch_size > 1 else 0, 8])
    return num_workers


def RemoveInboxes(boxes):
    """
    中心点的方法,去除重叠的框,并去掉相对较小的框
    :param boxes: [minc, minr, maxc, maxr]
    :return: 过滤的boxes
    """
    final_boxes = []
    for box in boxes:
        x1, y1, x2, y2 = box
        cx = (x1 + x2) / 2
        cy = (y1 + y2) / 2
        box_area = (x2 - x1) * (y2 - y1)
        is_inner = False
        for fb in final_boxes:
            fx1, fy1, fx2, fy2 = fb
            if fx1 <= cx <= fx2 and fy1 <= cy <= fy2:
                fb_area = (fx2 - fx1) * (fy2 - fy1)
                if box_area >= fb_area:
                    final_boxes.remove(fb)
                else:
                    is_inner = True
                    break
        if not is_inner:
            final_boxes.append(box)
    return final_boxes

def cvt2Center(boxes):
    """
    Convert from (upper left, lower right) to (center, width, height).
    :param boxes: [x1, y1, x2, y2]
    :return: [cx, cy, w, h]
    """
    converted_boxes = []
    for box in boxes:
        x1, y1, x2, y2 = box[0], box[1], box[2], box[3]
        cx = (x1 + x2) / 2
        cy = (y1 + y2) / 2
        w = x2 - x1
        h = y2 - y1
        converted_box = [cx, cy, w, h]
        converted_boxes.append(converted_box)
    return converted_boxes

def cvt2Corner(boxes):
    """Convert from (center, width, height) to (upper left, lower right).
    :param boxes: [cx, cy, w, h]
    :return: [x1, y1, x2, y2]
    """
    converted_boxes = []
    for box in boxes:
        cx, cy, w, h = box[0], box[1], box[2], box[3]
        x1 = cx - 0.5 * w
        y1 = cy - 0.5 * h
        x2 = cx + 0.5 * w
        y2 = cy + 0.5 * h
        converted_box = [x1, y1, x2, y2]
        converted_boxes.append(converted_box)
    return converted_boxes


def putBoxText(background_image, bbox, text, mode=1, rect=False, bboxcolor=(0, 255, 0), textcolor=(0, 0, 255),
               fontsize=1, thickness=2, bboxthickness=2,font=cv2.FONT_HERSHEY_SIMPLEX):
    """
    在给定的背景图像上绘制一个框，并在框的中心位置添加文本。获取文本的text_size，使文本居中。

    :param background_image: 背景图像，要在其上绘制框和文本的图像
    :param bbox: 框的边界框，表示为 [x1, y1, x2, y2]
    :param text: 要添加到框中心的文本
    :param mode: 模式0 表示框的中心位置，1 表示框的左上角
    :param rect: 是否绘制角标记框，导入的为 cornerRect 函数
    :param bboxcolor: 框的颜色，以 BGR 格式表示，例如 (0, 255, 0) 表示绿色
    :param textcolor: 文本的颜色，以 BGR 格式表示，例如 (0, 0, 255) 表示红色
    :param fontsize: 文本的字体大小
    :param thickness: 文本的线宽
    :param bboxthickness: 框的线宽，默认为 2
    :return: 绘制了框和文本的图像
    """
    text_x = text_y =None
    x1, y1, x2, y2 = bbox
    text_size, _ = cv2.getTextSize(text, font, fontsize, thickness)
    if rect:
        background_image = cornerRect(background_image, bbox)
    else:
        cv2.rectangle(background_image, (x1, y1), (x2, y2), bboxcolor, bboxthickness)
    if mode == 0:
        text_x = int((x1 + x2) / 2 - text_size[0] / 2)
        text_y = int((y1 + y2) / 2 + text_size[1] / 2)
    elif mode == 1:
        text_x = int(x1)
        text_y = int(y1 - text_size[1])
    cv2.putText(background_image, text, (text_x, text_y), font, fontsize, textcolor, thickness)


def cornerRect(drawimg, bbox, length=30, lthickness=5, rthickness=1, bboxcolor=Z.purple, colorCorner=Z.green):
    """
    在图像上绘制带有角标记的矩形框
    :param drawimg: 需要绘制的图像
    :param bbox: 边界框, 表示为 [x, y, x1, y1]
    :param length: 角标记的长度
    :param lthickness: 角标记的线条宽度
    :param rthickness: 矩形框的线条宽度
    :param bboxcolor: 矩形框的颜色
    :param colorCorner: 角标记的颜色
    :return: 绘制了角标记的图像副本
    """
    x, y, x1, y1 = bbox
    cv2.rectangle(drawimg, (x, y), (x1, y1), bboxcolor, rthickness)
    # 左上角  (x,y)
    cv2.line(drawimg, (x, y), (x + length, y), colorCorner, lthickness)
    cv2.line(drawimg, (x, y), (x, y + length), colorCorner, lthickness)
    # 右上角  (x1,y)
    cv2.line(drawimg, (x1, y), (x1 - length, y), colorCorner, lthickness)
    cv2.line(drawimg, (x1, y), (x1, y + length), colorCorner, lthickness)
    # 左下角  (x,y1)
    cv2.line(drawimg, (x, y1), (x + length, y1), colorCorner, lthickness)
    cv2.line(drawimg, (x, y1), (x, y1 - length), colorCorner, lthickness)
    # 右下角  (x1,y1)
    cv2.line(drawimg, (x1, y1), (x1 - length, y1), colorCorner, lthickness)
    cv2.line(drawimg, (x1, y1), (x1, y1 - length), colorCorner, lthickness)

    return drawimg

def cvt8png(pngpath, bit_depth=False, target=Z.white, convert=(128, 0, 0)):
    """
    Voc: RGB png彩图转换
    :param bit_depth: 默认转为8位,需要使用out_png.save可以正确保存
    True:                      False:
        plt.imshow(img)             cv2.imshow("img",img)
        plt.axis('off')             cv2.waitKey(0)
        plt.show()
    :param pngpath: 为保证是按照cv2的方式读入,所以传入路径即可
    :param target: 目标颜色,RGB方式
    :param convert: 转换颜色,同RGB格式
    :return:
    """
    png = cv2.imread(pngpath)
    png = cv2.cvtColor(png, cv2.COLOR_BGR2RGB)
    h, w = png.shape[:2]

    mask = np.all(png == target, axis=-1)
    out_png = np.zeros((h, w, 3), dtype=np.uint8)
    out_png[mask] = convert
    out_png[~mask] = png[~mask]

    if bit_depth:
        out_png_pil = Image.fromarray(out_png)
        out_png_pil = out_png_pil.convert("P", palette=Image.ADAPTIVE, colors=256)
        return out_png_pil
    else:
        out_png_cv = cv2.cvtColor(out_png, cv2.COLOR_RGB2BGR)
        return out_png_cv

def cvtMask(pngpath, lower ,upper, convert=(255, 255, 255)):
    """
    去除标签中的杂色,并对其进行颜色的转换

    :param pngpath: 输入图像的文件路径
    :param lower: 颜色范围的下限值，作为一个包含(B, G, R)值的元组
    :param upper: 颜色范围的上限值，作为一个包含(B, G, R)值的元组
    :param convert: 需要转换为的颜色，作为一个包含(B, G, R)值的元组
    :return: 处理后的图像
    """
    image = cv2.imread(pngpath)
    lower = np.array(lower, dtype=np.uint8)
    upper = np.array(upper, dtype=np.uint8)

    mask = cv2.inRange(image, lower, upper)
    image[mask > 0] = convert

    return image