# Copyright (C) 2022 Intel Corporation
#
# SPDX-License-Identifier: MIT

from collections import OrderedDict
from enum import Enum, auto

from datumaro.components.annotation import AnnotationType, LabelCategories, MaskCategories
from datumaro.util import parse_json_file
from datumaro.util.mask_tools import generate_colormap


def parse_config_file(config_path):
    label_map = OrderedDict()
    config = parse_json_file(config_path)
    for label in config["labels"]:
        label_map[label["name"]] = tuple(map(int, label["color"]))
    return label_map


def make_mapillary_instance_categories(label_map):
    label_cat = LabelCategories()
    for label, _ in label_map.items():
        label_cat.add(label)

    has_colors = any([len(color) == 3 for color in label_map.values()])
    if not has_colors:
        colormap = generate_colormap(len(label_map))
    else:
        colormap = {label_cat.find(label)[0]: color for label, color in label_map.items()}

    mask_cat = MaskCategories(colormap)
    mask_cat.inverse_colormap  # pylint: disable=pointless-statement

    return {AnnotationType.label: label_cat, AnnotationType.mask: mask_cat}


class MapillaryVistasTask(Enum):
    instances = auto()
    panoptic = auto()


class MapillaryVistasPath:
    IMAGES_DIR = "images"
    INSTANCES_DIR = "instances"
    PANOPTIC_DIR = "panoptic"
    POLYGON_DIR = "polygons"
    CLASS_DIR = "labels"
    MASK_EXT = ".png"

    ANNOTATION_DIRS = {
        "v1.2": [CLASS_DIR, INSTANCES_DIR],
        "v2.0": [CLASS_DIR, INSTANCES_DIR, PANOPTIC_DIR, POLYGON_DIR],
    }

    CONFIG_FILES = {"v1.2": "config_v1.2.json", "v2.0": "config_v2.0.json"}
    PANOPTIC_CONFIG = "panoptic_2020.json"

    CLASS_BY_DIR = {
        INSTANCES_DIR: MapillaryVistasTask.instances,
        CLASS_DIR: MapillaryVistasTask.instances,
        POLYGON_DIR: MapillaryVistasTask.instances,
        PANOPTIC_DIR: MapillaryVistasTask.panoptic,
    }


MapillaryVistasLabelMap_V12 = OrderedDict(
    [
        ("animal--bird", (165, 42, 42)),
        ("animal--ground-animal", (0, 192, 0)),
        ("construction--barrier--curb", (196, 196, 196)),
        ("construction--barrier--fence", (190, 153, 153)),
        ("construction--barrier--guard-rail", (180, 165, 180)),
        ("construction--barrier--other-barrier", (90, 120, 150)),
        ("construction--barrier--wall", (102, 102, 156)),
        ("construction--flat--bike-lane", (128, 64, 255)),
        ("construction--flat--crosswalk-plain", (140, 140, 200)),
        ("construction--flat--curb-cut", (170, 170, 170)),
        ("construction--flat--parking", (250, 170, 160)),
        ("construction--flat--pedestrian-area", (96, 96, 96)),
        ("construction--flat--rail-track", (230, 150, 140)),
        ("construction--flat--road", (128, 64, 128)),
        ("construction--flat--service-lane", (110, 110, 110)),
        ("construction--flat--sidewalk", (244, 35, 232)),
        ("construction--structure--bridge", (150, 100, 100)),
        ("construction--structure--building", (70, 70, 70)),
        ("construction--structure--tunnel", (150, 120, 90)),
        ("human--person", (220, 20, 60)),
        ("human--rider--bicyclist", (255, 0, 0)),
        ("human--rider--motorcyclist", (255, 0, 100)),
        ("human--rider--other-rider", (255, 0, 200)),
        ("marking--crosswalk-zebra", (200, 128, 128)),
        ("marking--general", (255, 255, 255)),
        ("nature--mountain", (64, 170, 64)),
        ("nature--sand", (230, 160, 50)),
        ("nature--sky", (70, 130, 180)),
        ("nature--snow", (190, 255, 255)),
        ("nature--terrain", (152, 251, 152)),
        ("nature--vegetation", (107, 142, 35)),
        ("nature--water", (0, 170, 30)),
        ("object--banner", (255, 255, 128)),
        ("object--bench", (250, 0, 30)),
        ("object--bike-rack", (100, 140, 180)),
        ("object--billboard", (220, 220, 220)),
        ("object--catch-basin", (220, 128, 128)),
        ("object--cctv-camera", (222, 40, 40)),
        ("object--fire-hydrant", (100, 170, 30)),
        ("object--junction-box", (40, 40, 40)),
        ("object--mailbox", (33, 33, 33)),
        ("object--manhole", (100, 128, 160)),
        ("object--phone-booth", (142, 0, 0)),
        ("object--pothole", (70, 100, 150)),
        ("object--street-light", (210, 170, 100)),
        ("object--support--pole", (153, 153, 153)),
        ("object--support--traffic-sign-frame", (128, 128, 128)),
        ("object--support--utility-pole", (0, 0, 80)),
        ("object--traffic-light", (250, 170, 30)),
        ("object--traffic-sign--back", (192, 192, 192)),
        ("object--traffic-sign--front", (220, 220, 0)),
        ("object--trash-can", (140, 140, 20)),
        ("object--vehicle--bicycle", (119, 11, 32)),
        ("object--vehicle--boat", (150, 0, 255)),
        ("object--vehicle--bus", (0, 60, 100)),
        ("object--vehicle--car", (0, 0, 142)),
        ("object--vehicle--caravan", (0, 0, 90)),
        ("object--vehicle--motorcycle", (0, 0, 230)),
        ("object--vehicle--on-rails", (0, 80, 100)),
        ("object--vehicle--other-vehicle", (128, 64, 64)),
        ("object--vehicle--trailer", (0, 0, 110)),
        ("object--vehicle--truck", (0, 0, 70)),
        ("object--vehicle--wheeled-slow", (0, 0, 192)),
        ("void--car-mount", (32, 32, 32)),
        ("void--ego-vehicle", (120, 10, 10)),
        ("void--unlabeled", (0, 0, 0)),
    ]
)

