#!/usr/bin/env python
from __future__ import print_function

import pyexotica as exo
import numpy as np
import math
from pyexotica.publish_trajectory import *
import exotica_core_task_maps_py
import exotica_scipy_solver
from time import sleep, time
import signal


def figure_eight(t):
    return np.array([0.6, -0.1 + math.sin(t * 2.0 * math.pi * 0.5) * 0.1, 0.5 + math.sin(t * math.pi * 0.5) * 0.2])


use_scipy_solver = True

exo.Setup.init_ros()

problem = None
solver = None
if use_scipy_solver:
    problem = exo.Setup.load_problem('{exotica_examples}/resources/configs/example_ik_joint_torque_minimization_proxy.xml')
    solver = exotica_scipy_solver.SciPyEndPoseSolver(problem, method='trust-constr')
    solver.max_iterations = 10
    solver.debug = False
else:
    # Requires to uncomment one of the C++ solvers in the XML
    solver = exo.Setup.load_solver(
        '{exotica_examples}/resources/configs/example_ik_joint_torque_minimization_proxy.xml')
    problem = solver.get_problem()

# Get the task so we can modify the h
joint_torque_minimization_proxy_task = problem.get_task_maps()["JointTorqueMinimizationProxy"]

dt = 0.002
t = 0.0
q = np.array([0.2]*7)
print('Publishing IK')
signal.signal(signal.SIGINT, sig_int_handler)
while True:
    try:
        start_time = time()

        # Update h
        joint_torque_minimization_proxy_task.h = np.array([1,1,1,0,0,0])

        # Update goal position for end-effector
        problem.set_goal_eq('Position', figure_eight(t))

        # Set last solution as initialisation for optimisation problem
        problem.start_state = q

        # Solve the problem
        s = time()
        q = solver.solve()[0]
        e = time()

        # Extract the positions of the individual links
        problem.update(q)
        for i in range(7):
            name = "lwr_arm_" + str(i) + "_link"
            position = problem.get_scene().fk(name).get_translation()
            print(name, position)

        # Publish to RViz (visualisation)
        publish_pose(q, problem)

        # Sleep if there is remaining time in the control loop
        end_time = time()
        if dt > end_time - start_time:
            sleep(dt - (end_time - start_time))

        print("Optimisation took", e-s, "\tTotal loop time:", end_time - start_time)
        t = t + dt
    except KeyboardInterrupt:
        break
