tasks_voc = {
    "offline": {
        0: list(range(21)),
    },
    "19-1": {
        0: list(range(20)),
        1: [20],
    },
    "15-5": {
        0: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15],
        1: [16, 17, 18, 19, 20]
    },
    "15-1":
        {
            0: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15],
            1: [16],
            2: [17],
            3: [18],
            4: [19],
            5: [20]
    },
    "10-1":
        {
            0: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10],
            1: [11],
            2: [12],
            3: [13],
            4: [14],
            5: [15],
            6: [16],
            7: [17],
            8: [18],
            9: [19],
            10: [20]
    },
    "5-3":
        {
            0: [0, 1, 2, 3, 4, 5],
            1: [6, 7, 8],
            2: [9, 10, 11],
            3: [12, 13, 14],
            4: [15, 16, 17],
            5: [18, 19, 20]
    },
    "5-1":
        {
            0: [0, 1, 2, 3, 4, 5],
            1: [6],
            2: [7],
            3: [8],
            4: [9],
            5: [10],
            6: [11],
            7: [12],
            8: [13],
            9: [14],
            10: [15],
            11: [16],
            12: [17],
            13: [18],
            14: [19],
            15: [20],
    },
    "15-1a":
        {
            0: [0, 12, 9, 20, 7, 15, 8, 14, 16, 5, 19, 4, 1, 13, 2, 11],
            1: [17], 2: [3], 3: [6], 4: [18], 5: [10]
        },
    "15-1b":
        {
            0: [0, 13, 19, 15, 17, 9, 8, 5, 20, 4, 3, 10, 11, 18, 16, 7],
            1: [12], 2: [14], 3: [6], 4: [1], 5: [2]
        },
    "15-1c":
        {
            0: [0, 15, 3, 2, 12, 14, 18, 20, 16, 11, 1, 19, 8, 10, 7, 17],
            1: [6], 2: [5], 3: [13], 4: [9], 5: [4]
        },
    "15-1d":
        {
            0: [0, 7, 5, 3, 9, 13, 12, 14, 19, 10, 2, 1, 4, 16, 8, 17],
            1: [15], 2: [18], 3: [6], 4: [11], 5: [20]
        },
    "15-1e":
        {
            0: [0, 7, 13, 5, 11, 9, 2, 15, 12, 14, 3, 20, 1, 16, 4, 18],
            1: [8], 2: [6], 3: [10], 4: [19], 5: [17]
        },
    "15-1f":
        {
            0: [0, 7, 5, 9, 1, 15, 18, 14, 3, 20, 10, 4, 19, 11, 17, 16],
            1: [12], 2: [8], 3: [6], 4: [2], 5: [13]
        },
    "15-1g":
        {
            0: [0, 12, 9, 19, 6, 4, 10, 5, 18, 14, 15, 16, 3, 8, 7, 11],
            1: [13], 2: [2], 3: [20], 4: [17], 5: [1]
        },
    "15-1h":
        {
            0: [0, 13, 10, 15, 8, 7, 19, 4, 3, 16, 12, 14, 11, 5, 20, 6],
            1: [2], 2: [18], 3: [9], 4: [17], 5: [1]
        },
    "15-1i":
        {
            0: [0, 1, 14, 9, 5, 2, 15, 8, 20, 6, 16, 18, 7, 11, 10, 19],
            1: [3], 2: [4], 3: [17], 4: [12], 5: [13]
        },
    "15-1j":
        {
            0: [0, 16, 13, 1, 11, 12, 18, 6, 14, 5, 3, 7, 9, 20, 19, 15],
            1: [4], 2: [2], 3: [10], 4: [8], 5: [17]
        }
}


tasks_cityscapes = {
    "offline": {0: list(range(20))},
    "14-1":  # 14-1 setting (6 tasks)
        {
            0: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14],
            1: [15],
            2: [16],
            3: [17],
            4: [18],
            5: [19]
        },
    "14-5":  # 14-5 setting (2 tasks)
        {
            0: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14],
            1: [15, 16, 17, 18, 19],
        },
    "18-1":
        {
            0: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18],
            1: [19]
        },
    "10-1":
        {
            0: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10],
            1: [11],
            2: [12],
            3: [13],
            4: [14],
            5: [15],
            6: [16],
            7: [17],
            8: [18],
            9: [19]
       },
}


tasks_ade = {
    "offline": {
        0: [x for x in range(151)]
    },
    "100-50": {
        0: [x for x in range(0, 101)],
        1: [x for x in range(101, 151)]
    },
    "100-10":
        {
            0: [x for x in range(0, 101)],
            1: [x for x in range(101, 111)],
            2: [x for x in range(111, 121)],
            3: [x for x in range(121, 131)],
            4: [x for x in range(131, 141)],
            5: [x for x in range(141, 151)]
        },
    "100-5":
        {
            0: [x for x in range(0, 101)],
            1: [x for x in range(101, 106)],
            2: [x for x in range(106, 111)],
            3: [x for x in range(111, 116)],
            4: [x for x in range(116, 121)],
            5: [x for x in range(121, 126)],
            6: [x for x in range(126, 131)],
            7: [x for x in range(131, 136)],
            8: [x for x in range(136, 141)],
            9: [x for x in range(141, 146)],
            10: [x for x in range(146, 151)]
        },
    "50":
        {
            0: [x for x in range(0, 51)],
            1: [x for x in range(51, 101)],
            2: [x for x in range(101, 151)]
        },
}


