import numpy as np
from minecraft_pb2 import *
from utils_novelty import *
from utils_deap_tools import *
from utils_evolution_tools import *
from utils_coordinates import *
from utils_server_client import *
from utils_workspace_manger import *
from deap import base, creator, tools, algorithms
import utils_deap_tools as u_tools

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

def rotiational_equality(ind, other_ind):
    return np.all([np.array_equal(rot_ind,other_ind) for rot_ind in NumpyTools.rotate_all_3D(ind)])

def runs_of_value(array, run_value):
    num_run = 0
    prev_value = None

    for value in array:
        if value != prev_value:
            prev_value = value
            if value == run_value:
                num_run+=1
    return num_run

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 objective_func(ind:np.ndarray, empty_block=AIR):
    result = 0
    for indices, value in np.ndenumerate(ind):
        neighbours = 0
        coor = Coordinates.from_tuple(indices)
        for dir in CARDINAL_DIRECTIONS:
            adj = coor + Coordinates.get_cardinal(dir)
            if (adj>=Coordinates.zero() and 
                adj<Coordinates.from_tuple(np.shape(ind)) and 
                (ind[adj.x][adj.y][adj.z])==value):
                neighbours+=1
        result+=min(neighbours,len(CARDINAL_DIRECTIONS)-neighbours)
    return result

def objective_func_2(ind:np.ndarray, empty_block=AIR):
    result = 0
    for value in NumpyTools.columns_along_all_axes(ind):
        result += (1 if runs_of_value(value,OAK_PLANKS)<2 else 0)
    return result

def building_blocks_evolution_novelty(ngen=100):
    toolbox = base.Toolbox()
    wm = WorkspaceManager(SC)
    eq = rotiational_equality # np.array_equal
    wm.create_grid_workspace(
        DEFAULT_POINT,
        Coordinates(10,1,10),
        Coordinates(4,4,4),
        Coordinates(2,2,2)
    )
    pop_size = wm.get_workspace_size()
    archive = UniqueArchive(
        3*pop_size,
        lambda ind, arch_ind: Metrics.hamming_distance(ind.flatten(), arch_ind.flatten()),
        eq
    )
    
    # Create tools
    creator.create("FitnessMax", base.Fitness, weights=(1.0,))
    creator.create("Individual", np.ndarray, 
        fitness=creator.FitnessMax,
        objective=creator.FitnessMax,
        coordinates=Coordinates.zero(),
    )
    
    toolbox.register("create_individual", NumpyTools.create_individual, BLOCKS, wm.get_shape().to_tuple(), p=[0.8,0.2])
    toolbox.register("individual", lambda ind_c :creator.Individual(ind_c()), toolbox.create_individual)
    toolbox.register("population", tools.initRepeat, list, toolbox.individual)

    toolbox.register("execute", execute, wm=wm)
    toolbox.register("fitness", lambda ind: (archive.novelty(ind),))
    toolbox.register("evaluate", toolbox.map, toolbox.fitness)
    toolbox.register("evaluate_obj", toolbox.map, lambda ind: (objective_func(ind),))
    toolbox.register("mate", NumpyTools.random_cross)
    toolbox.register("mutate", NumpyTools.flip_mutate, p=0.05, options=BLOCKS)
    toolbox.register("select", tools.selTournament, tournsize=2)
    toolbox.register("update_archive", archive.add)

    # Setup stats
    fstats = GeneralStatistics.create_novelty_stats()
    stats_pop = Statistics()
    stats_pop.register("unique", lambda pop: len(np.unique(pop,axis=0)))
    stats = MultiStatistics(performance=fstats, population=stats_pop)
        
    # Call Evolution Algorithm
    pop = toolbox.population(n=pop_size)
    hof = u_tools.HallOfFameGeneral(50, fit_attr="objective", similar=eq)
    pop,logbook = Algorithms.ea_novelty(pop, toolbox, 0.8, 0.1, ngen, stats=stats, halloffame=hof, verbose=True)

    # Final results
    best_individual_fit = tools.selBest(pop, 1)[0]
    best_individual_obj = tools.selBest(pop, 1, fit_attr='objective')[0]
    
    print(np.unique(pop,axis=0))
    print("Best Individual:")
    print(f"By Fitness:\n{best_individual_fit}\n o={best_individual_obj.objective.values[0]}")
    print(f"By Objective:\n{best_individual_obj}\n o={best_individual_obj.objective.values[0]}")
    
    print("Hall Of Fame:")
    for i,best_ind in enumerate(hof):
        print(f"Indilidual {i}:\n{best_ind}\no={best_ind.objective.values[0]}")
    
    # Display HOF
    hof_np =np.array(hof)
    hof_np[hof_np==BLOCKS[1]]=STONE
    
    toolbox.execute(hof_np, None)
    SC.ask_for_response("Send \"/response\" if u want to end showcase")
    wm.delete_workspace_grid()

if __name__ == "__main__":
    building_blocks_evolution_novelty()