import numpy as np
from minecraft_pb2 import *
from utils_coordinates import *
from utils_server_client import *

class ToyExampleTools:

    blocks = [AIR,REDSTONE_WIRE]
    
    def __init__(self, server_client:ServerClient) -> None:
        self.sc =server_client

    def redstone_line_experiment_setup(self, start_coor, pop_number, **kwargs):
        height = 2
        length = kwargs.get("length",15)
        width_single =  kwargs.get("width",3)
        width = pop_number*(width_single+1)+1
        unit_space_shape = Coordinates(width_single,height-1,length)
        experiment_space_shape = Coordinates(width,height,length)
        starting_coordinates = [start_coor + Coordinates(i,1,0) for i in range(1+width_single//2,width,1+width_single)]
        goal_offset =Coordinates(width_single-1,0,length-1)
        goal_coordinates = [coor + goal_offset  for coor in starting_coordinates]
        self.sc.place_blocks_from_loc(start_coor, experiment_space_shape, AIR)
        self.sc.place_blocks_from_loc(start_coor, Coordinates.move_down(experiment_space_shape,height-1), STONE)
        self.sc.place_blocks_from_list([ServerClient.create_block(coor, REDSTONE_BLOCK) for coor in starting_coordinates])
        self.sc.place_blocks_from_list([ServerClient.create_block(coor- Coordinates(0,0,1), REPEATER) for coor in goal_coordinates])
        self.sc.place_blocks_from_list([ServerClient.create_block(coor, REDSTONE_LAMP) for coor in goal_coordinates])
        
        return (experiment_space_shape,
                unit_space_shape,
                starting_coordinates,
                goal_coordinates)

    def assign_coordinates(self, c, s_c, g_c):
        def decorator(func):
            def wrapper(*args, **kargs):
                idx = args[1]
                for ind,c_i,s_c_i,g_c_i in [(args[0][i],c[i],s_c[i],g_c[i]) for i in idx]:
                    ind.coordinates =c_i
                    ind.s_coordinates =s_c_i
                    ind.g_coordinates =g_c_i
                pop = func(*args, **kargs)
                return pop
            return wrapper
        return decorator

    def execute_individual(self, ind):
        self.sc.place_block_types_from_np_array(ind.coordinates, ind)

    def execute(self, pop, toolbox):
        self.sc.place_blocks_from_list([ServerClient.create_block(ind.s_coordinates, AIR) for ind in pop])
        list(toolbox.map(self.execute_individual,pop))
        self.sc.place_blocks_from_list([ServerClient.create_block(ind.s_coordinates, REDSTONE_BLOCK) for ind in pop])

    def objective_func(self, ind):
        fit = 0
        goal_block = self.sc.read_block(ind.g_coordinates)
        b = (goal_block.type == REDSTONE_LAMP) and (goal_block.is_powered)
        fit += 100 if b else 0 
        fit += (-1 if b else 1)*np.count_nonzero(ind==REDSTONE_WIRE)
        return fit
