from matplotlib.colors import Normalize
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.collections import LineCollection
from mpl_toolkits.axes_grid1 import make_axes_locatable
import warnings
import os

from math import floor
from scipy.spatial import Delaunay
import scipy.stats as st 
from uuid import uuid4
from sklearn.cluster import KMeans

if os.path.basename(os.getcwd()) == 'sar_moe8':
    os.chdir('../')

WORK_DIR = '//math.aau.dk/ProjectGroups/f21moe8or'
MAP_DIR = os.path.join(WORK_DIR, 'maps')

class Map:
    """ 
    Generate a Map object by providing a path to the folder containing the score map.

    NOTE KGN: Maybe some parameters should be fixed when calling the class?
        e.g. min_score, type, nghbr_lvl
        Removed the methods solve_cplex, plot_cplex_solutions, _plot_all_nodes, solve_greedy and the function _dijkstra.
        If needed the previous file can be found in Old.

    Parameters
    ----------
    path_to_folder : str
        Path to the folder generated by MapGenerator.save.
    """
    def __init__(self, path_to_folder, base_node=0):
        locations = {'map': None, 'targets': None}
        for file in os.listdir(path_to_folder):
            if 'map' in file:
                locations['map'] = os.path.join(path_to_folder,file)
            elif 'targets' in file:
                locations['targets'] = os.path.join(path_to_folder,file)
            elif 'difficulty' in file:
                self.difficulty = int(file.split('_')[1])
            else:
                # warnings.warn(f"Unknown file located in folder {path_to_folder} that are neither a map or targets file")
                continue

        self.map = pd.read_csv(locations['map'],sep=' ',header=None)
        self.targets = pd.read_csv(locations['targets'],sep=' ',header=None)
        self.id = os.path.basename(path_to_folder)

        max_x = max(self.map.columns.values)
        max_y = max(self.map.index.values)
        start_x = [0,floor(max_x/2),max_x]
        start_y = [0,floor(max_y/2),max_y]

        self.start = []
        cnt = 0
        iter = 0
        for i in start_x:
            for j in start_y:
                if iter == 4:
                    iter += 1
                    continue
                self.start.append(np.array([i,j], dtype='int64'))
                iter += 1
                cnt += 1
        self.base_node = base_node
        self.base_coord = self.start[base_node]

    def plot(self, size = (7,7), type='score', show=True, add_targets=False):
        '''
        Plot a heatmap of the map instance, optionally with target locations marked.

        Parameters
        ----------
        type : {'score', 'targets'}
            Generate heatmap of either score or targets
        
        show : bool
            if figure should be plotted. If show=False, the function only returns the fig object
        
        add_targets : bool
            If targets should be superimposed on the heatmap. Only useful with type='score'.
        
        Returns
        -------
        Matplotlib figure object.
        '''
        fig, ax = plt.subplots(1,1,figsize=size) 
        if type == 'score':
            im = ax.imshow(self.map, cmap='plasma', zorder=0)
            divider = make_axes_locatable(ax)
            cax = divider.append_axes("right", size = "5%", pad = 0.05)
            fig.colorbar(im, cax=cax)
            # X = self.map.columns.values
            # Y = self.map.index.values
            # Z = self.map.values
            # x,y = np.meshgrid(X,Y)
            # ax.contourf(x,y,Z,cmap='plasma',norm=Normalize(vmin=0,vmax=3))
            # ax.invert_yaxis()

        elif type == 'targets':
            ax.imshow(self.targets, cmap='plasma', zorder=0)
        else:
            raise NotImplementedError(f'The plot type {plot_type} is not recognized')

        if add_targets:
            x, y = [], []
            for i in range(len(self.targets)):
                for j in range(len(self.targets[0])):
                    if self.targets[i][j] == 1:
                        x.append(i)
                        y.append(j)
            # print(list(zip(x,y)))
              
            # find vertices that include targets and plot these in red color
            ax.scatter(x, y, color='lime', marker='o', linewidths=1, zorder=1)

        # Superimpose targets on heatmap
        # points = np.array(list(self.nodes().values())[0:-2])

        # # find vertices that include targets and plot these in red color
        # target_marker = ['.' if self.targets.iloc[point[0], point[1]] == 0 else 'x' for point in points if point is not None]
        # target_ids = [i for i in range(len(points) - 2) if target_marker[i] == 'x']
        # ax.scatter(points[target_ids].T[0], points[target_ids].T[1],color='black',marker='o', linewidths=1, zorder=1)

        return fig

    def nodes(self, type='score', min_score=1, add_start=True, add_end=True):
        '''
        Generate the nodes of the network in a dictionary based on score/target and min value

        Parameters
        ----------
        add_start, add_end : bool
            indicates if artifical start and/or end should be added to the dict of nodes.

        Returns
        -------
        A dictionary with node_id as key and coordinates as value.
        '''
        if type == 'score':
            loc_x, loc_y = np.where(self.map >= min_score)
        elif type == 'targets':
            loc_x, loc_y = np.where(self.targets >= min_score)
        else:
            raise NotImplementedError(f'The type {type} is not recognized')
        
        points = np.array((loc_x, loc_y)).T
        ids = range(0,len(points))
        node_dict = dict(zip(ids, points))
        if add_start:
            node_dict['start'] = None
        if add_end:
            node_dict['end'] = None

        # add possible facility locations
        node_dict['base'] = self.base_coord
        # cnt = 0
        # iter = 0
        # for i in self.start_x:
        #     for j in self.start_y:
        #         if iter == 4:
        #             iter += 1
        #             continue
        #         node_dict['base' + str(cnt)] = np.array([i,j], dtype='int64')
        #         iter += 1
        #         cnt += 1
        
        return node_dict
    
    def scores(self, type='score', min_score=1, add_start=True, add_end=True):
        '''
        Generate a dictionary of scores with nodes_id

        Parameters
        ----------
        add_start, add_end : bool
            True if start and end nodes are instantiated, in that case we add score 0 to these.

        Returns
        -------
        A dictionary with node_id as key and score as value.
        '''
        nodes = self.nodes(type=type, min_score=min_score, add_start=False, add_end=False)
        node_ids = list(nodes.keys())
        scores = dict()
        for node_id in node_ids:
            scores[node_id] = self.map.iloc[nodes[node_id][0], nodes[node_id][1]]
        
        if add_start:
            scores['start'] = 0
        if add_end:
            scores['end'] = 0

        # add possible facility locations
        scores['base'] = 0
        # cnt = 0
        # iter = 0
        # for i in self.start_x:
        #     for j in self.start_y:
        #         if iter == 4:
        #             iter += 1
        #             continue
        #         scores['base' + str(cnt)] = 0
        #         iter += 1
        #         cnt += 1
        return scores

    def A(self, type = 'score', min_score=1, nghbr_lvl=0):
        '''
        Create incidence matrix of flows based on the delaunay triangulation

        Parameters
        ----------
        nghbr_lvl : int 
            The number of neighbors that should be included from the Delaunay triangulation.
            If 0, no neighbors are included, if 1 only the first neighbors are included etc.
        
        base_node : int
            Indicates from which base the UAV should start and end, can be between 0 and 7.

        Returns
        -------
        A dictionary with arc_id as key and (tail, head) as value.
        '''
        @memoize
        def recursive_nghbr(trngl_parent, current_lvl):
            nghbr_ids = [nghbr_id for nghbr_id in trngls.neighbors[trngl_parent] if nghbr_id != -1]
            children_nghbrs = set(nghbr_ids)
            if current_lvl > 1:
                for trngl_child in nghbr_ids:
                    trngl_child_nghbrs = recursive_nghbr(trngl_child, current_lvl - 1)
                    children_nghbrs.update(trngl_child_nghbrs)
            return children_nghbrs
        # points = list(self.nodes(type=type, min_score=min_score).values())[:-2]
        nodes_dict = self.nodes(type=type, min_score=min_score, add_start = False, add_end = False)
        points = list(nodes_dict.values())
        index_max = len(list(nodes_dict.keys())[:-1]) - 1

        trngls = Delaunay(points)
        arcs = set()
        for trngl_id, trngl in enumerate(trngls.simplices):
            arcs.update({(i,j) for i in trngl for j in trngl if i!=j})
        
            if nghbr_lvl > 0:
                current_lvl = nghbr_lvl
                connected_trngls = set(recursive_nghbr(trngl_id, current_lvl))
                arcs.update({(i,j)[::order] for i in trngl for nghbr_id in connected_trngls for j in trngls.simplices[nghbr_id] if (i!=j) for order in [1,-1]})

        ids = range(0, len(arcs))
        A = dict(zip(ids, arcs))

        # replace delaunay indexes with base indexes
        for arc in A:
            if (A[arc][0] > index_max):
                first = 'base'
                second = A[arc][1]
                A[arc] = (first, second)
            if (A[arc][1] > index_max):
                first = A[arc][0]
                second = 'base'
                A[arc] = (first, second)

        
        # Make all nodes connected to start/end
        new_id = len(A)
        A_vals_end = [item for item in A.values() if item[1] == 'base']
        for node_id in nodes_dict:
            if node_id == 'base':
                continue
            elif (node_id, 'base') not in A_vals_end:
                A[new_id] = (node_id, 'base')
                new_id += 1
            
            A['end'] = ('base', 'end')
        
        # if base_node != None:
        A['start'] = ('start', 'base')

        # if end_node != None:
        #     if end_node == -1:
        #         end_node = len(points) - 1
        #     # Make all nodes connected to A
        #     new_id = len(A)
        #     A_vals_end = [item for item in A.values() if item[1] == end_node]
        #     for node_id in nodes_dict:
        #         if (node_id, end_node) not in A_vals_end:
        #             A[new_id] = (node_id, end_node)
        #             new_id += 1
            
        #     A['end'] = (end_node, 'end')
        # if start_node != None:
        #     A['start'] = ('start', start_node)
        
        return A
    
    def lengths(self, type='score', min_score=1, nghbr_lvl=0):
        """
        Calculate the length of each arc.

        Parameters
        ----------
        All passed to nodes() and A() methods.

        Returns
        -------
        A dictionary with arc_id as key and length as value.
        """
        nodes = self.nodes(type = type, min_score = min_score, add_start=False, add_end=False)
        A = self.A(type = type, min_score = min_score, nghbr_lvl = nghbr_lvl)
        A.pop('end')
        A.pop('start')
        self.test_arcs = A
        self.test_nodes = nodes
        # A = dict(filter(lambda elem: elem[0] not in ['start', 'end'], A.items()))
        # for arc_id in A:
        #     try:
        #         np.linalg.norm(nodes[A[arc_id][1]]-nodes[A[arc_id][0]])
        #     except:
        #         print(arc_id)
        length_dict = {arc_id: np.linalg.norm(nodes[A[arc_id][1]]-nodes[A[arc_id][0]]) for arc_id in A}
        # if start_node != None:
        length_dict['start'] = 0
        # if end_node != None:
        length_dict['end'] = 0
        return length_dict

    def delta(self, type = 'score', min_score = 1, nghbr_lvl = 0):
        """
        Calculate the delta minus and delta plus arcs for each node

        Parameters
        ----------
        nghbr_lvl : int 
            The number of neighbors that should be included from the Delaunay triangulation.
            If 0, no neighbors are included, if 1 only the first neighbors are included etc.

        
        start_node : {int, None}
            index of the node that should be connected with start_node.
            Set to None if no start_node.
        
        end_node : {int, None}
            index of the node that should be connected with end_node.
            Set to None if no end_node.
            If end_node = -1, it will connect to the last node of the network.

        Returns
        -------
        two dictionaries containing delta minus and delta plus for all nodes
        """
        # TODO: add_start and add_end should probably always be True
        # if start_node != None:
        #     add_start = True
        # else: add_start = False
        # if end_node != None:
        #     add_end = True
        # else: add_end = False

        add_start = True
        add_end = True

        nodes = self.nodes(type = type, min_score = min_score, add_start = add_start, add_end = add_end)
        A = self.A(type = type, min_score = min_score, nghbr_lvl = nghbr_lvl)

        delta_minus = {i: [] for i in nodes.keys()}
        delta_plus = {i: [] for i in nodes.keys()}

        for (arc_id, arc) in A.items():
            tail, head = arc[0], arc[1]
            delta_plus[tail].append(arc_id)
            delta_minus[head].append(arc_id)

        delta_plus_nodes = dict()
        for node in nodes:
            delta_plus_nodes[node] = []
            for arc in delta_plus[node]:
                delta_plus_nodes[node].append(A[arc][1])
        
        delta_minus_nodes = dict()
        for node in nodes:
            delta_minus_nodes[node] = []
            for arc in delta_minus[node]:
                delta_minus_nodes[node].append(A[arc][0])

        return delta_minus, delta_plus, delta_minus_nodes, delta_plus_nodes

    def plot_flows(self, plot_type = 'score', min_score=1, size=(7,7), nghbr_lvl=0, show_nodes=False, plot_base=True):
        '''
        Plots the generated flows from the incidence matrix A.

        Parameters
        ----------
        nghbr_lvl : int 
            The number of neighbors that should be included from the Delaunay triangulation.
            If 0, no neighbors are included, if 1 only the first neighbors are included etc.


        show_nodes : bool
            if True, superimposes position of nodes on the map. 

        Returns
        -------
        Matplotlib figure object.
        '''
        points = list(self.nodes(type=plot_type, min_score=min_score, add_start = False, add_end = False).values())
        points = np.array(points)
        trngls = Delaunay(points)

        fig, ax = plt.subplots(1,1,figsize=size)
        ax.invert_yaxis()
        ax.triplot(points.T[1], points.T[0], trngls.simplices, color='0.9', zorder=1)
        
        if nghbr_lvl >= 0:
            A = self.A(type=plot_type, min_score=min_score, nghbr_lvl = nghbr_lvl)
            A.pop('end')
            A.pop('start')
            index_max = len(points[:-1]) - 1            
            
            blacklist = []
            for arc in A:
                if type(A[arc][1]) == str:
                    if A[arc][1] == 'base':
                        blacklist.append(arc)
                        continue
            
            A_base = {}
            for arc in blacklist:
                if plot_base:
                    A_base[arc] = A[arc]
                del A[arc]

            if plot_base:
                for arc in A_base:
                    # print(arc)
                    # self.A_base = A_base
                    if type(A_base[arc][0]) == str:
                        first = index_max + 1 # int(A_base[arc][0][-1])
                        second = A_base[arc][1]
                        A_base[arc] = (first, second)
                    if type(A_base[arc][1]) == str:
                        first = A_base[arc][0]
                        second = index_max + 1 #int(A_base[arc][1][-1]) +  
                        A_base[arc] = (first, second)

            for arc in A: 
                if type(A[arc][0]) == str:
                    first = index_max + 1 # int(A[arc][0][-1]) + 
                    second = A[arc][1]
                    A[arc] = (first, second)
                if type(A[arc][1]) == str:
                    first = A[arc][0]
                    second = index_max + 1 # int(A[arc][1][-1]) + 
                    A[arc] = (first, second)
                
            A = list(A.values())
            points_swapxy = np.array([point[::-1] for point in points])
            # self.points_swapxy = points_swapxy
            arc_collection = LineCollection(points_swapxy[A], color='0.9', zorder=1)
            ax.add_collection(arc_collection)

            if plot_base:
                A_base = list(A_base.values())
                arc_collection_base = LineCollection(points_swapxy[A_base], color='lightsteelblue', alpha = 0.1, zorder=0)
                ax.add_collection(arc_collection_base)

        if show_nodes:
            ax.scatter(points.T[1], points.T[0], color ='black', marker ='.', zorder=3)
        
        return fig

    def get_centroids(self, type = 'score', min_score = 1, num_vehicles = 2):
        '''
        This algortihm makes k-means clustering on the provided points.
        The purpose is to generate k clusters with centroids and find the 
        a node closest to each centroid. 
        This way we can set the greedy algorithm to fly directly to a centroid
        with each uav and act greedy from here.

        The KMeans function runs 10 times by default. I don't see a reason to tune this.

        Parameters
        ----------
        num_vehicles : int >=1
            The number of vehicles/clusters to generate.

        Returns
        -------
        node_ids of the nodes closest to the centroids.
        '''
        points_dict = self.nodes(type=type, min_score=min_score, add_start = False, add_end = False)
        # points_dict_no_base = points_dict.copy()
        del points_dict['base']

        #n_init is 10 by default.
        kmeans = KMeans(n_clusters = num_vehicles).fit(list(points_dict.values()))

        center_nodes = []
        for k in range(num_vehicles):
            length_dict = {node_id: np.linalg.norm(kmeans.cluster_centers_[k] - points_dict[node_id]) if (kmeans.labels_[node_id] == k) else np.inf for node_id in points_dict}
            # id of node with smallest dist to centroid.
            center_nodes.append(min(length_dict, key=length_dict.get))
        return center_nodes

    def get_properties(self, min_score = 1, num_vehicles = 1):
        '''
        This method calculates the chosen properties of the map instance.
        is_corner : bool
        avg_cent_dist : average distance to centroids
        avg_score_dist : average distance to the highest scoring nodes of the map.
        min_cent_dist : distance to the nearest centroid
        '''
        nodes_dict = self.nodes(type='score', min_score=min_score, add_start = False, add_end = False)
        scores_dict = self.scores(type='score', min_score=min_score, add_start = False, add_end = False)
        max_score = max(scores_dict.values())
        max_score_nodes = [node_id for node_id in nodes_dict if scores_dict[node_id] == max_score]
        # max_score_coord = [nodes_dict[node_id] for node_id in max_score_nodes]
        centroid_nodes = self.get_centroids(type='score', min_score = min_score, num_vehicles = 1)
        # centroid_nodes_coord = [nodes_dict[node_id] for node_id in centroid_nodes]

        max_dists = [np.linalg.norm(nodes_dict['base'] - nodes_dict[item]) for item in max_score_nodes]
        centroid_dists = [np.linalg.norm(nodes_dict['base'] - nodes_dict[item]) for item in centroid_nodes]

        self.avg_score_dist = np.mean(max_dists)
        self.avg_cent_dist = np.mean(centroid_dists)
        self.is_corner = True if self.base_node in [0, 2, 5, 7] else False
        self.min_cent_dist = min(centroid_dists)


