import numpy as np
from minecraft_pb2 import *
from utils_novelty import *
from utils_deap_tools import *
from utils_block_color 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 *
from evolution_roof import BuildingExecution
from deap import base, creator, tools
from typing import Any

GRADIENT_BLOCKS = (set(ALL_FULL_BLOCKS)
    .difference(CORAL_BLOCKS+LEAVES) # CORAL lose color on air - LEAVES color is not corect
    .difference(DEEPSLATE_ORES+STONE_ORES) # ORES have distracting texture
    .difference(GLAZED_TERRACOTTAS) # GLAZED TERRACOTTAS have distractive patters
)

ROOF_BLOCKS = PLANKS

# Initialization

def initialize_gradient(blocks:list[BlockType], color_reps:list, gradient_funcs:Callable, steps_range=(3,10)):
    return (
        random.choice(blocks),
        random.choice(blocks),
        random.randint(*steps_range),
        random.choice(color_reps),
        random.choice(gradient_funcs),
        GeneralUtils.random_3Dvector()
    )

class GradientExecution(Execution):
    wall_block=STONE
    roof_block=OAK_PLANKS

    def __init__(self, wm:WorkspaceManager, block_colors:dict[Any, BlockColors], building:list[RoofArguments], offset=1) -> None:
        super().__init__(wm)
        self.block_colors = block_colors
        self.set_building(building)
        
        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 calculate_heights(coordinates:list[Coordinates], direction_vector):
        heights = []
        
        for coor in coordinates:
            dot_product = sum([v1*v2 for v1,v2 in zip(coor.to_tuple(), direction_vector)])
            height = dot_product
            heights.append(height)
        
        return heights

    @staticmethod
    def filter_blocks(blocks:list[(Coordinates,BlockType)], blocks_to_select:set):
        selected = []
        rest = []
        for coor, block in blocks:
            if block in blocks_to_select:
                selected.append((coor, block))
            else:
                rest.append((coor, block))
        return selected, rest

    @staticmethod
    def blocks_to_request(blocks:list[(Coordinates,BlockType)], origin_coor:Coordinates):
        request = []
        for coor, block in blocks:
            request.append(ServerClient.create_block(origin_coor + coor, block))
        return request
    
    def set_building(self, building:list[RoofArguments]):
        blocks = BuildingExecution.create_building(
            building, Coordinates.zero(),
            roof_block=GradientExecution.roof_block,
            wall_block=GradientExecution.wall_block
        )
        
        blocks_walls, blocks_rest = GradientExecution.filter_blocks(
            blocks, set([GradientExecution.wall_block])
        )
        
        blocks_roofs, blocks_rest = GradientExecution.filter_blocks(
            blocks_rest, set([GradientExecution.roof_block])
        )
        
        self.blocks_walls = blocks_walls
        self.blocks_roofs =blocks_roofs
        self.blocks_rest = blocks_rest
    
    def apply_gradient(self, gradient:Gradient, blocks_to_substitude:list[(Coordinates,BlockType)]):
        block_color:BlockColors = self.block_colors[gradient.representation]
        start_color = block_color.get_block_color(gradient.start_block)
        end_color = block_color.get_block_color(gradient.end_block)
        block_gradient = getattr(block_color, gradient.gradient_func)(start_color, end_color, gradient.steps)
        block_gradient = np.unique(block_gradient)
        
        heights = GradientExecution.calculate_heights([coor for coor, _ in blocks_to_substitude], gradient.direction)
        floor, ceiling = GeneralUtils.get_range(heights)

        gradient_length = len(block_gradient) - 1
        blocks = []
        for (coor, block), height in zip(blocks_to_substitude, heights):
            gradient_value = gradient_length*(height - floor)/(ceiling - floor)
            prob, index = math.modf(gradient_value)
            index = int(index)
            if random.random()>prob:
                block = block_gradient[index]
            else:
                block = block_gradient[min(index+1,gradient_length)]
            blocks.append((coor, block))
        
        return blocks, block_gradient 
    
    def color_contrastly(self, gradient:Gradient, gradient_blocks, blocks_to_substitude:list[(Coordinates,BlockType)]):
        block_color:BlockColors = self.block_colors[gradient.representation]
        coef = (gradient.direction[1]+1)/2
        
        # g_block = gradient.start_block if y>=0 else gradient.end_block
        # sorted_blocks = sorted(ROOF_BLOCKS, key= lambda block : block_color.get_block_distance(block, g_block), reverse=True)
        # best_block = sorted_blocks[int(abs(y)*len(sorted_blocks))]
        
        sorted_blocks = sorted(ROOF_BLOCKS,
            key = lambda block : 
                coef*block_color.get_block_distance(block, gradient.start_block)+
                (1-coef)*block_color.get_block_distance(block, gradient.end_block),
        )
        best_block = sorted_blocks[0]
        
        blocks = []
        if best_block is not None:
            for coor, _ in blocks_to_substitude:
                blocks.append((coor, best_block))
        else:
            blocks = blocks_to_substitude
        
        return blocks
   
    def create_gradient_description(self, gradient:Gradient):
        description = []
        description.append(f"from= {BlockType.Name(gradient.start_block)}\nto= {BlockType.Name(gradient.end_block)}")
        description.append(f"representation= {gradient.representation.__name__}\n")
        description.append(f"direction= {', '.join([f'{v:.2f}' for v in gradient.direction])}\nsteps= {gradient.steps}\n")
        description.append(f"function= {gradient.gradient_func.replace('_',' ')}\n")
        return description
   
    def simulate(self, pop, no_placing=False):
        blocks = []
        gradients = []
        self.wm.server_client.place_fill_requests(self.clean_up_request)
        
        for gradient, coor in zip(pop, self.wm.get_coordinates_list()):
            blocks_applied_gradient, gradient_blocks = self.apply_gradient(gradient, self.blocks_walls)
            blocks_applied_contrast = self.color_contrastly(gradient, gradient_blocks, self.blocks_roofs)
            blocks += GradientExecution.blocks_to_request(blocks_applied_gradient + blocks_applied_contrast + self.blocks_rest, coor)
            gradients.append(gradient_blocks)
        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_gradient_description(gradient) for gradient in pop],
                return_clean_request=True
            )
            self.wm.server_client.place_signs(request)
            self.ask_debug()
            self.wm.server_client.place_blocks_from_list(clean_up)

        return (gradients,)
    
    def showcase(self, buildings, gradients):
        blocks = []
        self.wm.server_client.place_fill_requests(self.clean_up_request)
        
        for building, gradient, coor in zip(buildings, gradients, self.wm.get_coordinates_list()):
            self.set_building(building)
            blocks_applied_gradient, gradient_blocks = self.apply_gradient(gradient, self.blocks_walls)
            blocks_applied_contrast = self.color_contrastly(gradient, gradient_blocks, self.blocks_roofs)
            blocks += GradientExecution.blocks_to_request(blocks_applied_gradient + blocks_applied_contrast + self.blocks_rest, coor)
            
        self.wm.server_client.place_blocks_from_list(blocks)
        
        if(self.debug):
            request, clean_up = self.wm.create_sign_describtion_request(
                [self.create_gradient_description(gradient) for gradient in gradients],
                return_clean_request=True
            )
            self.wm.server_client.place_signs(request)
            self.ask_debug()
            self.wm.server_client.place_blocks_from_list(clean_up)
            
