import time
import copy
import numpy as np
import networkx as nx
import scipy.stats as sc_st
import scipy.ndimage as sc_nd
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
from deap import base, creator, tools
import utils_deap_tools as u_tools

# Initialization

class LayoutConfig:
    
    def __init__(self, 
            space:Coordinates,
            min_side=0,
            min_offset=0,
            big_size=1/2,
            small_size=1/8,
        ) -> None:
        
        self.space = space
        self.min_side = max(min_side, ShapeUtils.min_shape_side)
        self.min_offset = max(min_offset, self.min_side)
        
        size = space.x*space.z
        self.big_size = big_size*size
        self.small_size = small_size*size

    def side_a_range(self):
        return (self.min_side, self.space.x-self.min_offset)

    def side_b_range(self):
        return (self.min_side, self.space.z-self.min_offset)

    def max_x_coor(self, a):
        return self.space.x-a

    def max_y_coor(self, b):
        return self.space.z-b
    
    def offset_range(self, side):
        return (-side+1, side-1)
    
    def separate_shapes_by_size(self, layout):
        categories = ([],[],[])
        for shape in layout:
            size = shape.area()
            if size > self.big_size:
                categories[0].append(shape)
            elif size > self.small_size:
                categories[1].append(shape)
            else:
                categories[2].append(shape)
        return categories

def initialize_shape(l_config:LayoutConfig):
    a = random.randint( *l_config.side_a_range())
    b = random.randint( *l_config.side_b_range())
    return Shape(
        random.choice(list(Shape.Shapes)),
        random.randint(0, l_config.max_x_coor(a)),
        random.randint(0, l_config.max_y_coor(b)),
        a,b
    )

def shape_to_pixels(shape:Shape):
    match shape.shape:
        case Shape.Shapes.SQUARE : return ShapeUtils.square(shape.get_shape_a())
        case Shape.Shapes.RECTANGLE : return ShapeUtils.rectangle(shape.get_shape_a(), shape.get_shape_b())
        case Shape.Shapes.RHOMBOID : return ShapeUtils.rhomboid(shape.get_shape_a(), shape.get_shape_b(), shape.a_off, shape.b_off)
        case Shape.Shapes.CIRCLE : return ShapeUtils.circle(shape.get_shape_a())
        case Shape.Shapes.ELLIPSE : return ShapeUtils.ellipse(shape.get_shape_a(), shape.get_shape_b())
        case default : raise ValueError("Unsupported shape")

def add_pixels(pixels, block, offset=Coordinates.zero()):
    blocks = []
    for coor in ShapeUtils.pixels_to_coordinates(pixels):
        absolute_coor = offset + coor
        blocks.append(ServerClient.create_block(absolute_coor, block))
    return blocks

# Execution & Fitness

