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

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."""
    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')