import math
import pickle

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable
from operations import *
import model
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer

class LSTM_bimodal(nn.Module):
    def __init__(self, vocab_size1, vocab_size2,  d_model = 256, dropout=0.5, dropout_emb=0.5, length=48, pretrain = False):
        super().__init__()
        self.embbedding1 = nn.Sequential(nn.Linear(vocab_size1, d_model), nn.ReLU())
        self.linear = nn.Linear(vocab_size2, d_model)
        self.dropout = nn.Dropout(dropout)
        self.emb_dropout = nn.Dropout(dropout_emb)
        self.output_mlp = nn.Sequential(nn.Linear(d_model, 1))
        self.pooler = MaxPoolLayer()
        if pretrain:
            self.rnns = nn.LSTM(d_model, d_model, 1, bidirectional=False, batch_first=True)
        else:
            self.rnns = nn.LSTM(vocab_size1, d_model, 1, bidirectional=False, batch_first=True)
        self.sig = nn.Sigmoid()
        self.ts_encoder = model.LSTM_Encoder(length, vocab_size1, d_model)
        self.linear_2 = nn.Linear(32, d_model)
        self.pretrain = pretrain


    def forward(self, x):
        if self.pretrain == True:
            x = self.ts_encoder(x)[0]
            x = self.emb_dropout(x)

        rnn_output, _ = self.rnns(x)
        x = self.pooler(rnn_output)
        x = self.output_mlp(x)
        x = self.sig(x)
        return x


class Transformer(nn.Module):
    def __init__(self, vocab_size1, vocab_size2, d_model, dropout=0.5, dropout_emb=0.5, length=48, pretrain = False):
        super().__init__()
        self.embbedding1 = nn.Sequential(nn.Linear(vocab_size1, d_model), nn.ReLU())
        self.linear = nn.Linear(vocab_size2, d_model)
        self.dropout = nn.Dropout(dropout)
        self.emb_dropout = nn.Dropout(dropout_emb)
        self.output_mlp = nn.Sequential(nn.Linear(d_model, 1))
        self.pooler = MaxPoolLayer()
        self.attention = SelfAttention(d_model)
        self.ffn = FFN(d_model)
        self.sig = nn.Sigmoid()
        self.ts_encoder = model.LSTM_Encoder(length, vocab_size1, d_model)
        self.linear_2 = nn.Linear(32, d_model)
        self.pretrain = pretrain
    def forward(self, x):
        if self.pretrain == True:
            x = self.ts_encoder(x)[0]
            x = self.emb_dropout(x)
        else:
            x = self.embbedding1(x)

        x = self.attention(x, None, None)
        x = self.ffn(x, None, None)
        x = self.dropout(x)
        x = self.pooler(x)
        x = self.output_mlp(x)
        x = self.sig(x)
        return x
    
class ClinicalT5(nn.Module):
    def __init__(self, d_model = 256):
      super().__init__()

      self.sig = nn.Sigmoid()
      self.t5 =  AutoModelForSeq2SeqLM.from_pretrained("LLM/physionet.org/files/clinical-t5/1.0.0/Clinical-T5-Base").encoder
      self.fc2 = nn.Linear(768, 1)
      self.pooler = MaxPoolLayer()
      self.relu1 = nn.ReLU()
    def forward(self, ts_x, tb_x, input_ids, attention_mask):
      
      text = self.t5(input_ids=input_ids,attention_mask=attention_mask, return_dict=True).last_hidden_state
      sent_emb = torch.mean(text, dim=1)
      sent_emb = self.fc2(sent_emb)
      x = self.sig(sent_emb)
      return x

