import pygame, sys
from pygame.locals import *
import time
import random
from operator import *
import traceback
import os
import getpass
import cPickle as pickle
 
class GA:
    def __init__(self,
                 num_samples = 1,
                 num_selected = 1,
                 mutation_factor = 4):
        pygame.init()
        pygame.display.init()
        self.screen = pygame.display.set_mode((152, 200), 0, 32)
        while True:
            try:
                print "128 and 256 are good choices. 256 is better but also slower."
                self.ncircles = input("Enter number of circles you want to use: ")
                break
            except NameError:
                print "Must enter an integer!"
            except SyntaxError:
                print "Must enter an integer!"
        grayscale = raw_input("Is this picture grayscale? Y/N\n")
        accept = ["y", "Y", "n", "N"]
        while grayscale not in accept:
            grayscale = raw_input("Is this picture grayscale? Y/N\n")
        if grayscale == "y" or grayscale == "Y":
            self.grayscale = True
        else:
            self.grayscale = False
        self.imgpath = raw_input("Enter full path of file, including file name and extension.\n")
        while os.path.exists(self.imgpath) is False:
            print "Path/file does not exist!"
            self.imgpath = raw_input("Enter full path of file, including file name and extension.\n")
        self.bg = pygame.image.load(self.imgpath).convert()
        self.width = self.bg.get_width()
        self.height = self.bg.get_height()
        self.path = "C:/users/" + str(getpass.getuser()) + "/desktop/imgs/"
        self.num_samples = num_samples
        self.num_selected = num_selected
        self.mutation_factor = mutation_factor
        self.target = self.get_target()
    def get_target(self):
        x = self.width
        y = self.height
        pixelData = []
        for ypixel in xrange(0, y):
            for xpixel in xrange(0, x):
                pixelData.append(self.bg.get_at((xpixel, ypixel)))
                for event in pygame.event.get():
                    if event.type == QUIT:
                        pygame.quit()
                        sys.exit()
                    pygame.display.update()
        pygame.quit()
        return pixelData
    def get_state(self, screen):
        x = self.width
        y = self.height
        pixelData = []
        for ypixel in xrange(0, y):
            for xpixel in xrange(0, x):
                pixelData.append(screen.get_at((xpixel, ypixel)))
        return pixelData
    def generate_random_chromosome(self):
        chromo = []
        for n in xrange(1, self.ncircles+1):
            a = random.randint(0, 255)
            if self.grayscale:
                randcolor = (a,a,a,
                             random.randint(0, 255))
            else:
                randcolor = (random.randint(0, 255),
                             random.randint(0, 255),
                             random.randint(0, 255),
                             random.randint(0, 255))
            randpos = (random.randint(0, self.height),
                     random.randint(0, self.width))
            randradius = random.randint(1, self.height/3)
            chromo.append((n, randcolor, randpos, randradius))
        print "Random chromosome generated."
        return chromo
    def fitness(self, chromo):
        pygame.init()
        s = pygame.Surface((self.width, self.height), flags=pygame.SRCALPHA)
        total_fitness = 0
        for args in chromo:
            s.lock()
            pygame.draw.circle(s, args[1], args[2], args[3])
            s.unlock()
        state = self.get_state(s)
        s.fill((0, 0, 0))
        cpix = 0
        for pixel in self.target:
            if self.grayscale:
                current_fitness = abs(pixel[0]-state[cpix][0])
            else:
                current_fitness = abs(pixel[0]-state[cpix][0])+abs(pixel[1]-state[cpix][1])+abs(pixel[2]-state[cpix][2])
            cpix += 1
            total_fitness += current_fitness
        return total_fitness
    def mutate(self, chromo):
        mutations = 0
        mutateTo = random.randint(1, self.mutation_factor)
        mutateChoice = ["color", "opacity", "pos", "radius", "swap"]
        while mutations < mutateTo:
            argchange = random.choice(mutateChoice)
            if argchange == "swap":
                swap1 = random.randint(0, len(chromo)-1)
                swap2 = random.randint(0, len(chromo)-1)
                while swap1 == swap2:
                    swap2 = random.randint(0, len(chromo)-1)
                swapvalue1 = chromo[swap1]
                swapvalue2 = chromo[swap2]
                chromo[swap1] = swapvalue2
                chromo[swap2] = swapvalue1
                mutations += 1
            else:
                argint = random.randint(0, len(chromo)-1)
                args = list(chromo[argint])
                args[1] = list(args[1])
                args[2] = list(args[2])
                a = random.randint(0, 255)
                if argchange == "color":
                    if self.grayscale:
                        args[1][0] = a
                        args[1][1] = a
                        args[1][2] = a
                    else:
                        args[1][0] = random.randint(0, 255)
                        args[1][1] = random.randint(0, 255)
                        args[1][2] = random.randint(0, 255)
                elif argchange == "opacity":
                    args[1][3] = random.randint(0, 255)
                elif argchange == "pos":
                    args[2][0] = random.randint(0, self.width)
                    args[2][1] = random.randint(0, self.height)
                elif argchange == "radius":
                    args[3] = random.randint(1, self.height/3)
                args[2] = tuple(args[2])
                args[1] = tuple(args[1])
                args = tuple(args)
                chromo[argint] = args
                mutations += 1
        return chromo
    def display(self, data, gen):
        bif = "C:/users/justin/pictures/charles-darwin.jpg"
        mif = "C:/users/justin/pictures/mouse.png"
        path = self.path + str(gen) + "-" + str(self.fitness(data)/10000) + ".png"
        pygame.init()
        screen = pygame.display.set_mode((self.width, self.height), 0, 32)
        s = pygame.Surface((self.width, self.height), flags=pygame.SRCALPHA)
        s.fill((0, 0, 0))
        a = 0
        while True:
            for event in pygame.event.get():
                if event.type == QUIT:
                    pygame.quit()
                    sys.exit()
                   
            for args in data:
                s.lock()
                pygame.draw.circle(s, args[1], args[2], args[3])
                s.unlock()
                screen.blit(s, (0, 0))
                pygame.display.update()
               
            pygame.image.save(screen, path)
            screen.fill((0, 0, 0))
            pygame.quit()
            break
    def run(self):
        #Creates a random chromosome
        sample = self.generate_random_chromosome()
        load_pickle = raw_input("Would you like to load the previous genome? Y/N\n")
        accept = ["y", "Y", "n", "N"]
        while load_pickle not in accept:
            load_pickle = raw_input("Would you like to load the previous genome? Y/N\n")
        if load_pickle == "y" or load_pickle == "Y":
            generation = pickle.load(open("ge.neration", "rb"))-1
            mother = pickle.load(open("ge.nome", "rb"))
        else:
            generation = -1
            mother = sample
        try:
            while self.fitness(sample) != 0:
                generation += 1
                os.system("title Genetic Algorithm: Generation " + str(generation))
                #Mutate the chromo
                temp = []
                for args in mother:
                    temp.append(args)
                daughter = self.mutate(mother)
                if self.fitness(temp) <= self.fitness(daughter):
                    mother = temp
                    new = False
                else:
                    mother = daughter
                    new = True
                print "The best chromosome of generation %s is %s" %(
                    generation, self.fitness(mother))
                if new:
                    self.display(mother, generation)
                    new = False
                if self.fitness(mother) <= 800000:
                    self.mutation_factor = 1
        except:
            pickle.dump(mother, open("ge.nome", "w"))
            pickle.dump(generation, open("ge.neration", "w"))
            raw_input("Press enter to quit...\n")
 
if __name__ == "__main__":
    os.system("title Genetic Algorithm")
    print "You MUST have a folder called \"imgs\" without the quotes on your desktop."
    raw_input("Press enter to start the program.\n")
    ga = GA()
    ga.run()