#! /usr/bin/env python
"""The script for calibrating the arm lengths and 'joystick' extents.

You will normally not call this script directly but rather with
`roslaunch teleoperation calibration.launch`
"""

import os
import json

import rospy
import tf
import actionlib
import numpy as np
from naoqi import ALProxy

from teleoperation.msg import RequestSpeechAction, RequestSpeechGoal


def calib_field(sentences, src_frames, targ_frames):
    """Retrieve a relative transform of some marker.

    Parameters
    ----------

    sentences : list of str
        What to say.
    src_frames : list of str
        The names of tf source frames (from where to transform)
    targ_frames : list of str
        The names of tf target frames (to where to transform)

    Returns
    -------
    list of transforms

    """
    ok = False
    t_and_q = []

    while not ok:
        for s in sentences:
            tts.say(s)
        tts.say('Then say start')
        client.send_goal(RequestSpeechGoal(['start']))
        client.wait_for_result()
        for src, targ in zip(src_frames, targ_frames):
            try:
                t_and_q.append(ll.lookupTransform(src, targ, rospy.Time(0)))
                ok = True
            except:
                ok = False
                tts.say('Step failed. Try again')
                break

    return t_and_q


def calib_center():
    """Retrieve the coordinates of the 'joystick' zero state."""
    trans, q = calib_field([
        'Stand in front of camera',
        'Far enough to see your arms'
    ], ['odom'], ['Aruco_0_frame'])[0]
    rot = tf.transformations.euler_from_quaternion(q)
    return trans, rot


def calib_arms():
    """Retrieve the hand positions relative to the chest."""
    ts_and_qs = calib_field(
        ['Now stretch your arms', 'Make sure the markers are detected'],
        ['odom'] * 3,
        ['Aruco_{}_frame'.format(i) for i in range(3)]
    )
    a0 = ts_and_qs[0][0]
    a1 = ts_and_qs[1][0]
    a2 = ts_and_qs[2][0]
    return (
        tuple(c1 - c0 for c0, c1 in zip(a0, a1)),
        tuple(c2 - c0 for c0, c2 in zip(a0, a2))
    )


def calib_shoulders():
    """Retrieve the shoulder positions relative to the chest."""
    ts_and_qs = calib_field(
        ['Now place the markers on the corresponding shoulders',
         'Make sure the markers are detected'],
        ['odom'] * 3,
        ['Aruco_{}_frame'.format(i) for i in range(3)]
    )
    a0 = ts_and_qs[0][0]
    a1 = ts_and_qs[1][0]
    a2 = ts_and_qs[2][0]
    return (
        tuple(c1 - c0 for c0, c1 in zip(a0, a1)),
        tuple(c2 - c0 for c0, c2 in zip(a0, a2))
    )


def calib_rotation():
    """Retrieve the limits of the z-axis of the 'joystick'."""
    rots = []
    for side in ('left', 'right'):
        _, q = calib_field(
            ['Now rotate to your {}'.format(side),
             'So that marker on the chest is still detected'],
            ['odom'], ['Aruco_0_frame']
        )[0]
        rots.append(tf.transformations.euler_from_quaternion(q))
    return tuple(rots)


def calib_extremes():
    """Retrieve the limits of the x and y axes of the 'joystick'."""
    transs = []
    for direction in ('forward', 'backward', 'to the left', 'to the right'):
         trans, _ = calib_field(
            ['Now return to the initial position',
             'Take a big step {}'.format(direction)],
            ['odom'],
            ['Aruco_0_frame']
         )[0]
         transs.append(trans)

    return tuple(transs)


def calibration():
    """Run full calibration and dump the results to the config file."""
    center = calib_center()
    larm, rarm = calib_arms()
    lsh, rsh = calib_shoulders()
    larm, rarm, lsh, rsh = larm[1:], rarm[1:], lsh[1:], rsh[1:]
    lr, rr = calib_rotation()
    fw, bk, lt, rt = calib_extremes()

    tts.say('Well done')

    calib = {
        'cr': [center[0][0], center[0][1], center[1][2]],
        'arm': (np.linalg.norm(np.array(larm) - np.array(lsh)) +
                np.linalg.norm(np.array(rarm) - np.array(rsh))) / 2,
        'lsh': [0.0] + list(lsh),
        'rsh': [0.0] + list(rsh),
        'rr': rr[2],
        'lr': lr[2],
        'fw': fw[0],
        'bk': bk[0],
        'lt': lt[1],
        'rt': rt[1],
    }

    here = os.path.dirname(os.path.realpath(__file__))
    with open('{}/../config/default.yaml'.format(here), 'w') as f:
        json.dump(calib, f, indent=4)


if __name__ == '__main__':

    rospy.init_node('calibrator', disable_signals=True)

    try:
        tts = ALProxy('ALTextToSpeech', os.environ['NAO_IP'], 9559)
        client = actionlib.SimpleActionClient('speech_server',
                                              RequestSpeechAction)
        client.wait_for_server()

        ll = tf.TransformListener()

        rospy.on_shutdown(lambda: client.cancel_goal())

        calibration()

    except KeyboardInterrupt:
        rospy.signal_shutdown('Interrupted')