import copy
import numpy as np
import scipy.stats as sc_st
from minecraft_pb2 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 utils_blocks_categories import *
from utils_individuals import Shape, RoofArguments, Structure
from evolution_layout import shape_to_pixels
from deap import base, creator, tools

# Initialization

def initialize_roof(side_x, side_z, coef_mu=0, coef_sigma=1, threshold_h=0, scaling_h=6):
    
    def _get_coef(coef, value):
            factor = scaling_h/value
            threshold = threshold_h/value
            if factor>1 and coef>threshold:
                coef = (coef-threshold)*factor+threshold
            return coef

    return RoofArguments(
        random.choice(list(ShapeUtils.FuncMode)),
        _get_coef(random.gauss(coef_mu,coef_sigma), side_x),
        _get_coef(random.gauss(coef_mu,coef_sigma), side_z),
        random.uniform(-1,1),
        random.uniform(-1,1),
    )

def initialize_structure(shape:Shape, roof_init, mu=2, zero_height_p=0.05):
    mask = [False, False, False]
    mask[random.randint(0,len(mask)-1)] = True
    level = 1
    # if zero_height_p<random.random():
    #     level += sc_st.poisson.rvs(mu-LEVEL_MU_OFFSET)
    # else:
    #     level =0
    a,b = shape.get_shape_a(), shape.get_shape_b()
    return Structure(
        shape,
        level,
        random.choice(list(ShapeUtils.FuncMode)),
        mask,
        roof_init(a,b),
        roof_init(a,b),
        roof_init(a,b),
        1, #random.randint(1,2),
    )

# Execution & Fitness

def inflate_shape(shape:Shape, amount):
    new_shape = copy.copy(shape)
    new_shape.a += amount*2
    new_shape.b += amount*2
    new_shape.x -= amount
    new_shape.y -= amount
    new_shape.a_off += amount
    new_shape.b_off += amount
    return new_shape
    
def inflate_coor_offseting(coor, amount): 
    return coor + Coordinates(-amount,0,-amount)

class BuildingExecution(Execution):

    def __init__(self, wm:WorkspaceManager, offset=1) -> None:
        super().__init__(wm)
        
        offset= Coordinates(offset,0,offset)
        self.clean_up_request = self.wm.create_fill_shape_request(self.wm.get_shape() + offset*2, offset*(-1))
    
    @staticmethod
    def create_building_description(building:list[Structure]):
        description = []
        for s in building:
            mask_str =  str([ "T" if b else "F" for b in s.mask])
            description.append(f"mode={s.mode.name}\nmask=\n{mask_str}\nl={s.levels} o={s.overhang}")
            description.append(str(s.arg_2))
            description.append(str(s.arg_1))
            description.append(str(s.arg_f1))
        return description
    
    @staticmethod
    def create_building(building:list[Structure], origin_coor:Coordinates, return_roof_array=False, **kwargs):
        roof_block = kwargs.get("roof_block", DARK_OAK_PLANKS)
        wall_block = kwargs.get("wall_block", DEEPSLATE_BRICKS)
        
        blocks = []
        roofs = []
        for structure in building:
        
            selected_roof_block = roof_block
            if structure.levels==0:
                selected_roof_block = wall_block
            
            shape_offset = origin_coor + Coordinates(structure.shape.x, 0, structure.shape.y)
            walls = shape_to_pixels(structure.shape)
            
            level_inflation = 1
            level_offset = inflate_coor_offseting(shape_offset, level_inflation)
            levels = ShapeUtils.fill_shape(
                shape_to_pixels(inflate_shape(structure.shape, level_inflation)),
                ShapeUtils.line_fill_pixels
            )
            
            roof_offset = inflate_coor_offseting(shape_offset, structure.overhang)
            roof_base_shape = inflate_shape(structure.shape, structure.overhang)
            mask_roof = ShapeUtils.fill_shape(shape_to_pixels(roof_base_shape))
            mask_layout = np.pad(ShapeUtils.fill_shape(walls), pad_width=((1,1),(1,1)), mode='constant', constant_values=(0))
            mask = np.logical_or(mask_roof, mask_layout)
            # r_a,r_b = mask.shape #bug fix for squares and circles
            roof = ShapeUtils.roof(
                roof_base_shape.get_shape_a(),
                roof_base_shape.get_shape_b(),
                *structure.get_args()
            )
            masked_roof = ShapeUtils.apply_mask(roof, mask)
            
            if return_roof_array:
                roofs.append(masked_roof)
            
            height = structure.levels*building.level_height + 1
            
            # walls
            for pixel in walls:
                index = (
                    pixel[0] + structure.overhang,
                    pixel[1] + structure.overhang,
                )
                for y in range(int(height+masked_roof[index])):
                    absolute_coor = shape_offset + Coordinates(pixel[0],y,pixel[1])
                    blocks.append((absolute_coor, wall_block))
            
            # levels
            for l in range(structure.levels+1):
                for coor in ShapeUtils.pixels_to_coordinates(levels, l*building.level_height):
                    absolute_coor = level_offset + coor
                    blocks.append((absolute_coor, wall_block))
            
            # roof
            roof_offset += Coordinates.cardinal_up()*height
            for coor in ShapeUtils.np_array_to_coordinates(masked_roof):
                absolute_coor = roof_offset + coor
                blocks.append((absolute_coor, selected_roof_block))
        
        results = []
        if return_roof_array:
            results.append(roofs)
        return blocks if not results else tuple([blocks]+results)

    def simulate(self, pop, no_placing=False):
        blocks = []
        roofs = []
        self.wm.server_client.place_fill_requests(self.clean_up_request)

        for building, coor in zip(pop, self.wm.get_coordinates_list()):
            building_blocks, building_roofs = BuildingExecution.create_building(building, coor, return_roof_array=True)
            for coor, block in building_blocks:
                blocks.append(ServerClient.create_block(coor, block))
            roofs.append(building_roofs)
        if not no_placing:
            self.wm.server_client.place_blocks_from_list(blocks)
        
        if(self.debug):
            request, clean_up = self.wm.create_sign_describtion_request(
                [self.create_building_description(building) for building in pop],
                return_clean_request=True
            )
            self.wm.server_client.place_signs(request)
            self.wm.server_client.send_message(f"Placed {len(blocks)} blocks.")
            self.ask_debug()
            self.wm.server_client.place_blocks_from_list(clean_up)
        
        return (roofs,)

