import copy
import heapq
import re
import sys
import time
from math import isinf, inf
import networkx as nx
import matplotlib.pyplot as plt
from typing import Dict, Generator, List, Optional
import random


def main():
    # print("Welcome to Lab1")
    print("功能5中输入@停止功能")
    tu = None

    while True:

        i = input("请输入操作（q:退出, 0:读入文件, 1:展示图形, 2:查找桥接词, 3:生成新文本, 4:计算最短路径, 5:随机游走）：")
        if i == "q":
            sys.exit()
        if i == "0":
            tu = Tu()

            file_name = input("请输入文件名称：")

            tu.generate_directed_dict(tu.read_text_file(file_name))
            tu.generate_directed_graph()
            continue
        if tu is None:
            print("未找到有效的图")

            continue

        if i == "1":
            tu.drawDirectedGraph()
        elif i == "2":
            word1 = input("Enter the word1:")
            word2 = input("Enter the word2:")

            tu.queryBridgeWords(word1, word2)
        elif i == "3":
            text = input("Enter the text:")
            tu.generateNewText(text)

        elif i == "4":
            word1 = input("Enter the word1:")
            word2 = input("Enter the word2,just '' is all")
            tu.calcShortestPath(word1, word2)

        elif i == "5":
            tu.randomWalk()

        else:
            print("Invalid input")


