import sys
import torch
import os
from torchvision import models
from PyQt5.QtCore import pyqtSlot, QDir, pyqtSignal
from PyQt5.QtWidgets import QWidget, QApplication, QMessageBox,QFileDialog
from PyQt5.QtGui import QPixmap
from gui import Ui_Form
import simplified_train
import json
from prediction import Prediction
import simplified_export
import threading
from torchvision import models


class ClassificationGui(QWidget):

    def __init__(self, parent=None):
        super().__init__(parent)

        self.ui = Ui_Form()
        self.ui.setupUi(self)
        self.ui.rb_pretrained.toggled.connect(lambda: self.on_rb_clicked(self.ui.rb_pretrained))
        self.ui.rb_retrained.toggled.connect(lambda: self.on_rb_clicked(self.ui.rb_retrained))
        self.ui.rb_retrained.setChecked(True)
        self.ui.cb_model_prediction.currentIndexChanged.connect(self.on_cb_model_prediction)

        self.model_names = self.get_model_names()
        # train
        self.load_arg('args.txt')
        self.controller_ini()

        # predict
        self.predict_model = None
        self.classes = []
        self.predict_image = None
        self.prediction = Prediction(self.predict_model, self.classes)

    def controller_ini(self):
        self.ui.le_data.setText(self.dict['data'])
        self.ui.le_moder_dir.setText(self.dict['model_dir'])
        self.ui.le_resume.setText(self.dict['resume'])
        self.ui.le_epochs.setText(str(self.dict['epochs']))
        self.ui.le_learning_rate.setText(str(self.dict['lr']))
        self.ui.le_batch_size.setText(str(self.dict['batch_size']))
        self.ui.cb_model_names.addItems(self.model_names)
        self.ui.cb_model_prediction.addItems(self.model_names)
        for index in range(self.ui.cb_model_names.count()):
            if self.dict['arch'] == self.ui.cb_model_names.itemText(index):
                self.ui.cb_model_names.setCurrentIndex(index)

        self.ui.cb_workers.setCurrentIndex(self.dict['workers'])

    def get_avg(self):
        self.dict['data'] = self.ui.le_data.text() + ''
        self.dict['model_dir'] = self.ui.le_moder_dir.text() + ''
        if self.ui.ch_checkpoint.isChecked():
            self.dict['resume'] = self.ui.le_resume.text() + ''
        else:
            self.dict['resume'] = ''
        self.dict['epochs'] = int(self.ui.le_epochs.text())
        self.dict['lr'] = float(self.ui.le_learning_rate.text())
        self.dict['batch_size'] = int(self.ui.le_batch_size.text())
        self.dict['arch'] = self.ui.cb_model_names.currentText()
        self.dict['workers'] = self.ui.cb_workers.currentIndex()

    def load_arg(self, file_name):
        with open(file_name, 'r') as f:
            self.dict = json.load(f)

    def save_age(self):
        self.get_avg()
        with open('args.txt', 'w') as f:
            json.dump(self.dict, f, indent=2)

    def get_model_names(self):
        return sorted(name for name in models.__dict__
                      if name.islower() and not name.startswith("__")
                      and callable(models.__dict__[name]))

    def load_pretrained_model(self):
        if self.ui.cb_model_prediction.currentIndex() > 0:
            self.prediction.load_pretrained_model(self.ui.cb_model_prediction.currentText(), "imagenet_classes.txt")

    def load_retrained_model(self):
        model_path = self.ui.le_model.text().strip()
        classes_path = self.ui.le_classes.text().strip()
        if len(model_path) > 0 and len(classes_path) > 0:
            self.prediction.load_retrained_model(model_path=model_path, classes=classes_path)

    def do_predict(self, image):
        self.ui.lab_image.setPixmap(QPixmap(image))
        result = self.prediction.predict(image)
        text = ""
        if result:
            for name, pre in result:
                text += "{}({:0.2f}%)<br/>".format(name, pre)
            self.ui.lab_result.setText(text)

    def do_train(self):
        simplified_train.main()
        #t = threading.Thread(target=simplified_train.main)
        #.setDaemon(True)
        #t.start()
        #t.join()

    # 以下为事件处理函数

    # train
    @pyqtSlot()
    def on_btn_data_dir_clicked(self):
        cur_path = QDir.currentPath()
        selected_dir = QFileDialog.getExistingDirectory(self, "选择训练图像的路径", cur_path, QFileDialog.ShowDirsOnly)
        self.ui.le_data.setText(selected_dir + '')

    @pyqtSlot()
    def on_btn_model_dir_clicked(self):
        cur_path = QDir.currentPath()
        selected_dir = QFileDialog.getExistingDirectory(self, "选择模型保存的路径", cur_path, QFileDialog.ShowDirsOnly)
        self.ui.le_moder_dir.setText(selected_dir + '')

    @pyqtSlot()
    def on_btn_checkpoint_clicked(self):
        cur_path = QDir.currentPath()
        filt = "pth.tar文件(*.tar)"
        filename, filtUsed = QFileDialog.getOpenFileName(self, "请选择监测点模型文件", cur_path, filt)
        if filtUsed:
            self.ui.le_resume.setText(filename + '')
        print(filtUsed)

    @pyqtSlot()
    def on_btn_load_default_clicked(self):
        dlgtitle = "消息框"
        strInfo = "导入默认参数将覆盖当前设置的参数!，是否继续导入？"
        defaultBtn = QMessageBox.No
        result = QMessageBox.question(self, dlgtitle, strInfo, QMessageBox.Yes | QMessageBox.No, defaultBtn)
        if result == QMessageBox.Yes:
            self.load_arg('default_args.txt')
            self.controller_ini()


    @pyqtSlot()
    def on_btn_save_clicked(self):
        self.save_age()
        dlgtitle = "消息框"
        strInfo = "训练参数保存成功!，是否开始训练？"
        defaultBtn = QMessageBox.No
        result = QMessageBox.question(self, dlgtitle, strInfo, QMessageBox.Yes | QMessageBox.No, defaultBtn)
        if result == QMessageBox.Yes:
            self.do_train()

    @pyqtSlot()
    def on_btn_train_clicked(self):
        self.do_train()

    # predict

    @pyqtSlot()
    def on_btn_model_clicked(self):
        cur_path = QDir.currentPath()
        filt = "pth.tar文件(*.tar)"
        filename, filtUsed = QFileDialog.getOpenFileName(self, "请选择模型文件", cur_path, filt)
        if filtUsed:
            self.ui.le_model.setText(filename)
            self.load_retrained_model()

    @pyqtSlot()
    def on_btn_classes_clicked(self):
        cur_path = QDir.currentPath()
        filt = "txt文件(*.txt)"
        filename, filtUsed = QFileDialog.getOpenFileName(self, "类名文件", cur_path, filt)
        if filtUsed:
            self.ui.le_classes.setText(filename)
            self.load_retrained_model()

    @pyqtSlot()
    def on_btn_image_clicked(self):
        if self.dict['temp_dir'] == "":
            cur_path = QDir.currentPath()
        else:
            cur_path = self.dict['temp_dir']
        filt = "图片文件(*.jpg) ;; 图片文件(*.png);;All(*.*)"
        filename, filtUsed = QFileDialog.getOpenFileName(self, "请选择预测图片", cur_path, filt)
        if filtUsed:
            self.dict['temp_dir'] = os.path.dirname(os.path.realpath(filename))
            self.ui.le_image.setText(filename)
            self.do_predict(filename)

    @pyqtSlot()
    def on_rb_clicked(self, rb):
        if rb.objectName() == "rb_pretrained":
            self.ui.cb_model_prediction.setEnabled(rb.isChecked())
        else:
            self.ui.le_model.setEnabled(rb.isChecked())
            self.ui.btn_classes.setEnabled(rb.isChecked())
            self.ui.le_classes.setEnabled(rb.isChecked())
            self.ui.btn_model.setEnabled(rb.isChecked())

    @pyqtSlot()
    def on_cb_model_prediction(self):
        self.load_pretrained_model()

    # export
    @pyqtSlot()
    def on_btn_export_clicked(self):
        model_path = self.ui.le_export_model.text()
        export_dir = self.ui.le_onnx_dir.text()
        if len(model_path) > 0 and len(export_dir) > 0:
            simplified_export.export(model_path, export_dir)

    @pyqtSlot()
    def on_btn_export_model_clicked(self):
        cur_path = QDir.currentPath()
        filt = "pth.tar文件(*.tar)"
        filename, filtUsed = QFileDialog.getOpenFileName(self, "请选择模型文件", cur_path, filt)
        if filtUsed:
            self.ui.le_export_model.setText(filename)

    @pyqtSlot()
    def on_btn_export_dir_clicked(self):
        cur_path = QDir.currentPath()
        selected_dir = QFileDialog.getExistingDirectory(self, "请选择onnx模型的保存路径", cur_path, QFileDialog.ShowDirsOnly)
        self.ui.le_onnx_dir.setText(selected_dir + '')


if __name__ == '__main__':
    app = QApplication(sys.argv)
    gui = ClassificationGui()
    gui.show()
    sys.exit(app.exec_())