import pygame
import traceback
import numpy
import random


class Monster:
    def __init__(self):
        self.color = (random.randint(64, 255),
                      random.randint(64, 255),
                      random.randint(64, 255))
        # A bunch of random booleans
        self.pixels = numpy.random.randint(0, 2, size=(8, 4))  # 8 rows by 4 columns

    def to_rgb(self):
        # A bunch of numpy-fu to turn 8x4 into 8x4x3 wherever colors is 1
        full = numpy.hstack((self.pixels, self.pixels[:, ::-1]))
        colors = full[:, :, numpy.newaxis]
        colors = self.color * colors
        # This would also work:
        # colors = numpy.zeros(shape=(8,8,3))
        # colors[numpy.nonzeros(full)] = color
        return colors

    def draw_on(self, screen, x, y):
        screen.lock()
        pxs = self.to_rgb()
        for row in range(pxs.shape[0]):
            for col in range(pxs.shape[1]):
                screen.set_at((x + col * 2 + 0, y + row * 2 + 0), pxs[row, col])
                screen.set_at((x + col * 2 + 0, y + row * 2 + 1), pxs[row, col])
                screen.set_at((x + col * 2 + 1, y + row * 2 + 0), pxs[row, col])
                screen.set_at((x + col * 2 + 1, y + row * 2 + 1), pxs[row, col])

        screen.unlock()

    def mutate(self):
        # each pixel has a random chance of toggling
        for y in range(8):
            for x in range(4):
                if random.random() < 0.01:
                    self.pixels[y, x] = 1 - self.pixels[y, x]
        # also a small chance of flipping vertically or horizontally
        if random.random() < 0.01:
            self.pixels = self.pixels[::-1, :]
        if random.random() < 0.01:
            self.pixels = self.pixels[:, ::-1]

    def crossover_with(self, mon2):
        ind = Monster()
        # decide whether to go by columns or rows
        if random.random() < 0.5:
            split_point = random.randint(1, 3)
            # take every column before split_point from me
            ind.pixels[:, :split_point] = self.pixels[:, :split_point]
            # and every column after from mon2
            ind.pixels[:, split_point:] = mon2.pixels[:, split_point:]
        else:
            split_point = random.randint(1, 7)
            # take every row before split_point from me
            ind.pixels[:split_point, :] = self.pixels[:split_point, :]
            # and every row after from mon2
            ind.pixels[split_point:, :] = mon2.pixels[split_point:, :]
        return ind

def breed_monsters(mon1pos, mon2pos, monsters):
    mon1 = monsters[mon1pos[0]*8+mon1pos[1]]
    mon2 = monsters[mon2pos[0]*8+mon2pos[1]]
    new_monsters = [mon1.crossover_with(mon2) for _ in range(len(monsters)-2)]
    for mon in new_monsters:
        mon.mutate()
    return [mon1,mon2]+new_monsters

def play():
    pygame.init()
    screen = pygame.display.set_mode((320, 240))  # "Surface"
    # We'll fit 10 monsters across and 8 vertically
    mon_count = 10 * 8
    monsters = [Monster() for _ in range(mon_count)]
    done = False
    clock = pygame.time.Clock()
    parent1 = None
    parent2 = None
    selected = None
    was_pressed = False
    pressed = False
    while not done:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                done = True
        keys = pygame.key.get_pressed()
        if keys[pygame.K_ESCAPE]:
            done = True
        screen.fill((0, 0, 0))
        was_pressed = pressed
        pressed = pygame.mouse.get_pressed()[0]
        loc = pygame.mouse.get_pos()
        # click to select one then the other
        if pressed and not was_pressed:
            loc_mon = (loc[1]//32, loc[0]//32)
            if selected is not None:
                # mate!
                selected = None
                parent2 = loc_mon
                monsters = breed_monsters(parent1, parent2, monsters)
                parent1 = (0,0)
                parent2 = (0,1)
            else:
                selected = loc_mon
                parent1 = loc_mon
        for row in range(8):
            for col in range(10):
                mon = monsters[row * 8 + col]
                mon.draw_on(screen, col * 32, row * 32)
                if selected == (row, col):
                    pygame.draw.rect(screen,
                                     (0,255,0),
                                     (col*32-2,row*32-2,16+4,16+4),
                                     2)
        # b to breed again from same parents
        if keys[pygame.K_b]:
            monsters = breed_monsters(parent1, parent2, monsters)
        # r to randomize population completely
        if keys[pygame.K_r]:
            monsters = [Monster() for _ in range(mon_count)]
        pygame.display.flip()
        clock.tick(10)
    pygame.quit()


try:
    play()
except Exception:
    traceback.print_exc()
    pygame.quit()