from mmdet.registry import DATASETS
from mmdet.datasets.base_video_dataset import BaseVideoDataset


@DATASETS.register_module()
class YouTubeVISDatasetV2(BaseVideoDataset):
    """YouTube VIS dataset for video instance segmentation.

    Args:
        dataset_version (str): Select dataset year version.
    """

    def __init__(self, dataset_version: str, *args, **kwargs):
        self.set_dataset_classes(dataset_version)
        self.dataset_name = f'YouTubeVISDataset_{dataset_version}'
        super().__init__(*args, **kwargs)

    @classmethod
    def set_dataset_classes(cls, dataset_version: str) -> None:
        """Pass the category of the corresponding year to metainfo.

        Args:
            dataset_version (str): Select dataset year version.
        """
        classes_2019_version = ('person', 'giant_panda', 'lizard', 'parrot',
                                'skateboard', 'sedan', 'ape', 'dog', 'snake',
                                'monkey', 'hand', 'rabbit', 'duck', 'cat',
                                'cow', 'fish', 'train', 'horse', 'turtle',
                                'bear', 'motorbike', 'giraffe', 'leopard',
                                'fox', 'deer', 'owl', 'surfboard', 'airplane',
                                'truck', 'zebra', 'tiger', 'elephant',
                                'snowboard', 'boat', 'shark', 'mouse', 'frog',
                                'eagle', 'earless_seal', 'tennis_racket')

        classes_2021_version = ('airplane', 'bear', 'bird', 'boat', 'car',
                                'cat', 'cow', 'deer', 'dog', 'duck',
                                'earless_seal', 'elephant', 'fish',
                                'flying_disc', 'fox', 'frog', 'giant_panda',
                                'giraffe', 'horse', 'leopard', 'lizard',
                                'monkey', 'motorbike', 'mouse', 'parrot',
                                'person', 'rabbit', 'shark', 'skateboard',
                                'snake', 'snowboard', 'squirrel', 'surfboard',
                                'tennis_racket', 'tiger', 'train', 'truck',
                                'turtle', 'whale', 'zebra')

        classes_ovis_version = ('Person', 'Bird', 'Cat', 'Dog', 'Horse',
                                'Sheep', 'Cow', 'Elephant', 'Bear', 'Zebra',
                                'Giraffe', 'Poultry', 'Giant_panda', 'Lizard', 'Parrot',
                                'Monkey', 'Rabbit', 'Tiger', 'Fish', 'Turtle',
                                'Bicycle', 'Motorcycle', 'Airplane', 'Boat', 'Vehical')

        if dataset_version == '2019':
            cls.METAINFO = dict(classes=classes_2019_version)
        elif dataset_version == '2021':
            cls.METAINFO = dict(classes=classes_2021_version)
        elif dataset_version == 'ovis':
            cls.METAINFO = dict(classes=classes_ovis_version)
        else:
            raise NotImplementedError('Not supported YouTubeVIS dataset'
                                      f'version: {dataset_version}')