{-

    Bowntz -- an audio-visual pseudo-physical simulation of colliding circles
    Copyright (C) 2010,2013,2015,2016 Claude-Heiland-Allen <claude@mathr.co.uk>

    This program is free software: you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with this program.  If not, see <http://www.gnu.org/licenses/>.

-}

import Graphics.UI.GLUT hiding (position)
import Sound.OSC
import Sound.SC3 hiding (decay, line, ID)
import qualified Sound.SC3 as SC

import Control.Monad(forM_, when)
import Data.IORef
import System.Random

import qualified Data.Map as M
import Data.Map(Map)
import qualified Data.Set as S
import Data.Set(Set)
import Debug.Trace

type ID = Integer

firstID :: ID
firstID = 1000

type R = GLdouble

data V = V { x :: R, y :: R }
  deriving (Eq, Ord, Read, Show)

infixl 8 `dot`
infixl 7 *^,^*,^/
infixl 6 ^+^,^-^

(^/) :: V -> R -> V
u ^/ t = V (x u / t) (y u / t)

(^*) :: V -> R -> V
u ^* t = V (x u * t) (y u * t)

(*^) :: R -> V -> V
s *^ v = V (s * x v) (s * y v)

(^+^) :: V -> V -> V
u ^+^ v = V (x u + x v) (y u + y v)

(^-^) :: V -> V -> V
u ^-^ v = V (x u - x v) (y u - y v)

dot :: V -> V -> R
dot u v = x u * x v + y u * y v

data Ball = Ball { mass :: R, radius :: R, position :: V, velocity :: V, ring :: R }
  deriving (Eq, Ord, Read, Show)

data Collision = Collision { ctime :: R, ball1 :: ID, ball2 :: ID }
  deriving (Eq, Ord, Read, Show)

kineticEnergy :: Ball -> R
kineticEnergy b = 0.5 * mass b * velocity b `dot` velocity b

involves :: ID -> Collision -> Bool
involves i c = i == ball1 c || i == ball2 c

offset :: R -> Collision -> Collision
offset t c = c { ctime = ctime c + t }

collideOutside :: Ball -> Ball -> (Ball, Ball)
collideOutside b1 b2 = collideXside (radius b2 + radius b1) b1 b2

collideInside :: Ball -> Ball -> (Ball, Ball)
collideInside b1 b2 = collideXside (abs $ radius b2 - radius b1) b1 b2

collideXside :: R -> Ball -> Ball -> (Ball, Ball)
collideXside r b1 b2 = (b1 { velocity = v1, ring = ring b1 + abs e1 }, b2 { velocity = v2, ring = ring b2 + abs e2 })
  where n = (position b1 ^-^ position b2) ^/ r     -- normal vector
        t = V (-y n) (x n)                         -- tangent vector
        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
        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
        e1 = 0.5 * mass b1 * v1 `dot` v1 - kineticEnergy b1
        e2 = 0.5 * mass b2 * v2 `dot` v2 - kineticEnergy b2
        e = elasticity

outside :: Ball -> Ball -> Bool
outside b1 b2 =                         v `dot` v > (radius b2 + radius b1) ^ (2::Int)
  where v = position b2 ^-^ position b1

inside :: Ball -> Ball -> Bool
inside b1 b2 = radius b1 < radius b2 && v `dot` v < (radius b2 - radius b1) ^ (2::Int)
  where v = position b2 ^-^ position b1

intersects :: Ball -> Ball -> Bool
intersects b1 b2 = not (b1 `outside` b2) && not (b1 `inside` b2) && not (b2 `inside` b1)

approaching :: Ball -> Ball -> Bool
approaching b1 b2 = (position b2 ^-^ position b1) `dot` (velocity b1 ^-^ velocity b2) > 0

data World = World { nextID :: ID, balls :: Map ID Ball, now :: R, collisions :: Set Collision }
  deriving (Eq, Ord, Read, Show)

worldEnergy :: World -> R
worldEnergy w = sum . map kineticEnergy . M.elems $ balls w

initialWorld :: R -> World
initialWorld t = World { nextID = firstID, balls = M.empty, now = t, collisions = S.empty }

insertBall :: Ball -> World -> IO World
insertBall b w = do
  let i = nextID w
  createBall i b
  return $ insertBallAs i b w{ nextID = i + 1 }

insertBallAs :: ID -> Ball -> World -> World
insertBallAs i b w = w { balls = M.insert i b (balls w), collisions = S.union (collisions w) cs }
  where cs = S.fromList . map (offset (now w)) . concatMap (collides (i, b)) . M.assocs $ balls w

