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

'''
Plot normalized vecs inside a 3D unit sphere.
'''


import math

import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

from Vec import Vec


class VecPlot:
  def __init__(self, title='3D Unit Vectors'):
    self.vecs = []
    self.labels = []
    self.colors = []
    self.title = title


  def addVector(self, vec, label='v', color='blue'):
    """Add a unit vector with label and color."""
    self.vecs.append(vec.normalize())
    self.labels.append(label)
    self.colors.append(color)


  def show(self):
    fig = plt.figure(figsize=(8, 8))
    ax = fig.add_subplot(111, projection='3d')

    # Plot vecs
    for vec, label, color in zip(self.vecs, self.labels, self.colors):
      ax.quiver(0, 0, 0, vec.x, vec.y, vec.z, 
                color=color, arrow_length_ratio=0.1)
      ax.text(vec.x * 1.05, vec.y * 1.05, vec.z * 1.05, label, 
                color=color)

    self._drawAxes(ax)
    self._drawWireSphere(ax)

    # Place a Vector info box below the graph
    Vec.setPrintPrecision(3)
    vectorInfo = '\n'.join(f"{lbl} = {vec}" 
          for lbl, vec in zip(self.labels, self.vecs))
    props = dict(boxstyle='round', facecolor='white', alpha=0.9)
    fig.text(0.5, 0.02, vectorInfo, ha='center', va='bottom', 
                            fontsize=10, bbox=props)

    # format graph
    ax.set_xlim([-1.3, 1.3])
    ax.set_ylim([-1.3, 1.3])
    ax.set_zlim([-1.3, 1.3])
    ax.set_xlabel('X')
    ax.set_ylabel('Y')
    ax.set_zlabel('Z')
    ax.set_box_aspect([1, 1, 1])
    ax.set_title(self.title)
    plt.tight_layout(rect=[0, 0.05, 1, 1])

    plt.show()


  def _drawAxes(self, ax, length=1.2):
    # Draw dashed X, Y, Z axes with labels
    ax.plot([0, length], [0, 0], [0, 0], ls='--', color='black')
    ax.text(length, 0, 0, 'X', color='black')
    ax.plot([0, 0], [0, length], [0, 0], ls='--', color='black')
    ax.text(0, length, 0, 'Y', color='black')
    ax.plot([0, 0], [0, 0], [0, length], ls='--', color='black')
    ax.text(0, 0, length, 'Z', color='black')


  def _drawWireSphere(self, ax, radius=1.0, segs=24):
    # Draw a unit sphere with latitude and longitude lines
    for i in range(segs + 1):
      theta = i * math.pi / segs  # vary x
      latCircle = []   # build a latitude circle
      for j in range(segs + 1):
        phi = j * 2 * math.pi / segs
        x = radius * math.sin(theta) * math.cos(phi)
        y = radius * math.sin(theta) * math.sin(phi)
        z = radius * math.cos(theta)
        latCircle.append((x, y, z))
      xs, ys, zs = zip(*latCircle)
      ax.plot(xs, ys, zs, color='gray', linewidth=0.5, alpha=0.3)

    for j in range(segs + 1):
      phi = j * 2 * math.pi / segs   # vary y
      lonCircle = []  # build a longitude circle
      for i in range(segs + 1):
        theta = i * math.pi / segs
        x = radius * math.sin(theta) * math.cos(phi)
        y = radius * math.sin(theta) * math.sin(phi)
        z = radius * math.cos(theta)
        lonCircle.append((x, y, z))
      xs, ys, zs = zip(*lonCircle)
      ax.plot(xs, ys, zs, color='gray', linewidth=0.5, alpha=0.3)



# --- Example usage ---
if __name__ == "__main__":
  plotter = VecPlot("Rotate 'a' by 90 degs around Z axis")

  a = Vec(3, 4, 5).normalize()
  print("a:", a)
  z = Vec(0, 0, 1)   # z-axis

  print("Angle between a and z-axis (deg) =", a.angleToDeg(z))
  rotAngle = 90
  print("Rotation angle =", rotAngle)
  b = a.rotateAround(z, math.radians(rotAngle))
  print("Result b:", b)

  plotter.addVector(a, label='a', color='blue')
  plotter.addVector(b, label='b', color='red')

  plotter.show()