import sys
import math
import numpy as np
from datetime import datetime

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable


class DecAtt(nn.Module):

    def __init__(self, num_units, num_classes, embedding_size, dropout, device=0, 
                 training=True, project_input=True,
                 use_intra_attention=False, distance_biases=10, max_sentence_length=30):
        """
        Create the model based on MLP networks.

        :param num_units: size of the networks
        :param num_classes: number of classes in the problem
        :param embedding_size: size of each word embedding
        :param use_intra_attention: whether to use intra-attention model
        :param training: whether to create training tensors (optimizer)
        :p/word_embeddingaram project_input: whether to project input embeddings to a
            different dimensionality
        :param distance_biases: number of different distances with biases used
            in the intra-attention model
        """
        super().__init__()
        self.arch = "DecAtt"
        self.num_units = num_units
        self.num_classes = num_classes
        self.project_input = project_input
        self.embedding_size = embedding_size
        self.distance_biases = distance_biases
        self.intra_attention = False
        self.max_sentence_length = max_sentence_length
        self.device = device

        self.bias_embedding = nn.Embedding(max_sentence_length,1)
        self.linear_layer_project = nn.Linear(embedding_size, num_units, bias=False)
        #self.linear_layer_intra = nn.Sequential(nn.Linear(num_units, num_units), nn.ReLU(), nn.Linear(num_units, num_units), nn.ReLU())

        self.linear_layer_attend = nn.Sequential(nn.Dropout(p=dropout), nn.Linear(num_units, num_units), nn.ReLU(),
                                                 nn.Dropout(p=dropout), nn.Linear(num_units, num_units), nn.ReLU())

        self.linear_layer_compare = nn.Sequential(nn.Dropout(p=dropout), nn.Linear(num_units*2, num_units), nn.ReLU(),
                                                  nn.Dropout(p=dropout), nn.Linear(num_units, num_units), nn.ReLU())

        self.linear_layer_aggregate = nn.Sequential(nn.Dropout(p=dropout), nn.Linear(num_units*2, num_units), nn.ReLU(),
                                                    nn.Dropout(p=dropout), nn.Linear(num_units, num_units), nn.ReLU(),
                                                    nn.Linear(num_units, num_classes), nn.LogSoftmax())
        self.init_weight()

    def init_weight(self):
        self.linear_layer_project.weight.data.normal_(0, 0.01)
        self.linear_layer_attend[1].weight.data.normal_(0, 0.01)
        self.linear_layer_attend[1].bias.data.fill_(0)
        self.linear_layer_attend[4].weight.data.normal_(0, 0.01)
        self.linear_layer_attend[4].bias.data.fill_(0)
        self.linear_layer_compare[1].weight.data.normal_(0, 0.01)
        self.linear_layer_compare[1].bias.data.fill_(0)
        self.linear_layer_compare[4].weight.data.normal_(0, 0.01)
        self.linear_layer_compare[4].bias.data.fill_(0)
        self.linear_layer_aggregate[1].weight.data.normal_(0, 0.01)
        self.linear_layer_aggregate[1].bias.data.fill_(0)
        self.linear_layer_aggregate[4].weight.data.normal_(0, 0.01)
        self.linear_layer_aggregate[4].bias.data.fill_(0)
        #self.word_embedding.weight.data.copy_(torch.from_numpy(self.pretrained_emb))

    def attention_softmax3d(self, raw_attentions):
        reshaped_attentions = raw_attentions.view(-1, raw_attentions.size(2))
        out = nn.functional.softmax(reshaped_attentions, dim=1)
        return out.view(raw_attentions.size(0),raw_attentions.size(1),raw_attentions.size(2))

    def _transformation_input(self, embed_sent):
        embed_sent = self.linear_layer_project(embed_sent)
        result = embed_sent
        if self.intra_attention:
            f_intra = self.linear_layer_intra(embed_sent)
            f_intra_t = torch.transpose(f_intra, 1, 2)
            raw_attentions = torch.matmul(f_intra, f_intra_t)
            time_steps = embed_sent.size(1)
            r = torch.arange(0, time_steps)
            r_matrix = r.view(1,-1).expand(time_steps,time_steps)
            raw_index = r_matrix-r.view(-1,1)
            clipped_index = torch.clamp(raw_index,0,self.distance_biases-1)
            clipped_index = Variable(clipped_index.long())
            if torch.cuda.is_available():
                clipped_index = clipped_index.to(self.device)
            bias = self.bias_embedding(clipped_index)
            bias = torch.squeeze(bias)
            raw_attentions += bias
            attentions = self.attention_softmax3d(raw_attentions)
            attended = torch.matmul(attentions, embed_sent)
            result = torch.cat([embed_sent,attended],2)
        return result

    def attend(self, sent1, sent2, lsize_list, rsize_list):
        """
        Compute inter-sentence attention. This is step 1 (attend) in the paper

        :param sent1: tensor in shape (batch, time_steps, num_units),
            the projected sentence 1
        :param sent2: tensor in shape (batch, time_steps, num_units)
        :return: a tuple of 3-d tensors, alfa and beta.
        """
        repr1 = self.linear_layer_attend(sent1)
        repr2 = self.linear_layer_attend(sent2)
        repr2 = torch.transpose(repr2,1,2)
        raw_attentions = torch.matmul(repr1, repr2)

        #self.mask = generate_mask(lsize_list, rsize_list)
        # masked = mask(self.raw_attentions, rsize_list)
        #masked = raw_attentions * self.mask
        att_sent1 = self.attention_softmax3d(raw_attentions)
        beta = torch.matmul(att_sent1, sent2) #input2_soft

        raw_attentions_t = torch.transpose(raw_attentions,1,2).contiguous()
        #self.mask_t = torch.transpose(self.mask, 1, 2).contiguous()
        # masked = mask(raw_attentions_t, lsize_list)
        #masked = raw_attentions_t * self.mask_t
        att_sent2 = self.attention_softmax3d(raw_attentions_t)
        alpha = torch.matmul(att_sent2,sent1) #input1_soft

        return alpha, beta

    def compare(self, sentence, soft_alignment):
        """
        Apply a feed forward network to compare o   ne sentence to its
        soft alignment with the other.

        :param sentence: embedded and projected sentence,
            shape (batch, time_steps, num_units)
        :param soft_alignment: tensor with shape (batch, time_steps, num_units)
        :return: a tensor (batch, time_steps, num_units)
        """
        sent_alignment = torch.cat([sentence, soft_alignment],2)
        out = self.linear_layer_compare(sent_alignment)
        #out, (state, _) = self.lstm_compare(out)
        return out

    def aggregate(self, v1, v2):
        """
        Aggregate the representations induced from both sentences and their
        representations

        :param v1: tensor with shape (batch, time_steps, num_units)
        :param v2: tensor with shape (batch, time_steps, num_units)
        :return: logits over classes, shape (batch, num_classes)
        """
        v1_sum = torch.sum(v1,1)
        v2_sum = torch.sum(v2,1)
        out = self.linear_layer_aggregate(torch.cat([v1_sum,v2_sum],1))
        return out

    def forward(self, sent1, sent2, ext_feats=None, word_to_doc_count=None, raw_sent1=None, raw_sent2=None):
        lsize_list = [len(s.split(" ")) for s in raw_sent1]
        rsize_list = [len(s.split(" ")) for s in raw_sent2]
        sent1 = sent1.permute(0, 2, 1)
        sent2 = sent2.permute(0, 2, 1)
        sent1 = self._transformation_input(sent1)
        sent2 = self._transformation_input(sent2)
        alpha, beta = self.attend(sent1, sent2, lsize_list, rsize_list)
        v1 = self.compare(sent1, beta)
        v2 = self.compare(sent2, alpha)
        logits = self.aggregate(v1, v2)
        return logits


