#!/usr/bin/env python
# Solve the "Primal Squares" problem from:
#  http://www.recmath.org/contest/description.php
# Adam Sampson <ats@offog.org>
#
# The approach is to use a simple genetic algorithm to breed good
# squares. The bottleneck is the code that tests numbers for primality:
# it would probably help to add one of the probabilistic tests for
# compositeness first. It would probably also be worth introducing a few
# new random individuals each generation.

from __future__ import division

import math
from random import randrange
from copy import deepcopy

primes_list = [2, 3, 5]
primes_limit = 5
is_prime_memo = {}
def is_prime(n):
	global primes_limit

	if is_prime_memo.has_key(n):
		return is_prime_memo[n]

	v = (n >= primes_list[0])
	limit = int(math.sqrt(n))
	if primes_limit < limit:
		print "[computing primes from", primes_limit + 1, "up to", limit, "...",
		for i in xrange(primes_limit + 1, limit + 1):
			if is_prime(i):
				primes_list.append(i)
		print "now", len(primes_list), "primes]"
		primes_limit = limit
	for prime in primes_list:
		if prime > limit:
			break
		if n % prime == 0:
			v = False
			break

	is_prime_memo[n] = v
	return v

def find_in_line(line, out):
	s = 0
	n = len(line)
	for l in range(1, n + 1):
		for i in range(0, n - l + 1):
			v = 0
			for digit in line[i:i+l]:
				v = (10 * v) + digit
			if is_prime(v):
				out[v] = True

def rev(l):
	ll = [] + list(l)
	ll.reverse()
	return ll

def find_in_rect(grid, out):
	for line in grid:
		find_in_line(line, out)
		find_in_line(rev(line), out)
	for line in apply(zip, grid):
		find_in_line(line, out)
		find_in_line(rev(line), out)

def score_grid(grid):
	out = {}
	find_in_rect(grid, out)

	diag = []
	n = len(grid)
	for i in range(n):
		diag.append([0] * i + list(grid[i]) + [0] * (n - i - 1))
	find_in_rect(diag, out)

	diag = []
	n = len(grid)
	for i in range(n):
		diag.append([0] * (n - i - 1) + list(grid[i]) + [0] * i)
	find_in_rect(diag, out)

	p2_score = sum(map(lambda x: len(str(x)), out.keys()))
	digits = {}
	for line in grid:
		for x in line:
			digits[x] = True
	p2_score += len([x for x in digits.keys() if x % 2 == 0])
	return (len(out.keys()), p2_score)

grid_size = 6

def random_grid():
	g = []
	for i in xrange(grid_size):
		g.append(map(lambda x: randrange(10), xrange(grid_size)))
	return g

def print_grid(grid):
	print "\n".join(map(lambda l: " ".join(map(str, l)), grid))

def main():
	population = []
	pop_size = 50

	for i in xrange(pop_size):
		population.append(random_grid())

	for gen_number in xrange(100):
		# Sort the population by fitness.
		scores = {}
		for i in xrange(pop_size):
			scores[i] = score_grid(population[i])[0]
		ns = range(pop_size)
		ns.sort(lambda a, b: cmp(scores[a], scores[b]))
		print "generation", gen_number, "mean", sum(scores.values()) / pop_size, "best", scores[ns[-1]]
		print_grid(population[ns[-1]])

		new_pop = []
		for i in xrange(pop_size // 4):
			# Interbreed two randomly-chosen individuals by
			# swapping a block of numbers between them.
			a = randrange(pop_size)
			b = randrange(pop_size)
			a_grid = deepcopy(population[a])
			b_grid = deepcopy(population[b])

			w = randrange(grid_size - 1)
			h = randrange(grid_size - 1)
			x = randrange(grid_size - w)
			y = randrange(grid_size - h)
			for xi in xrange(x, x + w):
				for yi in xrange(y, y + h):
					t = a_grid[yi][xi]
					a_grid[yi][xi] = b_grid[yi][xi]
					b_grid[yi][xi] = t

			if randrange(100) < 20:
				n = randrange((grid_size * grid_size) // 4)
				for j in xrange(n):
					a_grid[randrange(grid_size)][randrange(grid_size)] = randrange(10)
					b_grid[randrange(grid_size)][randrange(grid_size)] = randrange(10)

			new_pop.append(a_grid)
			new_pop.append(b_grid)

		# Keep some of the old population.
		population = new_pop + [population[ns[i]] for i in xrange(len(new_pop), pop_size)]

main()