def novelty_roof_args(roof_args:RoofArguments, other_roof_args:RoofArguments, mode_coef=1):
    novelty = int(roof_args.mode != other_roof_args.mode)/mode_coef
    novelty += abs(roof_args.x - other_roof_args.x)
    novelty += abs(roof_args.z - other_roof_args.z)
    novelty += abs(roof_args.cx - other_roof_args.cx)
    novelty += abs(roof_args.cz - other_roof_args.cz)
    return novelty
    
def novelty_structures(structure:Structure, other_structure:Structure):
    roof_args_count = len(structure.mask)
    novelty = int(structure.mode != other_structure.mode)
    # the true Roof fenotipe difference
    # no genome-potential novelty
    
    novelty += novelty_roof_args(structure.arg_2, other_structure.arg_2, roof_args_count)
    novelty += novelty_roof_args(structure.arg_1, other_structure.arg_1, roof_args_count)
    novelty += novelty_roof_args(structure.arg_f1, other_structure.arg_f1, roof_args_count)
    novelty += Metrics.hamming_distance(structure.mask, other_structure.mask)/roof_args_count
    # novelty += int(structure.levels != other_structure.levels)
    # novelty += abs(structure.overhang - other_structure.overhang)
    return novelty
    
def novelty_building(building:list[Structure], other_building:list[Structure]):
    novelty = 0
    for structure, other_structure in zip(building, other_building):
        novelty += novelty_structures(structure, other_structure)
    return novelty/len(building)

def novlety_roofs(building_roofs:np.ndarray, other_building_roofs:np.ndarray, use_weights=False):
    # It assumes same shape
    results = []
    weights = []
    for roof, other_roof in zip(building_roofs, other_building_roofs):
        
        distances = np.abs(roof - other_roof)
        
        # valid_indices = distances > 0
        # distances[valid_indices] = np.log2(distances[valid_indices])
        
        # distances = np.abs(
        #     GeneralUtils.log_threshold(roof) - 
        #     GeneralUtils.log_threshold(other_roof)
        # )
        
        weights.append(np.sum(~np.isnan(distances)))
        results.append(np.nanmean(distances))
    if use_weights:
        return sum(val * w for val, w in zip(results, weights))/sum(weights)
    else:
        return np.mean(results)       


# Operators

def combine_roof_args(roof_args_1:RoofArguments, roof_args_2:RoofArguments):
    pass
    # structure_1.arg_2,structure_2.arg_2 = structure_2.arg_2, structure_1.arg_2
    # return (layout_1, layout_2)