def gradient_quality(gradient:Gradient, block_gradient:list[BlockType], block_colors:dict[Any, BlockColors], enlargement_coef= 0.5):
    #REM: compute in representation or in CIELAB?
    
    if len(block_gradient)<2:
        return (-float("inf"),)
    
    result = 0
    bc = block_colors[gradient.representation]
    direct_dist = bc.get_block_distance(block_gradient[0],block_gradient[-1])
    distances = np.array([bc.get_block_distance(b_from,b_to) for b_from,b_to in zip(block_gradient,block_gradient[1:])])

    full_dist = np.sum(distances)
    distances /=full_dist
    mean = np.mean(distances)
    std_dev = np.std(distances)
    enalrgment = abs((full_dist/direct_dist)-1)
    
    result -= max(enalrgment-enlargement_coef,0) # how much larger is the gradient to direct
    result -= std_dev/mean
    result -= abs(gradient.steps-len(block_gradient))/(gradient.steps-2)
    if bc.get_mean_distance()<mean:
        result -= math.log(mean+1, bc.get_mean_distance()+1)
    
    return (result,)

#TODO: itertools.pairwise
def gradient_novelty(block_gradient:list[BlockType], block_gradient_other:list[BlockType], bc:BlockColors):
    novelty = 0

    # Each with each - supports "optimal gradient that is far from its self"    
    # num = 0 
    # for block, other_block in itertools.product(block_gradient, block_gradient_other):
    #     novelty += bc.get_block_distance(block, other_block)
    #     num +=1
    
    
    len_bg, len_bgo = len(block_gradient),len(block_gradient_other)
    
    dist = (
        bc.get_block_distance(block_gradient[0],block_gradient_other[0]) +
        bc.get_block_distance(block_gradient[-1],block_gradient_other[-1])
    )
    r_dist = (
        bc.get_block_distance(block_gradient[0],block_gradient_other[-1]) +
        bc.get_block_distance(block_gradient[-1],block_gradient_other[0])
    )
    
    use_reversed = r_dist<dist

    num = max(len_bg, len_bgo)
    if num==1:
        return dist
    
    for i in range(num):
        i_ratio = i/(num-1)
        block = block_gradient[round(i_ratio*(len_bg-1))]
        if not use_reversed:
            block_other = block_gradient_other[round(i_ratio*(len_bgo-1))]
        else:
            block_other = block_gradient_other[len_bgo - round(i_ratio*(len_bgo-1)) - 1]
        novelty += bc.get_block_distance(block, block_other)
        
    return novelty/num

