from __future__ import print_function
import neat
import math
import numpy as np
from minecraft_pb2 import *
from utils_neat import *
from utils_novelty import *
from utils_deap_tools import *
from utils_coordinates import *
from utils_server_client import *
from utils_evolution_tools import *
from utils_workspace_manger import *
from deap import base, creator, tools, algorithms
from functools import partial

SC = ServerClient()
DEFAULT_POINT = Coordinates(420,0,420)
BLOCKS = [AIR,OAK_PLANKS]
NAME = "example_building_blocks"

def execute(pop, wm:WorkspaceManager):
    blocks = []
    for ind,coor in zip(pop,wm.get_coordinates_list()):
        blocks += ServerClient.create_blocks_from_np_array(coor, ind)
    SC.place_blocks_from_list(blocks)

def fitness(ind):
    # neig = EvolutionTools.count_neighbours(ind, EvolutionTools.cardinal_neighbours)
    # neig[neig>4] =0
    # return np.sum(neig)
    
    result = 0
    for col in NumpyTools.columns_along_all_axes(ind):
        count = 0
        max_len = 0
        for (val,len) in EvolutionTools.run_encoding(col):
            if val == BLOCKS[1]:
                count += 1
                max_len = max(max_len,len)
        if count==1 and max_len!=3:
            result += 1
    return result

def eval_genomes_novelty(genomes, config, archive:Archive, wm:WorkspaceManager, halloffame:HallOfFame=None, obj=False):
    results = []
    spawn = []
    for (genome_id, genome),coor in zip(genomes, wm.get_coordinates_list()):
        net = neat.nn.FeedForwardNetwork.create(genome, config)
        shape = wm.get_shape().to_tuple()
        blocks = creator.Individual(np.zeros(shape, int))
        for place in np.ndindex(shape):
            y,n = net.activate(place)
            blocks[place] = BLOCKS[1] if y>n else BLOCKS[0]
        blocks.fitness = fitness(blocks)
        results.append(blocks)
        genome.fitness = blocks.fitness if obj else archive.novelty(blocks)
        spawn += ServerClient.create_blocks_from_np_array(coor, blocks)

    SC.place_blocks_from_list(spawn)
    archive.add(results)
    if halloffame is not None:
        halloffame.update(results)
    # SC.get_message()

def run(experiment_name, ngen=100):
    # Prepare neat
    folders = SetUp(experiment_name)
    
    config = neat.Config(
        neat.DefaultGenome,
        neat.DefaultReproduction,
        neat.DefaultSpeciesSet,
        neat.DefaultStagnation,
        folders.get_config_file()
    )
    pop = neat.Population(config)
    pop_size = config.pop_size

    # Prepare Evaluation
    sq_pop = math.ceil(pop_size**(1/2))
    wm = WorkspaceManager(SC)
    eq = np.array_equal
    wm.create_grid_workspace(
        DEFAULT_POINT,
        Coordinates(sq_pop,1,sq_pop),
        Coordinates(4,4,4),
        Coordinates(2,2,2)
    )
    archive = UniqueArchive(pop_size*3, lambda ind, arch_ind: Metrics.hamming_distance(ind.flatten(), arch_ind.flatten()), eq)
    hof = HallOfFameGeneral(pop_size, similar=eq)
    eval = partial(eval_genomes_novelty, archive=archive, wm=wm, halloffame=hof)
    eval_obj =  partial(eval, obj=True)
    creator.create("FitnessMax", base.Fitness, weights=(1.0,))
    creator.create("Individual", np.ndarray, fitness=creator.FitnessMax)

    # Statistics
    pop.add_reporter(neat.StdOutReporter(True))
    pop.add_reporter(neat.StatisticsReporter())
    pop.add_reporter(neat.Checkpointer(ngen//10, filename_prefix=folders.get_checkpoint_prefix()))

    # Run evolution
    # winner = pop.run(eval, ngen)
    winner = pop.run(eval_obj, ngen)
    # pop.run(eval, 80)
    # pop.run(eval_obj, 20)
    # pop.run(eval, 80)
    # winner =pop.run(eval_obj, 20)
    
    Visualize.draw_net(config,winner, True, filename=folders.get_graph_prefix())
    
    hof_np =np.array(hof)
    hof_np[hof_np==BLOCKS[1]]=STONE
    execute(hof_np, wm)
    
    SC.get_message()
    wm.delete_workspace_grid()

if __name__ == '__main__':
    run(NAME)