tasks_attr = {
    "19-1": {
            0: list(range(20)),
            1: [20],
        },
    "15-5": {
            0: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15],
            1: [16, 17, 18, 19, 20]
        },
    "15-1":{0: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15],
            1: [16],
            2: [17],
            3: [18],
            4: [19],
            5: [20]
        },
    "14-1":{0: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14],
            1: [15],
            2: [16],
            3: [17],
            4: [18],
            5: [19]
        },
    "5-3":{
            0: [0, 1, 2, 3, 4, 5],
            1: [6, 7, 8],
            2: [9, 10, 11],
            3: [12, 13, 14],
            4: [15, 16, 17],
            5: [18, 19, 20]
        },
    "10-1":{0: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10],
            1: [11],
            2: [12],
            3: [13],
            4: [14],
            5: [15],
            6: [16],
            7: [17],
            8: [18],
            9: [19],
            10: [20],
        },
    "100-10":
        {
            0: [x for x in range(0, 101)],
            1: [x for x in range(101, 111)],
            2: [x for x in range(111, 121)],
            3: [x for x in range(121, 131)],
            4: [x for x in range(131, 141)],
            5: [x for x in range(141, 151)]
        },
    "100-5":
        {
            0: [x for x in range(0, 101)],
            1: [x for x in range(101, 106)],
            2: [x for x in range(106, 111)],
            3: [x for x in range(111, 116)],
            4: [x for x in range(116, 121)],
            5: [x for x in range(121, 126)],
            6: [x for x in range(126, 131)],
            7: [x for x in range(131, 136)],
            8: [x for x in range(136, 141)],
            9: [x for x in range(141, 146)],
            10: [x for x in range(146, 151)]
        },
    "50":
        {
            0: [x for x in range(0, 51)],
            1: [x for x in range(51, 101)],
            2: [x for x in range(101, 151)]
        },
    "100-50": {
        0: [x for x in range(0, 101)],
        1: [x for x in range(101, 151)]
        },
    "15-1a":
        {
            0: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15],
            1: [16], 2: [17], 3: [18], 4: [19], 5: [20]
        },
    "15-1b":
        {
            0: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15],
            1: [16], 2: [17], 3: [18], 4: [19], 5: [20]
        },
    "15-1c":
        {
            0: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15],
            1: [16], 2: [17], 3: [18], 4: [19], 5: [20]
        },
    "15-1d":
        {
            0: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15],
            1: [16], 2: [17], 3: [18], 4: [19], 5: [20]
        },
    "15-1e":
        {
            0: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15],
            1: [16], 2: [17], 3: [18], 4: [19], 5: [20]
        },
    "15-1f":
        {
            0: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15],
            1: [16], 2: [17], 3: [18], 4: [19], 5: [20]
        },
    "15-1g":
        {
            0: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15],
            1: [16], 2: [17], 3: [18], 4: [19], 5: [20]
        },
    "15-1h":
        {
            0: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15],
            1: [16], 2: [17], 3: [18], 4: [19], 5: [20]
        },
    "15-1i":
        {
            0: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15],
            1: [16], 2: [17], 3: [18], 4: [19], 5: [20]
        },
    "15-1j":
        {
            0: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15],
            1: [16], 2: [17], 3: [18], 4: [19], 5: [20]
        }
}


def get_task_list():
    return list(tasks_cityscapes.keys()) + list(tasks_voc.keys()) + list(tasks_ade.keys())


def get_task_labels(dataset, name, step):
    print(dataset)
    if dataset == 'voc':
        task_dict = tasks_voc[name]
    elif dataset == 'ade':
        task_dict = tasks_ade[name]
    elif dataset == 'cityscapes':
        task_dict = tasks_cityscapes[name]
    else:
        raise NotImplementedError
    assert step in task_dict.keys(), f"You should provide a valid step! [{step} is out of range]"

    labels = list(task_dict[step])
    print('labels ',labels)
    labels_old = [label for s in range(step) for label in task_dict[s]]
    return labels, labels_old, f'data/{dataset}/{name}'


def get_per_task_classes(dataset, name, step):
    if dataset == 'voc':
        task_dict = tasks_voc[name]
    elif dataset == 'ade':
        task_dict = tasks_ade[name]
    elif dataset == 'cityscapes':
        task_dict = tasks_cityscapes[name]
    else:
        raise NotImplementedError
    assert step in task_dict.keys(), f"You should provide a valid step! [{step} is out of range]"

    classes = [len(task_dict[s]) for s in range(step + 1)]
    return classes


def get_task_labels_attr(dataset, name, step):
    print(dataset, name, step)
    labels = list(tasks_attr[name][step])
    return labels
