import numpy as np
import pyqtgraph as pg
from PyQt5 import QtWidgets, QtGui, QtCore

class WidgetPose(QtWidgets.QWidget):

    video_label_signal = QtCore.pyqtSignal(list)

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

        layout = QtWidgets.QGridLayout()
        self.poseGraph = QtWidgets.QWidget()
        self.measures = QtWidgets.QWidget()
        self.measuresLayout = QtWidgets.QHBoxLayout()

        # Setup Movement Graph Plot Widget
        self.movementGraph = pg.PlotWidget()
        self.movementGraph.setBackground('w')
        self.movementGraph.setYRange(0, 400, padding=0)
        self.movementGraph.getPlotItem().getAxis('bottom').setTickSpacing(minor=50, major=100)
        self.movementGraph.getPlotItem().setTitle(title='Body Movement over Time')
        self.movementPlots = []

        layout.addWidget(self.movementGraph, 1, 0, 1, 1)
        layout.addWidget(self.measures, 1, 1, 1, 1)
        layout.setColumnStretch(0, 1)
        layout.setColumnStretch(1, 1)

        self.setLayout(layout)
        self.colors = None
        self.labels = dict()

    @QtCore.pyqtSlot(list)
    def onSelectedID(self, lst):
        """Change color of movement graph plot if ID should not be visible"""
        for i, button in enumerate(lst):
            if not button.isChecked():
                self.movementPlots[i].setPen(None)
            elif self.colors is not None:
                color = tuple([int(a * 255) for a in self.colors[i]])
                pen = pg.mkPen(color=color)
                self.movementPlots[i].setPen(pen)

    @QtCore.pyqtSlot(dict, list, int)
    def updateMovementGraph(self, data, colors, numberIDs):
        """Plot ID specific movement data from processing class
        data[id: (movements, frames)]
        """
        # handle NaN https://github.com/pyqtgraph/pyqtgraph/issues/1057
        # downgrade to 5.13 fixes the issue
        for id_no in range(numberIDs):
            if data.get(id_no):
                if not np.all(np.isnan(data.get(id_no)[0])):
                    self.movementPlots[id_no].setData(data.get(id_no)[1], data.get(id_no)[0])  # Update the data.

        self.movementGraph.setXRange(np.min(data.get(id_no)[1]), np.max(data.get(id_no)[1]))

    @QtCore.pyqtSlot(dict, int)
    def updateHandVelocity(self, data, numberIDs):
        """Update Velocity Label
        data[id: velocity for frame]"""
        for id_no in range(numberIDs):
            if data.get(id_no) is not None:
                self.labels['Velocity'][id_no].setText('%.2f' % data[id_no])
            else:
                self.labels['Velocity'][id_no].setText(' ')

    @QtCore.pyqtSlot(list, int)
    def initMovementGraph(self, colors, numberIDs):
        """Initialize plot lines with 0
        colors: plot color for each ID
        """
        for i in range(numberIDs):
            x = list(range(-200, 0))  # 100 time points
            y = [0 for _ in range(200)]  # 100 data points
            color = tuple([int(a * 255) for a in colors[i]])
            pen = pg.mkPen(color=color)
            dataLine = self.movementGraph.plot(x, y, pen=pen)
            self.movementPlots.append(dataLine)

    @QtCore.pyqtSlot(dict, dict, list, int)
    def setInit(self, mostActivity, hand, colors, numberIDs):
        self.colors = colors
        idLayout = QtWidgets.QVBoxLayout()

        labelID = QtWidgets.QLabel('      ')
        labelID.setFixedWidth(60)
        labelID.setFixedHeight(20)
        label = QtWidgets.QLabel('Most body activity in frame: ')
        labelA = QtWidgets.QLabel('Hands above table (relative): ')
        labelG = QtWidgets.QLabel('Gestures (relative): ')
        labelVel = QtWidgets.QLabel('Hand velocity: ')
        idLayout.addWidget(labelID)
        idLayout.addWidget(label)
        idLayout.addWidget(labelA)
        idLayout.addWidget(labelG)
        idLayout.addWidget(labelVel)
        self.measuresLayout.insertLayout(-1, idLayout)

        activityLabel = []
        aboveLabel = []
        gestureLabel = []
        velLabel = []
        for id_no in range(numberIDs):
            idLayout = QtWidgets.QVBoxLayout()

            [total, tracked, high_vel] = hand.get(id_no)

            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)
            labelID.setFixedHeight(20)

            label = QtWidgets.QLabel('%i' % mostActivity.get(id_no))
            activityLabel.append(label)
            labelG = QtWidgets.QLabel('%.2f' % (high_vel / total))
            gestureLabel.append(labelG)
            labelA = QtWidgets.QLabel('%.2f' % (tracked / total))
            aboveLabel.append(labelA)
            labelVel = QtWidgets.QLabel(' ')
            velLabel.append(labelVel)
            idLayout.addWidget(labelID)
            idLayout.addWidget(label)
            idLayout.addWidget(labelA)
            idLayout.addWidget(labelG)
            idLayout.addWidget(labelVel)

            self.measuresLayout.insertLayout(-1, idLayout)

        # Velocity will be updated each frame, rest is updated in _updateLabels
        self.labels['Velocity'] = velLabel
        self.labels['Above'] = aboveLabel
        self.labels['Gesture'] = gestureLabel
        self.labels['Activity'] = activityLabel
        self.measures.setLayout(self.measuresLayout)

    @QtCore.pyqtSlot(dict, dict, int)
    def updateLables(self, mostActivity, hand, numberIDs):
        """ Update above hands, gestures and most activity labels when segment was changed"""
        for id_no in range(numberIDs):
            [total, tracked, high_vel] = hand[id_no]
            self.labels['Activity'][id_no].setText('%i' % mostActivity[id_no])
            self.labels['Above'][id_no].setText('%.2f' % (tracked / total))
            self.labels['Gesture'][id_no].setText('%.2f' % (high_vel / total))