import evaluate
from evaluate.evaluation_suite import SubTask


class Suite(evaluate.EvaluationSuite):

    def __init__(self, name):
        super().__init__(name)

"""
{
          "data": "glue",
          "name": "cola",
          "split": "test[:10]",
          "args_for_task": {
            "metric": "accuracy",
            "input_column": "sentence",
            "label_column": "label",
            "label_mapping": {
              "LABEL_0": 0.0,
              "LABEL_1": 1.0
            }
          }
        },
        {
          "data": "glue",
          "name": "sst2",
          "split": "validation[:10]",
          "args_for_task": {
            "metric": "accuracy",
            "input_column": "sentence",
            "label_column": "label",
            "label_mapping": {
              "LABEL_0": 0.0,
              "LABEL_1": 1.0
            }
          }
        },
        {
          "data": "glue",
          "name": "mnli",
          "split": "validation_mismatched[:10]",
          "args_for_task": {
            "metric": "accuracy",
            "input_column": "premise",
            "second_input_column": "hypothesis",
            "label_mapping": {
              "LABEL_0": 0,
              "LABEL_1": 1,
              "LABEL_2": 2
            },
            "label_column": "label"
          }
        },
        {
          "data": "glue",
          "name": "mrpc",
          "split": "validation[:10]",
          "args_for_task": {
            "metric": "accuracy",
            "input_column": "sentence1",
            "second_input_column": "sentence2",
            "label_mapping": {
              "LABEL_0": 0,
              "LABEL_1": 1
            },
            "label_column": "label"
          }
        },
        {
          "data": "glue",
          "name": "qqp",
          "split": "validation[:10]",
          "args_for_task": {
            "metric": "accuracy",
            "input_column": "question1",
            "second_input_column": "question2",
            "label_mapping": {
              "LABEL_0": 0,
              "LABEL_1": 1
            },
            "label_column": "label"
          }
        },
        {
          "data": "glue",
          "name": "qnli",
          "split": "validation[:10]",
          "args_for_task": {
            "metric": "accuracy",
            "input_column": "question",
            "second_input_column": "sentence",
            "label_mapping": {
              "LABEL_0": 0,
              "LABEL_1": 1
            },
            "label_column": "label"
          }
        },
        {
          "data": "glue",
          "name": "rte",
          "split": "validation[:10]",
          "args_for_task": {
            "metric": "accuracy",
            "input_column": "sentence1",
            "second_input_column": "sentence2",
            "label_mapping": {
              "LABEL_0": 0,
              "LABEL_1": 1
            },
            "label_column": "label"
          }
        },
        {
          "data": "glue",
          "name": "wnli",
          "split": "validation[:10]",
          "args_for_task": {
            "metric": "accuracy",
            "input_column": "sentence1",
            "second_input_column": "sentence2",
            "label_mapping": {
              "LABEL_0": 0,
              "LABEL_1": 1
            },
            "label_column": "label"
          }
        }
      ]
    }
"""
    def setup(self):
        self.preprocessor = lambda x: {"text": x["text"].lower()}
        self.suite = [
            SubTask(
                task_type="text-classification",
                data="glue",
                subset="wnli",
                split="validation[:10]",
                args_for_task={
                    "metric": "accuracy",
                    "input_column": "sentence1",
                    "second_input_column": "sentence2",
                    "label_column": "label",
                    "label_mapping": {
                        "LABEL_0": 0,
                        "LABEL_1": 1
                    }
                }
            )
        ]