import copy
from typing import List

from mmdet.registry import DATASETS
from mmdet.datasets.coco_panoptic import CocoPanopticDataset
from mmengine import get_local_path


@DATASETS.register_module()
class CocoPanopticOVDataset(CocoPanopticDataset):
    """Coco Open Vocabulary dataset for Panoptic segmentation.
    The class names are changed.
    """

    METAINFO = {
        'classes':
            ('person,child,girl,boy,woman,man,people,children,girls,boys,women,men,lady,guy,ladies,guys,clothes',
             'bicycle,bicycles,bike,bikes',
             'car,cars',
             'motorcycle,motorcycles',
             'airplane,airplanes',
             'bus,buses',
             'train,trains,locomotive,locomotives,freight train',
             'truck,trucks',
             'boat,boats',
             'traffic light',
             'fire hydrant',
             'stop sign',
             'parking meter',
             'bench,benches',
             'bird,birds',
             'cat,cats,kitties,kitty',
             'dog,dogs,puppy,puppies',
             'horse,horses,foal',
             'sheep',
             'cow,cows,calf',
             'elephant,elephants',
             'bear,bears',
             'zebra,zebras',
             'giraffe,giraffes',
             'backpack,backpacks',
             'umbrella,umbrellas',
             'handbag,handbags',
             'tie',
             'suitcase,suitcases',
             'frisbee',
             'skis',
             'snowboard',
             'sports ball',
             'kite,kites',
             'baseball bat',
             'baseball glove',
             'skateboard',
             'surfboard',
             'tennis racket',
             'bottle,bottles,water bottle',
             'wine glass,wine glasses,wineglass',
             'cup,cups,water cup,water glass',
             'fork,forks',
             'knife,knives',
             'spoon,spoons',
             'bowl,bowls',
             'banana,bananas',
             'apple,apples,apple fruit',
             'sandwich,sandwiches',
             'orange fruit',
             'broccoli',
             'carrot,carrots',
             'hot dog',
             'pizza',
             'donut,donuts',
             'cake,cakes',
             'chair,chairs',
             'couch,sofa,sofas',
             'potted plant,potted plants,pottedplant,pottedplants,planter,planters',
             'bed,beds',
             'dining table,dining tables,diningtable,diningtables,plate,plates,diningtable tablecloth',
             'toilet',
             'tv',
             'laptop',
             'mouse',
             'tv remote,remote control',
             'keyboard',
             'cell phone,mobile',
             'microwave',
             'oven,ovens',
             'toaster',
             'sink,sinks',
             'refrigerator,fridge',
             'book,books',
             'clock',
             'vase,vases',
             'scissor,scissors',
             'teddy bear,teddy bears',
             'hair drier',
             'toothbrush,toothbrushes',
             'banner,banners',
             'blanket,blankets',
             'bridge',
             'cardboard',
             'counter',
             'curtain,curtains',
             'door,doors',
             'wood floor',
             'flower,flowers',
             'fruit,fruits',
             'gravel',
             'house',
             'lamp,bulb,lamps,bulbs',
             'mirror',
             'tennis net',
             'pillow,pillows',
             'platform',
             'playingfield,tennis court,baseball field,soccer field,tennis field',
             'railroad',
             'river',
             'road',
             'roof',
             'sand',
             'sea,sea wave,wave,waves',
             'shelf',
             'snow',
             'stairs',
             'tent',
             'towel',
             'brick wall',
             'stone wall',
             'tile wall',
             'wood wall',
             'water',
             'window blind',
             'window',
             'tree,trees,palm tree,bushes',
             'fence,fences',
             'ceiling',
             'sky,clouds',
             'cabinet,cabinets',
             'table',
             'floor,flooring,tile floor',
             'pavement',
             'mountain,mountains',
             'grass',
             'dirt',
             'paper',
             'food',
             'building,buildings',
             'rock',
             'wall,walls',
             'rug',
             ),
        'thing_classes':
            ('person,child,girl,boy,woman,man,people,children,girls,boys,women,men,lady,guy,ladies,guys,clothes',
             'bicycle,bicycles,bike,bikes',
             'car,cars',
             'motorcycle,motorcycles',
             'airplane,airplanes',
             'bus,buses',
             'train,trains,locomotive,locomotives,freight train',
             'truck,trucks',
             'boat,boats',
             'traffic light',
             'fire hydrant',
             'stop sign',
             'parking meter',
             'bench,benches',
             'bird,birds',
             'cat,cats,kitties,kitty',
             'dog,dogs,puppy,puppies',
             'horse,horses,foal',
             'sheep',
             'cow,cows,calf',
             'elephant,elephants',
             'bear,bears',
             'zebra,zebras',
             'giraffe,giraffes',
             'backpack,backpacks',
             'umbrella,umbrellas',
             'handbag,handbags',
             'tie',
             'suitcase,suitcases',
             'frisbee',
             'skis',
             'snowboard',
             'sports ball',
             'kite,kites',
             'baseball bat',
             'baseball glove',
             'skateboard',
             'surfboard',
             'tennis racket',
             'bottle,bottles,water bottle',
             'wine glass,wine glasses,wineglass',
             'cup,cups,water cup,water glass',
             'fork,forks',
             'knife,knives',
             'spoon,spoons',
             'bowl,bowls',
             'banana,bananas',
             'apple,apples,apple fruit',
             'sandwich,sandwiches',
             'orange fruit',
             'broccoli',
             'carrot,carrots',
             'hot dog',
             'pizza',
             'donut,donuts',
             'cake,cakes',
             'chair,chairs',
             'couch,sofa,sofas',
             'potted plant,potted plants,pottedplant,pottedplants,planter,planters',
             'bed,beds',
             'dining table,dining tables,diningtable,diningtables,plate,plates,diningtable tablecloth',
             'toilet',
             'tv',
             'laptop',
             'mouse',
             'tv remote,remote control',
             'keyboard',
             'cell phone,mobile',
             'microwave',
             'oven,ovens',
             'toaster',
             'sink,sinks',
             'refrigerator,fridge',
             'book,books',
             'clock',
             'vase,vases',
             'scissor,scissors',
             'teddy bear,teddy bears',
             'hair drier',
             'toothbrush,toothbrushes',
             ),
        'stuff_classes':
            ('banner,banners',
             'blanket,blankets',
             'bridge',
             'cardboard',
             'counter',
             'curtain,curtains',
             'door,doors',
             'wood floor',
             'flower,flowers',
             'fruit,fruits',
             'gravel',
             'house',
             'lamp,bulb,lamps,bulbs',
             'mirror',
             'tennis net',
             'pillow,pillows',
             'platform',
             'playingfield,tennis court,baseball field,soccer field,tennis field',
             'railroad',
             'river',
             'road',
             'roof',
             'sand',
             'sea,sea wave,wave,waves',
             'shelf',
             'snow',
             'stairs',
             'tent',
             'towel',
             'brick wall',
             'stone wall',
             'tile wall',
             'wood wall',
             'water',
             'window blind',
             'window',
             'tree,trees,palm tree,bushes',
             'fence,fences',
             'ceiling',
             'sky,clouds',
             'cabinet,cabinets',
             'table',
             'floor,flooring,tile floor',
             'pavement',
             'mountain,mountains',
             'grass',
             'dirt',
             'paper',
             'food',
             'building,buildings',
             'rock',
             'wall,walls',
             'rug'
             ),
        'palette':
            [(220, 20, 60), (119, 11, 32), (0, 0, 142), (0, 0, 230), (106, 0, 228),
             (0, 60, 100), (0, 80, 100), (0, 0, 70), (0, 0, 192), (250, 170, 30),
             (100, 170, 30), (220, 220, 0), (175, 116, 175), (250, 0, 30),
             (165, 42, 42), (255, 77, 255), (0, 226, 252), (182, 182, 255),
             (0, 82, 0), (120, 166, 157), (110, 76, 0), (174, 57, 255),
             (199, 100, 0), (72, 0, 118), (255, 179, 240), (0, 125, 92),
             (209, 0, 151), (188, 208, 182), (0, 220, 176), (255, 99, 164),
             (92, 0, 73), (133, 129, 255), (78, 180, 255), (0, 228, 0),
             (174, 255, 243), (45, 89, 255), (134, 134, 103), (145, 148, 174),
             (255, 208, 186), (197, 226, 255), (171, 134, 1), (109, 63, 54),
             (207, 138, 255), (151, 0, 95), (9, 80, 61), (84, 105, 51),
             (74, 65, 105), (166, 196, 102), (208, 195, 210), (255, 109, 65),
             (0, 143, 149), (179, 0, 194), (209, 99, 106), (5, 121, 0),
             (227, 255, 205), (147, 186, 208), (153, 69, 1), (3, 95, 161),
             (163, 255, 0), (119, 0, 170), (0, 182, 199), (0, 165, 120),
             (183, 130, 88), (95, 32, 0), (130, 114, 135), (110, 129, 133),
             (166, 74, 118), (219, 142, 185), (79, 210, 114), (178, 90, 62),
             (65, 70, 15), (127, 167, 115), (59, 105, 106), (142, 108, 45),
             (196, 172, 0), (95, 54, 80), (128, 76, 255), (201, 57, 1),
             (246, 0, 122), (191, 162, 208), (255, 255, 128), (147, 211, 203),
             (150, 100, 100), (168, 171, 172), (146, 112, 198), (210, 170, 100),
             (92, 136, 89), (218, 88, 184), (241, 129, 0), (217, 17, 255),
             (124, 74, 181), (70, 70, 70), (255, 228, 255), (154, 208, 0),
             (193, 0, 92), (76, 91, 113), (255, 180, 195), (106, 154, 176),
             (230, 150, 140), (60, 143, 255), (128, 64, 128), (92, 82, 55),
             (254, 212, 124), (73, 77, 174), (255, 160, 98), (255, 255, 255),
             (104, 84, 109), (169, 164, 131), (225, 199, 255), (137, 54, 74),
             (135, 158, 223), (7, 246, 231), (107, 255, 200), (58, 41, 149),
             (183, 121, 142), (255, 73, 97), (107, 142, 35), (190, 153, 153),
             (146, 139, 141), (70, 130, 180), (134, 199, 156), (209, 226, 140),
             (96, 36, 108), (96, 96, 96), (64, 170, 64), (152, 251, 152),
             (208, 229, 228), (206, 186, 171), (152, 161, 64), (116, 112, 0),
             (0, 114, 143), (102, 102, 156), (250, 141, 255)]
    }

    def load_data_list(self) -> List[dict]:
        """Load annotations from an annotation file named as ``self.ann_file``

        Returns:
            List[dict]: A list of annotation.
        """  # noqa: E501
        with get_local_path(
                self.ann_file, backend_args=self.backend_args) as local_path:
            self.coco = self.COCOAPI(local_path)
        # use all classes, cannot use self.metainfo anymore.
        self.cat_ids = self.coco.get_cat_ids()
        self.cat2label = {cat_id: i for i, cat_id in enumerate(self.cat_ids)}
        self.cat_img_map = copy.deepcopy(self.coco.cat_img_map)

        img_ids = self.coco.get_img_ids()
        data_list = []
        total_ann_ids = []
        for img_id in img_ids:
            raw_img_info = self.coco.load_imgs([img_id])[0]
            raw_img_info['img_id'] = img_id

            ann_ids = self.coco.get_ann_ids(img_ids=[img_id])
            raw_ann_info = self.coco.load_anns(ann_ids)
            total_ann_ids.extend(ann_ids)

            parsed_data_info = self.parse_data_info({
                'raw_ann_info':
                    raw_ann_info,
                'raw_img_info':
                    raw_img_info
            })
            data_list.append(parsed_data_info)
        if self.ANN_ID_UNIQUE:
            assert len(set(total_ann_ids)) == len(
                total_ann_ids
            ), f"Annotation ids in '{self.ann_file}' are not unique!"

        del self.coco

        return data_list