searchUI.py

import sys
from PyQt5.QtWidgets import QApplication, QMainWindow, QWidget, QVBoxLayout, QPushButton, QSpacerItem, QSizePolicy, QHBoxLayout, QLabel, QFileDialog, QInputDialog
from PyQt5.QtGui import QPainter, QPen, QBrush, QColor, QFont
from PyQt5.QtCore import Qt, QPoint, QTimer
from search import city_coordinates, city_data, bfs, dfs, astar, dijkstra, euclidianDistance
from time import sleep




class DrawingWidget(QWidget):
    def __init__(self, width, height):
        super().__init__()
        self.r = 10
        self.setGeometry(0, 0, width, height)
        self.cityColorDict = {}
        self.connectColorDict = {}
        self.cityCooridates = city_coordinates.copy()
        self.cityConn = []
        self.cost = 0
        #initialize the city connection data.
        for k, v in city_data.items():
            for i in v:
                self.cityConn.append((k, i[0]))
        keys = set()
        for i in self.cityConn:
            t = [i[0], i[1]]
            t.sort(key = lambda x: x[0])
            keys.add(tuple(t))
        self.cityConn = list(keys)
                
        # print(self.cityConn)
        # existedV = set(self.cityConn.values())
        # n = {}
        # for k, v in self.cityConn.items():
        #     if k not in existedV:
        #         n[k] = v
                
        # self.cityConn = n

    def getpainter(self, color):
        painter = QPainter(self)
        painter.setRenderHint(QPainter.Antialiasing)  # 抗锯齿
        painter.setPen(QPen(color, 6, Qt.SolidLine))  # 设置画笔颜色和样式
        painter.setBrush(QBrush(color, Qt.SolidPattern))  # 设置画刷颜色和样式
        return painter

    def drawPlainCircle(self, x, y, color):
        p = self.getpainter(color)
        ct = QPoint(x, y)
        p.drawEllipse(ct, self.r, self.r)
        return ct
    
    def connect(self, ct1, ct2, color):
        # if color != QColor('black'):
        #     print(color)
        painter = self.getpainter(color)
        painter.drawLine(ct1, ct2)
        midx = (ct1.x() + ct2.x()) / 2
        midy = (ct1.y() + ct2.y()) / 2
        dist = int(((ct1.x() - ct2.x()) ** 2 + (ct1.y() - ct2.y()) ** 2) ** 0.5)
        self.text(midx, midy, str(dist), fontsize=10)

    def paintCost(self, cost):
        self.text(10, 20, f"Cost: {cost}", fontsize=10)
    
    def paintEvent(self,event):
        self.paintCost(self.cost)
        #find the justified position of city dots
        maxX = 0
        maxY = 0
        minX = 2147483647
        minY = 2147483647
        for i in city_coordinates.keys():
            x = city_coordinates[i][0]
            y = city_coordinates[i][1]
            if (x > maxX):
                maxX = x
            if (y > maxY):
                maxY = y
            if (x < minX):
                minX = x
            if (y < minY):
                minY = y
        scaledInfo = {}
        scaleRate = 0.9
        for i in self.cityCooridates.keys():
            cinfo = ((self.cityCooridates[i][0] - minX + 55) * (self.width() / ((maxX - minX) * 1.2)) * scaleRate,
            (self.cityCooridates[i][1] - minY + 55) * (self.height() / ((maxY - minY) * 1.2)) * scaleRate)
            # print(self.width(), self.height(), maxX, maxY)
            x = cinfo[1]
            y = cinfo[0]
            scaledInfo[i] = cinfo
            #draw the city dots with color(color is in dict)
            if (self.cityColorDict.get(i) == None):
                # auto: gray
                ct = self.drawPlainCircle(int(y), int(x), QColor('gray'))
            else:
                ct = self.drawPlainCircle(int(y), int(x), self.cityColorDict[i])
            self.text(int(y) + 10, int(x) + 15, i, fontsize = 10)
        # draw connect lines
        for cityA, cityB in self.cityConn:
            # print(self.cityConn)
            cinfoA = scaledInfo[cityA]
            cinfoB = scaledInfo[cityB]
            # print(self.connectColorDict)
            if (self.connectColorDict.get((cityA, cityB)) == None):
                # auto: black
                color = QColor('gray')
            else:
                # print(color)
                color = self.connectColorDict[(cityA, cityB)]

            self.connect(QPoint(int(cinfoA[0]), int(cinfoA[1])) 
            ,QPoint(int(cinfoB[0]), int(cinfoB[1])), color)
    
    def updateCityColor(self, city, color):
        self.cityColorDict[city] = color
        self.update()
    
    def clearPainted(self):
        self.cityColorDict = {}
        self.connectColorDict = {}
        self.update()

    def updateConnectColor(self, cityA, cityB, color):
        t = [cityA, cityB]
        t.sort(key=lambda x: x[0])
        t = tuple(t)
        self.connectColorDict [t] = color
        self.update()

    def text(self, x, y, text, fontsize = 16):
        if (type(x) != int):
            x = int(x)
        if (type(y) != int):
            y = int(y)
        painter = self.getpainter(QColor('black'))
        font = QFont("Arial", fontsize, QFont.Bold)
        painter.setFont(font)
        painter.drawText(QPoint(x, y), text)
        return 
        
