;; Pair of bouncing balls, bouncing against each other, in Scheme
;; Base on Lagrangian Mechanics
;; By Piotr Mitros
;; Distributed under GPL

(define w1 (frame 0 50 -pi pi))
(define w2 (frame 0 50 -pi pi))
(graphics-enable-buffering w1)

(define (clear) 
  (graphics-clear w1)
  (graphics-clear w2))

(define (close) 
  (graphics-close w1)
  (graphics-close w2))

(define ((step1 el) x)
  (/ (+ (atan (* el x)) pi/2) pi))

(define ((step2 el) x)
  (* (+ (tanh (* el x)) 1) 1/2))

(define (spring k z)
  (* (/ k 2) (square z))
)

; We can also use the half-spring for the interaction between
; two particles. We'll also include gravity to be nice. 
(define ((two-ball r1 r2 m1 m2 g k el) x1 y1 z1 x2 y2 z2)
  (let ((distance (sqrt (+ (square (- x1 x2))
			   (square (- y1 y2))
			   (square (- z1 z2))))))
    (+ 0
       (*
	1
	(spring k (- distance (+ r1 r2)))
	((step2 -10) (- distance (+ r1 r2)))
	)
       (/ (* m1 m2 g) distance)
       )))

; Now, we write the Lagrangian for the particle
(define ((L-particle m1 k1 k2 r1 r2 g) local)
  (let ((q (coordinate local))
	 (v (velocity local)))
    (let ((z1 (ref q 0))
	  (z2 (ref q 1)))
      (- (* 1/2 m1 (+ (square (ref v 0)) (square (ref v 1))))
	 ((two-ball r1 r2 m1 m1 g k1 1) 0 0 z1 0 0 z2)))))

; We pretty-print the Lagrangian. Use 'se' for LaTeX output. 
(pe ((L-particle 'm1 'k1 'k2 'r1 'r2 'g)
     (up 't (up 'x 'y) (up 'xdot 'ydot))))

; We computer the equations of motion from the Lagrangian
(define (sysder m1 k1 k2 r1 r2 g)
  (Lagrangian->state-derivative
   (L-particle m1 k1 k2 r1 r2 g)))

; And we simplify and print those
(se ((sysder 'm1 'k1 'k2 'g)
     (up 't (up 'x 'y) (up 'xdot 'ydot))))

; We omit this, but we could compile the system derivative to be faster. 
(define double-sysder-compiled
  (compile-parametric 5 double-sysder))

(clear)

(define ((monitor-coords w1 w2) state)
  (let ((theta1 ((principal-value pi) (ref (coordinate state) 0)))
	(theta2 ((principal-value pi) (ref (coordinate state) 1))))
    ;(graphics-clear w1)
    ;(graphics-operation w1 'fill-circle (* .01 theta2) (* .01 theta1) 5)
    (write (time state))
    (write " ")
    (write theta1)
    (write " ")
    (write theta2)
    (write ":")
    (plot-point w1 (time state) theta1)
    (plot-point w1 (time state) 0)
    (plot-point w2 (time state) theta2)
    (plot-point w2 (time state) 0)
    state))

(clear)

; m k k r r g
((evolve sysder  1. 5 5 .7 .7 -.3)
 (up 0. (up 3 -3) (up 0. 0.)) 
 (monitor-coords w1 w2)
 .1
 50
 1.e-13)