from .matrix import sample, sample_gaussian, prob, rms_error_accu, Mat, np
from .entity import Entity
from .status import Status

def cd_jens(entity: Entity, v_states: Mat):
	params = entity.training_params
	v_probs = prob(v_states)
	h_states = entity.forward(v_states)
	h_probs = h_states

	if entity.params.do_gaussian_hidden:
		h_states += sample_gaussian(h_states)
	else:
		h_probs = entity.forward(v_states)
		if params.do_rao_blackwell:
			h_states = h_probs
		else:
			h_states = sample(h_probs)

	# Update weights (positive phase)
	dw = np.dot(np.transpose(v_states), h_states)
	dbv = np.sum(v_states, 0)
	dbh = np.sum(h_states, 0)

	# Gibbs	sampling with training params
	for i in range(params.num_gibbs_samples):
		if params.do_gibbs_sample_hidden:
			v_probs = entity.reconstruct(sample(h_probs))
		else:
			v_probs = entity.reconstruct(h_probs)

		# Create hidden representation given v
		if params.do_gibbs_sample_visible:
			h_probs = entity.forward(sample(v_probs))
		else:
			h_probs = entity.forward(v_probs)

	# Update weights (negative phase)
	dw -= np.dot(np.transpose(v_probs), h_probs)
	dbv -= np.sum(v_probs, 0)
	dbh -= np.sum(h_probs, 0)

	return dw, dbv, dbh

# %%%%%%%%% START POSITIVE PHASE %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
#   data = batchdata(:,:,batch);
#   poshidprobs = 1./(1 + exp(-data*vishid - repmat(hidbiases,numcases,1)));
#   batchposhidprobs(:,:,batch)=poshidprobs;
#   posprods    = data' * poshidprobs;
#   poshidact   = sum(poshidprobs);
#   posvisact = sum(data);
#
# %%%%%%%%% END OF POSITIVE PHASE  %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
#   poshidstates = poshidprobs > rand(numcases,numhid);
#
# %%%%%%%%% START NEGATIVE PHASE  %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
#   negdata = 1./(1 + exp(-poshidstates*vishid' - repmat(visbiases,numcases,1)));
#   neghidprobs = 1./(1 + exp(-negdata*vishid - repmat(hidbiases,numcases,1)));
#   negprods  = negdata'*neghidprobs;
#   neghidact = sum(neghidprobs);
#   negvisact = sum(negdata);
#
# %%%%%%%%% END OF NEGATIVE PHASE %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
def cd_binary_binary(entity: Entity, data_pos: Mat):
	params = entity.training_params

	# Positive phase
	h_probs_pos = prob(entity.h_given_v(data_pos))

	# Sample hidden states
	if not params.do_rao_blackwell:
		h_probs_pos = sample(h_probs_pos)

	# Update weights (positive phase)
	dw = np.dot(np.transpose(data_pos), h_probs_pos)
	dbh = np.sum(h_probs_pos, 0)
	dbv = np.sum(data_pos, 0)

	# Negative phase
	data_neg = prob(entity.v_given_h(h_probs_pos))
	h_probs_neg = prob(entity.h_given_v(data_neg))

	# Gibbs	sampling
	for _ in range(params.num_gibbs_samples-1):
		data_neg = prob(entity.v_given_h(h_probs_neg))
		h_probs_neg = prob(entity.h_given_v(data_neg))

	# Update weights (negative phase)
	dw -= np.dot(np.transpose(data_neg), h_probs_neg)
	dbh -= np.sum(h_probs_neg, 0)
	dbv -= np.sum(data_neg, 0)

	return dw, dbv, dbh

def cd_gaussian_binary(entity: Entity, data_pos: Mat):
	# Positive phase
	h_probs_pos = prob(entity.h_given_v(data_pos))

	# Update weights (positive phase)
	dw = np.dot(np.transpose(data_pos), h_probs_pos)
	dbh = np.sum(h_probs_pos, 0)
	dbv = np.sum(data_pos-entity.state.b_v, 0)

	# Negative phase — sample binary hidden states to get a sharper fantasy particle
	data_neg = entity.v_given_h(sample(h_probs_pos))
	h_probs_neg = prob(entity.h_given_v(data_neg))

	# Update weights (negative phase)
	dw -= np.dot(np.transpose(data_neg), h_probs_neg)
	dbh -= np.sum(h_probs_neg, 0)
	dbv -= np.sum(data_neg-entity.state.b_v, 0)

	return dw, dbv, dbh