class Raim(nn.Module):
    def __init__(self, vocab_size1, vocab_size2, vocab_size3, d_model, dropout=0.1, dropout_emb=0.1, length=48):
        super().__init__()
        self.embbedding1 = nn.Sequential(nn.Linear(vocab_size1, d_model), nn.ReLU())
        self.embbedding2 = nn.Sequential(nn.Linear(vocab_size2, d_model), nn.ReLU())
        self.linear = nn.Linear(vocab_size3, d_model)
        self.dropout = nn.Dropout(dropout)
        self.emb_dropout = nn.Dropout(dropout_emb)
        self.output_mlp = nn.Sequential(nn.Linear(d_model, 2))
        self.pooler = MaxPoolLayer()

        self.hidden_size = d_model

        self.rnn = nn.LSTM(d_model, d_model, 2, dropout=0.5)
        self.attn = nn.Linear(10, 10)
        self.attn1 = nn.Linear(60, 10)

        self.dense_h = nn.Linear(d_model, 1)
        self.softmax = nn.Softmax(dim=1)
        self.hidden2label = nn.Linear(d_model, 1)
        self.grucell = nn.GRUCell(d_model, d_model)

        self.mlp_for_x = nn.Linear(d_model, 1, bias=False)
        self.mlp_for_hidden = nn.Linear(d_model, length, bias=True)
        
        self.sigmoid = nn.Sigmoid()


    def init_hidden(self, batch_size):
        return Variable(torch.zeros(batch_size, self.hidden_size))

    def forward(self, x1, x2, s):
        x1 = self.embbedding1(x1)
        x2 = self.embbedding2(x2)
        s = self.linear(s)
        input_seqs = x1 + x2
        x = input_seqs
        self.hidden = self.init_hidden(x.size(0)).to(x.device)
        for i in range(x.size(1)):
            tt = x[:, 0:i + 1, :].reshape(x.size(0), (i + 1) * x[:, 0:i + 1, :].shape[2])
            if i < x.size(1) - 1:
                padding = torch.zeros(x.size(0), x.size(1)*x.size(2) - tt.shape[1]).to(x.device)
                self.temp1 = torch.cat((tt, padding), 1)
            else:
                self.temp1 = tt

            self.input_padded = self.temp1.reshape(x.size(0), x.size(1), x.size(-1))

            #### multuply with guidance #######
            temp_guidance = torch.zeros(x.size(0), x.size(1), 1).to(x.device)

            # temp_guidance[:, 0:i + 1, :] = x2[:, 0:i + 1, 0].unsqueeze(-1)

            if i > 0:

                zero_idx = torch.where(torch.sum(x2[:, :i, 0], dim=1) == 0)
                if len(zero_idx[0]) > 0:
                    temp_guidance[zero_idx[0], :i, 0] = 1

            temp_guidance[:, i, :] = 1

            self.guided_input = torch.mul(self.input_padded, temp_guidance)

            ######### MLP ###########
            self.t1 = self.mlp_for_x(self.guided_input) + self.mlp_for_hidden(self.hidden).reshape(x.size(0), x.size(1), 1)

            ######### softmax-> multiply->  context vector ###########
            self.t1_softmax = self.softmax(self.t1)
            final_output = torch.mul(self.input_padded, self.t1_softmax)

            context_vec = torch.sum(final_output, dim=1)

            self.hx = self.grucell(context_vec, self.hidden)
            self.hidden = self.hx

        y = self.hidden2label(self.hidden + s)
        return self.sigmoid(y)



