#!/usr/bin/env python
# -*- coding: UTF-8 -*-

# Copyright 2017 Timothy Dozat
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import six

import numpy as np
import tensorflow as tf

from .base_vocabs import BaseVocab
from . import conllu_vocabs as cv

from parser.neural import nn, nonlin, classifiers
import pdb
#***************************************************************
class SecondOrderLBPVocab(BaseVocab):
	""""""

	#=============================================================
	def __init__(self, *args, **kwargs):
		""""""

		super(SecondOrderLBPVocab, self).__init__(*args, **kwargs)

		self.PAD_STR = '_'
		self.PAD_IDX = -1
		self.ROOT_STR = '0'
		self.ROOT_IDX = 0
		return

	#=============================================================
	def add(self, token):
		""""""

		return self.index(token)

	#=============================================================
	def token(self, index):
		""""""

		if index > -1:
			return str(index)
		else:
			return '_'

	#=============================================================
	def index(self, token):
		""""""

		if token != '_':
			return int(token)
		else:
			return -1

	#=============================================================
	def get_root(self):
		""""""

		return self.ROOT_STR

	#=============================================================
	def get_bilinear_classifier(self, layer, token_weights, variable_scope=None, reuse=False):
		""""""

		recur_layer = layer
		hidden_keep_prob = 1 if reuse else self.hidden_keep_prob
		hidden_func = self.hidden_func
		hidden_size = self.hidden_size
		add_linear = self.add_linear
		linearize = self.linearize
		distance = self.distance
		n_splits = 2*(1+linearize+distance)
		with tf.variable_scope(variable_scope or self.field):
			for i in six.moves.range(0, self.n_layers-1):
				with tf.variable_scope('FC-%d' % i):
					layer = classifiers.hidden(layer, n_splits*hidden_size,
																		hidden_func=hidden_func,
																		hidden_keep_prob=hidden_keep_prob)
			with tf.variable_scope('FC-top'):
				layers = classifiers.hiddens(layer, n_splits*[hidden_size],
																		hidden_func=hidden_func,
																		hidden_keep_prob=hidden_keep_prob)
			layer1, layer2 = layers.pop(0), layers.pop(0)
			if linearize:
				lin_layer1, lin_layer2 = layers.pop(0), layers.pop(0)
			if distance:
				dist_layer1, dist_layer2 = layers.pop(0), layers.pop(0)

			with tf.variable_scope('Attention'):
				if self.diagonal:
					logits, _ = classifiers.diagonal_bilinear_attention(
						layer1, layer2,
						hidden_keep_prob=hidden_keep_prob,
						add_linear=add_linear)
					if linearize:
						with tf.variable_scope('Linearization'):
							lin_logits = classifiers.diagonal_bilinear_discriminator(
								lin_layer1, lin_layer2,
								hidden_keep_prob=hidden_keep_prob,
								add_linear=add_linear)
					if distance:
						with tf.variable_scope('Distance'):
							dist_lamda = 1+tf.nn.softplus(classifiers.diagonal_bilinear_discriminator(
								dist_layer1, dist_layer2,
								hidden_keep_prob=hidden_keep_prob,
								add_linear=add_linear))
				else:
					logits, _ = classifiers.bilinear_attention(
						layer1, layer2,
						hidden_keep_prob=hidden_keep_prob,
						add_linear=add_linear)
					if linearize:
						with tf.variable_scope('Linearization'):
							lin_logits = classifiers.bilinear_discriminator(
								lin_layer1, lin_layer2,
								hidden_keep_prob=hidden_keep_prob,
								add_linear=add_linear)
					if distance:
						with tf.variable_scope('Distance'):
							dist_lamda = 1+tf.nn.softplus(classifiers.bilinear_discriminator(
								dist_layer1, dist_layer2,
								hidden_keep_prob=hidden_keep_prob,
								add_linear=add_linear))

				#-----------------------------------------------------------
				# Process the targets
				targets = self.placeholder
				shape = tf.shape(layer1)
				batch_size, bucket_size = shape[0], shape[1]
				# (1 x m)
				ids = tf.expand_dims(tf.range(bucket_size), 0)
				# (1 x m) -> (1 x 1 x m)
				head_ids = tf.expand_dims(ids, -2)
				# (1 x m) -> (1 x m x 1)
				dep_ids = tf.expand_dims(ids, -1)
				if linearize:
					# Wherever the head is to the left
					# (n x m), (1 x m) -> (n x m)
					lin_targets = tf.to_float(tf.less(targets, ids))
					# cross-entropy of the linearization of each i,j pair
					# (1 x 1 x m), (1 x m x 1) -> (n x m x m)
					lin_ids = tf.tile(tf.less(head_ids, dep_ids), [batch_size, 1, 1])
					# (n x 1 x m), (n x m x 1) -> (n x m x m)
					lin_xent = -tf.nn.softplus(tf.where(lin_ids, -lin_logits, lin_logits))
					# add the cross-entropy to the logits
					# (n x m x m), (n x m x m) -> (n x m x m)
					logits += tf.stop_gradient(lin_xent)
				if distance:
					# (n x m) - (1 x m) -> (n x m)
					dist_targets = tf.abs(targets - ids)
					# KL-divergence of the distance of each i,j pair
					# (1 x 1 x m) - (1 x m x 1) -> (n x m x m)
					dist_ids = tf.to_float(tf.tile(tf.abs(head_ids - dep_ids), [batch_size, 1, 1]))+1e-12
					# (n x m x m), (n x m x m) -> (n x m x m)
					#dist_kld = (dist_ids * tf.log(dist_lamda / dist_ids) + dist_ids - dist_lamda)
					dist_kld = -tf.log((dist_ids - dist_lamda)**2/2 + 1)
					# add the KL-divergence to the logits
					# (n x m x m), (n x m x m) -> (n x m x m)
					logits += tf.stop_gradient(dist_kld)

				#-----------------------------------------------------------
				# Compute probabilities/cross entropy
				# (n x m) + (m) -> (n x m)
				non_pads = tf.to_float(token_weights) + tf.to_float(tf.logical_not(tf.cast(tf.range(bucket_size), dtype=tf.bool)))
				# (n x m x m) o (n x 1 x m) -> (n x m x m)
				probabilities = tf.nn.softmax(logits) * tf.expand_dims(non_pads, -2)
				# (n x m), (n x m x m), (n x m) -> ()
				loss = tf.losses.sparse_softmax_cross_entropy(
					targets,
					logits,
					weights=token_weights)
				# (n x m) -> (n x m x m x 1)
				one_hot_targets = tf.expand_dims(tf.one_hot(targets, bucket_size), -1)
				# (n x m) -> ()
				n_tokens = tf.to_float(tf.reduce_sum(token_weights))
				if linearize:
					# (n x m x m) -> (n x m x 1 x m)
					lin_xent_reshaped = tf.expand_dims(lin_xent, -2)
					# (n x m x 1 x m) * (n x m x m x 1) -> (n x m x 1 x 1)
					lin_target_xent = tf.matmul(lin_xent_reshaped, one_hot_targets)
					# (n x m x 1 x 1) -> (n x m)
					lin_target_xent = tf.squeeze(lin_target_xent, [-1, -2])
					# (n x m), (n x m), (n x m) -> ()
					loss -= tf.reduce_sum(lin_target_xent*tf.to_float(token_weights)) / (n_tokens + 1e-12)
				if distance:
					# (n x m x m) -> (n x m x 1 x m)
					dist_kld_reshaped = tf.expand_dims(dist_kld, -2)
					# (n x m x 1 x m) * (n x m x m x 1) -> (n x m x 1 x 1)
					dist_target_kld = tf.matmul(dist_kld_reshaped, one_hot_targets)
					# (n x m x 1 x 1) -> (n x m)
					dist_target_kld = tf.squeeze(dist_target_kld, [-1, -2])
					# (n x m), (n x m), (n x m) -> ()
					loss -= tf.reduce_sum(dist_target_kld*tf.to_float(token_weights)) / (n_tokens + 1e-12)

				#-----------------------------------------------------------
				# Compute predictions/accuracy
				# (n x m x m) -> (n x m)
				predictions = tf.argmax(logits, axis=-1, output_type=tf.int32)
				# (n x m) (*) (n x m) -> (n x m)
				correct_tokens = nn.equal(targets, predictions) * token_weights
				# (n x m) -> (n)
				tokens_per_sequence = tf.reduce_sum(token_weights, axis=-1)
				# (n x m) -> (n)
				correct_tokens_per_sequence = tf.reduce_sum(correct_tokens, axis=-1)
				# (n), (n) -> (n)
				correct_sequences = nn.equal(tokens_per_sequence, correct_tokens_per_sequence)

		#-----------------------------------------------------------
		# Populate the output dictionary
		outputs = {}
		outputs['recur_layer'] = recur_layer
		outputs['unlabeled_targets'] = self.placeholder
		outputs['probabilities'] = probabilities
		outputs['unlabeled_loss'] = loss
		outputs['loss'] = loss

		outputs['unlabeled_predictions'] = predictions
		outputs['predictions'] = predictions
		outputs['correct_unlabeled_tokens'] = correct_tokens
		outputs['n_correct_unlabeled_tokens'] = tf.reduce_sum(correct_tokens)
		outputs['n_correct_unlabeled_sequences'] = tf.reduce_sum(correct_sequences)
		outputs['n_correct_tokens'] = tf.reduce_sum(correct_tokens)
		outputs['n_correct_sequences'] = tf.reduce_sum(correct_sequences)
		return outputs

	#=============================================================
	def __getitem__(self, key):
		if isinstance(key, six.string_types):
			if key == '_':
				return -1
			else:
				return int(key)
		elif isinstance(key, six.integer_types + (np.int32, np.int64)):
			if key > -1:
				return str(key)
			else:
				return '_'
		elif hasattr(key, '__iter__'):
			return [self[k] for k in key]
		else:
			raise ValueError('key to IndexVocab.__getitem__ must be (iterable of) string or integer')
		return
	#=============================================================
	@property
	def distance(self):
		return self._config.getboolean(self, 'distance')
	@property
	def linearize(self):
		return self._config.getboolean(self, 'linearize')
	@property
	def decomposition_level(self):
		return self._config.getint(self, 'decomposition_level')
	@property
	def diagonal(self):
		return self._config.getboolean(self, 'diagonal')
	@property
	def add_linear(self):
		return self._config.getboolean(self, 'add_linear')
	@property
	def n_layers(self):
		return self._config.getint(self, 'n_layers')
	@property
	def hidden_size(self):
		return self._config.getint(self, 'hidden_size')
	@property
	def hidden_keep_prob(self):
		return self._config.getfloat(self, 'hidden_keep_prob')
	@property
	def hidden_keep_prob_tri(self):
		try:
			return self._config.getfloat(self, 'hidden_keep_prob_tri')
		except:
			return self._config.getfloat(self, 'hidden_keep_prob')

	@property
	def hidden_func(self):
		hidden_func = self._config.getstr(self, 'hidden_func')
		if hasattr(nonlin, hidden_func):
			return getattr(nonlin, hidden_func)
		else:
			raise AttributeError("module '{}' has no attribute '{}'".format(nonlin.__name__, hidden_func))
	@property
	def num_iteration(self):
		return self._config.getfloat(self, 'num_iteration')
	@property
	def discriminator2(self):
		return self._config.getboolean(self,'discriminator2')
	@property
	def sibling_only(self):
		return self._config.getboolean(self,'sibling_only')
	@property
	def self_minus(self):
		return self._config.getboolean(self,'self_minus')
	@property
	def use_sib(self):
		return self._config.getboolean(self,'use_sib')
	@property
	def use_gp(self):
		return self._config.getboolean(self,'use_gp')
	@property
	def use_cop(self):
		return self._config.getboolean(self,'use_cop')
	@property
	def transposed(self):
		return self._config.getboolean(self,'transposed')
	@property
	def unary_weight(self):
		if self._config.getboolean(self,'unary_weight'):
			return int(self.use_cop)+int(self.use_sib)+2*int(self.use_gp)
		else:
			return 1
	@property
	def new_potential(self):
		return self._config.getboolean(self,'new_potential')
	@property
	def separate_embed(self):
		try:
			return self._config.getboolean(self,'separate_embed')
		except:
			return False
	@property
	def old_trilin(self):
		try:
			return self._config.getboolean(self,'old_trilin')
		except:
			return False
	@property
	def remove_loop(self):
		try:
			return self._config.getboolean(self,'remove_loop')
		except:
			return False
	@property
	def combine_loss(self):
		try:
			return self._config.getboolean(self,'combine_loss')
		except:
			return False
	@property
	def loss_weight(self):
		try:
			return self._config.getfloat(self,'loss_weight')
		except:
			return 0.5
	@property
	def loss_weight_unary(self):
		try:
			return self._config.getfloat(self,'loss_weight_unary')
		except:
			return 0.5
	@property
	def test_new_potential(self):
		try:
			return self._config.getboolean(self,'test_new_potential')
		except:
			return False
	@property
	def layer_mask(self):
		try:
			return self._config.getboolean(self,'layer_mask')
		except:
			return False
	@property
	def normalize(self):
		try:
			return self._config.getboolean(self,'normalize')
		except:
			return False
	@property
	def use_unary_hidden(self):
		try:
			return self._config.getboolean(self,'use_unary_hidden')
		except:
			return False
	@property
	def unary_hidden(self):
		try:
			return self._config.getint(self,'unary_hidden')
		except:
			return self._config.getint(self,'hidden_size')
	@property
	def tri_std(self):
		try:
			return self._config.getfloat(self,'tri_std')
		except:
			return 0.01
	@property
	def tri_std_unary(self):
		try:
			return self._config.getfloat(self,'tri_std_unary')
		except:
			return 0.5
	@property
	def hidden_k(self):
		try:
			return self._config.getint(self,'hidden_k')
		except:
			return 200
	@property
	def two_gpu(self):
		try:
			if self._config.get('GraphParserNetwork','two_gpu')=='True':
				return True
			else:
				return False
		except:
			return False
	@property
	def remove_root_child(self):
		try:
			return self._config.getboolean(self,'remove_root_child')
		except:
			return False
	@property
	def compare_precision(self):
		try:
			if self._config.get('DEFAULT', 'tb')=='ptb' or self._config.get('DEFAULT', 'tb')=='ctb':
				return True
			else:
				return False
		except:
			return False