def combine_structures(structure_1:Structure, structure_2:Structure, p=0.5):
    if random.random()<p:
        structure_1.arg_2,structure_2.arg_2 = structure_2.arg_2, structure_1.arg_2
    if random.random()<p:
        structure_1.arg_1,structure_2.arg_1 = structure_2.arg_1, structure_1.arg_1
    if random.random()<p:
        structure_1.arg_f1,structure_2.arg_f1 = structure_2.arg_f1, structure_1.arg_f1
    if random.random()<p:
        structure_1.mode,structure_2.mode = structure_2.mode, structure_1.mode

    structure_1.mask, structure_2.mask = (
        [b1 ^ b2 for b1, b2 in zip(structure_1.mask, structure_2.mask)],
        [b1 == b2 for b1, b2 in zip(structure_1.mask, structure_2.mask)]
    )
    
    l_min = min(structure_1.levels, structure_2.levels)
    l_max = max(structure_1.levels, structure_2.levels)
    structure_2.levels = random.randint(l_min, l_max)
    structure_2.levels = random.randint(l_min, l_max)

def combine_buildings(building_1:list[Structure], building_2:list[Structure], p=0.5):
    new_building_1 = copy.deepcopy(building_1)
    new_building_2 = copy.deepcopy(building_2)
    for structure_1, structure_2 in zip(new_building_1, new_building_2):
        combine_structures(structure_1, structure_2, p)
    building_1[:] = new_building_1
    building_2[:] = new_building_2
    return (building_1, building_2)

def combine_buildings_randomly(building_1:list[Structure], building_2:list[Structure], p=0.5):
    return tools.cxUniform(
        copy.deepcopy(building_1),
        copy.deepcopy(building_2),
        indpb=p
    )

def change_roof_args(shape:Shape, roof_args:RoofArguments, toolbox, p=0.20):
    new_roof_args = toolbox.roof(shape.get_shape_a(), shape.get_shape_b())
    if(random.random()<p):
        roof_args.cx = new_roof_args.cx
    if(random.random()<p):
        roof_args.cz = new_roof_args.cz
    if(random.random()<p):
        roof_args.x = new_roof_args.x
    if(random.random()<p):
        roof_args.z = new_roof_args.z
    if(random.random()<p):
        roof_args.mode = new_roof_args.mode

def change_structure_mask(structure:Structure, p=0.33):
    if(random.random()<p):
        mask_idx = random.randint(0,len(structure.mask)-1)
        structure.mask[mask_idx] = not structure.mask[mask_idx]

def change_structure_level(structure:Structure, p=0.1, mu=2, zero_height_p=0.05):
    if(random.random()<p):
        inc_prob = EvolutionTools.biased_step_increase_prob(structure.levels, mu, 0)
        inc_prob *=(1 - zero_height_p)
        
        if(structure.levels==0 or random.random()<inc_prob):
            structure.levels += 1
        else:
            structure.levels -= 1

def change_structure_mode(structure:Structure, p=0.1):
    if(random.random()<p):
        structure.mode = random.choice(list(ShapeUtils.FuncMode))

def change_structure_overhang(structure:Structure, p=0.1):
        if(random.random()<p):
            structure.overhang = ...

def change_building(building:list[Structure], toolbox):
    new_building = copy.deepcopy(building)
    for structure in new_building:
        toolbox.mutate_level(structure),
        toolbox.mutate_mode(structure),
        toolbox.mutate_mask(structure), 
        toolbox.mutate_roof_args(structure.shape, structure.arg_2, toolbox),
        toolbox.mutate_roof_args(structure.shape, structure.arg_1, toolbox),
        toolbox.mutate_roof_args(structure.shape, structure.arg_f1, toolbox),
        # toolbox.mutate_overhang(structure)
    building[:] = new_building
    return building,

# Evolution

