from vtk import vtkTransform, vtkMatrix4x4
import numpy as np
import sys
import math

def rodrigues_vec_to_rotation_mat(rodrigues_vec):
  theta = np.linalg.norm(rodrigues_vec)
  if theta < sys.float_info.epsilon:              
    rotation_mat = np.eye(3, dtype=float)
  else:
    r = rodrigues_vec / theta
    I = np.eye(3, dtype=float)
    r_rT = np.array([
      [r[0]*r[0], r[0]*r[1], r[0]*r[2]],
      [r[1]*r[0], r[1]*r[1], r[1]*r[2]],
      [r[2]*r[0], r[2]*r[1], r[2]*r[2]]
    ])
    r_cross = np.array([
      [0, -r[2], r[1]],
      [r[2], 0, -r[0]],
      [-r[1], r[0], 0]
    ])
    rotation_mat = math.cos(theta) * I + (1 - math.cos(theta)) * r_rT + math.sin(theta) * r_cross
  return rotation_mat

def vtk_transform_from_2d_array(array):
  vtk_transform_matrix = vtkMatrix4x4()
  for i in range(4):
    for j in range(4):
      vtk_transform_matrix.SetElement(i, j, array[i][j])

  t = vtkTransform()
  t.SetMatrix(vtk_transform_matrix)
  return t

def transformation_matrix_from_s_r_t(s, r, t):
  R = rodrigues_vec_to_rotation_mat(r)
  transform_matrix = np.eye(4)
  transform_matrix[0:3,0:3] = s * R
  transform_matrix[0:3,3] = t
  return transform_matrix