#!/usr/bin/env python
# -*- coding: utf-8 -*-

import sys
import os
import time
import copy

from math import *

from set_env import *
from armpos_data import *

save_path = sys.path[:]
sys.path.append(Rh2LibDir)
from rtc_handle import *
sys.path.append(Rh2LibDir+'/tools/geo')
from geo import *

sys.path.append(Rh2ARToolsDir)
import Rh2ARrecog_ctrl
sys.path.append(Rh2ArmToolsDir)
import Rh2Arm_ctrl

# 把持したドリンク（缶）の中心が手先の移動目標座標と一致するようオフセットがいる
offx =  3.0
offy = -27.5

# ドリンクの高さ
drink_hight = 120.0

# ドリンク設置（荷台）位置（ベース座標系）
set_pos = list()
set_pos.append([ pcan1[0]+offx, pcan1[1]+offy,  135.5, 0.0, 180.0, 0.0])   # can1 マーカID:0
set_pos.append([ pcan2[0]+offx, pcan2[1]+offy,  135.5, 0.0, 180.0, 0.0])   # can2 マーカID:1
set_pos.append([ pcan3[0]+offx, pcan3[1]+offy,  135.5, 0.0, 180.0, 0.0])   # can3 マーカID:2


# マーカ検索用の位置・姿勢
G_search_point1   = [0.0, 350.0, 300.0, 0.0, 180.0, -90.0]
G_search_point2   = [-50.0, 350.0, 300.0, 0.0, 180.0, -90.0]
G_search_list     = [G_search_point1, G_search_point2]
ABC_search_point1 = [0.0, 280.0, 350.0, 0.0, 180.0, -90.0]
ABC_search_point2 = [-120.0, 280.0, 350.0, 0.0, 180.0, -90.0]
ABC_search_list   = [ABC_search_point1, ABC_search_point2]
# ドリンク設置（テーブル）位置：目印ARマーカ（G）からのオフセット（目印ARマーカ座標系）
put_pos = list()
for i in range(3):
  put_pos.append(VECTOR())
#   マーカA
put_pos[0][0] = -100.0       # x
put_pos[0][1] =  100.0       # y 
put_pos[0][2] =   0.0        # z
#   マーカB
put_pos[1][0] =   0.0        # x
put_pos[1][1] =  100.0       # y 
put_pos[1][2] =   0.0        # z
#   マーカC
put_pos[2][0] =  100.0       # x
put_pos[2][1] =  100.0       # y 
put_pos[2][2] =   0.0        # z

# 6軸から取り付けたカメラまでのオフセット（手先座標系）
cam_offset = [59.5, 26.5, -95]

# AR認識で取得したマーカ座標に対するオフセット
#off_AR = [-35, 55, 0.0]
off_AR = [-25.0, 15.0, 55.0]

# ドリンク（または空き缶）を把持する際、go_to()に渡すオフセット
go_to_offset = [0.0, 0.0, 100.0]


def deactivate():
  Rh2ARrecog_ctrl.ARMarkerDetectSimple.deactivate()
  Rh2ARrecog_ctrl.CaptureTimedCameraImage.deactivate()
  Rh2Arm_ctrl.deactivate_all()

#
# ARマーカ（G）認識
#
def marker_search():
  # ARマーカ（G）認識結果取得
  Rh2ARrecog_ctrl.start_ARrecog(Rh2ARrecog_ctrl.env)
  for i in range(3):
    time.sleep(2)
    ARflg = Rh2ARrecog_ctrl.get_specified_data()
    if(not ARflg):
      break
  Rh2ARrecog_ctrl.stop_ARrecog(Rh2ARrecog_ctrl.env)
  # マーカが検出できなかった場合
  if(ARflg == 1):
    print"Marker was not detected!"
  return ARflg

