import pandas as pd
import math


def split_list(lst, n):
    """Split a list into n (roughly) equal-sized chunks"""
    chunk_size = math.ceil(len(lst) / n)  # integer division
    return [lst[i:i+chunk_size] for i in range(0, len(lst), chunk_size)]


def get_chunk(lst, n, k):
    chunks = split_list(lst, n)
    return chunks[k]


def load_dataset(dataset_name):
    if dataset_name == 'mmaad_base':
        dataset_path = 'data/mmaad_20240303_base.tsv'
    elif dataset_name == 'mmiasd_base':
        dataset_path = 'data/mmiasd_20240303_base.tsv'
    elif dataset_name == 'mmivqd_base':
        dataset_path = 'data/mmivqd_20240303_base.tsv'
    elif dataset_name == 'mmaad_option':
        dataset_path = 'data/mmaad_20240303_option.tsv'
    elif dataset_name == 'mmiasd_option':
        dataset_path = 'data/mmiasd_20240303_option.tsv'
    elif dataset_name == 'mmivqd_option':
        dataset_path = 'data/mmivqd_20240303_option.tsv'
    else:
        raise ValueError('Invalid dataset name: {}'.format(dataset_name))
    # huggingfaceのdatasetの形式でデータを読み込む
    # dataset = pd.read_table(dataset_path)
    dataset = pd.read_csv(dataset_path, sep='\t')
    dataset = get_chunk(dataset, 1, 0)
    return dataset
