import os
import pickle
import warnings
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_block_color import BlockColors
from evolution_layout import layout_evolution, LayoutConfig, LayoutExecution
from evolution_roof import building_evolution, BuildingExecution
from evolution_gradient import gradient_evolution, setup_block_color_dict, GradientExecution, GRADIENT_BLOCKS

VERBOSE = False

DATA_FOLDER = os.path.join("py_scripts","experiment_data")
FILE_FORMAT = ".pkl"

LAYOUT_FILE = "layout_experiment"
BUILDING_FILE = "building_experiment"
GRADIENT_FILE = "gradient_experiment"

def pickle_class(obj, filename):
    with open(filename + FILE_FORMAT, 'wb') as f:
        pickle.dump(obj, f)

def load_pickled_class(filename):
    with open(filename + FILE_FORMAT, 'rb') as f:
        return pickle.load(f)

def create_experimental_wm():
    sc = ServerClient()
    wm = WorkspaceManager(sc)
    wm.create_grid_workspace(
        Coordinates(-1024,0,-1024),
        Coordinates(10,1,10),
        Coordinates(16,100,16),
        Coordinates(14,16,14),
        create_platform = True,
        margin = Margin(1,1,1,1)
    )
    return wm

def run_experiment(evolution, repeates, filename = None):
    pops =[]
    archive_graphs = []
    logbooks = []
    start_time = time.time()
    
    warnings.filterwarnings("ignore", category=RuntimeWarning)
    
    for i in range(1,repeates+1):
        print(f"Starting run {i}")
        
        pop, archive_graph, logbook = evolution(i-1)
        pops.append(pop)
        archive_graphs.append(archive_graph)
        logbooks.append(logbook)
        
        elapsed_time = time.time()-start_time
        predicted_time = (elapsed_time/(i))*(repeates)
        print(f"Elapsed time = {elapsed_time:.2f}s, Estimanted time = {predicted_time:.2f}s")

        # safety write
        if filename is not None:
            temp_filename = f"temp_{i}"
            print(f"Savign to file {temp_filename}")
            pickle_class(
                (pops[-1], archive_graphs[-1], logbooks[-1]),
                os.path.join(DATA_FOLDER, temp_filename)
            )
        
    if filename is not None:
        print(f"Savign to file {filename}")
        pickle_class(
            (pops, archive_graphs, logbooks),
            os.path.join(DATA_FOLDER, filename)
        )
    
    return pops, archive_graphs, logbooks

def layout_experiment(wm):
    ngen = 250
    repeates = 10
    
    archive_type = "lc"
    novelty_k = 5
    threshold = 50

    results = run_experiment(
        lambda i: layout_evolution(
            wm, ngen = ngen, mode = EvolutionMode.EXPERIMENT,
            archive_type = archive_type, #"lc"/"bn"
            clustering = "agglomerative", #"spectral"/"agglomerative"
            
            novelty_k = novelty_k,
            threshold = threshold,
            cxpb = 0.5,
            mutpb = 0.5,
            verbose = VERBOSE,
            archive_size = 2500,
            
            min_side = 4,
            len_mu = 4,
            mutate_shape_p = 0.15,
            mutate_xy_p = 0.05,
            mutate_ab_p = 0.2,
            mutate_ab_off_p = 0.1,
            mutate_count_p = 0.2,
            tournsize = 2,
        ),
        repeates,
        f"{LAYOUT_FILE}_arch={archive_type}_k={novelty_k}_t={threshold}"
    )
    
    return results
    
def select_layouts(wm, archive_graphs, diversity_weight=0.5):
    layouts = []

    layout_config = LayoutConfig(wm.get_shape())
    layout_execution = LayoutExecution(wm, layout_config)
    compare = lambda ind, other_ind: ind.objective>other_ind.objective

    for graph in archive_graphs:
        archive_pop = sorted(
            nx.get_node_attributes(graph, 'data').values(),
            key= lambda ind: ind.objective, reverse=True
        )
        archive_subgraph = get_diversity_subgraph(archive_pop, graph, wm.get_workspace_size(), compare, diversity_weight)
        reps = get_representatives(archive_subgraph, wm.get_workspace_size(), compare)
        selected_idx = layout_execution.select_ind(reps)
        layouts.append(reps[selected_idx])

    layout_execution.set_debug(True)
    layout_execution.simulate(layouts ,fill=False)

    return layouts
 