#
# ドリンク（空き缶）検出
#   pos_list：ドリンク（空き缶）設置位置のリスト
#   drink_pos：ドリンク（空き缶）が実際に設置されている位置のみを格納したリスト
def search_drink(pos_list):
  # ARマーカ(A,B,C)認識結果取得
  # 挿入口4(マーカD)は使用しない
  Rh2ARrecog_ctrl.start_ARrecog(Rh2ARrecog_ctrl.env)
  time.sleep(2)
  Rh2ARrecog_ctrl.get_all_data()
  Rh2ARrecog_ctrl.stop_ARrecog(Rh2ARrecog_ctrl.env)
  # ドリンク設置状況の判断（マーカを認識した場所には置かれていない）
  #   認識マーカIDの抽出
  # ドリンク設置情報作成
  #   pos = [[ドリンク1の有無,ドリンク1の位置],[ドリンク2の有無,],[]]  
  pos = [[1], [1], [1]]   # 1:ドリンクあり 0:ドリンクなし
  # マーカA,B,Cの座標をコピー
  for i in range(3):
    pos[i].append(copy.deepcopy(pos_list[i]))
  print"--- recognized ID ----"
  for i in range(len(Rh2ARrecog_ctrl.all_data)):
    if(Rh2ARrecog_ctrl.all_data[i].id < 3):
      print("ID = %d" % Rh2ARrecog_ctrl.all_data[i].id)
      pos[Rh2ARrecog_ctrl.all_data[i].id][0] = 0
  print"---- marker info -----"
  markdic = ["A", "B", "C"]
  for i in range(len(pos)):
    print(markdic[i]),
    print(pos[i])
  # マーカが認識できた位置は移動先から除外
  drink_pos = list()
  for i in range(3):
    if(pos[i][0] != 0):
      drink_pos.append(copy.deepcopy(pos[i][1]))
  print"---- drink position -----"
  for i in range(len(drink_pos)):
    print(drink_pos[i])
  return drink_pos

#
# ドリンク設置または回収のためのアーム移動先座標を取得する
#   ドリンク設置位置は、最初マーカからオフセット(x,y,z)の位置と考え、
#   これを最終的にベース座標系における(X,Y,Z)に変換して返す
#   put_pos：目印ARマーカからのオフセット（目印ARマーカ座標系）
#   戻り値はput_posをベース座標系へ変換したもの
def calcPosition(put_pos, search_point):
  # ARマーカ座標（カメラ座標系）
  MarkerPos = VECTOR()
  MarkerPos[0] = Rh2ARrecog_ctrl.data[3]     # MarkerPos[0]: x
  MarkerPos[1] = Rh2ARrecog_ctrl.data[7]     # MarkerPos[1]: y
  MarkerPos[2] = Rh2ARrecog_ctrl.data[11]    # MarkerPos[2]: z
  # ARマーカ姿勢（カメラ座標系に対して）
  MarkerAxis = MATRIX()
  for i in range(3):
    MarkerAxis[i][0] = Rh2ARrecog_ctrl.data[4*i]
    MarkerAxis[i][1] = Rh2ARrecog_ctrl.data[4*i+1]
    MarkerAxis[i][2] = Rh2ARrecog_ctrl.data[4*i+2]
  print"---- Marker Axis  &  Pos -----"
  for i in range(3):
    print" %.3f %.3f %.3f" % (MarkerAxis[i][0], MarkerAxis[i][1], MarkerAxis[i][2]), 
    print" & %.3f" % (MarkerPos[i])
  # ARマーカ座標系 → カメラ座標系
  put_pos_cam = list()
  for i in range(len(put_pos)):
    put_pos_cam.append(VECTOR())
    put_pos_cam[i] = MarkerAxis*put_pos[i]
  for i in range(len(put_pos)):
    print""
    print"put_pos from marker[%d] = " % (i),
    print(put_pos_cam[i])
  for i in range(len(put_pos)):
    for j in range(3):
      put_pos_cam[i][j] = MarkerPos[j] + put_pos_cam[i][j]
  # カメラ座標系 → 手先座標系
  #   手先から設置カメラまでのオフセット、
  #   把持したドリンクと設置目標座標を一致させるためのオフセットを加味
  put_pos_tool = list()
  for i in range(len(put_pos)):
    put_pos_tool.append(VECTOR())
    put_pos_tool[i][0] =  put_pos_cam[i][1]
    put_pos_tool[i][1] = -put_pos_cam[i][0]
    put_pos_tool[i][2] =  put_pos_cam[i][2]
    for j in range(3):
      # 設置カメラのオフセットを加味
      put_pos_tool[i][j] = put_pos_tool[i][j] - cam_offset[j]
      # AR認識で取得したマーカ座標に対するオフセットを加味
      put_pos_tool[i][j] = put_pos_tool[i][j] - off_AR[j]
  print""
  print"put_pos_tool"
  for i in range(len(put_pos)):
    print"put_pos_tool[%d] = " % (i),
    print(put_pos_tool[i])
  # 設置位置のベース座標系への座標変換
  #   手先座標系 → ベース座標系
  rot_x = MATRIX(a=(search_point[3]*pi/180.0))
  rot_y = MATRIX(b=(search_point[4]*pi/180.0))
  rot_z = MATRIX(c=(search_point[5]*pi/180.0))
  put_pos_base = list()
  for i in range(len(put_pos)):
    put_pos_base.append(VECTOR())
    put_pos_base[i] = rot_z*put_pos_tool[i] 
    put_pos_base[i] = rot_y*put_pos_base[i] 
    put_pos_base[i] = rot_x*put_pos_base[i] 
  for i in range(len(put_pos)):
    for j in range(3):
      # 手先ツール長設定した分だけ補正
      put_pos_base[i][j] = search_point[j] + put_pos_base[i][j] + Rh2Arm_ctrl.tool_data[j]
      if(j == 2):
        # ドリンク把持位置（高さ）分だけ補正
        put_pos_base[i][j] = put_pos_base[i][j] + drink_hight
  print""
  print"put_pos_base"
  for i in range(len(put_pos)):
    print"put_pos_base[%d] = " % (i),
    print(put_pos_base[i])
  # listに変換してから返す
  put_pos_base_list = list()
  for i in range(len(put_pos)):
    put_pos_base_list.append(list())
    for j in range(3):
      put_pos_base_list[i].append(put_pos_base[i][j]) 
  print""
  print"put_pos_base_list"
  for i in range(len(put_pos)):
    print"put_pos_base_list[%d] = " % (i),
    print(put_pos_base_list[i])
  return put_pos_base_list

