import json


def parse_sthv2_splits(level):
    """Parse Something-Something dataset V2 into "train", "val" splits.

    Args:
        level (int): Directory level of data. 1 for the single-level directory,
            2 for the two-level directory.

    Returns:
        list: "train", "val", "test" splits of Something-Something V2 dataset.
    """
    # Read the annotations
    # yapf: disable
    class_index_file = 'data/sthv2/annotations/something-something-v2-labels.json'  # noqa
    # yapf: enable
    train_file = 'data/sthv2/annotations/something-something-v2-train.json'
    val_file = 'data/sthv2/annotations/something-something-v2-validation.json'
    test_file = 'data/sthv2/annotations/something-something-v2-test.json'

    with open(class_index_file, 'r') as fin:
        class_mapping = json.loads(fin.read())

    def line_to_map(item, test_mode=False):
        vid = item['id']
        vid = '/'.join(vid.split('/')[-level:])
        if test_mode:
            return vid
        else:
            template = item['template'].replace('[', '')
            template = template.replace(']', '')
            label = class_mapping[template]
            return vid, label

    with open(train_file, 'r') as fin:
        items = json.loads(fin.read())
        train_list = [line_to_map(item) for item in items]

    with open(val_file, 'r') as fin:
        items = json.loads(fin.read())
        val_list = [line_to_map(item) for item in items]

    with open(test_file, 'r') as fin:
        items = json.loads(fin.read())
        test_list = [line_to_map(item, test_mode=True) for item in items]
    return ((train_list, val_list, test_list), )