--- Imports ---


-- Goal --

import Goal.Core
import Goal.Geometry

import Goal.Simulation

import qualified Data.Vector.Storable as C

--- Script ---


main = do

    -- Generation --

    -- We can simulate sin either as non-autonomous or second order autonomous
    let sin' t _ = C.singleton $ cos t
        vsin' x = C.fromList [x C.! 1,-x C.! 0]
        exp' = id

        t0 = 0
        tf = 10
        dt1 = 2
        dt2 = 1
        dt3 = 0.1
        ts1 = [t0,t0+dt1..tf]
        ts2 = [t0,t0+dt2..tf]
        ts3 = [t0,t0+dt3..tf]

        sx0 = euclideanPoint [0]
        vx0 = euclideanPoint [0,1]
        ex0 = euclideanPoint [1]

        sinMealy = nonAutonomousODE sin' t0 sx0
        sinMealyEuler = nonAutonomousODEEuler sin' t0 sx0
        vsinMealy = autonomousODE vsin' vx0
        expMealy = autonomousODE exp' ex0
        expMealyEuler = autonomousODEEuler exp' ex0

        ssimulator ts mly = zip ts $ coordinate 0 <$> stream mly ts
        vsimulator ts mly = zip ts $ coordinate 0 <$> stream mly ts
        esimulator ts mly = zip ts $ coordinate 0 <$> stream mly ts

    -- Plots --

    -- Sin

    let sinrnbl = toRenderable . execEC $ do

            layout_title .= "Sin Wave (dt = {1,0.1})"

            plot . liftEC $ do
                plot_lines_style .= dashedLine 3 [10,5] (opaque black)
                plot_lines_title .= "True"
                plot_lines_values .= [zip ts3 $ sin <$> ts3]

            plot . liftEC $ do
                plot_lines_style .= solidLine 2 (opaque blue)
                plot_lines_title .= "RK4"
                plot_lines_values .=
                  [ ssimulator ts2 sinMealy, ssimulator ts3 sinMealy ]

            plot . liftEC $ do
                plot_lines_style .= solidLine 2 (opaque red)
                plot_lines_title .= "Euler"
                plot_lines_values .=
                  [ ssimulator ts2 sinMealyEuler, ssimulator ts3 sinMealyEuler ]

            plot . liftEC $ do
                plot_lines_style .= solidLine 2 (opaque purple)
                plot_lines_title .= "RK4 2nd Order"
                plot_lines_values .=
                  [ vsimulator ts2 vsinMealy , vsimulator ts3 vsinMealy ]

    -- Exponential

    let exprnbl = toRenderable . execEC $ do

            layout_title .= "Exponential Function (dt = {2,1,0.1})"

            plot . liftEC $ do
                plot_lines_style .= dashedLine 3 [10,5] (opaque black)
                plot_lines_title .= "True"
                plot_lines_values .= [zip ts3 $ exp <$> ts3]

            plot . liftEC $ do
                plot_lines_style .= solidLine 3 (opaque blue)
                plot_lines_title .= "RK4"
                plot_lines_values .= [ esimulator ts1 expMealy, esimulator ts2 expMealy, esimulator ts3 expMealy ]

            plot . liftEC $ do
                plot_lines_style .= solidLine 3 (opaque red)
                plot_lines_title .= "Euler"
                plot_lines_values .=
                    [ esimulator ts1 expMealyEuler, esimulator ts2 expMealyEuler, esimulator ts3 expMealyEuler ]

    -- IO

    renderableToAspectWindow False 800 600 . gridToRenderable . weights (1,1) . tallBeside sinrnbl $ tval exprnbl


--- Extra Functions ---


autonomousODEEuler f' p0 =
    accumulateFunction accumulator (0,p0)
      where accumulator t' (t,p) =
                let dt = t' - t
                    p' = p <+> fromCoordinates (manifold p) (stepEuler f' dt $ coordinates p)
                 in (p',(t',p'))

nonAutonomousODEEuler f' t0 p0 =
    accumulateFunction accumulator (t0,p0)
      where accumulator t' (t,p) =
                let dt = t' - t
                    p' = p <+> fromCoordinates (manifold p) (stepEuler' f' t dt $ coordinates p)
                 in (p',(t',p'))