import math

import torch
import torch.nn as nn

from models.neural import MultiHeadedAttention, PositionwiseFeedForward
from models.rnn import LayerNormLSTM


class Classifier(nn.Module):
    def __init__(self, hidden_size):
        super(Classifier, self).__init__()
        self.linear1 = nn.Linear(hidden_size, 1)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x, mask_cls):
        h = self.linear1(x).squeeze(-1)
        sent_scores = self.sigmoid(h) * mask_cls.float()
        return sent_scores


class PositionalEncoding(nn.Module):

    def __init__(self, dropout, dim, max_len=5000):
        pe = torch.zeros(max_len, dim)
        position = torch.arange(0, max_len).unsqueeze(1)
        div_term = torch.exp((torch.arange(0, dim, 2, dtype=torch.float) *
                              -(math.log(10000.0) / dim)))
        pe[:, 0::2] = torch.sin(position.float() * div_term)
        pe[:, 1::2] = torch.cos(position.float() * div_term)
        pe = pe.unsqueeze(0)
        super(PositionalEncoding, self).__init__()
        self.register_buffer('pe', pe)
        self.dropout = nn.Dropout(p=dropout)
        self.dim = dim

    def forward(self, emb, step=None):
        emb = emb * math.sqrt(self.dim)
        if (step):
            emb = emb + self.pe[:, step][:, None, :]

        else:
            emb = emb + self.pe[:, :emb.size(1)]
        emb = self.dropout(emb)
        return emb

    def get_emb(self, emb):
        return self.pe[:, :emb.size(1)]


class TransformerEncoderLayer(nn.Module):
    def __init__(self, d_model, heads, d_ff, dropout):
        super(TransformerEncoderLayer, self).__init__()

        self.self_attn = MultiHeadedAttention(
            heads, d_model, dropout=dropout)
        self.feed_forward = PositionwiseFeedForward(d_model, d_ff, dropout)
        self.layer_norm = nn.LayerNorm(d_model, eps=1e-6)
        self.dropout = nn.Dropout(dropout)

    def forward(self, iter, query, inputs, mask):
        if (iter != 0):
            input_norm = self.layer_norm(inputs)
        else:
            input_norm = inputs

        mask = mask.unsqueeze(1)
        context = self.self_attn(input_norm, input_norm, input_norm,
                                 mask=mask)
        out = self.dropout(context) + inputs
        return self.feed_forward(out)


class TransformerInterEncoder(nn.Module):
    def __init__(self, d_model, d_ff, heads, dropout, num_inter_layers=0):
        super(TransformerInterEncoder, self).__init__()
        self.d_model = d_model
        self.num_inter_layers = num_inter_layers
        self.pos_emb = PositionalEncoding(dropout, d_model)
        self.transformer_inter = nn.ModuleList(
            [TransformerEncoderLayer(d_model, heads, d_ff, dropout)
             for _ in range(num_inter_layers)])
        self.dropout = nn.Dropout(dropout)
        self.layer_norm = nn.LayerNorm(d_model, eps=1e-6)
        self.wo = nn.Linear(d_model, 1, bias=True)
        self.sigmoid = nn.Sigmoid()

    def forward(self, top_vecs, mask):
        """ See :obj:`EncoderBase.forward()`"""

        batch_size, n_sents = top_vecs.size(0), top_vecs.size(1)
        pos_emb = self.pos_emb.pe[:, :n_sents]
        x = top_vecs * mask[:, :, None].float()
        x = x + pos_emb

        for i in range(self.num_inter_layers):
            x = self.transformer_inter[i](i, x, x, ~mask)  # all_sents * max_tokens * dim

        x = self.layer_norm(x)
        sent_scores = self.sigmoid(self.wo(x))
        sent_scores = sent_scores.squeeze(-1) * mask.float()

        return sent_scores

