refactor ball creation and collision
[bowntz:bowntz.git] / bowntz.hs
1 import Graphics.UI.GLUT hiding (position)
2 import Sound.OpenSoundControl
3 import Sound.SC3 hiding (decay, ID)
4 import qualified Sound.SC3 as SC
5
6 import Control.Monad(forM_, when)
7 import Data.IORef
8 import System.Random
9
10 import qualified Data.Map as M
11 import Data.Map(Map)
12 import qualified Data.Set as S
13 import Data.Set(Set)
14 import Debug.Trace
15
16 type ID = Integer
17
18 firstID :: ID
19 firstID = 1000
20
21 type R = GLdouble
22
23 data V = V { x :: R, y :: R }
24   deriving (Eq, Ord, Read, Show)
25
26 infixl 8 `dot`
27 infixl 7 *^,^*,^/
28 infixl 6 ^+^,^-^
29
30 (^/) :: V -> R -> V
31 u ^/ t = V (x u / t) (y u / t)
32
33 (^*) :: V -> R -> V
34 u ^* t = V (x u * t) (y u * t)
35
36 (*^) :: R -> V -> V
37 s *^ v = V (s * x v) (s * y v)
38
39 (^+^) :: V -> V -> V
40 u ^+^ v = V (x u + x v) (y u + y v)
41
42 (^-^) :: V -> V -> V
43 u ^-^ v = V (x u - x v) (y u - y v)
44
45 dot :: V -> V -> R
46 dot u v = x u * x v + y u * y v
47
48 data Ball = Ball { mass :: R, radius :: R, position :: V, velocity :: V, ring :: R }
49   deriving (Eq, Ord, Read, Show)
50
51 data Collision = Collision { time :: R, ball1 :: ID, ball2 :: ID }
52   deriving (Eq, Ord, Read, Show)
53
54 kineticEnergy :: Ball -> R
55 kineticEnergy b = 0.5 * mass b * velocity b `dot` velocity b
56
57 involves :: ID -> Collision -> Bool
58 involves i c = i == ball1 c || i == ball2 c
59
60 offset :: R -> Collision -> Collision
61 offset t c = c { time = time c + t }
62
63 collideOutside :: Ball -> Ball -> (Ball, Ball)
64 collideOutside b1 b2 = collideXside (radius b2 + radius b1) b1 b2
65
66 collideInside :: Ball -> Ball -> (Ball, Ball)
67 collideInside b1 b2 = collideXside (abs $ radius b2 - radius b1) b1 b2
68
69 collideXside :: R -> Ball -> Ball -> (Ball, Ball)
70 collideXside r b1 b2 = (b1 { velocity = v1, ring = ring b1 + abs e1 }, b2 { velocity = v2, ring = ring b2 + abs e2 })
71   where n = (position b1 ^-^ position b2) ^/ r     -- normal vector
72         t = V (-y n) (x n)                         -- tangent vector
73         v1 = (e * velocity b1 `dot` t) *^ t ^+^ (mass b1 * velocity b1 `dot` n + mass b2 * velocity b2 `dot` n + e * mass b2 * (velocity b2 ^-^ velocity b1) `dot` n) / (mass b1 + mass b2) *^ n
74         v2 = (e * velocity b2 `dot` t) *^ t ^+^ (mass b2 * velocity b2 `dot` n + mass b1 * velocity b1 `dot` n + e * mass b1 * (velocity b1 ^-^ velocity b2) `dot` n) / (mass b1 + mass b2) *^ n
75         e1 = 0.5 * mass b1 * v1 `dot` v1 - kineticEnergy b1
76         e2 = 0.5 * mass b2 * v2 `dot` v2 - kineticEnergy b2
77         e = elasticity
78
79 outside :: Ball -> Ball -> Bool
80 outside b1 b2 =                         v `dot` v > (radius b2 + radius b1) ^ 2
81   where v = position b2 ^-^ position b1
82
83 inside :: Ball -> Ball -> Bool
84 inside b1 b2 = radius b1 < radius b2 && v `dot` v < (radius b2 - radius b1)^2
85   where v = position b2 ^-^ position b1
86
87 intersects b1 b2 = not (b1 `outside` b2) && not (b1 `inside` b2) && not (b2 `inside` b1)
88
89 approaching :: Ball -> Ball -> Bool
90 approaching b1 b2 = (position b2 ^-^ position b1) `dot` (velocity b1 ^-^ velocity b2) > 0
91
92 data World = World { nextID :: ID, balls :: Map ID Ball, now :: R, collisions :: Set Collision }
93   deriving (Eq, Ord, Read, Show)
94
95 worldEnergy :: World -> R
96 worldEnergy w = sum . map kineticEnergy . M.elems $ balls w
97
98 initialWorld :: World
99 initialWorld = World { nextID = firstID, balls = M.empty, now = 0, collisions = S.empty }
100
101 insertBall :: Ball -> World -> IO World
102 insertBall b w = do
103   let i = nextID w
104   createBall i b
105   return $ insertBallAs i b w{ nextID = i + 1 }
106
107 insertBallAs :: ID -> Ball -> World -> World
108 insertBallAs i b w = w { balls = M.insert i b (balls w), collisions = S.union (collisions w) cs }
109   where cs = S.fromList . map (offset (now w)) . concatMap (collides (i, b)) . M.assocs $ balls w
110
111 collides :: (ID, Ball) -> (ID, Ball) -> [Collision]
112 collides (i, bi) (j, bj)
113   | bi `outside` bj = let
114    o| not (bi `approaching` bj) = []
115     | disc1 >  0 = [Collision t1 i j, Collision t2 i j]
116     | disc1 == 0 = [Collision t1 i j]
117     | otherwise = []
118    in o
119   | bi `inside` bj || bj `inside` bi = let
120    o| disc2 >  0 = [Collision t3 i j, Collision t4 i j]
121     | disc2 == 0 = [Collision t3 i j]
122     | otherwise = []
123    in o
124   | bi `approaching` bj = let -- overlapped
125    o| disc2 >  0 = [Collision t3 i j, Collision t4 i j]
126     | disc2 == 0 = [Collision t3 i j]
127     | otherwise = []
128    in o
129   | otherwise = []
130   where
131     ri = radius bi
132     rj = radius bj
133     pi = position bi
134     pj = position bj
135     vi = velocity bi
136     vj = velocity bj
137     a = vi `dot` vi + vj `dot` vj - 2 * vi `dot` vj
138     b = vi `dot` pj + pi `dot` vj - pi `dot` vi - pj `dot` vj
139     c = a
140     d1 = pi `dot` pi + pj `dot` pj - 2 * pi `dot` pj - (ri + rj)^2
141     d2 = pi `dot` pi + pj `dot` pj - 2 * pi `dot` pj - (ri - rj)^2
142     disc1 = (-2 * b)^2 - 4 * c * d1
143     t1 = 0.5 * (2 * b - sqrt disc1) / a
144     t2 = 0.5 * (2 * b + sqrt disc1) / a
145     disc2 = (-2 * b)^2 - 4 * c * d2
146     t3 = 0.5 * (2 * b - sqrt disc2) / a
147     t4 = 0.5 * (2 * b + sqrt disc2) / a
148
149 update :: R -> World -> IO World
150 update t w
151   | t <= now w = return w
152   | otherwise  = update t =<< update1 t w
153
154 update1 :: R -> World -> IO World
155 update1 t w
156   | S.null future  = return $ w { now = t, balls = newBalls t, collisions = S.empty }
157   | t < time event = return $ w { now = t, balls = newBalls t, collisions = future }
158   | otherwise      = do
159       triggerBall (ball1 event) (realToFrac . sqrt $ 10 * ring b1 * radius b1^4)
160       triggerBall (ball2 event) (realToFrac . sqrt $ 10 * ring b2 * radius b2^4)
161       return $ w'{ now = now w' + et, balls = M.map (line et) (balls w'), collisions = S.filter (\e -> time e /= time event) (collisions w') }
162   where
163     et = 1e-9
164     w' = insertBallAs (ball1 event) b1 . insertBallAs (ball2 event) b2 $ w { now = time event, balls = M.delete (ball1 event) . M.delete (ball2 event) $ newBallsE, collisions = S.filter (not . invalid event) future' }
165     (past, future) = S.split (Collision (now w) (firstID-1) (firstID-1)) (collisions w)
166     (event, future') = S.deleteFindMin future
167     invalid c d = involves (ball1 c) d || involves (ball2 c) d
168     line dt b = b { position = position b ^+^ velocity b ^* dt, ring = ring b * decay ** dt }
169     newBalls tt = M.map (line (tt - now w)) (balls w)
170     newBallsE = newBalls (time event)
171     (b1, b2)
172       | (balls w M.! ball1 event) `outside` (balls w M.! ball2 event) = collideOutside (newBallsE M.! ball1 event) (newBallsE M.! ball2 event)
173       | (balls w M.! ball1 event) `inside`  (balls w M.! ball2 event) = collideInside  (newBallsE M.! ball1 event) (newBallsE M.! ball2 event)
174       | (balls w M.! ball2 event) `inside`  (balls w M.! ball1 event) = collideInside  (newBallsE M.! ball2 event) (newBallsE M.! ball1 event)
175       | otherwise = error $ "spurious collision event: " ++ show event
176
177 reshape worldRef vp@(Size w h) = do
178   let s = 1.75
179       (x0,x1,y0,y1) = if w > h then let v = s * fromIntegral h / fromIntegral w in (-s,s,-v,v)
180                                else let v = s * fromIntegral w / fromIntegral h in (-v,v,-s,s)
181   viewport $= (Position 0 0, vp)
182   matrixMode $= Projection
183   loadIdentity
184   ortho x0 x1 y0 y1 (-1) 1
185   matrixMode $= Modelview 0
186   loadIdentity
187   postRedisplay Nothing
188
189 display worldRef = do
190   w <- readIORef worldRef
191   clear [ColorBuffer]
192   lineWidth $= 1
193   renderPrimitive Lines $ do
194     mapM_ drawBall $ M.elems (balls w)
195   swapBuffers
196   
197 drawBall b = do
198   let n = 144 :: R
199       r = radius b
200       p = position b
201       d = mass b / radius b ^ 2
202       a = (((d - minDensity) / (maxDensity - minDensity))`min` 1) * (pi / 2)
203       m k = 1 + 5 * ring b * cos (16 * 2 * pi * k / n) / radius b
204   color $ Color3 (sin a `max` 0) (cos a `max` 0) (1 :: R)
205   forM_ [1 .. n] (\i -> do
206     vertex $ Vertex2 (r * m (i - 1) * cos (2 * pi * (i - 1) / n) + x p) (r * m (i - 1) * sin (2 * pi * (i - 1) / n) + y p)
207     vertex $ Vertex2 (r * m  i      * cos (2 * pi *  i      / n) + x p) (r * m  i      * sin (2 * pi *  i      / n) + y p))
208
209 physics dt worldRef = do
210   w <- readIORef worldRef
211   writeIORef worldRef =<< update (now w + fromIntegral dt / 1000) w
212   addTimerCallback dt $ physics dt worldRef
213   postRedisplay Nothing
214
215 spawn dt worldRef = do
216   w <- readIORef worldRef
217   let e = worldEnergy w
218   when (e < realToFrac spawnThreshold) $ do
219     d <- randomRIO (realToFrac minDensity, realToFrac maxDensity :: Double)
220     den <- randomRIO (2, 16 :: Int)
221     num <- randomRIO (1, den)
222     let r = fromIntegral num / fromIntegral den
223     pr <- randomRIO (0.0, 1.0)
224     pa <- randomRIO (-pi, pi)
225     va <- randomRIO (-pi, pi)
226     let m = d * r^2
227         px = pr * cos pa
228         py = pr * sin pa
229         vr = sqrt $ 2 * spawnEnergy / m
230         vx = vr * cos va
231         vy = vr * sin va
232         f = realToFrac :: Double -> R
233         b = Ball { mass = f m, radius = f r, position = V (f px) (f py), velocity = V (f vx) (f vy), ring = 0 }
234     when (not (any (intersects b) (M.elems (balls w)))) $ do
235       w' <- insertBall b w
236       writeIORef worldRef w'
237   addTimerCallback dt $ spawn dt worldRef
238   postRedisplay Nothing
239
240 main :: IO ()
241 main = do
242   initialWindowSize $= Size 788 576
243   initialDisplayMode $= [RGBAMode, DoubleBuffered]
244   (_,args) <- getArgsAndInitialize
245   createWindow "bowntz"
246   withSC3 (\fd -> send fd (g_new [(1, AddToTail, 0)])) -- new group 1 under root 0 group
247   withSC3 ballSynth
248   w <- insertBall (Ball{ mass = 1e6, radius = 1, position = V 0 0, velocity = V 0 0, ring = 0 }) initialWorld
249   worldRef <- newIORef w
250   displayCallback $= display worldRef
251   reshapeCallback $= Just (reshape worldRef)
252   addTimerCallback (floor$1000/60) $ physics (floor$1000/60) worldRef
253   addTimerCallback 125 $ spawn 125 worldRef
254   mainLoop
255
256 createBall id b = withSC3 (\fd -> send fd (s_new "Ball" (fromIntegral id) AddToTail 1 [("freq", realToFrac $ 50 / radius b)]))
257 triggerBall id t = withSC3 (\fd -> send fd (n_set1 (fromIntegral id) "t_amp" t))
258
259 ballSynth fd = do let f = control IR "freq" 0
260                       a = tr_control "t_amp" 0
261                       d = out 0 (fSinOsc AR (mce [f, f]) 0 * (SC.decay a 0.7))
262                   send fd (d_recv (synthdef "Ball" d))
263                   wait fd "/done"
264 minDensity = 2
265 maxDensity = 5
266 spawnThreshold = 0.02
267 spawnEnergy = spawnThreshold / 2
268 decay = 0.0001
269 elasticity = 0.995
270