"""
Class to encapsulate the input data in an object. It has added functionality
of convert input text to bert tensors.

Bert Tensors creation steps:
-Tokenize the text
-Convert the sequence of tokens into numbers
-Pad the sequences so each one has the same length

The maxlength of the text is 512 tokens.


"""

from bert.tokenization.bert_tokenization import FullTokenizer
from tqdm import tqdm
import numpy as np
class PartisanDetectionData:
    DATA_COLUMN = "content"
    LABEL_COLUMN = "truth"

    def __init__(self, train, test, tokenizer: FullTokenizer, classes, max_seq_len=512):
        self.tokenizer = tokenizer
        self.max_seq_len = 0
        self.classes = classes

        ((self.train_x, self.train_y), (self.test_x, self.test_y)) = map(self._prepare, [train, test])

        print("max seq_len", self.max_seq_len)
        self.max_seq_len = min(self.max_seq_len, max_seq_len)
        self.train_x, self.test_x = map(self._pad, [self.train_x, self.test_x])

    def _prepare(self, df):
        x, y = [], []

        for _, row in tqdm(df.iterrows()):
            text, label = row[PartisanDetectionData.DATA_COLUMN], row[PartisanDetectionData.LABEL_COLUMN]
            tokens = self.tokenizer.tokenize(text) # tokenisation
            tokens = ["[CLS]"] + tokens + ["[SEP]"]
            token_ids = self.tokenizer.convert_tokens_to_ids(tokens) #Convert the sequence of tokens into numbers
            self.max_seq_len = max(self.max_seq_len, len(token_ids))
            x.append(token_ids)
            y.append(self.classes.index(label))

        return np.array(x), np.array(y)

    def _pad(self, ids):
        x = []
        for input_ids in ids:
            input_ids = input_ids[:min(len(input_ids), self.max_seq_len - 2)]
            input_ids = input_ids + [0] * (self.max_seq_len - len(input_ids))
            x.append(np.array(input_ids))
        return np.array(x)