import logging

import os

import pandas as pd

import torch

from torch.utils.data import TensorDataset

from tqdm import tqdm

from datareader.FeatureSetBuilder import FeatureSetBuilder

from datareader.abstract_data_reader import AbstractDataReader, InputExample

class InputFeatures(object):

    """A single set of features of data."""

    def __init__(self, input_ids, input_mask, segment_ids, label_id):

        self.input_ids = input_ids

        self.input_mask = input_mask

        self.segment_ids = segment_ids

        self.label_id = label_id

class ChunkedDataReader(AbstractDataReader):

    def __init__(self, config, tokenizer):

        self.tokenizer = tokenizer

        self.max_sequence_length = config['max_sequence_length']

        self.config = config

        self.train = None

        self.valid = None

        self.test = None

        self.num_sections = config['num_sections']

    @staticmethod

    def chunks(lst, n):

        """Yield successive n-sized chunks from lst."""

        for i in range(0, len(lst), n):

            yield lst[i:i + n]

    def build_fresh_dataset(self, dataset):

        logging.info("Building fresh dataset...")

        df = pd.read_csv(os.path.join(self.config['data_dir'], dataset))

        return self.build_fresh_dataset(df)

    def _convert_rows_to_list_of_feature(self, df):

        input_features = []

        for _, row in tqdm(df.iterrows(), total=df.shape[0]):

            text = row['text']

            lbl = row[self.config['target']]

            input_example = InputExample(None, text, None, self.config['target'])

            feature = self.convert_example_to_feature(input_example, lbl)

            input_features.append(feature)

        return input_features

    def build_fresh_dataset(self, dataset):

        df = pd.read_csv(os.path.join(self.config['data_dir'], dataset))

        return self.build_fresh_dataset_from_df(df)

    def build_fresh_dataset_from_df(self, df):

        logging.info("Building fresh dataset...")

        features = self._convert_rows_to_list_of_feature(df)

        # Now parse them out into the proper parts.

        all_input_ids = torch.tensor([f.input_ids for f in features], dtype=torch.long)

        all_input_mask = torch.tensor([f.input_mask for f in features], dtype=torch.long)

        all_segment_ids = torch.tensor([f.segment_ids for f in features], dtype=torch.long)

        all_label_ids = torch.tensor([f.label_id for f in features], dtype=torch.long)

        return TensorDataset(all_input_ids, all_input_mask, all_segment_ids, all_label_ids)

    def convert_example_to_feature(self, example, label):

        # create a new feature set builder for this example

        inputFeatureBuilder = FeatureSetBuilder(label)

        # tokenize the text into a list

        tokens_a = self.tokenizer.tokenize(example.text_a)

        # chunk the list of tokens

        generator = self.chunks(tokens_a, self.max_sequence_length - 2)

        for section in generator:

            # convert the section to a feature

            section_feature = self.convert_section_to_feature(section, label)

            inputFeatureBuilder.add(section_feature)

        inputFeatureBuilder.resize(self.num_sections, self.convert_section_to_feature([0], label))

        assert len(inputFeatureBuilder.get()) == self.num_sections

        # We return the builder

        input_ids = [feature.input_ids for feature in inputFeatureBuilder.features ]

        input_masks = [feature.input_mask for feature in inputFeatureBuilder.features]

        segment_ids = [feature.segment_ids for feature in inputFeatureBuilder.features]

        # Now create a new 'type' of inputfeature

        return InputFeatures(input_ids, input_masks, segment_ids, label)

    def convert_section_to_feature(self, tokens_a, label):

        # Truncate the section if needed

        if len(tokens_a) > (self.max_sequence_length - 2):

            tokens_a = tokens_a[-(self.max_sequence_length - 2):]

        tokens = []

        segment_ids = []

        tokens.append("[CLS]")

        segment_ids.append(0)

        for token in tokens_a:

            tokens.append(token)

            segment_ids.append(0)

        tokens.append("[SEP]")

        segment_ids.append(0)

        input_ids = self.tokenizer.convert_tokens_to_ids(tokens)

        # The mask has 1 for real tokens and 0 for padding tokens. Only real

        # tokens are attended to.

        input_mask = [1] * len(input_ids)

        # Zero-pad up to the sequence length.

        while len(input_ids) < self.max_sequence_length:

            input_ids.append(0)

            input_mask.append(0)

            segment_ids.append(0)

        assert len(input_ids) == self.max_sequence_length

        assert len(input_mask) == self.max_sequence_length

        assert len(segment_ids) == self.max_sequence_length

        return InputFeatures(input_ids, input_mask, segment_ids, label)