import torch
import numpy as np
import os
from torch.utils.data import TensorDataset, DataLoader
from .utils import collate_sequences


NORMALIZER = { # (mu, std) per class computed on the concatenation of both features (discarding the binary feature)
    'hot dog': (1.3554527691145501, 55.15028414343622),
    'palm tree': (-0.7063322505461493, 49.02700047706162),
    'moon': (0.8036297226693486, 41.345375756324735),
    'envelope': (4.900210171097034, 69.4196392054246),
    'dumbbell': (2.0407119932197504, 47.695996108391235),
    'microwave': (2.9699868328411974, 66.93104801889736),
    'onion': (1.2284401861051968, 45.21653229074296),
    'nail': (1.6172277953943177, 62.21706638258232),
    'paper clip': (0.6436449511025123, 33.32139677497804),
    'soccer ball': (0.6708907017116656, 39.52546034271634),
    'drill': (1.1185769827821401, 48.722276882610934),
    'telephone': (1.3498681969396034, 40.76261400934935),
    'airplane': (1.0251489388319772, 53.19602656498733),
    'dishwasher': (2.2191710394266084, 61.508456849155735),
    'chair': (4.016188671169509, 62.53028260788498),
    'grass': (3.376122464598659, 69.31003265138725),
    'rhinoceros': (1.2215264458448767, 48.80834840225656),
    'octopus': (0.8146148966002359, 40.89244147955804),
    'cloud': (0.35621641218733063, 29.409365110585483),
    'bicycle': (0.46389958129146036, 48.99489500128756),
    'swan': (1.049680167563788, 42.94216649535794),
    'picture frame': (1.5065118700885085, 63.253773000519175),
    'shorts': (1.9229859470161206, 52.22414095061445),
    'flying saucer': (1.5281540318557478, 55.091686319872025),
    'basketball': (1.6894072481767088, 57.72410547176846),
    'harp': (2.906289433329914, 78.44568624724737),
    'beard': (0.9775846803866044, 40.10763763299041),
    'binoculars': (0.808846681416587, 47.034367710374035),
    'tiger': (0.9438875155470355, 50.66921493109194),
    'book': (2.5133103399080188, 65.60820901501357),
    'scissors': (0.6647841622339, 40.07199951607956),
    'raccoon': (0.7915126973835728, 43.36239169880799),
    'peanut': (0.5509739234166906, 30.524261515788336),
    'wheel': (0.7686023820927692, 47.607169012136815),
    'trombone': (1.0112428613309314, 52.411164111718705),
    'diamond': (2.395434500084604, 59.73484161759297),
    'parachute': (2.056040072916103, 60.77205525434674),
    'tractor': (0.5855071624918279, 50.5522849539403),
    'windmill': (0.24800006974356498, 52.12209721342569),
    'alarm clock': (0.391438978240927, 41.44493991046053),
    'clarinet': (1.2795783017970905, 49.905620294236705),
    'spider': (0.6505395210719399, 51.18743252881025),
    'violin': (0.8724565090414226, 52.533813964768754),
    'clock': (1.6654012441543409, 37.33134444355261),
    'tent': (3.3092329281631137, 79.47572994387069),
    'belt': (2.3132169051670886, 71.13105919924993),
    'map': (2.6555302484638714, 61.11029370697819),
    'toe': (0.5471757615653022, 43.27192861865762),
    'bread': (1.3686935317654665, 47.6839114556787),
    'kangaroo': (0.6918159454175237, 35.99155678225584),
    'camera': (1.4527253130110975, 49.336211227235296),
    'duck': (1.598900790833744, 41.45077993986563),
    'lipstick': (0.41066758960159977, 41.786372987299615),
    'snowman': (0.14670998509400804, 31.624590642386174),
    'pickup truck': (1.6892820330935685, 54.644954488199524),
    'radio': (1.1157698308056494, 56.49502963911298),
    'truck': (1.814018865712332, 55.76992610437815),
    'train': (1.470463028668502, 77.63271694640828),
    'teapot': (0.9014336302825292, 37.48169241933444),
    'tree': (0.13337780967798976, 42.97342154517355),
    'hourglass': (2.39448903480218, 52.622226862393084),
    'eyeglasses': (1.379818588483046, 52.57994649640675),
    'church': (0.9630982672082059, 69.99862099910592),
    'submarine': (0.9138035290673335, 56.77613283220326),
    'couch': (2.283752727373058, 68.77383224311272),
    'umbrella': (0.5170775226020248, 47.83678678400117),
    'whale': (0.3497782843722267, 52.43513503159438),
    'cooler': (2.19778540728888, 61.74770130955316),
    'sword': (0.5910085971920176, 48.46440617862079),
    'table': (4.542251159698462, 87.48848948789511),
    'skull': (0.6570416475324352, 36.02607871443743),
    'house': (2.654399100012597, 66.81448281800678),
    'blackberry': (0.3455386008050086, 29.1600574174796),
    'bush': (0.7558370448198207, 41.04289142315455),
    'giraffe': (1.4011522715876905, 46.32335477059355),
    'rainbow': (4.702348561309779, 82.07143165031478),
    'yoga': (1.1423119096294918, 50.79902795898255),
    'mailbox': (0.3511077577743624, 55.61495444057362),
    'wristwatch': (1.0924273980760375, 49.96303380973288),
    'The Eiffel Tower': (1.2008260944995623, 73.04798400687072),
    'syringe': (1.6277984132013836, 56.22798342770764),
    'bulldozer': (1.060340370028316, 50.385030079706375),
    'door': (3.36173249421909, 66.2933191994613),
    'zebra': (1.132649710524639, 52.459089632246396),
    'beach': (2.263430578427388, 106.91064036288513),
    'crown': (0.7879512551564102, 50.206077053610386),
    'screwdriver': (1.1442268550285573, 46.07164154904856),
    'bear': (0.9395651847291722, 36.986932048274426),
    'sink': (0.9422909049727696, 47.54959424164138),
    'teddy-bear': (0.7913359738933313, 33.36477019938705),
    'square': (1.3275239412907043, 65.89863242901156),
    'cruise ship': (2.15931097599974, 78.36615495337965),
    'waterslide': (3.4486527614397833, 83.83125777723943),
    'elbow': (4.092940205383508, 61.758770494053785),
    'stereo': (1.8269654619223092, 58.475208066649714),
    'sweater': (1.067301637554828, 44.59577281486724),
    'bandage': (1.4032828202796717, 49.86169830574158),
    'bat': (0.8121797269039484, 37.69212883824029),
    'The Mona Lisa': (1.6676774598611082, 58.162407907625514),
    'sea turtle': (0.7386565039500725, 46.86228560536563),
    'butterfly': (0.4342721164650034, 37.484845221008726),
    'mosquito': (0.6493471316616555, 40.10938957349605),
    'tennis racquet': (0.015468185574502526, 62.24923783294656),
    'tornado': (0.7822439181013964, 45.2077352961338),
    'computer': (3.0811423630717107, 60.20403781306317),
    'bridge': (3.679091358194862, 120.07641800536442),
    'toothbrush': (1.2915788907996562, 53.22105425492547),
    'baseball bat': (0.410479892106175, 39.02003924116569),
    'bench': (4.462592927663926, 85.6302327587279),
    'finger': (-0.6637118888775841, 49.09874846625699),
    'canoe': (2.9733556427417493, 68.6835039501244),
    'baseball': (1.6959011615443051, 45.45130310748645),
    'circle': (-0.39852915378672893, 31.77419572565473),
    'banana': (1.3562427804512358, 42.94349204337924),
    'bathtub': (2.3570421946852544, 65.3192157735626),
    'axe': (1.0293065652442999, 54.84964062528346),
    'lantern': (1.1476541043730428, 53.67189040723054),
    'birthday cake': (0.15146259492252578, 55.89568012892327),
    'castle': (1.804799071214271, 63.20589225473029),
    'wine bottle': (1.034291851931799, 44.04598147244387),
    'ant': (0.9303194592264448, 34.43552547266363),
    'The Great Wall of China': (4.285709330181438, 131.23951199298298),
    'bee': (0.43095116254566934, 40.56855963179127),
    'apple': (-0.44780125973592305, 30.208033668691396),
    'arm': (1.7757119621507091, 46.692967793920644),
    'asparagus': (-0.2421902384249924, 48.97218720603324),
    'angel': (0.5489444327750316, 41.66381961171915),
    'cup': (2.673919370605991, 43.54406248924784),
    'carrot': (0.6945175048056408, 46.104020850556616),
    'bucket': (1.7396654767172537, 48.828570427954205),
    'animal migration': (2.6285542168388782, 61.28180224245095),
    'cell phone': (1.9267526020713346, 49.38973568488984),
    'van': (1.9173825658872794, 54.8721828825201),
    'dolphin': (0.9714616061928398, 42.83044052150523),
    'bowtie': (1.168151585565935, 37.61503592501492),
    'campfire': (0.2534712087647997, 42.286814756524535),
    'ceiling fan': (1.0603067359693852, 40.52738328774831),
    'boomerang': (0.5759666273292099, 39.78957492087158),
    'aircraft carrier': (1.5172469688772912, 78.7478229402662),
    'cactus': (-0.1281463623029328, 42.27573114632624),
    'cake': (0.31108565857076187, 56.322115673526696),
    'anvil': (1.471075424663743, 48.99321880248113),
    'toothpaste': (1.8461911264030182, 51.53740072123023),
    'swing set': (3.971684529151281, 98.99892200987023),
    'feather': (-0.42952206263561854, 53.55639949373167),
    'flashlight': (1.9317251715822668, 62.79624045193533),
    'garden hose': (1.5452934595615202, 53.713569777275175),
    'camel': (1.5165348305653266, 35.07846843003865),
    'calculator': (1.991161645112966, 50.142844727554575),
    'diving board': (1.7300484119947224, 75.61560569527323),
    'chandelier': (1.991040877029286, 50.65396442677625),
    'helmet': (1.9722019205320098, 45.87749730234627),
    'squirrel': (0.729042851521045, 35.3641639039348),
    'ambulance': (1.0598312283596059, 55.875842882074),
    'bottlecap': (1.5970585109209756, 40.01592713375047),
    'hospital': (1.7313904919786411, 72.37806984815816),
    'coffee cup': (1.32151623967879, 41.665383540075005),
    'watermelon': (1.8482342559051477, 59.31958622930048),
    'dresser': (2.396722754292761, 79.1225545952145),
    'bed': (2.588378888585306, 78.08870505568636),
    'bird': (1.5906829218142842, 41.059856184169284),
    'cookie': (0.7048879723978447, 34.29958258051739),
    'underwear': (3.027964069514147, 54.355597943207094),
    'drums': (1.1572575727426198, 54.68602043565278),
    'cat': (0.9684180598296738, 43.19493215282525),
    'calendar': (2.495118096854286, 82.1800159400022),
    'bracelet': (0.4661401948292038, 31.127130949231766),
    'eraser': (2.3490401085702235, 56.643670114244784),
    'dog': (0.8907946320439043, 38.505287852990726),
    'barn': (2.2770830828592583, 77.75086163641558),
    'spoon': (0.5421543550003102, 37.016180276322515),
    'sun': (-0.2008690561101928, 57.798300559005334),
    'toilet': (1.291036016063847, 40.288417166228925),
    'backpack': (1.3079276772602353, 46.33461078978928),
    'trumpet': (1.233316766684717, 47.840050217395266),
    'frying pan': (1.1317137754492954, 42.27197781360748),
    'blueberry': (0.3428165650102726, 29.923143234478975),
    'toaster': (1.3159036268033921, 59.46381954781093),
    'floor lamp': (-1.4045719348973986, 52.73112796615196),
    'crocodile': (1.2462846638010021, 51.83360295588419),
    'police car': (0.6314716475098945, 51.402397657264785),
    'cow': (0.6487350495428166, 44.82200063524666),
    'basket': (1.781348034990179, 61.40405101602184),
    'cello': (1.4380096549620986, 59.481368251629206),
    'golf club': (2.935274820103259, 47.944997493610416),
    'school bus': (1.3202131289388477, 61.70753264839142),
    'hockey puck': (0.725588239742589, 48.55963543134594),
    'fence': (3.8660243770815614, 92.36222788620427),
    'donut': (0.431402194475543, 32.222374599013726),
    'goatee': (1.2211961416317247, 39.81077215140121),
    'traffic light': (1.269260032432163, 44.30942006032888),
    'hamburger': (1.4103828007350085, 49.04022894395681),
    'ear': (1.9563928536834947, 34.3747704500531),
    'compass': (0.8636275744036599, 38.906947603746346),
    'broccoli': (-0.08805269427735608, 30.880695648320078),
    'skyscraper': (1.3477313197584702, 87.73974365488579),
    'fan': (0.5595090068208328, 42.26975493031441),
    'hot air balloon': (1.0010255829235684, 45.01559229242698),
    'mountain': (5.349497596465423, 69.73739652862577),
    'fork': (0.21995268515715857, 43.66291957421616),
    'face': (1.1847102417517064, 41.81747854722619),
    'crab': (0.5500211063457824, 48.30558365265961),
    'ice cream': (0.5645385757011395, 41.72357855932428),
    'foot': (1.6352285029716924, 40.86466847411941),
    'hat': (2.1269765754849357, 53.181061994837336),
    'candle': (-0.9566338163648833, 46.30537462785261),
    'flip flops': (1.1195172002513105, 45.28787295602699),
    'hammer': (0.40690889202283986, 45.31354440860368),
    'guitar': (0.9118308447368665, 58.627968076179016),
    'brain': (0.5667801625156502, 39.94893006675094),
    'stove': (1.2695451153311437, 56.13115551721316),
    'headphones': (1.7442579010033754, 38.05663003234409),
    'flamingo': (1.3329066566304946, 44.20478550977875),
    'flower': (0.014788800722293086, 28.686447255310085),
    'bus': (1.5110163683563511, 65.58525727312637),
    'hot tub': (0.9262199087425361, 63.37602990315963),
    'elephant': (1.0286360401485168, 42.29328387209706),
    'fire hydrant': (0.4353600099500871, 48.49174159770318),
    'laptop': (2.5420362830209355, 63.093568635534155),
    'leaf': (-0.07888685459428697, 51.531397540382116),
    'potato': (0.7248796777877287, 36.04373128693473),
    'hockey stick': (2.560198275283893, 47.75516557446046),
    'lighter': (-0.10645657100081249, 38.600148168238576),
    'hexagon': (2.7288170353096675, 50.79748328406929),
    'garden': (0.881398058547382, 59.301002560708866),
    'marker': (1.4847281646438588, 50.021490600998504),
    'keyboard': (2.8496015722739236, 81.38936435354776),
    'camouflage': (0.8524647599170719, 65.65432278791238),
    'knee': (5.3541853695693575, 60.225209719801974),
    'sheep': (1.2468686657122494, 35.19467195151128),
    'microphone': (0.3006266208385552, 46.783442715555715),
    'mushroom': (0.28405131561550195, 40.671965829362236),
    'light bulb': (0.3093205629583717, 35.25819445171456),
    'hand': (0.7429242999868996, 34.70475212985948),
    'key': (0.7406380633244096, 34.13758650534517),
    'house plant': (-0.4396176672108764, 40.515632771810296),
    'eye': (0.8606006296728399, 44.889207403048424),
    'matches': (0.3485948924638904, 47.42024782911991),
    'broom': (2.9233557704577193, 49.52062851559808),
    'knife': (1.4292202706092547, 51.01808033723662),
    'crayon': (1.467668727844571, 51.82316360295973),
    'ocean': (7.872452229036218, 89.99111246191521),
    'dragon': (0.8266093687680877, 49.41364315921484),
    'leg': (5.117580228531927, 54.01711580361819),
    'horse': (0.9246973774561026, 48.65827974249926),
    'zigzag': (9.770917367360767, 61.744673036996616),
    'car': (1.1106827823007763, 47.60058589694208),
    'grapes': (0.6046526027097275, 27.16306317679192),
    'lightning': (4.090431090680993, 57.03172069825947),
    'moustache': (1.7875824399413591, 37.731677498630745),
    'mouth': (2.76090978291076, 57.20373539326289),
    'vase': (0.5528729482101566, 36.996243257356014),
    'fish': (0.8878609523273818, 44.34932077221152),
    'string bean': (1.346485501392611, 54.7312484146683),
    'lighthouse': (0.4274423658693314, 75.81546755799378),
    'ladder': (5.90632648589332, 110.16555003310557),
    'television': (1.3151946885305383, 62.90537952277926),
    'helicopter': (0.7111156159770702, 56.6546344981718),
    'pillow': (2.0992806701392936, 55.274535278488294),
    'pencil': (2.0345830706124053, 62.90446034037889),
    'rollerskates': (2.0053135688983006, 39.31457668947572),
    'jail': (5.661515872939487, 115.47255551897983),
    'mermaid': (0.3187352763659362, 39.8221589482459),
    'jacket': (2.0929497013270537, 50.6087533539712),
    'megaphone': (1.8135468059177202, 54.66219701027781),
    'nose': (4.435118108240006, 36.01419720778613),
    'pants': (1.4927018991320877, 55.47801110072461),
    'octagon': (2.0144474110553916, 49.61164954802588),
    'pizza': (0.9106006910867426, 49.75334623210151),
    'passport': (2.09209268126368, 53.80930291521799),
    'pool': (2.06494328488252, 67.72608882496336),
    'motorbike': (0.4038001637130562, 46.94203574972685),
    'snake': (1.5154800788642753, 49.350623204522535),
    'pond': (0.7752730687547197, 47.62409950756826),
    'frog': (0.8874821595962438, 39.61840650901404),
    'pig': (0.47576581658267675, 39.5924494951546),
    'penguin': (1.0164857569517498, 40.88730060878002),
    'cannon': (0.8927868329478531, 53.019935221920896),
    'parrot': (1.6070485082553567, 43.38710309821747),
    'lobster': (0.5829596663716866, 42.78511651754868),
    'saw': (1.6178343188617499, 43.19552103419441),
    'strawberry': (0.6209739512011668, 32.08453043500838),
    'firetruck': (1.125536302973774, 65.91057171556372),
    'speedboat': (2.0848389958987257, 76.42986457816829),
    'popsicle': (0.4813981088599971, 42.49229183387311),
    'hurricane': (0.7079895622313991, 61.715710494552816),
    'see saw': (1.8249850934378673, 70.89383197689017),
    'saxophone': (0.9072383454069756, 36.470803735437904),
    'mug': (2.5296236017401257, 42.26283334121334),
    'piano': (2.6469411517060806, 73.27448246359215),
    'mouse': (0.8020204927469491, 43.836228689128035),
    'power outlet': (2.071476276483809, 46.822370189887785),
    'hedgehog': (0.4703510415238984, 45.92192258266138),
    'oven': (1.8548425634903463, 62.43067850281194),
    'shoe': (1.297356215372919, 41.93847714957883),
    'rifle': (2.5223233995449474, 60.73555429659974),
    'roller coaster': (2.6065332991832584, 86.95567387367467),
    'peas': (0.7749159834453123, 42.94847025647788),
    'lion': (0.4463371384240275, 34.510210963204415),
    'rake': (3.442498762575747, 57.38005406297777),
    'postcard': (3.700937086574, 69.8261672011201),
    'sock': (1.9223557134218592, 43.2687682421636),
    'purse': (1.6872172724499956, 48.85082993380252),
    'sleeping bag': (1.2484033851490541, 52.138238912603775),
    'skateboard': (2.4819607493229663, 53.19362309156825),
    'necklace': (2.392666309866489, 41.3064841582455),
    'stairs': (5.195938168639603, 47.15470516213574),
    'lollipop': (0.10920444361594842, 38.89025105370695),
    'snowflake': (2.3415722082063986, 68.96721342968107),
    'rabbit': (0.9078200152038035, 34.75862482451542),
    'owl': (1.2457620241823235, 42.73803624793326),
    'shovel': (1.970015817486029, 45.419236670608626),
    'pear': (-0.45220059964010495, 30.843347488001527),
    'remote control': (1.1358869210694837, 44.83511889242796),
    'star': (0.3626996748657054, 52.65011227641426),
    'scorpion': (0.4161827603069684, 38.88321413686467),
    'washing machine': (1.5518183007862645, 51.91417194144562),
    'monkey': (0.9748166731813579, 38.60787650590758),
    'pineapple': (0.562007915664679, 43.7000843939721),
    'sandwich': (1.6847535599541337, 57.542891294247035),
    'shark': (1.272828952833183, 49.334895742299615),
    'sailboat': (1.309450897368411, 66.09322028103158),
    'steak': (0.8908929135892587, 46.82398060648129),
    'stethoscope': (2.300526882061146, 43.63511505624273),
    'wine glass': (2.1753360642422095, 42.95333738304328),
    'smiley face': (1.4208837631558537, 43.864342591767816),
    'streetlight': (-1.4343035375659503, 57.70810758721286),
    'squiggle': (5.131557337201909, 48.02532522224354),
    'stop sign': (1.3327274061718097, 42.78360537094287),
    'line': (40.59167311123959, 112.02341955570965),
    'pliers': (0.796279030471497, 45.67250007650481),
    'paint can': (1.3512234721466652, 47.35796924253278),
    'panda': (0.5475608600999033, 33.69643785103632),
    'paintbrush': (0.20347385695100456, 47.341806442823824),
    't-shirt': (0.9831120778329658, 42.21114938247829),
    'fireplace': (1.3117628588460688, 61.01045131707993),
    'river': (5.610367142280469, 117.56790294876312),
    'snorkel': (1.2366543753832537, 43.709326082973476),
    'rain': (3.6461954118834403, 61.31247784406768),
    'triangle': (1.1218274781431306, 64.34926695455631),
    'suitcase': (1.9098774305372213, 57.805580971303506),
    'stitches': (4.142179481399166, 79.85573072340479),
    'tooth': (0.7350361072423909, 34.97655580063578),
    'snail': (0.3764966115255877, 34.91367713422217),
    'spreadsheet': (4.333452826793876, 134.8852997594341)
}

