# Copyright (c) Alibaba, Inc. and its affiliates.

import os

# os.environ["CUDA_VISIBLE_DEVICES"] = "1"
import subprocess
import unittest
from evalscope.utils import test_level_list, is_module_installed
from evalscope.utils.logger import get_logger
from evalscope.run import run_task

logger = get_logger()


class TestCLIPBenchmark(unittest.TestCase):

    def setUp(self) -> None:
        self._check_env("webdataset")

    def tearDown(self) -> None:
        pass

    @staticmethod
    def _check_env(module_name: str):
        if is_module_installed(module_name):
            logger.info(f"{module_name} is installed.")
        else:
            raise ModuleNotFoundError(f"run: pip install {module_name}")

    @unittest.skipUnless(0 in test_level_list(), "skip test in current test level")
    def test_run_task(self):
        task_cfg = {
            "eval_backend": "RAGEval",
            "eval_config": {
                "tool": "clip_benchmark",
                "eval": {
                    "models": [
                        {
                            "model_name": "AI-ModelScope/chinese-clip-vit-large-patch14-336px",
                        }
                    ],
                    "dataset_name": ["muge", "mnist"],
                    "split": "test",
                    "batch_size": 128,
                    "num_workers": 1,
                    "verbose": True,
                    "skip_existing": False,
                    "output_dir": "outputs",
                    "cache_dir": "cache",
                    "limit": 1000,
                },
            },
        }

        run_task(task_cfg)

    @unittest.skipUnless(0 in test_level_list(), "skip test in current test level")
    def test_run_custom(self):
        task_cfg = {
            "eval_backend": "RAGEval",
            "eval_config": {
                "tool": "clip_benchmark",
                "eval": {
                    "models": [
                        {
                            "model_name": "AI-ModelScope/chinese-clip-vit-large-patch14-336px",
                        }
                    ],
                    "dataset_name": ["custom"],
                    "data_dir": "custom_eval/multimodal/text-image-retrieval",
                    "split": "test",
                    "batch_size": 128,
                    "num_workers": 1,
                    "verbose": True,
                    "skip_existing": False,
                    "limit": 1000,
                },
            },
        }

        run_task(task_cfg)


if __name__ == "__main__":
    unittest.main(buffer=False)
