import cv2
from PyQt5.QtCore import pyqtSignal, QThread
from PyQt5.QtGui import *
from PyQt5.QtWidgets import *
from PyQt5 import QtWidgets, QtGui
from PyQt5 import QtCore, QtGui, QtWidgets
from nst import Models
import time
import numpy as np

# 子窗口布局
from sub_windows import ui_sub_window_9


class SubWindow(QMainWindow):
    def __init__(self):
        super().__init__(parent=None)
        self.ui = ui_sub_window_9.Ui_Form()
        self.ui.setupUi(self)
        self.ui_init()
        self.file_path = None
        self.style_path = None

    def ui_init(self):
        # print("子窗口9初始化")
        self.ui.pushButton_open_style.clicked.connect(self.open_style)
        self.ui.pushButton_open_file.clicked.connect(self.open_file)
        self.ui.pushButton_generate.clicked.connect(self.generate)

    def open_file(self):
        # Open a file dialog and display the selected image
        options = QtWidgets.QFileDialog.Options()
        options |= QtWidgets.QFileDialog.ReadOnly
        fileName, _ = QtWidgets.QFileDialog.getOpenFileName(
            self,
            "QFileDialog.getOpenFileName()",
            "",
            "All Files (*);;Python Files (*.py)",
            options=options,
        )
        if fileName:
            self.file_path = fileName
            pixmap = QtGui.QPixmap(fileName)
            scaled_pixmap = pixmap.scaled(
                self.ui.label_image_1.size(), QtCore.Qt.KeepAspectRatio
            )
            self.ui.label_image_1.setPixmap(scaled_pixmap)

    def open_style(self):
        # Open a file dialog and display the selected style image
        options = QtWidgets.QFileDialog.Options()
        options |= QtWidgets.QFileDialog.ReadOnly
        fileName, _ = QtWidgets.QFileDialog.getOpenFileName(
            self,
            "QFileDialog.getOpenFileName()",
            "",
            "All Files (*);;Python Files (*.py)",
            options=options,
        )
        if fileName:
            self.style_path = fileName
            pixmap = QtGui.QPixmap(fileName)
            scaled_pixmap = pixmap.scaled(
                self.ui.label_image_2.size(), QtCore.Qt.KeepAspectRatio
            )
            self.ui.label_image_2.setPixmap(scaled_pixmap)

    def generate(self):
        print("generating")
        ph = "nst/weights/vgg19-dcbb9e9d.pth"

        transfer = Models.Transfer(self.file_path, self.style_path, ph)

        t = time.time()
        dt, img = transfer.fit()
        print("time:")
        print(time.time() - t)
        # print(dt,img)

        img = np.array(img)[:, :, ::-1]
        img = np.require(img, np.uint8, 'C')

        qimage = QImage(img.data, img.shape[1], img.shape[0], QImage.Format_RGB888)
        pixmap = QPixmap.fromImage(qimage)
        self.ui.label_image_3.setPixmap(pixmap)