import os
from typing import List

from mmengine.dataset import Compose
from torch.utils.data import Dataset

from opencompass.registry import DATASETS


@DATASETS.register_module()
class MMEDataset(Dataset):
    """Dataset to load MME dataset.

    Args:
        data_dir (str): The path of the dataset.
        pipeline (List[dict]): The data augmentation.
    """
    tasks = [
        'artwork', 'celebrity', 'code_reasoning', 'color',
        'commonsense_reasoning', 'count', 'existence', 'landmark',
        'numerical_calculation', 'OCR', 'position', 'posters', 'scene',
        'text_translation'
    ]
    sub_dir_name = ('images', 'questions_answers_YN')

    def __init__(self, data_dir: str, pipeline: List[dict]) -> None:
        self.pipeline = Compose(pipeline)
        self.load_data(data_dir)

    def load_data(self, data_dir: str):
        self.data_list = []
        image_dir, question_dir = self.sub_dir_name
        for task in self.tasks:
            if os.path.exists(os.path.join(data_dir, task, question_dir)):
                q_list = os.listdir(os.path.join(data_dir, task, question_dir))
                i_list = os.listdir(os.path.join(data_dir, task, image_dir))
                q_prefix = os.path.join(data_dir, task, question_dir)
                i_prefix = os.path.join(data_dir, task, image_dir)
            else:
                fn_list = os.listdir(os.path.join(data_dir, task))
                q_list = [fn for fn in fn_list if '.txt' in fn]
                i_list = [fn for fn in fn_list if fn not in q_list]
                q_prefix = i_prefix = os.path.join(data_dir, task)

            q_list.sort()
            i_list.sort()
            assert len(q_list) == len(i_list)
            for q_fn, i_fn in zip(q_list, i_list):
                assert q_fn.split('.')[0] == i_fn.split('.')[0]
                q_path = os.path.join(q_prefix, q_fn)
                image_path = os.path.join(i_prefix, i_fn)
                with open(q_path, 'r') as f:
                    q1, a1 = f.readline().strip().split('\t')
                    q2, a2 = f.readline().strip().split('\t')
                self.data_list.append({
                    'img_path': image_path,
                    'question': q1,
                    'answer': a1,
                    'task': task
                })
                self.data_list.append({
                    'img_path': image_path,
                    'question': q2,
                    'answer': a2,
                    'task': task
                })

    def __len__(self) -> None:
        return len(self.data_list)

    def __getitem__(self, idx: int) -> dict:
        data_sample = self.data_list[idx]
        data_sample = self.pipeline(data_sample)
        return data_sample