import os

from .o365_ovd import register_objects365_ovd_instances

OBEJECTS365V1_CATEGORIES = [
    {"name": "person", "id": 1},
    {"name": "sneakers", "id": 2},
    {"name": "chair", "id": 3},
    {"name": "hat", "id": 4},
    {"name": "lamp", "id": 5},
    {"name": "bottle", "id": 6},
    {"name": "cabinet/shelf", "id": 7},
    {"name": "cup", "id": 8},
    {"name": "car", "id": 9},
    {"name": "glasses", "id": 10},
    {"name": "picture/frame", "id": 11},
    {"name": "desk", "id": 12},
    {"name": "handbag", "id": 13},
    {"name": "street lights", "id": 14},
    {"name": "book", "id": 15},
    {"name": "plate", "id": 16},
    {"name": "helmet", "id": 17},
    {"name": "leather shoes", "id": 18},
    {"name": "pillow", "id": 19},
    {"name": "glove", "id": 20},
    {"name": "potted plant", "id": 21},
    {"name": "bracelet", "id": 22},
    {"name": "flower", "id": 23},
    {"name": "tv", "id": 24},
    {"name": "storage box", "id": 25},
    {"name": "vase", "id": 26},
    {"name": "bench", "id": 27},
    {"name": "wine glass", "id": 28},
    {"name": "boots", "id": 29},
    {"name": "bowl", "id": 30},
    {"name": "dining table", "id": 31},
    {"name": "umbrella", "id": 32},
    {"name": "boat", "id": 33},
    {"name": "flag", "id": 34},
    {"name": "speaker", "id": 35},
    {"name": "trash bin/can", "id": 36},
    {"name": "stool", "id": 37},
    {"name": "backpack", "id": 38},
    {"name": "couch", "id": 39},
    {"name": "belt", "id": 40},
    {"name": "carpet", "id": 41},
    {"name": "basket", "id": 42},
    {"name": "towel/napkin", "id": 43},
    {"name": "slippers", "id": 44},
    {"name": "barrel/bucket", "id": 45},
    {"name": "coffee table", "id": 46},
    {"name": "suv", "id": 47},
    {"name": "toy", "id": 48},
    {"name": "tie", "id": 49},
    {"name": "bed", "id": 50},
    {"name": "traffic light", "id": 51},
    {"name": "pen/pencil", "id": 52},
    {"name": "microphone", "id": 53},
    {"name": "sandals", "id": 54},
    {"name": "canned", "id": 55},
    {"name": "necklace", "id": 56},
    {"name": "mirror", "id": 57},
    {"name": "faucet", "id": 58},
    {"name": "bicycle", "id": 59},
    {"name": "bread", "id": 60},
    {"name": "high heels", "id": 61},
    {"name": "ring", "id": 62},
    {"name": "van", "id": 63},
    {"name": "watch", "id": 64},
    {"name": "sink", "id": 65},
    {"name": "horse", "id": 66},
    {"name": "fish", "id": 67},
    {"name": "apple", "id": 68},
    {"name": "camera", "id": 69},
    {"name": "candle", "id": 70},
    {"name": "teddy bear", "id": 71},
    {"name": "cake", "id": 72},
    {"name": "motorcycle", "id": 73},
    {"name": "wild bird", "id": 74},
    {"name": "laptop", "id": 75},
    {"name": "knife", "id": 76},
    {"name": "traffic sign", "id": 77},
    {"name": "cell phone", "id": 78},
    {"name": "paddle", "id": 79},
    {"name": "truck", "id": 80},
    {"name": "cow", "id": 81},
    {"name": "power outlet", "id": 82},
    {"name": "clock", "id": 83},
    {"name": "drum", "id": 84},
    {"name": "fork", "id": 85},
    {"name": "bus", "id": 86},
    {"name": "hanger", "id": 87},
    {"name": "nightstand", "id": 88},
    {"name": "pot/pan", "id": 89},
    {"name": "sheep", "id": 90},
    {"name": "guitar", "id": 91},
    {"name": "traffic cone", "id": 92},
    {"name": "tea pot", "id": 93},
    {"name": "keyboard", "id": 94},
    {"name": "tripod", "id": 95},
    {"name": "hockey", "id": 96},
    {"name": "fan", "id": 97},
    {"name": "dog", "id": 98},
    {"name": "spoon", "id": 99},
    {"name": "blackboard/whiteboard", "id": 100},
    {"name": "balloon", "id": 101},
    {"name": "air conditioner", "id": 102},
    {"name": "cymbal", "id": 103},
    {"name": "mouse", "id": 104},
    {"name": "telephone", "id": 105},
    {"name": "pickup truck", "id": 106},
    {"name": "orange", "id": 107},
    {"name": "banana", "id": 108},
    {"name": "airplane", "id": 109},
    {"name": "luggage", "id": 110},
    {"name": "skis", "id": 111},
    {"name": "soccer", "id": 112},
    {"name": "trolley", "id": 113},
    {"name": "oven", "id": 114},
    {"name": "remote", "id": 115},
    {"name": "baseball glove", "id": 116},
    {"name": "paper towel", "id": 117},
    {"name": "refrigerator", "id": 118},
    {"name": "train", "id": 119},
    {"name": "tomato", "id": 120},
    {"name": "machinery vehicle", "id": 121},
    {"name": "tent", "id": 122},
    {"name": "shampoo/shower gel", "id": 123},
    {"name": "head phone", "id": 124},
    {"name": "lantern", "id": 125},
    {"name": "donut", "id": 126},
    {"name": "cleaning products", "id": 127},
    {"name": "sailboat", "id": 128},
    {"name": "tangerine", "id": 129},
    {"name": "pizza", "id": 130},
    {"name": "kite", "id": 131},
    {"name": "computer box", "id": 132},
    {"name": "elephant", "id": 133},
    {"name": "toiletries", "id": 134},
    {"name": "gas stove", "id": 135},
    {"name": "broccoli", "id": 136},
    {"name": "toilet", "id": 137},
    {"name": "stroller", "id": 138},
    {"name": "shovel", "id": 139},
    {"name": "baseball bat", "id": 140},
    {"name": "microwave", "id": 141},
    {"name": "skateboard", "id": 142},
    {"name": "surfboard", "id": 143},
    {"name": "surveillance camera", "id": 144},
    {"name": "gun", "id": 145},
    {"name": "life saver", "id": 146},
    {"name": "cat", "id": 147},
    {"name": "lemon", "id": 148},
    {"name": "liquid soap", "id": 149},
    {"name": "zebra", "id": 150},
    {"name": "duck", "id": 151},
    {"name": "sports car", "id": 152},
    {"name": "giraffe", "id": 153},
    {"name": "pumpkin", "id": 154},
    {"name": "piano", "id": 155},
    {"name": "stop sign", "id": 156},
    {"name": "radiator", "id": 157},
    {"name": "converter", "id": 158},
    {"name": "tissue ", "id": 159},
    {"name": "carrot", "id": 160},
    {"name": "washing machine", "id": 161},
    {"name": "vent", "id": 162},
    {"name": "cookies", "id": 163},
    {"name": "cutting/chopping board", "id": 164},
    {"name": "tennis racket", "id": 165},
    {"name": "candy", "id": 166},
    {"name": "skating and skiing shoes", "id": 167},
    {"name": "scissors", "id": 168},
    {"name": "folder", "id": 169},
    {"name": "baseball", "id": 170},
    {"name": "strawberry", "id": 171},
    {"name": "bow tie", "id": 172},
    {"name": "pigeon", "id": 173},
    {"name": "pepper", "id": 174},
    {"name": "coffee machine", "id": 175},
    {"name": "bathtub", "id": 176},
    {"name": "snowboard", "id": 177},
    {"name": "suitcase", "id": 178},
    {"name": "grapes", "id": 179},
    {"name": "ladder", "id": 180},
    {"name": "pear", "id": 181},
    {"name": "american football", "id": 182},
    {"name": "basketball", "id": 183},
    {"name": "potato", "id": 184},
    {"name": "paint brush", "id": 185},
    {"name": "printer", "id": 186},
    {"name": "billiards", "id": 187},
    {"name": "fire hydrant", "id": 188},
    {"name": "goose", "id": 189},
    {"name": "projector", "id": 190},
    {"name": "sausage", "id": 191},
    {"name": "fire extinguisher", "id": 192},
    {"name": "extension cord", "id": 193},
    {"name": "facial mask", "id": 194},
    {"name": "tennis ball", "id": 195},
    {"name": "chopsticks", "id": 196},
    {"name": "electronic stove and gas stove", "id": 197},
    {"name": "pie", "id": 198},
    {"name": "frisbee", "id": 199},
    {"name": "kettle", "id": 200},
    {"name": "hamburger", "id": 201},
    {"name": "golf club", "id": 202},
    {"name": "cucumber", "id": 203},
    {"name": "clutch", "id": 204},
    {"name": "blender", "id": 205},
    {"name": "tong", "id": 206},
    {"name": "slide", "id": 207},
    {"name": "hot dog", "id": 208},
    {"name": "toothbrush", "id": 209},
    {"name": "facial cleanser", "id": 210},
    {"name": "mango", "id": 211},
    {"name": "deer", "id": 212},
    {"name": "egg", "id": 213},
    {"name": "violin", "id": 214},
    {"name": "marker", "id": 215},
    {"name": "ship", "id": 216},
    {"name": "chicken", "id": 217},
    {"name": "onion", "id": 218},
    {"name": "ice cream", "id": 219},
    {"name": "tape", "id": 220},
    {"name": "wheelchair", "id": 221},
    {"name": "plum", "id": 222},
    {"name": "bar soap", "id": 223},
    {"name": "scale", "id": 224},
    {"name": "watermelon", "id": 225},
    {"name": "cabbage", "id": 226},
    {"name": "router/modem", "id": 227},
    {"name": "golf ball", "id": 228},
    {"name": "pine apple", "id": 229},
    {"name": "crane", "id": 230},
    {"name": "fire truck", "id": 231},
    {"name": "peach", "id": 232},
    {"name": "cello", "id": 233},
    {"name": "notepaper", "id": 234},
    {"name": "tricycle", "id": 235},
    {"name": "toaster", "id": 236},
    {"name": "helicopter", "id": 237},
    {"name": "green beans", "id": 238},
    {"name": "brush", "id": 239},
    {"name": "carriage", "id": 240},
    {"name": "cigar", "id": 241},
    {"name": "earphone", "id": 242},
    {"name": "penguin", "id": 243},
    {"name": "hurdle", "id": 244},
    {"name": "swing", "id": 245},
    {"name": "radio", "id": 246},
    {"name": "CD", "id": 247},
    {"name": "parking meter", "id": 248},
    {"name": "swan", "id": 249},
    {"name": "garlic", "id": 250},
    {"name": "french fries", "id": 251},
    {"name": "horn", "id": 252},
    {"name": "avocado", "id": 253},
    {"name": "saxophone", "id": 254},
    {"name": "trumpet", "id": 255},
    {"name": "sandwich", "id": 256},
    {"name": "cue", "id": 257},
    {"name": "kiwi fruit", "id": 258},
    {"name": "bear", "id": 259},
    {"name": "fishing rod", "id": 260},
    {"name": "cherry", "id": 261},
    {"name": "tablet", "id": 262},
    {"name": "green vegetables", "id": 263},
    {"name": "nuts", "id": 264},
    {"name": "corn", "id": 265},
    {"name": "key", "id": 266},
    {"name": "screwdriver", "id": 267},
    {"name": "globe", "id": 268},
    {"name": "broom", "id": 269},
    {"name": "pliers", "id": 270},
    {"name": "volleyball", "id": 271},
    {"name": "hammer", "id": 272},
    {"name": "eggplant", "id": 273},
    {"name": "trophy", "id": 274},
    {"name": "dates", "id": 275},
    {"name": "board eraser", "id": 276},
    {"name": "rice", "id": 277},
    {"name": "tape measure/ruler", "id": 278},
    {"name": "dumbbell", "id": 279},
    {"name": "hamimelon", "id": 280},
    {"name": "stapler", "id": 281},
    {"name": "camel", "id": 282},
    {"name": "lettuce", "id": 283},
    {"name": "goldfish", "id": 284},
    {"name": "meat balls", "id": 285},
    {"name": "medal", "id": 286},
    {"name": "toothpaste", "id": 287},
    {"name": "antelope", "id": 288},
    {"name": "shrimp", "id": 289},
    {"name": "rickshaw", "id": 290},
    {"name": "trombone", "id": 291},
    {"name": "pomegranate", "id": 292},
    {"name": "coconut", "id": 293},
    {"name": "jellyfish", "id": 294},
    {"name": "mushroom", "id": 295},
    {"name": "calculator", "id": 296},
    {"name": "treadmill", "id": 297},
    {"name": "butterfly", "id": 298},
    {"name": "egg tart", "id": 299},
    {"name": "cheese", "id": 300},
    {"name": "pig", "id": 301},
    {"name": "pomelo", "id": 302},
    {"name": "race car", "id": 303},
    {"name": "rice cooker", "id": 304},
    {"name": "tuba", "id": 305},
    {"name": "crosswalk sign", "id": 306},
    {"name": "papaya", "id": 307},
    {"name": "hair drier", "id": 308},
    {"name": "green onion", "id": 309},
    {"name": "chips", "id": 310},
    {"name": "dolphin", "id": 311},
    {"name": "sushi", "id": 312},
    {"name": "urinal", "id": 313},
    {"name": "donkey", "id": 314},
    {"name": "electric drill", "id": 315},
    {"name": "spring rolls", "id": 316},
    {"name": "tortoise/turtle", "id": 317},
    {"name": "parrot", "id": 318},
    {"name": "flute", "id": 319},
    {"name": "measuring cup", "id": 320},
    {"name": "shark", "id": 321},
    {"name": "steak", "id": 322},
    {"name": "poker card", "id": 323},
    {"name": "binoculars", "id": 324},
    {"name": "llama", "id": 325},
    {"name": "radish", "id": 326},
    {"name": "noodles", "id": 327},
    {"name": "yak", "id": 328},
    {"name": "mop", "id": 329},
    {"name": "crab", "id": 330},
    {"name": "microscope", "id": 331},
    {"name": "barbell", "id": 332},
    {"name": "bread/bun", "id": 333},
    {"name": "baozi", "id": 334},
    {"name": "lion", "id": 335},
    {"name": "red cabbage", "id": 336},
    {"name": "polar bear", "id": 337},
    {"name": "lighter", "id": 338},
    {"name": "seal", "id": 339},
    {"name": "mangosteen", "id": 340},
    {"name": "comb", "id": 341},
    {"name": "eraser", "id": 342},
    {"name": "pitaya", "id": 343},
    {"name": "scallop", "id": 344},
    {"name": "pencil case", "id": 345},
    {"name": "saw", "id": 346},
    {"name": "table tennis paddle", "id": 347},
    {"name": "okra", "id": 348},
    {"name": "starfish", "id": 349},
    {"name": "eagle", "id": 350},
    {"name": "monkey", "id": 351},
    {"name": "durian", "id": 352},
    {"name": "game board", "id": 353},
    {"name": "rabbit", "id": 354},
    {"name": "french horn", "id": 355},
    {"name": "ambulance", "id": 356},
    {"name": "asparagus", "id": 357},
    {"name": "hoverboard", "id": 358},
    {"name": "pasta", "id": 359},
    {"name": "target", "id": 360},
    {"name": "hotair balloon", "id": 361},
    {"name": "chainsaw", "id": 362},
    {"name": "lobster", "id": 363},
    {"name": "iron", "id": 364},
    {"name": "flashlight", "id": 365},
]

