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

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, 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)
		#actual transparency
		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]#[f[TRA] * x for x in f[COL]]
		figures[i][PIX] = (rr, cc)
	return image

def figure_fitness(img, im):
	score = 0
	no_problem = False

	N, M = img.shape[0], img.shape[1]
	for i in range(N):
		for j in range(M):
			score += sum(abs(img[i][j] - im[i][j]))
	if score != 0:
		return score
	return -1


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(ind1, ind2):
	rez1, rez2 = [], []
	chance = np.random.randint(2,size=len(ind2[0]))
	for z in zip(ind1[0], ind2[0], chance):
		rez1.append(z[z[2]])
		rez2.append(z[1 - z[2]])
	bks = ind1[1], ind2[1]
	return (rez1, bks[1 - chance[0]]), (rez2, bks[chance[0]])

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_figs(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.03
	for h in humans:
		mutants.append(mutate_ind(h.copy(), chance, maxes))
	return mutants

def mutate_img(indiv, maxes, chances=None, backchance = 0.05):
	bak = indiv[1].copy()
	if random.random() < backchance:
		bak[COL] = np.random.uniform(size=3) * maxes[COL]
	return (mutate_figs(indiv[0], maxes=maxes), bak)


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

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

figs_num = 15
POP_SIZE = 50




maxes = [0, imgsize[0]/2, imgsize[1]/2, np.array([im.shape[0], im.shape[1]]), np.array([1, 1, 1]), 1, 1, 0]

MAX_GEN = 25
ATTEMPTS = 1

for a in range(ATTEMPTS):
	dirname = 'attempt{}'.format(a)
	if not os.path.exists(dirname):
		os.makedirs(dirname)

	individuals = []
	for i in range(POP_SIZE):
		figs = get_random_figures(figs_num, maxh=imgsize[0]/2, maxl=imgsize[1]/2, c=imgsize, ftype=1)
		background =[1, im.shape[0], im.shape[1], [im.shape[0]/2, im.shape[1]/2], figs[0][COL], 1, 0, None]
		individuals.append((figs, background))
	for g in range(MAX_GEN):

		inds = []
		cnt = 0
		fit_imgs = []
		for ind in individuals:
			img = generate_image(ind[0] + [ind[1]], imgsize)
			fit = figure_fitness(img, im)
			if fit != 0:
				cnt += 1
				fit_imgs.append((ind, fit))

		fit_imgs = sorted(fit_imgs, key=itemgetter(1))
		print(str(g) + ' : ' +str(cnt) + '[' + str(fit_imgs[0][1]) + ']')

		pop = list(zip(*fit_imgs))[0]
		mid = len(pop)//2
		quart = mid // 2
		fst, snd, trd, fth = pop[:quart], pop[quart:mid], pop[mid:mid+quart], pop[mid+quart:]

		elites = []
		for s, t in zip(fst, snd):
			#img = generate_image(fst[0][0] + [fst[0][1]], imgsize)
			#print(str(g) + ' : ' +str(cnt) + '[' + str(figure_fitness(img, im)) + ']')
			i1, i2 = cross_over(s, t)
			#img = generate_image(fst[0][0] + [fst[0][1]], imgsize)
			#print(str(g) + ' : ' +str(cnt) + '[' + str(figure_fitness(img, im)) + ']')
			elites.append(i1)
			elites.append(mutate_img(i2, maxes))

		#img = generate_image(fst[0][0] + [fst[0][1]], imgsize)
		#print(str(g) + ' : ' +str(cnt) + '[' + str(figure_fitness(img, im)) + ']')
		boomers = []
		for s, t in zip(snd, trd):
			i1, i2 = cross_over(s, t)
			boomers.append(mutate_img(i1, maxes))
			#boomers.append(mutate_img(i2, maxes))

		individuals = []
		individuals += fst
		individuals += elites
		#individuals += newborns
		individuals += boomers
		#figs += doomers


		img = generate_image(fst[0][0] + [fst[0][1]], imgsize)
		scipy.misc.imsave(dirname + '/outfile{}.png'.format(g), img)
		#print(str(g) + ' : ' +str(cnt) + '[' + str(figure_fitness(img, im)) + ']')
	#ax1.imshow(img)
	#ax2.imshow(im)

	#plt.show()

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