import numpy as np
import torch
import pinocchio as pin
from common.np_math import (quat_from_angle_axis, quat_mul,
                            xyzw2wxyz, index_map)

class ActToDof:
    def __call__(self, obs, action):
        hands_command = obs[..., 119:225]
        non_arm_joint_pos = action[..., :22]
        left_arm_residual = action[..., 22:29]

        q_pin = np.zeros_like(self.ikctrl.cfg.q)
        for i_mot in range(len(self.config.motor_joint)):
            i_pin = self.pin_from_mot[i_mot]
            q_pin[i_pin] = self.low_state.motor_state[i_mot].q

        d_quat = quat_from_angle_axis(
            torch.from_numpy(hands_command[..., 3:])
        ).detach().cpu().numpy()

        source_pose = self.ikctrl.fk(q_pin)
        source_xyz = source_pose.translation
        source_quat = xyzw2wxyz(pin.Quaternion(source_pose.rotation).coeffs())

        target_xyz = source_xyz + hands_command[..., :3]
        target_quat = quat_mul(
            torch.from_numpy(d_quat),
            torch.from_numpy(source_quat)).detach().cpu().numpy()
        target = np.concatenate([target_xyz, target_quat])

        res_q_ik = self.ikctrl(
            q_pin,
            target
        )
        print('res_q_ik', res_q_ik)

        target_dof_pos = np.zeros(29)
        for i_arm in range(len(res_q_ik)):
            i_mot = self.mot_from_arm[i_arm]
            i_pin = self.pin_from_mot[i_mot]
            target_q = (
                self.low_state.motor_state[i_mot].q
                + res_q_ik[i_arm]
                + np.clip(0.3 * left_arm_residual[i_arm],
                          -0.2, 0.2)
            )
            target_q = np.clip(target_q,
                               self.lim_lo_pin[i_pin],
                               self.lim_hi_pin[i_pin])
            target_dof_pos[i_mot] = target_q
        return target_dof_pos

def main():
    import yaml

    with open('configs/g1_eetrack.yaml', 'r') as fp:
        d = yaml.safe_load(fp)

    act_to_dof = ActToDof()
    obs = np.load('/tmp/eet4/obs001.npy')[0]
    act = np.load('/tmp/eet4/act001.npy')[0]
    dof_lab = np.load('/tmp/eet4/dof001.npy')[0]
    mot_from_lab = index_map(d['motor_joint'], d['lab_joint'])
    target_dof_pos = np.zeros_like(dof_lab)
    target_dof_pos[mot_from_lab] = dof_lab


    dof = act_to_dof(obs, act)

if __name__ == '__main__':
    main()