import numpy as np
import pyqtgraph
import matplotlib

matplotlib.use("Qt5Agg")
import matplotlib.animation as animation
from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as FigureCanvas
from matplotlib.figure import Figure
from PyQt5 import QtWidgets, QtGui
from PyQt5 import QtCore
import pyqtgraph as pg
import pyqtgraph.exporters

from utils.colors import random_colors
from utils.util import sperical2equirec


class WidgetGaze(QtWidgets.QWidget):

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

        layout = QtWidgets.QGridLayout()

        # Setup gaze graph
        self.gazeGraph = pg.PlotWidget()
        self.gazeGraph.setBackground('w')
        self.gazeGraph.setYRange(-1.25, 1.25, padding=0)
        self.gazeGraph.setXRange(-1.25, 1.25, padding=0)
        self.gazeGraph.hideAxis('left')
        self.gazeGraph.hideAxis('bottom')
        self.gazeGraph.setAspectLocked()
        self.gazeGraph.getPlotItem().setTitle(title='Top-down View of Gaze')
        self.gazeGraphPlots = []

        self.measures = QtWidgets.QWidget()
        self.measuresLayout = QtWidgets.QHBoxLayout()

        # self.gazeMap = QtWidgets.QWidget()
        # self.heatmapSlider = HeatmapSlider()
        # self.heatmapSlider.signalSetThreshold.connect(self.setThreshold)
        # self.heatmapSlider.signalSaveImage.connect(self.gazeMap.saveImage)

        # row, column, row span, column span
        layout.addWidget(self.measures, 0, 1, 2, 1)
        layout.addWidget(self.gazeGraph, 0, 0, 2, 1)
        layout.setColumnStretch(0, 1)
        layout.setColumnStretch(1, 1)

        self.setLayout(layout)
        # layout.addWidget(self.gazeMap, 0, 0, 3, 1)
        # layout.addWidget(self.heatmapSlider, 3, 0, 1, 1)
        self.gazeLabels = []
        self.colors = None

    @QtCore.pyqtSlot(dict, list, int)
    def setInit(self, measures, colors, numberIDs):
        """Initialize measure widget with labels for all IDs"""
        self.colors = colors  # Necessary for ID updates
        idLayout = QtWidgets.QVBoxLayout()
        labelID = QtWidgets.QLabel('      ')
        labelID.setFixedWidth(60)
        labelID.setFixedHeight(20)
        labelA = QtWidgets.QLabel('LookSomeone: ')
        labelNoLook = QtWidgets.QLabel('TotalNoLook: ')
        labelG = QtWidgets.QLabel('TotalWatched: ')
        labelRatio = QtWidgets.QLabel('RatioWatcherLookSOne: ')
        label = QtWidgets.QLabel('Tracked: ')
        # labelVel = QtWidgets.QLabel('totNoLook: ')
        idLayout.addWidget(labelID)
        idLayout.addWidget(labelA)
        idLayout.addWidget(labelNoLook)
        idLayout.addWidget(labelG)
        idLayout.addWidget(labelRatio)
        idLayout.addWidget(label)
        # idLayout.addWidget(labelVel)
        self.measuresLayout.insertLayout(-1, idLayout)

        for id_no in range(numberIDs):
            idLayout = QtWidgets.QVBoxLayout()

            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)
            # Look Someone
            labelA = QtWidgets.QLabel('{:.2%}'.format(measures[id_no][1] / measures[id_no][2]))
            labelNoLook = QtWidgets.QLabel('{:.2%}'.format((measures[id_no][2] - measures[id_no][1]) / measures[id_no][2]))
            # Total Watched
            labelG = QtWidgets.QLabel('{:.2%}'.format(measures[id_no][0] / measures[id_no][2]))
            # ratio totWatcher / lookSomeone
            labelRatio = QtWidgets.QLabel('{:.2}'.format(measures[id_no][0] / measures[id_no][1]))
            label = QtWidgets.QLabel('%i frames' % measures[id_no][2])
            # labelVel = QtWidgets.QLabel('%.2f' % np.random.uniform(0, 1))
            idLayout.addWidget(labelID)
            idLayout.addWidget(labelA)
            idLayout.addWidget(labelNoLook)
            idLayout.addWidget(labelG)
            idLayout.addWidget(labelRatio)
            idLayout.addWidget(label)
            # idLayout.addWidget(labelVel)
            # self.gazeLabels.append(labelVel)
            self.measuresLayout.insertLayout(-1, idLayout)
        self.measures.setLayout(self.measuresLayout)

    @QtCore.pyqtSlot(list, int)
    def initGazeGraph(self, colors, numberIDs):
        """ initialize gaze graph """
        # Big circle
        x1, y1 = self.get_circle(radius=1)
        self.gazeGraph.addItem(self.gazeGraph.plot(x1, y1, pen=pg.mkPen(0.5)))

        # Camera
        x2, y2 = self.get_circle(radius=0.02)
        self.gazeGraph.addItem(self.gazeGraph.plot(x2, y2, pen=pg.mkPen(color=(0, 0, 0), width=3)))

        for id_no in range(numberIDs):
            color = tuple([int(a * 255) for a in colors[id_no]])
            plt = self.gazeGraph.plot(x=[], y=[], pen=pg.mkPen(color=color, width=2))
            self.gazeGraphPlots.append(plt)

    @QtCore.pyqtSlot(dict, int)
    def updateGazeGraph(self, data, numberIDs):
        """ frame updates for gaze graph """
        for id_no in range(numberIDs):
            if data and id_no in data:
                self.gazeGraphPlots[id_no].setData(data[id_no][0], data[id_no][1])

    def get_circle(self, radius):
        """ helper function returns circle to x, y coordinates"""
        theta = np.linspace(0, 2 * np.pi, 100)
        x = radius * np.cos(theta)
        y = radius * np.sin(theta)
        return np.array(x), np.array(y)

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