import os

import torch
from torch.utils.data import DataLoader, Dataset, random_split


class VQADataset(Dataset):
    def __init__(self, mode, root="data/coco/cocoqa"):
        assert mode in {"train", "test"}
        super().__init__()
        data_path = os.path.join(root, mode)
        self.image_ids = fetch_file(os.path.join(data_path, "img_ids.txt"))
        self.answers = fetch_file(os.path.join(data_path, "answers.txt"))
        self.questions = fetch_file(os.path.join(data_path, "questions.txt"))
        self.types = fetch_file(os.path.join(data_path, "types.txt"))

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

    def __getitem__(self, idx):
        image_id = int(self.image_ids[idx])
        image = torch.load(os.path.join("data/coco/bin", f"{image_id}.pt"))
        question = self.questions[idx]
        answer = answer2idx[self.answers[idx]]
        type_ = self.types[idx]
        return question, image, answer, type_


def fetch_file(path):
    with open(path) as f:
        return f.read().splitlines()


def make_loaders(data_config):
    dataset = VQADataset("train")
    num_total = len(dataset)
    num_valid = int(data_config.valid_split * num_total)
    train_dataset, valid_dataset = random_split(
        dataset, [num_total - num_valid, num_valid]
    )
    test_dataset = VQADataset("test")
    train_loader = DataLoader(
        train_dataset, batch_size=data_config.train_batch_size, shuffle=True,
    )
    valid_loader = DataLoader(
        valid_dataset, batch_size=data_config.valid_batch_size, shuffle=True,
    )
    test_loader = DataLoader(
        test_dataset, batch_size=data_config.valid_batch_size, shuffle=True,
    )
    return train_loader, valid_loader, test_loader