class DCMN(nn.Module):

    def __init__(self, vocab_size1, vocab_size2, vocab_size3, d_model, dropout=0.1, dropout_emb=0.1, length=48):
        super().__init__()
        self.embbedding1 = nn.Sequential(nn.Conv1d(in_channels=1, out_channels=1, kernel_size=10, stride=5),
                                         nn.ReLU(),
                                         nn.Linear((vocab_size1 - 10) // 5 + 1, d_model))
        self.embbedding2 = nn.Sequential(nn.Conv1d(in_channels=1, out_channels=1, kernel_size=10, stride=5),
                                         nn.ReLU(),
                                         nn.Linear((vocab_size2 - 10) // 5 + 1, d_model))
        self.linear = nn.Linear(vocab_size3, d_model)
        self.batchnorm1 = nn.BatchNorm1d(d_model)
        self.batchnorm2 = nn.BatchNorm1d(d_model)
        self.conv = nn.Conv1d(d_model, d_model, 3, padding=1)
        self.dropout = nn.Dropout(dropout)
        self.emb_dropout = nn.Dropout(dropout_emb)
        self.output_mlp = nn.Sequential(nn.Linear(d_model, 1))
        self.c_emb = nn.LSTM(d_model, d_model, 1, bidirectional=False, batch_first=True)
        self.c_out = nn.LSTM(d_model, d_model, 1, bidirectional=False, batch_first=True)
        self.w_emb = nn.LSTM(d_model, d_model, 1, bidirectional=False, batch_first=True)
        self.w_out = nn.LSTM(d_model, d_model, 1, bidirectional=False, batch_first=True)
        self.linear1 = nn.Linear(d_model, d_model)
        self.linear2 = nn.Linear(d_model, d_model)
        self.linear3 = nn.Linear(d_model, d_model)
        self.linear4 = nn.Linear(d_model, d_model)
        self.gate_linear = nn.Linear(d_model, d_model)
        self.gate_linear2 = nn.Linear(d_model, d_model)
        self.pooler = MaxPoolLayer()
        self.sigmoid = nn.Sigmoid()

    def forward(self, x1, x2, s):
        bs, l, fdim = x1.size()
        x1 = x1.view(bs * l, -1).unsqueeze(1)
        x2 = x2.view(bs * l, -1).unsqueeze(1)
        x1 = self.embbedding1(x1)
        x2 = self.embbedding2(x2)
        x1 = x1.squeeze().view(bs, l, -1)
        x2 = x2.squeeze().view(bs, l, -1)
        s = self.dropout(self.linear(s))
        x1 = self.batchnorm1(x1.permute(0, 2, 1)).permute(0, 2, 1)
        x2 = self.batchnorm2(x2.permute(0, 2, 1)).permute(0, 2, 1)
        wm_embedding_memory, _ = self.w_emb(x1)
        wm_out_query, _ = self.w_out(x1)
        cm_embedding_memory, _ = self.c_emb(x2)
        cm_out_query, _ = self.c_out(x2)
        wm_in = cm_out_query[:, -1]
        cm_in = wm_out_query[:, -1]
        w_embedding_E = self.linear1(wm_embedding_memory)
        w_embedding_F = self.linear2(wm_embedding_memory)
        wm_out = torch.matmul(wm_in.unsqueeze(1), w_embedding_E.permute(0, 2, 1))
        wm_prob = torch.softmax(wm_out, dim=-1)
        wm_contex = torch.matmul(wm_prob, w_embedding_F)
        wm_gate_prob = torch.sigmoid(self.gate_linear(wm_in)).unsqueeze(1)
        wm_dout = wm_contex * wm_gate_prob + wm_in.unsqueeze(1) * (1 - wm_gate_prob)

        c_embedding_E = self.linear3(cm_embedding_memory)
        c_embedding_F = self.linear4(cm_embedding_memory)
        cm_out = torch.matmul(cm_in.unsqueeze(1), c_embedding_E.permute(0, 2, 1))
        cm_prob = torch.softmax(cm_out, dim=-1)
        cm_contex = torch.matmul(cm_prob, c_embedding_F)
        cm_gate_prob = torch.sigmoid(self.gate_linear2(cm_in)).unsqueeze(1)
        cm_dout = cm_contex * cm_gate_prob + cm_in.unsqueeze(1) * (1 - cm_gate_prob)
        output = wm_dout + cm_dout
        output = self.output_mlp(output.squeeze() + s)
        return self.sigmoid(output)


class Mufasa(nn.Module):

    def __init__(self, vocab_size1, vocab_size2, vocab_size3, d_model, dropout=0.1, dropout_emb=0.1, length=48):
        super().__init__()
        self.embbedding1 = nn.Sequential(nn.Conv1d(in_channels=1, out_channels=1, kernel_size=10, stride=5),
                                         nn.ReLU(),
                                         nn.Linear((vocab_size1 - 10) // 5 + 1, d_model))
        self.embbedding2 = nn.Sequential(nn.Conv1d(in_channels=1, out_channels=1, kernel_size=10, stride=5),
                                         nn.ReLU(),
                                         nn.Linear((vocab_size2 - 10) // 5 + 1, d_model))
        self.linear = nn.Linear(vocab_size3, d_model)
        self.linear_conti = nn.Linear(d_model, d_model)
        self.linear_cate = nn.Linear(2*d_model, d_model)
        self.linears = nn.Linear(2 * d_model, d_model)
        self.linear_late = nn.Sequential(nn.Linear(d_model, d_model), nn.ReLU(inplace=False))
        self.dense = nn.Sequential(nn.Linear(d_model, 4*d_model), nn.ReLU(inplace=False), nn.Linear(4*d_model, d_model))
        self.relu = nn.ReLU(inplace=False)
        self.layernorm = nn.LayerNorm(d_model)
        self.layernorm2 = nn.LayerNorm(d_model)
        self.layernorm3 = nn.LayerNorm(d_model)
        self.self_att = SelfAttention(d_model)
        self.self_att2 = SelfAttention(d_model)
        self.conv = nn.Conv1d(d_model, d_model, 3, padding=1)
        self.leaky = nn.LeakyReLU(inplace=False)
        self.pooler = MaxPoolLayer()
        self.output_mlp = nn.Sequential(nn.Linear(d_model, 1))
        self.sigmoid = nn.Sigmoid()


    def forward(self, x1, x2, s):
        bs, l, fdim = x1.size()
        x1 = x1.view(bs * l, -1).unsqueeze(1).clone()
        x2 = x2.view(bs * l, -1).unsqueeze(1).clone()
        x1 = self.embbedding1(x1)
        x2 = self.embbedding2(x2)
        x1 = x1.squeeze().view(bs, l, -1)
        x2 = x2.squeeze().view(bs, l, -1)
        s = self.linear(s)
        continues_res = x2
        continues_hs = self.layernorm(x2)
        continues_hs = self.self_att(continues_hs, None, None)
        continues_hs = self.leaky(continues_hs)
        continues_hs = continues_res + continues_hs
        continuous_res = continues_hs
        continues_hs = self.layernorm(continues_hs)
        continues_hs = self.linear_conti(continues_hs)
        continues_hs = self.relu(continues_hs)
        continues_hs = continuous_res + continues_hs
        categorical_res = x1
        categorical_hs = self.layernorm2(x1)
        categorical_hs = self.self_att2(categorical_hs, None, None)
        categorical_hs = torch.cat((categorical_hs, categorical_res), dim=-1)
        categorical_res = categorical_hs.clone()
        categorical_hs = self.linear_cate(categorical_hs)
        categorical_hs = self.relu(categorical_hs)
        categorical_res = self.linears(categorical_res)
        categorical_hybrid_point = categorical_hs + categorical_res
        categorical_late_point = self.linear_late(categorical_res)
        temp = s.unsqueeze(1).clone()
        fusion_hs = temp.expand_as(categorical_hybrid_point) + categorical_hybrid_point
        fusion_res = fusion_hs
        fusion_hs = self.layernorm3(fusion_hs)
        fusion_branch = self.conv(fusion_hs.permute(0, 2, 1)).permute(0, 2, 1)
        out = fusion_res + fusion_hs + fusion_branch + categorical_late_point + continues_hs
        out = self.pooler(out)
        out = self.output_mlp(out)
        return self.sigmoid(out)

if __name__ == '__main__':
    model = Transformer(1318, 73, 256)
    x1 = torch.randn((32, 48, 1318))
    s = torch.randn((32, 73))
    print(model(x1, s).size())