# load all of the libraries we'll need for the model
from ipythonblocks import BlockGrid
from IPython.display import clear_output
import random

# environment variables
world_width = 150
world_height = 150
num_patches = 20
patch_size = 5

# genetic algorithm variables
population_size = 30
mutation_rate = 0.01
generations = 100
evaluation_ticks = 1000
number_evaluations_per_generation = 1
visualization_display_frequency = 1

# load all of the functions we'll need for the model

def setup_world():
    """
    Sets up a grid world with randomly-placed food patches
    and returns the grid.
    """
    
    # setup world visualization
    grid = BlockGrid(world_width, world_height, block_size=4)
    grid.lines_on = False
    
    # choose random (x, y) positions to place food patches
    patch_centers = []
    
    for i in range(num_patches):
        patch_x = random.randrange(0, world_width)
        patch_y = random.randrange(0, world_height)
        
        patch_centers.append((patch_x, patch_y))
        
    # mark the centers of the food patches
    for (x, y) in patch_centers:
        grid[x, y].green = 255
        
    # fill in the food patches
    for iteration in range(patch_size):
        blocks_to_fill = []
        
        for x in range(world_width):
            for y in range(world_height):
                # if any of the non-diagonal blocks next to this block have food
                if (grid[(x - 1) % world_width, y].green == 255 or
                    grid[x, (y - 1) % world_height].green == 255 or
                    grid[(x + 1) % world_width, y].green == 255 or
                    grid[x, (y + 1) % world_height].green == 255):
                    blocks_to_fill.append([x, y])
                    
        for (x, y) in blocks_to_fill:
            grid[x, y].green = 255
        
    return grid

def convert_bit_list_to_int(bit_list):
    """
    Takes a list of bits (binary representation of a number) and converts it into the corresponding integer.
    """
    return int(str(bit_list).strip("[]").replace(", ", ""), 2)

def fitness_evaluation(world, genotype, display=False):
    """
    Evaluates the given genome of a forager and returns its fitness.
    """
    
    # map the genotype into its corresponding phenotype
    # first 8 bits encode the angle the forager turns when it doesn't find any food on the current grid
    forager_turn_no_food = convert_bit_list_to_int(genotype[0:8])
    
    # next 8 bits encode the angle the forager turns when it finds food on the current grid
    forager_turn_food = convert_bit_list_to_int(genotype[8:16])
    
    # last 8 bits encode how many ticks the forager stays in the "food" state when it doesn't find any food
    # on the current grid before moving into the "no food" state
    forager_memory = convert_bit_list_to_int(genotype[16:24])
    
    # forager starting location and facing
    forager_x = random.randrange(0, world_width)
    forager_y = random.randrange(0, world_height)
    forager_facing = random.randrange(0, 360)
    
    ticks_since_last_food = forager_memory
    
    for tick in range(evaluation_ticks):
        ticks_since_last_food += 1
        
        # if food is present at current location OR memory limit has not been exceeded yet
        #     block.green == 255 and block.red == 0 means there is food and the predator hasn't eaten it
        #     block.green == 255 and block.red == 255 means there is food and the predator has eaten it
        #     block.green == 0 and block.red == 0 means there is no food and the predator hasn't visited that block
        #     block.green == 0 and block.red == 255 means there is no food and the predator has visited that block
        if (world[forager_x, forager_y].green == 255 and world[forager_x, forager_y].red == 0):
            ticks_since_last_food = 0
        
        # mark the location where the predator has visited
        world[forager_x, forager_y].red = 255
        
        # forager is in "food" mode
        if ticks_since_last_food < forager_memory:
            forager_facing = random.triangular(forager_facing - 2.0 * forager_turn_food, forager_facing + 2.0 * forager_turn_food)
        # forager is in "no food" mode
        else:
            forager_facing = random.triangular(forager_facing - 2.0 * forager_turn_no_food, forager_facing + 2.0 * forager_turn_no_food)

        forager_facing %= 360
        
        # move the forager in the direction it's facing (diagonal moves allowed)
        direction_to_move = int(((forager_facing + 22.5) % 360.0) / 45.0)
        
        offset_dict = { 0:(1, 0),    # to the right
                        1:(1, 1),    # up and to the right
                        2:(0, 1),    # up
                        3:(-1, 1),   # up and to the left
                        4:(-1, 0),   # to the left
                        5:(-1, -1),  # down and to the left
                        6:(0, -1),   # down
                        7:(1, -1) }  # down and to the right
        
        xoff, yoff = offset_dict[direction_to_move]
        
        forager_x += xoff
        forager_y += yoff
            
        # keep the forager within the bounds of the world
        forager_x %= world_width
        forager_y %= world_height
        
    # display the forager's activity if the display flag is active
    if display:
        clear_output()
        world.show()
        print("forager_turn_no_food: " + str(forager_turn_no_food))
        print("forager_turn_food: " + str(forager_turn_food))
        print("forager_memory: " + str(forager_memory))
            
    # evaluate the fitness by counting how many food blocks the forager picked up
    fitness = 0
    for block in world:
        if block.red == 255 and block.green == 255:
            fitness += 1
            
        # also unmark the predator's movement to "clean up" for the next evaluation
        block.red = 0
            
    # disallow <= 0 fitness
    if fitness <= 0:
        fitness = 1
            
    return fitness