collides :: (ID, Ball) -> (ID, Ball) -> [Collision]
collides (i, bi) (j, bj)
  | bi `outside` bj = let
   o| not (bi `approaching` bj) = []
    | disc1 >  0 = [Collision t1 i j, Collision t2 i j]
    | disc1 == 0 = [Collision t1 i j]
    | otherwise = []
   in o
  | bi `inside` bj || bj `inside` bi = let
   o| disc2 >  0 = [Collision t3 i j, Collision t4 i j]
    | disc2 == 0 = [Collision t3 i j]
    | otherwise = []
   in o
  | bi `approaching` bj = let -- overlapped
   o| disc2 >  0 = [Collision t3 i j, Collision t4 i j]
    | disc2 == 0 = [Collision t3 i j]
    | otherwise = []
   in o
  | otherwise = []
  where
    ri = radius bi
    rj = radius bj
    qi = position bi
    qj = position bj
    vi = velocity bi
    vj = velocity bj
    a = vi `dot` vi + vj `dot` vj - 2 * vi `dot` vj
    b = vi `dot` qj + qi `dot` vj - qi `dot` vi - qj `dot` vj
    c = a
    d1 = qi `dot` qi + qj `dot` qj - 2 * qi `dot` qj - (ri + rj)^(2::Int)
    d2 = qi `dot` qi + qj `dot` qj - 2 * qi `dot` qj - (ri - rj)^(2::Int)
    disc1 = (-2 * b)^(2::Int) - 4 * c * d1
    t1 = 0.5 * (2 * b - sqrt disc1) / a
    t2 = 0.5 * (2 * b + sqrt disc1) / a
    disc2 = (-2 * b)^(2::Int) - 4 * c * d2
    t3 = 0.5 * (2 * b - sqrt disc2) / a
    t4 = 0.5 * (2 * b + sqrt disc2) / a

update :: R -> World -> IO World
update t w
  | t <= now w = return w
  | otherwise  = update t =<< update1 t w

update1 :: R -> World -> IO World
update1 t w
  | S.null future  = return $ w { now = t, balls = newBalls t, collisions = S.empty }
  | t < ctime event = return $ w { now = t, balls = newBalls t, collisions = future }
  | otherwise      = do
      triggerBall (ctime event) (ball1 event) (realToFrac $ ring b1 * mass b1)
      triggerBall (ctime event) (ball2 event) (realToFrac $ ring b2 * mass b2)
      return $ w'{ collisions = S.filter (\e -> ctime e /= ctime event) (collisions w') }
  where
    w' = insertBallAs (ball1 event) b1 . insertBallAs (ball2 event) b2 $ w { now = ctime event, balls = M.delete (ball1 event) . M.delete (ball2 event) $ newBallsE, collisions = S.filter (not . invalid event) future' }
    (_past, future) = S.split (Collision (now w) (firstID-1) (firstID-1)) (collisions w)
    (event, future') = S.deleteFindMin future
    invalid c d = involves (ball1 c) d || involves (ball2 c) d
    line dt b = b { position = position b ^+^ velocity b ^* dt, ring = ring b * decay ** dt }
    newBalls tt = M.map (line (tt - now w)) (balls w)
    newBallsE = newBalls (ctime event)
    (b1, b2)
      | (balls w M.! ball1 event) `outside` (balls w M.! ball2 event) = collideOutside (newBallsE M.! ball1 event) (newBallsE M.! ball2 event)
      | (balls w M.! ball1 event) `inside`  (balls w M.! ball2 event) = collideInside  (newBallsE M.! ball1 event) (newBallsE M.! ball2 event)
      | (balls w M.! ball2 event) `inside`  (balls w M.! ball1 event) = collideInside  (newBallsE M.! ball2 event) (newBallsE M.! ball1 event)
      | otherwise = trace ("spurious collision event: " ++ show event) (newBallsE M.! ball1 event, newBallsE M.! ball2 event)

reshape :: IORef World -> Size -> IO ()
reshape _worldRef vp@(Size w h) = do
  let s = 1.75
      (x0,x1,y0,y1) = if w > h then let v = s * fromIntegral h / fromIntegral w in (-s,s,-v,v)
                               else let v = s * fromIntegral w / fromIntegral h in (-v,v,-s,s)
  viewport $= (Position 0 0, vp)
  matrixMode $= Projection
  loadIdentity
  ortho x0 x1 y0 y1 (-1) 1
  matrixMode $= Modelview 0
  loadIdentity
  postRedisplay Nothing

display :: IORef World -> IO ()
display worldRef = do
  w <- readIORef worldRef
  clear [ColorBuffer]
  lineWidth $= 1
  renderPrimitive Lines $ do
    mapM_ drawBall $ M.elems (balls w)
  swapBuffers

drawBall :: Ball -> IO ()
drawBall b = do
  let n = 144 :: R
      r = radius b
      p = position b
      d = mass b / radius b ^ (2::Int)
      a = (((d - minDensity) / (maxDensity - minDensity))`min` 1) * (pi / 2)
      m k = 1 + 5 * ring b * cos (16 * 2 * pi * k / n) / radius b
  color $ if d > maxDensity then Color3 0 0 1 else Color3 (sin a `max` 0) (cos a `max` 0) 0
  forM_ [1 .. n] (\i -> do
    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)
    vertex $ Vertex2 (r * m  i      * cos (2 * pi *  i      / n) + x p) (r * m  i      * sin (2 * pi *  i      / n) + y p))

