'''
Realtime Display 3D Human Reconstrction
3D image drawing by pygtagrph based on OpenGL
speed about 25 FPS
'''
import math
import sys
from argparse import ArgumentParser

import cv2
import pyqtgraph as pg
import pyqtgraph.opengl as gl
from pyqtgraph.Qt import QtCore, QtGui
from pyqtgraph.opengl import *

from joints_detectors.openpose.main import load_model as Model2Dload

model2D = Model2Dload()
from joints_detectors.openpose.main import generate_frame_kpt as OpenPoseInterface

interface2D = OpenPoseInterface
from tools.utils import videopose_model_load as Model3Dload

model3D = Model3Dload()
from tools.utils import interface as VideoPoseInterface

interface3D = VideoPoseInterface
from tools.utils import draw_2Dimg, resize_img, common

common = common()
item = 0
item_num = 0
pos_init = np.zeros(shape=(17, 3))


class Visualizer(object):
    def __init__(self, input_video):
        self.traces = dict()
        self.app = QtGui.QApplication(sys.argv)
        self.w = gl.GLViewWidget()
        self.w.opts['distance'] = 45.0  ## distance of camera from center
        self.w.opts['fov'] = 60  ## horizontal field of view in degrees
        self.w.opts['elevation'] = 10  ## camera's angle of elevation in degrees 仰俯角
        self.w.opts['azimuth'] = 90  ## camera's azimuthal angle in degrees 方位角
        self.w.setWindowTitle('pyqtgraph example: GLLinePlotItem')
        self.w.setGeometry(450, 700, 980, 700)  # 原点在左上角
        self.w.show()

        # create the background grids
        gx = gl.GLGridItem()
        gx.rotate(90, 0, 1, 0)
        gx.translate(-10, 0, 0)
        self.w.addItem(gx)
        gy = gl.GLGridItem()
        gy.rotate(90, 1, 0, 0)
        gy.translate(0, -10, 0)
        self.w.addItem(gy)
        gz = gl.GLGridItem()
        gz.translate(0, 0, -10)
        self.w.addItem(gz)

        # special setting
        self.cap = cv2.VideoCapture(input_video)
        self.video_name = input_video.split('/')[-1].split('.')[0]
        self.kpt2Ds = []
        pos = pos_init

        for j, j_parent in enumerate(common.skeleton_parents):
            if j_parent == -1:
                continue

            x = np.array([pos[j, 0], pos[j_parent, 0]]) * 10
            y = np.array([pos[j, 1], pos[j_parent, 1]]) * 10
            z = np.array([pos[j, 2], pos[j_parent, 2]]) * 10 - 10
            pos_total = np.vstack([x, y, z]).transpose()
            self.traces[j] = gl.GLLinePlotItem(pos=pos_total, color=pg.glColor((j, 10)), width=6, antialias=True)
            self.w.addItem(self.traces[j])

    def start(self):
        if (sys.flags.interactive != 1) or not hasattr(QtCore, 'PYQT_VERSION'):
            QtGui.QApplication.instance().exec_()

    def set_plotdata(self, name, points, color, width):
        self.traces[name].setData(pos=points, color=color, width=width)

    def update(self):
        global item
        global item_num
        num = item / 2
        azimuth_value = abs(num % 120 + math.pow(-1, int((num / 120))) * 120) % 120
        self.w.opts['azimuth'] = azimuth_value
        print(item, '  ')
        _, frame = self.cap.read()
        if item % 2 != 1:
            frame, W, H = resize_img(frame)
            joint2D = interface2D(frame, model2D)
            img2D = draw_2Dimg(frame, joint2D, 1)
            if item == 0:
                for _ in range(30):
                    self.kpt2Ds.append(joint2D)
            elif item < 30:
                self.kpt2Ds.append(joint2D)
                self.kpt2Ds.pop(0)
            else:
                self.kpt2Ds.append(joint2D)
                self.kpt2Ds.pop(0)

            item += 1
            joint3D = interface3D(model3D, np.array(self.kpt2Ds), W, H)
            pos = joint3D[-1]  # (17, 3)

            for j, j_parent in enumerate(common.skeleton_parents):
                if j_parent == -1:
                    continue
                x = np.array([pos[j, 0], pos[j_parent, 0]]) * 10
                y = np.array([pos[j, 1], pos[j_parent, 1]]) * 10
                z = np.array([pos[j, 2], pos[j_parent, 2]]) * 10 - 10
                pos_total = np.vstack([x, y, z]).transpose()
                self.set_plotdata(
                    name=j, points=pos_total,
                    color=pg.glColor((j, 10)),
                    width=6)

            # save
            if item_num < 10:
                name = '000' + str(item_num)

            elif item_num < 100:
                name = '00' + str(item_num)

            elif item_num < 1000:
                name = '0' + str(item_num)

            else:
                name = str(item_num)
            im3Dname = 'VideoSave/' + '3D_' + name + '.png'
            d = self.w.renderToArray((img2D.shape[1], img2D.shape[0]))  # (W, H)
            print('Save 3D image: ', im3Dname)
            pg.makeQImage(d).save(im3Dname)

            im2Dname = 'VideoSave/' + '2D_' + name + '.png'
            print('Save 2D image: ', im2Dname)
            cv2.imwrite(im2Dname, img2D)

            item_num += 1
        else:
            item += 1

    def animation(self):
        timer = QtCore.QTimer()
        timer.timeout.connect(self.update)
        timer.start(1)
        self.start()


def main():
    parser = ArgumentParser()
    parser.add_argument("-i", "--video", help="input video file name", default="/home/xyliu/Videos/sports/dance.mp4")
    args = parser.parse_args()
    print(args.video)
    v = Visualizer(args.video)
    v.animation()
    cv2.destroyAllWindows()


if __name__ == '__main__':
    main()