import torch
from torch import nn
import math

#0 left hip
#1 left knee
#2 left foot
#3 right hip
#4 right knee
#5 right foot
#6 middle hip
#7 neck
#8 nose
#9 head
#10 left shoulder
#11 left elbow
#12 left wrist
#13 right shoulder
#14 right elbow
#15 right wrist

def random_rotation(J3d):
  J = J3d # need copy????
  batch_size = J.shape[0]
  theta = torch.rand(batch_size).cuda() * 2*torch.tensor(math.pi).cuda() # random theta
  root = J[:,:,8] # joint 8 = nose is root
  J3d_R = rotation(J.cuda(), theta.cuda(), root.unsqueeze(-1).cuda(), False)
  return J3d_R, theta, root # need these values in the code

def rotation(J, theta, root, is_reversed): # rotation over y axis by theta
  D =  root[:,2].cuda() # absolute depth of the root joint
  batch_size = root.shape[0]
  v_t = torch.zeros((batch_size, 3, 1)).cuda()
  v_t[:, 2, :] = D.cuda() # translation vector
  if is_reversed:
    root, v_t = v_t, root # swap
    theta = -theta
  # R = torch.tensor([[torch.cos(theta), -torch.sin(theta), 0], [torch.sin(theta), torch.cos(theta), 0], [0, 0, 1]]) # rotation matrix over z by theta degrees
  R = torch.zeros((batch_size, 3, 3)).cuda() # rotation matrix over y by theta degrees
  R[:, 0, 0] = torch.cos(theta)
  R[:, 0, 2] = torch.sin(theta)
  R[:, 1, 1] = torch.ones(batch_size)
  R[:, 2, 0] = -torch.sin(theta)
  R[:, 2, 2] = torch.cos(theta)
  # R = torch.tensor([[torch.cos(theta), 0, torch.sin(theta)], [0, 1, 0], [-torch.sin(theta), 0, torch.cos(theta)]]) # rotation matrix over y by theta degrees
  # R = torch.tensor([[1, 0, 0], [0, torch.cos(theta), -torch.sin(theta)], [0, torch.sin(theta), torch.cos(theta)]]) # rotation matrix over x by theta degrees
  
  J_R = torch.matmul(R, J - root) + v_t # rotation
  return J_R

def reverse_rotation(J3d_R, theta, root):
  J = J3d_R # need copy????
  return rotation(J.cuda(), theta.cuda(), root.unsqueeze(-1).cuda(), True)

def temporal_loss(J, K, J_R, K_R): # J is J3d at time t and K is J3d at time t+k. J_R means the reversed rotation of J
  
  #print(torch.norm(J.reshape(J.shape[0], 3, 16) - K.reshape(J.shape[0], 3, 16) - J_R.reshape(J.shape[0], 3, 16) + K_R.reshape(J.shape[0], 3, 16), dim=1).shape)
  #stop
  mse_fn = nn.MSELoss()
  return mse_fn(J.reshape(J.shape[0], 3, 16) - K.reshape(J.shape[0], 3, 16) - J_R.reshape(J.shape[0], 3, 16) + K_R.reshape(J.shape[0], 3, 16), torch.zeros(J.shape[0], 3, 16).cuda())
  #return torch.norm(J.reshape(J.shape[0], 3, 16) - K.reshape(J.shape[0], 3, 16) - J_R.reshape(J.shape[0], 3, 16) + K_R.reshape(J.shape[0], 3, 16), dim=1)**2
'''
def temporal_loss(J, K, J_R, K_R): # J is J3d at time t and K is J3d at time t+k. J_R means the reversed rotation of J
  return torch.norm(J - K - J_R + K_R, dim=1)**2
'''

