OpenDynamicalVis.hs

{-# LANGUAGE Arrows #-}
{-# LANGUAGE FlexibleContexts #-}
module OpenDynamicalVis where

import OpenDynamical
import Control.Arrow
import Linear.Vector
import Linear.Affine
-- import Control.Concurrent
import Data.Time.Clock
import Diagrams.Prelude hiding ((*^), (^+^))
import Diagrams.Backend.Cairo
import Diagrams.Backend.Gtk
import Data.IORef
import Data.Maybe
import Graphics.UI.Gtk hiding (Size)
import Control.Monad
import Control.Monad.Trans

-- mainG :: SF (V B (N B), Event (V B (N B))) (Diagram B) -> IO ()
mainG :: OD Double (Bool, V2 Double) (Diagram B) -> IO ()
mainG mainOD = do
  void initGUI

  w <- windowNew
  da <- drawingAreaNew
  w `containerAdd` da
  da `widgetAddEvents` [PointerMotionMask]
  widgetShowAll w

  odRef <- newIORef (mempty, mainOD)
  let step a h = do
        modifyIORef' odRef (odStep rk4 a h . snd)
        return False
      draw = do
        (dia, _) <- readIORef odRef
        defaultRender da dia

  void . (w `on` deleteEvent) . liftIO $ do
    widgetHideAll w
    mainQuit
    return True

  void . (da `on` exposeEvent) . liftIO $ True <$ draw

  timeRef <- newIORef Nothing
  posRef <- newIORef zero
  downRef <- newIORef False
  let tick = 5
  void . flip timeoutAdd tick $ do
    -- TODO properly use transform from rendering
    V2 xRaw yRaw <- readIORef posRef
    (dw, dh) <- liftIO $ windowGetSize w
    (dia, _) <- readIORef odRef
    let (P (V2 loX loY), P (V2 hiX hiY)) =
          fromMaybe (P zero, P zero) . getCorners . boundingBox $ dia
        V1 xDia = lerp (xRaw / fromIntegral dw) (V1 hiX) (V1 loX)
        V1 yDia = lerp (yRaw / fromIntegral dh) (V1 loY) (V1 hiY)
        pos = V2 xDia yDia
    down <- readIORef downRef
    t <- getCurrentTime
    prev <- readIORef timeRef
    writeIORef timeRef (Just t)
    let dt0 = realToFrac (maybe 0 (diffUTCTime t) prev)
        dt  = min dt0 (fromIntegral tick / 250)
    -- when (dt /= dt0) . liftIO $ print "skip"
    shouldQuit <- step (down, pos) dt
    when shouldQuit $ do
      widgetHideAll w
      mainQuit
    return (not shouldQuit)
  void . flip timeoutAdd 15 $ True <$ draw
  {-
  let updateInput :: (HasCoordinates e, HasModifier e) => EventM e Bool
      updateInput = do
        (x, y) <- eventCoordinates
        -- liftIO $ print (x, y)
        mods <- eventModifierMouse
        liftIO $ print mods
        liftIO $ writeIORef downRef (not (null mods))
        liftIO $ writeIORef posRef (V2 x y)
        return True
        -}
  void . (da `on` motionNotifyEvent) $ do
    (x, y) <- eventCoordinates
    liftIO $ writeIORef posRef (V2 x y)
    return True
  -- void . (da `on` enterNotifyEvent) $ updateInput
  void . (da `on` leaveNotifyEvent) $ do
    liftIO $ writeIORef downRef False
    return True
  void . (da `on` buttonPressEvent) $ do
    liftIO $ writeIORef downRef True
    return True
  void . (da `on` buttonReleaseEvent) $ do
    liftIO $ writeIORef downRef False
    return True

  mainGUI


{-
spring' :: (Num s, Add v) => s -> v s -> v s -> v s -> OD s a (v s)
spring' k g r0 v0 = OD (P (Pair r0 v0)) d o
  where d _ (P (Pair r v)) = Pair v ((-k) *^ r ^+^ g)
        o _ (P (Pair r _)) = r
        -}

spring ::
  (Num s, Add v, Finite v, KnownNat (Size v)) =>
  s -> v s -> v s -> v s -> OD s (Bool, v s) (v s, v s)
spring k g r0 v0 = proc (down, _) -> do
  rec
    let a = (-k) *^ r ^+^ if down then zero else g
    v <- integralV v0 -< a
    r <- integralV r0 -< v
  returnA -< (r, v)

springGrab ::
  (Num s, Add v, Finite v, KnownNat (Size v)) =>
  s -> v s -> v s -> v s -> OD s (Bool, v s) (v s, v s)
springGrab k g r0 v0 = proc (down, pos) -> do
  rec
    -- let a = (-k) *^ r ^+^ if down then zero else g
    let a = (-k) *^ r ^+^ g
    v <- integralV v0 -< a
    r <- if down
       then returnA -< pos
       else integralV r0 -< v
  returnA -< (r, v)

mainSpring :: IO ()
mainSpring = mainG $ proc inp@(down, _) -> do
  (r, v) <- spring 5.0 (V2 0 (-5)) (V2 0 (-2)) (V2 1 0) -< inp
  returnA -< showOrigin $
    square 5
    <>
    moveTo (P r) (circle 0.1 # fc (if down then blue else red))
    <>
    p2 (0, 0) ~~ P r
    <>
    arrowAt (P r) (v ^/ 2.5)
    {-
    <>
    moveTo (P foo) (circle 0.1 # fc orange)
    -}

doublePendulum :: Floating s =>
  s -> s -> s -> s -> s -> s -> s -> s -> s -> OD s a (V2 s)
doublePendulum l1 l2 m1 m2 g theta1_0 theta2_0 omega1_0 omega2_0 =
  proc _ -> do
    rec
      let d = theta2 - theta1
          accel1 = (m2 * l1 * omega1**2 * sin d * cos d + m2 * g * sin theta2 * cos d + m2 * l2 * omega2**2 * sin d - (m1 + m2) * g * sin theta1) / ((m1 + m2) * l1 - m2 * l1 * (cos d)**2)
          accel2 = (-m2 * l2 * omega2**2 * sin d * cos d + (m1 + m2) * (g * sin theta1 * cos d - l1 * omega1**2 * sin d - g * sin theta2)) / ((m1 + m2) * l2 - m2 * l2 * (cos d)**2)
      omega1 <- integralN omega1_0 -< accel1
      omega2 <- integralN omega2_0 -< accel2
      theta1 <- integralN theta1_0 -< omega1
      theta2 <- integralN theta2_0 -< omega2
    returnA -< V2 theta1 theta2

mainPend :: IO ()
mainPend = mainG (drawPend <$> pendOD)
  where (l1, l2) = (1, 1)
        (m1, m2) = (1, 1)
        pendOD = doublePendulum l1 l2 m1 m2 9.8 2 2 0 0
        drawPend :: V2 Double -> Diagram B
        drawPend (V2 theta1 theta2) = showOrigin $
          let tq = halfTurn <> quarterTurn
              point1 = origin .+^ ((l1, (theta1 @@ rad) <> tq) @@ r2PolarIso)
              point2 = point1 .+^ ((l2, (theta2 @@ rad) <> tq) @@ r2PolarIso)
          in mconcat [
            square 5,
            fromVertices [origin, point1, point2],
            moveTo point1 (circle (sqrt m1 / 10) # fc blue),
            moveTo point2 (circle (sqrt m2 / 10) # fc blue)
          ]


main :: IO ()
main = mainSpring