# -*- coding:utf-8 -*-

import torch
import torch.nn as nn
import numpy as np
from models.BasicModule import BasicModule
from layers.DynamicLSTM import DynamicLSTM
class TD_LSTM(BasicModule):
	
	def __init__(self, embedding_matrix, opt):
		super(TD_LSTM, self).__init__()
		self.model_name = "td_lstm"
		self.embed = nn.Embedding.from_pretrained(torch.FloatTensor(embedding_matrix))
		self.lstm_l = DynamicLSTM(opt.embed_dim, opt.hidden_dim, num_layer=1, batch_first=True)
		self.lstm_r = DynamicLSTM(opt.embed_dim, opt.hidden_dim, num_layer=1, batch_first=True)
		self.dense = nn.Linear(opt.hidden_dim * 2, opt.polarity_dim)
		
	def forward(self, inputs):
		x_l, x_r = inputs[0], inputs[1]
		# x_r = torch.cuda.LongTensor(list(map(lambda x: x[::-1], x_r.data.tolist())))
		x_l_len, x_r_len = torch.sum(x_l != 0, dim=-1), torch.sum(x_r != 0, dim=-1)
		x_l, x_r = self.embed(x_l), self.embed(x_r)
		_, (h_n_l,_) = self.lstm_l(x_l, x_l_len)
		_, (h_n_r,_) = self.lstm_r(x_r, x_r_len)
		h_n = torch.cat((h_n_l[0], h_n_r[0]), dim=-1)
		out =self.dense(h_n)
		return out
		