import sys
import os
import fastText
import csv
import pandas as pd
from sklearn import preprocessing
import xml.etree.ElementTree as ET
from tqdm import tqdm_notebook as tqdm

#from pytorch_pretrained_bert.tokenization import BertTokenizer, WordpieceTokenizer
from pytorch_pretrained_bert import BertTokenizer
from pytorch_pretrained_bert.modeling import BertForPreTraining, BertPreTrainedModel, BertModel, BertConfig, BertForMaskedLM, BertForSequenceClassification
from pathlib import Path
import torch
import re
from torch import Tensor
from torch.nn import BCEWithLogitsLoss
from fastai.text import Tokenizer, Vocab
import pandas as pd
import collections
import os
import pdb
from tqdm import tqdm, trange
import sys
import random
import numpy as np
import apex
import nltk
import spacy
from nltk.tokenize.treebank import TreebankWordDetokenizer
import string
import argparse

from sklearn.model_selection import train_test_split
module_path = os.path.abspath(os.path.join('..'))
if module_path not in sys.path:
    sys.path.append(module_path)

from sklearn.metrics import roc_curve, auc


from torch.utils.data import TensorDataset, DataLoader, RandomSampler, SequentialSampler
from torch.utils.data.distributed import DistributedSampler
from pytorch_pretrained_bert.optimization import BertAdam

import logging
logging.basicConfig(format='%(asctime)s - %(levelname)s - %(name)s -   %(message)s',
                    datefmt='%m/%d/%Y %H:%M:%S',
                    level=logging.INFO)
logger = logging.getLogger(__name__)

model_state_dict = None
BERT_PRETRAINED_PATH = Path('/media/lfdharo/97481d74-4cb5-4983-9a69-a748c32711ba/Data/BERT/multi_cased_L-12_H-768_A-12/')
PYTORCH_PRETRAINED_BERT_CACHE = BERT_PRETRAINED_PATH/'cache/'
PYTORCH_PRETRAINED_BERT_CACHE.mkdir(exist_ok=True)

data_path = "../TASS2019/DATASETS/public_data/"
DATA_PATH=Path(data_path+'bert/')
DATA_PATH.mkdir(exist_ok=True)
PATH=Path(data_path + 'bert/tmp')
PATH.mkdir(exist_ok=True)


args = {
    "train_size": -1,
    "val_size": -1,
    "full_data_dir": data_path + '/bert',
    "data_dir": PATH,
    "task_name": "twitter_classification",
    "no_cuda": False,
    "bert_model": BERT_PRETRAINED_PATH,
    "output_dir": DATA_PATH/'output',
    "max_seq_length": 32,
    "do_train": True,
    "do_eval": True,
    "do_lower_case": False,
    "train_batch_size": 32,
    "eval_batch_size": 32,
    "learning_rate": 3e-5,
    "num_train_epochs": 10.0,
    "warmup_proportion": 0.1,
    "no_cuda": False,
    "local_rank": -1,
    "seed": 42,
    "gradient_accumulation_steps": 1,
    "optimize_on_cpu": True,
    "fp16": False,
    "loss_scale": 128
}


# Append script path to import path to locate tass_eval
sys.path.append(os.path.realpath(__file__))

# Import evalTask1 function fro tass_eval module
from tass_eval import evalTask1

# model = fastText.load_model('model.bin')


ALL_LANGUAGES = ['es', 'cr', 'mx', 'pe', 'uy']
LANGUAGE_CODE = ['uy']
CROSS_LINGUAL = [True]
bTestPhase = False  # If we are doing test, then concatenate train + dev, if not use dev as test
bReduceLabels = False
replace_reducedLabel_perLanguage = {'cr': 'NONE', 'es': 'NEU', 'mx': 'NEU', 'pe': 'NONE', 'uy': 'NEU'}

LABEL_ENCODER = preprocessing.LabelEncoder()
TERNARY_LABEL_ENCODER = preprocessing.LabelEncoder()


LABEL_ENCODER = preprocessing.LabelEncoder()
TERNARY_LABEL_ENCODER = preprocessing.LabelEncoder()

print("Loading Spacy Model")
lemmatizer = spacy.load("es_core_news_sm")  # GLOBAL to avoid loading the model several times

emoji_pattern = re.compile("[" 
                           u"\U0001F600-\U0001F64F"  # emoticons
                           u"\U0001F300-\U0001F5FF"  # symbols & pictographs
                           u"\U0001F680-\U0001F6FF"  # transport & map symbols
                           u"\U0001F1E0-\U0001F1FF"  # flags (iOS)
                           u"\U0001F910-\U0001F93F"  # emoticos (iOS)
                           u"\U00002702-\U000027B0"
                           u"\U000024C2-\U0001F251"
                           u"\U00002639-\U0000263B"  # more emoticons
                           "]+", flags=re.UNICODE)


def camel_case_split(identifier):
    clean_identifier = re.sub('[#]', '', identifier)
    matches = re.finditer(".+?(?:(?<=[a-z])(?=[A-Z])|(?<=[A-Z])(?=[A-Z][a-z])|$)", clean_identifier)
    return ' '.join([m.group(0) for m in matches])


def text_preprocessing(data):
    result = data
    result = [tweet.replace('\n', '').strip() for tweet in result]  # Newline and leading/trailing spaces
    result = [emoji_pattern.sub(r'', tweet) for tweet in result]
    result = [tweet.replace(u'\u2018', "'").replace(u'\u2019', "'") for tweet in result]  # Quotes replace by general
    result = [re.sub(r"\B#\w+", lambda m: camel_case_split(m.group(0)), tweet) for tweet in result]  # Hashtag
    result = [tweet.lower() for tweet in result]
    result = [re.sub(r"^.*http.*$", 'http', tweet) for tweet in result]  # Remove all http contents
    result = [re.sub(r"\B@\w+", 'username', tweet) for tweet in result]  # Remove all usernames
    result = [re.sub(r"(\w)(\1{2,})", r"\1", tweet) for tweet in result]  # Remove all letter repetitions
    result = [re.sub(r"[a-zA-Z]*jaj[a-zA-Z]*", 'jajaja', tweet) for tweet in result]  # Normalize laughs
    result = [re.sub(r"\d+", '', tweet) for tweet in result]  # Remove all numbers
    result = [tweet.translate(str.maketrans('', '', string.punctuation + '¡')) for tweet in result]  # Remove punctuation

    return result


def tokenize_list(datalist):
    print("Tokenizing")
    return [nltk.word_tokenize(row) for row in datalist]


def lemmatize_list(datalist):
    print("Lemmatizing data...")
    return [[token.lemma_ for token in lemmatizer(row)] for row in datalist]