'''
def random_rotation(J3d):
  # J = torch.transpose(J3d, 1, 2)
  J = J3d
  root = torch.zeros(J.shape[0:2])
  for i in range(J.shape[0]):
    theta = torch.rand(1).cuda() * 2*torch.tensor(math.pi).cuda() # random theta
    root[i] = J[i,:,8] # joint 8 = nose is root
    temp = rotation(J[i,:,:], theta, root[i].unsqueeze(1), False)
    # print(temp.shape)
    J[i,:,:] = temp
  return J, theta, root # need these values in the code

def rotation(J, theta, root, is_reversed): # rotation over y axis by theta
  D =  root[2] # absolute depth of the root joint
  v_t = torch.tensor([[0], [0], [D]]).cuda() # translation vector
  if is_reversed:
    root, v_t = v_t, root # swap
    theta = -theta
  # R = torch.tensor([[torch.cos(theta), -torch.sin(theta), 0], [torch.sin(theta), torch.cos(theta), 0], [0, 0, 1]]) # rotation matrix over z by theta degrees
  R = torch.tensor([[torch.cos(theta), 0, torch.sin(theta)], [0, 1, 0], [-torch.sin(theta), 0, torch.cos(theta)]]).cuda() # rotation matrix over y by theta degrees
  # R = torch.tensor([[1, 0, 0], [0, torch.cos(theta), -torch.sin(theta)], [0, torch.sin(theta), torch.cos(theta)]]) # rotation matrix over x by theta degrees
  
  J_R = torch.matmul(R, J.cuda() - root.cuda()) + v_t # rotation
  return J_R

def reverse_rotation(J3d_R, theta, root):
  # J = torch.transpose(J3d_R, 1, 2)
  J = J3d_R
  for i in range(J.shape[0]):
    J[i,:,:] = rotation(J[i,:,:].cuda(), theta.cuda(), root[i].unsqueeze(1).cuda(), True)
  return J
'''
[((40, 15, 40, 31), 'torch.cos', 'torch.cos', ({(40, 25, 40, 30): 'theta'}, {}), '(theta)', False, 'import torch\n'), ((41, 15, 41, 31), 'torch.sin', 'torch.sin', ({(41, 25, 41, 30): 'theta'}, {}), '(theta)', False, 'import torch\n'), ((42, 15, 42, 37), 'torch.ones', 'torch.ones', ({(42, 26, 42, 36): 'batch_size'}, {}), '(batch_size)', False, 'import torch\n'), ((44, 15, 44, 31), 'torch.cos', 'torch.cos', ({(44, 25, 44, 30): 'theta'}, {}), '(theta)', False, 'import torch\n'), ((59, 11, 59, 23), 'torch.nn.MSELoss', 'nn.MSELoss', ({}, {}), '()', False, 'from torch import nn\n'), ((43, 16, 43, 32), 'torch.sin', 'torch.sin', ({(43, 26, 43, 31): 'theta'}, {}), '(theta)', False, 'import torch\n'), ((48, 8, 48, 33), 'torch.matmul', 'torch.matmul', ({(48, 21, 48, 22): 'R', (48, 24, 48, 32): '(J - root)'}, {}), '(R, J - root)', False, 'import torch\n'), ((33, 8, 33, 39), 'torch.zeros', 'torch.zeros', ({(33, 20, 33, 38): '(batch_size, 3, 1)'}, {}), '((batch_size, 3, 1))', False, 'import torch\n'), ((39, 6, 39, 37), 'torch.zeros', 'torch.zeros', ({(39, 18, 39, 36): '(batch_size, 3, 3)'}, {}), '((batch_size, 3, 3))', False, 'import torch\n'), ((25, 44, 25, 65), 'torch.tensor', 'torch.tensor', ({(25, 57, 25, 64): 'math.pi'}, {}), '(math.pi)', False, 'import torch\n'), ((60, 143, 60, 173), 'torch.zeros', 'torch.zeros', ({(60, 155, 60, 165): 'J.shape[0]', (60, 167, 60, 168): '(3)', (60, 170, 60, 172): '(16)'}, {}), '(J.shape[0], 3, 16)', False, 'import torch\n'), ((25, 10, 25, 32), 'torch.rand', 'torch.rand', ({(25, 21, 25, 31): 'batch_size'}, {}), '(batch_size)', False, 'import torch\n')]