import math
import numpy as np
import matplotlib.pyplot as plt
import itertools
import random
import scipy.misc

from PIL import Image
from operator import add, itemgetter
from skimage.draw import (line, polygon, circle,
                          circle_perimeter,
                          ellipse, ellipse_perimeter,
                          bezier_curve, set_color, rectangle)

TYPE = 0 # 0 elipse, 1 square, 2 else
LEN = 1 #length
HEI = 2 #height
CEN = 3 #center
COL = 4 #color
TRA = 5 #transparency
ZIN = 6 #z index
PIX = 7 #pixel coordinates


def new_color(cf, cb, t):
	cf1 = [x * t for x in cf]
	cb1 = [x * (1 - t) for x in cb]
	return list(map(add, cf1, cb1))


def get_poligon(f):
	cx, cy = f[CEN]
	l = f[LEN]
	h = f[HEI]
	p1 = [cx - l, cx - l, cx + l, cx + l]
	p2 = [cy - h, cy + h, cy + h, cy - h]
	return p1, p2

def get_random_figures(POP, maxl=120, maxh=120, c=(500, 500), ftype=None):
	FIGS = []
	types = None
	if ftype is not None:
		if ftype == 1:
			types = np.ones(POP)
		elif ftype == 0:
			types = np.zeros(POP)
	else:
		types = np.random.randint(2, size=POP)
	lengths = np.random.randint(low=1, high=maxl, size=POP)
	heights = np.random.randint(low=1, high=maxh, size=POP)
	ycenters = np.random.randint(c[0], size=POP)
	xcenters = np.random.randint(c[1], size=POP)
	centers = np.column_stack((ycenters, xcenters))
	colors = np.random.uniform(size=(POP, 3))
	transps = np.random.uniform(size=(POP))
	zs = np.random.uniform(size=(POP))

	for i in range(POP):
		FIGS.append([types[i], lengths[i], heights[i], centers[i], colors[i], transps[i], zs[i], None])
	return FIGS


def generate_image(figures, size, front_pixels, transp=False):
	image = np.zeros((size[0], size[1], 3), dtype=np.double)
	figures = sorted(figures, key=itemgetter(ZIN))
	for i in range(len(figures)):
		rr = 0
		cc = 0
		f  = figures[i]
		if f[TYPE] == 0:
			rr, cc = ellipse(f[CEN][0], f[CEN][1], f[LEN] / 2, f[HEI] / 2, image.shape)
		elif f[TYPE] == 1:
			p1, p2 = get_poligon(f)
			rr, cc = polygon(p1, p2, image.shape)
		if transp:
			image[rr, cc] = [new_color(f[COL], x, f[TRA]) for x in image[rr, cc]]
		else:
			image[rr, cc] = f[TRA] * f[COL]
		figures[i][PIX] = (rr, cc)
		front_pixels[rr, cc] = i
	return image, front_pixels

def figure_fitness(figure, img, im):
	rr, cc = figure[PIX]
	score = 0
	no_problem = False
	clr = [figure[TRA] * x for x in figure[COL]]

	for r, c in zip(rr, cc):
		score += sum(abs(im[r][c] - clr))

	return score / max(1, len(rr))



def get_children(parents, count):
	children = []
	for i in range(count):
		p1, p2 = random.choice(parents), random.choice(parents)
		chance = np.random.randint(2,size=8)
		child = []
		for p in zip(p1, p2, chance):
			child.append(p[p[2]])
		children.append(child)
	return children

def cross_over(men, women):
	aphrodites = []
	for x, y in zip(men, women):
		x1, y1 = [], []
		chance = np.random.randint(2,size=8)
		for z in zip(x, y, chance):
			x1.append(z[z[2]])
			y1.append(z[1 - z[2]])
		aphrodites.append(x1)
		aphrodites.append(y1)
	return aphrodites

def mutate_ind(f, chance, maxes):
	for i in range(1, len(f) - 1):
		if random.random() < chance[i]:
			f[i] = random.random() * maxes[i]
			if i == CEN:
				f[CEN] = np.random.uniform(size=2) * maxes[CEN]
			if i == COL:
				f[COL] = np.random.uniform(size=3) * maxes[COL]
	return f

def mutate_pop(humans, chance=None, maxes=[0, 100, 100, np.array([500, 500]), np.array([1, 1, 1]), 1, 1, 0]):
	mutants = []
	if chance is None:
		chance = np.random.uniform(size=8)
	chance *= 0.07
	for h in humans:
		mutants.append(mutate_ind(h.copy(), chance, maxes))
	return mutants



fig, (ax1, ax2) = plt.subplots(ncols=2, nrows=1, figsize=(10, 6))


#im = Image.open("mona_lisa.png")
im = Image.open("windcl.png")
im=np.asarray(im)
imgsize = (im.shape[0], im.shape[1])

population = 45
MAX_GEN = 50
ATTEMPTS = 1
MAX_DIM = 0.5

front_pixels = np.zeros((im.shape[0], im.shape[1]))
maxes = [0, imgsize[0] * MAX_DIM, imgsize[1] * MAX_DIM, np.array([im.shape[0], im.shape[1]]), np.array([1, 1, 1]), 1, 1, 0]


def visible_figs(figures, front_pixels):
	goodfigs = []
	print(np.unique(front_pixels.flat))
	for idx in np.unique(front_pixels.flat):
		goodfigs.append(figures[int(idx)])
	return goodfigs

for a in range(ATTEMPTS):
	figs = get_random_figures(population, maxh=imgsize[0] * MAX_DIM, maxl=imgsize[1] * MAX_DIM, c=imgsize, ftype=1)
	img, front_pixels = generate_image(figs, imgsize, front_pixels)
	for g in range(MAX_GEN):

		figs1 = []
		cnt = 0

		for f in figs:
			fit = figure_fitness(f, img, im)
			cnt += 1
			figs1.append((f, fit))

		print('cnt: ' + str(cnt) + ' fitness: '  + str(sum([x[1] for x in figs1])))
		figs1 = sorted(figs1, key=itemgetter(1))
		figs1 = list(zip(*figs1))[0]
		mid = len(figs1)//2
		quart = mid // 2
		fst, snd, trd, fth = figs1[:quart], figs1[quart:mid], figs1[mid:mid+quart], figs1[mid+quart:]

		newborns = get_children(fst, (population - cnt)//2)
		newborns += get_children(fth, (population - cnt)//2)
		newborns = mutate_pop(newborns, maxes=maxes)
		boomers = cross_over(fst, snd)
		doomers = mutate_pop(trd, maxes=maxes)

		figs = []
		figs += fst
		figs += newborns
		figs += boomers
		figs += doomers
		print(len(figs))
		img, front_pixels = generate_image(figs, imgsize, front_pixels)
		figs = visible_figs(figs, front_pixels)
		front_pixels = np.zeros((im.shape[0], im.shape[1]))
		scipy.misc.imsave('outfile{}.png'.format(g), img)
	#ax1.imshow(img)
	#ax2.imshow(im)

	#plt.show()

	#scipy.misc.imsave('outfile{}.png'.format(a), img)