import random
import matplotlib.pyplot as plt

from model.model import Model
from rbm.entity import Entity, EntityParams, TrainingParams
from image.sub_image import normalize
from rbm.matrix import Mat, np

N = 8
N_VIS = N * N
N_HID = 32
N_CASES = 400


class GaussianAEModel(Model):
	def __init__(self, name: str, work_dir: str = '.'):
		super().__init__(name, work_dir)
		self.unit1 = Entity(
			(N_VIS, N_HID),
			EntityParams(do_gaussian_visible=True, do_gaussian_hidden=False),
			TrainingParams(learning_rate=0.001, momentum=0.9, num_epochs=3000)
		)

	def forward(self, x: Mat) -> Mat:
		return self.unit1.forward(x)

	def reconstruct(self, x: Mat) -> Mat:
		return self.unit1.reconstruct(x)


def make_gaussian_blobs(n_cases: int, n: int = N) -> Mat:
	xs = np.arange(n, dtype=np.float64)
	ys = np.arange(n, dtype=np.float64)
	xx, yy = np.meshgrid(xs, ys)
	data = np.zeros((n_cases, n * n), dtype=np.float64)
	for i in range(n_cases):
		cx = random.uniform(1.0, n - 2.0)
		cy = random.uniform(1.0, n - 2.0)
		img = np.exp(-((xx - cx) ** 2 + (yy - cy) ** 2) / 4.0)
		data[i] = img.flatten()
	return data


def test_gaussian_autoencoder():
	model = GaussianAEModel("gaussian_autoencoder_pytest", "results")
	model.unit1.training_params.num_epochs = 1500
	model.init(0.01)
	train_batch = normalize(make_gaussian_blobs(N_CASES))
	model.train(train_batch)
	for x in train_batch[:10]:
		recon = model.reconstruct(model.forward(x))
		assert recon.size == x.size


def main():
	model = GaussianAEModel("gaussian_autoencoder", "results")
	model.init(0.01)

	train_batch = normalize(make_gaussian_blobs(N_CASES))
	model.train(train_batch)
	model.save()

	test_batch = normalize(make_gaussian_blobs(10))
	n_show = len(test_batch)

	fig, axes = plt.subplots(2, n_show, figsize=(n_show * 1.5, 3))
	axes[0, 0].set_ylabel("Original")
	axes[1, 0].set_ylabel("Recon")

	total_error = 0.0
	for i, inp in enumerate(test_batch):
		h = model.forward(inp)
		recon = model.reconstruct(h)
		total_error += float(np.mean(np.abs(inp - recon)))

		axes[0, i].imshow(np.asnumpy(np.reshape(inp, (N, N))), cmap='viridis')
		axes[0, i].axis('off')
		axes[1, i].imshow(np.asnumpy(np.reshape(recon, (N, N))), cmap='viridis')
		axes[1, i].axis('off')

	print(f"Mean reconstruction error: {total_error / n_show:.4f}")
	plt.tight_layout()
	plt.show()


if __name__ == "__main__":
	main()