class QuickDrawDataset(torch.utils.data.Dataset):
    def __init__(self, data_dict, normalizers, task_vector=None):

        self.data_dict = data_dict
        self.normalizers = normalizers
        self.task_vector = task_vector

        self.patterns_per_class = [len(v) for k,v in self.data_dict.items()]

        self.min_class_id = min(list(self.data_dict.keys()))

    def __getitem__(self, idx):
        # select class based on idx
        class_id = None
        curr_idx = idx
        ppc =  [0] + self.patterns_per_class

        for i in range(1, len(ppc)):
            if curr_idx < ppc[i]:
                class_id = self.min_class_id + (i - 1)
                break
            elif curr_idx == ppc[i]:
                curr_idx -= ppc[i]
                class_id = self.min_class_id + i
                break
            else:
                curr_idx -= ppc[i]

        if class_id is None:
            raise IndexError('Out of range when indexing QuickDraw!')

        # normalize
        x_cur = torch.from_numpy(self.data_dict[class_id][curr_idx]).float() #/ self.normalizers[class_id][1]
        y_cur = torch.tensor(class_id).long()

        if self.task_vector is not None:
            x_cur = torch.cat((self.task_vector.unsqueeze(0).repeat(x_cur.size(0),1), x_cur), dim=1)

        return x_cur, y_cur

    def __len__(self):
        return sum(self.patterns_per_class)


