import zmq
import sys
import os

sys.path.append(os.getcwd() + '/build')
sys.path.append(os.getcwd() + 'plot')
sys.path.append(os.getcwd())
import time

import json
from ruamel.yaml import YAML
import numpy as np
# from pinocchio.visualize import MeshcatVisualizer
# import pinocchio as pin

from utils.python_utils import util

from plot.data_saver import DataSaver

from messages.draco_pb2 import *

import argparse

parser = argparse.ArgumentParser()
parser.add_argument("--b_visualize", type=bool, default=False)
args = parser.parse_args()

with open("config/draco/pnc.yaml", 'r') as stream:
    config = YAML().load(stream)
    addr = config["ip_addr"]

context = zmq.Context()
pnc_socket = context.socket(zmq.SUB)
pnc_socket.connect(addr)
pnc_socket.setsockopt_string(zmq.SUBSCRIBE, "")

pj_context = zmq.Context()
pj_socket = pj_context.socket(zmq.PUB)
pj_socket.bind("tcp://*:9872")

data_saver = DataSaver()

if args.b_visualize:
    model, collision_model, visual_model = pin.buildModelsFromUrdf(
        "robot_model/draco/draco.urdf", "robot_model/draco",
        pin.JointModelFreeFlyer())
    viz = MeshcatVisualizer(model, collision_model, visual_model)
    try:
        viz.initViewer(open=True)
    except ImportError as err:
        print(
            "Error while initializing the viewer. It seems you should install python meshcat"
        )
        print(err)
        exit()
    viz.loadViewerModel()
    vis_q = pin.neutral(model)

    icp_model, icp_collision_model, icp_visual_model = pin.buildModelsFromUrdf(
        "robot_model/ground/sphere.urdf", "robot_model/ground",
        pin.JointModelFreeFlyer())
    icp_viz = MeshcatVisualizer(icp_model, icp_collision_model,
                                icp_visual_model)
    icp_viz.initViewer(viz.viewer)
    icp_viz.loadViewerModel(rootNodeName="icp", color=[0., 0., 1., 0.5])
    icp_viz_q = pin.neutral(icp_model)

    icp_des_model, icp_des_collision_model, icp_des_visual_model = pin.buildModelsFromUrdf(
        "robot_model/ground/sphere.urdf", "robot_model/ground",
        pin.JointModelFreeFlyer())
    icp_des_viz = MeshcatVisualizer(icp_des_model, icp_des_collision_model,
                                    icp_des_visual_model)
    icp_des_viz.initViewer(viz.viewer)
    icp_des_viz.loadViewerModel(rootNodeName="icp_des",
                                color=[1., 0., 0., 0.5])
    icp_des_viz_q = pin.neutral(icp_des_model)

    cmp_model, cmp_col, cmp_vis = pin.buildModelsFromUrdf(
        "robot_model/ground/sphere.urdf", "robot_model/ground",
        pin.JointModelFreeFlyer())
    cmp_viz = MeshcatVisualizer(cmp_model, cmp_col, cmp_vis)
    cmp_viz.initViewer(viz.viewer)
    cmp_viz.loadViewerModel(rootNodeName="cmp_des", color=[0., 1., 0., 0.5])
    cmp_des_q = pin.neutral(cmp_model)

msg = pnc_msg()