answer2idx = {
    "sheep": 0,
    "blender": 1,
    "hangar": 2,
    "toothbrushes": 3,
    "vehicle": 4,
    "door": 5,
    "chocolate": 6,
    "eight": 7,
    "tie": 8,
    "donuts": 9,
    "paw": 10,
    "flower": 11,
    "fence": 12,
    "pipe": 13,
    "taxi": 14,
    "kites": 15,
    "cap": 16,
    "airplanes": 17,
    "cupcakes": 18,
    "tent": 19,
    "jar": 20,
    "jet": 21,
    "zebras": 22,
    "skateboard": 23,
    "tray": 24,
    "crib": 25,
    "suits": 26,
    "umbrella": 27,
    "brick": 28,
    "wheelchair": 29,
    "dog": 30,
    "goats": 31,
    "hat": 32,
    "buses": 33,
    "bat": 34,
    "birds": 35,
    "bag": 36,
    "glass": 37,
    "calf": 38,
    "kitten": 39,
    "shelf": 40,
    "station": 41,
    "snowboard": 42,
    "ship": 43,
    "garage": 44,
    "coat": 45,
    "cabinet": 46,
    "machine": 47,
    "bananas": 48,
    "sailboat": 49,
    "store": 50,
    "bicycles": 51,
    "library": 52,
    "backpack": 53,
    "wall": 54,
    "snowboards": 55,
    "dryer": 56,
    "bathtub": 57,
    "beer": 58,
    "paddle": 59,
    "sinks": 60,
    "helicopter": 61,
    "doors": 62,
    "sandwich": 63,
    "freight": 64,
    "clocks": 65,
    "graffiti": 66,
    "scooter": 67,
    "bicycle": 68,
    "painting": 69,
    "pizzas": 70,
    "giraffe": 71,
    "helmets": 72,
    "bread": 73,
    "dock": 74,
    "case": 75,
    "tracks": 76,
    "mountain": 77,
    "two": 78,
    "zebra": 79,
    "platter": 80,
    "subway": 81,
    "apartment": 82,
    "appliances": 83,
    "poles": 84,
    "fish": 85,
    "knife": 86,
    "dishes": 87,
    "road": 88,
    "dessert": 89,
    "pigeon": 90,
    "bathroom": 91,
    "guitar": 92,
    "train": 93,
    "slices": 94,
    "cigarette": 95,
    "holder": 96,
    "airliner": 97,
    "vegetables": 98,
    "street": 99,
    "container": 100,
    "skies": 101,
    "surfboard": 102,
    "jacket": 103,
    "outdoors": 104,
    "locomotive": 105,
    "window": 106,
    "salad": 107,
    "trains": 108,
    "banana": 109,
    "pen": 110,
    "restaurant": 111,
    "sun": 112,
    "turkey": 113,
    "couch": 114,
    "candle": 115,
    "chair": 116,
    "cage": 117,
    "terminal": 118,
    "leaves": 119,
    "ten": 120,
    "yellow": 121,
    "glasses": 122,
    "hose": 123,
    "remote": 124,
    "stove": 125,
    "horses": 126,
    "sink": 127,
    "frame": 128,
    "rail": 129,
    "toilet": 130,
    "laptop": 131,
    "apron": 132,
    "tree": 133,
    "bar": 134,
    "entree": 135,
    "basket": 136,
    "boards": 137,
    "ingredients": 138,
    "side": 139,
    "carrot": 140,
    "owl": 141,
    "counter": 142,
    "ties": 143,
    "white": 144,
    "airplane": 145,
    "sky": 146,
    "hallway": 147,
    "pictures": 148,
    "motorcycle": 149,
    "purple": 150,
    "curtain": 151,
    "elephant": 152,
    "orange": 153,
    "wine": 154,
    "oven": 155,
    "umbrellas": 156,
    "garden": 157,
    "gray": 158,
    "suit": 159,
    "vegetable": 160,
    "dinner": 161,
    "walls": 162,
    "jets": 163,
    "walkway": 164,
    "fridge": 165,
    "seagulls": 166,
    "vehicles": 167,
    "toilets": 168,
    "five": 169,
    "bottle": 170,
    "bank": 171,
    "urinal": 172,
    "vest": 173,
    "warehouse": 174,
    "four": 175,
    "bottles": 176,
    "marina": 177,
    "flag": 178,
    "hamburger": 179,
    "tools": 180,
    "sign": 181,
    "plate": 182,
    "mirror": 183,
    "mask": 184,
    "scooters": 185,
    "pastry": 186,
    "trucks": 187,
    "pole": 188,
    "sofa": 189,
    "brown": 190,
    "pans": 191,
    "cat": 192,
    "suitcases": 193,
    "hay": 194,
    "fork": 195,
    "beverage": 196,
    "duck": 197,
    "shirt": 198,
    "luggage": 199,
    "apple": 200,
    "colors": 201,
    "scissors": 202,
    "six": 203,
    "phones": 204,
    "room": 205,
    "restroom": 206,
    "sailboats": 207,
    "urinals": 208,
    "bull": 209,
    "picture": 210,
    "lambs": 211,
    "vases": 212,
    "kitchen": 213,
    "motorcycles": 214,
    "chicken": 215,
    "grass": 216,
    "cakes": 217,
    "nine": 218,
    "bridge": 219,
    "tub": 220,
    "hillside": 221,
    "doorway": 222,
    "plants": 223,
    "plant": 224,
    "device": 225,
    "rack": 226,
    "meal": 227,
    "hydrant": 228,
    "sandwiches": 229,
    "screen": 230,
    "frisbee": 231,
    "milk": 232,
    "airport": 233,
    "drink": 234,
    "tablet": 235,
    "pool": 236,
    "statues": 237,
    "cargo": 238,
    "one": 239,
    "surfboards": 240,
    "beach": 241,
    "highway": 242,
    "boat": 243,
    "seat": 244,
    "boxes": 245,
    "shower": 246,
    "plates": 247,
    "clock": 248,
    "suitcase": 249,
    "pin": 250,
    "trunks": 251,
    "black": 252,
    "glove": 253,
    "meat": 254,
    "candles": 255,
    "goat": 256,
    "carriage": 257,
    "skis": 258,
    "foil": 259,
    "mountains": 260,
    "eagle": 261,
    "house": 262,
    "slice": 263,
    "trailer": 264,
    "backyard": 265,
    "elephants": 266,
    "trunk": 267,
    "freezer": 268,
    "cows": 269,
    "cats": 270,
    "barn": 271,
    "clothes": 272,
    "puppy": 273,
    "engine": 274,
    "bats": 275,
    "stadium": 276,
    "gear": 277,
    "benches": 278,
    "phone": 279,
    "coffee": 280,
    "furniture": 281,
    "television": 282,
    "bags": 283,
    "shelves": 284,
    "cart": 285,
    "driveway": 286,
    "trail": 287,
    "beds": 288,
    "lunch": 289,
    "computers": 290,
    "ram": 291,
    "biplane": 292,
    "pizza": 293,
    "laptops": 294,
    "keyboard": 295,
    "bird": 296,
    "toothbrush": 297,
    "dish": 298,
    "flowers": 299,
    "remotes": 300,
    "horse": 301,
    "lane": 302,
    "tram": 303,
    "fruits": 304,
    "pug": 305,
    "cow": 306,
    "river": 307,
    "bench": 308,
    "equipment": 309,
    "floor": 310,
    "carrots": 311,
    "ball": 312,
    "office": 313,
    "classroom": 314,
    "boats": 315,
    "pie": 316,
    "photograph": 317,
    "cup": 318,
    "statue": 319,
    "bucket": 320,
    "bath": 321,
    "ocean": 322,
    "shop": 323,
    "seven": 324,
    "refrigerators": 325,
    "drinks": 326,
    "ski": 327,
    "parrot": 328,
    "wagon": 329,
    "sidewalk": 330,
    "tv": 331,
    "mouse": 332,
    "towel": 333,
    "flags": 334,
    "stick": 335,
    "cupcake": 336,
    "trolley": 337,
    "uniform": 338,
    "cellphone": 339,
    "lake": 340,
    "wheel": 341,
    "helmet": 342,
    "cake": 343,
    "green": 344,
    "drawer": 345,
    "chairs": 346,
    "museum": 347,
    "breakfast": 348,
    "cattle": 349,
    "panda": 350,
    "spoon": 351,
    "bears": 352,
    "camera": 353,
    "cars": 354,
    "monkey": 355,
    "desk": 356,
    "purse": 357,
    "fruit": 358,
    "toy": 359,
    "apples": 360,
    "vase": 361,
    "pot": 362,
    "pony": 363,
    "trees": 364,
    "oranges": 365,
    "pastries": 366,
    "grizzly": 367,
    "sunglasses": 368,
    "pan": 369,
    "outside": 370,
    "skateboards": 371,
    "frisbees": 372,
    "containers": 373,
    "bed": 374,
    "donut": 375,
    "blue": 376,
    "broccoli": 377,
    "lamb": 378,
    "fireplace": 379,
    "seagull": 380,
    "three": 381,
    "red": 382,
    "roses": 383,
    "hats": 384,
    "bowl": 385,
    "dogs": 386,
    "shuttle": 387,
    "track": 388,
    "cabinets": 389,
    "computer": 390,
    "lamp": 391,
    "hill": 392,
    "box": 393,
    "cheese": 394,
    "scene": 395,
    "ramp": 396,
    "pots": 397,
    "snack": 398,
    "sculpture": 399,
    "table": 400,
    "giraffes": 401,
    "racquet": 402,
    "toys": 403,
    "bus": 404,
    "car": 405,
    "dress": 406,
    "kite": 407,
    "bun": 408,
    "can": 409,
    "buildings": 410,
    "shorts": 411,
    "canoe": 412,
    "building": 413,
    "surface": 414,
    "baskets": 415,
    "towels": 416,
    "cups": 417,
    "refrigerator": 418,
    "truck": 419,
    "ducks": 420,
    "bear": 421,
    "booth": 422,
    "signs": 423,
    "crosswalk": 424,
    "bedroom": 425,
    "zoo": 426,
    "tower": 427,
    "stall": 428,
    "kitty": 429,
}