#***************************************************************
class GraphSecondLBPVocab(SecondOrderLBPVocab): #second order trilinear classifier
	""""""

	_depth = -1

	#=============================================================
	def __init__(self, *args, **kwargs):
		""""""

		kwargs['placeholder_shape'] = [None, None, None]
		super(GraphSecondLBPVocab, self).__init__(*args, **kwargs)
		return

	#=============================================================
	def get_bilinear_discriminator(self, layer, token_weights, variable_scope=None, reuse=False, debug=False, token_weights4D=None):
		""""""
		#in fact here is get_trilinear_discriminator
		outputs = {}
		recur_layer = layer
		hidden_keep_prob = 1 if reuse else self.hidden_keep_prob
		hidden_keep_prob_tri = 1 if reuse else self.hidden_keep_prob_tri
		add_linear = self.add_linear
		#here set n_splits to be three
		if self.separate_embed:
			n_splits = 9*(1+self.linearize+self.distance)
		else:
			n_splits = 3*(1+self.linearize+self.distance)
		with tf.variable_scope(variable_scope or self.field):
			for i in six.moves.range(0, self.n_layers-1):#number of layers of FNN? what is this?
				with tf.variable_scope('FC-%d' % i):#here is FNN? did not run
					layer = classifiers.hidden(layer, n_splits*self.hidden_size,
												 hidden_func=self.hidden_func,
												 hidden_keep_prob=hidden_keep_prob)
			with tf.variable_scope('FC-top'):#FNN output and split two layer? FNN+split. Linear transform for a sentence, n_splits is number of features you want
				#this linear transformation contains word information
				if self.use_unary_hidden:
					print('separate unary and binary hidden size')
					if self.separate_embed:
						hidden_list=2*[self.unary_hidden]+(n_splits-2)*[self.hidden_size]
					else:
						hidden_list=2*[self.unary_hidden]+n_splits*[self.hidden_size]
				else:
					hidden_list=n_splits*[self.hidden_size]
				#pdb.set_trace()
				layers = classifiers.hiddens(layer, hidden_list,
											 hidden_func=self.hidden_func,
											 hidden_keep_prob=hidden_keep_prob)
			if self.separate_embed:
				# unary_head + unary_dep + sib_head + sib_dep + gp_head + gp_dep + gp_(head+dep) + cop_head + cop_dep
				unary_layer1, unary_layer2, sib_head, sib_dep, gp_head, gp_dep, gp_headdep, cop_head, cop_dep = layers.pop(0), layers.pop(0), layers.pop(0), layers.pop(0), layers.pop(0)\
																																													, layers.pop(0), layers.pop(0), layers.pop(0), layers.pop(0)
			else:
				# head + dep + (head+dep)
				if self.use_unary_hidden:
					unary_layer1, unary_layer2, layer1, layer2, layer3 = layers.pop(0), layers.pop(0), layers.pop(0), layers.pop(0), layers.pop(0)
				else:
					layer1, layer2, layer3 = layers.pop(0), layers.pop(0), layers.pop(0)
					#pdb.set_trace()
					unary_layer1=layer1
					unary_layer2=layer2
			if self.linearize:#false
				lin_layer1, lin_layer2 = layers.pop(0), layers.pop(0), layers.pop(0)
			if self.distance:#false in graph
				dist_layer1, dist_layer2 = layers.pop(0), layers.pop(0), layers.pop(0)
			#pdb.set_trace()
			if self.layer_mask:
				#pdb.set_trace()
				sentence_mask=tf.expand_dims(tf.cast(tf.transpose(token_weights,[0,2,1])[:,0],dtype=tf.float32),-1)

				unary_layer1=unary_layer1*sentence_mask
				unary_layer2=unary_layer2*sentence_mask
				if self.separate_embed:
					sib_head, sib_dep, gp_head, gp_dep, gp_headdep, cop_head, cop_dep = sib_head*sentence_mask, sib_dep*sentence_mask,\
																						gp_head*sentence_mask, gp_dep*sentence_mask,\
																						gp_headdep*sentence_mask, cop_head*sentence_mask,\
																						cop_dep*sentence_mask
				else:
					if not self.separate_embed:
						layer1=layer1*sentence_mask
						layer2=layer2*sentence_mask
					layer3=layer3*sentence_mask
				pass
			with tf.variable_scope('Discriminator'):
				if self.diagonal:
					logits = classifiers.diagonal_bilinear_discriminator(
						layer1, layer2,
						hidden_keep_prob=hidden_keep_prob,
						add_linear=add_linear)
					if self.linearize:
						with tf.variable_scope('Linearization'):
							lin_logits = classifiers.diagonal_bilinear_discriminator(
								lin_layer1, lin_layer2,
								hidden_keep_prob=hidden_keep_prob,
								add_linear=add_linear)
					if self.distance:
						with tf.variable_scope('Distance'):
							dist_lamda = 1+tf.nn.softplus(classifiers.diagonal_bilinear_discriminator(
								dist_layer1, dist_layer2,
								hidden_keep_prob=hidden_keep_prob,
								add_linear=add_linear))
				else:
					#only run here
					#First order potential and second order potential
					#change it to two label discriminator, because we need a score for the condition of 0 and 1
					#(n x m x 2 x m)

					#with tf.device('/device:GPU:1'):
					'''
					unary, unary_weights = classifiers.bilinear_classifier(
						layer1, layer2, 2,
						hidden_keep_prob=hidden_keep_prob,
						add_linear=add_linear)
					'''
					unary = classifiers.bilinear_classifier(
						unary_layer1, unary_layer2, 2,
						hidden_keep_prob=hidden_keep_prob,
						add_linear=add_linear,target_model='LBP', tri_std=self.tri_std_unary)
					#'''
					'''
					unary = classifiers.bilinear_discriminator(
						layer1, layer2,
						hidden_keep_prob=hidden_keep_prob,
						add_linear=add_linear)
					#'''

					if self.separate_embed:
						print('separate')
						# head dep dep
						if self.test_new_potential:
							print('testing new potential')
							if self.use_sib:
								with tf.variable_scope('Sibling'):
									layer_sib = classifiers.trilinear_discriminator_outer(
										sib_head, sib_dep, sib_dep,
										hidden_keep_prob=hidden_keep_prob_tri,
										add_linear=add_linear, tri_std=self.tri_std, hidden_k=self.hidden_k)
							# head head+dep dep
							if self.use_gp:
								with tf.variable_scope('GrandParents'):
									layer_gp = classifiers.trilinear_discriminator_outer(
										gp_head, gp_headdep, gp_dep,
										hidden_keep_prob=hidden_keep_prob_tri,
										add_linear=add_linear, tri_std=self.tri_std, hidden_k=self.hidden_k)
							# head dep head
							if self.use_cop:
								with tf.variable_scope('CoParents'):
									layer_cop = classifiers.trilinear_discriminator_outer(
										cop_head, cop_dep, cop_head,
										hidden_keep_prob=hidden_keep_prob_tri,
										add_linear=add_linear, tri_std=self.tri_std, hidden_k=self.hidden_k)
						else:
							if self.use_sib:
								with tf.variable_scope('Sibling'):
									layer_sib = classifiers.trilinear_discriminator_new(
										sib_head, sib_dep, sib_dep,
										hidden_keep_prob=hidden_keep_prob_tri,
										add_linear=add_linear,target_model='LBP',tri_std=self.tri_std)
							if self.use_gp:
								with tf.variable_scope('GrandParents'):
									layer_gp= classifiers.trilinear_discriminator_new(
										gp_head, gp_headdep, gp_dep,
										hidden_keep_prob=hidden_keep_prob_tri,
										add_linear=add_linear,target_model='LBP',tri_std=self.tri_std)
									#'''
							# head dep head
							if self.use_cop:
								with tf.variable_scope('CoParents'):
									layer_cop = classifiers.trilinear_discriminator_new(
										cop_head, cop_dep, cop_head,
										hidden_keep_prob=hidden_keep_prob_tri,
										add_linear=add_linear,target_model='LBP',tri_std=self.tri_std)
									#'''
					elif not self.old_trilin:
						print('head dep')
						if self.test_new_potential:
							print('testing new potential')
							# head dep dep
							if self.use_sib:
								with tf.variable_scope('Sibling'):
									layer_sib = classifiers.trilinear_discriminator_test(
										layer1, layer2, layer2,
										hidden_keep_prob=hidden_keep_prob_tri,
										add_linear=add_linear,target_model='LBP',tri_std=self.tri_std)
							# head head+dep dep
							if self.use_gp:
								with tf.variable_scope('GrandParents'):
									layer_gp = classifiers.trilinear_discriminator_test(
										layer1, layer3, layer2,
										hidden_keep_prob=hidden_keep_prob_tri,
										add_linear=add_linear,target_model='LBP',tri_std=self.tri_std)
							# head dep head
							if self.use_cop:
								with tf.variable_scope('CoParents'):
									layer_cop = classifiers.trilinear_discriminator_test(
										layer1, layer2, layer1,
										hidden_keep_prob=hidden_keep_prob_tri,
										add_linear=add_linear,target_model='LBP',tri_std=self.tri_std)
						#=================================================
						else:
							# head dep dep
							if self.use_sib:
								with tf.variable_scope('Sibling'):
									layer_sib = classifiers.trilinear_discriminator_new(
										layer1, layer2, layer2,
										hidden_keep_prob=hidden_keep_prob_tri,
										add_linear=add_linear,target_model='LBP',tri_std=self.tri_std)
							# head head+dep dep
							if self.use_gp:
								with tf.variable_scope('GrandParents'):
									layer_gp = classifiers.trilinear_discriminator_new(
										layer1, layer3, layer2,
										hidden_keep_prob=hidden_keep_prob_tri,
										add_linear=add_linear,target_model='LBP',tri_std=self.tri_std)
							# head dep head
							if self.use_cop:
								with tf.variable_scope('CoParents'):
									layer_cop = classifiers.trilinear_discriminator_new(
										layer1, layer2, layer1,
										hidden_keep_prob=hidden_keep_prob_tri,
										add_linear=add_linear,target_model='LBP',tri_std=self.tri_std)
					#-----------------------------------------------------------
					if self.linearize:
						with tf.variable_scope('Linearization'):
							lin_logits = classifiers.bilinear_discriminator(
								lin_layer1, lin_layer2,
								hidden_keep_prob=hidden_keep_prob,
								add_linear=add_linear)
					if self.distance:
						with tf.variable_scope('Distance'):
							dist_lamda = 1+tf.nn.softplus(classifiers.bilinear_discriminator(
								dist_layer1, dist_layer2,
								hidden_keep_prob=hidden_keep_prob,
								add_linear=add_linear))

				binary_shape = layer_sib.shape.as_list()
				if debug:
					outputs['printdata']={}
					if self.new_potential:
						outputs['printdata']['unary_old']=unary
					if self.use_sib:
						outputs['printdata']['layer_sib_old']=layer_sib
					if self.use_cop:
						outputs['printdata']['layer_cop_old']=layer_cop
					if self.use_gp:
						outputs['printdata']['layer_gp_old']=layer_gp

				#======================LBP==========================

				#pdb.set_trace()
				#update: use log space for LBP
				if self.two_gpu:
					print('two gpu training for LBP')
					#pdb.set_trace()
					GPUID=1
				else:
					GPUID=0
				with tf.device('/device:GPU:'+str(GPUID)):
					if self.use_sib:
						layer_type=layer_sib
					elif self.use_cop:
						layer_type=layer_cop
					elif self.use_gp:
						layer_type=layer_gp
					# (n x 2 x ma x mb x mc) ab <- ac
					log_message_sib=tf.stack([tf.zeros_like(layer_type),tf.zeros_like(layer_type)],1)
					# (n x 2 x ma x mb x mc) ab <- cb
					log_message_cop=tf.stack([tf.zeros_like(layer_type),tf.zeros_like(layer_type)],1)
					# (n x 2 x ma x mb x mc) ab <- bc
					log_message_gp1=tf.stack([tf.zeros_like(layer_type),tf.zeros_like(layer_type)],1)
					# (n x 2 x ma x mb x mc) ab <- ca
					log_message_gp2=tf.stack([tf.zeros_like(layer_type),tf.zeros_like(layer_type)],1)
					#pdb.set_trace()
					#'''
					if self.layer_mask:
						print('use layer mask')
						unary=unary*tf.cast(tf.expand_dims(tf.transpose(token_weights,[0,2,1]),-2),dtype=tf.float32)
						if self.remove_root_child:
							# abc -> ab,ac
							if self.use_sib:
								layer_sib=layer_sib*self.token_weights_sib
								log_message_sib=log_message_sib*self.token_weights_sib[:,None,:,:,:]
							if self.use_cop:
								layer_cop=layer_cop*self.token_weights_cop
								log_message_cop=log_message_cop*self.token_weights_cop[:,None,:,:,:]
							if self.use_gp:
								layer_gp=layer_gp*self.token_weights_gp
								log_message_gp1=log_message_gp1*self.token_weights_gp[:,None,:,:,:]
								log_message_gp2=log_message_gp2*self.token_weights_gp2[:,None,:,:,:]
						else:
							if self.use_sib:
								layer_sib=layer_sib*token_weights4D
							if self.use_cop:
								layer_cop=layer_cop*token_weights4D
							if self.use_gp:
								layer_gp=layer_gp*token_weights4D
							token_weights4D=tf.expand_dims(token_weights4D,1)
							if self.use_sib:
								log_message_sib=log_message_sib*token_weights4D
							if self.use_cop:
								log_message_cop=log_message_cop*token_weights4D
							if self.use_gp:
								log_message_gp1=log_message_gp1*token_weights4D
								log_message_gp2=log_message_gp2*token_weights4D
					#'''
					#(n x m x 2 x m) -> (n x 2 x m x m)
					batch_size=nn.get_sizes(layer_sib)[0]
					unary=tf.transpose(unary,[0,2,1,3])
					if self.use_sib:
						layer_sib = layer_sib-tf.linalg.band_part(layer_sib,-1,0) + tf.transpose(tf.linalg.band_part(layer_sib,0,-1),perm=[0,1,3,2])
					if self.use_gp:
						layer_gp2 = tf.transpose(layer_gp,perm=[0,2,3,1])
					if self.use_cop:
						#(n x ma x mb x mc) -> (n x mb x ma x mc)
						#in order to create a symmtric tensor on ma and mc
						layer_cop = tf.transpose(layer_cop,perm=[0,2,1,3])
						# first set lower triangle part to be zero, then assign the upper triangle part transposed to lower triangle part
						layer_cop = layer_cop - tf.linalg.band_part(layer_cop,-1,0) + tf.transpose(tf.linalg.band_part(layer_cop,0,-1),perm=[0,1,3,2])
						# Finally (n x mb x ma x mc) -> (n x ma x mb x mc)
						layer_cop = tf.transpose(layer_cop,perm=[0,2,1,3])
					for i in range(1,int(self.num_iteration)+1):
						#(n x ma x mb x mc x 2)
						prev_sib=log_message_sib
						prev_cop=log_message_cop
						prev_gp1=log_message_gp1
						prev_gp2=log_message_gp2
						#first gp is a->b->c second is c->a->b
						#(n x 2 x ma x mb x mc) -> (n x 2 x ma x mb)
						#TODO: add mask when message passing
						FP=tf.reduce_sum(prev_gp1,-1)+tf.reduce_sum(prev_gp2,-1)+tf.reduce_sum(prev_cop,-1)+tf.reduce_sum(prev_sib,-1)
						#FP=tf.reduce_sum(prev_gp1,-1)+tf.reduce_sum(prev_gp2,-1)
						if debug:
							outputs['printdata']['FP0'+str(i)]=FP
							outputs['printdata']['prev_sib'+str(i)]=prev_sib
							outputs['printdata']['prev_cop'+str(i)]=prev_cop
							outputs['printdata']['prev_gp1'+str(i)]=prev_gp1
							outputs['printdata']['prev_gp2'+str(i)]=prev_gp2
						# remove self loop and self repeat (a b a) (a b b)
						#'''
						FP=FP-tf.transpose(tf.linalg.diag_part(tf.transpose(prev_sib,perm=[0,1,3,2,4])),perm=[0,1,3,2])-tf.linalg.diag_part(prev_sib)
						FP=FP-tf.transpose(tf.linalg.diag_part(tf.transpose(prev_cop,perm=[0,1,3,2,4])),perm=[0,1,3,2])-tf.linalg.diag_part(prev_cop)
						FP=FP-tf.transpose(tf.linalg.diag_part(tf.transpose(prev_gp1,perm=[0,1,3,2,4])),perm=[0,1,3,2])-tf.linalg.diag_part(prev_gp1)
						FP=FP-tf.transpose(tf.linalg.diag_part(tf.transpose(prev_gp2,perm=[0,1,3,2,4])),perm=[0,1,3,2])-tf.linalg.diag_part(prev_gp2)
						#'''
						#TODO: softmax form
						FP_potential=tf.expand_dims(unary+FP,-1)
						if self.use_sib:
							#update sibling
							#pdb.set_trace()
							#(n x 2 x ma x mc) -> (n x 2 x ma x mc x 1) - (n x 2 x ma x mc x mb)-> (n x 2 x ma x mc x mb)-> (n x 2 x ma x mb x mc)
							sib_FP=tf.transpose(FP_potential-prev_sib,perm=[0,1,2,4,3])
							log_message_sib_0=tf.reduce_logsumexp(sib_FP,axis=1)
							#(n x ma x mb x mc) + (n x ma x mb x mc) -> (n x ma x mb x mc)
							log_message_sib_1=tf.reduce_logsumexp(tf.stack([sib_FP[:,0],sib_FP[:,1]+layer_sib],1),axis=1)
							#(n x ma x mb x mc) -> (n x 2 x ma x mb x mc)
							log_message_sib=tf.stack([log_message_sib_0,log_message_sib_1],axis=1)
						if self.use_cop:
							#update coparents
							#(n x 2 x mc x mb) -> (n x 2 x mc x mb x 1) - (n x 2 x mc x mb x ma)-> (n x 2 x mc x mb x ma)-> (n x 2 x ma x mb x mc)
							cop_FP=tf.transpose(FP_potential-prev_cop,perm=[0,1,4,3,2])
							log_message_cop_0=tf.reduce_logsumexp(cop_FP,axis=1)
							#(n x ma x mb x mc) + (n x ma x mb x mc) -> (n x ma x mb x mc)
							log_message_cop_1=tf.reduce_logsumexp(tf.stack([cop_FP[:,0],cop_FP[:,1]+layer_cop],1),axis=1)
							#(n x ma x mb x mc) -> (n x 2 x ma x mb x mc)
							log_message_cop=tf.stack([log_message_cop_0,log_message_cop_1],axis=1)
						if self.use_gp:
							#update gp1 (a->b->c)
							#(n x 2 x ma x mc) -> (n x 2 x mb x mc x 1) - (n x 2 x mb x mc x ma)-> (n x 2 x mb x mc x ma)-> (n x 2 x ma x mb x mc)
							# Note that here gp1 should minus prev_gp2 as we need to keep ab->bc is still representing a->b->c
							gp_FP=tf.transpose(FP_potential-prev_gp2,perm=[0,1,4,2,3])
							log_message_gp_0=tf.reduce_logsumexp(gp_FP,axis=1)
							#(n x ma x mb x mc) + (n x ma x mb x mc) -> (n x ma x mb x mc)
							log_message_gp_1=tf.reduce_logsumexp(tf.stack([gp_FP[:,0],gp_FP[:,1]+layer_gp],1),axis=1)
							#(n x ma x mb x mc) -> (n x 2 x ma x mb x mc)
							log_message_gp1=tf.stack([log_message_gp_0,log_message_gp_1],axis=1)

							#update gp2 (c->a->b)
							#(n x 2 x ma x mc) -> (n x 2 x mc x ma x 1) - (n x 2 x mc x ma x mb)-> (n x 2 x mc x ma x mb)-> (n x 2 x ma x mb x mc)
							gp2_FP=tf.transpose(FP_potential-prev_gp1,perm=[0,1,3,4,2])
							log_message_gp2_0=tf.reduce_logsumexp(gp2_FP,axis=1)
							#(n x ma x mb x mc) + (n x ma x mb x mc) -> (n x ma x mb x mc)
							log_message_gp2_1=tf.reduce_logsumexp(tf.stack([gp2_FP[:,0],gp2_FP[:,1]+layer_gp2],1),axis=1)
							#(n x ma x mb x mc) -> (n x 2 x ma x mb x mc)
							log_message_gp2=tf.stack([log_message_gp2_0,log_message_gp2_1],axis=1)
						if self.normalize:
							if i==1:
								print('normalize each LBP iteration')
							if self.use_sib:
								log_message_sib=tf.nn.log_softmax(log_message_sib,1)
							if self.use_cop:
								log_message_cop=tf.nn.log_softmax(log_message_cop,1)
							if self.use_gp:
								log_message_gp1=tf.nn.log_softmax(log_message_gp1,1)
								log_message_gp2=tf.nn.log_softmax(log_message_gp2,1)
						if self.layer_mask:
							if self.remove_root_child:
								# abc -> ab,ac
								if self.use_sib:
									layer_sib=layer_sib*self.token_weights_sib
								if self.use_cop:
									layer_cop=layer_cop*self.token_weights_cop
								if self.use_gp:
									layer_gp=layer_gp*self.token_weights_gp
								log_message_sib=log_message_sib*self.token_weights_sib[:,None,:,:,:]
								log_message_cop=log_message_cop*self.token_weights_cop[:,None,:,:,:]
								log_message_gp1=log_message_gp1*self.token_weights_gp[:,None,:,:,:]
								log_message_gp2=log_message_gp2*self.token_weights_gp2[:,None,:,:,:]
							else:
								log_message_sib=log_message_sib*token_weights4D
								log_message_cop=log_message_cop*token_weights4D
								log_message_gp1=log_message_gp1*token_weights4D
								log_message_gp2=log_message_gp2*token_weights4D
						if debug:
							'''
							outputs['printdata']['gp2_FP'+str(i)]=gp2_FP
							outputs['printdata']['FP_potential'+str(i)]=FP_potential
							outputs['printdata']['FP4'+str(i)]=FP
							outputs['printdata']['sib_FP'+str(i)]=sib_FP
							outputs['printdata']['cop_FP'+str(i)]=cop_FP
							outputs['printdata']['gp1_FP'+str(i)]=gp_FP
							outputs['printdata']['gp2_FP'+str(i)]=gp2_FP
							outputs['printdata']['message_sib_0'+str(i)]=log_message_sib_0
							outputs['printdata']['message_sib_1'+str(i)]=log_message_sib_1
							outputs['printdata']['message_sib'+str(i)]=log_message_sib
							outputs['printdata']['message_cop_0'+str(i)]=log_message_cop_0
							outputs['printdata']['message_cop_1'+str(i)]=log_message_cop_1
							outputs['printdata']['message_cop'+str(i)]=log_message_cop
							outputs['printdata']['message_gp1'+str(i)]=log_message_gp1
							outputs['printdata']['message_gp2'+str(i)]=log_message_gp2
							outputs['printdata']['message_gp1_0'+str(i)]=log_message_gp_0
							outputs['printdata']['message_gp1_1'+str(i)]=log_message_gp_1
							outputs['printdata']['message_gp2_0'+str(i)]=log_message_gp2_0
							outputs['printdata']['message_gp2_1'+str(i)]=log_message_gp2_1
							#'''
							pass

					log_belief=tf.reduce_sum(log_message_gp1,-1)+tf.reduce_sum(log_message_gp2,-1)+tf.reduce_sum(log_message_cop,-1)+tf.reduce_sum(log_message_sib,-1)
					#log_belief=tf.reduce_sum(log_message_gp1,-1)+tf.reduce_sum(log_message_gp2,-1)
					# remove self loop and self repeat (a b a) (a b b)
					#'''
					log_belief=log_belief-tf.transpose(tf.linalg.diag_part(tf.transpose(log_message_sib,perm=[0,1,3,2,4])),perm=[0,1,3,2])-tf.linalg.diag_part(log_message_sib)
					log_belief=log_belief-tf.transpose(tf.linalg.diag_part(tf.transpose(log_message_cop,perm=[0,1,3,2,4])),perm=[0,1,3,2])-tf.linalg.diag_part(log_message_cop)
					log_belief=log_belief-tf.transpose(tf.linalg.diag_part(tf.transpose(log_message_gp1,perm=[0,1,3,2,4])),perm=[0,1,3,2])-tf.linalg.diag_part(log_message_gp1)
					log_belief=log_belief-tf.transpose(tf.linalg.diag_part(tf.transpose(log_message_gp2,perm=[0,1,3,2,4])),perm=[0,1,3,2])-tf.linalg.diag_part(log_message_gp2)
					#'''
					log_belief=log_belief+unary
					#q_value=belief/tf.reduce_sum(belief,axis=1,keepdims=True)
					#calculate softmax loss later
					q_value=log_belief
					if debug:
						outputs['printdata']['message_sib_fin']=tf.reduce_sum(log_message_sib,-1)-tf.transpose(tf.linalg.diag_part(tf.transpose(log_message_sib,perm=[0,1,3,2,4])),perm=[0,1,3,2])-tf.linalg.diag_part(log_message_sib)
						outputs['printdata']['message_cop_fin']=tf.reduce_sum(log_message_cop,-1)-tf.transpose(tf.linalg.diag_part(tf.transpose(log_message_cop,perm=[0,1,3,2,4])),perm=[0,1,3,2])-tf.linalg.diag_part(log_message_cop)
						outputs['printdata']['message_gp1_fin']=tf.reduce_sum(log_message_gp1,-1)-tf.transpose(tf.linalg.diag_part(tf.transpose(log_message_gp1,perm=[0,1,3,2,4])),perm=[0,1,3,2])-tf.linalg.diag_part(log_message_gp1)
						outputs['printdata']['message_gp2_fin']=tf.reduce_sum(log_message_gp2,-1)-tf.transpose(tf.linalg.diag_part(tf.transpose(log_message_gp2,perm=[0,1,3,2,4])),perm=[0,1,3,2])-tf.linalg.diag_part(log_message_gp2)
					#======================LBP==========================

					#-----------------------------------------------------------
					# Process the targetsb
					# (n x m x m) -> (n x m x m)
					#here in fact is a graph, which is m*m representing the connection between each edge
					unlabeled_targets = self.placeholder#ground truth graph, what is self.placeholder?
					#USELESS
					shape = tf.shape(unary_layer1)
					batch_size, bucket_size = shape[0], shape[1]
					# (1 x m)
					ids = tf.expand_dims(tf.range(bucket_size), 0)
					# (1 x m) -> (1 x 1 x m)
					head_ids = tf.expand_dims(ids, -2)
					# (1 x m) -> (1 x m x 1)
					dep_ids = tf.expand_dims(ids, -1)
					#pdb.set_trace()
					if debug:
						#outputs['printdata']['logits']=logits
						outputs['printdata']['q_value']=q_value
						outputs['printdata']['unary']=unary
						#outputs['printdata']['binary']=binary
						outputs['printdata']['belief']=log_belief
						if self.use_sib:
							outputs['printdata']['message_sib']=log_message_sib
							outputs['printdata']['message_sib_0']=log_message_sib_0
							outputs['printdata']['layer_sib']=layer_sib
						if self.use_gp:
							#outputs['printdata']['second_temp_gp']=second_temp_gp
							outputs['printdata']['message_gp1']=log_message_gp1
							outputs['printdata']['message_gp2']=log_message_gp2
							outputs['printdata']['layer_gp']=layer_gp
						if self.use_cop:
							outputs['printdata']['layer_cop']=layer_cop
							outputs['printdata']['message_cop']=log_message_cop
						#outputs['printdata']['layer1']=layer1
						#outputs['printdata']['layer2']=layer2
						if not self.separate_embed:
							outputs['printdata']['layer3']=layer3
						outputs['printdata']['targets']=unlabeled_targets
						outputs['printdata']['token_weights']=token_weights
						if self.sibling_only:
							outputs['printdata']['binary_weights']=binary_weights
							outputs['printdata']['binary']=layer_sib
						if self.new_potential:
								#outputs['printdata']['layer_sib2']=layer_sib2
								outputs['printdata']['layer_gp2']=layer_gp2
								#outputs['printdata']['layer_cop2']=layer_cop2
								pass
						#outputs['printdata']['binary_weights']=binary_weights
						'''
						outputs['printdata']['unary_weights']=unary_weights
						outputs['printdata']['binary_weights']=binary_weights
						outputs['printdata']['binary_weights_cop']=binary_weights_cop
						outputs['printdata']['binary_weights_gp']=binary_weights_gp
						#'''

					#no running here
					if self.linearize:
						# Wherever the head is to the left
						# (n x m x m), (1 x m x 1) -> (n x m x m)
						lin_targets = tf.to_float(tf.less(unlabeled_targets, dep_ids))
						# cross-entropy of the linearization of each i,j pair
						# (1 x 1 x m), (1 x m x 1) -> (n x m x m)
						lin_ids = tf.tile(tf.less(head_ids, dep_ids), [batch_size, 1, 1])
						# (n x 1 x m), (n x m x 1) -> (n x m x m)
						lin_xent = -tf.nn.softplus(tf.where(lin_ids, -lin_logits, lin_logits))
						# add the cross-entropy to the logits
						# (n x m x m), (n x m x m) -> (n x m x m)
						logits += tf.stop_gradient(lin_xent)
					if self.distance:
						# (n x m x m) - (1 x m x 1) -> (n x m x m)
						dist_targets = tf.abs(unlabeled_targets - dep_ids)
						# KL-divergence of the distance of each i,j pair
						# (1 x 1 x m) - (1 x m x 1) -> (n x m x m)
						dist_ids = tf.to_float(tf.tile(tf.abs(head_ids - dep_ids), [batch_size, 1, 1]))+1e-12
						# (n x m x m), (n x m x m) -> (n x m x m)
						#dist_kld = (dist_ids * tf.log(dist_lamda / dist_ids) + dist_ids - dist_lamda)
						dist_kld = -tf.log((dist_ids - dist_lamda)**2/2 + 1)
						# add the KL-divergence to the logits
						# (n x m x m), (n x m x m) -> (n x m x m)
						logits += tf.stop_gradient(dist_kld)
					#pdb.set_trace()
					#-----------------------------------------------------------
					# Note: here need a transpose as the target is the transpose graph(or opposite direction of adjacency graph)
					# (n x 2 x ma x mb) -> (n x 2 x mb x ma)
					if self.transposed:
						q_value=tf.transpose(q_value, [0,1,3,2])
					# Compute probabilities/cross entropy
					# (n x 2 x m x m) -> (n x m x m x 2)
					transposed_logits = tf.transpose(q_value, [0,2,3,1])
					probabilities=tf.nn.softmax(transposed_logits) * tf.to_float(tf.expand_dims(token_weights, axis=-1))
					#TODO: what I want is still a probability of label 1, compared to the origin sigmoid(logits)? check later
					probabilities=probabilities[:,:,:,1]
					#label_probabilities = tf.nn.softmax(transposed_logits) * tf.to_float(tf.expand_dims(token_weights, axis=-1))
					# (n x m x m), (n x m x m x c), (n x m x m) -> ()
					# change sparse_softmax_cross_entropy to softmax_cross_entropy? It is the same in this situation
					loss = tf.losses.sparse_softmax_cross_entropy(unlabeled_targets, transposed_logits, weights=token_weights)
					#loss = tf.losses.sparse_sigmoid_cross_entropy(unlabeled_targets, transposed_logits, weights=token_weights)
					#pdb.set_trace()
					'''
					transposed_logits=tf.nn.softmax(transposed_logits,-1)
					L2_target=tf.cast(unlabeled_targets,dtype=tf.float32)
					loss = tf.reduce_sum((tf.pow(L2_target-transposed_logits[:,:,:,1],2)+tf.pow(1-L2_target-transposed_logits[:,:,:,0],2))\
							*tf.cast(token_weights,dtype=tf.float32))/tf.cast(batch_size,dtype=tf.float32)
					if debug:
						outputs['printdata']['L2_target']=L2_target
						outputs['printdata']['transposed_logits']=transposed_logits
					#'''
					if self.combine_loss:
						print('use combined loss')
						unary_probs=-unary
						# similar to q_value
						if self.transposed:
							unary_probs=tf.transpose(unary_probs,[0,3,2,1])
						unary_probs=tf.transpose(unary_probs, [0,1,3,2])
						#loss /= 2
						loss = loss*self.loss_weight + tf.losses.sparse_softmax_cross_entropy(unlabeled_targets, unary_probs, weights=token_weights)*self.loss_weight_unary
					'''
					# (n x m x m) -> (n x m x m)
					probabilities = tf.nn.sigmoid(logits) * tf.to_float(token_weights)#token weights is sentence length?
					# (n x m x m), (n x m x m), (n x m x m) -> ()
					loss = tf.losses.sigmoid_cross_entropy(unlabeled_targets, logits, weights=token_weights)#here label_smoothing is 0, the sigmoid XE have any effect?
					'''
					n_tokens = tf.to_float(tf.reduce_sum(token_weights))
					if self.linearize:
						lin_target_xent = lin_xent * unlabeled_targets
						loss -= tf.reduce_sum(lin_target_xent * tf.to_float(token_weights)) / (n_tokens + 1e-12)
					if self.distance:
						dist_target_kld = dist_kld * unlabeled_targets
						loss -= tf.reduce_sum(dist_target_kld * tf.to_float(token_weights)) / (n_tokens + 1e-12)

					#-----------------------------------------------------------
					# Compute predictions/accuracy
					# precision/recall
					# (n x m x m) -> (n x m x m)
					#predictions = nn.greater(logits, 0, dtype=tf.int32) * token_weights#edge that predicted
					predictions = tf.argmax(transposed_logits, axis=-1, output_type=tf.int32) * token_weights
					# if self.compare_precision:
					# 	cond = tf.equal(transposed_logits[:,:,:,1], tf.expand_dims(tf.reduce_max(transposed_logits[:,:,:,1],-1),-1))
					# 	predictions = tf.where(cond, tf.cast(cond,tf.float32), tf.zeros_like(transposed_logits[:,:,:,1]))
					# 	predictions = tf.cast(predictions,tf.int32) * token_weights

					# (n x m x m) (*) (n x m x m) -> (n x m x m)
					true_positives = predictions * unlabeled_targets
					# (n x m x m) -> ()
					n_predictions = tf.reduce_sum(predictions)
					n_targets = tf.reduce_sum(unlabeled_targets)
					n_true_positives = tf.reduce_sum(true_positives)
					# () - () -> ()
					n_false_positives = n_predictions - n_true_positives
					n_false_negatives = n_targets - n_true_positives
					# (n x m x m) -> (n)
					n_targets_per_sequence = tf.reduce_sum(unlabeled_targets, axis=[1,2])
					n_true_positives_per_sequence = tf.reduce_sum(true_positives, axis=[1,2])
					# (n) x 2 -> ()
					n_correct_sequences = tf.reduce_sum(nn.equal(n_true_positives_per_sequence, n_targets_per_sequence))
			#-----------------------------------------------------------
			# Populate the output dictionary
			if debug:
				outputs['printdata']['logits']=transposed_logits
				outputs['printdata']['loss']=loss
			#print(123)
			outputs['unlabeled_targets'] = unlabeled_targets
			outputs['probabilities'] = probabilities
			outputs['unlabeled_loss'] = loss
			outputs['loss'] = loss

			outputs['unlabeled_predictions'] = predictions
			outputs['n_unlabeled_true_positives'] = n_true_positives
			outputs['n_unlabeled_false_positives'] = n_false_positives
			outputs['n_unlabeled_false_negatives'] = n_false_negatives
			outputs['n_correct_unlabeled_sequences'] = n_correct_sequences
			outputs['predictions'] = predictions
			outputs['n_true_positives'] = n_true_positives
			outputs['n_false_positives'] = n_false_positives
			outputs['n_false_negatives'] = n_false_negatives
			outputs['n_correct_sequences'] = n_correct_sequences
			return outputs

	#=============================================================
	# token should be: 1:rel|2:acl|5:dep or 1|2|5
	def index(self, token):
		""""""

		nodes = []
		if token != '_':
			token = token.split('|')
			for edge in token:
				head = edge.split(':')[0]
				nodes.append(int(head))
		return nodes

	#=============================================================
	# index should be [1, 2, 5]
	def token(self, index):
		""""""

		return [str(head) for head in index]

	#=============================================================
	def get_root(self):
		""""""

		return '_'

	#=============================================================
	def __getitem__(self, key):
		if isinstance(key, six.string_types):
			nodes = []
			if key != '_':
				token = key.split('|')
				for edge in token:
					head = edge.split(':')[0]
					nodes.append(int(head))
			return nodes
		elif hasattr(key, '__iter__'):
			if len(key) > 0:
				if isinstance(key[0], six.integer_types + (np.int32, np.int64)):
					return '|'.join([str(head) for head in key])
				else:
					return [self[k] for k in key]
			else:
				return '_'
		else:
			raise ValueError('Key to GraphIndexVocab.__getitem__ must be (iterable of) strings or iterable of integers')

