# rollingSquare.py
# Andrew Davison, Nov 2025, ad@coe.psu.ac.th

'''
Rolling square visualization on catenary curves.
Animates a square rolling along a series of catenary curves,
showing the path traced by one of its vertices.
After animation completes, plots angle vs x position.

The original Mathematica program is from "Mathematica in Action" 
by Stan Wagon , W. H. Freeman and Co., New York, 1991
https://link.springer.com/book/10.1007/978-0-387-75477-2
https://extras.springer.com/?query=978-0-387-75366-9
  -- Product Achive File
'''

import math
import matplotlib.pyplot as plt
from matplotlib.patches import Polygon as MplPolygon
from matplotlib.animation import FuncAnimation


def catenaryPts(i, cr, ir, hArch, numPts=100):
  # generate i-th catenary arch made up of numPts points
  xs = []
  ys = []
  xMin = -hArch + 2*hArch*i
  xMax = hArch + 2*hArch*i
  for j in range(numPts):
    x = xMin + (xMax - xMin)*j/(numPts - 1)
    y = cr - ir*math.cosh((x - 2*hArch*i)/ir)
        # catenary equation, offset by 2*hArch*i
    xs.append(x)
    ys.append(y)
  return xs, ys


def createSquare(cr):
  # Square has 4 corners at angles: -45, 45, 135, 225 degrees
  startSqr = []
  for i in range(4):
    angle = -math.pi/4 + i*math.pi/2
    x = cr*math.cos(angle)
    y = cr*math.sin(angle) + cr
    startSqr.append([x, y])
  return startSqr


def genPositions(numFrames, hArch, startSqr, 
                          cr, ir, rangeMult):
  # generate all the animation info (times, squares, trace dots, angles)
  ds = []
  for i in range(numFrames):
    d = rangeMult*hArch*i/(numFrames - 1)  \
                   if numFrames > 1 else rangeMult*hArch
    ds.append(d)
  squares = []
  traceDots = []
  angles = []
  for d in ds:
    currSqr, angle = squareAtT(d, startSqr, cr, ir, hArch)
    squares.append(currSqr)
    traceDots.append(currSqr[3])
    angles.append(angle)
  return ds, squares, traceDots, angles


def squareAtT(d, startSqr, cr, ir, hArch):
  ''' calculates the position and orientation 
      of the square at a x-distance d
      as it rolls over the catenary curves floor
  '''
  loop = int(math.floor((d + hArch)/(2*hArch)))
  angle = loop*math.pi/2 + \
             math.atan(math.sinh((d - loop*2*hArch)/ir))
  center = [0, cr]  # of the square above x-axis
  
  nSquare = []
  for pt in startSqr:
    newPt = [d, 0]  # offset from origin
    rots = rotatePoint(pt, angle, center)
    nSquare.append([newPt[0] + rots[0], newPt[1] + rots[1]])
  return nSquare, angle


def rotatePoint(pt, angle, center):
  cosA = math.cos(-angle)
  sinA = math.sin(-angle)
  px = pt[0] - center[0]
  py = pt[1] - center[1]
  return [ cosA*px - sinA*py + center[0],
           sinA*px + cosA*py + center[1] ]


def animate(frameNum):
  sqr = squares[frameNum]
  
  sqrPatch.set_xy(sqr)   # Update square

  # Update diagonals
  diagonal1.set_data([sqr[0][0], sqr[2][0]], [sqr[0][1], sqr[2][1]])
  diagonal2.set_data([sqr[1][0], sqr[3][0]], [sqr[1][1], sqr[3][1]])
  
  # Update center point
  centerPt.set_data([ds[frameNum]], [cr])
  
  # Update corner point
  cornerPt.set_data([sqr[3][0]], [sqr[3][1]])
  
  # Update trace points (show all points up to current frame)
  traceX = [traceDots[i][0] for i in range(frameNum + 1)]
  traceY = [traceDots[i][1] for i in range(frameNum + 1)]
  tracePts.set_data(traceX, traceY)
  
  return sqrPatch, diagonal1, diagonal2, centerPt, cornerPt, tracePts


# ----------------------------
# Square geometry (hardwired for n=4)
cr = math.sqrt(2)      # circumradius
ir = 1.0               # inradius
hArch = math.asinh(1)  # half-width of catenary curve

# Animation settings
numFrames = 60  # number of animation frames
rangeMult = 18  
   # how far the square rolls == rangeMult * hArch
numCats = 10    # number of catenary curves on floor


startSqr = createSquare(cr)
ds, squares, traceDots, angles = genPositions(
          numFrames, hArch, startSqr, cr, ir, rangeMult)


fig, ax = plt.subplots(figsize=(16, 4))

# Draw catenary curves floor
for i in range(numCats):
  xs, ys = catenaryPts(i, cr, ir, hArch)
  ax.plot(xs, ys, 'k-', lw=2)

ax.plot([-cr, rangeMult*hArch], [cr, cr], 'b-', lw=2)  
              # center of square's path

# Initialize plot elements
sqrPatch = MplPolygon([[0,0]], 
       facecolor='lightgray', edgecolor='black', lw=1.5)
ax.add_patch(sqrPatch)   # the square
diagonal1, = ax.plot([], [], 'k-', lw=1) # two diagonals
diagonal2, = ax.plot([], [], 'k-', lw=1)
centerPt, = ax.plot([], [], 'ko', ms=8)  # two pts on square
cornerPt, = ax.plot([], [], 'ko', ms=8)
tracePts, = ax.plot([], [], 'o',  color='red', ms=6)   # trace dots
                      

ax.set_xlim(-cr, rangeMult*hArch + 0.2)
ax.set_ylim(-0.2, 2*cr + 0.15)
ax.set_aspect('equal')
ax.grid(True, alpha=0.3)
ax.set_xlabel('x')
ax.set_ylabel('y')
ax.set_title('Rolling Square on Catenary Curve')

# Create animation without repeat
anim = FuncAnimation(fig, animate, frames=numFrames, 
                   interval=50, blit=True, repeat=False)

plt.tight_layout()
plt.show()

# Plot angle vs x position
fig2, ax2 = plt.subplots(figsize=(12, 6))
anglesDeg = [a * 180 / math.pi for a in angles]
ax2.plot(ds, anglesDeg, 'b-', lw=2)
# ax2.plot([ds[0], ds[-1]], [anglesDeg[0], anglesDeg[-1]], 
         # 'r--', lw=3, alpha=0.5)
ax2.set_xlabel('x position')
ax2.set_ylabel('angle (degrees)')
ax2.set_title('Square Angle vs X Position')
ax2.grid(True, alpha=0.3)
plt.tight_layout()
plt.show()