import sys, time
from PyQt5.QtWidgets import (QMainWindow, QWidget, QGridLayout, QApplication,
                             QPushButton, QStatusBar, QProgressBar, QLabel,
                             QTreeView, QListView, QFileSystemModel, QLineEdit,
                             QInputDialog, QFileDialog, QTextEdit, QMessageBox)
from PyQt5.QtCore import QDir
import os
from PyQt5.QtGui import QPixmap, QFont
from PyQt5.Qt import QSize, QImageReader
import qdarkstyle


class SecondWindow(QWidget):

def __init__(self, msg):
        super(SecondWindow, self).__init__()
        self.resize(400, 400)
        self.move(200, 200)

layout = QGridLayout()
        self.info = QTextEdit()
        layout.addWidget(self.info)
        self.info.setText(msg)
        self.setLayout(layout)
        self.setWindowTitle('详细信息')


class mainwindow(QMainWindow):

def __init__(self):
        super(mainwindow, self).__init__()

self.setWindowTitle("天然草地类型识别系统")

# 显示正在加载
        self.status = QStatusBar()
        self.status.setStyleSheet('QStatusBar::item {border: none;}')
        self.setStatusBar(self.status)

self.progressBar = QProgressBar()
        self.label = QLabel()
        self.label.setText("加载中，请稍后... ")
        self.status.addPermanentWidget(self.label, stretch=2)
        self.status.addPermanentWidget(self.progressBar, stretch=4)
        self.progressBar.setRange(0, 100)
        self.progressBar.setMinimum(0)
        self.progressBar.setMaximum(0)
        self.statusBar().setVisible(False)

self.setMinimumSize(1500, 720)

layout = QGridLayout()
        w = QWidget()
        w.setLayout(layout)
        self.setCentralWidget(w)

# 文件树
        self.treeview = QTreeView()
        self.listview = QListView()
        layout.addWidget(self.treeview, 0, 0, 7, 2)
        layout.addWidget(self.listview, 0, 2, 7, 2)

path = QDir.rootPath()

self.dirModel = QFileSystemModel()
        self.dirModel.setRootPath(QDir.rootPath())
        self.dirModel.setFilter(QDir.NoDotAndDotDot | QDir.AllDirs)

self.fileModel = QFileSystemModel()
        self.fileModel.setFilter(QDir.NoDotAndDotDot | QDir.Files)

self.treeview.setModel(self.dirModel)
        self.listview.setModel(self.fileModel)

self.treeview.setRootIndex(self.dirModel.index(path))
        self.listview.setRootIndex(self.fileModel.index(path))

self.treeview.clicked.connect(self.on_clicked)
        self.listview.clicked.connect(self.run_model)

self.treeview.hideColumn(1)
        self.treeview.hideColumn(2)
        self.treeview.hideColumn(3)

# 显示图片
        self.image_label = QLabel()
        self.image_label.setMinimumSize(800, 700)
        layout.addWidget(self.image_label, 0, 4, 5, 10)

# 加载模型
        btn = QPushButton("加载模型")
        layout.addWidget(btn, 7, 0, 1, 2)
        btn.clicked.connect(self.load_model)

more_btn = QPushButton("详细信息")
        layout.addWidget(more_btn, 7, 14, 1, 1)
        more_btn.clicked.connect(self.show_info)

reco_btn = QPushButton("识别")
        layout.addWidget(reco_btn, 7, 2, 1, 1)
        reco_btn.clicked.connect(self.predict)

clear_btn = QPushButton("清空")
        layout.addWidget(clear_btn, 7, 3, 1, 1)
        clear_btn.clicked.connect(self.clear)

# 预测结果
        self.res = QLineEdit()
        self.res.setReadOnly(True)
        layout.addWidget(self.res, 7, 4, 1, 10)

self.model = None
        self.fname = None

def clear(self):
        self.res.clear()
        self.image_label.setVisible(False)

def show_info(self):
        if self.fname is not None:
            # 调用你的图片信息代码，我这里没对应的 excel
            # 详细信息赋值给 msg 即可
            msg = "详细信息"
            self.child = SecondWindow(msg)
            self.child.show()

# 当选中图片的时候，直接进行预测
    def run_model(self):
        self.image_label.setVisible(True)
        idx = self.listview.currentIndex()
        fname = self.fileModel.filePath(idx)
        img_type = fname.split('.')[-1]
        # 判断下是不是图片
        if img_type in ["png", "jpg"]:
            img = QImageReader(fname)
            scale = self.image_label.width() / img.size().width()
            height = int(img.size().height() * scale)
            img.setScaledSize(QSize(self.image_label.width(), height))
            img = img.read()
            pixmap = QPixmap(img)
            self.image_label.setPixmap(pixmap)
            self.fname = fname

# 测试代码放在这里
    def predict(self):
        if self.fname is not None:
            # y = self.model(x)
            # y = str(y)
            self.res.setText("cls 1: 0.9, cls 2: 0.8, cls 3: 0.9")
        # res = self.model(img)
        # self.res.setText(res)
        else:
            msgBox = QMessageBox()
            msgBox.setIcon(QMessageBox.Warning)
            msgBox.warning(self, "警告", "请选择图片后再预测")

def on_clicked(self, index):
        path = self.dirModel.fileInfo(index).absoluteFilePath()
        self.listview.setRootIndex(self.fileModel.setRootPath(path))

# 加载模型的代码
    def load_model(self):
        self.status.setVisible(True)
        file_filter = 'PKL File (*.pt *.pth *.pkl)'
        response = QFileDialog.getOpenFileName(
            parent=self,
            caption='Select a data file',
            directory=os.getcwd(),
            filter=file_filter,
        )
        pth_file, _ = response
        # torch load model
        # self.model.load_state_dict(torch.load(pth_file))
        self.status.setVisible(False)


if __name__ == '__main__':
    app = QApplication([])
    dark_stylesheet = qdarkstyle.load_stylesheet_pyqt5()
    # 如果想美化就取消注释
    # app.setStyleSheet(dark_stylesheet)
    m = mainwindow()
    m.show()
    sys.exit(app.exec())