class GRUEncoder_attn(nn.Module):

    def __init__(self,bidirectional, num_layers, input_size, hidden_size,dropout=0.0):
        super(GRUEncoder_attn,self).__init__()


class RNNEncoder_attn(nn.Module):

    def __init__(self, bidirectional, num_layers, input_size,
                 hidden_size, dropout=0.0):
        super(RNNEncoder_attn, self).__init__()
        num_directions = 2 if bidirectional else 1
        assert hidden_size % num_directions == 0
        hidden_size = hidden_size // num_directions
        self.relu = nn.ReLU()

        self.rnn = LayerNormLSTM(
            input_size=input_size,
            hidden_size=hidden_size,
            num_layers=num_layers,
            bidirectional=bidirectional)

        self.wo = nn.Linear(num_directions * hidden_size, 1, bias=True)

        self.dropout = nn.Dropout(dropout)
        self.softmax = nn.Softmax()
        print('this is dropout',dropout)

    def forward(self, x, mask):
        """See :func:`EncoderBase.forward()`"""
        batch, layer, seq, hidden = x.size()
        x1=x.contiguous().view(batch * layer, -1, hidden)
        x1 = torch.transpose(x1, 1, 0)
        memory_bank, _ = self.rnn(x1)
        memory_bank = self.dropout(memory_bank) + x1
        memory_bank = torch.transpose(memory_bank, 1, 0)
        # sent_scores = self.softmax(self.relu(self.wo(memory_bank)).squeeze(dim=-1)).unsqueeze(-1)

        sent_scores = self.softmax(self.relu(self.wo(memory_bank[:,-1,:])).squeeze(dim=-1).view(-1,layer)).unsqueeze(-1)
        x=x.transpose(1,2)
        sent_vec = torch.matmul(sent_scores.transpose(1,2).unsqueeze(dim = 1).expand(batch,seq,1,layer),x)

        return sent_vec.squeeze(dim = 2)

class TransformerDecoderLayer(nn.Module):
    def __init__(self, d_model, heads, d_ff, dropout):
        super(TransformerDecoderLayer, self).__init__()

        self.self_attn = MultiHeadedAttention(
            heads, d_model, dropout=dropout)
        self.feed_forward = PositionwiseFeedForward(d_model, d_ff, dropout)
        self.layer_norm = nn.LayerNorm(d_model, eps=1e-6)

    def forward(self, iter, ent_enc, inputs, self_attn_mask=None,context_attn_mask=None):
        context = self.self_attn(inputs, inputs, inputs,
                                 mask=self_attn_mask)
        dec_output = self.self_attn(
            ent_enc, ent_enc, context, mask=context_attn_mask)
        dec_output = self.feed_forward(dec_output)
        return dec_output

class TransformerInterDecoder(nn.Module):
    def __init__(self, d_model, d_ff, heads, dropout, d_hidden, num_inter_layers=0):
        super(TransformerInterDecoder, self).__init__()
        self.d_model = d_model
        self.num_inter_layers = num_inter_layers
        self.pos_emb = PositionalEncoding(dropout, d_model)
        self.transformer_inter = nn.ModuleList(
            [TransformerDecoderLayer(d_model, heads, d_ff, dropout)
             for _ in range(num_inter_layers)])
        self.dropout = nn.Dropout(dropout)
        self.layer_norm = nn.LayerNorm(d_model, eps=1e-6)
        self.wo = nn.Linear(d_model, d_hidden , bias=True)
        self.wi = nn.Linear(d_model, d_hidden, bias=True)
        self.v = nn.Linear(d_hidden, 1, bias=True)
        self.LR = nn.LeakyReLU()
        self.softmax = nn.Softmax(dim=-1)

    def forward(self, top_vecs, inputs, mask, label_mask=None):
        """ See :obj:`EncoderBase.forward()`"""
        n_out = inputs.size(1)
        pos_emb = self.pos_emb.pe[:, :n_out]
        seq_mask=subsequent_mask(inputs)
        self_attn_mask = torch.gt((~label_mask.unsqueeze(1).expand(-1, n_out, -1) + seq_mask), 0)
        inputs=inputs+pos_emb
        for i in range(self.num_inter_layers):
            inputs = self.transformer_inter[i](i, top_vecs, inputs,self_attn_mask,~ mask.unsqueeze(1).expand(-1, n_out,-1))
        scores=self.v(self.LR(
            self.wo(inputs.unsqueeze(2)).expand(-1, -1, top_vecs.size(1), -1) + self.wi(top_vecs).unsqueeze(
                1))).squeeze(-1)
        sent_scores = self.softmax(scores)
        return sent_scores


