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

class TrainingParams:
	def __init__(self,
				 learning_rate: float = 0.1,
				 momentum: float = 0.5,
				 weight_decay: float = 0.0,
				 num_epochs: int = 1000,
				 num_gibbs_samples: int = 1,
				 mini_batch_size: int = 0,
				 do_rao_blackwell: bool = False,
				 do_gibbs_sample_visible: bool = False,
				 do_gibbs_sample_hidden: bool = False,
				 do_batch_sample: bool = False
				 ):
		# Training parameters
		self.learning_rate = learning_rate
		self.momentum = momentum
		self.weight_decay = weight_decay
		self.num_epochs = num_epochs
		self.num_gibbs_samples = num_gibbs_samples
		self.mini_batch_size = mini_batch_size
		self.do_rao_blackwell = do_rao_blackwell
		self.do_gibbs_sample_visible = do_gibbs_sample_visible
		self.do_gibbs_sample_hidden = do_gibbs_sample_hidden
		self.do_batch_sample = do_batch_sample

	@classmethod
	def from_dict(cls, params: dict):
		obj = TrainingParams()
		obj.learning_rate = params["learningRate"]
		obj.momentum = params["momentum"]
		obj.weight_decay = params["weightDecay"]
		obj.num_epochs = params["numEpochs"]
		obj.num_gibbs_samples = params["numGibbs"]
		obj.mini_batch_size = params["miniBatchSize"]
		obj.do_rao_blackwell = params["doRaoBlackwell"]
		obj.do_gibbs_sample_visible = params["gibbsDoSampleVisible"]
		obj.do_gibbs_sample_hidden = params["gibbsDoSampleHidden"]
		obj.do_batch_sample = params["doSampleBatch"]

		return obj

def cd_jens(entity: Entity, v_states: Mat, params: TrainingParams):
	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

def cd_binary_binary(entity: Entity, data_pos: Mat, params: TrainingParams):
	# 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 with training params
	# ToDo

	# 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, params: TrainingParams):
	# 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
	h_states_pos = sample(h_probs_pos)

	# Negative phase
	data_neg = entity.v_given_h(h_states_pos)
	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 cd_gaussian_gaussian(entity: Entity, data_pos: Mat, params: TrainingParams):
	# Positive phase
	h_probs_pos = entity.h_given_v(data_pos)

	# Sample hidden states
	h_states_pos = h_probs_pos + sample_gaussian(h_probs_pos)

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

	# Negative phase
	data_neg = entity.v_given_h(h_states_pos)
	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, params: TrainingParams, status: Status):
	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()
	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

	for mini_batch in to_mini_batch(batch, mini_batch_size):
		if not keep_running:
			break

		entity.grad_zero()
		for epochs in range(params.num_epochs):
			# Contrastive divergence learning: calculate gradients
			dwhv, dbv, dbh = cd_func(entity, mini_batch, params)

			# Adjust weight and biases
			grad = entity.grad_compute(dbv, dbh, dwhv, learning_rate=params.learning_rate/batch.shape[0], momentum=params.momentum, weight_decay=params.weight_decay)
			entity.state_adjust(grad)

			# 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({"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({"progress": {"value": round(progress), "unit": "%"},
					  "err_rms_total": {"value": err_rms, "unit": ""}})

	return None