class Tu:
    def __init__(self):
        self.dict = {}
        self.graph = nx.MultiDiGraph()

    @staticmethod
    def read_text_file(filename):
        try:
            with open(filename, 'r', encoding="utf-8") as file:
                text = file.read()
                # 用正则表达式将非字母字符替换为空格，并将换行符也替换为空格
                text = re.sub(r'[^a-zA-Z\n]+', ' ', text)
                # 将文本转换为小写，并按空格分割成单词列表
                word_sequence = text.lower().split()
                print(word_sequence)
                return word_sequence
        except FileNotFoundError:
            print("File not found.")
            return []

    def generate_directed_dict(self, word_sequence):
        graph = self.dict
        for i in range(len(word_sequence) - 1):
            current_word = word_sequence[i]
            next_word = word_sequence[i + 1]
            # 忽略空字符串
            if current_word == '' or next_word == '':
                continue
            # 将当前单词和下一个单词添加到图中
            if current_word not in graph:
                graph[current_word] = {}
            if next_word not in graph[current_word]:
                graph[current_word][next_word] = 1
            else:
                graph[current_word][next_word] += 1

        # draw_directed_graph(graph)
        return graph

    def generate_directed_graph(self):
        G = self.graph  # wufuzhi,yiqigai
        graph = self.dict
        # 添加节点和边
        for node, neighbors in graph.items():
            for neighbor, weight in neighbors.items():
                G.add_edge(node, neighbor, weight=weight)
                # 如果存在相反方向的边，则再画一个箭头
                # if neighbor in graph and node in graph[neighbor]:
                # G.add_edge(neighbor, node, weight=graph[neighbor][node])
        return G

    def draw_directed_graph(self):
        G = self.graph
        pos = nx.spring_layout(G)
        plt.rcParams['figure.figsize'] = (12.8, 7.2)
        nx.draw_networkx(G, pos, with_labels=True, node_size=1000,
                         node_color="skyblue", edge_color="black", font_size=10,
                         font_weight="bold", arrows=True, connectionstyle='arc3,rad=0.2')
        edge_labels = nx.get_edge_attributes(G, 'weight')
        # 调整箭头位置
        for edge, weight in edge_labels.items():
            if G.has_edge(*edge[::-1]):  # 检查反向边是否存在于图中
                dx = pos[edge[0]][0] - pos[edge[1]][0]
                dy = pos[edge[0]][1] - pos[edge[1]][1]
                plt.annotate(weight, (
                    (pos[edge[0]][0] + pos[edge[1]][0]) / 2 + dy * 0.1,
                    (pos[edge[0]][1] + pos[edge[1]][1]) / 2 - dx * 0.1))
            else:  # 反向边不存在
                plt.annotate(weight, ((pos[edge[0]][0] + pos[edge[1]][0]) / 2,
                                      (pos[edge[0]][1] + pos[edge[1]][1]) / 2))
        plt.show()

    def drawDirectedGraph(self):
        self.draw_directed_graph()

    def find_bridge_words(self, word1, word2):

        graph = self.dict
        if word1 not in graph or word2 not in graph:
            print("No", word1, "or", word2, "in the graph!")
            return []

        bridge_words = [
            word for word in graph[word1]
            if word2 in graph.get(word, {})
        ]

        return bridge_words

    @staticmethod
    def print_bridge_words(bridge_words, word1, word2):

        if not bridge_words:
            print("No bridge words from", word1, "to", word2, "!")
            return []
        else:
            bridge_words_str = ", ".join(bridge_words)
            print("The bridge words from", word1, "to", word2, "are:", bridge_words_str + ".")

            return bridge_words

    def queryBridgeWords(self, word1, word2):
        word1 = self.input_check(word1)
        word2 = self.input_check(word2)
        bridge_words = self.find_bridge_words(word1, word2)
        if bridge_words:
            self.print_bridge_words(bridge_words, word1, word2)
            return bridge_words
        else:
            return bridge_words

    def insert_bridge_words(self, text):

        words = text.split()

        new_text = []

        for i in range(len(words) - 1):

            word1 = words[i].lower()
            word2 = words[i + 1].lower()

            new_text.append(words[i])
            bridge_words = self.find_bridge_words(word1, word2)

            if len(bridge_words) > 0:
                random_bridge_word = random.choice(bridge_words)
                new_text.append(random_bridge_word)

        new_text.append(words[-1])

        return ' '.join(new_text)

    def generateNewText(self, text):

        (text_re, text_ren) = re.subn(r'[^a-zA-Z\n]+', ' ', text)
        if text_ren > 0:
            print("There are illegal signs in your input")
        if not text_re.islower():
            print("There are upper characters in your input")
        print(self.insert_bridge_words(text_re))

    def find_shortest_path_old(self, word1, word2):

        graph = self.graph
        try:
            shortest_path_generator = (
                nx.all_shortest_paths(graph, source=word1, target=word2, weight='weight'))
            shortest_length = (
                nx.shortest_path_length(graph, source=word1, target=word2, weight='weight'))
            return shortest_path_generator, shortest_length
        except nx.NetworkXNoPath:
            return None, None

    def find_shortest_path(self, word1, word2):

        try:

            (shortest_path_list, shortest_length) = self.calc_shortest_path_len(word1, word2)
            return shortest_path_list, shortest_length
        except nx.NetworkXNoPath:
            return None, None

    def find_shortest_path_ex(self, word1, word2):

        try:
            shortest_path_generator = self.all_simple_paths_graph(word1, word2)
            shortest_length = self.calc_shortest_path_len(word1, word2)

            shortest_path_list = []

            for shortest_path in shortest_path_generator:

                if shortest_path.__len__() == shortest_length:
                    shortest_path_list.append(shortest_path)
            return shortest_path_list, shortest_length
        except nx.NetworkXNoPath:
            return None, None

    def draw_shortest_path(self, shortest_path_list):
        graph = self.graph
        plt.figure(figsize=(10, 6))

        # !!! HUATU ROLL YANSE

        pos = nx.spring_layout(graph)
        nx.draw_networkx_nodes(graph, pos, node_size=1000, node_color="skyblue")
        nx.draw_networkx_labels(graph, pos, font_size=10, font_weight="bold")

        for shortest_path in shortest_path_list:
            colorslist = '0123456789ABCDEF'
            num = "#"
            for i in range(6):
                num += random.choice(colorslist)

            for edge in graph.edges():

                # print(shortest_path)

                if (edge in zip(shortest_path[:-1], shortest_path[1:]) or
                        edge in zip(shortest_path[1:], shortest_path[:-1])):

                    nx.draw_networkx_edges(graph, pos, edgelist=[edge], width=2.0, edge_color=num,
                                           arrows=True, arrowstyle="->", arrowsize=30,
                                           connectionstyle='arc3,rad=0.2')
                else:
                    nx.draw_networkx_edges(graph, pos, edgelist=[edge], width=1.0,
                                           edge_color="black",
                                           arrows=True, arrowstyle="->", arrowsize=30,
                                           connectionstyle='arc3,rad=0.2')
        plt.show()
        time.sleep(1)

    def calcShortestPath(self, word1, word2):
        if word1 not in self.graph or word2 not in self.graph:
            print("No", word1, "or", word2, "in the graph!")
            return [],0

        if word2 == "":
            word1 = self.input_check(word1)
            for wordtmp in self.dict:
                if wordtmp == word1:
                    continue
                else:
                    result=self.calc_shortest_path(word1, wordtmp)
                    if result == []:
                        print("No path exists between", word1, "and", word2)
                    else:
                        print(result)

            pass
        else:
            word1 = self.input_check(word1)
            word2 = self.input_check(word2)
            result = self.calc_shortest_path(word1, word2)
            pass
        if result==[] :
            print("No path exists between", word1, "and", word2)
        else:
            print(result)
        result_len = 0 if result==[] or result is None else len(result[0])-1
        return result,result_len

    def calc_shortest_path_old(self, word1, word2):

        shortest_path_generator, shortest_length = self.find_shortest_path_old(word1, word2)
        shortest_path_list = []

        for shortest_path in shortest_path_generator:

            if shortest_path:
                print("Shortest path:", '→'.join(shortest_path))
                print("Length of shortest path:", shortest_length)
                shortest_path_list.append(shortest_path)
        if shortest_path_list.__len__() > 0:
            self.draw_shortest_path(shortest_path_list)
            return shortest_path_list
        else:
            return []
            #print("No path exists between", word1, "and", word2)

    def calc_shortest_path(self, word1, word2):

        shortest_path_list, shortest_length = self.find_shortest_path_ex(word1, word2)

        for shortest_path in shortest_path_list:

            if shortest_path:
                print("Shortest path:", '→'.join(shortest_path))
                print("Length of shortest path:", shortest_length)

        if shortest_path_list.__len__() > 0:
            self.draw_shortest_path(shortest_path_list)
            return shortest_path_list
        else:
            return []

            #print("No path exists between", word1, "and", word2)

    def random_traversal(self):

        # 用户输入文件名
        filename = input("Enter the filename to save traversal results: ")
        filename += ".txt"

        graph = copy.deepcopy(self.dict)
        start_node = random.choice(list(graph.keys()))
        visited_nodes = set()
        visited_edges = []
        current_node = start_node

        with open(filename, 'w') as file:
            while True:
                visited_nodes.add(current_node)
                file.write(current_node + '\n')
                print(current_node)
                if input() == "@":
                    file.close()
                    break
                if graph.get(current_node) is not None:
                    pass
                else:
                    break
                neighbors = list(graph[current_node].keys())
                if not neighbors:
                    break
                next_node = random.choice(neighbors)
                visited_edges.append((current_node, next_node))
                graph[current_node].pop(next_node)
                current_node = next_node

        file.close()
        print("Visited nodes:", visited_nodes)
        print("Visited edges:", visited_edges)
        return visited_nodes, visited_edges

    def randomWalk(self):
        self.random_traversal()

    @staticmethod
    def input_check(input_word):
        (input_word_re, input_word_ren) = re.subn(r'[^a-zA-Z\n]', ' ', input_word)

        if input_word_ren > 0:
            print("There are illegal signs in your input,the amount is " + str(input_word_ren))

        if not input_word_re.islower():
            print("There are upper characters in your input")
        tmp = input_word_re.lower().split()
        if len(tmp) > 1:
            print("There are more than one word in your input")

        return tmp[0]

    def all_simple_paths_graph(self, source: str, targets: str) -> Generator[List[str], None, None]:
        G = self.graph
        cutoff = len(G) - 1  # 设置路径的最大深度，防止无限循环。
        visited = dict.fromkeys([source])  # 使用字典跟踪访问过的节点
        stack = [iter(G[source])]  # 使用栈保存当前路径中的节点迭代器
        while stack:
            children = stack[-1]  # 获取当前栈顶节点的迭代器
            child = next(children, None)

            if child is None:
                stack.pop()
                visited.popitem()  # 移除最近访问的节点
            elif len(visited) < cutoff:
                if child in visited:
                    continue
                if child == targets:
                    yield list(visited) + [child]
                visited[child] = None
                if {targets} - set(visited.keys()):  # 扩展栈直到找到目标
                    stack.append(iter(G[child]))
                else:
                    visited.popitem()  # 可能有到达子节点的其他路径
            else:  # 达到最大深度:
                for target in ({targets} & (set(children) | {child})) - set(visited.keys()):
                    yield list(visited) + [target]
                stack.pop()
                visited.popitem()

    def calc_shortest_path_len(self, word1: str, word2: str) -> int:

        if word1 not in self.graph.nodes or word2 not in self.graph.nodes:
            return 0

        distances = {node: inf for node in self.graph.nodes}
        # 存储word1到所有结点的距离，初始化为无穷大
        previous_nodes: Dict[str, Optional[str]] = {node: None for node in
                                                    self.graph.nodes}  # 存储每个节点最短路径中的前一个节点，初始化为None
        distances[word1] = 0  # 距离初始化为0
        priority_queue = [(0, word1)]  # 优先级队列，用于按照从小到大获取节点

        while priority_queue:
            current_distance, current_node = heapq.heappop(priority_queue)  # 优先弹出距离最小的节点

            if current_node == word2:
                break

            if current_distance > distances[current_node]:
                continue

            for neighbor, attributes in self.graph[current_node].items():  # 遍历所有邻居节点及属性
                weight = attributes.get('weight', 1)
                distance = current_distance + weight  # 更新距离，当前距离加入权重

                if distance < distances[neighbor]:
                    distances[neighbor] = distance
                    previous_nodes[neighbor] = current_node
                    heapq.heappush(priority_queue, (distance, neighbor))

        path = []
        current_node = word2
        while prev := previous_nodes[current_node]:  # 从目标节点回溯直到起始节点

            path.insert(0, current_node)
            current_node = prev
        if path:
            path.insert(0, current_node)

        if isinf(distances[word2]):  # 目标节点不可达
            return 0
        return path.__len__()
        # return (' '.join(path),path.__len__())


if __name__ == '__main__':
    main()