physics :: Timeout -> IORef World -> IO ()
physics dt worldRef = do
  reallyNow <- time
  w <- readIORef worldRef
  writeIORef worldRef =<< update (realToFrac reallyNow + fromIntegral dt / 1000) w
  addTimerCallback dt $ physics dt worldRef
  postRedisplay Nothing

spawn :: Timeout -> IORef World -> IO ()
spawn dt worldRef = do
  w <- readIORef worldRef
  let e = worldEnergy w
  when (e < realToFrac spawnThreshold) $ do
    d  <- realToFrac `fmap` randomRIO (realToFrac minDensity, realToFrac maxDensity :: Double)
    r  <- realToFrac `fmap` randomRIO (realToFrac minSize, realToFrac maxSize :: Double)
    pr <- realToFrac `fmap` randomRIO (0.0, 1.0   :: Double)
    pa <- realToFrac `fmap` randomRIO (-pi, pi    :: Double)
    va <- realToFrac `fmap` randomRIO (-pi, pi    :: Double)
    let m = d * r^(2::Int)
        px = pr * cos pa
        py = pr * sin pa
        vr = sqrt $ 2 * spawnEnergy / m
        vx = vr * cos va
        vy = vr * sin va
        b = Ball { mass = m, radius = r, position = V px py, velocity = V vx vy, ring = 0 }
    when (not (any (intersects b) (M.elems (balls w)))) $ do
      w' <- insertBall b w
      writeIORef worldRef w'
  addTimerCallback dt $ spawn dt worldRef
  postRedisplay Nothing

main :: IO ()
main = do
  (_,_args) <- getArgsAndInitialize
  initialWindowSize $= Size 788 576
  initialDisplayMode $= [RGBAMode, DoubleBuffered]
  _ <- createWindow "bowntz"
  withSC3 (sendMessage (g_new [(1, AddToTail, 0)])) -- new group 1 under root 0 group
  _ <- withSC3 ballSynth
  t0 <- time
  w <- insertBall (Ball{ mass = 1e6, radius = 1, position = V 0 0, velocity = V 0 0, ring = 0 }) $ initialWorld (realToFrac t0)
  worldRef <- newIORef w
  displayCallback $= display worldRef
  reshapeCallback $= Just (reshape worldRef)
  let mspf = floor (1000 / 60 :: Double)
  addTimerCallback mspf $ physics mspf worldRef
  addTimerCallback 125 $ spawn 125 worldRef
  mainLoop

createBall :: ID -> Ball -> IO ()
createBall i b = withSC3 (sendMessage (s_new "Ball" (fromIntegral i) AddToTail 1 [("freq", realToFrac $ 50 / radius b)]))

triggerBall :: R -> ID -> Double -> IO ()
triggerBall t i v = withSC3 (sendBundle (bundle (realToFrac (t + latency)) [n_set1 (fromIntegral i) "t_amp" v]))

ballSynth :: Connection UDP Packet
ballSynth = do let f = control IR "freq" 0
                   a = tr_control "t_amp" 0
                   d = out 0 (fSinOsc AR (mce [f, f]) 0 * (SC.decay a 0.7))
               sendMessage (d_recv (synthdef "Ball" d))
               waitAddress "/done"

minSize, maxSize, minDensity, maxDensity, spawnThreshold, spawnEnergy, decay, elasticity, latency :: R
minSize = 0.02
maxSize = 0.8
minDensity = 40
maxDensity = 60
spawnThreshold = 0.1
spawnEnergy = spawnThreshold / 8
decay = 0.0001
elasticity = 0.99
latency = 1 / 30