class CLQuickDraw():
    def __init__(self, root, train_batch_size, test_batch_size,
                 len_task_vector=0, task_vector_at_test=False):

        self.root = root
        self.train_batch_size = train_batch_size
        self.test_batch_size = test_batch_size

        self.len_task_vector = len_task_vector
        self.task_vector_at_test = task_vector_at_test

        self.dataloaders = []

        self.current_class_id = 0

    def _load_data(self, classes):
        train_dict, test_dict, normalizer = {}, {}, {}
        for classname in classes:
            feature = np.load(os.path.join(self.root, f"{classname}.npz"), encoding='latin1', allow_pickle=True)
            train, test = feature['train'], feature['test'] # discard feature['valid'] because we don't need it
            train_dict[self.current_class_id] = train
            test_dict[self.current_class_id] = test
            normalizer[self.current_class_id] = NORMALIZER[classname]
            self.current_class_id += 1

        return train_dict, test_dict, normalizer

    def get_task_loaders(self, classes=None, task_id=None):

        if classes is not None:
            train, test, normalizer = self._load_data(classes)

            if self.len_task_vector > 0:
                task_vector = torch.zeros(self.len_task_vector).float()
                task_vector[len(self.dataloaders)] = 1.
            else:
                task_vector = None

            train_dataset = QuickDrawDataset(train, normalizer,
                                             task_vector=task_vector)
            test_dataset = QuickDrawDataset(test, normalizer,
                                            task_vector=task_vector if self.task_vector_at_test else None)

            train_batch_size = len(train_dataset) if self.train_batch_size == 0 else self.train_batch_size
            test_batch_size = len(test_dataset) if self.test_batch_size == 0 else self.test_batch_size

            train_loader = DataLoader(train_dataset, batch_size=train_batch_size, shuffle=True, drop_last=True,
                                      collate_fn=collate_sequences)
            test_loader = DataLoader(test_dataset, batch_size=test_batch_size, shuffle=False, drop_last=True,
                                     collate_fn=collate_sequences)

            self.dataloaders.append([train_loader, test_loader])

            return train_loader, test_loader

        elif task_id is not None:
            return self.dataloaders[task_id]