#***************************************************************

class SecondOrderGraphLBPVocab(GraphSecondLBPVocab, cv.SemheadVocab):
	pass


# -*- coding: utf-8 -*-
"""
Created on Mon May  3 15:59:09 2021

This file is is and adaptation of the WorkMATe source code, from the article:
    "Flexible Working Memory through Selective Gating and Attentional Tagging"
Wouter Kruijne, Sander M Bohte, Pieter R Roelfsema, Christian N L Olivers

-------------------
The class defined in this file implements the WorkMATe model;
Entry point for the model is the 'step()' function which defined what WorkMATe
does during a single time step; this entails:
1-  the feedforward sweep to determine its  next action,
2-  Learning: update the weights in response to any potentially obtained reward
3-  recurrent attentional feedback to place synaptic tags on connections
    responsible for the selected action.
4-  Executing the desired action, including possibly storing
    an observation in memory.

Several extensions are added to the original code --> making it a Deep WorkMATe
with CIFAAR-10 inputs and a locally learned matching function.
1- An extra layer was added (latent)
2- Timestamp is added at latent level, learning rate is set to 0.05
3- Only weight updates are performed between trials
4- Both CIFAR-10 latent features and the alphabet can be used as a stimulus dictionary

@author: Lieke Ceton
"""
#%% Dependencies

