import numpy as np

import pyqtgraph as pg
from PyQt5 import QtWidgets, QtCore, QtGui
from PyQt5.QtGui import QColor

lst_aus = ['AU1: Inner Brow Raiser', 'AU2: Outer Brow Raiser', 'AU4: Brow Lowerer', 'AU5: Upper Lid Raiser',
              'AU6: Cheek Raiser', 'AU9: Nose Wrinkler', 'AU12: Lip Corner Puller', 'AU15: Lip Corner Depressor',
              'AU17: Chin Raiser',  'AU20: Lip Stretcher', 'AU25: Lips Part',  'AU26: Jaw Drop']

class WidgetFacialExpression(QtWidgets.QWidget):
    def __init__(self, parent=None):
        super(WidgetFacialExpression, self).__init__(parent)

        self.faceLayout = QtWidgets.QHBoxLayout()
        self.setLayout(self.faceLayout)

        self.numberIDs = None
        self.valueLabels = dict()
        self.imgPlots = dict()

    @QtCore.pyqtSlot(list, int)
    def setInit(self, colors, numberIDs):
        self.numberIDs = numberIDs

        for id_no in range(numberIDs):
            idLayout = QtWidgets.QHBoxLayout()
            labelNameLayout = QtWidgets.QVBoxLayout()
            labelValueLayout = QtWidgets.QVBoxLayout()
            imageLayout = QtWidgets.QVBoxLayout()
            imageWidget = pg.PlotWidget(background=QColor(53, 53, 53))
            imageWidget.invertY()
            imageWidget.hideAxis('bottom'), imageWidget.hideAxis('left')
            imageWidget.setMaximumHeight(150), imageWidget.setMaximumWidth(150)
            imageWidget.setAspectLocked(True)
            self.imgPlots[id_no] = imageWidget

            color = tuple([int(a * 255) for a in colors[id_no]])

            labelID = QtWidgets.QLabel('ID%i' % id_no)
            labelID.setStyleSheet('font: bold 12px; color: black; background-color: rgb(%i,%i,%i)' % color)



            #labelID.setFixedWidth(60)

            labelNameLayout.addWidget(labelID)
            labelID = QtWidgets.QLabel(' ')

            #labelID.setStyleSheet('background-color: rgb(%i,%i,%i)' % color)
            labelValueLayout.addWidget(labelID)
            lst = []
            for au in lst_aus:
                nLabel = QtWidgets.QLabel(au)
                labelNameLayout.addWidget(nLabel)

                vLabel = QtWidgets.QLabel(' ')
                labelValueLayout.addWidget(vLabel)
                lst.append(vLabel)

            self.valueLabels[id_no] = lst
            idLayout.addWidget(imageWidget)
            idLayout.addLayout(labelNameLayout)
            idLayout.addLayout(labelValueLayout)
            self.faceLayout.addLayout(idLayout)

    @QtCore.pyqtSlot(dict, int)
    def updateImages(self, imgs, id_no):
        if imgs[id_no] is not None:
            img = np.moveaxis(imgs[id_no], 0, 1)
            img = pg.ImageItem(img)
            self.imgPlots[id_no].addItem(img)


    @QtCore.pyqtSlot(dict)
    def updateFrame(self, aus):
        if self.numberIDs is None:
            return

        for id_no in range(self.numberIDs):
            if len(aus[id_no]) > 0:
                for i, label in enumerate(self.valueLabels[id_no]):
                    if not np.any(np.isnan(np.array(aus[id_no].flatten()[0], dtype=np.float64))):
                        label.setText('%.2f' % aus[id_no].flatten()[0][i])