[((41, 30, 41, 65), 'torch.nn.Embedding', 'nn.Embedding', ({(41, 43, 41, 62): 'max_sentence_length', (41, 63, 41, 64): '1'}, {}), '(max_sentence_length, 1)', True, 'import torch.nn as nn\n'), ((42, 36, 42, 84), 'torch.nn.Linear', 'nn.Linear', (), '', True, 'import torch.nn as nn\n'), ((74, 14, 74, 63), 'torch.nn.functional.softmax', 'nn.functional.softmax', (), '', True, 'import torch.nn as nn\n'), ((111, 16, 111, 42), 'torch.transpose', 'torch.transpose', ({(111, 32, 111, 37): 'repr2', (111, 38, 111, 39): '1', (111, 40, 111, 41): '2'}, {}), '(repr2, 1, 2)', False, 'import torch\n'), ((112, 25, 112, 51), 'torch.matmul', 'torch.matmul', ({(112, 38, 112, 43): 'repr1', (112, 45, 112, 50): 'repr2'}, {}), '(repr1, repr2)', False, 'import torch\n'), ((118, 15, 118, 45), 'torch.matmul', 'torch.matmul', ({(118, 28, 118, 37): 'att_sent1', (118, 39, 118, 44): 'sent2'}, {}), '(att_sent1, sent2)', False, 'import torch\n'), ((125, 16, 125, 45), 'torch.matmul', 'torch.matmul', ({(125, 29, 125, 38): 'att_sent2', (125, 39, 125, 44): 'sent1'}, {}), '(att_sent2, sent1)', False, 'import torch\n'), ((139, 25, 139, 64), 'torch.cat', 'torch.cat', ({(139, 35, 139, 61): '[sentence, soft_alignment]', (139, 62, 139, 63): '2'}, {}), '([sentence, soft_alignment], 2)', False, 'import torch\n'), ((153, 17, 153, 32), 'torch.sum', 'torch.sum', ({(153, 27, 153, 29): 'v1', (153, 30, 153, 31): '1'}, {}), '(v1, 1)', False, 'import torch\n'), ((154, 17, 154, 32), 'torch.sum', 'torch.sum', ({(154, 27, 154, 29): 'v2', (154, 30, 154, 31): '1'}, {}), '(v2, 1)', False, 'import torch\n'), ((45, 49, 45, 70), 'torch.nn.Dropout', 'nn.Dropout', (), '', True, 'import torch.nn as nn\n'), ((45, 72, 45, 103), 'torch.nn.Linear', 'nn.Linear', ({(45, 82, 45, 91): 'num_units', (45, 93, 45, 102): 'num_units'}, {}), '(num_units, num_units)', True, 'import torch.nn as nn\n'), ((45, 105, 45, 114), 'torch.nn.ReLU', 'nn.ReLU', ({}, {}), '()', True, 'import torch.nn as nn\n'), ((46, 49, 46, 70), 'torch.nn.Dropout', 'nn.Dropout', (), '', True, 'import torch.nn as nn\n'), ((46, 72, 46, 103), 'torch.nn.Linear', 'nn.Linear', ({(46, 82, 46, 91): 'num_units', (46, 93, 46, 102): 'num_units'}, {}), '(num_units, num_units)', True, 'import torch.nn as nn\n'), ((46, 105, 46, 114), 'torch.nn.ReLU', 'nn.ReLU', ({}, {}), '()', True, 'import torch.nn as nn\n'), ((48, 50, 48, 71), 'torch.nn.Dropout', 'nn.Dropout', (), '', True, 'import torch.nn as nn\n'), ((48, 73, 48, 106), 'torch.nn.Linear', 'nn.Linear', ({(48, 83, 48, 94): 'num_units * 2', (48, 96, 48, 105): 'num_units'}, {}), '(num_units * 2, num_units)', True, 'import torch.nn as nn\n'), ((48, 108, 48, 117), 'torch.nn.ReLU', 'nn.ReLU', ({}, {}), '()', True, 'import torch.nn as nn\n'), ((49, 50, 49, 71), 'torch.nn.Dropout', 'nn.Dropout', (), '', True, 'import torch.nn as nn\n'), ((49, 73, 49, 104), 'torch.nn.Linear', 'nn.Linear', ({(49, 83, 49, 92): 'num_units', (49, 94, 49, 103): 'num_units'}, {}), '(num_units, num_units)', True, 'import torch.nn as nn\n'), ((49, 106, 49, 115), 'torch.nn.ReLU', 'nn.ReLU', ({}, {}), '()', True, 'import torch.nn as nn\n'), ((51, 52, 51, 73), 'torch.nn.Dropout', 'nn.Dropout', (), '', True, 'import torch.nn as nn\n'), ((51, 75, 51, 108), 'torch.nn.Linear', 'nn.Linear', ({(51, 85, 51, 96): 'num_units * 2', (51, 98, 51, 107): 'num_units'}, {}), '(num_units * 2, num_units)', True, 'import torch.nn as nn\n'), ((51, 110, 51, 119), 'torch.nn.ReLU', 'nn.ReLU', ({}, {}), '()', True, 'import torch.nn as nn\n'), ((52, 52, 52, 73), 'torch.nn.Dropout', 'nn.Dropout', (), '', True, 'import torch.nn as nn\n'), ((52, 75, 52, 106), 'torch.nn.Linear', 'nn.Linear', ({(52, 85, 52, 94): 'num_units', (52, 96, 52, 105): 'num_units'}, {}), '(num_units, num_units)', True, 'import torch.nn as nn\n'), ((52, 108, 52, 117), 'torch.nn.ReLU', 'nn.ReLU', ({}, {}), '()', True, 'import torch.nn as nn\n'), ((53, 52, 53, 85), 'torch.nn.Linear', 'nn.Linear', ({(53, 62, 53, 71): 'num_units', (53, 73, 53, 84): 'num_classes'}, {}), '(num_units, num_classes)', True, 'import torch.nn as nn\n'), ((53, 87, 53, 102), 'torch.nn.LogSoftmax', 'nn.LogSoftmax', ({}, {}), '()', True, 'import torch.nn as nn\n'), ((82, 24, 82, 54), 'torch.transpose', 'torch.transpose', ({(82, 40, 82, 47): 'f_intra', (82, 49, 82, 50): '1', (82, 52, 82, 53): '2'}, {}), '(f_intra, 1, 2)', False, 'import torch\n'), ((83, 29, 83, 61), 'torch.matmul', 'torch.matmul', ({(83, 42, 83, 49): 'f_intra', (83, 51, 83, 60): 'f_intra_t'}, {}), '(f_intra, f_intra_t)', False, 'import torch\n'), ((85, 16, 85, 43), 'torch.arange', 'torch.arange', ({(85, 29, 85, 30): '0', (85, 32, 85, 42): 'time_steps'}, {}), '(0, time_steps)', False, 'import torch\n'), ((88, 28, 88, 75), 'torch.clamp', 'torch.clamp', ({(88, 40, 88, 49): 'raw_index', (88, 50, 88, 51): '0', (88, 52, 88, 74): 'self.distance_biases - 1'}, {}), '(raw_index, 0, self.distance_biases - 1)', False, 'import torch\n'), ((90, 15, 90, 40), 'torch.cuda.is_available', 'torch.cuda.is_available', ({}, {}), '()', False, 'import torch\n'), ((93, 19, 93, 38), 'torch.squeeze', 'torch.squeeze', ({(93, 33, 93, 37): 'bias'}, {}), '(bias)', False, 'import torch\n'), ((96, 23, 96, 59), 'torch.matmul', 'torch.matmul', ({(96, 36, 96, 46): 'attentions', (96, 48, 96, 58): 'embed_sent'}, {}), '(attentions, embed_sent)', False, 'import torch\n'), ((97, 21, 97, 55), 'torch.cat', 'torch.cat', ({(97, 31, 97, 52): '[embed_sent, attended]', (97, 53, 97, 54): '2'}, {}), '([embed_sent, attended], 2)', False, 'import torch\n'), ((155, 42, 155, 70), 'torch.cat', 'torch.cat', ({(155, 52, 155, 67): '[v1_sum, v2_sum]', (155, 68, 155, 69): '1'}, {}), '([v1_sum, v2_sum], 1)', False, 'import torch\n'), ((120, 27, 120, 62), 'torch.transpose', 'torch.transpose', ({(120, 43, 120, 57): 'raw_attentions', (120, 58, 120, 59): '1', (120, 60, 120, 61): '2'}, {}), '(raw_attentions, 1, 2)', False, 'import torch\n')]