import numpy as np

from stimuli import obs_time, obs_cifar, match_cifar, obs_alpha, match_alpha
from matchnet import MatchNet, cosine_similarity

#%% Define WorkMATe architecture

#np.random.seed(20)
class WorkMATe(object):
    """
    Architecture schematic:
    x -- l -- S
          \/
          h
          |
    [q_int|q_ext]
    """
    def __init__(self, env=None, nhidden=25, nblocks=2, block_size=20, lr=0.05, wh=0.5):
        super(WorkMATe, self).__init__()

        #np.random.seed(20)
        assert env is not None
        self.env=env

        ## learning params (adopted from Rombouts et al., 2015)
        self.beta = lr
        self.gamma = 0.90
        self.L = 0.8
        # exploration:
        self.epsilon = 0.025
        self.bias = 1
        print('Agent WorkMATe_match initialising')

        ## member lambda functions:
        # sigmoid transfer function, offset at 2.5
        sigmoid_offset = 2.5
        self.transfer = lambda x: 1 / ( 1. + np.exp(sigmoid_offset - x) )
        self.dtransfer= lambda x: x * (1. - x) # derivative

        # softmax normalization; for action selection - boltzmann controller
        self.softmaxnorm = lambda x: (
            np.exp( x - x.max() ) / np.exp( x - x.max() ) .sum() )

        ## init network architecture -- inputs and output shape from env
        # input and hidden
        #nx = inputs.get_obs('a').size
        nl = block_size
        nh = nhidden
        # memory cell properties:
        self.nblocks = nblocks
        self.block_size = block_size
        nS = nblocks * block_size

        # output -- q layer consisting of 2 modules
        # module for n external actions, internal actions for nblocks + 1 (null)
        mod_sz = env.n_actions, nblocks + 1
        nq = np.sum(mod_sz)
        # indices of module for each node:
        self.zmods = np.hstack( [ [i] * sz for i, sz in enumerate( mod_sz ) ] )

        ## init network layers (activations 0)
        # (x will be constructed when processing 'new_obs')
        time_input = obs_time().size
        if self.env.stim_type == 'alpha':
            letter_input = obs_alpha().size
        if self.env.stim_type == 'cifar':
            letter_input = obs_cifar().size

        self.S = np.zeros( nS )
        self.l = np.zeros( nl )
        self.h = np.zeros( nh )
        self.q = np.zeros( nq )

        ## init weights, tags traces, (+1 indicates projection from bias node)
        wl, wh = -wh, wh

        #Characteristics used as attributes for step_match()
        self.ch = [letter_input, nl, nh, nS, wh, wl]

        # Memory projection (x > S)
        #self.W_Sx  = np.random.sample( (nS, nx) )  * (wh-wl) + wl
        self.W_Sl  = np.random.sample( (nS, nl) )  * (wh-wl) + wl
        # Note that time and sensory input cells are not separated in memory

        # PLASTIC CONNECTIONS (all except memory projection)

        # Input projection with bias node
        self.W_lx = np.random.sample((nl, letter_input + 1))*(wh-wl) + wl
        #self.W_lx = np.array(W_lx)
        self.W_lx_start = np.copy(self.W_lx)

        # connections l -> h; nl + match nodes + bias
        self.W_hl  = np.random.sample( (nh, time_input + nl + nblocks + 1) ) * (wh-wl) + wl
        self.W_hl_start = np.copy(self.W_hl)

        # connections S-> h:
        self.W_hS  = np.random.sample( (nh, nS    ) ) * (wh-wl) + wl
        # connections h-> q:
        self.W_qh  = np.random.sample( (nq, nh + 1) ) * (wh-wl) + wl

        # tags are shaped like weights but initialized at 0:
        zeros_ = np.zeros_like
        # W_lx is only updated between trials to keep memory trace stable
        # W_lx_trial accumulates all changes made within a trial
        self.W_lx_trial = zeros_(self.W_lx)

        #initialise tags
        self.Tag_W_lx, self.Trace_W_lx = zeros_(self.W_lx), zeros_(self.W_lx)
        self.Tag_W_hl, self.Trace_W_hl = zeros_(self.W_hl), zeros_(self.W_hl)
        self.Tag_W_hS, self.Trace_W_hS = zeros_(self.W_hS), zeros_(self.W_hS)
        self.Tag_W_qh, self.Trace_W_qh = zeros_(self.W_qh), zeros_(self.W_qh)

        ## Init matchnet:
        dummy_inp = np.zeros(block_size) #dummy input of length S == Sproj and Memory
        self.matchnet = [MatchNet(dummy_inp, dummy_inp) for item in range(self.nblocks)]

        # Init action state
        self.action = -1
        # (prev) predicted reward:
        self.qat_1 = self.qat = None
        self.t   = 0
        return

    def _intertrial_reset(self):
        """
        Reset time, memory, tags and traces
        """
        # update weights W_lx with the accumulated weight updates from the whole trial
        self.W_lx += self.W_lx_trial

        # reset time and memory
        self.t = 0
        self.S *= 0

        # previous action = zeros
        self.z *= 1 #I set this from 0 to 1, it does not seem to change anything
        # reset tags/traces for each Wmat
        zeros_ = np.zeros_like
        self.Tag_W_lx, self.Trace_W_lx = zeros_(self.W_lx), zeros_(self.W_lx)
        self.Tag_W_hl, self.Trace_W_hl = zeros_(self.W_hl), zeros_(self.W_hl)
        self.Tag_W_hS, self.Trace_W_hS = zeros_(self.W_hS), zeros_(self.W_hS)
        self.Tag_W_qh, self.Trace_W_qh = zeros_(self.W_qh), zeros_(self.W_qh)
        #reset accumulated weight updates W_lx
        self.W_lx_trial = zeros_(self.W_lx)

        # reset 'current action', and the like
        self.action = -1
        self.qat = None
        return

    def step_match(self):
        if self.env.stim_type == 'cifar':
            self.x_sens = match_cifar() #get random input
        elif self.env.stim_type == 'alpha':
            self.x_sens = match_alpha() #get random input

        self.x = np.r_[self.x_sens, self.bias]
        Sproj = self.compute_latent()   #Compute the memory projection
        Sproj = Sproj.reshape((self.nblocks, self.block_size))

        #Decide on trial type and thus on memory content
        #For each memory block there is a match network
        #tt is the trial type: 0 is match, 1 is mismatch
        #If tt = 0, the memory is updated to the current input
        #The chance is around equal to save or not
        tt = np.random.choice([0,1], self.nblocks, p=[0.5, 0.5])
        S_ = self.S.reshape( (self.nblocks, self.block_size) )
        for i in range(self.nblocks):
            if tt[i] == 0:
                S_[i,:] = Sproj[i,:]#the S projection should be here!

        #Compute the match values
        m, tt, corr = self.compute_match(Sproj)

        #Noise/random changes in weights W_lx and W_Sl
        self.W_lx = np.random.sample((self.ch[1], self.ch[0] + 1))*(0.5--0.5) - 0.5
        self.W_Sl  = np.random.sample((self.ch[3], self.ch[1]))*(self.ch[4]-self.ch[5]) + self.ch[5]

        return m, tt, corr #return match values and trial_types

    def step(self):
        # get observation and reward from env
        self.obs, self.r = self.env.step(self.action)
        # do feedforward:
        #matches, tt_match = self._feedforward()
        matches, tt = self._feedforward()
        # learn from the obtained reward
        self._learn()
        # end of trial?
        if 'RESET' == self.obs: #change in to == because the observation is now an integer not a string
            self._intertrial_reset()
            return self.r, None, None #fake matches
        # do feedback (tag placement)
        self._feedback()
        # act (internal, external)
        self._act()
        self.t += 1
        #return self.r, matches #, tt_match
        return self.r, matches, tt

    def _feedforward(self):
        # shift previous action
        self.qat_1 = self.qat
        if 'RESET' == self.obs: #change in to == because the observation is now an integer not a string
            # however, we do not expect RESET to ever only be part of the observation so the equal to boolean should be good enough
            # no meaningful feedforward sweep:  qat is not computed
            self.qat = None
            return None, None #return fake match
        # else:
        # compute input, hidden, output:
        self.construct_input()
        Sproj = self.compute_latent()
        matches, tt, _ = self.compute_match(Sproj)
        self.compute_hidden(matches)
        self.compute_output()
        # determine z from q (action selection)
        self.action_selection()
        # determine new qat
        self.qat = (self.z * self.q).sum()
        return matches, tt

    def _learn(self):
        """
        Learn from the reward; compute RPE and update weights
        general form form delta = r + gamma * qat - qat_1
        ...but there are edge cases
        """
        r = self.r
        if self.qat and self.qat_1: # regular
            delta = r + (self.gamma * self.qat) - self.qat_1
        elif self.qat_1 is None: # first step
            delta = r + (self.gamma * self.qat) - self.qat
        else: # self.qa(t) is None (final step):
            delta = r - self.qat_1
        self.delta = delta
        self.update_weights()
        return

    def _feedback(self):
        # updates traces and tags, based on action selection:
        # traces and tags
        self.update_traces()
        self.update_tags()
        return

    def _act(self):
        # external and internal actions:
        zext =  self.z[self.zmods == 0]
        zint =  self.z[self.zmods == 1]
        self.action = np.argmax(zext)
        self.update_memory( zint )
        return

    def construct_input(self):
        """
        Turn obs into a vector; uses coding defined in 'inputs.py'
        """
        # input consists of: observation and time t
        self.x_time = obs_time(self.t) #time-part
        #letter-part
        if self.env.stim_type == 'alpha':
           self.x_sens = obs_alpha(self.obs)
           #self.x_sens = inputs.obs_alpha_with_noise(self.obs)
        if self.env.stim_type == 'cifar':
           self.x_sens = obs_cifar(self.obs)
           #self.x_sens = inputs.obs_orthogonal(self.obs)
        #self.x_sens = inputs.get_obs(self.obs, self.t) #only letter-part
        self.x = np.r_[self.x_sens, self.bias] #only bias included
        return

    def compute_latent(self):
        # x -> l
        self.l_in = self.W_lx.dot(self.x)
        # Compute l activities
        self.l_sens = self.transfer(np.r_[self.l_in])
        # Compute match value:
        Sproj = self.W_Sl.dot(self.l_in).reshape( (self.nblocks, self.block_size) )
        return Sproj

    def compute_match(self, Sproj):
        #output = self.matchnet.step(x,y)
        matches = []
        corr = []
        tt = []
        S_ = self.S.reshape( (self.nblocks, self.block_size) )
        for i in range(self.nblocks):
            m, _ = self.matchnet[i].step(Sproj[i,:], S_[i,:])
            cos_corr = cosine_similarity(Sproj[i,:], S_[i,:])
            corr.append(cos_corr)
            matches.append(m)
            if np.all(Sproj[i,:] == S_[i,:]):
                 tt_match = 0
            else:
                 tt_match = 1
            tt.append(tt_match)
        return matches, tt, corr

    def compute_hidden(self, matches):
        # add match nodes + bias to input vector
        #add time-part here
        self.l_out = np.r_[self.x_time, self.l_sens, matches, self.bias] #the matches are added here now

        # x->h  +  S->h
        self.S_out  = self.S

        # Compute Ha and h
        h_in = self.W_hl.dot(self.l_out) + self.W_hS.dot(self.S_out)
        self.h_out  = np.r_[self.transfer(h_in), self.bias] # bias added
        return

    def compute_output(self):
        # hidden output (has bias added)
        self.q = self.W_qh.dot(self.h_out)
        # (no transfer, q nodes are linear)
        return

    def action_selection(self):
        #np.random.seed(self.seed)
        # using q, per module, determine z (based on argmax or exploration)
        self.z = np.zeros_like(self.q)
        # action selection for both modules separately
        for mod_idx in np.unique(self.zmods):
            qvec = self.q[self.zmods == mod_idx] # get the module's qvalues
            # check exploration; if not just take argmax:
            if ( np.random.sample() >= self.epsilon ):
                action = np.argmax(qvec)
            else: # compute softmax over Q and explore:
                pvec = self.softmaxnorm(qvec)
                action = np.random.choice( list(range(qvec.size)), p = pvec)
            # set zvec: 1-hot code of actions:
            zvec = np.zeros_like(qvec)
            zvec[ action ] = 1.0
            # place zvec into z at the right indices:
            self.z[self.zmods == mod_idx] = zvec
        return

    def update_weights(self):
        """
        all Weight-Trace pairs are updated with the same rule:
        w += beta * delta * tag
        """
        #Only update W_lx during intertrial resets
        #Accumulate changes during trial and update later
        self.W_lx_trial += self.beta * self.delta * self.Tag_W_lx

        #Change the weights between h and S slower to have memory..
        self.W_hS += self.beta * self.delta * self.Tag_W_hS

        wt_pairs = (
                     ( self.W_hl, self.Tag_W_hl ),
                     ( self.W_qh, self.Tag_W_qh )) #added by LJC
        for W, Tag in wt_pairs:
            W += self.beta * self.delta * Tag

        return

    def update_traces(self):
        """
        Traces are the intermediate layers' markers
        The Traces are a relic from old AuGMEnT code, have no 'meaning' here
        """
        # Regulars, are replaced by new input:
        self.Trace_W_lx *= 0.0
        self.Trace_W_hl *= 0.0
        self.Trace_W_hS *= 0.0
        # add 1 x X vec to H x X matrix yields H copies of 1 x X vec
        self.Trace_W_lx += self.x.reshape(1, self.x.size) # this includes trace for bias
        self.Trace_W_hl += self.l_out.reshape(1, self.l_out.size) # this includes trace for bias
        self.Trace_W_hS += self.S_out.reshape(1, self.S.size)
        return

    def update_tags(self):

        # 1. old tag decay:
        alltags = (self.Tag_W_lx, self.Tag_W_hl, self.Tag_W_hS, self.Tag_W_qh)
        for Tag in alltags:
            Tag *= (self.L * self.gamma)

        # 2. form new tags:
        # tags onto output units: selected action.
        self.Tag_W_qh[self.z.astype('bool'), :] += self.h_out

        # feedback to hidden
        dh = self.dtransfer(self.h_out[:-1])
        self.fbh = self.W_qh[self.z.astype('bool'), :-1] # excluding the bias node
        self.fbh = self.fbh.sum(axis = 0 ) # summed contribution of all actions
        self.feedbackh = self.fbh * dh
        self.Tag_W_hS +=  np.expand_dims(self.feedbackh, 1) *  self.Trace_W_hS
        #update tag hl
        self.Tag_W_hl +=  np.expand_dims(self.feedbackh, 1) *  self.Trace_W_hl

        #feedback to latent
        dl = self.dtransfer(self.l_sens)
        self.fbh_transfer = self.dtransfer(self.fbh) #the feedback onto h is transferred through the hidden layer
        #W_hl_no_m = self.W_hl[:,:-3] #no match no bias, no time-part
        #W_hl_no_m = self.W_hl[:,10:30]

        W_hl_no_m = self.W_hl[:,obs_time().size:-(self.nblocks + 1)] #only learn back onto the stimulus input, not the time inputs (exclude the bias node)
        #inputs.obs_time().size

        self.W_hl_transpose = W_hl_no_m.T
        self.fhl = self.W_hl_transpose.dot(self.fbh_transfer) #for each node in l, all active h units are summed
        self.feedbackl = self.fhl * dl
        self.Tag_W_lx += np.expand_dims(self.feedbackl,1) * self.Trace_W_lx
        return

    def update_memory(self, zvec):
        # final z is 'do not gate'; nothing happens then
        if not zvec[-1] == 1:
            # else:
            gate_idx = np.argmax(zvec)
            l = self.l_in   #this excludes bias- and  match
            S_ = self.S.reshape( (self.nblocks, self.block_size) )
            W_ = self.W_Sl.reshape( S_.shape +  (l.size, ) )
            # project x->S (encode)
            Sproj = W_.dot(l)
            # store @ gated 'stripe'
            S_[gate_idx,:] = Sproj[gate_idx,:]
            # transform S back to its flat representation
            self.S = S_.reshape(self.S.shape)
        return


