import numpy as np
import cv2  # pylint: disable=import-error

def rot_matrix(roll, pitch, yaw):
  cr, sr = np.cos(roll), np.sin(roll)
  cp, sp = np.cos(pitch), np.sin(pitch)
  cy, sy = np.cos(yaw), np.sin(yaw)
  rr = np.array([[1,0,0],[0, cr,-sr],[0, sr, cr]])
  rp = np.array([[cp,0,sp],[0, 1,0],[-sp, 0, cp]])
  ry = np.array([[cy,-sy,0],[sy, cy,0],[0, 0, 1]])
  return ry.dot(rp.dot(rr))

def draw_pose(img, pose, loc, W=160, H=320, xyoffset=(0,0), faceprob=0):
  rcmat = np.zeros((3,4))
  rcmat[:,:3] = rot_matrix(*pose[0:3]) * 0.5
  rcmat[0,3] = (loc[0]+0.5) * W
  rcmat[1,3] = (loc[1]+0.5) * H
  rcmat[2,3] = 1.0
  # draw nose
  p1 = np.dot(rcmat, [0,0,0,1])[0:2]
  p2 = np.dot(rcmat, [0,0,100,1])[0:2]
  tr = tuple([int(round(x + xyoffset[i])) for i,x in enumerate(p1)])
  pr = tuple([int(round(x + xyoffset[i])) for i,x in enumerate(p2)])
  if faceprob > 0.4:
    color = (255,255,0)
    cv2.line(img, tr, pr, color=(255,255,0), thickness=3)
  else:
    color = (64,64,64)
  cv2.circle(img, tr, 7, color=color)