import random
import math

class Neural_net:
	"""docstring for Neural_net"""
	def __init__(self, complexity):
		self.complexity = complexity
		self.nodes = self.generate_nodes()
		self.genes = self.generate_genes(32)

	def generate_genes(self, amount):
		genes = []
		while len(genes) < amount:
			trans = self.random_transformation()
			performance = 0
			for a in range(1000):
				number_1 = random.randint(0, 10)
				number_2 = random.randint(0, 10)
				result = self.get_answer(trans, [number_1, number_2])
				if (number_1 + number_2) >= 5:
					if result == 1:
						performance += 1
					else:
						performance -= 1
				else:
					if result == 1:
						performance -= 1
					else:
						performance += 1

			if performance > 500:
				genes.append([trans, performance])

		return genes

	def generate_nodes(self):
		nodes = []
		for i in range(len(self.complexity)):
			node_row = list(x for x in range(self.complexity[i]))
			nodes.append(node_row)

		return nodes

	def random_transformation(self):
		transformation = []
		for i in range(1, len(self.complexity)):
			transformation.append([])
			connections = self.complexity[i] * self.complexity[i-1]
			for _ in range(connections):
				transformation[i-1].append(random.uniform(-1.0, 1.0))

		return transformation

	def get_answer(self, transformation, param):
		# The first nodes are input data
		for i in range(len(self.nodes[0])):
			self.nodes[0][i] = param[i]


		for a in range(1, len(self.nodes)):
			for b in range(len(self.nodes[a])):
				self.nodes[a][b] = 0
				for c in range(len(self.nodes[a-1])):
					self.nodes[a][b] += self.nodes[a-1][c] * transformation[a-1][b+1]

				# Normalize the value
				self.nodes[a][b] = math.tanh(self.nodes[a][b])

		# Choose output by looking at which of the two output
		# nodes has the highest value
		result = 0.0
		output = 0
		for i in range(len(self.nodes[-1])):
			if self.nodes[-1][i] > result:
				output = i
				result = self.nodes[-1][i]

		return output

	def train(self, iterations):
		for a in range(iterations):
			# Kill the worst genes
			self.genes = sorted(self.genes, key=lambda x: x[1], reverse=True)
			self.genes = self.genes[:(len(self.genes) // 2)]

			print("AAA")
			for gene in self.genes:
				print(gene[1], end=", ")

			print("\n")
			# Breed new genes
			random.shuffle(self.genes)
			for g in range(len(self.genes)):
				new_gene = self.combine_transformations(self.genes[g][0], self.genes[g+1][0])
				self.genes.append([new_gene, 0])

			# Test the genes
			for b in range(len(self.genes)):
				performance = 0
				for c in range(1000):
					number_1 = random.randint(0, 10)
					number_2 = random.randint(0, 10)
					result = self.get_answer(self.genes[b][0], [number_1, number_2])
					if (number_1 + number_2) >= 5:
						if result == 1:
							performance += 1
						else:
							performance -= 1
					else:
						if result == 1:
							performance -= 1
						else:
							performance += 1
				self.genes[b][1] = performance

	def combine_transformations(self, trans_one, trans_two):
		new_trans = []
		for a in range(len(trans_one)):
			new_trans.append([])
			# TODO Use zip
			for b in range(len(trans_one[a])):
				# Alternate choosing connections between trans_one and trans_two
				if b % 2 == 0:
					connection = trans_one[a][b]
				else:
					connection = trans_two[a][b]
				new_trans[a].append(connection)
		return new_trans


nn = Neural_net([2,8,8,2])
nn.train(1000)