#!/usr/bin/env python
# -*- coding: utf-8 -*-
# -----------------------------------------------------------------------------
# glumpy is an OpenGL framework for the fast visualization of numpy arrays.
# Copyright (C) 2009-2011  Nicolas P. Rougier. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice,
#    this list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY NICOLAS P. ROUGIER ''AS IS'' AND ANY EXPRESS OR
# IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF
# MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO
# EVENT SHALL NICOLAS P. ROUGIER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT,
# INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
# (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
# ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF
# THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#
# The views and conclusions contained in the software and documentation are
# those of the authors and should not be interpreted as representing official
# policies, either expressed or implied, of Nicolas P. Rougier.
# -----------------------------------------------------------------------------
'''
A filter is a shader that transform the current displayed texture. Since
shaders cannot be easily serialized within the GPU, they have to be well
structured on the python side such that we can possibly merge them into a
single source code for both vertex and fragment. Consequently, there is a
default code for both vertex and fragment with specific entry points such that
filter knows where to insert their specific code (declarations, functions and
call (or code) to be inserted in the main function).

Spatial interpolation filter classes for OpenGL textures.

Each filter generates a one-dimensional lookup table (weights value from 0 to
ceil(radius)) that is uploaded to video memory (as a 1d texture) and is then
read by the shader when necessary. It avoids computing weight values for each
pixel. Furthemore, each 2D-convolution filter is separable and can be computed
using 2 1D-convolution with same 1d-kernel (= the lookup table values).

Available filters:

  - Nearest  (radius 0.5)
  - Bilinear (radius 1.0)
  - Hanning (radius 1.0)
  - Hamming (radius 1.0)
  - Hermite (radius 1.0)
  - Kaiser (radius 1.0)
  - Quadric (radius 1.5)
  - Bicubic (radius 2.0)
  - CatRom (radius 2.0)
  - Mitchell (radius 2.0)
  - Spline16 (radius 2.0)
  - Spline36 (radius 4.0)
  - Gaussian (radius 2.0)
  - Bessel (radius 3.2383)
  - Sinc (radius 4.0)
  - Lanczos (radius 4.0)
  - Blackman (radius 4.0)

Note::

  Weights code has been translated from the antigrain geometry library
  available at http://www.antigrain.com/
'''

