import numpy as np

from geometry.vector import Vector


def reflect(direction: Vector, normal: Vector) -> Vector:
  reflected = direction - 2 * (direction @ normal) * normal
  return reflected.normalize()


def refract(direction: Vector, normal: Vector, n1: float, n2: float) -> Vector:
  if n1 > n2:
    return refract(direction, -1 * normal, n2, n1)

  proj = normal @ (n1 * direction)
  k = (n2 ** 2 - n1 ** 2) / proj ** 2 + 1
  refracted = (1 / n2) * (n1 * direction + normal * proj
                          * (np.sqrt(k) - 1))

  return refracted.normalize()
