#! /usr/bin/env python
from argparse import ArgumentParser
from math import radians

import rospy
import tf
import numpy as np

from masterloop import inform_masterloop_factory, mp
from controller import nao_cart_movement, movement
from controller import dumb, our_cartesian, our_diff_cartesian


_inform_masterloop = inform_masterloop_factory('imitator')


TORSO = False


def controller_factory(ctrl):
    if ctrl == 'nao_cartesian':
        return lambda my_arm_xyz, side: nao_cart_movement(my_arm_xyz, side)

    if ctrl == 'dumb':
        cfunc = dumb
    elif ctrl == 'our_cartesian':
        cfunc = our_cartesian
    elif ctrl == 'our_diff_cartesian':
        cfunc = our_diff_cartesian
    return lambda my_arm_xyz, side: movement(my_arm_xyz, side, cfunc)


if __name__ == '__main__':
    rospy.init_node('imitator')
    ap = ArgumentParser()
    ap.add_argument('--controller', default='dumb', const='dumb', nargs='?',
                    choices=['dumb', 'nao_cartesian', 'our_cartesian',
                             'our_diff_cartesian'],
                    help='Choose the controller for arm motions')
    args, _ = ap.parse_known_args()
    imitation_cycle = controller_factory(args.controller)

    rospy.wait_for_service('inform_masterloop')

    ll = tf.TransformListener()

    while not rospy.is_shutdown():
        rospy.Rate(20).sleep()
        if not _inform_masterloop('imitate'):
            continue
        rospy.logdebug('IMITATOR: ACTIVE')

        # if TORSO:
            # try:
                # _, q = ll.lookupTransform('odom',
                                          # 'Aruco_0_frame',
                                          # rospy.Time(0))
                # rot = tf.transformations.euler_from_quaternion(q)
                # mp.setAngles(['LHipYawPitch', 'RHipYawPitch'],
                             # [-rot[1], -rot[1]], 0.3)
            # except Exception as e:
                # rospy.logwarn(e)

        for i, side in enumerate(['L', 'R'], 1):
            try:
                a0, _ = ll.lookupTransform(
                    'odom',
                    'Aruco_0_frame',
                    rospy.Time(0)
                )
                arm, _ = ll.lookupTransform(
                    'odom',
                    'Aruco_{}_frame'.format(i),
                    rospy.Time(0)
                )
            except Exception as e:
                rospy.logwarn(e)
                continue
            my_arm_xyz = np.array(arm) - np.array(a0)
            # rospy.loginfo('{}'.format(my_arm_xyz))
            # rospy.loginfo('{}'.format(dumb(my_arm_xyz, side)))
            imitation_cycle(my_arm_xyz, side)
            mp.setAngles(('LWristYaw', 'RWristYaw'),
                         (radians(-70), radians(70)), 0.3)