import heapq
import sys
import pandas as pd
import matplotlib.pyplot as plt
import numpy as np


class Graph:

    def __init__(self):
        self.vertices = {}

    def add_vertex(self, name, edges):
        self.vertices[name] = edges

    def shortest_path(self, start, finish):
        distances = {}  # Distance from start to node
        previous = {}  # Previous node in optimal path from source
        nodes = []  # Priority queue of all nodes in Graph

        for vertex in self.vertices:
            if vertex == start:  # Set root node as distance of 0
                distances[vertex] = 0
                heapq.heappush(nodes, [0, vertex])
            else:
                distances[vertex] = sys.maxsize
                heapq.heappush(nodes, [sys.maxsize, vertex])
            previous[vertex] = None

        while nodes:
            smallest = heapq.heappop(nodes)[1]  # Vertex in nodes with smallest distance in distances
            if smallest == finish:  # If the closest node is our target we're done so print the path
                path = []
                while previous[smallest]:  # Traverse through nodes til we reach the root which is 0
                    path.append(smallest)
                    smallest = previous[smallest]
                return path
            if distances[smallest] == sys.maxsize:  # All remaining vertices are inaccessible from source
                break

            for neighbor in self.vertices[smallest]:  # Look at all the nodes that this vertex is attached to
                alt = distances[smallest] + self.vertices[smallest][neighbor]  # Alternative path distance
                if alt < distances[neighbor]:  # If there is a new shortest path update our priority queue (relax)
                    distances[neighbor] = alt
                    previous[neighbor] = smallest
                    for n in nodes:
                        if n[1] == neighbor:
                            n[0] = alt
                            break
                    heapq.heapify(nodes)
        return distances

    def __str__(self):
        return str(self.vertices)


if __name__ == '__main__':
    excel_data = pd.read_excel('rate.xlsx')
    data_matrix = excel_data.to_numpy()[:8, :]
    data_matrix = np.rot90(data_matrix)

    data_matrix[1, 5] = 0.14

    start_list = ['1_3', '1_4', '1_6', '2_4']
    zero_finish = []
    for i in range(data_matrix.shape[0]):
        for j in range(data_matrix.shape[1]):
            if type(data_matrix[i, j]) is not str:
                data_matrix[i, j] = float(data_matrix[i, j])
            else:
                data_matrix[i, j] = np.finfo(np.float32).max
            if data_matrix[i, j] == 0:
                zero_finish.append('%s_%s' % (i, j))
    data_matrix = data_matrix.astype(np.float32)

    g = Graph()
    for i in range(data_matrix.shape[0]):
        for j in range(data_matrix.shape[1]):

            if (i, j) == (1, 5):
                continue

            if i == 0 and j == 0:
                g.add_vertex('%s_%s' % (i, j), {'%s_%s' % (i, j + 1): data_matrix[i, j + 1],
                                                '%s_%s' % (i + 1, j): data_matrix[i + 1, j]})
            elif i == 0 and j == data_matrix.shape[1] - 1:
                g.add_vertex('%s_%s' % (i, j), {'%s_%s' % (i + 1, j): data_matrix[i + 1, j],
                                                '%s_%s' % (i, j - 1): data_matrix[i, j - 1]})
            elif i == data_matrix.shape[0] - 1 and j == 0:
                g.add_vertex('%s_%s' % (i, j), {'%s_%s' % (i - 1, j): data_matrix[i - 1, j],
                                                '%s_%s' % (i, j + 1): data_matrix[i, j + 1]})
            elif i == data_matrix.shape[0] - 1 and j == data_matrix.shape[1] - 1:
                g.add_vertex('%s_%s' % (i, j), {'%s_%s' % (i - 1, j): data_matrix[i - 1, j],
                                                '%s_%s' % (i, j - 1): data_matrix[i, j - 1]})
            elif i == 0:
                g.add_vertex('%s_%s' % (i, j), {'%s_%s' % (i, j + 1): data_matrix[i, j + 1],
                                                '%s_%s' % (i + 1, j): data_matrix[i + 1, j],
                                                '%s_%s' % (i, j - 1): data_matrix[i, j - 1]})
            elif j == 0:
                g.add_vertex('%s_%s' % (i, j), {'%s_%s' % (i - 1, j): data_matrix[i - 1, j],
                                                '%s_%s' % (i, j + 1): data_matrix[i, j + 1],
                                                '%s_%s' % (i + 1, j): data_matrix[i + 1, j]})
            elif i == data_matrix.shape[0] - 1:
                g.add_vertex('%s_%s' % (i, j), {'%s_%s' % (i - 1, j): data_matrix[i - 1, j],
                                                '%s_%s' % (i, j + 1): data_matrix[i, j + 1],
                                                '%s_%s' % (i, j - 1): data_matrix[i, j - 1]})
            elif j == data_matrix.shape[1] - 1:
                g.add_vertex('%s_%s' % (i, j), {'%s_%s' % (i - 1, j): data_matrix[i - 1, j],
                                                '%s_%s' % (i + 1, j): data_matrix[i + 1, j],
                                                '%s_%s' % (i, j - 1): data_matrix[i, j - 1]})
            else:
                g.add_vertex('%s_%s' % (i, j), {'%s_%s' % (i - 1, j): data_matrix[i - 1, j],
                                                '%s_%s' % (i, j + 1): data_matrix[i, j + 1],
                                                '%s_%s' % (i + 1, j): data_matrix[i + 1, j],
                                                '%s_%s' % (i, j - 1): data_matrix[i, j - 1]})

    for node in g.vertices:
        if '1_5' in g.vertices[node].keys():
            del g.vertices[node]['1_5']

    fig, ax = plt.subplots()
    ax.matshow(data_matrix / data_matrix.max(), cmap=plt.cm.Reds)
    for (i, j), z in np.ndenumerate(data_matrix):
        if z == np.finfo(np.float32).max:
            ax.text(j, i, 'source', ha='center', va='center')
        else:
            ax.text(j, i, '{:.2f}'.format(z), ha='center', va='center')
    for i_num, start_loc in enumerate(start_list):

        distance_list = []
        for zero_loc in zero_finish:
            distance = g.shortest_path(start_loc, zero_loc)
            distance.reverse()
            distance_list.append(distance)

        distance = min(distance_list, key=len)

        x, y = start_loc.split('_')
        dist = [[int(x), int(y)]]

        for item in distance:
            x, y = item.split('_')
            dist.append([int(x), int(y)])

        # 绘制向量
        for i in range(len(dist) - 1):
            ax.quiver(dist[i][1], dist[i][0], dist[i + 1][1] - dist[i][1], dist[i + 1][0] - dist[i][0],
                      angles='xy', scale_units='xy', scale=1)
        # ax.set_xlim([-1, 10])
        # ax.set_ylim([-1, 10])
    plt.savefig("figure_night.png", dpi=1200)
    # plt.show()