def building_experiment(wm, layouts):
    ngen = 100
    repeates = 10
    
    novelty_k = 5
    threshold = 0.50
    
    results = run_experiment(
        lambda i: building_evolution(
            wm, layout=layouts[i],
            ngen = ngen, mode = EvolutionMode.EXPERIMENT,
            clustering = "agglomerative", #'spectral'
            novelty_k = novelty_k,
            threshold = threshold,
            cxpb = 0.7,
            mutpb = 0.2,
            verbose = VERBOSE,
            archive_size = 1000,

            slope_mu = 0.5,
            slope_sigma = 0.5,
            threshold_height = 0,
            scaling_height = 6,
            floors_mu = 2,
            no_floors_p = 0.05,
            cx_ind_p = 0.5,
            mutate_roof_args_p = 0.2,
            mutate_mask_p = 0.33,
            mutate_level_p = 0.1,
            mutate_mode_p = 0.1,
            tournsize = 2,
        ),
        repeates,
        f"{BUILDING_FILE}_k={novelty_k}_t={threshold}"
    )
    
    return results

def select_building(wm, archive_graphs):
    buildings = []

    building_execution = BuildingExecution(wm)
    compare = lambda ind, other_ind: True

    for graph in archive_graphs:
        reps = get_representatives(graph, wm.get_workspace_size(), compare)
        selected_idx = building_execution.select_ind(reps)
        buildings.append(reps[selected_idx])
        print(selected_idx)

    building_execution.set_debug(True)
    building_execution.simulate(buildings)

    return buildings

def gradient_experiment(wm, buildings):
    
    ngen = 100
    repeates = 10

    archive_type = "lc"
    novelty_k = 5
    threshold = 15

    results = run_experiment(
        lambda i : gradient_evolution(
            wm, building=buildings[i],
            ngen = ngen, mode = EvolutionMode.EXPERIMENT,
            archive_type = archive_type, #"lc"
            clustering = "agglomerative", #'spectral'

            novelty_k = novelty_k,
            threshold = threshold,
            cxpb = 0.5,
            mutpb = 0.2,
            verbose = VERBOSE,
            archive_size = 1000,
            
            mutate_block_p = 0.5,
            mutate_steps_p = 0.2,
            mutate_representation_p = 0.2,
            mutate_gradient_func_p = 0.2,
            mutate_direction_p = 0.3,
            tournsize = 2,
        ),
        repeates,
        f"{GRADIENT_FILE}_arch={archive_type}_k={novelty_k}_t={threshold}"
    )
    
    return results

def select_gradient(wm, archive_graphs, buildings, diversity_weight=0.5):
    gradients = []

    block_color = BlockColors.load_from_default_file()
    block_color.set_block_pallete(GRADIENT_BLOCKS)
    block_colors = setup_block_color_dict(block_color)
    gradient_execution = GradientExecution(wm, block_colors, buildings[0])
    compare = lambda ind, other_ind: ind.objective>other_ind.objective

    for graph,building in zip(archive_graphs, buildings):
        gradient_execution.set_building(building)
        archive_pop = sorted(
            nx.get_node_attributes(graph, 'data').values(),
            key= lambda ind: ind.objective, reverse=True
        )
        archive_subgraph = get_diversity_subgraph(archive_pop, graph, wm.get_workspace_size(), compare, diversity_weight)
        reps = get_representatives(archive_subgraph, wm.get_workspace_size(), compare)
        selected_idx = gradient_execution.select_ind(reps)
        gradients.append(reps[selected_idx])

    gradient_execution.set_debug(True)
    gradient_execution.showcase( buildings, gradients)

    return gradients

if __name__ == "__main__":
    
    wm = create_experimental_wm()
    
    print("\n___ LAYOUT EXPERIMENTS ___")
    _,archive_graphs,_ = layout_experiment(wm)
    layouts = select_layouts(wm, archive_graphs)
    pickle_class(layouts, os.path.join(DATA_FOLDER, "layouts"))
    
    print("\n___ BUILDING EXPERIMENTS ___")
    layouts = load_pickled_class(os.path.join(DATA_FOLDER, "layouts"))
    _,archive_graphs,_ = building_experiment(wm, layouts)
    bluildings = select_building(wm, archive_graphs)
    pickle_class(bluildings, os.path.join(DATA_FOLDER, "buildings"))
    
    print("\n___ GRADIENT EXPERIMENTS ___")
    bluildings = load_pickled_class(os.path.join(DATA_FOLDER, "buildings"))
    _,archive_graphs,_ = gradient_experiment(wm, bluildings)
    gradient = select_gradient(wm, archive_graphs, bluildings)
    pickle_class((bluildings, gradient), os.path.join(DATA_FOLDER, "gradient"))

    wm.delete_workspace_grid()
