# coding=utf-8
# Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team.
# Copyright (c) 2018, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

""" GLUE processors and helpers """

import csv
import os
import logging

from .core import InputExample

logger = logging.getLogger(__name__)


class TaskType:
    REGRESSION = "REGRESSION"
    CLASSIFICATION = "CLASSIFICATION"


class DataProcessor(object):
    """Base class for data converters for sequence classification data sets."""

    def get_train_examples(self, data_dir):
        """Gets a collection of `InputExample`s for the train set."""
        raise NotImplementedError()

    def get_dev_examples(self, data_dir):
        """Gets a collection of `InputExample`s for the dev set."""
        raise NotImplementedError()

    def get_labels(self):
        """Gets the list of labels for this data set."""
        raise NotImplementedError()

    @classmethod
    def _read_tsv(cls, input_file, quotechar=None):
        """Reads a tab separated value file."""
        with open(input_file, "r", encoding='utf-8') as f:
            reader = csv.reader(f, delimiter="\t", quotechar=quotechar)
            lines = []
            for line in reader:
                lines.append(line)
            return lines


"""https://huggingface.co/transformers/_modules/transformers/data/processors/glue.html"""


class SstProcessor(DataProcessor):
    """Processor for the SST-2 data set (GLUE version)."""
    TASK_TYPE = TaskType.CLASSIFICATION

    def get_train_examples(self, data_dir):
        """See base class."""
        return self._create_examples(
            self._read_tsv(os.path.join(data_dir, "train.tsv")), "train")

    def get_dev_examples(self, data_dir):
        """See base class."""
        return self._create_examples(
            self._read_tsv(os.path.join(data_dir, "dev.tsv")), "dev")

    def get_test_examples(self, data_dir):
        """See base class."""
        return self._create_examples(
            self._read_tsv(os.path.join(data_dir, "test.tsv")), "test")

    def get_labels(self):
        """See base class."""
        return ["0", "1"]

    def _create_examples(self, lines, set_type):
        """Creates examples for the training and dev sets."""
        examples = []
        for (i, line) in enumerate(lines):
            if i == 0:
                continue
            guid = "%s-%s" % (set_type, i)
            if set_type == "test":
                text_a = line[1]
                label = "0"
            else:
                text_a = line[0]
                label = line[1]
            examples.append(
                InputExample(guid=guid, text_a=text_a, text_b=None, label=label))
        return examples



PROCESSORS = {
    "sst": SstProcessor,
}


DEFAULT_FOLDER_NAMES = {
    "sst": "SST-2",
}


class Task:
    def __init__(self, name, processor, data_dir):
        self.name = name
        self.processor = processor
        self.data_dir = data_dir
        self.task_type = processor.TASK_TYPE

    def get_train_examples(self):
        return self.processor.get_train_examples(self.data_dir)

    def get_dev_examples(self):
        return self.processor.get_dev_examples(self.data_dir)

    def get_test_examples(self):
        return self.processor.get_test_examples(self.data_dir)

    def get_labels(self):
        return self.processor.get_labels()


def get_task(task_name, data_dir):
    task_name = task_name.lower()
    task_processor = PROCESSORS[task_name]()
    if data_dir is None:
        data_dir = os.path.join(
            os.environ["GLUE_DIR"], DEFAULT_FOLDER_NAMES[task_name])
    return Task(task_name, task_processor, data_dir)
