import time
import random
import numpy as np
from minecraft_pb2 import *
from utils_coordinates import *
from utils_server_client import ServerClient
import itertools as it

DEFAULT_POINT = Coordinates(69,0,69)
TARGET_BLOCK = REDSTONE_LAMP
SC = ServerClient()

def redstone_line_experiment_setup(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]
    SC.place_blocks_from_loc(start_coor, experiment_space_shape, AIR)
    SC.place_blocks_from_loc(start_coor, Coordinates.move_down(experiment_space_shape,height-1), STONE)
    SC.place_blocks_from_list([ServerClient.create_block(coor, REDSTONE_BLOCK) for coor in starting_coordinates])
    SC.place_blocks_from_list([ServerClient.create_block(coor- Coordinates(0,0,1), REPEATER) for coor in goal_coordinates])
    SC.place_blocks_from_list([ServerClient.create_block(coor, TARGET_BLOCK) for coor in goal_coordinates])
    
    return (experiment_space_shape,
            unit_space_shape,
            starting_coordinates,
            goal_coordinates)

def execute( coordinate , array):
    SC.place_block_types_from_np_array(coordinate, array)
    
def fitness(goal_coordinates , array):
    fit = 0
    goal_block = SC.read_block(goal_coordinates)
    b = (goal_block.type == TARGET_BLOCK) and (goal_block.is_powered)
    fit += 100 if b else 0 
    fit += (-1 if b else 1)*np.count_nonzero(array==REDSTONE_WIRE)
    return fit
       
def tournament_selection(pop, fits, k):
    childrens =[]
    for _ in range(k):
        parent_1_id = random.randint(0, len(pop)-1)    
        parent_2_id = random.randint(0, len(pop)-1)
        result = random.random()<0.8
        if fits[parent_1_id]>fits[parent_2_id]:
            if result:
                childrens.append(pop[parent_1_id])
            else:
                childrens.append(pop[parent_2_id])
        
        else:
            if result:
                childrens.append(pop[parent_2_id])
            else:
                childrens.append(pop[parent_1_id])
    return childrens

def mutation(pop, mutate, mut_prob):
    return [mutate(p) if random.random() < mut_prob else p[:] for p in pop]

def flip_mutate(p):
    replace_rate = 0.1
    mask = np.random.choice([0, 1], size=p.shape, p=((1 - replace_rate), replace_rate)).astype(bool)
    return np.where(mask, p, random.choice([AIR,REDSTONE_WIRE]))

def crossover(pop, cross, cx_prob):
    off = []
    random.shuffle(pop)
    for p1, p2 in zip(pop[0::2], pop[1::2]):
        if random.random() < cx_prob:
            o1, o2 = cross(p1, p2)
        else:
            o1, o2 = p1[:], p2[:]
        off.append(o1)
        off.append(o2)
    return off

def random_cross(p1, p2):
    mask = np.random.choice([0, 1], size=p1.shape).astype(bool)
    return (np.where(mask, p1, p2),np.where(mask, p2, p1))

def evolution(generations=100, pop_number=20,  map_fn=map):
    length = 10
    experiment_space_shape,unit_space_shape,starting_coordinates,goal_coordinates = redstone_line_experiment_setup(DEFAULT_POINT, pop_number,length=length)
    population = [np.random.choice([AIR,REDSTONE_WIRE], size=(unit_space_shape - Coordinates(0,0,3)).to_tuple(), p=[0.8,0.2]) for i in range(pop_number)]
    coordinates = [c + Coordinates(0,0,1) for c in starting_coordinates]
    
    s_time = time.time()

    for G in range(generations):
        SC.place_blocks_from_list([ServerClient.create_block(coor, AIR) for coor in starting_coordinates])
        list(map_fn(execute,coordinates,population))
        SC.place_blocks_from_list([ServerClient.create_block(coor, REDSTONE_BLOCK) for coor in starting_coordinates])
        time.sleep(0.1)
        fits = list(map_fn(fitness,goal_coordinates,population))
        print(fits,  round(time.time() - s_time,2))
        population = tournament_selection(population, fits, pop_number)
        population = crossover(population, random_cross, 0.8)
        population = mutation(population, flip_mutate, 0.2)
        
    input()
    SC.place_blocks_from_loc(DEFAULT_POINT, experiment_space_shape, AIR)
    return population
    
if __name__ == "__main__":
    evolution() 