from __future__ import absolute_import, division, print_function

import csv
import logging
import os
import sys
from io import open

from scipy.stats import pearsonr, spearmanr
from sklearn.metrics import matthews_corrcoef, f1_score

from multiprocessing import Pool, cpu_count
from tqdm import tqdm
    
from transformers import BertPreTrainedModel, BertModel
import torch
from torch import nn

logger = logging.getLogger(__name__)
csv.field_size_limit(2147483647)

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

    def __init__(self, guid, text_a, text_b=None, label=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.
            label: (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.label = label


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 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-sig") as f:
            reader = csv.reader(f, delimiter="\t", quotechar=quotechar)
            lines = []
            for line in reader:
                if sys.version_info[0] == 2:
                    line = list(unicode(cell, 'utf-8') for cell in line)
                lines.append(line)
            return lines


class SameSideProcessor(DataProcessor):
    """Processor for the sameside data set"""
    
    def __init__(self, trainset, devset):
        self.trainset = trainset
        self.devset = devset        

    def get_train_examples(self, data_dir):
        """See base class."""
        return self._create_examples(
            self.trainset, "train")

    def get_dev_examples(self, data_dir):
        """See base class."""
        return self._create_examples(
            self.devset, "dev")

    def get_labels(self):
        """See base class."""
        return [False, True]

    def _create_examples(self, items, set_type):
        """Creates examples for the training and dev sets."""
        examples = []
        
        for (i, item) in enumerate(items):
            guid = "%s-%s" % (set_type, i)
            text_a = item[0]
            text_b = item[1]
            label = item[2]
            examples.append(
                InputExample(guid=guid, text_a=text_a, text_b=text_b, label=label))
        return examples
    

class UkpSamProcessor(SameSideProcessor):
    """Processor for the UKP SAM data set"""

    def _create_examples(self, items, set_type):
        """Creates examples for the training and dev sets."""
        examples = []
        
        for (i, item) in enumerate(items):
            guid = "%s-%s" % (set_type, i)
            text_a = item[0]
            label = item[1]
            examples.append(InputExample(guid=guid, text_a=text_a, label=label))
        return examples



class BertForBinaryClassification(BertPreTrainedModel):
    r"""
        **labels**: (`optional`) ``torch.LongTensor`` of shape ``(batch_size,)``:
            Labels for computing the sequence classification/regression loss.
            Indices should be in ``[0, ..., config.num_labels - 1]``.
            If ``config.num_labels == 1`` a regression loss is computed (Mean-Square loss),
            If ``config.num_labels > 1`` a classification loss is computed (Cross-Entropy).
    Outputs: `Tuple` comprising various elements depending on the configuration (config) and inputs:
        **loss**: (`optional`, returned when ``labels`` is provided) ``torch.FloatTensor`` of shape ``(1,)``:
            Classification (or regression if config.num_labels==1) loss.
        **logits**: ``torch.FloatTensor`` of shape ``(batch_size, config.num_labels)``
            Classification (or regression if config.num_labels==1) scores (before SoftMax).
        **hidden_states**: (`optional`, returned when ``config.output_hidden_states=True``)
            list of ``torch.FloatTensor`` (one for the output of each layer + the output of the embeddings)
            of shape ``(batch_size, sequence_length, hidden_size)``:
            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
        **attentions**: (`optional`, returned when ``config.output_attentions=True``)
            list of ``torch.FloatTensor`` (one for each layer) of shape ``(batch_size, num_heads, sequence_length, sequence_length)``:
            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads.
    Examples::
        tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
        model = BertForSequenceClassification.from_pretrained('bert-base-uncased')
        input_ids = torch.tensor(tokenizer.encode("Hello, my dog is cute")).unsqueeze(0)  # Batch size 1
        labels = torch.tensor([1]).unsqueeze(0)  # Batch size 1
        outputs = model(input_ids, labels=labels)
        loss, logits = outputs[:2]
    """
    def __init__(self, config):
        super(BertForBinaryClassification, self).__init__(config)
        self.num_labels = config.num_labels
        # print(config)
        self.bert = BertModel(config)
        self.dropout = nn.Dropout(config.hidden_dropout_prob)
        # self.classifier = nn.Sequential(
        #     nn.Linear(config.hidden_size, config.hidden_size),
        #     nn.ReLU(),
        #     nn.Dropout(p=0.1),
        #     nn.Linear(config.hidden_size, config.num_labels)
        # )
        self.classifier = nn.Linear(config.hidden_size, config.num_labels)

        self.init_weights()
        self.debug = True

    def forward(self, input_ids, token_type_ids=None, attention_mask=None, labels=None,
                position_ids=None, head_mask=None):
        outputs = self.bert(input_ids, position_ids=position_ids, token_type_ids=token_type_ids,
                            attention_mask=attention_mask, head_mask=head_mask)
        pooled_output = outputs[1]
        
        # debug
        # if self.debug:
        #     import pickle
        #     with open("pooled_output.pkl", "wb") as f:
        #         pickle.dump(outputs, f)
        #         pickle.dump(token_type_ids, f)
        #         pickle.dump(input_ids, f)
        #         pickle.dump(labels, f)
        #     self.debug = False

        pooled_output = self.dropout(pooled_output)
        logits = self.classifier(pooled_output)

        outputs = (logits,) + outputs[2:]  # add hidden states and attention if they are here
        
        if labels is not None:
            loss_fct = torch.nn.BCEWithLogitsLoss()
            loss = loss_fct(logits.view(-1), labels.view(-1))
            outputs = (loss,) + outputs

        return outputs  # (loss), logits, (hidden_states), (attentions)

    
class BertForSiameseClassification(BertPreTrainedModel):
    def __init__(self, config):
        super(BertForSiameseClassification, self).__init__(config)
        self.num_labels = config.num_labels
        # print(config)
        self.bert = BertModel(config)
        self.nd = config.hidden_size
        self.encoder = nn.Sequential(
            nn.Linear(self.nd, self.nd),
            nn.Tanh(),
            nn.Dropout(config.hidden_dropout_prob)
        )
        self.init_weights()
        

    def forward(self, input_ids, token_type_ids=None, attention_mask=None, labels=None,
                position_ids=None, head_mask=None):
        outputs = self.bert(input_ids, position_ids=position_ids, token_type_ids=token_type_ids,
                            attention_mask=attention_mask, head_mask=head_mask)
        
        # select along token_type_ids
        # [CLS] seq_a [SEP] seq_b
        # avg(seq_a), avg(seq_b)
        # dropout(tanh(linear(a))), dropout(tanh(linear(b)))
        # cosine loss (y = 1 or -1) | sigmoid loss (y = 0 or 1)
        
        last_hidden_state = outputs[0]
        bs = list(last_hidden_state.size())[0]
        
        selection_left = input_ids.bool() * ~token_type_ids.bool()
        selection_left[:,0] = False # do not use CLS token
        selection_right = input_ids.bool() * token_type_ids.bool()
        
        mean_left = torch.zeros((bs, self.nd), device="cuda")
        mean_right = torch.zeros((bs, self.nd), device="cuda")
        target_ones = torch.ones(bs, device="cuda")
        target_minus_ones = -torch.ones(bs, device="cuda")
        
        for i in range(bs):
            mean_left[i] = last_hidden_state[i][selection_left[i]].mean(axis=0)
            mean_right[i] = last_hidden_state[i][selection_right[i]].mean(axis=0)
            
        emb_left = self.encoder(mean_left)
        emb_right = self.encoder(mean_right)
        
        sim_fct = nn.CosineSimilarity()
        similarities = sim_fct(emb_left, emb_right)
        outputs = (similarities,) + outputs[2:] 
        
        if labels is not None:
            loss_fct = torch.nn.CosineEmbeddingLoss()
            # convert labels to 1 or -1
            cosine_labels = torch.where(labels > 0, target_ones, target_minus_ones)
            loss = loss_fct(emb_left, emb_right, cosine_labels)
            outputs = (loss,) + outputs

        return outputs  # (loss), similarities, (hidden_states), (attentions)

    
    
class BertForSimilarityClassification(BertPreTrainedModel):
    def __init__(self, config):
        super(BertForSimilarityClassification, self).__init__(config)
        self.num_labels = config.num_labels
        # print(config)
        self.bert = BertModel(config)
        self.nd = config.hidden_size
        self.encoder = nn.Sequential(
            nn.Linear(self.nd, self.nd),
            nn.Tanh(),
            nn.Dropout(config.hidden_dropout_prob)
        )
        self.classifier = torch.nn.Sequential(
            torch.nn.Linear(3 * self.nd, 1),
            torch.nn.Sigmoid()
        )
        self.init_weights()
        

    def forward(self, input_ids, token_type_ids=None, attention_mask=None, labels=None,
                position_ids=None, head_mask=None):
        outputs = self.bert(input_ids, position_ids=position_ids, token_type_ids=token_type_ids,
                            attention_mask=attention_mask, head_mask=head_mask)
        
        # select along token_type_ids
        # [CLS] seq_a [SEP] seq_b
        # avg(seq_a), avg(seq_b)
        # dropout(tanh(linear(a))), dropout(tanh(linear(b)))
        # cosine loss (y = 1 or -1) | sigmoid loss (y = 0 or 1)
        
        last_hidden_state = outputs[0]
        bs = list(last_hidden_state.size())[0]
        
        selection_left = input_ids.bool() * ~token_type_ids.bool()
        selection_left[:,0] = False # do not use CLS token
        selection_right = input_ids.bool() * token_type_ids.bool()
        
        mean_left = torch.zeros((bs, self.nd), device="cuda")
        mean_right = torch.zeros((bs, self.nd), device="cuda")
        
        for i in range(bs):
            mean_left[i] = last_hidden_state[i][selection_left[i]].mean(axis=0)
            mean_right[i] = last_hidden_state[i][selection_right[i]].mean(axis=0)
            
        emb_left = self.encoder(mean_left)
        emb_right = self.encoder(mean_right)
        
        combined_features = torch.cat((emb_left, emb_right, torch.abs(emb_left - emb_right)), 1)
        
        activations = self.classifier(combined_features)
        outputs = (activations,) + outputs[2:] 
        
        if labels is not None:
            loss_fct = nn.BCELoss()
            loss = loss_fct(activations.view(-1), labels.view(-1))
            outputs = (loss,) + outputs

        return outputs  # (loss), similarities, (hidden_states), (attentions)    
    
    
def convert_example_to_feature(example_row, pad_token=0,
sequence_a_segment_id=0, sequence_b_segment_id=1,
cls_token_segment_id=1, pad_token_segment_id=0,
mask_padding_with_zero=True):
    example, label_map, max_seq_length, tokenizer, output_mode, cls_token_at_end, cls_token, sep_token, cls_token_segment_id, pad_on_left, pad_token_segment_id = example_row

    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 unambiguously 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 = tokens_a + [sep_token]
    segment_ids = [sequence_a_segment_id] * len(tokens)

    if tokens_b:
        tokens += tokens_b + [sep_token]
        segment_ids += [sequence_b_segment_id] * (len(tokens_b) + 1)

    if cls_token_at_end:
        tokens = tokens + [cls_token]
        segment_ids = segment_ids + [cls_token_segment_id]
    else:
        tokens = [cls_token] + tokens
        segment_ids = [cls_token_segment_id] + segment_ids

    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 if mask_padding_with_zero else 0] * len(input_ids)

    # Zero-pad up to the sequence length.
    padding_length = max_seq_length - len(input_ids)
    if pad_on_left:
        input_ids = ([pad_token] * padding_length) + input_ids
        input_mask = ([0 if mask_padding_with_zero else 1] * padding_length) + input_mask
        segment_ids = ([pad_token_segment_id] * padding_length) + segment_ids
    else:
        input_ids = input_ids + ([pad_token] * padding_length)
        input_mask = input_mask + ([0 if mask_padding_with_zero else 1] * padding_length)
        segment_ids = segment_ids + ([pad_token_segment_id] * padding_length)

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

    if output_mode == "classification":
        label_id = label_map[example.label]
    elif output_mode == "regression":
        label_id = float(example.label)
    else:
        raise KeyError(output_mode)

    return InputFeatures(input_ids=input_ids,
                        input_mask=input_mask,
                        segment_ids=segment_ids,
                        label_id=label_id)
    

def convert_examples_to_features(examples, label_list, max_seq_length,
                                 tokenizer, output_mode,
                                 cls_token_at_end=False, pad_on_left=False,
                                 cls_token='[CLS]', sep_token='[SEP]', pad_token=0,
                                 sequence_a_segment_id=0, sequence_b_segment_id=1,
                                 cls_token_segment_id=1, pad_token_segment_id=0,
                                 mask_padding_with_zero=True):
    """ Loads a data file into a list of `InputBatch`s
        `cls_token_at_end` define the location of the CLS token:
            - False (Default, BERT/XLM pattern): [CLS] + A + [SEP] + B + [SEP]
            - True (XLNet/GPT pattern): A + [SEP] + B + [SEP] + [CLS]
        `cls_token_segment_id` define the segment id associated to the CLS token (0 for BERT, 2 for XLNet)
    """

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

    examples = [(example, label_map, max_seq_length, tokenizer, output_mode, cls_token_at_end, cls_token, sep_token, cls_token_segment_id, pad_on_left, pad_token_segment_id) for example in examples]

    process_count = cpu_count() - 2

    with Pool(process_count) as p:
        features = list(tqdm(p.imap(convert_example_to_feature, examples, chunksize=100), total=len(examples)))


    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()


processors = {
    "binary": SameSideProcessor,
    "ukp_sam": UkpSamProcessor
}

output_modes = {
    "binary": "classification",
    "ukp_sam": "classification"
}

GLUE_TASKS_NUM_LABELS = {
    "binary": "1",
    "ukp_sam": "2"
}
