import sys
import numpy as np
import torch
from PIL import Image, ImageQt
from PyQt5.QtWidgets import QApplication, QMainWindow, QVBoxLayout, QWidget, QPushButton, QLabel
from PyQt5.QtGui import QPainter, QPen, QImage, QPixmap, QColor
from PyQt5.QtCore import Qt, QPoint
from mm import CNN


class PaintBoard(QWidget):
    def __init__(self):
        super().__init__()
        self.image = QImage(224, 224, QImage.Format_RGB32)
        self.image.fill(Qt.black)
        self.last_point = QPoint()
        self.initUI()

    def initUI(self):
        self.setGeometry(0, 0, 224, 224)
        self.setFixedSize(224, 224)
        self.setStyleSheet("background-color: black;")

    def clear(self):
        self.image.fill(Qt.black)
        self.update()

    def getContentAsQImage(self):
        return self.image

    def mousePressEvent(self, event):
        if event.button() == Qt.LeftButton:
            self.last_point = event.pos()

    def mouseMoveEvent(self, event):
        if event.buttons() & Qt.LeftButton:
            painter = QPainter(self.image)
            pen = QPen(Qt.white, 20, Qt.SolidLine, Qt.RoundCap, Qt.RoundJoin)
            painter.setPen(pen)
            painter.drawLine(self.last_point, event.pos())
            self.last_point = event.pos()
            self.update()

    def paintEvent(self, event):
        canvas_painter = QPainter(self)
        canvas_painter.drawImage(self.rect(), self.image, self.image.rect())


class MainWindow(QMainWindow):
    def __init__(self, model):
        super().__init__()
        self.model = model
        self.model.eval()

        self.setWindowTitle("Handwritten Digit Recognition")
        self.setGeometry(100, 100, 300, 400)
        self.paint_board = PaintBoard()
        self.result_label = QLabel("Prediction: ", self)
        self.result_label.setGeometry(50, 250, 200, 50)
        self.result_label.setStyleSheet("font-size: 18px;")

        clear_button = QPushButton("Clear", self)
        clear_button.setGeometry(50, 300, 100, 50)
        clear_button.clicked.connect(self.paint_board.clear)

        predict_button = QPushButton("Predict", self)
        predict_button.setGeometry(150, 300, 100, 50)
        predict_button.clicked.connect(self.predict)

        #创建垂直布局管理器，把组件放到布局里展示出来
        layout = QVBoxLayout()
        layout.addWidget(self.paint_board)
        layout.addWidget(self.result_label)
        layout.addWidget(clear_button)
        layout.addWidget(predict_button)
        #创建一个QWidget对象作为主容器，然后将垂直布局管理器放在这个主容器中
        container = QWidget()
        container.setLayout(layout)
        self.setCentralWidget(container)

    def predict(self):
        try:
            print("Starting prediction...")
            image = self.paint_board.getContentAsQImage()
            pil_image = ImageQt.fromqimage(image)
            pil_image = pil_image.resize((28, 28), Image.ANTIALIAS)
            pil_image = pil_image.convert('L')

            img_array = np.array(pil_image).reshape(1, 1, 28, 28).astype(np.float32) / 255.0
            img_tensor = torch.tensor(img_array).to(device)

            with torch.no_grad():
                output = self.model(img_tensor)
                prediction = torch.argmax(output, dim=1).item()

            self.result_label.setText(f"Prediction: {prediction}")
            print(f"Prediction: {prediction}")
        except Exception as e:
            print(f"Error during prediction: {e}")


if __name__ == "__main__":
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = CNN().to(device)
    model.load_state_dict(torch.load('model.pth', map_location=device))

    app = QApplication(sys.argv)
    main_window = MainWindow(model)
    main_window.show()
    sys.exit(app.exec_())