import os
from enum import Enum
from pathlib import Path

import yaml

ImagePath = {
    "ImageDatasetPath": "",
    
    "MME": "MME",
    "MMBench": "mmbench",
    "GQA": "GQA",
    "COCO": "MSCOCO/val2014",
    "VQAv2": "MSCOCO",
}
QuestionPath = {
    "QuestionDatasetPath": "",
    
    "MME": "MME",
    "MMBench": "mmbench",
    "GQA": "GQA",
    "COCO": "MSCOCO",
    "VQAv2": "VQAv2",
}

ModelPath = {
    "llava15_7B": "liuhaotian/llava-v1.5-7b",
    "llava15_13B": "liuhaotian/llava-v1.5-13b",
    "qwen_7B": "Qwen/Qwen-VL",
    "instructblip_7B": "vicuna7b",
    "instructblip_13B": "vicuna13b",
}

def get_image_path(dataset: str) -> Path:
    if dataset not in ImagePath.keys(): raise ValueError(f"Dataset {dataset} not found")
    image_path = os.path.join(ImagePath['ImageDatasetPath'], ImagePath[dataset])
    print(f"Image path: {image_path}")
    return image_path

def get_question_path(path_name: str, dataset_type: str) -> Path:
    if path_name in QuestionPath.keys():
        question_path_name = QuestionPath[path_name]
        if type(question_path_name) != str and dataset_type in question_path_name.keys():
            question_path_name = question_path_name[dataset_type]
        question_path = os.path.join(QuestionPath['QuestionDatasetPath'], question_path_name)
        print(f"Question path: {question_path}")
        return question_path
    else:
        return Path(path_name)

def get_model_path(model: str) -> str:
    if model not in ModelPath.keys(): return model
    return ModelPath[model]
    
def get_save_path(dataset: str, exp_name: str) -> Path:
    save_path = os.path.join(os.getcwd(), 'output', dataset, exp_name)
    return save_path