_PREDEFINED_SPLITS = {
    # image_root, json_file, num_sampled_classes, template
    "o365v1_train_ovd_unipro": (
        "o365/train",
        "o365/annotations/objects365_train.json",
        150,
        "full",
    ),
    "o365v1_train_ovd": (
        "o365/train",
        "o365/annotations/objects365_train.json",
        150,
        "identity",
    ),
}


def _get_objects365_instances_meta():
    OBEJECTS365_CATEGORIES = OBEJECTS365V1_CATEGORIES

    thing_ids = [k["id"] - 1 for k in OBEJECTS365_CATEGORIES]
    assert len(thing_ids) == 365, len(thing_ids)
    # Mapping from the incontiguous Objects365 category id to an id in [0, 364]
    thing_dataset_id_to_contiguous_id = {k: i for i, k in enumerate(thing_ids)}
    thing_classes = [k["name"] for k in OBEJECTS365_CATEGORIES]
    ret = {
        "thing_dataset_id_to_contiguous_id": thing_dataset_id_to_contiguous_id,
        "thing_classes": thing_classes,
    }
    return ret


def register_all_objects365_instances(root):
    for key, (
        image_root,
        json_file,
        num_sampled_classes,
        template,
    ) in _PREDEFINED_SPLITS.items():
        # Assume pre-defined datasets live in `./datas`.
        register_objects365_ovd_instances(
            key,
            _get_objects365_instances_meta(),
            os.path.join(root, json_file) if "://" not in json_file else json_file,
            os.path.join(root, image_root),
            num_sampled_classes,
            template=template,
            test_mode=True if "val" in key else False,
        )


_root = os.getenv("DETECTRON2_DATASETS", "datasets")
register_all_objects365_instances(_root)
