from PIL import Image, ImageDraw
import imageio
import math
import scipy.optimize
import matplotlib.pyplot as plt

n = 8
image_size = 400
gif_writer = imageio.get_writer('b.gif', mode='I')

def draw_board(y, obstacle=None):
	image = Image.new('RGB', (image_size, image_size), color=(255, 255, 200))
	draw = ImageDraw.Draw(image)
	# 0..7  — удобные координаты → пиксельные
	cell_size = image_size / n
	if obstacle:
		x_obstacle, y_obstacle, r_obstacle = obstacle
		xc_pixel = (x_obstacle + 0.5) * cell_size
		yc_pixel = (y_obstacle + 0.5) * cell_size
		r_pixel = r_obstacle * cell_size
		draw.ellipse((xc_pixel - r_pixel, yc_pixel - r_pixel, xc_pixel + r_pixel, yc_pixel + r_pixel), fill=(255, 0, 0))
	for i in range(n):  # 0 .. n-1
		xc_pixel = (i + 0.5) * cell_size
		yc_pixel = (y[i] + 0.5) * cell_size
		r_pixel = 0.3 * cell_size
		draw.ellipse((xc_pixel - r_pixel, yc_pixel - r_pixel, xc_pixel + r_pixel, yc_pixel + r_pixel), fill=(0, 0, 240))
	image.save('a.png')
	gif_writer.append_data(imageio.imread('a.png'))


def solve(obstacle):
	def is_forbidden(x, y):
		return math.hypot(x - obstacle[0], y - obstacle[1]) < obstacle[2]
	penalty_matrix = [[is_forbidden(x, y) for y in range(n)] for x in range(n)]
	assignment = list(scipy.optimize.linear_sum_assignment(penalty_matrix)[1])
	draw_board(assignment, obstacle)
	return sum([is_forbidden(x, assignment[x]) for x in range(n)])


scores = []
for frame in range(50):
	score = solve((2 + frame * 0.1, 1 + frame * 0.05, 2 + frame * 0.05))
	scores.append(score)

plt.plot(scores)
plt.title('Lost rooks')
plt.savefig('scores.svg')