{-# LANGUAGE TypeFamilies #-}
import Graphics.Rendering.Chart.Easy
import Data.Default.Class
import Control.Lens
import Control.Monad
import Control.Arrow
import Data.VectorSpace

type ℝ = Double
type Time = Double

s :: Time -> ℝ
s t = fromIntegral (round t) - t

signalPlot :: String -> [(Time,ℝ)] -> EC (Layout ℝ ℝ) ()
signalPlot capt sig = do
     layout_x_axis . laxis_title .= "t"
     plot (line capt [takeWhile ((<3).fst) sig])

toRenderable $ signalPlot "s(t)" [(t,s t) | t <- [0,0.01..]]

rk₄ :: (VectorSpace y, Scalar y ~ ℝ)
    => (Time -> y -> y) -- the function r_t(t,r)
    -> Time             -- time-step for the solver
    -> [(Time, y)]      -- sequence of solution snapshots [r(t)]
rk₄ f h = go 0 zeroV
 where go t y = (t,y) : go (t+h) (y ^+^ (h/6)*^(k₁ ^+^ 2*^k₂ ^+^ 2*^k₃ ^+^ k₄))
        where k₁ = f t y
              k₂ = f (t + h/2) (y ^+^ (h/2)*^k₁)
              k₃ = f (t + h/2) (y ^+^ (h/2)*^k₂)
              k₄ = f (t + h) (y ^+^ h*^k₃)

toRenderable $ forM [0 .. 8] $ \η -> signalPlot ("η = "++show η)
          $ rk₄ (\t r -> η * (s t - r)) 0.01

toRenderable $ forM [0, 5 .. 30] $ \η -> signalPlot ("η = "++show η) $ second fst
          <$> rk₄ (\t (r,u) -> (u, η^2 * (s t - r) - η*u)) 0.005