while True:

    encoded_msg = pnc_socket.recv()
    msg.ParseFromString(encoded_msg)

    # save pkl
    data_saver.add('time', msg.time)
    data_saver.add('phase', msg.phase)

    data_saver.add('rf_sg', msg.rf_sg)
    data_saver.add('lf_sg', msg.lf_sg)

    data_saver.add('l_knee_int_frc_cmd', msg.l_knee_int_frc_cmd)
    data_saver.add('r_knee_int_frc_cmd', msg.r_knee_int_frc_cmd)

    data_saver.add('task_com_pos_des', list(msg.task_com_pos_des))
    data_saver.add('task_com_vel_des', list(msg.task_com_vel_des))
    data_saver.add('task_com_acc_des', list(msg.task_com_acc_des))
    data_saver.add('task_com_pos', list(msg.task_com_pos))
    data_saver.add('task_com_vel', list(msg.task_com_vel))

    data_saver.add('task_com_pos_des_local', list(msg.task_com_pos_des_local))
    data_saver.add('task_com_vel_des_local', list(msg.task_com_vel_des_local))
    data_saver.add('task_com_acc_des_local', list(msg.task_com_acc_des_local))
    data_saver.add('task_com_pos_local', list(msg.task_com_pos_local))
    data_saver.add('task_com_vel_local', list(msg.task_com_vel_local))

    data_saver.add('task_torso_ori_pos_des',
                   list(util.quat_to_rpy(msg.task_torso_ori_pos_des)))
    data_saver.add('task_torso_ori_vel_des', list(msg.task_torso_ori_vel_des))
    data_saver.add('task_torso_ori_acc_des', list(msg.task_torso_ori_acc_des))
    data_saver.add('task_torso_ori_pos',
                   list(util.quat_to_rpy(msg.task_torso_ori_pos)))
    data_saver.add('task_torso_ori_vel', list(msg.task_torso_ori_vel))

    # data_saver.add('task_torso_rpy_pos_des',
    # list(util.quat_to_rpy(msg.task_torso_ori_pos_des)))
    # data_saver.add('task_torso_rpy_pos',
    # list(util.quat_to_rpy(msg.task_torso_ori_pos)))

    data_saver.add('task_torso_ori_pos_des_local',
                   util.quat_to_rpy(msg.task_torso_ori_pos_des_local).tolist())
    data_saver.add('task_torso_ori_vel_des_local',
                   list(msg.task_torso_ori_vel_des_local))
    data_saver.add('task_torso_ori_acc_des_local',
                   list(msg.task_torso_ori_acc_des_local))
    data_saver.add('task_torso_ori_pos_local',
                   util.quat_to_rpy(msg.task_torso_ori_pos_local).tolist())
    data_saver.add('task_torso_ori_vel_local',
                   list(msg.task_torso_ori_vel_local))

    # data_saver.add('task_torso_rpy_pos_des_local',
    # list(util.quat_to_rpy(msg.task_torso_ori_pos_des_local)))
    # data_saver.add('task_torso_rpy_pos_local',
    # list(util.quat_to_rpy(msg.task_torso_ori_pos_local)))

    data_saver.add('task_rfoot_lin_pos_des', list(msg.task_rfoot_lin_pos_des))
    data_saver.add('task_rfoot_lin_vel_des', list(msg.task_rfoot_lin_vel_des))
    data_saver.add('task_rfoot_lin_acc_des', list(msg.task_rfoot_lin_acc_des))
    data_saver.add('task_rfoot_lin_pos', list(msg.task_rfoot_lin_pos))
    data_saver.add('task_rfoot_lin_vel', list(msg.task_rfoot_lin_vel))

    data_saver.add('task_rfoot_lin_pos_des_local',
                   list(msg.task_rfoot_lin_pos_des_local))
    data_saver.add('task_rfoot_lin_vel_des_local',
                   list(msg.task_rfoot_lin_vel_des_local))
    data_saver.add('task_rfoot_lin_acc_des_local',
                   list(msg.task_rfoot_lin_acc_des_local))
    data_saver.add('task_rfoot_lin_pos_local',
                   list(msg.task_rfoot_lin_pos_local))
    data_saver.add('task_rfoot_lin_vel_local',
                   list(msg.task_rfoot_lin_vel_local))

    data_saver.add('task_rfoot_ori_pos_des',
                   list(util.quat_to_rpy(msg.task_rfoot_ori_pos_des)))
    data_saver.add('task_rfoot_ori_vel_des', list(msg.task_rfoot_ori_vel_des))
    data_saver.add('task_rfoot_ori_acc_des', list(msg.task_rfoot_ori_acc_des))
    data_saver.add('task_rfoot_ori_pos',
                   list(util.quat_to_rpy(msg.task_rfoot_ori_pos)))
    data_saver.add('task_rfoot_ori_vel', list(msg.task_rfoot_ori_vel))

    # data_saver.add('task_rfoot_rpy_pos_des',
    # (util.quat_to_rpy(msg.task_rfoot_ori_pos_des)).tolist())
    # data_saver.add('task_rfoot_rpy_pos',
    # (util.quat_to_rpy(msg.task_rfoot_ori_pos)).tolist())

    data_saver.add('task_rfoot_ori_pos_des_local',
                   list(util.quat_to_rpy(msg.task_rfoot_ori_pos_des_local)))
    data_saver.add('task_rfoot_ori_vel_des_local',
                   list(msg.task_rfoot_ori_vel_des_local))
    data_saver.add('task_rfoot_ori_acc_des_local',
                   list(msg.task_rfoot_ori_acc_des_local))
    data_saver.add('task_rfoot_ori_pos_local',
                   list(util.quat_to_rpy(msg.task_rfoot_ori_pos_local)))
    data_saver.add('task_rfoot_ori_vel_local',
                   list(msg.task_rfoot_ori_vel_local))

    # data_saver.add('task_rfoot_rpy_pos_des_local', (util.quat_to_rpy(
    # msg.task_rfoot_ori_pos_des_local)).tolist())
    # data_saver.add('task_rfoot_rpy_pos_local',
    # (util.quat_to_rpy(msg.task_rfoot_ori_pos_local)).tolist())

    data_saver.add('task_lfoot_lin_pos_des', list(msg.task_lfoot_lin_pos_des))
    data_saver.add('task_lfoot_lin_vel_des', list(msg.task_lfoot_lin_vel_des))
    data_saver.add('task_lfoot_lin_acc_des', list(msg.task_lfoot_lin_acc_des))
    data_saver.add('task_lfoot_lin_pos', list(msg.task_lfoot_lin_pos))
    data_saver.add('task_lfoot_lin_vel', list(msg.task_lfoot_lin_vel))

    data_saver.add('task_lfoot_lin_pos_des_local',
                   list(msg.task_lfoot_lin_pos_des_local))
    data_saver.add('task_lfoot_lin_vel_des_local',
                   list(msg.task_lfoot_lin_vel_des_local))
    data_saver.add('task_lfoot_lin_acc_des_local',
                   list(msg.task_lfoot_lin_acc_des_local))
    data_saver.add('task_lfoot_lin_pos_local',
                   list(msg.task_lfoot_lin_pos_local))
    data_saver.add('task_lfoot_lin_vel_local',
                   list(msg.task_lfoot_lin_vel_local))

    data_saver.add('task_lfoot_ori_pos_des',
                   list(util.quat_to_rpy(msg.task_lfoot_ori_pos_des)))
    data_saver.add('task_lfoot_ori_vel_des', list(msg.task_lfoot_ori_vel_des))
    data_saver.add('task_lfoot_ori_acc_des', list(msg.task_lfoot_ori_acc_des))
    data_saver.add('task_lfoot_ori_pos',
                   list(util.quat_to_rpy(msg.task_lfoot_ori_pos)))
    data_saver.add('task_lfoot_ori_vel', list(msg.task_lfoot_ori_vel))

    # data_saver.add('task_lfoot_rpy_pos_des',
    # (util.quat_to_rpy(msg.task_lfoot_ori_pos_des)).tolist())
    # data_saver.add('task_lfoot_rpy_pos',
    # (util.quat_to_rpy(msg.task_lfoot_ori_pos)).tolist())

    data_saver.add('task_lfoot_ori_pos_des_local',
                   list(util.quat_to_rpy(msg.task_lfoot_ori_pos_des_local)))
    data_saver.add('task_lfoot_ori_vel_des_local',
                   list(msg.task_lfoot_ori_vel_des_local))
    data_saver.add('task_lfoot_ori_acc_des_local',
                   list(msg.task_lfoot_ori_acc_des_local))
    data_saver.add('task_lfoot_ori_pos_local',
                   list(util.quat_to_rpy(msg.task_lfoot_ori_pos_local)))
    data_saver.add('task_lfoot_ori_vel_local',
                   list(msg.task_lfoot_ori_vel_local))

    # data_saver.add('task_lfoot_rpy_pos_des_local', (util.quat_to_rpy(
    # msg.task_lfoot_ori_pos_des_local)).tolist())
    # data_saver.add('task_lfoot_rpy_pos_local',
    # (util.quat_to_rpy(msg.task_lfoot_ori_pos_local)).tolist())

    data_saver.add('task_upper_body_pos_des',
                   list(msg.task_upper_body_pos_des))
    data_saver.add('task_upper_body_vel_des',
                   list(msg.task_upper_body_vel_des))
    data_saver.add('task_upper_body_acc_des',
                   list(msg.task_upper_body_acc_des))
    data_saver.add('task_upper_body_pos', list(msg.task_upper_body_pos))
    data_saver.add('task_upper_body_vel', list(msg.task_upper_body_vel))

    data_saver.add('task_cam_vel_des', list(msg.task_cam_vel_des))
    data_saver.add('task_cam_vel', list(msg.task_cam_vel))
    data_saver.add('task_cam_acc_des', list(msg.task_cam_vel_des))

    data_saver.add('cmd_lfoot_rf', list(msg.cmd_lfoot_rf))
    data_saver.add('cmd_rfoot_rf', list(msg.cmd_rfoot_rf))
    data_saver.add('cmd_joint_positions', list(msg.cmd_joint_positions))
    data_saver.add('cmd_joint_velocities', list(msg.cmd_joint_velocities))
    data_saver.add('cmd_joint_torques', list(msg.cmd_joint_torques))
    data_saver.add('cmd_joint_accelerations',
                   list(msg.cmd_joint_accelerations))
    data_saver.add('cmd_joint_accelerations_fb',
                   list(msg.cmd_joint_accelerations_fb))

    data_saver.add('joint_positions', list(msg.joint_positions))
    data_saver.add('joint_velocities', list(msg.joint_velocities))

    data_saver.add('com_vel_est_exp', list(msg.com_vel_est_exp))
    data_saver.add('com_vel_est', list(msg.com_vel_est))
    data_saver.add('com_vel_raw', list(msg.com_vel_raw))
    data_saver.add('imu_ang_vel_est', list(msg.imu_ang_vel_est))
    data_saver.add('imu_ang_vel_raw', list(msg.imu_ang_vel_raw))
    data_saver.add('cam_est', list(msg.cam_est))
    data_saver.add('cam_raw', list(msg.cam_raw))

    data_saver.add('icp_des', list(msg.icp_des))
    data_saver.add('icp', list(msg.icp))
    data_saver.add('icp_dot_des', list(msg.icp_dot_des))
    data_saver.add('icp_dot', list(msg.icp_dot))

    data_saver.add('base_joint_position', list(msg.base_joint_pos))
    data_saver.add('base_joint_quat',
                   list(msg.base_joint_quat))  # (w, x, y, z) order
    data_saver.add('imu_accel', list(msg.imu_accel))

    data_saver.add('des_cmp', list(msg.des_cmp))

    data_saver.add('base_joint_lin_vel', list(msg.base_joint_lin_vel))

    ##hand
    data_saver.add('task_rhand_lin_pos_des', list(msg.task_rhand_lin_pos_des))
    data_saver.add('task_rhand_lin_vel_des', list(msg.task_rhand_lin_vel_des))
    data_saver.add('task_rhand_lin_acc_des', list(msg.task_rhand_lin_acc_des))
    data_saver.add('task_rhand_lin_pos', list(msg.task_rhand_lin_pos))
    data_saver.add('task_rhand_lin_vel', list(msg.task_rhand_lin_vel))

    data_saver.add('task_rhand_lin_pos_des_local',
                   list(msg.task_rhand_lin_pos_des_local))
    data_saver.add('task_rhand_lin_vel_des_local',
                   list(msg.task_rhand_lin_vel_des_local))
    data_saver.add('task_rhand_lin_acc_des_local',
                   list(msg.task_rhand_lin_acc_des_local))
    data_saver.add('task_rhand_lin_pos_local',
                   list(msg.task_rhand_lin_pos_local))
    data_saver.add('task_rhand_lin_vel_local',
                   list(msg.task_rhand_lin_vel_local))

    data_saver.add('task_rhand_ori_pos_des', list(msg.task_rhand_ori_pos_des))
    data_saver.add('task_rhand_ori_vel_des', list(msg.task_rhand_ori_vel_des))
    data_saver.add('task_rhand_ori_acc_des', list(msg.task_rhand_ori_acc_des))
    data_saver.add('task_rhand_ori_pos', list(msg.task_rhand_ori_pos))
    data_saver.add('task_rhand_ori_vel', list(msg.task_rhand_ori_vel))

    data_saver.add('task_rhand_rpy_pos_des',
                   (util.quat_to_rpy(msg.task_rhand_ori_pos_des)).tolist())
    data_saver.add('task_rhand_rpy_pos',
                   (util.quat_to_rpy(msg.task_rhand_ori_pos)).tolist())

    data_saver.add('task_rhand_ori_pos_des_local',
                   list(msg.task_rhand_ori_pos_des_local))
    data_saver.add('task_rhand_ori_vel_des_local',
                   list(msg.task_rhand_ori_vel_des_local))
    data_saver.add('task_rhand_ori_acc_des_local',
                   list(msg.task_rhand_ori_acc_des_local))
    data_saver.add('task_rhand_ori_pos_local',
                   list(msg.task_rhand_ori_pos_local))
    data_saver.add('task_rhand_ori_vel_local',
                   list(msg.task_rhand_ori_vel_local))

    data_saver.add('task_rhand_rpy_pos_des_local', (util.quat_to_rpy(
        msg.task_rhand_ori_pos_des_local)).tolist())
    data_saver.add('task_rhand_rpy_pos_local',
                   (util.quat_to_rpy(msg.task_rhand_ori_pos_local)).tolist())

    data_saver.add('task_lhand_lin_pos_des', list(msg.task_lhand_lin_pos_des))
    data_saver.add('task_lhand_lin_vel_des', list(msg.task_lhand_lin_vel_des))
    data_saver.add('task_lhand_lin_acc_des', list(msg.task_lhand_lin_acc_des))
    data_saver.add('task_lhand_lin_pos', list(msg.task_lhand_lin_pos))
    data_saver.add('task_lhand_lin_vel', list(msg.task_lhand_lin_vel))

    data_saver.add('task_lhand_lin_pos_des_local',
                   list(msg.task_lhand_lin_pos_des_local))
    data_saver.add('task_lhand_lin_vel_des_local',
                   list(msg.task_lhand_lin_vel_des_local))
    data_saver.add('task_lhand_lin_acc_des_local',
                   list(msg.task_lhand_lin_acc_des_local))
    data_saver.add('task_lhand_lin_pos_local',
                   list(msg.task_lhand_lin_pos_local))
    data_saver.add('task_lhand_lin_vel_local',
                   list(msg.task_lhand_lin_vel_local))

    data_saver.add('task_lhand_ori_pos_des', list(msg.task_lhand_ori_pos_des))
    data_saver.add('task_lhand_ori_vel_des', list(msg.task_lhand_ori_vel_des))
    data_saver.add('task_lhand_ori_acc_des', list(msg.task_lhand_ori_acc_des))
    data_saver.add('task_lhand_ori_pos', list(msg.task_lhand_ori_pos))
    data_saver.add('task_lhand_ori_vel', list(msg.task_lhand_ori_vel))

    data_saver.add('task_lhand_rpy_pos_des',
                   (util.quat_to_rpy(msg.task_lhand_ori_pos_des)).tolist())
    data_saver.add('task_lhand_rpy_pos',
                   (util.quat_to_rpy(msg.task_lhand_ori_pos)).tolist())

    data_saver.add('task_lhand_ori_pos_des_local',
                   list(msg.task_lhand_ori_pos_des_local))
    data_saver.add('task_lhand_ori_vel_des_local',
                   list(msg.task_lhand_ori_vel_des_local))
    data_saver.add('task_lhand_ori_acc_des_local',
                   list(msg.task_lhand_ori_acc_des_local))
    data_saver.add('task_lhand_ori_pos_local',
                   list(msg.task_lhand_ori_pos_local))
    data_saver.add('task_lhand_ori_vel_local',
                   list(msg.task_lhand_ori_vel_local))

    data_saver.add('task_lhand_rpy_pos_des_local', (util.quat_to_rpy(
        msg.task_lhand_ori_pos_des_local)).tolist())
    data_saver.add('task_lhand_rpy_pos_local',
                   (util.quat_to_rpy(msg.task_lhand_ori_pos_local)).tolist())

    data_saver.add('icp_err_integrator', list(msg.icp_err_integrator))
    data_saver.add('icp_err_leaky_integral', list(msg.icp_err_leaky_integral))

    data_saver.advance()

    # publish back for plot juggler
    pj_socket.send_string(json.dumps(data_saver.history, default=str))

    # publish joint positions for meshcat
    if args.b_visualize:
        vis_q[0:3] = np.array(msg.base_joint_pos)  # << base pos
        vis_q[3] = msg.base_joint_quat[1]  # << quaternion x
        vis_q[4] = msg.base_joint_quat[2]  # << quaternion y
        vis_q[5] = msg.base_joint_quat[3]  # << quaternion z
        vis_q[6] = msg.base_joint_quat[0]  # << quaternion w
        vis_q[7:] = np.array(msg.joint_positions)  # << joint pos

        icp_viz_q[0] = msg.icp[0]
        icp_viz_q[1] = msg.icp[1]
        icp_viz_q[2] = 0.

        icp_des_viz_q[0] = msg.icp_des[0]
        icp_des_viz_q[1] = msg.icp_des[1]
        icp_des_viz_q[2] = 0.

        cmp_des_q[0] = msg.des_cmp[0]
        cmp_des_q[1] = msg.des_cmp[1]
        cmp_des_q[2] = 0.

        viz.display(vis_q)
        icp_viz.display(icp_viz_q)
        icp_des_viz.display(icp_des_viz_q)
        cmp_viz.display(cmp_des_q)
