#!/usr/bin/env python

# ##########################################################################
# Copyright (c) 2014 Shadow Robot Company Ltd.
#
# This program is free software: you can redistribute it and/or modify it
# under the terms of the GNU General Public License as published by the Free
# Software Foundation, either version 2 of the License, or (at your option)
# any later version.
#
# This program is distributed in the hope that it will be useful, but WITHOUT
# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
# FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for
# more details.
#
# You should have received a copy of the GNU General Public License along
# with this program. If not, see <http://www.gnu.org/licenses/>.
#
# ###########################################################################

import os, unittest;
import rospy, rostest
from math import radians
from actionlib import SimpleActionClient, GoalStatus
from sr_robot_msgs.msg import GraspAction, GraspGoal
from sr_hand.shadowhand_ros import ShadowHand_ROS
from sr_grasp.utils import mk_grasp
from sr_grasp import Grasp, GraspStash

class QuickGraspNode(object):
    def __init__(self):
        self.stash = GraspStash()
        #self.stash.load_yaml_file("grasps.yaml")
        self.stash.load_all()
        if self.stash.size() == 0:
            rospy.logerr("No grasps found!")
            raise Exception("No Grasps")
        else:
            self.grasp = self.stash.get_all()[0]
        #print self.grasp
        rospy.loginfo("Looking for hand...")
        self.hand = ShadowHand_ROS()
        rospy.loginfo("Found")

    def zero_hand(self):
        for j in self.hand.allJoints:
            self.hand.sendupdate(j.name, 0.0)
        rospy.sleep(2)

    def run_grasp(self, pre=False):
        goal = GraspGoal()
        goal.grasp = self.grasp
        goal.pre_grasp = pre
        client = SimpleActionClient('grasp', GraspAction)
        client.wait_for_server()
        client.send_goal(goal)
        client.wait_for_result(rospy.Duration.from_sec(20.0))
        if client.get_state() == GoalStatus.SUCCEEDED:
            rospy.loginfo("Success!")
        else:
            rospy.logerr("Fail!")

    def new_grasp(self, name):
        grasp = Grasp()
        grasp.id = name
        positions = self.hand.read_all_current_positions()
        for k, v in positions.iteritems():
            positions[k] = radians(v)
        grasp.set_grasp_point(positions = positions)
        self.stash.put_grasp(grasp)
        node.grasp = grasp

if __name__ == '__main__':
    rospy.init_node("quick_grasp")
    node = QuickGraspNode()
    while node:
        print ""
        print "Grasps:"
        for i, g in enumerate(node.stash.get_all()):
            print str(i) + " - " + g.id
        print "Current grasp: %s"%node.grasp.id
        print "Number select grasp, z zero hand, g grasp, p pre-grasp, n new grasp, s save grasps, q quit"
        act = raw_input(": ")
        num = None
        try:
            num = int(act)
        except ValueError:
            pass
        if num is not None:
            node.grasp = node.stash.get_grasp_at(num)
            act = 'p' # auto go to pre grasp on grasp select

        if act == "z":
            node.zero_hand()
        elif act == "p":
            node.run_grasp(pre=True)
        elif act == "g":
            node.run_grasp()
        elif act == "n":
            name = raw_input("Enter name: ")
            node.new_grasp(name=name)
        elif act == "s":
            node.stash.save_yaml_file()
        elif act == "q":
            break
        else:
            print "Unknown action %s"%act