class RNNEncoder(nn.Module):

    def __init__(self, bidirectional, num_layers, input_size,
                 hidden_size, dropout=0.0):
        super(RNNEncoder, self).__init__()
        num_directions = 2 if bidirectional else 1
        assert hidden_size % num_directions == 0
        hidden_size = hidden_size // num_directions

        self.rnn = LayerNormLSTM(
            input_size=input_size,
            hidden_size=hidden_size,
            num_layers=num_layers,
            bidirectional=bidirectional)

        self.wo = nn.Linear(num_directions * hidden_size, 1, bias=True)
        self.dropout = nn.Dropout(dropout)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x, mask):
        """See :func:`EncoderBase.forward()`"""
        x = torch.transpose(x, 1, 0)
        memory_bank, _ = self.rnn(x)
        memory_bank = self.dropout(memory_bank) + x
        memory_bank = torch.transpose(memory_bank, 1, 0)

        sent_scores = self.sigmoid(self.wo(memory_bank))
        sent_scores = sent_scores.squeeze(-1) * mask.float()
        return sent_scores
class GCN(nn.Module):
    def __init__(self,in_channel,out_channel,hidden_dim,drop):
        super(GCN, self).__init__()
        self.in_channel=in_channel
        self.out_channel=out_channel
        self.hidden_dim=hidden_dim
        self.dropout = nn.Dropout(p=drop)
        self.gcn_x_11=GCNConv(self.in_channel,self.hidden_dim)
        self.gcn_x_12=GCNConv(self.hidden_dim,self.out_channel)#No.1-*2*2
        # self.gcn_x_21=GCNConv(self.in_channel,self.hidden_dim)
        # self.gcn_x_22=GCNConv(self.hidden_dim,self.out_channel)#No.2-*2
        # self.gcn_mix=GCNConv(self.hidden_dim*2,self.hidden_dim)#No.2-*2
        self.relu=nn.ReLU(inplace=True)
    def forward(self, x_1, edge_index_1,  edge_index_2=None,edge_weight_1=None,edge_weight_2=None):
        syn=self.gcn_x_11(x_1, edge_index_1, edge_weight_1)
        syn=self.relu(syn)
        syn=self.dropout(syn)
        syn = self.gcn_x_12(syn, edge_index_1, edge_weight_1)
        syn = self.relu(syn)
        syn = self.dropout(syn)
        # x2 = self.gcn_x_21(x_1, edge_index_2, edge_weight_2)
        # x2 = self.relu(x2)
        # x2 = self.dropout(x2)
        # mix = self.gcn_mix(torch.cat((syn,x2),-1), edge_index_2, edge_weight_2)
        # x2 = self.gcn_x_22(mix, edge_index_2, edge_weight_2)
        # syn=self.gcn_x_12(mix, edge_index_1, edge_weight_1)
        # syn=self.relu(syn)
        # syn=self.dropout(syn)
        # x2 = self.relu(x2)
        # x2 = self.dropout(x2)
        return syn