class LayoutExecution(Execution):
    
    def __init__(self, wm:WorkspaceManager, l_config:LayoutConfig,
            fill_offset = Coordinates.zero(),
            perimeter_offset = Coordinates.cardinal_up(),
            debug_offset = Coordinates.cardinal_up()*2) -> None:
        
        super().__init__(wm)
        self.l_config = l_config
        
        self.fill_offset = fill_offset
        self.perimeter_offset = perimeter_offset
        self.debug_offset = debug_offset 
        
        self.shape = Coordinates(self.wm.get_shape().x, 1, self.wm.get_shape().z)
        self.perimiter_read_request = self.wm.create_read_shape_request(self.shape, self.perimeter_offset)
        self.fill_read_request = self.wm.create_read_shape_request(self.shape, self.fill_offset)

    @staticmethod
    def create_layout_description(layout:list[Shape]):
        description = []
        for shape in layout:
            description.append(
                f"s={shape.shape.name}\n"+
                f"x={shape.x},y={shape.y}\n"+
                f"a={shape.a},b={shape.b}\n"+
                f"ao={shape.a_off},bo={shape.b_off}\n"
            )
        return description

    def read_layout(self, read_request):
        result = []
        for blocks in self.wm.server_client.read_cubes(read_request):            
            result.append(
                ServerClient.create_np_array_from_read_blocks(
                    blocks,
                    self.shape,
                    lambda bi: int(bi.type!=AIR)
                )
            )
        return result

    def create_layout(self, layout:list[Shape], origin_coor:Coordinates, colors=CONCRETES, fill=False, debug=False):
        blocks = []
        repeates = len(layout) // len(colors) + 1
        colors_adjusted = (colors*repeates)[:len(layout)]
        for shape, color in zip(layout, colors_adjusted):
            shape_offset = origin_coor + Coordinates(shape.x, 0, shape.y)
            pixels = shape_to_pixels(shape)
            blocks += add_pixels(
                pixels, color,
                shape_offset + self.perimeter_offset
            )
            if fill:
                fill_pixels = ShapeUtils.fill_shape(pixels, ShapeUtils.line_fill_pixels)
                blocks += add_pixels(
                    fill_pixels, color,
                    shape_offset + self.fill_offset,
                )
            if debug:
                blocks += add_pixels(
                    ShapeUtils.rectangle(shape.a, shape.b), color,
                    shape_offset + self.debug_offset
                )
        return blocks

    def simulate(self, pop, fill=True, debug=False):
        blocks = []
        self.wm.clean_workspace()
        for layout, coor in zip(pop, self.wm.get_coordinates_list()):
            blocks += self.create_layout(layout, coor, fill=fill, debug=debug)
        self.wm.server_client.place_blocks_from_list(blocks)
        
        time.sleep(0.1)
        
        # result = np.stack([
        #     self.read_layout(self.fill_read_request),
        #     self.read_layout(self.perimiter_read_request),
        # ], axis=1)
        result = (self.read_layout(self.fill_read_request),self.read_layout(self.perimiter_read_request))
        
        if(self.debug):
            request, clean_up = self.wm.create_sign_describtion_request(
                [self.create_layout_description(layout) for layout 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 result

def shape_rectangle_overlap(shapes:list[Shape]):
    # Calculate rightmost and leftmost x coordinates of overlap
    right_x = min([shape.x + shape.get_shape_a() for shape in shapes])
    left_x = max([shape.x for shape in shapes])
    
    # Calculate topmost and bottommost y coordinates of overlap
    top_y = min([shape.y + shape.get_shape_b() for shape in shapes])
    bottom_y = max([shape.y for shape in shapes])
    
    # Check for overlap
    if left_x >= right_x or bottom_y >= top_y:
        return (0,0)  # No overlap
    
    # Calculate overlap area
    overlap_width = right_x - left_x
    overlap_height = top_y - bottom_y
    return (overlap_width, overlap_height)

def is_continuous_shape(border_pixels):
    def coor_neighbors(coor):
        for neighbor in GeneralUtils.iterate_max_metric_neighbours(coor, 1):
            if neighbor in border_set:
                yield neighbor

    if not border_pixels:
        return False
    border_set = set(border_pixels)
    visited = set()
    GeneralUtils.dfs( border_pixels[0], visited, coor_neighbors)
    return len(visited) == len(border_set)

def is_continuous_rectangle_area(layout, overlap_condition):
    def shape_neighbors(shape):
        for neighbor in layout:
            overlap = shape_rectangle_overlap((shape, neighbor))
            if overlap_condition(overlap):
                yield neighbor

    if not layout:
        return False
    visited = set()
    GeneralUtils.dfs( layout[0], visited, shape_neighbors)
    return len(visited) == len(layout)

def is_continuous_area(array, neighbours=None):
    _,n = sc_nd.label(array, neighbours)
    return n==1

def layout_quality(layout_genotype:list[Shape], layout_fenotype_fill:np.ndarray, layout_fenotype_perim:np.ndarray, l_config:LayoutConfig, **kwargs):
    result = 0
    
    # shape sides
    side_points = 0
    for shape in layout_genotype:
        min_side = min(shape.get_shape_a(), shape.get_shape_b())
        max_side = max(shape.get_shape_a(), shape.get_shape_b())
        if max_side/min_side<=2:
            side_points +=1
    result += side_points/len(layout_genotype)
            
    # shape areas
    b,m,s = l_config.separate_shapes_by_size(layout_genotype)
    if 0<len(b)<=2:
        result += 1
    small_min = (len(layout_genotype)-1)//2
    if len(s)>small_min:
        result += small_min
    
    # overlap small
    for shapes in combinations(s,2):
        overlap = math.prod(shape_rectangle_overlap(shapes))
        area = min([shape.area() for shape in shapes])
        result -= overlap/area
    
    # overlap 
    for shapes in combinations(layout_genotype,3):
        overlap = math.prod(shape_rectangle_overlap(shapes))
        if overlap>0:
            result -= 1
    
    # continuous
    if is_continuous_area(layout_fenotype_fill):
        result += 2
    if is_continuous_area(layout_fenotype_perim, np.ones((3,3,3))):
        result += 2
    if is_continuous_rectangle_area(
        layout_genotype,
        lambda overlap: max(overlap)>=3
    ):
        result += 1

    
    return (result,)

# Operators

def combine_layouts_randomly(layout_1, layout_2):
    layout_genome = layout_1 + layout_2
    l_min = min(len(layout_1), len(layout_2))
    l_max = max(len(layout_1), len(layout_2))
    
    return (
        random.sample(layout_genome, random.randint(l_min, l_max)),
        random.sample(layout_genome, random.randint(l_min, l_max))
    )

def combine_layouts(layout_1, layout_2):
    layout_1[:], layout_2[:] = combine_layouts_randomly(layout_1, layout_2)
    return (layout_1, layout_2)

def combine_layouts_by_size(layout_1, layout_2, l_config:LayoutConfig):
    new_layout_1, new_layout_2 = [], []
    
    layout_sizes_1 = l_config.separate_shapes_by_size(layout_1)
    layout_sizes_2 = l_config.separate_shapes_by_size(layout_2)
    
    big_min = min(len(layout_sizes_1[0]),len(layout_sizes_2[0]))
    
    if big_min>0:
        for layout_size_1, layout_size_2 in zip(layout_sizes_1, layout_sizes_2):
            new_layout_size_1, new_layout_size_2 = combine_layouts_randomly(layout_size_1, layout_size_2)
            new_layout_1 += new_layout_size_1
            new_layout_2 += new_layout_size_2
    else:
        new_layout_1, new_layout_2 = combine_layouts_randomly(layout_1, layout_2)
    
    layout_1[:] = new_layout_1
    layout_2[:] = new_layout_2
    return (layout_1, layout_2)

def change_shape(shape:Shape, p=0.2):
    if(random.random()<p):
        shape.shape = random.choice(list(Shape.Shapes))

def change_shape_xy(shape:Shape, l_config:LayoutConfig, p=0.05):
    if(random.random()<p):
        shape.x += round(random.gauss())
        shape.y += round(random.gauss())
        shape.x = GeneralUtils.clip(shape.x, 0, l_config.max_x_coor(shape.a))
        shape.y = GeneralUtils.clip(shape.y, 0, l_config.max_y_coor(shape.b))

def change_shape_ab(shape:Shape, l_config:LayoutConfig, p=0.1):
    if(random.random()<p):
        new_a = shape.a + round(random.gauss())
        new_b = shape.b + round(random.gauss())
        new_a = GeneralUtils.clip(new_a, *l_config.side_a_range())
        new_b = GeneralUtils.clip(new_b, *l_config.side_b_range())
        shape.x = round(l_config.max_x_coor(new_a)*shape.x/l_config.max_x_coor(shape.a))
        shape.y = round(l_config.max_y_coor(new_b)*shape.y/l_config.max_y_coor(shape.b))
        shape.a,shape.b = new_a,new_b

def change_shape_ab_off(shape:Shape, l_config:LayoutConfig, p=0.1):
    if(random.random()<p):
        direction = -1 if random.random() < sc_st.norm.cdf(shape.a_off, scale=shape.a/2) else 1
        shape.a_off += direction * round(abs(random.gauss()))
        direction = -1 if random.random() < sc_st.norm.cdf(shape.b_off, scale=shape.b/2) else 1
        shape.b_off += direction * round(abs(random.gauss()))
        shape.a_off = GeneralUtils.clip(shape.a_off, *l_config.offset_range(shape.a))
        shape.b_off = GeneralUtils.clip(shape.b_off, *l_config.offset_range(shape.b))

def change_layout_count(layout:list[Shape], toolbox, p=0.2, mu=6, minimum=1):
    if(random.random()<p):
        inc_prob = EvolutionTools.biased_step_increase_prob(len(layout), mu, minimum)
        if(random.random()<inc_prob):
            layout.append(toolbox.shape())
        else:
            layout.remove(random.choice(layout))

def change_layout(layout:list[Shape], toolbox, l_config:LayoutConfig):
    new_layout = copy.deepcopy(layout)
    for shape in new_layout:
        toolbox.mutate_shape(shape)
        toolbox.mutate_xy(shape, l_config)
        toolbox.mutate_ab(shape, l_config)
        toolbox.mutate_ab_off(shape, l_config)
    toolbox.mutate_count(new_layout, toolbox)
    layout[:] = new_layout
    return layout, 

# Evolution

def layout_evolution(work_space_manager:WorkspaceManager, ngen=100, mode = EvolutionMode.RELEASE, **kwargs):
    # Args
    
    k = kwargs.get("novelty_k",5)
    threshold = kwargs.get("threshold", 50)
    cxpb = kwargs.get("cxpb",0.5)
    mutpb = kwargs.get("mutpb",0.5)
    verbose = kwargs.get("verbose",True)
    archive_size = kwargs.get("archive_size",1000)
    hof_size = kwargs.get("hof_size",100)
    
    min_side = kwargs.get("min_side", 4)
    mu = kwargs.get("len_mu", 4)
    m_shape_p = kwargs.get("mutate_shape_p", 0.15)
    m_xy_p = kwargs.get("mutate_xy_p", 0.05)
    m_ab_p = kwargs.get("mutate_ab_p", 0.2)
    m_ab_off_p = kwargs.get("mutate_ab_off_p", 0.1)
    m_count_p = kwargs.get("mutate_count_p", 0.2)
    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()
    layout_config = LayoutConfig(work_space_manager.get_shape(), min_side=min_side)
    layout_execution = LayoutExecution(work_space_manager, layout_config)
    layout_execution.set_debug(debug)
    
    # Archive
    metric = lambda ind1, ind2: Metrics.np_hamming_distance(ind1[1],ind2[1])
    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("LayoutFitnessMax", base.Fitness, weights=(1.0,))
    creator.create("LayoutObjectiveMax", base.Fitness, weights=(1.0,))
    creator.create("LayoutIndividual", list, 
        fitness=creator.LayoutFitnessMax,
        objective=creator.LayoutObjectiveMax,
    )
    
    toolbox.register("shape", initialize_shape, layout_config)
    toolbox.register("individual", tools.initRepeat, creator.LayoutIndividual, toolbox.shape, n=1)
    toolbox.register("population", tools.initRepeat, list, toolbox.individual)

    toolbox.register("simulate", layout_execution.simulate)
    toolbox.register("fitness", lambda g, f0, f1: (archive.novelty((g,f0)),))
    toolbox.register("objective", layout_quality, l_config=layout_config, mu=mu)

    # toolbox.register("mate", combine_layouts)
    toolbox.register("mate", combine_layouts_by_size, l_config=layout_config)
    toolbox.register("mutate_shape", change_shape, p=m_shape_p)
    toolbox.register("mutate_xy", change_shape_xy, p=m_xy_p)
    toolbox.register("mutate_ab", change_shape_ab, p=m_ab_p)
    toolbox.register("mutate_ab_off", change_shape_ab_off, p=m_ab_off_p)
    toolbox.register("mutate_count", change_layout_count, p=m_count_p, mu=mu)
    toolbox.register("mutate", change_layout, toolbox=toolbox, l_config=layout_config)
    toolbox.register("select", tools.selTournament, tournsize=tournsize)
    # toolbox.register("select", tools.selTournament, tournsize=2, fit_attr="objective")

    # Setup stats
    stats_fit = Statistics(key=lambda ind: ind.fitness.values)
    GeneralStatistics.register_general_stats(stats_fit)
    stats_obj = Statistics(key=lambda ind: ind.objective.values)
    GeneralStatistics.register_general_stats(stats_obj)
    stats_pop = Statistics(key=lambda ind: len(ind))
    GeneralStatistics.register_general_stats(stats_pop)
    stats_archive = Statistics()
    stats_archive.register("size", lambda pop: archive.num_archived_items())
    stats = MultiStatistics(fitness=stats_fit, objective=stats_obj, shape_count=stats_pop, archive=stats_archive)
    
    # Call Evolution Algorithm
    pop = toolbox.population(n=pop_size)
    hof = u_tools.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=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
    n_clusters = min(pop_size,archive_graph.number_of_nodes())
    
    match mode:
        case EvolutionMode.DEBUG:
            layout_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, 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:
            archive_subgraph = get_diversity_subgraph(archive_pop, archive_graph, pop_size, compare, diversity_weight)
            reps = get_representatives(archive_subgraph, n_clusters, compare, clust_alg = clustering)
            selected_idx = layout_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,5,16),
        Coordinates(14,16,14),
        create_platform = True,
        margin = Margin(1,1,1,1)
    )
    layout_evolution(wm, 10, mode=EvolutionMode.DEBUG)
    wm.delete_workspace_grid()