class MainWindow(QMainWindow):
    
    def __init__(self):
        self.renderInterval = 500
        self.NotOptimal = QColor('yellow')
        self.Optimal = QColor(0,255,0)
        self.fromCity = None
        self.toCity = None
        self.cities = set()
        for ct in city_coordinates.items():
            self.cities.add(ct[0])
        super().__init__()
        self.main_container = QWidget()
        width = 1200
        height = 1200
        self.setGeometry(100,100, width,height)
        self.setWindowTitle("My App")
        subwidget0 = DrawingWidget(int(width), int(height * 0.9))
        subwidget1 = QWidget()
        self.drawable = subwidget0
        
        self.main_container.setLayout(QVBoxLayout())
        self.main_container.layout().addWidget(subwidget0)
        self.main_container.layout().addWidget(subwidget1)
        
        subwidget0.setLayout(QVBoxLayout())
        subwidget0.layout().addStretch(1)
        subwidget1.setLayout(QHBoxLayout())
        btn1 = QPushButton("BFS遍历所有路径求最小路")
        btn2 = QPushButton("DFS遍历所有路径求最小路")
        btn3 = QPushButton("A*求最小路")
        btn4 = QPushButton("Dijkstra(启发函数为0的A*)")
        
        btn1.clicked.connect(self.onBtn1)
        btn2.clicked.connect(self.onBtn2)
        btn3.clicked.connect(self.onBtn3)
        btn4.clicked.connect(self.onBtn4)
        subwidget1.layout().addWidget(btn1)
        subwidget1.layout().addWidget(btn2)
        subwidget1.layout().addWidget(btn3)
        subwidget1.layout().addWidget(btn4)
        
        self.setCentralWidget(self.main_container)
    def setFromToCity(self, handleError = 0):
        if(handleError):
            self.fromCity = self.interactIn("未成功输入，请重新输入出发城市:", title="重新输入")
            self.toCity = self.interactIn("未成功输入，请重新输入到达城市:", title="重新输入")
        else:
            self.fromCity = self.interactIn("请输入出发城市:")
            self.toCity = self.interactIn("请输入到达城市:")
        if(self.fromCity == None or self.toCity == None):
            return False
        if(self.fromCity not in self.cities or self.toCity not in self.cities):
            self.setFromToCity(handleError=1)
        else:
            return True
        
    def handleResult(self, result):
        index = 0
        x = result
        def updateColors():
            nonlocal index
            nonlocal x

            print(f"{index}/{len(x) - 1}")
            if(index > len(x) - 1):
                self.drawable.cost = x[0][1]
                self.drawable.clearPainted()
                x.sort(key=lambda x: x[1])
                print(x)
                methods = x[0]

                route = methods[0]
                cost = methods[1]
                for city in route:
                    self.drawable.updateCityColor(city,self.Optimal)
                for i in range(len(route) - 1):
                    self.drawable.updateConnectColor(route[i], route[i + 1], self.Optimal)
                timer.stop()
                return
            self.drawable.cost = x[index][1]
            methods = x[index]
            index = index + 1
            self.drawable.clearPainted()
            route = methods[0]
            cost = methods[1]
            for city in route:
                self.drawable.updateCityColor(city,self.NotOptimal)
                
            for i in range(len(route) - 1):
                self.drawable.updateConnectColor(route[i], route[i + 1], self.NotOptimal)
        timer = QTimer()
        timer.timeout.connect(updateColors)
        timer.start(self.renderInterval)
        
    def onBtn1(self):
        print("Button 1 clicked")

        # self.drawable.updateCityColor("Arad", QColor(0, 0, 255))
        # self.drawable.updateConnectColor("Arad", "Zerind", QColor('blue'))
        suc = self.setFromToCity()
        if(not suc):
            return
        x = bfs(self.fromCity, self.toCity)
        # print(x)
        self.handleResult(x)
                
        # for i in range(len(x) - 1):
            
    def onBtn2(self):
        print("Button 2 clicked")

        # self.drawable.updateCityColor("Arad", QColor(0, 0, 255))
        # self.drawable.updateConnectColor("Arad", "Zerind", QColor('blue'))
        suc = self.setFromToCity()
        if(not suc):
            return
        x = dfs(self.fromCity, self.toCity)
        self.handleResult(x)

    def onBtn3(self):
        print("Button 3 clicked")

        self.setFromToCity()
        d = dijkstra(self.fromCity, self.toCity)
        suc = self.handleResult(d)
        if(not suc):
            return
        

    def onBtn4(self):
        print("Button 4 clicked")

        self.setFromToCity()
        a = astar(self.fromCity, self.toCity, estimator=euclidianDistance)
        suc = self.handleResult(a)
        if(not suc):
            return
        

    def interactIn(self, prompts, title='输入起止点'):
        text, ok = QInputDialog.getText(self, title, prompts)
        if ok:
            print((f'Hello, {text}!'))
            return text
        else:
            return None
if __name__ == "__main__":
    app = QApplication(sys.argv)
    window = MainWindow()
    window.show()
    sys.exit(app.exec_())