MapillaryVistasLabelMap_V20 = OrderedDict(
    [
        ("animal--bird", (165, 42, 42)),
        ("animal--ground-animal", (0, 192, 0)),
        ("construction--barrier--ambiguous", (250, 170, 31)),
        ("construction--barrier--concrete-block", (250, 170, 32)),
        ("construction--barrier--curb", (196, 196, 196)),
        ("construction--barrier--fence", (190, 153, 153)),
        ("construction--barrier--guard-rail", (180, 165, 180)),
        ("construction--barrier--other-barrier", (90, 120, 150)),
        ("construction--barrier--road-median", (250, 170, 33)),
        ("construction--barrier--road-side", (250, 170, 34)),
        ("construction--barrier--separator", (128, 128, 128)),
        ("construction--barrier--temporary", (250, 170, 35)),
        ("construction--barrier--wall", (102, 102, 156)),
        ("construction--flat--bike-lane", (128, 64, 255)),
        ("construction--flat--crosswalk-plain", (140, 140, 200)),
        ("construction--flat--curb-cut", (170, 170, 170)),
        ("construction--flat--driveway", (250, 170, 36)),
        ("construction--flat--parking", (250, 170, 160)),
        ("construction--flat--parking-aisle", (250, 170, 37)),
        ("construction--flat--pedestrian-area", (96, 96, 96)),
        ("construction--flat--rail-track", (230, 150, 140)),
        ("construction--flat--road", (128, 64, 128)),
        ("construction--flat--road-shoulder", (110, 110, 110)),
        ("construction--flat--service-lane", (110, 110, 110)),
        ("construction--flat--sidewalk", (244, 35, 232)),
        ("construction--flat--traffic-island", (128, 196, 128)),
        ("construction--structure--bridge", (150, 100, 100)),
        ("construction--structure--building", (70, 70, 70)),
        ("construction--structure--garage", (150, 150, 150)),
        ("construction--structure--tunnel", (150, 120, 90)),
        ("human--person--individual", (220, 20, 60)),
        ("human--person--person-group", (220, 20, 60)),
        ("human--rider--bicyclist", (255, 0, 0)),
        ("human--rider--motorcyclist", (255, 0, 100)),
        ("human--rider--other-rider", (255, 0, 200)),
        ("marking--continuous--dashed", (255, 255, 255)),
        ("marking--continuous--solid", (255, 255, 255)),
        ("marking--continuous--zigzag", (250, 170, 29)),
        ("marking--discrete--ambiguous", (250, 170, 28)),
        ("marking--discrete--arrow--left", (250, 170, 26)),
        ("marking--discrete--arrow--other", (250, 170, 25)),
        ("marking--discrete--arrow--right", (250, 170, 24)),
        ("marking--discrete--arrow--split-left-or-straight", (250, 170, 22)),
        ("marking--discrete--arrow--split-right-or-straight", (250, 170, 21)),
        ("marking--discrete--arrow--straight", (250, 170, 20)),
        ("marking--discrete--crosswalk-zebra", (255, 255, 255)),
        ("marking--discrete--give-way-row", (250, 170, 19)),
        ("marking--discrete--give-way-single", (250, 170, 18)),
        ("marking--discrete--hatched--chevron", (250, 170, 12)),
        ("marking--discrete--hatched--diagonal", (250, 170, 11)),
        ("marking--discrete--other-marking", (255, 255, 255)),
        ("marking--discrete--stop-line", (255, 255, 255)),
        ("marking--discrete--symbol--bicycle", (250, 170, 16)),
        ("marking--discrete--symbol--other", (250, 170, 15)),
        ("marking--discrete--text", (250, 170, 15)),
        ("marking-only--continuous--dashed", (255, 255, 255)),
        ("marking-only--discrete--crosswalk-zebra", (255, 255, 255)),
        ("marking-only--discrete--other-marking", (255, 255, 255)),
        ("marking-only--discrete--text", (255, 255, 255)),
        ("nature--mountain", (64, 170, 64)),
        ("nature--sand", (230, 160, 50)),
        ("nature--sky", (70, 130, 180)),
        ("nature--snow", (190, 255, 255)),
        ("nature--terrain", (152, 251, 152)),
        ("nature--vegetation", (107, 142, 35)),
        ("nature--water", (0, 170, 30)),
        ("object--banner", (255, 255, 128)),
        ("object--bench", (250, 0, 30)),
        ("object--bike-rack", (100, 140, 180)),
        ("object--catch-basin", (220, 128, 128)),
        ("object--cctv-camera", (222, 40, 40)),
        ("object--fire-hydrant", (100, 170, 30)),
        ("object--junction-box", (40, 40, 40)),
        ("object--mailbox", (33, 33, 33)),
        ("object--manhole", (100, 128, 160)),
        ("object--parking-meter", (20, 20, 255)),
        ("object--phone-booth", (142, 0, 0)),
        ("object--pothole", (70, 100, 150)),
        ("object--sign--advertisement", (250, 171, 30)),
        ("object--sign--ambiguous", (250, 172, 30)),
        ("object--sign--back", (250, 173, 30)),
        ("object--sign--information", (250, 174, 30)),
        ("object--sign--other", (250, 175, 30)),
        ("object--sign--store", (250, 176, 30)),
        ("object--street-light", (210, 170, 100)),
        ("object--support--pole", (153, 153, 153)),
        ("object--support--pole-group", (153, 153, 153)),
        ("object--support--traffic-sign-frame", (128, 128, 128)),
        ("object--support--utility-pole", (0, 0, 80)),
        ("object--traffic-cone", (210, 60, 60)),
        ("object--traffic-light--cyclists", (250, 170, 30)),
        ("object--traffic-light--general-horizontal", (250, 170, 30)),
        ("object--traffic-light--general-single", (250, 170, 30)),
        ("object--traffic-light--general-upright", (250, 170, 30)),
        ("object--traffic-light--other", (250, 170, 30)),
        ("object--traffic-light--pedestrians", (250, 170, 30)),
        ("object--traffic-sign--ambiguous", (192, 192, 192)),
        ("object--traffic-sign--back", (192, 192, 192)),
        ("object--traffic-sign--direction-back", (192, 192, 192)),
        ("object--traffic-sign--direction-front", (220, 220, 0)),
        ("object--traffic-sign--front", (220, 220, 0)),
        ("object--traffic-sign--information-parking", (0, 0, 196)),
        ("object--traffic-sign--temporary-back", (192, 192, 192)),
        ("object--traffic-sign--temporary-front", (220, 220, 0)),
        ("object--trash-can", (140, 140, 20)),
        ("object--vehicle--bicycle", (119, 11, 32)),
        ("object--vehicle--boat", (150, 0, 255)),
        ("object--vehicle--bus", (0, 60, 100)),
        ("object--vehicle--car", (0, 0, 142)),
        ("object--vehicle--caravan", (0, 0, 90)),
        ("object--vehicle--motorcycle", (0, 0, 230)),
        ("object--vehicle--on-rails", (0, 80, 100)),
        ("object--vehicle--other-vehicle", (128, 64, 64)),
        ("object--vehicle--trailer", (0, 0, 110)),
        ("object--vehicle--truck", (0, 0, 70)),
        ("object--vehicle--vehicle-group", (0, 0, 142)),
        ("object--vehicle--wheeled-slow", (0, 0, 192)),
        ("object--water-valve", (170, 170, 170)),
        ("void--car-mount", (32, 32, 32)),
        ("void--dynamic", (111, 74, 0)),
        ("void--ego-vehicle", (120, 10, 10)),
        ("void--ground", (81, 0, 81)),
        ("void--static", (111, 111, 0)),
        ("void--unlabeled", (0, 0, 0)),
    ]
)

MapillaryVistasLabelMaps = {
    "v1.2": MapillaryVistasLabelMap_V12,
    "v2.0": MapillaryVistasLabelMap_V20,
}