[((13, 23, 13, 48), 'torch.nn.Linear', 'nn.Linear', ({(13, 33, 13, 44): 'hidden_size', (13, 46, 13, 47): '1'}, {}), '(hidden_size, 1)', True, 'import torch.nn as nn\n'), ((14, 23, 14, 35), 'torch.nn.Sigmoid', 'nn.Sigmoid', ({}, {}), '()', True, 'import torch.nn as nn\n'), ((25, 13, 25, 38), 'torch.zeros', 'torch.zeros', ({(25, 25, 25, 32): 'max_len', (25, 34, 25, 37): 'dim'}, {}), '(max_len, dim)', False, 'import torch\n'), ((34, 23, 34, 44), 'torch.nn.Dropout', 'nn.Dropout', (), '', True, 'import torch.nn as nn\n'), ((55, 25, 56, 44), 'models.neural.MultiHeadedAttention', 'MultiHeadedAttention', (), '', False, 'from models.neural import MultiHeadedAttention, PositionwiseFeedForward\n'), ((57, 28, 57, 75), 'models.neural.PositionwiseFeedForward', 'PositionwiseFeedForward', ({(57, 52, 57, 59): 'd_model', (57, 61, 57, 65): 'd_ff', (57, 67, 57, 74): 'dropout'}, {}), '(d_model, d_ff, dropout)', False, 'from models.neural import MultiHeadedAttention, PositionwiseFeedForward\n'), ((58, 26, 58, 57), 'torch.nn.LayerNorm', 'nn.LayerNorm', (), '', True, 'import torch.nn as nn\n'), ((59, 23, 59, 42), 'torch.nn.Dropout', 'nn.Dropout', ({(59, 34, 59, 41): 'dropout'}, {}), '(dropout)', True, 'import torch.nn as nn\n'), ((83, 23, 83, 42), 'torch.nn.Dropout', 'nn.Dropout', ({(83, 34, 83, 41): 'dropout'}, {}), '(dropout)', True, 'import torch.nn as nn\n'), ((84, 26, 84, 57), 'torch.nn.LayerNorm', 'nn.LayerNorm', (), '', True, 'import torch.nn as nn\n'), ((85, 18, 85, 50), 'torch.nn.Linear', 'nn.Linear', (), '', True, 'import torch.nn as nn\n'), ((86, 23, 86, 35), 'torch.nn.Sigmoid', 'nn.Sigmoid', ({}, {}), '()', True, 'import torch.nn as nn\n'), ((119, 20, 119, 29), 'torch.nn.ReLU', 'nn.ReLU', ({}, {}), '()', True, 'import torch.nn as nn\n'), ((121, 19, 125, 40), 'models.rnn.LayerNormLSTM', 'LayerNormLSTM', (), '', False, 'from models.rnn import LayerNormLSTM\n'), ((127, 18, 127, 71), 'torch.nn.Linear', 'nn.Linear', (), '', True, 'import torch.nn as nn\n'), ((129, 23, 129, 42), 'torch.nn.Dropout', 'nn.Dropout', ({(129, 34, 129, 41): 'dropout'}, {}), '(dropout)', True, 'import torch.nn as nn\n'), ((130, 23, 130, 35), 'torch.nn.Softmax', 'nn.Softmax', ({}, {}), '()', True, 'import torch.nn as nn\n'), ((137, 13, 137, 38), 'torch.transpose', 'torch.transpose', ({(137, 29, 137, 31): 'x1', (137, 33, 137, 34): '1', (137, 36, 137, 37): '0'}, {}), '(x1, 1, 0)', False, 'import torch\n'), ((140, 22, 140, 56), 'torch.transpose', 'torch.transpose', ({(140, 38, 140, 49): 'memory_bank', (140, 51, 140, 52): '1', (140, 54, 140, 55): '0'}, {}), '(memory_bank, 1, 0)', False, 'import torch\n'), ((153, 25, 154, 44), 'models.neural.MultiHeadedAttention', 'MultiHeadedAttention', (), '', False, 'from models.neural import MultiHeadedAttention, PositionwiseFeedForward\n'), ((155, 28, 155, 75), 'models.neural.PositionwiseFeedForward', 'PositionwiseFeedForward', ({(155, 52, 155, 59): 'd_model', (155, 61, 155, 65): 'd_ff', (155, 67, 155, 74): 'dropout'}, {}), '(d_model, d_ff, dropout)', False, 'from models.neural import MultiHeadedAttention, PositionwiseFeedForward\n'), ((156, 26, 156, 57), 'torch.nn.LayerNorm', 'nn.LayerNorm', (), '', True, 'import torch.nn as nn\n'), ((175, 23, 175, 42), 'torch.nn.Dropout', 'nn.Dropout', ({(175, 34, 175, 41): 'dropout'}, {}), '(dropout)', True, 'import torch.nn as nn\n'), ((176, 26, 176, 57), 'torch.nn.LayerNorm', 'nn.LayerNorm', (), '', True, 'import torch.nn as nn\n'), ((177, 18, 177, 58), 'torch.nn.Linear', 'nn.Linear', (), '', True, 'import torch.nn as nn\n'), ((178, 18, 178, 57), 'torch.nn.Linear', 'nn.Linear', (), '', True, 'import torch.nn as nn\n'), ((179, 17, 179, 50), 'torch.nn.Linear', 'nn.Linear', (), '', True, 'import torch.nn as nn\n'), ((180, 18, 180, 32), 'torch.nn.LeakyReLU', 'nn.LeakyReLU', ({}, {}), '()', True, 'import torch.nn as nn\n'), ((181, 23, 181, 41), 'torch.nn.Softmax', 'nn.Softmax', (), '', True, 'import torch.nn as nn\n'), ((208, 19, 212, 40), 'models.rnn.LayerNormLSTM', 'LayerNormLSTM', (), '', False, 'from models.rnn import LayerNormLSTM\n'), ((214, 18, 214, 71), 'torch.nn.Linear', 'nn.Linear', (), '', True, 'import torch.nn as nn\n'), ((215, 23, 215, 42), 'torch.nn.Dropout', 'nn.Dropout', ({(215, 34, 215, 41): 'dropout'}, {}), '(dropout)', True, 'import torch.nn as nn\n'), ((216, 23, 216, 35), 'torch.nn.Sigmoid', 'nn.Sigmoid', ({}, {}), '()', True, 'import torch.nn as nn\n'), ((220, 12, 220, 36), 'torch.transpose', 'torch.transpose', ({(220, 28, 220, 29): 'x', (220, 31, 220, 32): '1', (220, 34, 220, 35): '0'}, {}), '(x, 1, 0)', False, 'import torch\n'), ((223, 22, 223, 56), 'torch.transpose', 'torch.transpose', ({(223, 38, 223, 49): 'memory_bank', (223, 51, 223, 52): '1', (223, 54, 223, 55): '0'}, {}), '(memory_bank, 1, 0)', False, 'import torch\n'), ((234, 23, 234, 41), 'torch.nn.Dropout', 'nn.Dropout', (), '', True, 'import torch.nn as nn\n'), ((240, 18, 240, 39), 'torch.nn.ReLU', 'nn.ReLU', (), '', True, 'import torch.nn as nn\n'), ((38, 20, 38, 39), 'math.sqrt', 'math.sqrt', ({(38, 30, 38, 38): 'self.dim'}, {}), '(self.dim)', False, 'import math\n'), ((26, 19, 26, 43), 'torch.arange', 'torch.arange', ({(26, 32, 26, 33): '0', (26, 35, 26, 42): 'max_len'}, {}), '(0, max_len)', False, 'import torch\n'), ((27, 30, 27, 72), 'torch.arange', 'torch.arange', (), '', False, 'import torch\n'), ((28, 32, 28, 49), 'math.log', 'math.log', ({(28, 41, 28, 48): '10000.0'}, {}), '(10000.0)', False, 'import math\n')]