def pipeline_processing_data(train_data, dev_data, tst_data):
    # PRE-PROCESSING
    preprocessed_train_content = text_preprocessing(train_data['content'])
    preprocessed_dev_content = text_preprocessing(dev_data['content'])
    preprocessed_tst_content = text_preprocessing(tst_data['content'])

    # # TOKENIZE
    # tokenized_train_content = tokenize_list(preprocessed_train_content)
    # tokenized_dev_content = tokenize_list(preprocessed_dev_content)
    # tokenized_tst_data = tokenize_list(preprocessed_tst_content)

    # clean_train_content = tokenized_train_content  # remove_stopwords(tokenized_train_content)
    # clean_dev_content = tokenized_dev_data  # remove_stopwords(tokenized_dev_data)
    # if bTestPhase is True:
    #     clean_tst_content = tokenized_tst_data  # remove_stopwords(tokenized_tst_data)

    # LIBRE OFFICE PROCESSING
    # libreoffice_train_tweets = [TreebankWordDetokenizer().detokenize(row)
    #                             for row in libreoffice_processing(tokenized_train_content)]
    # libreoffice_dev_tweets = [TreebankWordDetokenizer().detokenize(row)
    #                           for row in libreoffice_processing(tokenized_dev_data)]
    # libreoffice_tst_tweets = [TreebankWordDetokenizer().detokenize(row)
    #                           for row in libreoffice_processing(tokenized_tst_data)]

    # LEMMATIZING
    lemmatized_train_tweets = lemmatize_list(preprocessed_train_content)
    lemmatized_dev_tweets = lemmatize_list(preprocessed_dev_content)
    lemmatized_tst_tweets = lemmatize_list(preprocessed_tst_content)
    #
    # # REMOVING ACCENTS
    # without_accents_train = remove_accents(lemmatized_train_tweets)
    # without_accents_dev = remove_accents(lemmatized_dev_tweets)
    # if bTestPhase is True:
    #   without_accents_tst = remove_accents(lemmatized_tst_tweets)

    final_train_content = [TreebankWordDetokenizer().detokenize(row) for row in lemmatized_train_tweets]
    train_data['final_sentences'] = final_train_content

    final_dev_content = [TreebankWordDetokenizer().detokenize(row) for row in lemmatized_dev_tweets]
    dev_data['final_sentences'] = final_dev_content

    final_tst_content = [TreebankWordDetokenizer().detokenize(row) for row in lemmatized_tst_tweets]
    tst_data['final_sentences'] = final_tst_content

    return final_train_content, final_dev_content, final_tst_content


def read_files(sLang, bCross):
    train_data = pd.DataFrame()

    tree_tst = ET.parse(data_path + sLang + "/intertass_" + sLang + "_test.xml")
    tst_data = get_dataframe_from_xml(tree_tst)

    if bCross is True:
        dev_cross_data = pd.DataFrame()
        for sLangCross in [x for x in ALL_LANGUAGES if x != sLang]:
            tree_train_cross = ET.parse(data_path + sLangCross + "/intertass_" + sLangCross + "_train.xml")
            df_train_cross = get_dataframe_from_xml(tree_train_cross)

            tree_dev_cross = ET.parse(data_path + sLangCross + "/intertass_" + sLangCross + "_dev.xml")
            df_dev_cross = get_dataframe_from_xml(tree_dev_cross)

            train_data = pd.concat([train_data, df_train_cross], ignore_index=True)
            dev_cross_data = pd.concat([dev_cross_data, df_dev_cross], ignore_index=True)

        # We add train_cross + dev_cross (but never we can add dev_lang)
        train_data = pd.concat([train_data, dev_cross_data], ignore_index=True)

    else:
        tree_train = ET.parse(data_path + sLang + "/intertass_" + sLang + "_train.xml")
        train_data = get_dataframe_from_xml(tree_train)

    tree_dev = ET.parse(data_path + sLang + "/intertass_" + sLang + "_dev.xml")
    dev_data = get_dataframe_from_xml(tree_dev)

    return train_data, dev_data, tst_data


class BertForMultiLabelSequenceClassification(BertPreTrainedModel):
    """BERT model for classification.
    This module is composed of the BERT model with a linear layer on top of
    the pooled output.
    Params:
        `config`: a BertConfig class instance with the configuration to build a new model.
        `num_labels`: the number of classes for the classifier. Default = 2.
    Inputs:
        `input_ids`: a torch.LongTensor of shape [batch_size, sequence_length]
            with the word token indices in the vocabulary(see the tokens preprocessing logic in the scripts
            `extract_features.py`, `run_classifier.py` and `run_squad.py`)
        `token_type_ids`: an optional torch.LongTensor of shape [batch_size, sequence_length] with the token
            types indices selected in [0, 1]. Type 0 corresponds to a `sentence A` and type 1 corresponds to
            a `sentence B` token (see BERT paper for more details).
        `attention_mask`: an optional torch.LongTensor of shape [batch_size, sequence_length] with indices
            selected in [0, 1]. It's a mask to be used if the input sequence length is smaller than the max
            input sequence length in the current batch. It's the mask that we typically use for attention when
            a batch has varying length sentences.
        `labels`: labels for the classification output: torch.LongTensor of shape [batch_size]
            with indices selected in [0, ..., num_labels].
    Outputs:
        if `labels` is not `None`:
            Outputs the CrossEntropy classification loss of the output with the labels.
        if `labels` is `None`:
            Outputs the classification logits of shape [batch_size, num_labels].
    Example usage:
    ```python
    # Already been converted into WordPiece token ids
    input_ids = torch.LongTensor([[31, 51, 99], [15, 5, 0]])
    input_mask = torch.LongTensor([[1, 1, 1], [1, 1, 0]])
    token_type_ids = torch.LongTensor([[0, 0, 1], [0, 1, 0]])
    config = BertConfig(vocab_size_or_config_json_file=32000, hidden_size=768,
        num_hidden_layers=12, num_attention_heads=12, intermediate_size=3072)
    num_labels = 2
    model = BertForSequenceClassification(config, num_labels)
    logits = model(input_ids, token_type_ids, input_mask)
    ```
    """

    def __init__(self, config, num_labels=2):
        super(BertForMultiLabelSequenceClassification, self).__init__(config)
        self.num_labels = num_labels
        self.bert = BertModel(config)
        self.dropout = torch.nn.Dropout(config.hidden_dropout_prob)
        self.classifier = torch.nn.Linear(config.hidden_size, num_labels)
        self.apply(self.init_bert_weights)

    def forward(self, input_ids, token_type_ids=None, attention_mask=None, labels=None):
        _, pooled_output = self.bert(input_ids, token_type_ids, attention_mask, output_all_encoded_layers=False)
        pooled_output = self.dropout(pooled_output)
        logits = self.classifier(pooled_output)

        if labels is not None:
            loss_fct = BCEWithLogitsLoss()
            loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1, self.num_labels))
            # loss = loss_fct(logits.view(-1, 1), labels.view(-1, 1))
            return loss
        else:
            return logits

    def freeze_bert_encoder(self):
        for param in self.bert.parameters():
            param.requires_grad = False

    def unfreeze_bert_encoder(self):
        for param in self.bert.parameters():
            param.requires_grad = True


