-- Baed on: program: S.A.R.A.H. road simulator
--          author: Maurício C. Antunes
--          from Gtk2Hs library demos
-- license: public domain
-- author: Chuck Anderson  anderson@cs.colostate.edu

-- make executable and run:
--   ghc --make Pole.hs
--   ./Pole

import Graphics.UI.Gtk
import Graphics.UI.Gtk.Abstract.Widget
import Graphics.Rendering.Cairo 
import Data.IORef
import System.Time
import System.IO

main = do

  initGUI
  frame <- windowNew
  set frame [ containerBorderWidth := 10, 
              windowTitle := "Inverted Pendulum",
              windowWindowPosition := WinPosCenter]
  onDestroy frame mainQuit
  vbox <- vBoxNew True 5
  canvas <- drawingAreaNew

  canvas `onSizeRequest` return (Requisition 1000 500)
  frame `containerAdd` vbox
  boxPackStart vbox canvas PackGrow 0

  state <- newIORef [0,0,0,0 :: Double]
  
  animate canvas 10 state updateState drawState

  widgetShowAll frame
  mainGUI

animate canvas rateInMilliseconds state updateState drawState = do

  last_time <- (newIORef =<< getClockTime)
  buttonDown <- newIORef (Nothing :: Maybe MouseButton)

  canvas `onButtonPress`  
    \evt -> do writeIORef buttonDown (Just (eventButton evt))
               return True

  canvas `onButtonRelease` 
    \evt -> do writeIORef buttonDown Nothing
               return True

  (flip timeoutAdd) rateInMilliseconds $ do
    (TOD s1 ps1) <- readIORef last_time
    (TOD s2 ps2) <- getClockTime
    writeIORef last_time (TOD s2 ps2)
    -- how much time since last update?
    timestep <- return $ 1e-12 * fromInteger(10^12*(s2-s1)+ps2-ps1)

    but <- readIORef buttonDown

    bu <- return $ case but of
                Nothing -> "nothing"
                Just LeftButton -> "left"
                Just RightButton -> "right"
                Just MiddleButton -> "middle"

    -- Update pole state
     
    s <- readIORef state
    writeIORef state $ updateState s but timestep

    -- putStrLn $ "Time is " ++ (show timestep) ++ "  x is " ++ (show oldx)
    --           ++ " " ++ (show step) ++ " " ++ bu

    drawWindow <- widgetGetDrawWindow canvas
    (w,h) <- drawableGetSize drawWindow
    drawWindowBeginPaintRect drawWindow (Rectangle 0 0 w h)
    st <- readIORef state
    renderWithDrawable drawWindow (drawState st but w h)
    drawWindowEndPaint drawWindow

    return True

-- Two functions that define the physics, and draw the state

-- Parameters

cartMass = 0.2
poleMass = 2.0
poleLength = 5.0
forceMag = 10.0
tau = 0.005  -- time step
fricCart = 0.5
fricPole = 0.4


updateState state button timestep = case button of
                                      Just MiddleButton -> [0,0,0,0]
                                      _ -> newstate
    where newstate = [newX, newXDot, newAngle, newAngleDot]
          [x,xDot,angle,angleDot] = state
          action = case button of
                   Just LeftButton -> (-1)
                   Just RightButton -> 1
                   _ -> 0

          totalMass  = cartMass + poleMass
          halfPole   = 0.5 * poleLength
          poleMassLength   = halfPole * poleMass
          force = forceMag * action
          angleDotSq = angleDot * angleDot
          common = (force + poleMassLength * angleDotSq * (sin angle)
                    - fricCart * (signum xDot)) / totalMass
          angleDDot = (-9.8 * (sin angle) - (cos angle) * common
		       - fricPole * angleDot / poleMassLength) /
	              (halfPole * (4.0/3.0 - poleMass * (cos angle) * (cos angle) /
		                           totalMass))
          xDDot = common - poleMassLength * angleDDot * (cos angle) / totalMass

          newX = x + timestep * xDot
          newXDot = xDot + timestep * xDDot
          newAngle = angle + timestep * angleDot
          newAngleDot = angleDot + timestep * angleDDot

drawState state button w h = let centerx = (fromIntegral w) / 2
                                 centery = (fromIntegral h) / 2
                                 width = fromIntegral w
                                 [x,_,angle,_] = state
                                 action = case button of
                                            Just LeftButton -> (-1)
                                            Just RightButton -> 1
                                            _ -> 0
                                 text1 = "Left Button: push left    Right Button: push right"
                                 text2 = "Middle Button: PANIC"
                             in do
  -- track
  setSourceRGB 0.5 0.5 0.5
  rectangle 0 (centery-10) width 20
  rectangle 0 (centery-30) 20 60
  rectangle (width-20) (centery-30) 20 60
  fillPreserve
  stroke

  -- cart
  setSourceRGB 0.2 0.4 0.8
  rectangle (centerx + 10*x - 40) (centery-24) 80.0 48.0
  fillPreserve
  stroke

  -- arrow
  setSourceRGB 0.8 0.8 0.0
  setLineWidth 3
  moveTo (centerx + 10*x) (centery + 0)
  relLineTo (35 * action) 0
  relLineTo (-8 * action) (6 * action)
  relMoveTo (8 * action) (-6 * action)
  relLineTo (-8 * action) (-6 * action)
  stroke

  -- pole
  setSourceRGB 0.8 0.3 0.1
  setLineWidth 10
  setLineCap LineCapRound
  moveTo (centerx + 10*x ) (centery-22)
  relLineTo (sin(angle) * poleLength * 40) (cos(angle) * poleLength * 40)
  stroke

  -- text
  setSourceRGB 0.1 0.1 0.0
  setFontSize 12
  (TextExtents xb _ wtext htext _ _) <- textExtents text1
  moveTo (centerx - wtext/2) (centery + centery/2 + htext/2)
  showText text1
  (TextExtents xb _ wtext htext _ _) <- textExtents text2
  moveTo (centerx - wtext/2) (centery + centery/2 + htext*4)
  setSourceRGB 0.8 0.1 0.0
  showText text2
  stroke