# brachAnim.py
# Andrew Davison, ad@coe.psu.ac.th, Dec. 2025

"""
Brachistochrone Problem Animation

This program visualizes and compares three different paths for a bead sliding under gravity from point (0,0) to point (xEnd, yEnd):

1. Brachistochrone (Cycloid) - The optimal path that minimizes travel time
2. Straight Line - The shortest distance path along an inclined slope
3. Circular Arc - A curved path forming part of a circle

The animation demonstrates that the brachistochrone curve, despite being longer, allows the bead to reach the endpoint faster than both the straight line and circular arc paths.

- No friction or air resistance

Controls:
- Spacebar: Pause/Resume animation


Brachistochrone:
  - Uses parametric cycloid equations: 
         x = a(theta - sin(theta)),  y = a(1 - cos(theta))
  - Direct time-parameter relationship: theta = t/sqrt(a/g)

Straight Line:
  - Constant acceleration down an inclined plane
  - Travel time: t = sqrt(2L / (g sin(alpha)))

Circular Arc:
  - Numerical integration to calculate total travel time
  - Approximates non-uniform velocity along the arc

"""

import math
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from scipy.optimize import brentq


aVal = 1.0
G = 9.81
numFrames = 300

xEnd = 3.16; yEnd = 2.0
# xEnd = 4.5; yEnd = 1.7596


isRunning = True
anim = None

beads = {
  'brachistochrone': {
    'x': 0, 'y': 0, 'time': 0, 'finished': False, 'color': 'red',
    'marker': None, 'text': None
  },
  'straight': {
    'x': 0, 'y': 0, 'time': 0, 'finished': False, 'color': 'green',
    'marker': None, 'text': None
  },
  'circular': {
    'x': 0, 'y': 0, 'time': 0, 'finished': False, 'color': 'blue',
    'marker': None, 'text': None
  }
}


def cycloid(theta, a):
  x = a * (theta - math.sin(theta))
  y = a * (1 - math.cos(theta))
  return x, y


def cycloidT(t, x, a):
  return a * (t - math.sin(t)) - x


def cycloidYTfromX(x, a):
  archWidth = 2*math.pi * a
  
  # Determine arch index and x within that arch
  archIndex = math.floor(x / archWidth)
  xArch = x - archIndex * archWidth
      # handles negative and too large x values
  
  # cusp handling
  if abs(xArch) < 1e-12 or abs(xArch - archWidth) < 1e-12:
    return 0, 0
  else:
    tMin = 0
    tMax = 2*math.pi
    t = brentq(cycloidT, tMin, tMax, args=(xArch, a))
    y = a * (1 - math.cos(t))
    return y, t


def cycloidPts(a, thetaEnd, nPts):
  curveX = []
  curveY = []
  for i in range(nPts + 1):
    theta = thetaEnd * (i / nPts)
    x, y = cycloid(theta, a)
    curveX.append(x)
    curveY.append(y)
  return curveX, curveY


def calculateCircularArc(xEnd, yEnd, nPts):
  """ Calculate parameters and points for a circular arc 
      from (0,0) to (xEnd, yEnd).
      Returns the curve points and the parametric time 
      to traverse it.
  """
  # Midpoint of the chord
  midX = xEnd / 2.0
  midY = yEnd / 2.0
  
  chordLen = math.sqrt(xEnd**2 + yEnd**2)
  
  radius = chordLen * 0.8
  
  # Dist from midpoint to center along perpendicular
  halfChord = chordLen / 2.0
  if radius < halfChord:
    radius = halfChord * 1.1
  
  distToCenter = math.sqrt(radius**2 - halfChord**2)
  
  # Perpendicular dir (rotated 90 degrees clockwise for downward arc)
  perpX = yEnd / chordLen
  perpY = -xEnd / chordLen
  
  # Center pos (below the chord)
  centerX = midX + perpX * distToCenter
  centerY = midY + perpY * distToCenter
  
  startAngle = math.atan2(0 - centerY, 0 - centerX)
  endAngle = math.atan2(yEnd - centerY, xEnd - centerX)
  # Ensure we go clockwise from start to end
  if endAngle > startAngle:
    endAngle -= 2 * math.pi
  
  angles = [startAngle + (endAngle - startAngle) * i / nPts 
                                   for i in range(nPts + 1)]
  curveX = [centerX + radius * math.cos(a) for a in angles]
  curveY = [centerY + radius * math.sin(a) for a in angles]

  # Calculate time using time = distance / velocity
  totalTime = 0.0
  numSteps = 1000
  for i in range(numSteps):
    t = i / numSteps
    angle = startAngle + (endAngle - startAngle) * t
    x = centerX + radius * math.cos(angle)
    y = centerY + radius * math.sin(angle)
    
    # for distance, sum up small segments along the arc
    # totalTime = sum (ds/v)
    ds = abs(endAngle - startAngle) * radius / numSteps
    # ds = |dTheta|*r
    height = y
    # v = sqrt(2gh) where h is the vertical drop from the starting point
    v = math.sqrt(2 * G * height) if height > 0 else 0.1
    if v > 0:
      totalTime += ds / v

  return centerX, centerY, radius, startAngle, endAngle, curveX, curveY, totalTime



def initAnim():
  for bead in beads.values():
    bead['marker'].set_data([], [])
    bead['text'].set_text('')
  return tuple(bead['marker'] for bead in beads.values()) + \
         tuple(bead['text'] for bead in beads.values())



