import numpy as np

from geometry.vector import Vector


def rotate2d(v: Vector, alpha: float) -> Vector:
  R = np.array([
    [np.cos(alpha), -np.sin(alpha)],
    [np.sin(alpha), np.cos(alpha)],
  ])

  return Vector(R @ v.a)


def rotate3d(v: Vector, alpha: float, axis: int) -> Vector:
  R = np.eye(3)

  if axis == 0:
    R[1, 1] = np.cos(alpha)
    R[1, 2] = -np.sin(alpha)
    R[2, 1] = np.sin(alpha)
    R[2, 2] = np.cos(alpha)

  elif axis == 1:
    R[0, 0] = np.cos(alpha)
    R[0, 2] = np.sin(alpha)
    R[2, 0] = -np.sin(alpha)
    R[2, 2] = np.cos(alpha)

  elif axis == 2:
    R[0, 0] = np.cos(alpha)
    R[0, 1] = -np.sin(alpha)
    R[1, 0] = np.sin(alpha)
    R[1, 1] = np.cos(alpha)

  else:
    raise ValueError

  return Vector(R @ v.a)


def rotate_rodrigues(v: Vector, k: Vector, alpha: float) -> Vector:
  k = k.normalize()

  return v * np.cos(alpha) + (k | v) * np.sin(alpha) + k * (k @ v) * (1 - np.cos(alpha))