# run the genetic algorithm to evolve the prey foragers

# best forager data
best_fitness_over_time = []
best_forager_turn_no_food_over_time = []
best_forager_turn_food_over_time = []
best_forager_memory_over_time = []

# population averages
population_fitness_over_time = []
population_turn_no_food_over_time = []
population_turn_food_over_time = []
population_memory_over_time = []

# seed the GA populaton with blank genotypes
population = []
for forager in range(population_size):
    genotype = []
    for bit in range(24):
        genotype.append(0)
    
    population.append(genotype)
    
for generation in range(1, generations + 1):
    # evaluate the fitnesses of the population of foragers
    world = setup_world()
    population_fitnesses = []
    highest_fitness = 0
    fittest_genotype = None
    for genotype in population:
        fitness = 0
        for evaluation_number in range(number_evaluations_per_generation):
            fitness += fitness_evaluation(world, genotype) / float(number_evaluations_per_generation)
            
        population_fitnesses.append(fitness)
        
        if fitness > highest_fitness:
            highest_fitness = fitness
            fittest_genotype = list(genotype)
            
    # keep track of best forager over time
    best_fitness_over_time.append(highest_fitness)
    best_forager_turn_no_food_over_time.append(convert_bit_list_to_int(fittest_genotype[0:8]))
    best_forager_turn_food_over_time.append(convert_bit_list_to_int(fittest_genotype[8:16]))
    best_forager_memory_over_time.append(convert_bit_list_to_int(fittest_genotype[16:24]))
    
    # keep track of population averages over time
    population_fitness = 0.0
    population_turn_no_food = 0.0
    population_turn_food = 0.0
    population_memory = 0.0
    
    for index, genotype in enumerate(population):
        population_fitness += population_fitnesses[index]
        population_turn_no_food += convert_bit_list_to_int(genotype[0:8])
        population_turn_food += convert_bit_list_to_int(genotype[8:16])
        population_memory += convert_bit_list_to_int(genotype[16:24])
        
    population_fitness_over_time.append(population_fitness / float(population_size))
    population_turn_no_food_over_time.append(population_turn_no_food / float(population_size))
    population_turn_food_over_time.append(population_turn_food / float(population_size))
    population_memory_over_time.append(population_memory / float(population_size))
    
    # display the performance of the fittest forager at a regular interval
    if generation % visualization_display_frequency == 0:
        fitness_evaluation(world, fittest_genotype, display=True)
        print("generation " + str(generation) + " fittest genotype")
    
    # apply fitness proportionate selection to determine the genotypes that reproduce into the next generation
    fitness_sum = sum(population_fitnesses)
    new_population = []
    
    for new_offspring_counter in range(population_size):
        roll = random.random() * fitness_sum
        sum_so_far = 0
        for index, genotype in enumerate(population):
            sum_so_far += population_fitnesses[index]
            
            # if the genotype is chosen to reproduce
            if sum_so_far >= roll:
                parent = list(genotype)
                offspring = []
                
                # copy the offspring's genotype from the parent's genotype
                for bit in parent:
                    # mutations sometimes occur
                    if random.random() <= mutation_rate:
                        offspring.append(random.randrange(0, 2))
                    else:
                        offspring.append(bit)
                        
                new_population.append(offspring)
                break
                
    population = list(new_population)

# run this code after the model finishes to see graphs of the forager's traits evolving over time
%pylab inline

figure()
title("Fitness over time")
plot(best_fitness_over_time, label="Best")
plot(population_fitness_over_time, label="Average")
legend()

figure()
title("Turn angle (no food) over time")
plot(best_forager_turn_no_food_over_time, label="Best")
plot(population_turn_no_food_over_time, label="Average")
legend()

figure()
title("Turn angle (w/ food) over time")
plot(best_forager_turn_food_over_time, label="Best")
plot(population_turn_food_over_time, label="Average")
legend()

figure()
title("Memory length over time")
plot(best_forager_memory_over_time, label="Best")
plot(population_memory_over_time, label="Average")
legend()