"""Random Retriever."""

from typing import List, Optional

from tqdm import trange

from opencompass.openicl.icl_retriever import BaseRetriever
from opencompass.openicl.utils.logging import get_logger
from opencompass.registry import ICL_RETRIEVERS

logger = get_logger(__name__)


@ICL_RETRIEVERS.register_module()
class FixKRetriever(BaseRetriever):
    """Fix-K Retriever. Each in-context example of the test prompts is
    retrieved as the same K examples from the index set.

    Args:
        dataset (`BaseDataset`): Any BaseDataset instances.
            Attributes of ``reader``, ``train`` and ``test`` will be used.
        fix_id_list (List[int]): List of in-context example indices for every
            test prompts.
        ice_separator (`Optional[str]`): The separator between each in-context
            example template when origin `PromptTemplate` is provided. Defaults
            to '\n'.
        ice_eos_token (`Optional[str]`): The end of sentence token for
            in-context example template when origin `PromptTemplate` is
            provided. Defaults to '\n'.
        ice_num (`Optional[int]`): The number of in-context example template
            when origin `PromptTemplate` is provided. Defaults to 1.
    """

    def __init__(self,
                 dataset,
                 fix_id_list: List[int],
                 ice_separator: Optional[str] = '\n',
                 ice_eos_token: Optional[str] = '\n',
                 ice_num: Optional[int] = 1) -> None:
        super().__init__(dataset, ice_separator, ice_eos_token, ice_num)
        self.fix_id_list = fix_id_list

    def retrieve(self):
        """Retrieve the in-context example index for each test example."""
        num_idx = len(self.index_ds)
        for idx in self.fix_id_list:
            assert idx < num_idx, f'Index {idx} is out of range of {num_idx}'
        rtr_idx_list = []
        for _ in trange(len(self.test_ds), disable=not self.is_main_process):
            rtr_idx_list.append(self.fix_id_list)
        return rtr_idx_list