# Operators

def random_color_jitter(rgb_color, max_jitter=50):
    new_color = []
    for val in rgb_color.get_value_tuple():
        new_val = val + random.randint(-max_jitter, max_jitter)
        new_val = GeneralUtils.clip(new_val,0,255)
        new_color.append(new_val)
    return type(rgb_color)(tuple(new_color)) 

def change_block(ind, new_ind, p):
    if random.random() < p:
        if random.randint(0,1)==0:
            ind.start_block = new_ind.start_block   
        else:
            ind.end_block = new_ind.end_block

def change_steps(ind, new_ind, p):
    if random.random() < p:
        ind.steps = random.randint(*GeneralUtils.get_range([ind.steps,new_ind.steps]))

def change_representation(ind, new_ind, p):
    if random.random() < p:
        ind.representation = new_ind.representation

def change_gradient_func(ind, new_ind, p):
    if random.random() < p:
        ind.gradient_func = new_ind.gradient_func
        
def change_direction(ind, new_ind, p, ratio =None):    
    # direction attribute update
    if random.random() < p:
        ind.direction =  combine_direction(ind.direction, new_ind.direction, ratio = ratio)
    
def change_gradient(ind, toolbox):
    new_ind = toolbox.individual()
    toolbox.mutate_block(ind, new_ind)
    toolbox.mutate_steps(ind, new_ind)
    toolbox.mutate_representation(ind, new_ind)
    toolbox.mutate_gradient_func(ind, new_ind)
    toolbox.mutate_direction(ind, new_ind)
    return ind,

def combine_direction(direction1, direction2, ratio =None):
    if ratio is None:
        ratio = random.random()
    ratio = GeneralUtils.clip(ratio,0,1)
    
    new_direction = tuple(ratio*v1+(1-ratio)*v2 for v1,v2 in zip(direction1,direction2))
    norm = GeneralUtils.get_l_norm(2,*new_direction)
    return tuple(v/norm for v in new_direction)

