import numpy as np

import torch

import torch.nn as nn

from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence

import sklearn

OPS1 = {

    'identity': lambda d_model: Identity(d_model),

    'ffn': lambda d_model: FFN(d_model),

    'interaction_1': lambda d_model: Attention_s1(d_model),

    'interaction_2': lambda d_model: Attention_s2(d_model),

}

OPS2 = {

    'identity': lambda d_model: Identity(d_model),

    'conv': lambda d_model: Conv(d_model),

    'attention': lambda d_model: SelfAttention(d_model),

    'rnn': lambda d_model: RNN(d_model),

    'ffn': lambda d_model: FFN(d_model),

    'interaction_1': lambda d_model: CatFC(d_model),

    'interaction_2': lambda d_model: Attention_x(d_model)

}

OPS3 = {

    'identity': lambda d_model: Identity(d_model),

    'zero': lambda d_model: Zero(d_model),

}

OPS4 = {

    'sum': lambda d_model: Sum(d_model),

    'mul': lambda d_model: Mul(d_model),

}

class Zero(nn.Module):

    def __init__(self, d_model):

        super(Zero, self).__init__()

    def forward(self, x, masks, lengths):

        return torch.mul(x, 0)

class Sum(nn.Module):

    def __init__(self, d_model):

        super(Sum, self).__init__()

    def forward(self, all_x):

        out = all_x[0]

        for x in all_x[1:]:

            out += x

        return out

class Mul(nn.Module):

    def __init__(self, d_model):

        super(Mul, self).__init__()

    def forward(self, all_x):

        out = all_x[0]

        for x in all_x[1:]:

            out = out * x

        return out

class CatFC(nn.Module):

    def __init__(self, d_model):

        super(CatFC, self).__init__()

        self.ffn = nn.Sequential(nn.Linear(2*d_model, 4 * d_model), nn.ReLU(),

                                 nn.Linear(4 * d_model, d_model))

        self.layer_norm = nn.LayerNorm(d_model)

    def forward(self, current_x, s, other_x):

        s_ = s.unsqueeze(1).expand_as(current_x)

        x = torch.cat((current_x, s_), dim=-1)

        return self.layer_norm(self.ffn(x))

class Conv(nn.Module):

    def __init__(self, d_model):

        super(Conv, self).__init__()

        self.op = nn.Sequential(

            nn.ReLU(),

            nn.Conv1d(d_model, d_model, 3, padding=1),

            nn.BatchNorm1d(d_model, affine=True)

        )

        # self.batchnm = nn.BatchNorm1d(d_model, affine=True)

        # self.conv = nn.Conv1d(d_model, d_model, 3, padding=1)

    def forward(self, x, masks, lengths):

        x = self.op(x.permute(0, 2, 1))

        return x.permute(0, 2, 1)

class FFN(nn.Module):

  def __init__(self, d_model):

      super(FFN, self).__init__()

      self.ffn = nn.Sequential(nn.Linear(d_model, 4 * d_model), nn.ReLU(),

                                         nn.Linear(4 * d_model, d_model))

      self.layer_norm = nn.LayerNorm(d_model)

  def forward(self, x, masks, lengths):

      x = self.layer_norm(x + self.ffn(x))

      return x

class Identity(nn.Module):

  def __init__(self, d_model):

      super(Identity, self).__init__()

  def forward(self, x, masks, lengths):

      return x

