from torch.utils.data.dataset import T_co

from libs.vision_libs.utils import draw_keypoints
from models.base.base_dataset import BaseDataset

import glob
import json
import math
import os
import random
import cv2
import PIL
import imageio
import matplotlib.pyplot as plt
import matplotlib as mpl
from torchvision.utils import draw_bounding_boxes

import numpy as np
import numpy.linalg as LA
import torch
from skimage import io
from torch.utils.data import Dataset
from torch.utils.data.dataloader import default_collate

import matplotlib.pyplot as plt
from models.dataset_tool import read_masks_from_txt_wire, read_masks_from_pixels_wire, adjacency_matrix

def validate_keypoints(keypoints, image_width, image_height):
    for kp in keypoints:
        x, y, v = kp
        if not (0 <= x < image_width and 0 <= y < image_height):
            raise ValueError(f"Key point ({x}, {y}) is out of bounds for image size ({image_width}, {image_height})")

"""
直接读取xanlabel标注的数据集json格式

"""
class LineDataset(BaseDataset):
    def __init__(self, dataset_path, data_type, transforms=None, dataset_type=None,img_type='rgb', target_type='pixel'):
        super().__init__(dataset_path)

        self.data_path = dataset_path
        self.data_type = data_type
        print(f'data_path:{dataset_path}')
        self.transforms = transforms
        self.img_path = os.path.join(dataset_path, "images/" + dataset_type)
        self.lbl_path = os.path.join(dataset_path, "labels/" + dataset_type)
        self.imgs = os.listdir(self.img_path)
        self.lbls = os.listdir(self.lbl_path)
        self.target_type = target_type
        self.img_type=img_type
        # self.default_transform = DefaultTransform()

    def __getitem__(self, index) -> T_co:
        img_path = os.path.join(self.img_path, self.imgs[index])

        lbl_path = os.path.join(self.lbl_path, self.imgs[index][:-3] + 'json')
        img = PIL.Image.open(img_path).convert('RGB')
        w, h = img.size
        # wire_labels, target = self.read_target(item=index, lbl_path=lbl_path, shape=(h, w))
        target = self.read_target(item=index, lbl_path=lbl_path, shape=(h, w))
        if self.transforms:
            img, target = self.transforms(img, target)

        else:
            img = self.default_transform(img)

        # print(f'img:{img}')
        # print(f'img shape:{img.shape}')
        return img, target

    def __len__(self):
        return len(self.imgs)

    def read_target(self, item, lbl_path, shape, extra=None):
        # print(f'shape:{shape}')
        # print(f'lbl_path:{lbl_path}')
        with open(lbl_path, 'r') as file:
            lable_all = json.load(file)


        objs = lable_all["shapes"]
        point_pairs=objs[0]['points']


        # print(f'point_pairs:{point_pairs}')
        target = {}

        target["image_id"] = torch.tensor(item)

        target["boxes"], lines = get_boxes_lines(objs,shape)
        # print(f'lines:{lines}')
        target["labels"] = torch.ones(len(target["boxes"]), dtype=torch.int64)


        a = torch.full((lines.shape[0],), 2).unsqueeze(1)
        lines = torch.cat((lines, a), dim=1)

        target["lines"] = lines.to(torch.float32).view(-1,2,3)
        target["img_size"]=shape

        validate_keypoints(lines, shape[0], shape[1])
        return target

    def show(self, idx,show_type='all'):
        image, target = self.__getitem__(idx)

        cmap = plt.get_cmap("jet")
        norm = mpl.colors.Normalize(vmin=0.4, vmax=1.0)
        sm = plt.cm.ScalarMappable(cmap=cmap, norm=norm)
        sm.set_array([])

        img_path = os.path.join(self.img_path, self.imgs[idx])
        img = PIL.Image.open(img_path).convert('RGB')
        if show_type=='all':
            boxed_image = draw_bounding_boxes((self.default_transform(img) * 255).to(torch.uint8), target["boxes"],
                                                  colors="yellow", width=1)
            keypoint_img=draw_keypoints(boxed_image,target['lines'],colors='red',width=3)
            plt.imshow(keypoint_img.permute(1, 2, 0).numpy())
            plt.show()

        if show_type=='lines':
            keypoint_img=draw_keypoints((self.default_transform(img) * 255).to(torch.uint8),target['lines'],colors='red',width=3)
            plt.imshow(keypoint_img.permute(1, 2, 0).numpy())
            plt.show()

        if show_type=='boxes':
            boxed_image = draw_bounding_boxes((self.default_transform(img) * 255).to(torch.uint8), target["boxes"],
                                              colors="yellow", width=1)
            plt.imshow(boxed_image.permute(1, 2, 0).numpy())
            plt.show()




    def show_img(self, img_path):
        pass

def get_boxes_lines(objs,shape):
    boxes = []
    h,w=shape
    line_point_pairs = []

    for obj in objs:
        # plt.plot([a[1], b[1]], [a[0], b[0]], c="red", linewidth=1)  # a[1], b[1]无明确大小

        # print(f"points:{obj['points']}")

        a,b=obj['points'][0],obj['points'][1]

        line_point_pairs.append(a)
        line_point_pairs.append(b)

        xmin = max(0, (min(a[0], b[0]) - 6))
        xmax = min(w, (max(a[0], b[0]) + 6))
        ymin = max(0, (min(a[1], b[1]) - 6))
        ymax = min(h, (max(a[1], b[1]) + 6))

        boxes.append([ xmin,ymin,  xmax,ymax])

    boxes=torch.tensor(boxes)
    line_point_pairs=torch.tensor(line_point_pairs)

    # print(f'boxes:{boxes.shape},line_point_pairs:{line_point_pairs.shape}')
    return boxes,line_point_pairs

if __name__ == '__main__':
    path=r"\\192.168.50.222/share/rlq/datasets/0706_"
    dataset= LineDataset(dataset_path=path, dataset_type='train',data_type='jpg')
    dataset.show(1,show_type='lines')