# QuatPlot.py
# Andrew Davison, ad@coe.psu.ac.th, August 2025

'''
Use VecPlot to display quat rotations

'''


import math

from Vec import Vec
from Quat import Quat
from VecPlot import VecPlot


def vecs1(plotter):
  # vector rotation
  v1 = Vec(3, 4, 5).normalize()
  print("v1:", v1)
  
  print("Angle between v1 and z-axis (deg) =", v1.angleToDeg(z))
  rotAngle = 90
  print("Rotation angle =", rotAngle)
  v2 = v1.rotateAround(z, math.radians(rotAngle))
  print("Result v2:", v2)

  plotter.addVector(v1, label='v1', color='blue')
  plotter.addVector(z, label='Ax(90)', color='green')
  plotter.addVector(v2, label='v2', color='red')


def quats1(plotter):
  # quaternion rotation
  v1 = Vec(1, 0, 0)
  q = Quat.fromAxisAngle(z, math.pi/2)
  v2 = q.rotateVec(v1)
  print(f"Rotate {v1} around {q} to {v2}")
  
  plotter.addVector(v1, label='v1', color='blue')
  plotter.addVector(v2, label='v2', color='red')
  axis, angle =  Quat.toAxisAngle(q)
  ang = math.degrees(angle)
  plotter.addVector(axis, label='q('+str(ang) +')', color='green')


def quats2(plotter):
  # rotate v = 2i + 1j by theta = pi/3 radians about 
  # u = 1/sqrt(2)j + 1/sqrt(2)k
  v1 = Vec(2, 1, 0) 
  u = Vec(0, 1/math.sqrt(2), 1/math.sqrt(2))
  # q = Quat(math.cos(math.pi/6), 
             # u * math.sin(math.pi/6))
  q = Quat.fromAxisAngle(u, math.pi/3)
  v2 = q.rotateVec(v1)
  print(f"Rotate {v1} around {q} to {v2}")
  
  axis, angle = q.toAxisAngle()
  ang = round(math.degrees(angle), 1)  
  plotter.addVector(v1, label='v1', color='blue')
  plotter.addVector(v2, label='v2', color='red')
  plotter.addVector(axis, label='Ax:'+str(ang), color='green')


def slerpShow(plotter):
  print("\n--- SLERP ---")
  q1 = Quat.fromAxisAngle(Vec(0, 1, 0), 0)
  q2 = Quat.fromAxisAngle(Vec(0, 1, 0), math.pi)
  n = 10
  steps = [t/n for t in range(n+1)]
  for i, t in enumerate(steps):
    qt = Quat.slerp(q1, q2, t)
    v = qt.rotateVec(Vec(1, 0, 0))
    print(f"{i} t={t:.2f}: {v}")
    plotter.addVector(v, label='q'+str(i), color='blue')


def lerpShow(plotter):
  print("\n--- LERP ---")
  q1 = Quat.fromAxisAngle(Vec(0, 1, 0), 0)
  q2 = Quat.fromAxisAngle(Vec(0, 1, 0), math.pi)
  n = 10
  steps = [t/n for t in range(n+1)]
  for i, t in enumerate(steps):
    qt = Quat.lerp(q1, q2, t)
    v = qt.rotateVec(Vec(1, 0, 0))
    print(f"{i} t={t:.2f}: {v}")
    plotter.addVector(v, label='q'+str(i), color='blue')



def euler1(plotter, x, y, z):
  # quaternion rotation using Euler angles
  print("\nRoll (x), pitch (y), yaw (z):", x, y, z)
  q = Quat.fromEuler(math.radians(x), 
         math.radians(y), math.radians(z))
  v1 = Vec(0, 1, 0)
  v2 = q.rotateVec(v1)
  print(f"Rotate {v1} around {q} to {v2}")
  
  plotter.addVector(v1, label='v1', color='blue')
  plotter.addVector(v2, label='v2', color='red')
  plotter.addVector(q.v, label=f'q({x},{y},{z})', color='green')

# ----------------

plotter = VecPlot("Quat Tests")

z = Vec(0, 0, 1)   # z-axis

# uncoment one of these examples:
# vecs1(plotter)
# quats1(plotter)
# quats2(plotter)
slerpShow(plotter)
# lerpShow(plotter)
# euler1(plotter, 45, 0, 0)

plotter.show()