def cd_gaussian_gaussian(entity: Entity, data_pos: Mat):
	# Positive phase — use mean h for weight update, sample h to drive reconstruction
	h_probs_pos = entity.h_given_v(data_pos)
	h_sample_pos = h_probs_pos + sample_gaussian(h_probs_pos)

	dw = np.dot(np.transpose(data_pos), h_probs_pos)
	dbh = np.sum(h_probs_pos, 0)
	dbv = np.sum(data_pos-entity.state.b_v, 0)

	# Negative phase — reconstruct from sampled h, use means for weight update
	data_neg = entity.v_given_h(h_sample_pos)
	h_probs_neg = entity.h_given_v(data_neg)

	dw -= np.dot(np.transpose(data_neg), h_probs_neg)
	dbh -= np.sum(h_probs_neg, 0)
	dbv -= np.sum(data_neg-entity.state.b_v, 0)

	return dw, dbv, dbh

# %%%%%%%%%START POSITIVE PHASE %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
# data = batchdata(:,:, batch);
# poshidprobs = (data * vishid) + repmat(hidbiases, numcases, 1);
# batchposhidprobs(:,:, batch)=poshidprobs;
#
# posprods = data' * poshidprobs;
# poshidact = sum(poshidprobs);
# posvisact = sum(data);
#
# %%%%%%%%%END OF POSITIVE PHASE %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
# poshidstates = poshidprobs + randn(numcases, numhid);
#
# %%%%%%%%%START NEGATIVE PHASE %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
# negdata = 1. / (1 + exp(-poshidstates * vishid' - repmat(visbiases,numcases,1)));
# neghidprobs = (negdata * vishid) + repmat(hidbiases, numcases, 1);
# negprods = negdata'*neghidprobs;
# neghidact = sum(neghidprobs);
# negvisact = sum(negdata);
#
# %%%%%%%%%END OF NEGATIVE PHASE %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
def cd_binary_gaussian(entity: Entity, data_pos: Mat):
	params = entity.training_params
	# Positive phase
	h_probs_pos = entity.h_given_v(data_pos)

	# Update weights (positive phase)
	dw = np.dot(np.transpose(data_pos), h_probs_pos)
	dbh = np.sum(h_probs_pos, 0)
	dbv = np.sum(data_pos, 0)

	# Sample hidden states
	if not params.do_rao_blackwell:
		h_probs_pos += sample_gaussian(h_probs_pos)

	# Negative phase
	data_neg = prob(entity.v_given_h(h_probs_pos))
	h_probs_neg = entity.h_given_v(data_neg)

	# Gibbs	sampling
	for _ in range(params.num_gibbs_samples-1):
		data_neg = prob(entity.v_given_h(h_probs_neg))
		h_probs_neg = entity.h_given_v(data_neg)

	# Update weights (negative phase)
	dw -= np.dot(np.transpose(data_neg), h_probs_neg)
	dbh -= np.sum(h_probs_neg, 0)
	dbv -= np.sum(data_neg, 0)

	return dw, dbv, dbh

def to_mini_batch(batch: Mat, mini_batch_size: int):
	mini_batches = []
	remain = batch.shape[0]
	index = 0
	while remain > 0:
		batch_size = min(mini_batch_size, remain)
		mini_batches.append(batch[index:index + batch_size])
		remain -= batch_size
		index += batch_size

	return mini_batches

def train(entity: Entity, batch: Mat, status: Status):
	params = entity.training_params
	if params is None:
		return False

	mini_batch_size = min(params.mini_batch_size, batch.shape[0]) if params.mini_batch_size > 0 else batch.shape[0]
	d_progress = 100.0 / (batch.shape[0]*params.num_epochs)
	progress = 0
	keep_running = True
	status.on_change(entity)
	cd_func = None
	if entity.type == Entity.Type.GB_RBM:
		cd_func = cd_gaussian_binary
	if entity.type == Entity.Type.BB_RBM:
		cd_func = cd_binary_binary
	if entity.type == Entity.Type.GG_RBM:
		cd_func = cd_gaussian_gaussian
	if entity.type == Entity.Type.BG_RBM:
		cd_func = cd_binary_gaussian

	entity.grad_zero()
	for epochs in range(params.num_epochs):
		if not keep_running:
			break

		for mini_batch in to_mini_batch(batch, mini_batch_size):
			# Contrastive divergence learning: calculate gradients
			dwhv, dbv, dbh = cd_func(entity, mini_batch)

			# Adjust weight and biases
			grad = entity.grad_compute(dbv, dbh, dwhv)
			entity.state_adjust(grad, 1.0/mini_batch.shape[0])

			# check if status update is needed
			if status.want_report(round(progress)):
				# Calculate	error
				err_rms = rms_error_accu(batch - entity.reconstruct(entity.forward(batch)))
				if not status.on_change(entity, {"progress": {"value": round(progress), "unit": "%"},
										 "err_rms": {"value": err_rms, "unit": ""}}):
					keep_running = False
					break

			progress += d_progress*mini_batch.shape[0]

	# Update final status
	err_rms = rms_error_accu(batch - entity.reconstruct(entity.forward(batch)))
	status.on_change(entity, {"progress": {"value": round(progress), "unit": "%"},
					  "err_rms_total": {"value": err_rms, "unit": ""}})

	return True