# 
# 終了関数
#
def end():
  # アーム終了
  Rh2Arm_ctrl.arm_release(Rh2Arm_ctrl.env)
  Rh2ARrecog_ctrl.end(Rh2ARrecog_ctrl.env)

################################
#      空き缶の回収シナリオ      #
################################

# テーブル上目印ARマーカの検出
#   目印ARマーカ(G)認識位置へアームを移動
Rh2Arm_ctrl.gripper_open()
Rh2Arm_ctrl.set_joint_vel(0.8)
Rh2Arm_ctrl.move_joint(Rh2Arm_ctrl.j_table1)
#  目印ARマーカ認識
search_p = 0
ABC_search_point = 0
for i in range(len(G_search_list)):
  global search_p, ABC_search_point
  Rh2Arm_ctrl.go_to(G_search_list[i])
  marker_flg = marker_search()
  if(marker_flg == 0):
    # 目印ARマーカ認識を行ったカメラの座標を記憶
    search_p = G_search_list[i]
    ABC_search_point = ABC_search_list[i]
    break
if(search_p == 0):
  sys.exit("Marker(G) was not detected!")
#  ドリンク設置位置（テーブル）の算出
put_pos_list = calcPosition(put_pos, search_p)
attitude = [0.0, 180.0, -90.0]
print""
for j in range(len(put_pos_list)):
  put_pos_list[j] = put_pos_list[j] + attitude
  print"put_pos_list[%d] = " % j,
  print(put_pos_list[j])
# テーブル上の空き缶検出（空き缶検出用ARマーカ（A,B,C）の検出）
print"ABC_search_point = ",
print(ABC_search_point)
Rh2Arm_ctrl.go_to(ABC_search_point)
can_pos_list = search_drink(put_pos_list)
# 空き缶の回収
for i in range(len(can_pos_list)):
  if(i != 0):
    Rh2Arm_ctrl.move_joint(Rh2Arm_ctrl.j_table2)
  # 空き缶を把持
  Rh2Arm_ctrl.go_to(Rh2Arm_ctrl.offset(can_pos_list[i], go_to_offset[0], go_to_offset[1], go_to_offset[2]))
  Rh2Arm_ctrl.go_to(can_pos_list[i])
  Rh2Arm_ctrl.gripper_close()
  time.sleep(4)
  Rh2Arm_ctrl.go_to(Rh2Arm_ctrl.offset(can_pos_list[i], go_to_offset[0], go_to_offset[1], go_to_offset[2]))
  # 空き缶を荷台ホルダーへ積載
  Rh2Arm_ctrl.move_joint(Rh2Arm_ctrl.j_holder)
  Rh2Arm_ctrl.go_to(Rh2Arm_ctrl.offset(set_pos[i], go_to_offset[0], go_to_offset[1], go_to_offset[2]))
  Rh2Arm_ctrl.go_to(set_pos[i])
  Rh2Arm_ctrl.gripper_open()
  time.sleep(8)
  Rh2Arm_ctrl.go_to(Rh2Arm_ctrl.offset(set_pos[i], go_to_offset[0], go_to_offset[1], go_to_offset[2]))
# アームを走行姿勢へ移行
Rh2Arm_ctrl.gripper_close()
Rh2Arm_ctrl.set_joint_vel(0.5)
Rh2Arm_ctrl.move_joint(Rh2Arm_ctrl.j_squat)
Rh2Arm_ctrl.set_joint_vel(0.1)
time.sleep(2)
# 回収した空き缶の数を保存
strs = ["num = ", str(len(can_pos_list))]
f = open("../../../etc/env/collected_cans.py","w")
f.writelines(strs)
f.close()
# シナリオ終了
end()
