import tensorflow as tf

class Model(object):
	def __init__(self, input, is_training,hidden_size,vocab_size,num_layers,dropout=0.5,init_scale=0.05):
		self.is_training = is_training
		self.input_obj = input
		self.batch_size = input.batch_size
		self.num_steps = input.num_steps
		self.hidden_size = hidden_size

		with tf.device("/cpu:0"):
			embedding = tf.Variable(tf.random_uniform([vocab_size,self.hidden_size], -init_scale,init_scale))
			inputs = tf.nn.embedding_lookup(embedding,self.input_obj.input_data)

		if is_training and dropout <1:
			inputs = tf.nn.dropout(inputs, dropout)


		self.init_state = tf.placeholder(tf.float32,[num_layers,2,self.batch_size,self.hidden_size])
		state_per_layer_list = tf.unstack(self.init_state,axis=0)

		rnn_tuple_state = tuple([tf.contrib.rnn.LSTMStateTuple(state_per_layer_list[idx][0],
						state_per_layer_list[idx][1]) for idx in range(num_layers)])

		cell = tf.contrib.rnn.LSTMCell(hidden_size)

		if is_training and dropout <1:
			cell = tf.contrib.rnn.DropoutWrapper(cell,output_keep_prob=dropout)

		if num_layers > 1:
			cell = tf.contrib.rnn.MultiRNNCell([cell for _ in range(num_layers)],state_is_tuple=True)

		output, self.state = tf.nn.dynamic_rnn(cell, inputs,dtype=tf.float32, inital_state=rnn_tuple_state)

		output = tf.reshape(output, [-1, hidden_size])

		softmax_w = tf.Variable(tf.random_uniform([hidden_size, vocab_size], -init_state,init_state))
		softmax_b = tf.Variable(tf.random_uniform([vocab_size],-init_state,init_state))

		logits = tf.nn.xw_plus_b(output, softmax_w,softmax_b)

		logits = tf.reshape(logits,[self.batch_size, self.num_steps,vocab_size])
		loss = tf.contrib.seq2seq.sequence_loss(logits, self.input_obj.targets,tf.one([self.batch_size,self.num_steps],
							dtype=tf.float32,average_across_timesteps=False,
							average_across_batch=True))

		self.cost = tf.reduce_sum(loss)
		self.softmax_out = tf.nn.softmax(tf.reshape(logits,[-1,vocab_size]))

		self.predict = tf.cast(tf.argmax(self.softmax_out,axis=1),tf.int32)
		correct_prediction = tf.equal(self.predict, tf.reshape(self.input_obj.targets,[-1]))

		self.accuracy = tf.reduce_mean(tf.cast(correct_prediction,tf.float32))

		if not is_training:
			return

		self.learning_rate = tf.Variable(0.0, trainable=False)
		tvars = tf.trainable_variable()

		grads, _ = tf.clip_by_global_norm(tf.gradients(self.cost,tvars), 5)

		optimizer = tf.train.GradientDescentOptimizer(self.learning_rate)

		self.train_op = optimizer.apply_gradients(zip(grads, tvars),global_step= tf.train.get_or_create_global_step())

		self.new_lr = tf.placeholder(tf.float32, shape=[])
		self.lr_update = tf.assign(self.learning_rate,self.new_lr)
	def assign_lr(self, session,lr_value):
		session.run(self.lr_update, feed_dict ={self.new_lr:lr_value})	
		