import math
import numpy as np

class SpatialFilter(object):
    ''' '''

    def __init__(self, radius=1.0):
        self.radius = radius

    def weight(self, x):
        '''
        Return filter weight for a distance x.

        :Parameters:
            ``x`` : 0 < float < ceil(self.radius)
                Distance to be used to compute weight.
        '''
        raise NotImplemented

    def kernel(self, size=4*512):
        radius = self.radius
        r = int(max(1.0, math.ceil(radius)))
        samples = int(size / r)
        n = size  # r*samples
        kernel = np.zeros(n)
        X = np.linspace(0, r, n)
        for i in range(n):
            kernel[i] = self.weight(X[i])
        N = np.zeros(samples)
        for i in range(r):
            N += kernel[::+1][i*samples:(i+1)*samples]
            N += kernel[::-1][i*samples:(i+1)*samples]
        for i in range(r):
            kernel[i*samples:(i+1)*samples:+1] /= N
        return kernel

    def filter_code(self):

        n = int(math.ceil(self.radius))
        filter_1 = 'filter1D_radius%d' % n
        filter_2 = 'filter2D_radius%d' % n

        code = ''
        code += 'vec4\n'
        code += '%s( sampler2D kernel, float index, float x, ' % filter_1
        for i in range(2*n):
            if i == 2*n-1:
                code += 'vec4 c%d )\n' % i
            else:
                code += 'vec4 c%d, ' % i
        code += '{\n'
        code += '    float w, w_sum = 0.0;\n'
        code += '    vec4 r = vec4(0.0,0.0,0.0,0.0);\n'
        for i in range(n):
            code += '    w = unpack_interpolate(kernel, vec2(%f+(x/%.1f), index));\n' % (1.0 - (i + 1) / float(n), n)  # noqa
            code += '    w = w*kernel_scale + kernel_bias;\n'  # noqa
            # code += '   w_sum += w;'
            code += '    r += c%d * w;\n' % i
            code += '    w = unpack_interpolate(kernel, vec2(%f-(x/%.1f), index));\n' % ((i+1)/float(n), n)  # noqa
            code += '    w = w*kernel_scale + kernel_bias;\n'
            # code += '   w_sum += w;'
            code += '    r += c%d * w;\n' % (i + n)
        # code += '    return r/w_sum;\n'
        code += '    return r;\n'
        code += '}\n'
        code += "\n"
        code += 'vec4\n'
        code += '%s' % filter_2
        code += '(sampler2D texture, sampler2D kernel, float index, vec2 uv, vec2 pixel)\n'  # noqa
        code += '{\n'
        code += '    vec2 texel = uv/pixel - vec2(0.5, 0.5) ;\n'
        code += '    vec2 f = fract(texel);\n'
        code += '    texel = (texel-fract(texel) + vec2(0.001, 0.001)) * pixel;\n'  # noqa
        for i in range(2*n):
            code += '    vec4 t%d = %s(kernel, index, f.x,\n' % (i, filter_1)
            for j in range(2*n):
                x, y = (-n+1+j, -n+1+i)
                code += '        texture2D( texture, texel + vec2(%d, %d) * pixel),\n' % (x, y)  # noqa

            # Remove last trailing',' and close function call
            code = code[:-2] + ');\n'

        code += '    return %s(kernel, index, f.y, ' % filter_1
        for i in range(2*n):
            code += 't%d, ' % i

        # Remove last trailing',' and close function call
        code = code[:-2] + ');\n'
        code += '}\n'

        return code

    def call_code(self, index):
        code = ""
        n = int(math.ceil(self.radius))
        filter_1 = 'filter1D_radius%d' % n  # noqa
        filter_2 = 'filter2D_radius%d' % n

        code += 'vec4 %s(sampler2D texture, vec2 shape, vec2 uv)\n' % self.__class__.__name__  # noqa
        code += '{'
        code += ' return %s(texture, u_kernel, %f, uv, 1.0/shape); ' % (filter_2, index)  # noqa
        code += '}\n'
        return code

class Nearest(SpatialFilter):
    '''
    Nearest (=None) filter (radius = 0.5).

    Weight function::

      w(x) = 1

    '''

    def __init__(self):
        SpatialFilter.__init__(self, radius=.5)

    def weight(self, x):
        return 1.0

    def _get_code(self):
        self.build_LUT()
        code = 'vec4\n'
        code += 'interpolate(sampler2D texture, sampler1D kernel, vec2 uv, vec2 pixel)\n'  # noqa
        code += '{\n   return texture2D(texture, uv);\n}\n'
        return code
    code = property(_get_code, doc='''filter functions code''')

class Bilinear(SpatialFilter):
    '''
    Bilinear filter (radius = 1.0).

    Weight function::

      w(x) = 1 - x

    '''

    def __init__(self):
        SpatialFilter.__init__(self, radius=1.0)

    def weight(self, x):
        return 1.0 - x