def updateAnim(frame):
  global isRunning, anim

  if not isRunning:
    return tuple(bead['marker'] for bead in beads.values()) + \
           tuple(bead['text'] for bead in beads.values())

  currTime = totTime * (frame / numFrames)

  # Map path types to their end times
  pathTimes = {
    'brachistochrone': timeBr,
    'straight': timeSt,
    'circular': timeCirc
  }
  
  # Map path types to their display names
  pathNames = {
    'brachistochrone': 'Brachistochrone',
    'straight': 'Straight Line',
    'circular': 'Circular Arc'
  }

  # Update all beads
  for pathType, bead in beads.items():
    x, y, bead['finished'], bead['time'] = updateBeadPos(
      currTime, pathType, bead, pathTimes[pathType])
    bead['x'], bead['y'] = x, y
    bead['marker'].set_data([x], [y])
    bead['text'].set_text(f'{pathNames[pathType]}: {bead["time"]:.2f}s')

  # Stop animation when all beads finish
  if all(bead['finished'] for bead in beads.values()) and isRunning:
    anim.pause()
    isRunning = False

  return tuple(bead['marker'] for bead in beads.values()) + \
         tuple(bead['text'] for bead in beads.values())



def updateBeadPos(currTime, pathType, beadState, endTime):
  if beadState['finished'] or currTime >= endTime:
    return xEnd, yEnd, True, endTime
  
  if pathType == 'brachistochrone':
    currTheta = currTime / math.sqrt(aVal / G)
    ''' For a cycloid starting from rest, the relationship 
        between time 't' and the parameter 'theta' is:
          t = sqrt(a/g) × theta
    '''
    xPos, yPos = cycloid(currTheta, aVal)
  elif pathType == 'straight':
    distance = 0.5 * (G * math.sin(alphaSt)) * currTime**2
    xPos = distance * math.cos(alphaSt)
    yPos = distance * math.sin(alphaSt)
  elif pathType == 'circular':
    t = currTime / endTime
    xPos, yPos = getCircularPos(t, endTime)
  else:
    raise ValueError(f"Unknown path type: {pathType}")
    
  return xPos, yPos, False, currTime


def getCircularPos(t, timeCirc):
  ''' Unlike the brachistochrone, a circular arc doesn't have 
    a simple closed-form relationship between time and position. 
    The code uses a normalized time parameter approach:
       * t = currTime / endTime gives a value from 0 to 1
       * The angle along the arc is interpolated linearly
  '''
  if t >= 1.0:
    return xEnd, yEnd
  angle = startAngleCirc + (endAngleCirc - startAngleCirc) * t
  x = centerCircX + radiusCirc * math.cos(angle)
  y = centerCircY + radiusCirc * math.sin(angle)
  return x, y



def onKeyPress(event):
  global isRunning

  if event.key == ' ':  # Spacebar toggles pause/resume
    if isRunning:
      anim.pause()
      isRunning = False
    else:
      anim.resume()
      isRunning = True


def main():
  global timeBr, timeSt, timeCirc, alphaSt, totTime
  global centerCircX, centerCircY, radiusCirc, startAngleCirc, endAngleCirc
  global fig, ax, anim
  
  # Calculate brachistochrone path
  _, thetaBrEnd = cycloidYTfromX(xEnd, aVal)
  timeBr = thetaBrEnd * math.sqrt(aVal / G)
  brCurveX, brCurveY = cycloidPts(aVal, thetaBrEnd, 500)

  # Calculate straight slope path
  slopeLength = math.sqrt(xEnd**2 + yEnd**2)
  alphaSt = math.asin(yEnd / slopeLength)
  timeSt = math.sqrt(2 * slopeLength / (G * math.sin(alphaSt)))

  # Calculate circular arc path with pre-calculated time
  (centerCircX, centerCircY, radiusCirc, startAngleCirc, endAngleCirc, 
   circCurveX, circCurveY, timeCirc) = calculateCircularArc(xEnd, yEnd, 500)

  # Set total animation time
  totTime = max(timeBr, timeSt, timeCirc, 8.0)

  print(f"Travel times:")
  print(f"  Brachistochrone: {timeBr:.3f}s")
  print(f"  Straight Line: {timeSt:.3f}s")
  print(f"  Circular Arc: {timeCirc:.3f}s")

  # Set up the plot
  fig, ax = plt.subplots(figsize=(10, 6))
  ax.set_aspect('equal')
  ax.set_title('Brachistochrone Problem: Comparing Different Paths')
  ax.set_xlabel('x')
  ax.set_ylabel('y')
  ax.grid(True, alpha=0.3)
  ax.invert_yaxis()

  # Plot the curves
  ax.plot(brCurveX, brCurveY, 'r-', lw=2.5, alpha=0.8)
  ax.plot([0, xEnd], [0, yEnd], 'g-', lw=2.5, alpha=0.8)
  ax.plot(circCurveX, circCurveY, 'b-', lw=2.5, alpha=0.8)

  # Initialize bead markers and text objects
  text_positions = {'brachistochrone': 0.98, 'straight': 0.94, 'circular': 0.90}
  
  for path_type, bead in beads.items():
    color = bead['color']
    dark_color = f'dark{color}'
    
    # Create marker
    bead['marker'], = ax.plot([], [], 'o', color=color, markersize=14, 
                     markeredgecolor=dark_color, markeredgewidth=2, 
                     markerfacecolor=color, alpha=0.9)
    # Create text
    bead['text'] = ax.text(0.98, text_positions[path_type], '', 
                  transform=ax.transAxes, verticalalignment='top', 
                  horizontalalignment='right',
                  color=color)

  fig.canvas.mpl_connect('key_press_event', onKeyPress)
  
  anim = FuncAnimation(
             fig, updateAnim, frames=numFrames,
             init_func=initAnim, blit=True,
             interval=totTime * 1000 / numFrames,
             repeat=False
  )
  
  plt.tight_layout()
  plt.show()


if __name__ == "__main__":
  main()