class InputExample(object):
    """A single training/test example for simple sequence classification."""

    def __init__(self, guid, text_a, text_b=None, labels=None):
        """Constructs a InputExample.

        Args:
            guid: Unique id for the example.
            text_a: string. The untokenized text of the first sequence. For single
            sequence tasks, only this sequence must be specified.
            text_b: (Optional) string. The untokenized text of the second sequence.
            Only must be specified for sequence pair tasks.
            labels: (Optional) [string]. The label of the example. This should be
            specified for train and dev examples, but not for test examples.
        """
        self.guid = guid
        self.text_a = text_a
        self.text_b = text_b
        self.labels = labels


class InputFeatures(object):
    """A single set of features of data."""

    def __init__(self, input_ids, input_mask, segment_ids, label_ids):
        self.input_ids = input_ids
        self.input_mask = input_mask
        self.segment_ids = segment_ids
        self.label_ids = label_ids


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

    def get_train_examples(self, data_file_name, size=-1):
        """Gets a collection of `InputExample`s for the train set."""
        raise NotImplementedError()

    def get_dev_examples(self, data_file_name, size=-1):
        """Gets a collection of `InputExample`s for the dev set."""
        raise NotImplementedError()

    def get_test_examples(self, data_file_name, size=-1):
        """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()


class MultiLabelTextProcessor(DataProcessor):

    def __init__(self, data_dir):
        self.data_dir = data_dir
        self.labels = None

    def get_train_examples(self, filename, size=-1):
    #def get_train_examples(self, data_df, size=-1):
        #filename = 'train.csv'
        logger.info("LOOKING AT {}".format(filename))
        if size == -1:
            data_df = pd.read_csv(filename, sep="\t", header=None)
            #             data_df['comment_text'] = data_df['comment_text'].apply(cleanHtml)
            return self._create_examples(data_df, "train")
        else:
            data_df = pd.read_csv(filename)
            #             data_df['comment_text'] = data_df['comment_text'].apply(cleanHtml)
            return self._create_examples(data_df.sample(size), "train")

    def get_dev_examples(self, filename, size=-1):
        # def get_train_examples(self, data_df, size=-1):
        # filename = 'train.csv'
        logger.info("LOOKING AT {}".format(filename))
        if size == -1:
            data_df = pd.read_csv(filename, sep="\t", header=None)
            #             data_df['comment_text'] = data_df['comment_text'].apply(cleanHtml)
            return self._create_examples(data_df, "dev")
        else:
            data_df = pd.read_csv(filename)
            #             data_df['comment_text'] = data_df['comment_text'].apply(cleanHtml)
            return self._create_examples(data_df.sample(size), "dev")

    def get_test_examples(self, filename, size=-1):
        logger.info("LOOKING AT {}".format(filename))
        if size == -1:
            data_df = pd.read_csv(filename, sep="\t", header=None)
            #             data_df['comment_text'] = data_df['comment_text'].apply(cleanHtml)
            return self._create_examples(data_df, "test")
        else:
            data_df = pd.read_csv(filename)
            #             data_df['comment_text'] = data_df['comment_text'].apply(cleanHtml)
            return self._create_examples(data_df.sample(size), "test")

    def get_labels(self):
        """See base class."""
        if self.labels == None:
            self.labels = list(pd.read_csv(os.path.join(self.data_dir, "classes.txt"), header=None)[0].values)
        return self.labels

    def _create_examples(self, df, set_type, labels_available=True):
        """Creates examples for the training and dev sets."""
        examples = []
        for (i, row) in enumerate(df.values):
            guid = row[0]
            text_a = row[3]
            if labels_available:
                if bReduceLabels is True:
                    if row[1] == 0:
                        labels = [1, 0, 0]
                    elif row[1] == 1:
                        labels = [0, 1, 0]
                    elif row[1] == 2:
                        labels = [0, 0, 1]
                else:
                    if row[1] == 0:
                        labels = [1, 0, 0 ,0]
                    elif row[1] == 1:
                        labels =[0, 1, 0, 0]
                    elif row[1] == 2:
                        labels =[0, 0, 1, 0]
                    elif row[1] == 3:
                        labels =[0, 0, 0, 1]
            else:
                labels = []
            examples.append(
                InputExample(guid=guid, text_a=text_a, labels=labels))
        return examples


def convert_examples_to_features(examples, label_list, max_seq_length, tokenizer):
    """Loads a data file into a list of `InputBatch`s."""

    label_map = {label: i for i, label in enumerate(label_list)}

    features = []
    for (ex_index, example) in enumerate(examples):
        tokens_a = tokenizer.tokenize(example.text_a)

        tokens_b = None
        if example.text_b:
            tokens_b = tokenizer.tokenize(example.text_b)
            # Modifies `tokens_a` and `tokens_b` in place so that the total
            # length is less than the specified length.
            # Account for [CLS], [SEP], [SEP] with "- 3"
            _truncate_seq_pair(tokens_a, tokens_b, max_seq_length - 3)
        else:
            # Account for [CLS] and [SEP] with "- 2"
            if len(tokens_a) > max_seq_length - 2:
                tokens_a = tokens_a[:(max_seq_length - 2)]

        # The convention in BERT is:
        # (a) For sequence pairs:
        #  tokens:   [CLS] is this jack ##son ##ville ? [SEP] no it is not . [SEP]
        #  type_ids: 0   0  0    0    0     0       0 0    1  1  1  1   1 1
        # (b) For single sequences:
        #  tokens:   [CLS] the dog is hairy . [SEP]
        #  type_ids: 0   0   0   0  0     0 0
        #
        # Where "type_ids" are used to indicate whether this is the first
        # sequence or the second sequence. The embedding vectors for `type=0` and
        # `type=1` were learned during pre-training and are added to the wordpiece
        # embedding vector (and position vector). This is not *strictly* necessary
        # since the [SEP] token unambigiously separates the sequences, but it makes
        # it easier for the model to learn the concept of sequences.
        #
        # For classification tasks, the first vector (corresponding to [CLS]) is
        # used as as the "sentence vector". Note that this only makes sense because
        # the entire model is fine-tuned.
        tokens = ["[CLS]"] + tokens_a + ["[SEP]"]
        segment_ids = [0] * len(tokens)

        if tokens_b:
            tokens += tokens_b + ["[SEP]"]
            segment_ids += [1] * (len(tokens_b) + 1)

        input_ids = 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.
        padding = [0] * (max_seq_length - len(input_ids))
        input_ids += padding
        input_mask += padding
        segment_ids += padding

        assert len(input_ids) == max_seq_length
        assert len(input_mask) == max_seq_length
        assert len(segment_ids) == max_seq_length

        labels_ids = []
        for label in example.labels:
            labels_ids.append(float(label))

        #         label_id = label_map[example.label]
        if ex_index < 0:
            logger.info("*** Example ***")
            logger.info("guid: %s" % (example.guid))
            logger.info("tokens: %s" % " ".join(
                [str(x) for x in tokens]))
            logger.info("input_ids: %s" % " ".join([str(x) for x in input_ids]))
            logger.info("input_mask: %s" % " ".join([str(x) for x in input_mask]))
            logger.info(
                "segment_ids: %s" % " ".join([str(x) for x in segment_ids]))
            logger.info("label: %s (id = %s)" % (example.labels, labels_ids))

        features.append(
            InputFeatures(input_ids=input_ids,
                          input_mask=input_mask,
                          segment_ids=segment_ids,
                          label_ids=labels_ids))
    return features


def _truncate_seq_pair(tokens_a, tokens_b, max_length):
    """Truncates a sequence pair in place to the maximum length."""

    # This is a simple heuristic which will always truncate the longer sequence
    # one token at a time. This makes more sense than truncating an equal percent
    # of tokens from each, since if one sequence is very short then each token
    # that's truncated likely contains more information than a longer sequence.
    while True:
        total_length = len(tokens_a) + len(tokens_b)
        if total_length <= max_length:
            break
        if len(tokens_a) > len(tokens_b):
            tokens_a.pop()
        else:
            tokens_b.pop()


def accuracy(out, labels):
    outputs = np.argmax(out, axis=1)
    return np.sum(outputs == labels)


def accuracy_thresh(y_pred:Tensor, y_true:Tensor, thresh:float=0.5, sigmoid:bool=True):
    "Compute accuracy when `y_pred` and `y_true` are the same size."
    if sigmoid: y_pred = y_pred.sigmoid()
#     return ((y_pred>thresh)==y_true.byte()).float().mean().item()
    return np.mean(((y_pred>thresh)==y_true.byte()).float().cpu().numpy(), axis=1).sum()


def fbeta(y_pred:Tensor, y_true:Tensor, thresh:float=0.2, beta:float=2, eps:float=1e-9, sigmoid:bool=True):
    "Computes the f_beta between `preds` and `targets`"
    beta2 = beta ** 2
    if sigmoid: y_pred = y_pred.sigmoid()
    y_pred = (y_pred>thresh).float()
    y_true = y_true.float()
    TP = (y_pred*y_true).sum(dim=1)
    prec = TP/(y_pred.sum(dim=1)+eps)
    rec = TP/(y_true.sum(dim=1)+eps)
    res = (prec*rec)/(prec*beta2+rec+eps)*(1+beta2)
    return res.mean().item()


def warmup_linear(x, warmup=0.002):
    if x < warmup:
        return x/warmup
    return 1.0 - x


def get_dataframe_from_xml(data):
    print("Preparing data...")
    tweet_id, user, content, day_of_week, month, hour, lang, sentiment, ternary_sentiment = [], [], [], [], [], [], [], [], []
    for tweet in data.iter('tweet'):
        for element in tweet.iter():
            if element.tag == 'tweetid':
                tweet_id.append(element.text)
            elif element.tag == 'user':
                user.append(element.text)
            elif element.tag == 'content':
                content.append(element.text)
            elif element.tag == 'date':
                day_of_week.append(element.text[:3])
                month.append(element.text[4:7])
                hour.append(element.text[11:13])
            elif element.tag == 'lang':
                lang.append(element.text)
            elif element.tag == 'value':
                sentiment.append(element.text)
                if element.text == 'NONE' or element.text == 'NEU':
                    ternary_sentiment.append('N-N')
                else:
                    ternary_sentiment.append(element.text)

    result_df = pd.DataFrame()
    result_df['tweet_id'] = tweet_id
    # result_df['user'] = user
    result_df['content'] = content
    # result_df['lang'] = lang
    result_df['day_of_week'] = day_of_week
    result_df['month'] = month
    result_df['hour'] = hour

    if bReduceLabels is True:
        sentiment = ['O' if (x == 'NEU' or x == 'NONE') else x for x in sentiment]

    LABEL_ENCODER.fit(sentiment)
    TERNARY_LABEL_ENCODER.fit(ternary_sentiment)
    result_df['sentiment'] = LABEL_ENCODER.transform(sentiment)
    result_df['ternary_sentiment'] = TERNARY_LABEL_ENCODER.transform(ternary_sentiment)
    return result_df


from torch.optim.lr_scheduler import _LRScheduler, Optimizer

class CyclicLR(object):
    """Sets the learning rate of each parameter group according to
    cyclical learning rate policy (CLR). The policy cycles the learning
    rate between two boundaries with a constant frequency, as detailed in
    the paper `Cyclical Learning Rates for Training Neural Networks`_.
    The distance between the two boundaries can be scaled on a per-iteration
    or per-cycle basis.
    Cyclical learning rate policy changes the learning rate after every batch.
    `batch_step` should be called after a batch has been used for training.
    To resume training, save `last_batch_iteration` and use it to instantiate `CycleLR`.
    This class has three built-in policies, as put forth in the paper:
    "triangular":
        A basic triangular cycle w/ no amplitude scaling.
    "triangular2":
        A basic triangular cycle that scales initial amplitude by half each cycle.
    "exp_range":
        A cycle that scales initial amplitude by gamma**(cycle iterations) at each
        cycle iteration.
    This implementation was adapted from the github repo: `bckenstler/CLR`_
    Args:
        optimizer (Optimizer): Wrapped optimizer.
        base_lr (float or list): Initial learning rate which is the
            lower boundary in the cycle for eachparam groups.
            Default: 0.001
        max_lr (float or list): Upper boundaries in the cycle for
            each parameter group. Functionally,
            it defines the cycle amplitude (max_lr - base_lr).
            The lr at any cycle is the sum of base_lr
            and some scaling of the amplitude; therefore
            max_lr may not actually be reached depending on
            scaling function. Default: 0.006
        step_size (int): Number of training iterations per
            half cycle. Authors suggest setting step_size
            2-8 x training iterations in epoch. Default: 2000
        mode (str): One of {triangular, triangular2, exp_range}.
            Values correspond to policies detailed above.
            If scale_fn is not None, this argument is ignored.
            Default: 'triangular'
        gamma (float): Constant in 'exp_range' scaling function:
            gamma**(cycle iterations)
            Default: 1.0
        scale_fn (function): Custom scaling policy defined by a single
            argument lambda function, where
            0 <= scale_fn(x) <= 1 for all x >= 0.
            mode paramater is ignored
            Default: None
        scale_mode (str): {'cycle', 'iterations'}.
            Defines whether scale_fn is evaluated on
            cycle number or cycle iterations (training
            iterations since start of cycle).
            Default: 'cycle'
        last_batch_iteration (int): The index of the last batch. Default: -1
    Example:
        >>> optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9)
        >>> scheduler = torch.optim.CyclicLR(optimizer)
        >>> data_loader = torch.utils.data.DataLoader(...)
        >>> for epoch in range(10):
        >>>     for batch in data_loader:
        >>>         scheduler.batch_step()
        >>>         train_batch(...)
    .. _Cyclical Learning Rates for Training Neural Networks: https://arxiv.org/abs/1506.01186
    .. _bckenstler/CLR: https://github.com/bckenstler/CLR
    """

    def __init__(self, optimizer, base_lr=1e-3, max_lr=6e-3,
                 step_size=2000, mode='triangular', gamma=1.,
                 scale_fn=None, scale_mode='cycle', last_batch_iteration=-1):

#         if not isinstance(optimizer, Optimizer):
#             raise TypeError('{} is not an Optimizer'.format(
#                 type(optimizer).__name__))
        self.optimizer = optimizer

        if isinstance(base_lr, list) or isinstance(base_lr, tuple):
            if len(base_lr) != len(optimizer.param_groups):
                raise ValueError("expected {} base_lr, got {}".format(
                    len(optimizer.param_groups), len(base_lr)))
            self.base_lrs = list(base_lr)
        else:
            self.base_lrs = [base_lr] * len(optimizer.param_groups)

        if isinstance(max_lr, list) or isinstance(max_lr, tuple):
            if len(max_lr) != len(optimizer.param_groups):
                raise ValueError("expected {} max_lr, got {}".format(
                    len(optimizer.param_groups), len(max_lr)))
            self.max_lrs = list(max_lr)
        else:
            self.max_lrs = [max_lr] * len(optimizer.param_groups)

        self.step_size = step_size

        if mode not in ['triangular', 'triangular2', 'exp_range'] \
                and scale_fn is None:
            raise ValueError('mode is invalid and scale_fn is None')

        self.mode = mode
        self.gamma = gamma

        if scale_fn is None:
            if self.mode == 'triangular':
                self.scale_fn = self._triangular_scale_fn
                self.scale_mode = 'cycle'
            elif self.mode == 'triangular2':
                self.scale_fn = self._triangular2_scale_fn
                self.scale_mode = 'cycle'
            elif self.mode == 'exp_range':
                self.scale_fn = self._exp_range_scale_fn
                self.scale_mode = 'iterations'
        else:
            self.scale_fn = scale_fn
            self.scale_mode = scale_mode

        self.batch_step(last_batch_iteration + 1)
        self.last_batch_iteration = last_batch_iteration

    def batch_step(self, batch_iteration=None):
        if batch_iteration is None:
            batch_iteration = self.last_batch_iteration + 1
        self.last_batch_iteration = batch_iteration
        for param_group, lr in zip(self.optimizer.param_groups, self.get_lr()):
            param_group['lr'] = lr

    def _triangular_scale_fn(self, x):
        return 1.

    def _triangular2_scale_fn(self, x):
        return 1 / (2. ** (x - 1))

    def _exp_range_scale_fn(self, x):
        return self.gamma**(x)

    def get_lr(self):
        step_size = float(self.step_size)
        cycle = np.floor(1 + self.last_batch_iteration / (2 * step_size))
        x = np.abs(self.last_batch_iteration / step_size - 2 * cycle + 1)

        lrs = []
        param_lrs = zip(self.optimizer.param_groups, self.base_lrs, self.max_lrs)
        for param_group, base_lr, max_lr in param_lrs:
            base_height = (max_lr - base_lr) * np.maximum(0, (1 - x))
            if self.scale_mode == 'cycle':
                lr = base_lr + base_height * self.scale_fn(cycle)
            else:
                lr = base_lr + base_height * self.scale_fn(self.last_batch_iteration)
            lrs.append(lr)
        return lrs


def eval():
    args['output_dir'].mkdir(exist_ok=True)

    eval_features = convert_examples_to_features(
        eval_examples, label_list, args['max_seq_length'], tokenizer)
    logger.info("***** Running evaluation *****")
    logger.info("  Num examples = %d", len(eval_examples))
    logger.info("  Batch size = %d", args['eval_batch_size'])
    all_input_ids = torch.tensor([f.input_ids for f in eval_features], dtype=torch.long)
    all_input_mask = torch.tensor([f.input_mask for f in eval_features], dtype=torch.long)
    all_segment_ids = torch.tensor([f.segment_ids for f in eval_features], dtype=torch.long)
    all_label_ids = torch.tensor([f.label_ids for f in eval_features], dtype=torch.float)
    eval_data = TensorDataset(all_input_ids, all_input_mask, all_segment_ids, all_label_ids)
    # Run prediction for full data
    eval_sampler = SequentialSampler(eval_data)
    eval_dataloader = DataLoader(eval_data, sampler=eval_sampler, batch_size=args['eval_batch_size'])

    all_logits = None
    all_labels = None

    model.eval()
    eval_loss, eval_accuracy = 0, 0
    nb_eval_steps, nb_eval_examples = 0, 0
    for input_ids, input_mask, segment_ids, label_ids in eval_dataloader:
        input_ids = input_ids.to(device)
        input_mask = input_mask.to(device)
        segment_ids = segment_ids.to(device)
        label_ids = label_ids.to(device)

        with torch.no_grad():
            tmp_eval_loss = model(input_ids, segment_ids, input_mask, label_ids)
            logits = model(input_ids, segment_ids, input_mask)

        #         logits = logits.detach().cpu().numpy()
        #         label_ids = label_ids.to('cpu').numpy()
        #         tmp_eval_accuracy = accuracy(logits, label_ids)
        tmp_eval_accuracy = accuracy_thresh(logits, label_ids)
        if all_logits is None:
            all_logits = logits.detach().cpu().numpy()
        else:
            all_logits = np.concatenate((all_logits, logits.detach().cpu().numpy()), axis=0)

        if all_labels is None:
            all_labels = label_ids.detach().cpu().numpy()
        else:
            all_labels = np.concatenate((all_labels, label_ids.detach().cpu().numpy()), axis=0)

        eval_loss += tmp_eval_loss.mean().item()
        eval_accuracy += tmp_eval_accuracy

        nb_eval_examples += input_ids.size(0)
        nb_eval_steps += 1

    eval_loss = eval_loss / nb_eval_steps
    eval_accuracy = eval_accuracy / nb_eval_examples

    #     ROC-AUC calcualation
    # Compute ROC curve and ROC area for each class
    fpr = dict()
    tpr = dict()
    roc_auc = dict()

    for i in range(num_labels):
        fpr[i], tpr[i], _ = roc_curve(all_labels[:, i], all_logits[:, i])
        roc_auc[i] = auc(fpr[i], tpr[i])

    # Compute micro-average ROC curve and ROC area
    fpr["micro"], tpr["micro"], _ = roc_curve(all_labels.ravel(), all_logits.ravel())
    roc_auc["micro"] = auc(fpr["micro"], tpr["micro"])

    result = {'eval_loss': eval_loss,
              'eval_accuracy': eval_accuracy,
              #               'loss': tr_loss/nb_tr_steps,
              'roc_auc': roc_auc}

    output_eval_file = os.path.join(args['output_dir'], "eval_results.txt")
    with open(output_eval_file, "w") as writer:
        logger.info("***** Eval results *****")
        for key in sorted(result.keys()):
            logger.info("  %s = %s", key, str(result[key]))
    #             writer.write("%s = %s\n" % (key, str(result[key])))
    return result


def fit(num_epocs=args['num_train_epochs']):
    global_step = 0
    model.train()
    for i_ in tqdm(range(int(num_epocs)), desc="Epoch"):

        tr_loss = 0
        nb_tr_examples, nb_tr_steps = 0, 0
        for step, batch in enumerate(tqdm(train_dataloader, desc="Iteration")):

            batch = tuple(t.to(device) for t in batch)
            input_ids, input_mask, segment_ids, label_ids = batch
            loss = model(input_ids, segment_ids, input_mask, label_ids)
            if n_gpu > 1:
                loss = loss.mean() # mean() to average on multi-gpu.
            if args['gradient_accumulation_steps'] > 1:
                loss = loss / args['gradient_accumulation_steps']

            if args['fp16']:
                optimizer.backward(loss)
            else:
                loss.backward()

            tr_loss += loss.item()
            nb_tr_examples += input_ids.size(0)
            nb_tr_steps += 1
            if (step + 1) % args['gradient_accumulation_steps'] == 0:
    #             scheduler.batch_step()
                # modify learning rate with special warm up BERT uses
                lr_this_step = args['learning_rate'] * warmup_linear(global_step/t_total, args['warmup_proportion'])
                for param_group in optimizer.param_groups:
                    param_group['lr'] = lr_this_step
                optimizer.step()
                optimizer.zero_grad()
                global_step += 1

        logger.info('Loss after epoc {}'.format(tr_loss / nb_tr_steps))
        logger.info('Eval after epoc {}'.format(i_+1))
        eval()


def predict(model, path, filename):
    predict_processor = MultiLabelTextProcessor(path)
    # test_examples = predict_processor.get_test_examples(filename=args['full_data_dir'] + "/intertass_" + sLang + "_test.bert"
    #                                                     , size=-1)
    test_examples = predict_processor.get_dev_examples(filename, size=-1)


    # Hold input data for returning it
    input_data = [{'id': input_example.guid, 'comment_text': input_example.text_a} for input_example in test_examples]

    test_features = convert_examples_to_features(
        test_examples, label_list, args['max_seq_length'], tokenizer)

    logger.info("***** Running prediction *****")
    logger.info("  Num examples = %d", len(test_examples))
    logger.info("  Batch size = %d", args['eval_batch_size'])

    all_input_ids = torch.tensor([f.input_ids for f in test_features], dtype=torch.long)
    all_input_mask = torch.tensor([f.input_mask for f in test_features], dtype=torch.long)
    all_segment_ids = torch.tensor([f.segment_ids for f in test_features], dtype=torch.long)

    test_data = TensorDataset(all_input_ids, all_input_mask, all_segment_ids)

    # Run prediction for full data
    test_sampler = SequentialSampler(test_data)
    test_dataloader = DataLoader(test_data, sampler=test_sampler, batch_size=args['eval_batch_size'])

    all_logits = None

    model.eval()
    eval_loss, eval_accuracy = 0, 0
    nb_eval_steps, nb_eval_examples = 0, 0
    for step, batch in enumerate(tqdm(test_dataloader, desc="Prediction Iteration")):
        input_ids, input_mask, segment_ids = batch
        input_ids = input_ids.to(device)
        input_mask = input_mask.to(device)
        segment_ids = segment_ids.to(device)

        with torch.no_grad():
            logits = model(input_ids, segment_ids, input_mask)
            logits = logits.sigmoid()

        if all_logits is None:
            all_logits = logits.detach().cpu().numpy()
        else:
            all_logits = np.concatenate((all_logits, logits.detach().cpu().numpy()), axis=0)

        nb_eval_examples += input_ids.size(0)
        nb_eval_steps += 1

    return pd.merge(pd.DataFrame(input_data), pd.DataFrame(all_logits, columns=label_list), left_index=True,
                    right_index=True)


def write_for_bert(train_data, dev_data, tst_data, sLang, bCross):
    suffix = ''
    if bCross is True:
        suffix = '_cross'

    if bReduceLabels is True:
        suffix += '_reduced'
    else:
        suffix += '_full'

    filename_out = data_path + sLang + "/intertass_" + sLang + suffix + "_train.bert"
    if os.path.exists(filename_out) is False:
        with open(filename_out, 'w', encoding="utf-8") as f_out:
            print('Writing train file for Bert at: ' + filename_out)
            for tid, txt, lbl in zip(train_data['tweet_id'], train_data['final_sentences'], train_data['sentiment']):
                f_out.write(str(tid) + '\t' + str(lbl) + '\ta\t' + txt + '\n')
    else:
        print('Skipping ' + filename_out)

    filename_out = data_path + sLang + "/intertass_" + sLang + suffix + "_dev.bert"
    if os.path.exists(filename_out) is False:
        with open(filename_out , 'w', encoding="utf-8") as f_out:
            print('Writing dev file for Bert at: ' + filename_out)
            for tid, txt, lbl in zip(dev_data['tweet_id'], dev_data['final_sentences'], dev_data['sentiment']):
                f_out.write(str(tid) + '\t' + str(lbl) + '\ta\t' + txt + '\n')
    else:
        print('Skipping ' + filename_out)

    filename_out = data_path + sLang + "/intertass_" + sLang + suffix + "_test.bert"
    if os.path.exists(filename_out) is False:
        with open(filename_out, 'w', encoding="utf-8") as f_out:
            print('Writing tst file for Bert at: ' + filename_out)
            for tid, txt, lbl in zip(tst_data['tweet_id'], tst_data['final_sentences'], tst_data['sentiment']):
                f_out.write(str(tid) + '\t' + str(lbl) + '\ta\t' + txt + '\n')
    else:
        print('Skipping ' + filename_out)


if __name__ == '__main__':

    parser = argparse.ArgumentParser()
    parser.add_argument("-t", "--btest", help="if we want to work for test", action="store_true")
    parser.add_argument("-r", "--breduced", help="if we want to use reduced labels", action="store_true")
    parser.add_argument("-c", "--bcross", help="if we want to make the cross data", action="store_true")
    parser.add_argument("-l", "--lang", help="Language to process", required=True)
    cmd_args = parser.parse_args()

    if cmd_args.btest:
        bTestPhase = True
    if cmd_args.breduced:
        bReduceLabels = True
    if cmd_args.bcross:
        CROSS_LINGUAL = [True]
    else:
        CROSS_LINGUAL = [False]

    LANGUAGE_CODE = [cmd_args.lang]  # we can only process one language at time because of CUDA

    # GET THE DATA
    output_dir = data_path + "final_outputs/"
    for bCross in CROSS_LINGUAL:
        print('----------------- CROSS : ' + str(bCross) + ' ----------------')

        prefix = ''
        if bCross is True:
            output_dir = output_dir + '/cross/'
            prefix = '_cross'
        else:
            output_dir = output_dir + '/mono/'

        for sLang in LANGUAGE_CODE:
            print('** LANG: ' + sLang)

            output_dir = data_path + "final_outputs/"
            prefix = ''
            if bCross is True:
                output_dir = output_dir + '/cross/'
                prefix = '_cross'
            else:
                output_dir = output_dir + '/mono/'

            train_data, dev_data, tst_data = read_files(sLang, bCross)
            final_train_content, final_dev_content, final_tst_content = pipeline_processing_data(train_data, dev_data,
                                                                                                 tst_data)

            write_for_bert(train_data, dev_data, tst_data, sLang, bCross)
            if bReduceLabels is True:
                prefix += '_reduced'
                train_filename = data_path + sLang + "/intertass_" + sLang + prefix + "_train.bert"
                dev_filename = data_path + sLang + "/intertass_" + sLang + prefix + "_dev.bert"
                tst_filename = data_path + sLang + "/intertass_" + sLang + prefix + "_test.bert"
                label_list = ['N', 'O', 'P']
            else:
                prefix += '_full'
                train_filename = data_path + sLang + "/intertass_" + sLang + prefix + "_train.bert"
                dev_filename = data_path + sLang + "/intertass_" + sLang + prefix + "_dev.bert"
                tst_filename = data_path + sLang + "/intertass_" + sLang + prefix + "_test.bert"
                label_list = ['N', 'NEU', 'NONE', 'P']

            if bTestPhase is True:  # Combine train + dev
                import subprocess
                subprocess.call("cat " + train_filename + ' ' + dev_filename + " > " + data_path + sLang + "/intertass_" + sLang + prefix + "_train_dev.bert", shell=True)
                train_filename = data_path + sLang + "/intertass_" + sLang + prefix + "_train_dev.bert"

            processors = {
                "twitter_classification": MultiLabelTextProcessor
            }

            # Setup GPU parameters

            if args["local_rank"] == -1 or args["no_cuda"]:
                device = torch.device("cuda" if torch.cuda.is_available() and not args["no_cuda"] else "cpu")
                n_gpu = torch.cuda.device_count()
            #     n_gpu = 1
            else:
                torch.cuda.set_device(args['local_rank'])
                device = torch.device("cuda", args['local_rank'])
                n_gpu = 1
                # Initializes the distributed backend which will take care of sychronizing nodes/GPUs
                torch.distributed.init_process_group(backend='nccl')
            logger.info("device: {} n_gpu: {}, distributed training: {}, 16-bits training: {}".format(
                device, n_gpu, bool(args['local_rank'] != -1), args['fp16']))

            args['train_batch_size'] = int(args['train_batch_size'] / args['gradient_accumulation_steps'])

            random.seed(args['seed'])
            np.random.seed(args['seed'])
            torch.manual_seed(args['seed'])
            if n_gpu > 0:
                torch.cuda.manual_seed_all(args['seed'])

            task_name = args['task_name'].lower()

            if task_name not in processors:
                raise ValueError("Task not found: %s" % (task_name))

            processor = processors[task_name](args['data_dir'])
            # label_list = processor.get_labels()
            # label_list = train_data['sentiment'].tolist()
            # label_list = ['N', 'NEU', 'NONE', 'P']
            num_labels = len(label_list)

            tokenizer = BertTokenizer.from_pretrained(args['bert_model'], do_lower_case=args['do_lower_case'])

            train_examples = None
            num_train_steps = None
            if args['do_train']:
                # train_examples = processor.get_train_examples(args['full_data_dir'], size=args['train_size'])
                train_examples = processor.get_train_examples(filename=train_filename, size=args['train_size'])

                #     train_examples = processor.get_train_examples(args['data_dir'], size=args['train_size'])
                num_train_steps = int(
                    len(train_examples) / args['train_batch_size'] / args['gradient_accumulation_steps'] * args[
                        'num_train_epochs'])


            # Prepare model
            def get_model():
                #     pdb.set_trace()
                if model_state_dict:
                    model = BertForMultiLabelSequenceClassification.from_pretrained(args['bert_model'],
                                                                                    num_labels=num_labels,
                                                                                    state_dict=model_state_dict)
                else:
                    model = BertForMultiLabelSequenceClassification.from_pretrained(args['bert_model'],
                                                                                    num_labels=num_labels)
                return model


            model = get_model()

            if args['fp16']:
                model.half()
            model.to(device)
            if args['local_rank'] != -1:
                try:
                    from apex.parallel import DistributedDataParallel as DDP
                except ImportError:
                    raise ImportError(
                        "Please install apex from https://www.github.com/nvidia/apex to use distributed and fp16 training.")

                model = DDP(model)
            elif n_gpu > 1:
                model = torch.nn.DataParallel(model)

            # Prepare optimizer
            param_optimizer = list(model.named_parameters())
            no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight']
            optimizer_grouped_parameters = [
                {'params': [p for n, p in param_optimizer if not any(nd in n for nd in no_decay)],
                 'weight_decay': 0.01},
                {'params': [p for n, p in param_optimizer if any(nd in n for nd in no_decay)], 'weight_decay': 0.0}
            ]
            t_total = num_train_steps
            if args['local_rank'] != -1:
                t_total = t_total // torch.distributed.get_world_size()
            if args['fp16']:
                try:
                    from apex.optimizers import FP16_Optimizer
                    from apex.optimizers import FusedAdam
                except ImportError:
                    raise ImportError(
                        "Please install apex from https://www.github.com/nvidia/apex to use distributed and fp16 training.")

                optimizer = FusedAdam(optimizer_grouped_parameters,
                                      lr=args['learning_rate'],
                                      bias_correction=False,
                                      max_grad_norm=1.0)
                if args['loss_scale'] == 0:
                    optimizer = FP16_Optimizer(optimizer, dynamic_loss_scale=True)
                else:
                    optimizer = FP16_Optimizer(optimizer, static_loss_scale=args['loss_scale'])

            else:
                optimizer = BertAdam(optimizer_grouped_parameters,
                                     lr=args['learning_rate'],
                                     warmup=args['warmup_proportion'],
                                     t_total=t_total)

            scheduler = CyclicLR(optimizer, base_lr=2e-5, max_lr=5e-5, step_size=2500, last_batch_iteration=0)

            # Eval Fn
            eval_examples = processor.get_dev_examples(filename=dev_filename, size=args['val_size'])
            train_features = convert_examples_to_features(
                train_examples, label_list, args['max_seq_length'], tokenizer)

            logger.info("***** Running training *****")
            logger.info("  Num examples = %d", len(train_examples))
            logger.info("  Batch size = %d", args['train_batch_size'])
            logger.info("  Num steps = %d", num_train_steps)
            all_input_ids = torch.tensor([f.input_ids for f in train_features], dtype=torch.long)
            all_input_mask = torch.tensor([f.input_mask for f in train_features], dtype=torch.long)
            all_segment_ids = torch.tensor([f.segment_ids for f in train_features], dtype=torch.long)
            all_label_ids = torch.tensor([f.label_ids for f in train_features], dtype=torch.float)
            train_data = TensorDataset(all_input_ids, all_input_mask, all_segment_ids, all_label_ids)
            if args['local_rank'] == -1:
                train_sampler = RandomSampler(train_data)
            else:
                train_sampler = DistributedSampler(train_data)
            train_dataloader = DataLoader(train_data, sampler=train_sampler, batch_size=args['train_batch_size'])

            model.unfreeze_bert_encoder()
            fit()
            # Save a trained model
            model_to_save = model.module if hasattr(model, 'module') else model  # Only save the model it-self
            output_model_file = os.path.join(PYTORCH_PRETRAINED_BERT_CACHE, "finetuned_pytorch_model.bin")
            torch.save(model_to_save.state_dict(), output_model_file)

            # Load a trained model that you have fine-tuned
            model_state_dict = torch.load(output_model_file)
            model = BertForMultiLabelSequenceClassification.from_pretrained(args['bert_model'], num_labels=num_labels,
                                                                            state_dict=model_state_dict)
            model.to(device)

            eval()
            result = predict(model, DATA_PATH, filename=dev_filename)
            result.shape

            if bTestPhase is False:
                if os.path.exists(output_dir + sLang + "_dev_bert" + prefix + ".tsv") is False:
                    print("Writing " + output_dir + sLang + "_dev_bert" + prefix + ".tsv")
                    result.drop('comment_text', axis=1).to_csv(output_dir + sLang + "_dev_bert" + prefix + ".csv")

                    # with open(output_dir + sLang + "_dev_bert" + prefix + ".tsv", 'w', newline='') as out_file, \
                    #      open(output_dir + sLang + "_dev_bert"+ prefix + ".csv", 'w', newline='') as out_csv_file:
                    with open(output_dir + sLang + "_dev_bert" + prefix + ".tsv", 'w', newline='') as out_file:
                        tsv_writer = csv.writer(out_file, delimiter='\t')
                        # csv_writer = csv.writer(out_csv_file)
                        # csv_writer.writerow(['ID'] + label_list)

                        for row in result.iterrows():
                            # find the column with the largest value
                            lbl = row[1].drop(['id', 'comment_text']).astype(float).idxmax(axis=0)
                            if bReduceLabels is True:
                                if lbl == 'O':
                                    lbl = replace_reducedLabel_perLanguage[sLang]
                            tsv_writer.writerow(
                                [row[1]['id'], lbl])
                else:
                    print("Skipping " + output_dir + sLang + "_dev_bert" + prefix + ".tsv")

            run_file = output_dir + sLang + "_dev_bert" + prefix + ".tsv"
            gold_file = data_path + sLang + "/intertass_" + sLang + "_dev_gold.tsv"

            scores = evalTask1(gold_file, run_file)
            with open(output_dir + prefix + '_' + sLang + "_bert.res", 'w', newline='') as out_file:
                print("f1_score: %f\n" % scores['maf1'])
                out_file.write("f1_score: %f\n" % scores['maf1'])
                print("precision: %f\n" % scores['map'])
                out_file.write("precision: %f\n" % scores['map'])
                print("recall: %f\n" % scores['mar'])
                out_file.write("recall: %f\n" % scores['mar'])
                print("%f\t%f\t%f\n" % (scores['map'], scores['mar'], scores['maf1']))

            if bTestPhase is True:
                if os.path.exists(output_dir + sLang + "_test_bert" + prefix + ".tsv") is False:
                    result = predict(model, DATA_PATH, filename=tst_filename)
                    print("Writing " + output_dir + sLang + "_test_bert" + prefix + ".tsv")
                    result.drop('comment_text', axis=1).to_csv(output_dir + sLang + "_test_bert" + prefix + ".csv")

                    # with open(output_dir + sLang + "_dev_bert" + prefix + ".tsv", 'w', newline='') as out_file, \
                    #      open(output_dir + sLang + "_dev_bert"+ prefix + ".csv", 'w', newline='') as out_csv_file:
                    with open(output_dir + sLang + "_test_bert" + prefix + ".tsv", 'w', newline='') as out_file:
                        tsv_writer = csv.writer(out_file, delimiter='\t')
                        # csv_writer = csv.writer(out_csv_file)
                        # csv_writer.writerow(['ID'] + label_list)

                        for row in result.iterrows():
                            # find the column with the largest value
                            lbl = row[1].drop(['id', 'comment_text']).astype(float).idxmax(axis=0)
                            if bReduceLabels is True:
                                if lbl == 'O':
                                    lbl = replace_reducedLabel_perLanguage[sLang]
                            tsv_writer.writerow(
                                [row[1]['id'], lbl])

                else:
                    print("Skipping " + output_dir + sLang + "_test_bert" + prefix + ".tsv")

            torch.cuda.empty_cache()