def building_evolution(work_space_manager:WorkspaceManager, layout:list[Shape], ngen=100,  mode = EvolutionMode.RELEASE, **kwargs):
    # Args
    
    k = kwargs.get("novelty_k",5)
    threshold = kwargs.get("threshold",1.5)
    cxpb = kwargs.get("cxpb",0.7)
    mutpb = kwargs.get("mutpb",0.2)
    verbose = kwargs.get("verbose",True)
    archive_size = kwargs.get("archive_size",1000)
    
    slope_mu = kwargs.get("slope_mu", 0.5)
    slope_sigma = kwargs.get("slope_sigma", 0.5)
    thr_h = kwargs.get("threshold_height", 0)
    scal_h = kwargs.get("scaling_height", 6)
    mu = kwargs.get("floors_mu", 2)
    zero_p = kwargs.get("no_floors_p", 0.05)
    cx_ind_p = kwargs.get("cx_ind_p", 0.5)
    m_roof_args_p = kwargs.get("mutate_roof_args_p", 0.2)
    m_mask_p= kwargs.get("mutate_mask_p", 0.33)
    m_level_p= kwargs.get("mutate_level_p", 0.1)
    m_mode_p= kwargs.get("mutate_mode_p", 0.1)
    tournsize = kwargs.get("tournsize", 2)
    
    clustering = kwargs.get("clustering", "agglomerative")
    debug = mode == EvolutionMode.DEBUG 
    
    #Setup
    toolbox = base.Toolbox()
    pop_size = work_space_manager.get_workspace_size()
    roof_execution = BuildingExecution(work_space_manager)
    roof_execution.set_debug(debug)
    # metric = lambda ind1, ind2: novelty_building(ind1[0], ind2[0])
    metric = lambda ind1, ind2: novlety_roofs(ind1[1], ind2[1])
    archive = TrestholdArchive(archive_size, metric, k=k, threshold = threshold)

    # Create tools
    creator.create("BuildingFitnessMax", base.Fitness, weights=(1.0,))
    creator.create("BuildingObjectiveMax", base.Fitness, weights=(1.0,))
    creator.create("BuildingIndividual", list, 
        fitness=creator.BuildingFitnessMax,
        # objective=creator.BuildingObjectiveMax,
        level_height=5 #Part of wall evolution
    )

    toolbox.register("roof", initialize_roof,
        coef_mu = slope_mu, coef_sigma = slope_sigma,
        threshold_h=thr_h, scaling_h=scal_h)
    toolbox.register("structure", initialize_structure, roof_init=toolbox.roof, mu=mu, zero_height_p=zero_p)
    toolbox.register("building", toolbox.map, toolbox.structure, layout)
    toolbox.register("individual", lambda: creator.BuildingIndividual(toolbox.building()))
    toolbox.register("population", tools.initRepeat, list, toolbox.individual)

    toolbox.register("simulate", roof_execution.simulate, no_placing = mode == EvolutionMode.EXPERIMENT)
    toolbox.register("fitness", lambda g,f: (archive.novelty((g,f)),))
    # toolbox.register("objective", lambda ind: random.random())

    toolbox.register("mate", combine_buildings_randomly, p=cx_ind_p)
    # toolbox.register("mate", combine_buildings, p=0.5)
    toolbox.register("mutate_roof_args", change_roof_args, p=m_roof_args_p)
    toolbox.register("mutate_mask", change_structure_mask, p=m_mask_p)
    toolbox.register("mutate_level", change_structure_level, p=m_level_p,
        mu=mu, zero_height_p=zero_p)
    toolbox.register("mutate_mode", change_structure_mode, p=m_mode_p)
    # toolbox.register("mutate_overhang", change_structure_overhang, p=0.2)
    toolbox.register("mutate", change_building, toolbox=toolbox)
    toolbox.register("select", tools.selTournament, tournsize=tournsize)

    # Setup stats
    stats_fit = Statistics(key=lambda ind: ind.fitness.values)
    GeneralStatistics.register_general_stats(stats_fit)
    stats_archive = Statistics()
    stats_archive.register("size", lambda pop: archive.num_archived_items())
    stats = MultiStatistics(fitness=stats_fit, archive=stats_archive)
    
    # Call Evolution Algorithm
    pop = toolbox.population(n=pop_size)
    # hof = u_tools.HallOfFameGeneral(hof_size, similar=eq)
    pop, logbook = Algorithms.ea_general(pop, toolbox, Algorithms.evaluate_simple, cxpb, mutpb, ngen, stats=stats, verbose=verbose) # halloffame=hof

    archive_graph = archive.get_graph( lambda ind: ind[0])
    compare = lambda ind, other_ind: True
    n_clusters = min(pop_size,archive_graph.number_of_nodes())
    
    match mode:
        case EvolutionMode.DEBUG:
            roof_execution.set_debug(debug)
            work_space_manager.server_client.send_message("Last generation showcase")
            toolbox.simulate(pop)
            work_space_manager.server_client.send_message("Diverse archive showcase")
            reps, labels = get_representatives(archive_graph, n_clusters, compare, clust_alg = clustering, debug=True)
            toolbox.simulate(reps)
            # response = work_space_manager.server_client.ask_for_response("Showcase clusters? (Y/N)")
            
        case EvolutionMode.EXPERIMENT:
            return pop, archive_graph, logbook
        
        case default:
            reps = get_representatives(archive_graph, n_clusters, compare, clust_alg = clustering)
            selected_idx = roof_execution.select_ind(reps)
            return reps[selected_idx]

if __name__ == "__main__":
    sc = ServerClient()
    wm = WorkspaceManager(sc)
    wm.create_grid_workspace(
        Coordinates(-1024,0,-1024),
        Coordinates(10,1,10),
        Coordinates(16,70,16),
        Coordinates(14,16,14),
        create_platform = True,
        margin = Margin(1,1,1,1)
    )
    layout = [
        Shape(Shape.Shapes.RECTANGLE, 2, 2, 12, 8),
        Shape(Shape.Shapes.CIRCLE, 1, 1, 4, 4)
    ]
    building_evolution(wm, layout, 100, mode = EvolutionMode.DEBUG)
    wm.delete_workspace_grid()