class SelfAttention(nn.Module):

    def __init__(self, in_feature, num_head=4, dropout=0.1):

        super(SelfAttention, self).__init__()

        self.in_feature = in_feature

        self.num_head = num_head

        self.size_per_head = in_feature // num_head

        self.out_dim = num_head * self.size_per_head

        assert self.size_per_head * num_head == in_feature

        self.q_linear = nn.Linear(in_feature, in_feature, bias=False)

        self.k_linear = nn.Linear(in_feature, in_feature, bias=False)

        self.v_linear = nn.Linear(in_feature, in_feature, bias=False)

        self.fc = nn.Linear(in_feature, in_feature, bias=False)

        self.dropout = nn.Dropout(dropout)

        self.layer_norm = nn.LayerNorm(in_feature)

    def forward(self, x, attn_mask, lengths):

        batch_size = x.size(0)

        res = x

        query = self.q_linear(x)

        key = self.k_linear(x)

        value = self.v_linear(x)

        query = query.view(batch_size, self.num_head, -1, self.size_per_head)

        key = key.view(batch_size, self.num_head, -1, self.size_per_head)

        value = value.view(batch_size, self.num_head, -1, self.size_per_head)

        scale = np.sqrt(self.size_per_head)

        energy = torch.matmul(query, key.permute(0, 1, 3, 2)) / scale

        attention = torch.softmax(energy, dim=-1)

        x = torch.matmul(attention, value)

        x = x.permute(0, 2, 1, 3).contiguous()

        x = x.view(batch_size, -1, self.in_feature)

        x = self.fc(x)

        x = self.dropout(x)

        x += res

        x = self.layer_norm(x)

        return x

class Attention_s1(nn.Module):

    def __init__(self, in_feature, num_head=4, dropout=0.1):

        super(Attention_s1, self).__init__()

        self.in_feature = in_feature

        self.num_head = num_head

        self.size_per_head = in_feature // num_head

        self.out_dim = num_head * self.size_per_head

        assert self.size_per_head * num_head == in_feature

        self.q_linear = nn.Linear(in_feature, in_feature, bias=False)

        self.k_linear = nn.Linear(in_feature, in_feature, bias=False)

        self.v_linear = nn.Linear(in_feature, in_feature, bias=False)

        self.fc = nn.Linear(in_feature, in_feature, bias=False)

        self.dropout = nn.Dropout(dropout)

        self.layer_norm = nn.LayerNorm(in_feature)

    def forward(self, s, x1, x2):

        batch_size = x1.size(0)

        s = s.unsqueeze(1)

        res = s

        query = self.q_linear(s)

        key = self.k_linear(x1)

        value = self.v_linear(x1)

        query = query.view(batch_size, self.num_head, -1, self.size_per_head)

        key = key.view(batch_size, self.num_head, -1, self.size_per_head)

        value = value.view(batch_size, self.num_head, -1, self.size_per_head)

        scale = np.sqrt(self.size_per_head)

        energy = torch.matmul(query, key.permute(0, 1, 3, 2)) / scale

        attention = torch.softmax(energy, dim=-1)

        x = torch.matmul(attention, value)

        x = x.permute(0, 2, 1, 3).contiguous()

        x = x.view(batch_size, -1, self.in_feature)

        x = self.fc(x)

        x = self.dropout(x)

        x += res

        x = self.layer_norm(x)

        return x.squeeze()

class Attention_s2(nn.Module):

    def __init__(self, in_feature, num_head=4, dropout=0.1):

        super(Attention_s2, self).__init__()

        self.in_feature = in_feature

        self.num_head = num_head

        self.size_per_head = in_feature // num_head

        self.out_dim = num_head * self.size_per_head

        assert self.size_per_head * num_head == in_feature

        self.q_linear = nn.Linear(in_feature, in_feature, bias=False)

        self.k_linear = nn.Linear(in_feature, in_feature, bias=False)

        self.v_linear = nn.Linear(in_feature, in_feature, bias=False)

        self.fc = nn.Linear(in_feature, in_feature, bias=False)

        self.dropout = nn.Dropout(dropout)

        self.layer_norm = nn.LayerNorm(in_feature)

    def forward(self, s, x1, x2):

        batch_size = x2.size(0)

        s = s.unsqueeze(1)

        res = s

        query = self.q_linear(s)

        key = self.k_linear(x2)

        value = self.v_linear(x2)

        query = query.view(batch_size, self.num_head, -1, self.size_per_head)

        key = key.view(batch_size, self.num_head, -1, self.size_per_head)

        value = value.view(batch_size, self.num_head, -1, self.size_per_head)

        scale = np.sqrt(self.size_per_head)

        energy = torch.matmul(query, key.permute(0, 1, 3, 2)) / scale

        attention = torch.softmax(energy, dim=-1)

        x = torch.matmul(attention, value)

        x = x.permute(0, 2, 1, 3).contiguous()

        x = x.view(batch_size, -1, self.in_feature)

        x = self.fc(x)

        x = self.dropout(x)

        x += res

        x = self.layer_norm(x)

        return x.squeeze()