#TODO: this class could be a child of the map class in order to inherit the plot method
class MapGenerator:
    """ 
    Generate a score map.

    Parameters
    ----------
    map_dim : tuble of int with len 2
        Specifies the dimensions of the map.

    centers : int
        Specify the number of desired centers in the score map.
    """
    def __init__(self, map_dim=(100,100), centers=round(np.random.uniform(3,6)), difficulty = 1):
        mean = [
            map_dim[0]/2 - np.random.uniform(low=-map_dim[0]/5, high=map_dim[0]/5), 
            map_dim[1]/2 - np.random.uniform(low=-map_dim[1]/5, high=map_dim[1]/5)
        ]
        x_mean = np.random.uniform(low=mean[0], high=mean[1], size=centers)
        y_mean = np.random.uniform(low=mean[0], high=mean[1], size=centers)

        sd = [2, 6]
        x_sd = np.random.uniform(low=sd[0], high=sd[1], size=centers)
        y_sd = np.random.uniform(low=sd[0], high=sd[1], size=centers)

        x = []
        for i in range(centers):
            x.append(np.random.normal(x_mean[i], x_sd[i], 150*centers))
        x = np.concatenate(x)

        y = []
        for i in range(centers):
            y.append(np.random.normal(y_mean[i], y_sd[i], 150*centers))
        y = np.concatenate(y)

        map, xedges, yedges = np.histogram2d(x, y, bins=map_dim)
        data_mean = np.mean(map[map>0])

        targets = np.zeros(map_dim, dtype=int)
        k=0
        for l in map:
            for y in l:
                div = int(k/map_dim[0])
                mod = k%map_dim[0]
                k = k+1
                if (map[div,mod]>data_mean*difficulty and np.random.rand()<.1*difficulty**2.5):
                    targets[div, mod] = 1

        self.number_of_targets = int(np.sum(targets))
        self.map = map
        self.targets = targets
        self.map_dim = map_dim
        self.difficulty = difficulty
    
    def plot(self, plot_type='score', show=True, size=(7,7), add_targets=False):
        """
        Plot the map.

        Parameters
        ----------
        plot_type : ('score', 'targets')
            Whether the score map or the map of targets should be plotted.
        
        show : bool
            Whether the plot should be shown or not.
        
        size : (int, int)
            The size of the plot.

        Returns
        -------
        Matplotlib figure object.
        """
        fig, ax = plt.subplots(1,1,figsize=size)
        
        if plot_type == 'score':
            ax.imshow(self.map, cmap='plasma')
            # X = pd.DataFrame(self.map).columns.values
            # Y = pd.DataFrame(self.map).index.values
            # Z = pd.DataFrame(self.map).values
            # x,y = np.meshgrid(X,Y)
            # ax.contourf(x,y,Z,cmap='plasma',norm=Normalize(vmin=0,vmax=3))
            # ax.invert_yaxis()

        elif plot_type == 'targets':
            ax.imshow(self.targets, cmap='plasma')
        else:
            raise NotImplementedError(f'The plot type {plot_type} is not recognized')

        if add_targets:
            x, y = [], []
            for i in range(len(self.targets)):
                for j in range(len(self.targets[0])):
                    if self.targets[i][j] == 1:
                        y.append(i)
                        x.append(j)
            # print(list(zip(x,y)))
              
            # find vertices that include targets and plot these in red color
            ax.scatter(x, y, color='lime', marker='o', linewidths=1, zorder=1)
        
        return fig

    def save(self):
        """
        Save the given map in a subfolder of 'maps' in the working directory stored in the WORK_DIR variable.
        """
        instance_id = str(uuid4().hex[0:6])
        dest_path = os.path.join(WORK_DIR, 'maps', instance_id)
        try:
            os.mkdir(dest_path)
        except:
            print('Creation of the directory f{dest_path} failed')

        filename = instance_id + '_' + str(self.map_dim[0]) + 'x' + str(self.map_dim[0]) + '.txt'
        try:
            np.savetxt(os.path.join(dest_path, 'map_' + filename), self.map, fmt='%d')
            np.savetxt(os.path.join(dest_path,'targets_' + filename), self.targets, fmt='%d')
            np.savetxt(os.path.join(dest_path,'difficulty_' + str(self.difficulty)), self.targets, fmt='%d')
        except:
            print('Creation of txt files failed')