class Hanning(SpatialFilter):
    '''
    Hanning filter (radius = 1.0).

    Weight function::

      w(x) = 0.5 + 0.5 * cos(pi * x)

    '''

    def __init__(self):
        SpatialFilter.__init__(self, radius=1.0)

    def weight(self, x):
        return 0.5 + 0.5 * math.cos(math.pi * x)

class Hamming(SpatialFilter):
    '''
    Hamming filter (radius = 1.0).

    Weight function::

      w(x) = 0.54 + 0.46 * cos(pi * x)

    '''

    def __init__(self):
        SpatialFilter.__init__(self, radius=1.0)

    def weight(self, x):
        return 0.54 + 0.46 * math.cos(math.pi * x)

class Hermite(SpatialFilter):
    ''' Hermite filter (radius = 1.0).

    Weight function::

      w(x) = (2*x-3)*x^2 + 1

    '''

    def __init__(self):
        SpatialFilter.__init__(self, radius=1.0)

    def weight(self, x):
        return (2.0 * x - 3.0) * x * x + 1.0

class Quadric(SpatialFilter):
    '''
    Quadric filter (radius = 1.5).

    Weight function::

             |  0.0 ≤ x < 0.5: 0.75 - x*x
      w(x) = |  0.5 ≤ x < 1.5: 0.5 - (x-1.5)^2
             |  1.5 ≤ x      : 0

    '''

    def __init__(self):
        SpatialFilter.__init__(self, radius=1.5)

    def weight(self, x):
        if x < 0.75:
            return 0.75 - x * x
        elif x < 1.5:
            t = x - 1.5
            return 0.5 * t * t
        else:
            return 0.0

class Bicubic(SpatialFilter):
    '''
    Bicubic filter (radius = 2.0).

    Weight function::

      w(x) = 1/6((x+2)^3 - 4*(x+1)^3 + 6*x^3 -4*(x-1)^3)
    '''

    def __init__(self):
        SpatialFilter.__init__(self, radius=2.0)

    def pow3(self, x):
        if x <= 0:
            return 0
        else:
            return x * x * x

    def weight(self, x):
        return (1.0/6.0) * (self.pow3(x + 2) -
                            4 * self.pow3(x + 1) +
                            6 * self.pow3(x) -
                            4 * self.pow3(x - 1))

class Kaiser(SpatialFilter):
    '''
    Kaiser filter (radius = 1.0).

    Weight function::

      w(x) = bessel_i0(a sqrt(1-x^2)* 1/bessel_i0(b)

    '''

    def __init__(self, b=6.33):
        self.a = b
        self.epsilon = 1e-12
        self.i0a = 1.0 / self.bessel_i0(b)
        SpatialFilter.__init__(self, radius=1.0)

    def bessel_i0(self, x):
        s = 1.0
        y = x * x / 4.0
        t = y
        i = 2
        while t > self.epsilon:
            s += t
            t *= float(y) / (i * i)
            i += 1
        return s

    def weight(self, x):
        if x > 1:
            return 0
        return self.bessel_i0(self.a * math.sqrt(1.0 - x * x)) * self.i0a

class CatRom(SpatialFilter):
    '''
    Catmull-Rom filter (radius = 2.0).

    Weight function::

             |  0 ≤ x < 1: 0.5*(2 + x^2*(-5+x*3))
      w(x) = |  1 ≤ x < 2: 0.5*(4 + x*(-8+x*(5-x)))
             |  2 ≤ x    : 0

    '''

    def __init__(self, size=256*8):
        SpatialFilter.__init__(self, radius=2.0)

    def weight(self, x):
        if x < 1.0:
            return 0.5 * (2.0 + x * x * (-5.0 + x * 3.0))
        elif x < 2.0:
            return 0.5 * (4.0 + x * (-8.0 + x * (5.0 - x)))
        else:
            return 0.0

class Mitchell(SpatialFilter):
    '''
    Mitchell-Netravali filter (radius = 2.0).

    Weight function::

             |  0 ≤ x < 1: p0 + x^2*(p2 + x*p3)
      w(x) = |  1 ≤ x < 2: q0 + x*(q1 + x*(q2 + x*q3))
             |  2 ≤ x    : 0

    '''

    def __init__(self, b=1.0/3.0, c=1.0/3.0):
        self.p0 = (6.0 - 2.0 * b) / 6.0
        self.p2 = (-18.0 + 12.0 * b + 6.0 * c) / 6.0
        self.p3 = (12.0 - 9.0 * b - 6.0 * c) / 6.0
        self.q0 = (8.0 * b + 24.0 * c) / 6.0
        self.q1 = (-12.0 * b - 48.0 * c) / 6.0
        self.q2 = (6.0 * b + 30.0 * c) / 6.0
        self.q3 = (-b - 6.0 * c) / 6.0
        SpatialFilter.__init__(self, radius=2.0)

    def weight(self, x):
        if x < 1.0:
            return self.p0 + x * x * (self.p2 + x * self.p3)
        elif x < 2.0:
            return self.q0 + x * (self.q1 + x * (self.q2 + x * self.q3))
        else:
            return 0.0

class Spline16(SpatialFilter):
    '''
    Spline16 filter (radius = 2.0).

    Weight function::

             |  0 ≤ x < 1: ((x-9/5)*x - 1/5)*x + 1
      w(x) = |
             |  1 ≤ x < 2: ((-1/3*(x-1) + 4/5)*(x-1) - 7/15 )*(x-1)

    '''

    def __init__(self):
        SpatialFilter.__init__(self, radius=2.0)

    def weight(self, x):
        if x < 1.0:
            return ((x - 9.0/5.0) * x - 1.0/5.0) * x + 1.0
        else:
            return ((-1.0/3.0 * (x-1) + 4.0/5.0) * (x-1) - 7.0/15.0) * (x-1)

class Spline36(SpatialFilter):
    '''
    Spline36 filter (radius = 3.0).

    Weight function::

             |  0 ≤ x < 1: ((13/11*x - 453/209)*x -3/209)*x +1
      w(x) = |  1 ≤ x < 2: ((-6/11*(x-1) - 270/209)*(x-1) -156/209)*(x-1)
             |  2 ≤ x < 3: (( 1/11*(x-2) - 45/209)*(x-2) + 26/209)*(x-2)
    '''

    def __init__(self):
        SpatialFilter.__init__(self, radius=3.0)

    def weight(self, x):
        if x < 1.0:
            return ((13.0/11.0 * x - 453.0/209.0) * x - 3.0/209.0) * x + 1.0
        elif x < 2.0:
            return ((-6.0/11.0 * (x-1) + 270.0/209.0) * (x-1) - 156.0 / 209.0) * (x-1)  # noqa
        else:
            return ((1.0 / 11.0 * (x-2) - 45.0/209.0) * (x - 2) + 26.0/209.0) * (x-2)  # noqa

class Gaussian(SpatialFilter):
    '''
    Gaussian filter (radius = 2.0).

    Weight function::

      w(x) = exp(-2x^2) * sqrt(2/pi)

    Note::

      This filter does not seem to be correct since:

        x = np.linspace(0, 1.0, 100 )
        f = weight
        z = f(x+1)+f(x)+f(1-x)+f(2-x)

        z should be 1 everywhere but it is not the case and it produces "grid
        effects".
    '''
    def __init__(self):
        SpatialFilter.__init__(self, radius=2.0)

    def weight(self, x):
        return math.exp(-2.0 * x * x) * math.sqrt(2.0 / math.pi)

class Bessel(SpatialFilter):
    '''
    Bessel filter (radius = 3.2383).
    '''

    def __init__(self):
        SpatialFilter.__init__(self, radius=3.2383)

    def besj(self, x, n):
        '''
        Function BESJ calculates Bessel function of first kind of order n
        Arguments:
            x - value at which the Bessel function is required
            n - an integer (>=0), the order
        --------------------
        C++ Mathematical Library
        Converted from equivalent FORTRAN library
        Converted by Gareth Walker for use by course 392 computational project
        All functions tested and yield the same results as the corresponding
        FORTRAN versions.

        If you have any problems using these functions please report them to
        M.Muldoon@UMIST.ac.uk

        Documentation available on the web
        http://www.ma.umist.ac.uk/mrm/Teaching/392/libs/392.html
        Version 1.0   8/98
        29 October, 1999
        --------------------
        Adapted for use in AGG library by
                    Andy Wilk (castor.vulgaris@gmail.com)
        Adapted for use in vispy library by
                    Nicolas P. Rougier (Nicolas.Rougier@inria.fr)
        -----------------------------------------------------------------------
        '''
        if n < 0:
            return 0.0
        x = float(x)  # force float type

        d = 1e-6
        b = 0
        if math.fabs(x) <= d:
            if n != 0:
                return 0
            return 1

        b1 = 0  # b1 is the value from the previous iteration
        # Set up a starting order for recurrence
        m1 = int(math.fabs(x)) + 6
        if math.fabs(x) > 5:
            m1 = int(math.fabs(1.4 * x + 60 / x))

        m2 = int(n + 2 + math.fabs(x) / 4)
        if m1 > m2:
            m2 = m1

        # Apply recurrence down from current max order
        while True:
            c3 = 0
            c2 = 1e-30
            c4 = 0
            m8 = 1
            if m2 // 2 * 2 == m2:
                m8 = -1

            imax = m2 - 2
            for i in range(1, imax+1):
                c6 = 2 * (m2 - i) * c2 / x - c3
                c3 = c2
                c2 = c6
                if m2 - i - 1 == n:
                    b = c6
                m8 = -1 * m8
                if m8 > 0:
                    c4 = c4 + 2 * c6

            c6 = 2 * c2 / x - c3
            if n == 0:
                b = c6
            c4 += c6
            b /= c4
            if math.fabs(b - b1) < d:
                return b
            b1 = b
            m2 += 3

    def weight(self, x):
        if x == 0.0:
            return math.pi/4.0
        else:
            return self.besj(math.pi * x, 1) / (2.0 * x)

class Sinc(SpatialFilter):
    '''
    Sinc filter (radius = 4.0).

    Weight function::

    '''

    def __init__(self, size=256, radius=4.0):
        SpatialFilter.__init__(self, radius=max(radius, 2.0))

    def weight(self, x):
        if x == 0.0:
            return 1.0
        x *= math.pi
        return (math.sin(x) / x)

class Lanczos(SpatialFilter):
    '''
    Lanczos filter (radius = 4.0).

    Weight function::

    '''

    def __init__(self, size=256, radius=4.0):
        SpatialFilter.__init__(self, radius=max(radius, 2.0))

    def weight(self, x):
        if x == 0.0:
            return 1.0
        elif x > self.radius:
            return 0.0
        x *= math.pi
        xr = x / self.radius
        return (math.sin(x) / x) * (math.sin(xr)/xr)

class Blackman(SpatialFilter):
    '''
    Blackman filter (radius = 4.0).
    '''

    def __init__(self, size=256, radius=4.0):
        SpatialFilter.__init__(self, radius=max(radius, 2.0))

    def weight(self, x):
        if x == 0.0:
            return 1.0
        elif x > self.radius:
            return 0.0
        x *= math.pi
        xr = x / self.radius
        return (math.sin(x) / x) * (0.42 + 0.5*math.cos(xr) + 0.08*math.cos(2*xr))  # noqa

# Generate kernels texture (16 x 1024)
filters = [Bilinear(), Hanning(),  Hamming(),  Hermite(),
           Kaiser(),   Quadric(),  Bicubic(),  CatRom(),
           Mitchell(), Spline16(), Spline36(), Gaussian(),
           Bessel(),   Sinc(),     Lanczos(),  Blackman()]

n = 1024
K = np.zeros((16, n))
for i, f in enumerate(filters):
    K[i] = f.kernel(n)

bias = K.min()
scale = K.max()-K.min()
K = (K-bias)/scale
np.save("spatial-filters.npy", K.astype(np.float32))

print("// ------------------------------------")
print("// Automatically generated, do not edit")
print("// ------------------------------------")
print("")
print("const float kernel_bias  = %f;" % bias)
print("const float kernel_scale = %f;" % scale)
print("const float kernel_size = %f;" % n)
print("const vec4 bits = vec4(1.0, 1.0/256.0, 1.0/(256.0*256.0), 1.0/(256.0*256.0*256.0));")  # noqa
print("uniform sampler2D u_kernel;")
print("")

code = 'float\n'
code += 'unpack_unit(vec4 rgba)\n'
code += '{\n'
code += '\t// return rgba.r;  // uncomment this for r32f debugging\n'
code += '\treturn dot(rgba, bits);\n'
code += '}\n'
print(code.expandtabs(4))

code = 'float\n'
code += 'unpack_ieee(vec4 rgba)\n'
code += '{\n'
code += '\t// return rgba.r;  // uncomment this for r32f debugging\n'
code += '\trgba.rgba = rgba.abgr * 255.;\n'
code += '\tfloat sign = 1.0 - step(128.0,rgba[0])*2.0;\n'
code += '\tfloat exponent = 2.0 * mod(rgba[0],128.0) + ' \
        'step(128.0,rgba[1]) - 127.0;\n'