class Attention_x(nn.Module):

    def __init__(self, in_feature, num_head=4, dropout=0.1):

        super(Attention_x, self).__init__()

        self.in_feature = in_feature

        self.num_head = num_head

        self.size_per_head = in_feature // num_head

        self.out_dim = num_head * self.size_per_head

        assert self.size_per_head * num_head == in_feature

        self.q_linear = nn.Linear(in_feature, in_feature, bias=False)

        self.k_linear = nn.Linear(in_feature, in_feature, bias=False)

        self.v_linear = nn.Linear(in_feature, in_feature, bias=False)

        self.fc = nn.Linear(in_feature, in_feature, bias=False)

        self.dropout = nn.Dropout(dropout)

        self.layer_norm = nn.LayerNorm(in_feature)

    def forward(self, current_x, s, other_x):

        batch_size = current_x.size(0)

        res = current_x

        query = self.q_linear(current_x)

        key = self.k_linear(other_x)

        value = self.v_linear(other_x)

        query = query.view(batch_size, self.num_head, -1, self.size_per_head)

        key = key.view(batch_size, self.num_head, -1, self.size_per_head)

        value = value.view(batch_size, self.num_head, -1, self.size_per_head)

        scale = np.sqrt(self.size_per_head)

        energy = torch.matmul(query, key.permute(0, 1, 3, 2)) / scale

        attention = torch.softmax(energy, dim=-1)

        x = torch.matmul(attention, value)

        x = x.permute(0, 2, 1, 3).contiguous()

        x = x.view(batch_size, -1, self.in_feature)

        x = self.fc(x)

        x = self.dropout(x)

        x += res

        x = self.layer_norm(x)

        return x

class RNN(nn.Module):

    def __init__(self, d_model):

        super(RNN, self).__init__()

        self.rnn = nn.GRU(d_model, d_model, num_layers=1, batch_first=True)

    def forward(self, x, masks, lengths):

        rnn_input = x

        rnn_output, _ = self.rnn(rnn_input)

        return rnn_output

class MaxPoolLayer(nn.Module):

    """

    A layer that performs max pooling along the sequence dimension

    """

    def __init__(self):

        super().__init__()

    def forward(self, inputs, mask_or_lengths=None):

        """

        inputs: tensor of shape (batch_size, seq_len, hidden_size)

        mask_or_lengths: tensor of shape (batch_size) or (batch_size, seq_len)

        returns: tensor of shape (batch_size, hidden_size)

        """

        bs, sl, _ = inputs.size()

        if mask_or_lengths is not None:

            if len(mask_or_lengths.size()) == 1:

                mask = (torch.arange(sl, device=inputs.device).unsqueeze(0).expand(bs, sl) >= mask_or_lengths.unsqueeze(

                    1))

            else:

                mask = mask_or_lengths

            inputs = inputs.masked_fill(mask.unsqueeze(-1).expand_as(inputs), float('-inf'))

        max_pooled = inputs.max(1)[0]

        return max_pooled

def prroc(testy, probs):

    precision, recall, thresholds = sklearn.metrics.precision_recall_curve(testy, probs)

    auc = auc(recall, precision)

    return auc