
# Vec.py
# Andrew Davison, ad@coe.psu.ac.th, August 2025

'''
A basic class to support 3D vectors used in the Quat class.
'''

import math

class Vec:
  _printPrecision = 3

  def __init__(self, x=0.0, y=0.0, z=0.0):
    self.x = x
    self.y = y
    self.z = z


  def __repr__(self):
    p = Vec._printPrecision
    fmt = f"{{:.{p}f}}"
    return f"Vec({fmt.format(self.x)}, {fmt.format(self.y)}, {fmt.format(self.z)})"


  @classmethod
  def setPrintPrecision(cls, digits):
    cls._printPrecision = max(0, int(digits))


  def __eq__(self, other):
    if not isinstance(other, Vec):
      return NotImplemented
    return self.isclose(other)


  def isclose(self, other, rt= 1e-9, at=0.0):
    """ Compares if this vector is close to another vector 
        within a tolerance.
           rt: relative tolerance;  at: absolute tolerance.
    """
    return (math.isclose(self.x, other.x, rel_tol=rt, abs_tol=at) and
            math.isclose(self.y, other.y, rel_tol=rt, abs_tol=at) and
            math.isclose(self.z, other.z, rel_tol=rt, abs_tol=at))


  def copy(self):
    return Vec(self.x, self.y, self.z)


  def __add__(self, other):
    return Vec(self.x + other.x, self.y + other.y, self.z + other.z)

  def __iadd__(self, other):
    # In-place addition: self += other
    self.x += other.x
    self.y += other.y
    self.z += other.z
    return self


  def __sub__(self, other):
    return Vec(self.x - other.x, self.y - other.y, self.z - other.z)

  def __isub__(self, other):
    # In-place subtraction: self -= other
    self.x -= other.x
    self.y -= other.y
    self.z -= other.z
    return self


  def __mul__(self, scalar):
    return Vec(self.x*scalar, self.y*scalar, self.z*scalar)

  def __rmul__(self, scalar):
    return self.__mul__(scalar)

  def __imul__(self, scalar):
    # In-place multiplication by a scalar: self *= scalar
    self.x *= scalar
    self.y *= scalar
    self.z *= scalar
    return self


  def __truediv__(self, scalar):
    if scalar == 0:
      raise ZeroDivisionError("Cannot divide by zero")
    return Vec(self.x/scalar, self.y/scalar, self.z/scalar)

  def __itruediv__(self, scalar):
    # In-place division by a scalar: self /= scalar
    if scalar == 0:
      raise ZeroDivisionError("Cannot divide by zero")
    self.x /= scalar
    self.y /= scalar
    self.z /= scalar
    return self


  def __neg__(self):
    return Vec(-self.x, -self.y, -self.z)

  def __abs__(self):
    return self.magnitude()

  def normalize(self):
    mag = self.magnitude()
    if mag == 0:
      raise ValueError("Cannot normalize a zero vector")
    return self/mag

  def magnitude(self):
    return math.sqrt(self.x**2 + self.y**2 + self.z**2)


  def dot(self, other):
    return self.x*other.x + self.y*other.y + self.z*other.z


  def cross(self, other):
    return Vec(
      self.y*other.z - self.z*other.y,
      self.z*other.x - self.x*other.z,
      self.x*other.y - self.y*other.x)


  def normSquared(self):
    return self.dot(self)


  def isZero(self, tol=1e-8):
    return self.normSquared() < tol*tol


  def projectOnto(self, other):
    unitOther = other.normalize()
    scale = self.dot(unitOther)
    return unitOther*scale


  def angleTo(self, other):
    dot = self.dot(other)
    mag = self.magnitude()*other.magnitude()
    if mag == 0:
      raise ValueError("Cannot compute angle with zero-length vector")
    cosTheta = max(-1, min(1, dot/mag))
    return math.acos(cosTheta)


  def angleToDeg(self, other):
    return math.degrees(self.angleTo(other))


  def rotateAround(self, axis, theta):
    # Rodrigues Rotation Formula; v = self
    k = axis.normalize()
    cosA = math.cos(theta)
    sinA = math.sin(theta)
    vterm1 = self*cosA
    vterm2 = k.cross(self)*sinA
    vterm3 = k*(k.dot(self)*(1 - cosA))
    return vterm1 + vterm2 + vterm3  # rotated vector


  @staticmethod
  def fromList(lst):
    if len(lst) != 3:
      raise ValueError("List must have 3 elements")
    return Vec(lst[0], lst[1], lst[2])

  def toList(self):
    return [self.x, self.y, self.z]

  @staticmethod
  def fromTuple(tpl):
    if len(tpl) != 3:
      raise ValueError("Tuple must have 3 elements")
    return Vec(tpl[0], tpl[1], tpl[2])

  def toTuple(self):
    return (self.x, self.y, self.z)


# ----------- Test Rig ---------------

if __name__ == "__main__":
  Vec.setPrintPrecision(2)

  a = Vec(3, 4, 0)
  b = Vec(0, 5, 0)
  zAxis = Vec(0, 0, 1)

  print("a =", a)
  print("-a =", -a)
  print("|a| =", abs(a))
  print("b =", b)
  print("a + b =", a + b)
  # Test in-place addition
  a_copy = a.copy()
  a_copy += b
  print("a_copy += b =", a_copy)
  print("a . b =", a.dot(b))
  print("a × b =", a.cross(b))
  print("Normalized 'a' =", a.normalize())
  print("Projection of a onto b =", a.projectOnto(b))
  print(f"Angle between a and b (deg) = {a.angleToDeg(b):.3f}")
  print("Rotate 'a' by 90° around Z axis =", a.rotateAround(zAxis, math.pi/2))

  print("\n--- isclose and equality demo ---")
  v1 = Vec(1.0, 2.0, 3.0)
  v2 = Vec(1.0 + 1e-10, 2.0, 3.0)
  v3 = Vec(1.0, 2.0, 3.001)
  print(f"v1 = {v1}, v2 = {v2}, v3 = {v3}")
  print(f"v1 == v2 (default tolerance): {v1 == v2}")
  print(f"v1.isclose(v2, rt=1e-8): {v1.isclose(v2, rt=1e-8)}")
  print(f"v1 == v3 (default tolerance): {v1 == v3}")
  print(f"v1.isclose(v3, at=0.01): {v1.isclose(v3, at=0.01)}")