def combine_gradients(ind1:Gradient, ind2:Gradient):
    ind1.start_block,ind2.start_block = ind2.start_block,ind1.start_block
    
    steps_range = GeneralUtils.get_range([ind1.steps,ind2.steps])
    ind1.steps = random.randint(*steps_range)
    ind2.steps = random.randint(*steps_range)
    
    ind1.direction,ind2.direction = (
        combine_direction(ind1.direction, ind2.direction),
        combine_direction(ind1.direction, ind2.direction),
    )
    
    return ind1,ind2

# Evolution

def setup_block_color_dict(block_color:BlockColors):
    block_colors = {}
    for metrics in [
        (sRGBColor, ColorMetrics.rgb_distance, ColorMetrics.get_intermediate_color),
        (HSVColor, ColorMetrics.hsv_distance, ColorMetrics.get_intermediate_color),
        (LabColor, ColorMetrics.lab_distance, ColorMetrics.get_intermediate_color),
    ]:
        block_color.set_color_metrics(*metrics)
        block_colors[metrics[0]] = deepcopy(block_color)
    if LabColor in block_colors:
        block_colors[LabColor].set_approximation_metric(delta_e_cie1976)
    
    return block_colors

def gradient_evolution(work_space_manager:WorkspaceManager, building:list[Structure], ngen=100, gradient_blocks=GRADIENT_BLOCKS, mode = EvolutionMode.RELEASE, **kwargs):
    # Args

    k = kwargs.get("novelty_k",5)
    threshold = kwargs.get("threshold",15)
    cxpb = kwargs.get("cxpb",0.5)
    mutpb = kwargs.get("mutpb",0.2)
    verbose = kwargs.get("verbose",True)
    archive_size = kwargs.get("archive_size",1000)
    hof_size = kwargs.get("hof_size",100)
    
    m_block_p = kwargs.get("mutate_block_p", 0.5)
    m_steps_p = kwargs.get("mutate_steps_p", 0.2)
    m_rep_p = kwargs.get("mutate_representation_p", 0.2)
    m_grad_func_p = kwargs.get("mutate_gradient_func_p", 0.2)
    m_dir_p = kwargs.get("mutate_direction_p", 0.3)
    tournsize = kwargs.get("tournsize", 2)
        
    archive_type = kwargs.get("archive_type", "bn")
    clustering = kwargs.get("clustering", "agglomerative")
    diversity_weight = kwargs.get("diversity_weight", 0.5)
    diversity_weight = GeneralUtils.clip(diversity_weight,0,1)
    debug = mode == EvolutionMode.DEBUG
    
    #Setup
    toolbox = base.Toolbox()
    pop_size = work_space_manager.get_workspace_size()
    block_color = BlockColors.load_from_default_file()
    block_color.set_block_pallete(gradient_blocks)
    block_colors = setup_block_color_dict(block_color)

    gradient_funcs = [
        "generate_gradient",
        "generate_recurcive_gradient",
        "generate_gradient_via_graph",
    ]
    
    gradient_execution = GradientExecution(work_space_manager, block_colors, building)
    gradient_execution.set_debug(debug)
    
    metric = lambda ind1, ind2:gradient_novelty(ind1[1],ind2[1], block_colors[LabColor])
    get_fitness = lambda ind: ind[0].objective
    
    if archive_type == "lc":
        archive = LocalCompetitionArchive(archive_size, metric, get_fitness, k=k, threshold=threshold)
    elif archive_type == "bn":
        archive = BestNeighbourArchive(archive_size, metric, get_fitness, k=k, threshold=threshold)
    
    # Create tools
    creator.create("GradientFitnessMax", base.Fitness, weights=(1.0,))
    creator.create("GradientObjectiveMax", base.Fitness, weights=(1.0,))
    creator.create("GradientIndividual", Gradient,
        fitness=creator.GradientFitnessMax,
        objective=creator.GradientObjectiveMax,
    )

    toolbox.register("gradient", initialize_gradient,
        list(block_color.get_gradient_set()), list(block_colors.keys()), gradient_funcs)
    toolbox.register("individual", lambda: creator.GradientIndividual(*toolbox.gradient()))
    toolbox.register("population", tools.initRepeat, list, toolbox.individual)

    toolbox.register("simulate", gradient_execution.simulate, no_placing = mode == EvolutionMode.EXPERIMENT)
    toolbox.register("fitness", lambda g, f: (archive.novelty((g,f)),))
    toolbox.register("objective", gradient_quality, block_colors=block_colors)
    
    toolbox.register("mate", combine_gradients)
    toolbox.register("mutate_block", change_block, p = m_block_p)
    toolbox.register("mutate_steps",change_steps, p = m_steps_p)
    toolbox.register("mutate_representation",change_representation, p = m_rep_p)
    toolbox.register("mutate_gradient_func",change_gradient_func, p = m_grad_func_p) 
    toolbox.register("mutate_direction",change_direction, p = m_dir_p)
    toolbox.register("mutate", change_gradient, 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_obj = Statistics(key=lambda ind: tuple(np.nan if np.isinf(value) else value for value in ind.objective.values))
    GeneralStatistics.register_general_stats(stats_obj)
    stats_archive = Statistics()
    stats_archive.register("size", lambda pop: archive.num_archived_items())
    stats = MultiStatistics(fitness=stats_fit, objective=stats_obj, archive=stats_archive)
    
    # Call Evolution Algorithm
    pop = toolbox.population(n=pop_size)
    hof = HallOfFameGeneral(hof_size, similar=eq, fit_attr="objective") if debug else None
    pop, logbook = Algorithms.ea_general(pop, toolbox, Algorithms.evaluate_novelty, cxpb, mutpb, ngen, stats, halloffame=hof, verbose=verbose)
    
    archive_graph = archive.get_graph( lambda ind: ind[0])
    archive_pop = sorted(
        nx.get_node_attributes(archive_graph, 'data').values(),
        key= lambda ind: ind.objective, reverse=True
    )
    compare = lambda ind, other_ind: ind.objective>=other_ind.objective
    
    match mode:
        case EvolutionMode.DEBUG:
            gradient_execution.set_debug(debug)
            work_space_manager.server_client.send_message("Last generation showcase")
            toolbox.simulate(pop)
            work_space_manager.server_client.send_message("Hall of fame showcase")
            toolbox.simulate(hof[:pop_size])
            work_space_manager.server_client.send_message("Best archive showcase")
            toolbox.simulate(archive_pop[:pop_size])
            work_space_manager.server_client.send_message("Diverse best archive showcase")
            reps, labels = get_representatives(archive_graph, pop_size, 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:
            archive_subgraph = archive_subgraph = get_diversity_subgraph(archive_pop, archive_graph, pop_size, compare, diversity_weight)
            reps = get_representatives(archive_subgraph, pop_size, compare, clust_alg = clustering)
            selected_idx = gradient_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)
    )
    
    creator.create("Building", list, level_height=5)
    building = creator.Building(
        [
            Structure(
                Shape(Shape.Shapes.RECTANGLE, 2, 2, 12, 8),
                1,
                ShapeUtils.FuncMode.SUM,
                [True,True,True],
                RoofArguments(),
                RoofArguments(ShapeUtils.FuncMode.MAX,1,1),
                RoofArguments(),
            ),
            Structure(
                Shape(Shape.Shapes.CIRCLE, 1, 1, 4, 4),
                2,
                ShapeUtils.FuncMode.SUM,
                [True,True,True],
                RoofArguments(),
                RoofArguments(),
                RoofArguments(ShapeUtils.FuncMode.MIN,5,5),
            ),
        ]
    )
    
    gradient_evolution(wm, building, 100, mode = EvolutionMode.DEBUG)
    wm.delete_workspace_grid()