code += '\tfloat mantissa = mod(rgba[1],128.0)*65536.0 + rgba[2]*256.0 + ' \
        'rgba[3] + float(0x800000);\n'
code += '\treturn sign * exp2(exponent) * (mantissa * exp2(-23.));\n'
code += '}\n'
print(code.expandtabs(4))

code = 'float\n'
code += 'unpack_interpolate(sampler2D kernel, vec2 uv)\n'
code += '{\n'
code += '\t// return texture2D(kernel, uv).r; ' \
        '//uncomment this for r32f debug without interpolation\n'
code += '\tfloat kpixel = 1. / kernel_size;\n'
code += '\tfloat u = uv.x / kpixel;\n'
code += '\tfloat v = uv.y;\n'
code += '\tfloat uf = fract(u);\n'
code += '\tu = (u - uf) * kpixel;\n'
code += '\n'
code += '\tfloat d0 = unpack_unit(texture2D(kernel, vec2(u, v)));\n'
code += '\tfloat d1 = unpack_unit(texture2D(kernel, vec2(u + 1. * kpixel, v)));\n'  # noqa
code += '\treturn mix(d0, d1, uf);\n'
code += '}\n'
print(code.expandtabs(4))

F = SpatialFilter(1.0)
print(F.filter_code())
F = SpatialFilter(2.0)
print(F.filter_code())
F = SpatialFilter(3.0)
print(F.filter_code())
F = SpatialFilter(4.0)
print(F.filter_code())

# Generate filter functions
# Special case for nearest
print("""vec4 Nearest(sampler2D texture, vec2 shape, vec2 uv)""")
print("""{ return texture2D(texture,uv); }\n""")

for i, f in enumerate(filters):
    print(f.call_code((i+0.5)/16.0))


import torch
import numpy as np
# 19-4-12  pytorch

# input: coeff with shape [1,257]
def Split_coeff(coeff):
    id_coeff = coeff[:,:80] # identity(shape) coeff of dim 80
    ex_coeff = coeff[:,80:144] # expression coeff of dim 64
    tex_coeff = coeff[:,144:224] # texture(albedo) coeff of dim 80
    angles = coeff[:,224:227] # ruler angles(x,y,z) for rotation of dim 3
    gamma = coeff[:,227:254] # lighting coeff for 3 channel SH function of dim 27
    translation = coeff[:,254:] # translation coeff of dim 3

    return id_coeff,ex_coeff,tex_coeff,angles,gamma,translation

# CHJ_ADD
class _need_const:
    import numpy as np
    a0 = np.pi
    a1 = 2 * np.pi / np.sqrt(3.0)
    a2 = 2 * np.pi / np.sqrt(8.0)
    c0 = 1 / np.sqrt(4 * np.pi)
    c1 = np.sqrt(3.0) / np.sqrt(4 * np.pi)
    c2 = 3 * np.sqrt(5.0) / np.sqrt(12 * np.pi)
    d0 = 0.5/ np.sqrt(3.0)

    illu_consts=[a0, a1, a2, c0, c1, c2, d0]

    focal = 1015.0
    center = 112.0
    cam_pos = 10
    p_matrix = np.concatenate([[focal], [0.0], [center], [0.0], [focal], [center], [0.0], [0.0], [1.0]],
                              axis=0).astype(np.float32)  # projection matrix
    p_matrix = np.reshape(p_matrix, [1, 3, 3])
    p_matrix = torch.from_numpy(p_matrix)
    gpu_p_matrix = None

# compute face shape with identity and expression coeff, based on BFM model
# input: id_coeff with shape [1,80]
#         ex_coeff with shape [1,64]
# output: face_shape with shape [1,N,3], N is number of vertices
def Shape_formation(id_coeff,ex_coeff,facemodel):
    n_b = id_coeff.size(0)
    face_shape = torch.einsum('ij,aj->ai',facemodel.idBase,id_coeff) + \
                torch.einsum('ij,aj->ai',facemodel.exBase,ex_coeff) + \
                facemodel.meanshape

    face_shape = face_shape.view(n_b,-1,3)
    # re-center face shape
    face_shape = face_shape - facemodel.meanshape.view(1, -1, 3).mean(dim=1, keepdim=True)

    return face_shape

# compute vertex texture(albedo) with tex_coeff
# input: tex_coeff with shape [1,N,3]
# output: face_texture with shape [1,N,3], RGB order, range from 0-255
def Texture_formation(tex_coeff,facemodel):
    n_b = tex_coeff.size(0)
    face_texture = torch.einsum('ij,aj->ai',facemodel.texBase,tex_coeff) + facemodel.meantex

    face_texture = face_texture.view(n_b,-1,3)
    return face_texture

# compute vertex normal using one-ring neighborhood
# input: face_shape with shape [1,N,3]
# output: v_norm with shape [1,N,3]
def Compute_norm(face_shape,facemodel):

    face_id = facemodel.tri # vertex index for each triangle face, with shape [F,3], F is number of faces
    point_id = facemodel.point_buf # adjacent face index for each vertex, with shape [N,8], N is number of vertex
    shape = face_shape

    v1 = shape[:,face_id[:,0],:]
    v2 = shape[:,face_id[:,1],:]
    v3 = shape[:,face_id[:,2],:]
    e1 = v1 - v2
    e2 = v2 - v3
    face_norm = e1.cross(e2) # compute normal for each face
    empty = torch.zeros((face_norm.size(0), 1, 3), dtype=face_norm.dtype, device=face_norm.device)

    face_norm = torch.cat((face_norm, empty), 1) # concat face_normal with a zero vector at the end

    v_norm = face_norm[:,point_id,:].sum(2) # compute vertex normal using one-ring neighborhood
    # CHJ: not average, directly normalize
    v_norm = v_norm/v_norm.norm(dim=2).unsqueeze(2) # normalize normal vectors

    return v_norm

# compute rotation matrix based on 3 ruler angles
# input: angles with shape [1,3]
# output: rotation matrix with shape [1,3,3]
def Compute_rotation_matrix(angles):

    n_b = angles.size(0)
    sinx = torch.sin(angles[:, 0])
    siny = torch.sin(angles[:, 1])
    sinz = torch.sin(angles[:, 2])
    cosx = torch.cos(angles[:, 0])
    cosy = torch.cos(angles[:, 1])
    cosz = torch.cos(angles[:, 2])

    rotXYZ = torch.eye(3).view(1, 3, 3).repeat(n_b * 3, 1, 1).view(3, n_b, 3, 3)

    if angles.is_cuda: rotXYZ = rotXYZ.cuda(angles.get_device())

    rotXYZ[0, :, 1, 1] = cosx
    rotXYZ[0, :, 1, 2] = -sinx
    rotXYZ[0, :, 2, 1] = sinx
    rotXYZ[0, :, 2, 2] = cosx
    rotXYZ[1, :, 0, 0] = cosy
    rotXYZ[1, :, 0, 2] = siny
    rotXYZ[1, :, 2, 0] = -siny
    rotXYZ[1, :, 2, 2] = cosy
    rotXYZ[2, :, 0, 0] = cosz
    rotXYZ[2, :, 0, 1] = -sinz
    rotXYZ[2, :, 1, 0] = sinz
    rotXYZ[2, :, 1, 1] = cosz

    rotation = rotXYZ[2].bmm(rotXYZ[1]).bmm(rotXYZ[0])

    return rotation.permute(0, 2, 1)

# project 3D face onto image plane
# input: face_shape with shape [1,N,3]
#          rotation with shape [1,3,3]
#         translation with shape [1,3]
# output: face_projection with shape [1,N,2]
#           z_buffer with shape [1,N,1]
def Projection_layer(face_shape,rotation,translation,focal=1015.0,center=112.0): # we choose the focal length and camera position empirically

    n_b, nV, _ = face_shape.size()
    if face_shape.is_cuda:
        if _need_const.gpu_p_matrix is None:
            _need_const.gpu_p_matrix = _need_const.p_matrix.cuda(face_shape.get_device())

        p_matrix = _need_const.gpu_p_matrix.expand(n_b, 3, 3)
    else:
        p_matrix = _need_const.p_matrix.expand(n_b, 3, 3)

    face_shape_r = face_shape.bmm(rotation)  # CHJ: R has been transposed
    face_shape_t = face_shape_r + translation.view(n_b, 1, 3)

    face_shape_t[:, :, 2] = _need_const.cam_pos - face_shape_t[:, :, 2]

    aug_projection = face_shape_t.bmm(p_matrix.permute(0, 2, 1))

    #print(aug_projection)
    #exit()

    face_projection = aug_projection[:, :, 0:2] / aug_projection[:, :, 2:]

    # CHJ_WARN: I do this for visualization
    z_buffer = _need_const.cam_pos - aug_projection[:, :, 2:] # CHJ: same as the z of  face_shape_t

    return face_projection, z_buffer

# CHJ: It's different from what I knew.
# compute vertex color using face_texture and SH function lighting approximation
# input: face_texture with shape [1,N,3]
#          norm with shape [1,N,3]
#         gamma with shape [1,27]
# output: face_color with shape [1,N,3], RGB order, range from 0-255
#          lighting with shape [1,N,3], color under uniform texture

def Illumination_layer(face_texture, norm, gamma):

    n_b, num_vertex, _ = face_texture.size()
    n_v_full = n_b * num_vertex
    gamma = gamma.view(-1, 3, 9).clone()
    gamma[:, :, 0] += 0.8

    gamma = gamma.permute(0, 2, 1)

    a0, a1, a2, c0, c1, c2, d0 = _need_const.illu_consts

    Y0 = torch.ones(n_v_full).float() * a0*c0
    if gamma.is_cuda: Y0=Y0.cuda(gamma.get_device())
    norm = norm.view(-1, 3)
    nx, ny, nz = norm[:,0], norm[:,1], norm[:,2]
    arrH = []

    arrH.append( Y0 )
    arrH.append(-a1*c1*ny)
    arrH.append(a1*c1*nz)
    arrH.append(-a1*c1*nx)
    arrH.append(a2*c2*nx*ny)
    arrH.append(-a2*c2*ny*nz)
    arrH.append(a2*c2*d0*(3*nz.pow(2)-1))
    arrH.append(-a2*c2*nx*nz)
    arrH.append(a2*c2*0.5*(nx.pow(2)-ny.pow(2)))

    H = torch.stack(arrH, 1)
    Y = H.view(n_b, num_vertex, 9)

    # Y shape:[batch,N,9].

    # shape:[batch,N,3]
    lighting = Y.bmm(gamma)

    face_color = face_texture * lighting
    #lighting *= 128

    #print( face_color[0, 5] )
    return face_color,lighting

# face reconstruction with coeff and BFM model
def Reconstruction(coeff,facemodel):
    id_coeff,ex_coeff,tex_coeff,angles,gamma,translation = Split_coeff(coeff)

    # compute face shape
    face_shape = Shape_formation(id_coeff, ex_coeff, facemodel)
    # compute vertex texture(albedo)
    face_texture = Texture_formation(tex_coeff, facemodel)
    #print(face_texture[0, 2])
    # vertex normal
    face_norm = Compute_norm(face_shape,facemodel)
    # rotation matrix
    rotation = Compute_rotation_matrix(angles)
    face_norm_r = face_norm.bmm(rotation)

    # compute vertex projection on image plane (with image sized 224*224)
    face_projection,z_buffer = Projection_layer(face_shape,rotation,translation)

    # compute vertex color using SH function lighting approximation
    face_color,lighting = Illumination_layer(face_texture, face_norm_r, gamma)

    # vertex index for each face of BFM model
    tri = facemodel.tri

    face_shape = torch.matmul(face_shape,rotation)
    face_shape = face_shape + torch.reshape(translation,[-1,1,3])

    return face_shape, angles,face_color,tri,face_projection,z_buffer,gamma

def Reconstruction_for_render(coeff,facemodel):
    id_coeff,ex_coeff,tex_coeff,angles,gamma,translation = Split_coeff(coeff)
    face_shape = Shape_formation(id_coeff, ex_coeff, facemodel)
    face_texture = Texture_formation(tex_coeff, facemodel)
    face_norm = Compute_norm(face_shape,facemodel)
    rotation = Compute_rotation_matrix(angles)
    face_shape_r = torch.matmul(face_shape,rotation)
    face_shape_r = face_shape_r + torch.reshape(translation,[-1,1,3])
    face_norm_r = torch.matmul(face_norm,rotation)
    face_color,lighting = Illumination_layer(face_texture, face_norm_r, gamma)
    tri = facemodel.tri

    return face_shape_r,face_norm_r,face_color,tri