"""

Patient2Vec: a self-attentive representation learning framework

author: Jinghe Zhang

"""

import torch

import torch.nn as nn

import torch.optim as optim

from torch.autograd import Variable

import numpy as np

class Patient2Vec(nn.Module):

    """

    Self-attentive representation learning framework,

    including convolutional embedding layer,

    recurrent autoencoder with an encoder, recurrent module, and a decoder.

    In addition, a linear layer is on top of each decode step and the weights are shared at these step.

    """

    def __init__(self, input_size, hidden_size, n_layers, att_dim, initrange,

                 output_size, rnn_type, seq_len, pad_size, n_filters, bi, dropout_p=0.5):

        """

        Initilize a recurrent model

        :param input_size: int

        :param hidden_size: int

        :param n_layers: number of layers; int

        :param att_dim: dimension of the attention; int

        :param initrange: upper bound of the initial weights; symmetric

        :param output_size: int

        :param rnn_type: str, such as 'GRU'

        :param seq_len: length of the sequence; int

        :param pad_size: padding size; int

        :param n_filters: number of hops; int

        :param bi: bidirectional; bool

        :param dropout_p: dropout rate; float

        """

        super(Patient2Vec, self).__init__()

        self.initrange = initrange

        # convolution

        self.b = 1

        if bi:

            self.b = 2

        self.conv = nn.Conv1d(in_channels=1, out_channels=1, kernel_size=input_size, stride=2)

        self.conv2 = nn.Conv1d(in_channels=1, out_channels=n_filters, kernel_size=hidden_size * self.b, stride=2)

        # Bidirectional RNN

        self.rnn = getattr(nn, rnn_type)(input_size, hidden_size, n_layers, dropout=dropout_p,

                                         batch_first=True, bias=True, bidirectional=bi)

        # initialize 2-layer attention weight matrics

        self.att_w1 = nn.Linear(hidden_size * self.b, att_dim, bias=False)

        # final linear layer

        self.linear = nn.Linear(hidden_size * self.b * n_filters + 3, output_size, bias=True)

        self.func_softmax = nn.Softmax()

        self.func_sigmoid = nn.Sigmoid()

        self.func_tanh = nn.Hardtanh(0, 1)

        # Add dropout

        self.dropout_p = dropout_p

        self.dropout = nn.Dropout(p=self.dropout_p)

        self.init_weights()

        self.pad_size = pad_size

        self.input_size = input_size

        self.hidden_size = hidden_size

        self.output_size = output_size

        self.n_layers = n_layers

        self.seq_len = seq_len

        self.n_filters = n_filters

    def init_weights(self):

        """

        weight initialization

        """

        for param in self.parameters():

            param.data.uniform_(-self.initrange, self.initrange)

    def convolutional_layer(self, inputs):

        convolution_all = []

        conv_wts = []

        for i in range(self.seq_len):

            convolution_one_month = []

            for j in range(self.pad_size):

                convolution = self.conv(torch.unsqueeze(inputs[:, i, j], dim=1))

                convolution_one_month.append(convolution)

            convolution_one_month = torch.stack(convolution_one_month)

            convolution_one_month = torch.squeeze(convolution_one_month, dim=3)

            convolution_one_month = torch.transpose(convolution_one_month, 0, 1)

            convolution_one_month = torch.transpose(convolution_one_month, 1, 2)

            convolution_one_month = torch.squeeze(convolution_one_month, dim=1)

            convolution_one_month = self.func_tanh(convolution_one_month)

            convolution_one_month = torch.unsqueeze(convolution_one_month, dim=1)

            vec = torch.bmm(convolution_one_month, inputs[:, i])

            convolution_all.append(vec)

            conv_wts.append(convolution_one_month)

        convolution_all = torch.stack(convolution_all, dim=1)

        convolution_all = torch.squeeze(convolution_all, dim=2)

        conv_wts = torch.squeeze(torch.stack(conv_wts, dim=1), dim=2)

        return convolution_all, conv_wts

    def encode_rnn(self, embedding, batch_size):

        self.weight = next(self.parameters()).data

        init_state = (Variable(self.weight.new(self.n_layers * self.b, batch_size, self.hidden_size).zero_()))

        embedding = self.dropout(embedding)

        outputs_rnn, states_rnn = self.rnn(embedding, init_state)

        return outputs_rnn

    def add_beta_attention(self, states, batch_size):

        # beta attention

        att_wts = []

        for i in range(self.seq_len):

            m1 = self.conv2(torch.unsqueeze(states[:, i], dim=1))

            att_wts.append(torch.squeeze(m1, dim=2))

        att_wts = torch.stack(att_wts, dim=2)

        att_beta = []

        for i in range(self.n_filters):

            a0 = self.func_softmax(att_wts[:, i])

            att_beta.append(a0)

        att_beta = torch.stack(att_beta, dim=1)

        context = torch.bmm(att_beta, states)

        context = context.view(batch_size, -1)

        return att_beta, context

    def forward(self, inputs, inputs_other, batch_size):

        """

        the recurrent module

        """

        # Convolutional

        convolutions, alpha = self.convolutional_layer(inputs)

        # RNN

        states_rnn = self.encode_rnn(convolutions, batch_size)

        # Add attentions and get context vector

        beta, context = self.add_beta_attention(states_rnn, batch_size)

        # Final linear layer with demographic info added as extra variables

        context_v2 = torch.cat((context, inputs_other), 1)

        linear_y = self.linear(context_v2)

        out = self.func_softmax(linear_y)

        return out, alpha, beta

def get_loss(pred, y, criterion, mtr, a=0.5):

    """

    To calculate loss

    :param pred: predicted value

    :param y: actual value

    :param criterion: nn.CrossEntropyLoss

    :param mtr: beta matrix

    """

    mtr_t = torch.transpose(mtr, 1, 2)

    aa = torch.bmm(mtr, mtr_t)

    loss_fn = 0

    for i in range(aa.size()[0]):

        aai = torch.add(aa[i, ], Variable(torch.neg(torch.eye(mtr.size()[1]))))

        loss_fn += torch.trace(torch.mul(aai, aai).data)

    loss_fn /= aa.size()[0]

    loss = torch.add(criterion(pred, y), Variable(torch.FloatTensor([loss_fn * a])))

    return loss