#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Thu Oct 12 16:11:58 2023

@author: Pierre Lecomte
"""
from interbotix_xs_modules.xs_robot.arm import InterbotixManipulatorXS
import numpy as np
import nidaq

#%%########################################################################################################
# ROBOTIC ARM PARAMETERS
bot = InterbotixManipulatorXS(
        robot_model='wx250s_custom',
        group_name='arm',
        gripper_name = None
    )


# The arm is going to the initial pose
bot.arm.go_to_home_pose()
bot.arm.set_ee_pose_components(x=0,y=0,z=0.5)

# HP coordinates in EE frame are set here. One should align the directions with a laser
hp = np.array([1, 0, 0]) # in EE frame


# EE frame to SMA frame rotation
thetaA, phiA, psiA = 91.5*np.pi/180, -45*np.pi/180, 0

rotEA = np.array([
    [-np.cos(phiA) * np.sin(thetaA), np.cos(psiA) * np.sin(phiA) - np.cos(thetaA) * np.sin(psiA), -np.cos(thetaA) * np.cos(phiA) * np.cos(psiA) - np.sin(phiA) * np.sin(psiA)],
    [-np.sin(thetaA) * np.sin(phiA), -np.cos(phiA) * np.cos(psiA) - np.cos(thetaA) * np.sin(phiA) * np.sin(psiA), - np.cos(thetaA) * np.cos(psiA) * np.sin(phiA) + np.cos(phiA) * np.sin(psiA)],
    [-np.cos(thetaA), np.sin(thetaA) * np.sin(psiA), np.cos(psiA) * np.sin(thetaA)]
    ])

#%% xyzA: DOA to be measured in Cartesian Coordinates in SMA frame
xyzA = np.array([[0, 0, 1], [1, 0, 0], [0, 1, 0], [-1*1, 0, 0], [0, -1*1, 0], [0, 
  0, -1*1]])#, [1/np.sqrt(3), 1/np.sqrt(3), 1/np.sqrt(3)], [-1*(1/np.sqrt(3)), 1/np.sqrt(3),
  #  1/np.sqrt(3)], [-1*(1/np.sqrt(3)), -1*(1/np.sqrt(3)), 1/np.sqrt(3)], [1/np.sqrt(
  # 3), -1*(1/np.sqrt(3)), 1/np.sqrt(3)], [1/np.sqrt(3), 1/np.sqrt(
  # 3), -1*(1/np.sqrt(3))], [-1*(1/np.sqrt(3)), 1/np.sqrt(
  # 3), -1*(1/np.sqrt(3))], [-1*(1/np.sqrt(3)), -1*(1/np.sqrt(3)), -1*(1/np.sqrt(3))], [1/
  # np.sqrt(3), -1*(1/np.sqrt(3)), -1*(1/np.sqrt(3))], [1/np.sqrt(2), 0, 1/np.sqrt(2)], [0,
  #  1/np.sqrt(2), 1/np.sqrt(2)], [-1*(1/np.sqrt(2)), 0, 1/np.sqrt(
  # 2)], [0, -1*(1/np.sqrt(2)), 1/np.sqrt(2)], [1/np.sqrt(2), 1/np.sqrt(2), 
  # 0], [-1*(1/np.sqrt(2)), 1/np.sqrt(2), 0], [-1*(1/np.sqrt(2)), -1*(1/np.sqrt(2)), 
  # 0], [1/np.sqrt(2), -1*(1/np.sqrt(2)), 0], [1/np.sqrt(2), 0, -1*(1/np.sqrt(2))], [0, 
  # 1/np.sqrt(2), -1*(1/np.sqrt(2))], [-1*(1/np.sqrt(2)), 
  # 0, -1*(1/np.sqrt(2))], [0, -1*(1/np.sqrt(2)), -1*(1/np.sqrt(2))], [1/np.sqrt(11), 1/
  # np.sqrt(11), 3/np.sqrt(11)], [-1*(1/np.sqrt(11)), 1/np.sqrt(11), 3/np.sqrt(
  # 11)], [-1*(1/np.sqrt(11)), -1*(1/np.sqrt(11)), 3/np.sqrt(11)], [1/np.sqrt(
  # 11), -1*(1/np.sqrt(11)), 3/np.sqrt(11)], [3/np.sqrt(11), 1/np.sqrt(11), 1/np.sqrt(
  # 11)], [1/np.sqrt(11), 3/np.sqrt(11), 1/np.sqrt(11)], [-1*(1/np.sqrt(11)), 3/np.sqrt(
  # 11), 1/np.sqrt(11)], [-1*(3/np.sqrt(11)), 1/np.sqrt(11), 1/np.sqrt(
  # 11)], [-1*(3/np.sqrt(11)), -1*(1/np.sqrt(11)), 1/np.sqrt(
  # 11)], [-1*(1/np.sqrt(11)), -1*(3/np.sqrt(11)), 1/np.sqrt(11)], [1/np.sqrt(
  # 11), -1*(3/np.sqrt(11)), 1/np.sqrt(11)], [3/np.sqrt(11), -1*(1/np.sqrt(11)), 1/np.sqrt(
  # 11)], [3/np.sqrt(11), 1/np.sqrt(11), -1*(1/np.sqrt(11))], [1/np.sqrt(11), 3/np.sqrt(
  # 11), -1*(1/np.sqrt(11))], [-1*(1/np.sqrt(11)), 3/np.sqrt(
  # 11), -1*(1/np.sqrt(11))], [-1*(3/np.sqrt(11)), 1/np.sqrt(
  # 11), -1*(1/np.sqrt(11))], [-1*(3/np.sqrt(11)), -1*(1/np.sqrt(11)), -1*(1/np.sqrt(
  #  11))], [-1*(1/np.sqrt(11)), -1*(3/np.sqrt(11)), -1*(1/np.sqrt(11))], [1/np.sqrt(
  # 11), -1*(3/np.sqrt(11)), -1*(1/np.sqrt(11))], [3/np.sqrt(
  # 11), -1*(1/np.sqrt(11)), -1*(1/np.sqrt(11))], [1/np.sqrt(11), 1/np.sqrt(
  # 11), -1*(3/np.sqrt(11))], [-1*(1/np.sqrt(11)), 1/np.sqrt(
  # 11), -1*(3/np.sqrt(11))], [-1*(1/np.sqrt(11)), -1*(1/np.sqrt(11)), -1*(3/np.sqrt(
  #  11))], [1/np.sqrt(11), -1*(1/np.sqrt(11)), -1*(3/np.sqrt(11))]])
 
#%% xyzE: DOA to be measured in Cartesian Coordinates in EE frame
xyzE = np.linalg.inv(rotEA) @ np.transpose(xyzA)


sr = 44100
duration = 8
t = np.linspace(0, duration, sr*duration, endpoint=False)
sig = nidaq.sine_sweep(sr,duration)

for i in range(xyzE.shape[1]):
    xyzEi = xyzE[:,i]
    print(["DOA: " + str(xyzA[i,:])])
    gamma = np.arctan2(xyzEi[2], xyzEi[1])
    k =np.array([
        [1, 0, 0],
        [0, np.cos(gamma), np.sin(gamma)],
        [0, -np.sin(gamma), np.cos(gamma)]
        ]) @ xyzEi
    alpha = np.arccos(np.dot(k, hp/np.linalg.norm(hp)))
    bot.arm.set_single_joint_position(joint_name='wrist_rotate', position=-gamma)
    bot.arm.set_single_joint_position(joint_name='waist', position=-alpha)
    out = np.array(nidaq.playrec(sig, sr))
    np.save(str(i) + '.npy', out)
    
    print("Done")