from abc import ABC, abstractmethod
from dataclasses import dataclass
from typing import Callable, Iterable, List, Union

from lm_eval.api.instance import Instance


class Filter(ABC):
    """
    Filter classes operate on a per-task level.
    They take all model outputs (`instance.resps` for all `task.instances`)
    across all instances of a task, and perform operations.
    In a single run, one can configure any number of separate filters or lists of filters.

    """

    def __init__(self, **kwargs) -> None:
        """
        Can define custom behavior here, if an individual instantiation of a Filter class should have state.
        """

    @abstractmethod
    def apply(self, resps: Union[List, Iterable], docs: List[dict]) -> Iterable:
        """
        Defines the operation to perform on a list of the `inst.resps` properties of `Instance` objects.
        Should return the list of (filtered) response lists *in the same order as they were input*, e.g.
        if pass in [<inst.resps for instance 0>, <inst.resps for instance 1>] should return
        [<filtered resps for instance 0>, <filtered resps for instance 1>]
        """
        return resps


@dataclass
class FilterEnsemble:
    """
    FilterEnsemble creates a pipeline applying multiple filters.
    Its intended usage is to stack multiple post-processing steps in order.
    `task.apply_filters` should use a list of FilterEnsemble classes that it stores, to apply each
    pipeline separately.
    """

    name: str
    filters: List[Callable[[], Filter]]

    def apply(self, instances: List[Instance]) -> None:
        resps, docs = zip(*((inst.resps, inst.doc) for inst in instances))
        resps, docs = list(resps), list(docs)
        # print("resps", resps)
        # print("docs", docs)
        num_return_sequences = len(resps[0][0])
        start_flag = True

        # if(self.name=="create_test"):
        #     # 要单独处理
        #     for f in self.filters:
        #         # resps: List[List[List[str]]] [Batch, 1, num_return_sequences]
        #         if start_flag:
        #             start_flag = False
        #             new_resps = []  # new_resps: List[List[List[str]]] [num_return_sequences, Batch, 1]
        #             for i in range(num_return_sequences):
        #                 resps_i = [[inst_resps[0][i]] for inst_resps in resps]  # resps_i: List[List[str]]
        #                 new_resps.append(resps_i)
        #             resps = new_resps
        #         new_resps = []
        #         for i in range(num_return_sequences):
        #             new_resps.append(list(f().apply(resps[i], docs)))  # 过滤失败了也有返回值, First 过滤器就是为了去除那个多余维度的
        #         resps = new_resps
        #     # 再重塑回去,转置形状
        #     ret_resps = []  # ret_resps: List[List[str]] [Batch,  num_return_sequences]
        #     for i in range(len(resps[0])):
        #         ret_resps.append([resps[j][i] for j in range(num_return_sequences)])
        #     resps = ret_resps
        #
        #     # add the end results after filtering to filtered_requests of their respective source instances.
        #     # has key `self.name`: each FilterEnsemble applied in a given run should use a different name.
        #     for inst, resp in zip(instances, resps):
        #         inst.filtered_resps[self.name] = resp
        #     return


        for f in self.filters:
            # apply filters in sequence
            if type(resps[0][0]) is str:
                resps = f().apply(resps, docs)  # resps: List[List[str]]->List[str]
            else:
                # resps: List[List[List[str]]] [Batch, 1, num_return_sequences]
                if start_flag:
                    start_flag = False
                    new_resps = []  # new_resps: List[List[List[str]]] [num_return_sequences, Batch, 1]
                    for i in range(num_return_sequences):
                        resps_i = [[inst_resps[0][i]] for inst_resps in resps]  # resps_i: List[List[str]] 代表全集的一组采样
                        new_resps.append(resps_i)
                    resps = new_resps
                new_resps = []
                for i in range(num_return_sequences):
                    new_resps.append(list(f().apply(resps[i], docs)))  # 过滤失败了也有返回值, First 过滤器就是为了去除那个多余维度的
                resps = new_resps
        # 再重塑回去,转置形状
        ret_resps = []  # ret_resps: List[List[str]] [Batch,  num_return_sequences]
        for i in range(len(resps[0])):
            ret_resps.append([resps[j][i] for j in range(num_return_sequences)])
        resps = ret_resps

        # add the end results after filtering to filtered_requests of their respective source instances.
        # has key `self.name`: each FilterEnsemble applied in a given run should use a different name.
        for inst, resp in zip(instances, resps):
            inst.filtered_resps[self.name] = resp
        # print("self.name", self.name)
        # print("resps", resps)
        # exit()
