#!/usr/bin/env python3
import pygame, random, sys
from pygame.locals import *

from random import randint

from genevo.bitint import BitInt
from genevo.engine import Engine

GREEN   = Color (000, 255, 000)
WHITE   = Color (255, 255, 255)
BLACK   = Color (000, 000, 000)

WIDTH   = 600
HEIGHT  = 600
CELLXNR =  60 # number of cells
CELLYNR =  60
CELLWIDTH  = WIDTH // CELLXNR
CELLHEIGHT = HEIGHT // CELLYNR

FPS  = pygame.time.Clock ()
SURF = pygame.display.set_mode ((WIDTH, HEIGHT))

clade1 = BitInt (random.getrandbits (16))
clade2 = BitInt (random.getrandbits (16))
clade3 = BitInt (random.getrandbits (16))
clade4 = BitInt (random.getrandbits (16))

def make_init (cx, cy):
    return [ (x % CELLXNR, y % CELLYNR) for x in (cx - 1, cx, cx + 1)
                    for y in (cy - 1, cy, cy + 1) ]

init = {
        clade: sum (
            (make_init (randint (0, CELLXNR - 1), randint (0, CELLYNR - 1))
                for _ in range (randint (1, 10))), []
            )
            for clade in (clade1, clade2, clade3, clade4)
}

game_engine = Engine (CELLXNR, CELLYNR, init, debug = 0)
gens = 0


def cladehash (cell):
    c = [0, 0, 0]
    c [(cell.genom % 2) + 1] = cell.trait
    return c

def aggrohash (cell):
    return (cell.aggro << 2, 0, 0)

def attackhash (cell):
    return (0, 0, cell.attack << 2)

def agehash (cell):
    return (0, cell.age if cell.age < 256 else 255, 0)

def overlay (cell):
    # display traits in greyscale, to make the overlays easier to spot
    return (cell.trait,) * 3

def drawcell (x, y, color):
    pygame.draw.rect (
        SURF, color, 
        (x * CELLWIDTH, y * CELLHEIGHT, CELLWIDTH, CELLHEIGHT)
    )

def main ():
    cellhash = [cladehash, aggrohash, attackhash, agehash, overlay]
    hashno = 0

    while 1:

        dead, born = game_engine.step ()

        for event in pygame.event.get ():
            if event.type == QUIT:
                pygame.quit ()
                sys.exit ()
            if event.type == KEYDOWN:
                k = event.key - 48
                if 0 < k <= len (cellhash):
                    hashno = k - 1

        SURF.fill (WHITE)
        for x in range (CELLXNR):
            for y in range (CELLYNR):
                cell = game_engine [x] [y]
                if cell != None:
                    drawcell (x, y, cellhash [hashno] (cell))
                if hashno == len (cellhash) - 1:
                    # overlays
                    if   (x, y) in dead and dead [x, y] == "Was murdered.":
                        drawcell (x, y, (255, 0, 0, 200))
                    elif (x, y) in born and born [x, y] [1]:
                        drawcell (x, y, (0, 0, 255, 200))

        [pygame.draw.line (SURF, BLACK, (i, 0), (i, HEIGHT))
                for i in range (0, WIDTH, CELLWIDTH)]
        [pygame.draw.line (SURF, BLACK, (0, i), (WIDTH, i))
                for i in range (0, HEIGHT, CELLHEIGHT)]

        pygame.display.update ()
        FPS.tick (16)

try:
    main ()
except KeyboardInterrupt:
    pop = game_engine.count_clades ()
    print ("")
    print ("{} different genomes among {} individuals observed.".format (
        len (pop),
        sum (pop.values ())
    ))
    print ("{} died. {} born. {} born as mutants. Mutationrate {}.".format (
        game_engine.deaths, game_engine.births, game_engine.mutations,
        game_engine.mutations / game_engine.births
    ))
