from opencompass.openicl.icl_prompt_template import PromptTemplate
from opencompass.openicl.icl_retriever import ZeroRetriever
from opencompass.openicl.icl_inferencer import GenInferencer
from opencompass.datasets.NPHardEval import (
    hard_GCP_Dataset, hard_GCP_Evaluator,
    hard_TSP_Dataset, hard_TSP_Evaluator,
    hard_MSP_Dataset, hard_MSP_Evaluator,
    cmp_GCP_D_Dataset, cmp_GCP_D_Evaluator,
    cmp_TSP_D_Dataset, cmp_TSP_D_Evaluator,
    cmp_KSP_Dataset, cmp_KSP_Evaluator,
    p_BSP_Dataset, p_BSP_Evaluator,
    p_EDP_Dataset, p_EDP_Evaluator,
    p_SPP_Dataset, p_SPP_Evaluator,
)

NPHardEval_tasks = [
    ["hard_GCP", "GCP", hard_GCP_Dataset, hard_GCP_Evaluator],
    ["hard_TSP", "TSP", hard_TSP_Dataset, hard_TSP_Evaluator],
    ["hard_MSP", "MSP", hard_MSP_Dataset, hard_MSP_Evaluator],
    ["cmp_GCP_D", "GCP_Decision", cmp_GCP_D_Dataset, cmp_GCP_D_Evaluator],
    ["cmp_TSP_D", "TSP_Decision", cmp_TSP_D_Dataset, cmp_TSP_D_Evaluator],
    ["cmp_KSP", "KSP", cmp_KSP_Dataset, cmp_KSP_Evaluator],
    ["p_BSP", "BSP", p_BSP_Dataset, p_BSP_Evaluator],
    ["p_EDP", "EDP", p_EDP_Dataset, p_EDP_Evaluator],
    ["p_SPP", "SPP", p_SPP_Dataset, p_SPP_Evaluator],
]

NPHardEval_datasets = []
for name, path_name, dataset, evaluator in NPHardEval_tasks:
    NPHardEval_reader_cfg = dict(input_columns=["prompt", "level"], output_column="q")

    NPHardEval_infer_cfg = dict(
        ice_template=dict(
            type=PromptTemplate,
            template=dict(
                begin="</E>",
                round=[
                    dict(role="HUMAN", prompt="</E>{prompt}"),
                    dict(role="BOT", prompt=""),
                ],
            ),
            ice_token="</E>",
        ),
        retriever=dict(type=ZeroRetriever),
        inferencer=dict(type=GenInferencer),
    )

    NPHardEval_eval_cfg = dict(evaluator=dict(type=evaluator), pred_role="BOT")

    NPHardEval_datasets.append(
        dict(
            type=dataset,
            abbr=name,
            path=f"./data/NPHardEval/{path_name}/",
            reader_cfg=NPHardEval_reader_cfg,
            infer_cfg=NPHardEval_infer_cfg,
            eval_cfg=NPHardEval_eval_cfg,
        )
    )