import torch
import os
import wavmark
from wavmark.utils import file_reader
from PyQt5 import QtCore, QtGui, QtWidgets
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
global file_path_g


def funcNNWProof( input):
    '''
    Verify the inference
    '''

    payload = [0, 1, 1, 1, 1, 0, 0, 0, 0, 1, 1, 0, 1, 0, 1, 1]
    model = wavmark.load_model().to(device)
    signal = file_reader.read_as_single_channel(input, aim_sr=16000)
    # 5.decode watermark
    payload_decoded, _ = wavmark.decode_watermark(model, signal, show_progress=True)
    if isinstance(payload_decoded, type(None)): return False
    BER = (payload != payload_decoded).mean() * 100
    return BER == 0




class DragDropMainWindow(QtWidgets.QMainWindow):
    fileDropped = QtCore.pyqtSignal(str)

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

    def dragEnterEvent(self, event):
        if event.mimeData().hasUrls():
            event.accept()
        else:
            event.ignore()

    def dropEvent(self, event):
        for url in event.mimeData().urls():
            file_path = str(url.toLocalFile())
            self.fileDropped.emit(file_path)


class Ui_MainWindow(object):
    def setupUi(self, MainWindow):
        MainWindow.setObjectName("MainWindow")
        MainWindow.resize(800, 600)
        self.centralwidget = QtWidgets.QWidget(MainWindow)
        self.centralwidget.setObjectName("centralwidget")
        self.pushButton = QtWidgets.QPushButton(self.centralwidget)
        self.pushButton.setGeometry(QtCore.QRect(300, 470, 221, 61))
        self.pushButton.setObjectName("pushButton")
        self.pushButton_2 = QtWidgets.QPushButton(self.centralwidget)
        self.pushButton_2.setGeometry(QtCore.QRect(590, 60, 191, 81))
        self.pushButton_2.setObjectName("pushButton_2")
        self.label = QtWidgets.QLabel(self.centralwidget)
        self.label.setGeometry(QtCore.QRect(70, 110, 721, 391))
        self.label.setObjectName("label")
        MainWindow.setCentralWidget(self.centralwidget)
        self.statusbar = QtWidgets.QStatusBar(MainWindow)
        self.statusbar.setObjectName("statusbar")
        MainWindow.setStatusBar(self.statusbar)

        self.retranslateUi(MainWindow)
        QtCore.QMetaObject.connectSlotsByName(MainWindow)

        # Connect button clicks to functions
        self.pushButton.clicked.connect(self.run_UseCase)
        self.pushButton_2.clicked.connect(self.watermark_proof)

        self.file_path_g=None

    def retranslateUi(self, MainWindow):
        _translate = QtCore.QCoreApplication.translate
        MainWindow.setWindowTitle(_translate("MainWindow", "MainWindow"))
        self.pushButton.setText(_translate("MainWindow", "Running the UseCase"))
        self.pushButton_2.setText(_translate("MainWindow", "Watermarking proof"))
        self.label.setText(_translate("MainWindow", "<html><head/><body><p><img src=\"MPAI_NNW-MQA.png\"/></p></body></html>"))

    def run_UseCase(self):
        # Function to execute when the "Running the UseCase" button is clicked
        print("Openning new Window")
        os.system("gnome-terminal & disown")

    def watermark_proof(self):
        # Function to execute when the "WaterMarking proof" button is clicked
        if self.file_path_g==None:
            print("Please, first drag an audio file")
        else:
            print("Processing...")
            answer=funcNNWProof(self.file_path_g)
            if answer:
                print("This audio is watermarked")
            else:
                print("This audio is not watermarked")

if __name__ == "__main__":
    import sys

    app = QtWidgets.QApplication(sys.argv)
    MainWindow = DragDropMainWindow()
    ui = Ui_MainWindow()
    ui.setupUi(MainWindow)

    def file_dropped(file_path):
        ui.file_path_g=file_path


    MainWindow.fileDropped.connect(file_dropped)

    MainWindow.show()
    sys.exit(app.exec_())