### TESTING ###
# reps = 500
# seed = 123
# result = []
# difficulties = [1,1.25,1.5,1.75,2,2.25,2.5]#[1,1.1,1.2,1.3,1.4,1.5,1.6,1.7,1.8,1.9,2]
# total_iterations = reps*len(difficulties)
# iteration = 0
# PrintProgressBar(0, total_iterations, prefix = 'Progress:', suffix = 'Complete', length = 25)
# for i in range(reps):
#     np.random.seed(seed)
#     seed += 1    
#     for difficulty in difficulties:
#         inst = MapGenerator(difficulty=difficulty,centers=5)
#         num_targets = sum(sum(i) for i in inst.targets)
#         result.append([difficulty, num_targets])
#         iteration += 1
#         PrintProgressBar(iteration, total_iterations, prefix = 'Progress:', suffix = 'Complete', length = 25)

# df = pd.DataFrame(result, columns=['difficulty', 'num_targets'])
# df.groupby('difficulty').mean()
# df.pivot(columns='difficulty', values='num_targets').plot.hist(bins = 55, stacked=True, cmap='viridis', figsize=(15,6))
# plt.savefig('some_figure.png')

# for base in range(0,8):
#     inst = Map(os.path.join(MAP_DIR, os.listdir(MAP_DIR)[0]), base_node = base)
    # inst.get_properties(num_vehicles=3)
    # print('data for the base number:', base)
    # print('average centroid dist', inst.avg_cent_dist)
    # print('average high-node dist', inst.avg_score_dist)
    # print('base is in a corner: ', inst.is_corner)
    # print()


# path_to_sample = '//math.aau.dk/ProjectGroups/f21moe8or/maps/c46798'
# inst = Map(path_to_sample, base_node = 0)
# inst.plot(add_targets=True)
# inst.plot_flows(plot_base = False, show_nodes=True)
