{-# LANGUAGE TypeFamilies,FlexibleContexts,Arrows #-}

--- Imports ---


-- Goal --

import Pendulum

import Goal.Core
import Goal.Geometry
import Goal.Probability
import Goal.Simulation


--- Program ---


-- Globals --

qdq0 = fromList (Bundle pndl) [1,0]

-- Functions --

pathToRenderable tqdqs =

    let phslyt = execEC $ do

            vectorFieldLayout

            layout_title .= "Phase Space"

            plot . liftEC $ do
                plot_lines_title .= "Phase"
                plot_lines_values .= [[(coordinate 0 qdq,coordinate 1 qdq) | (_,qdq) <- tqdqs]]
                plot_lines_style .= solidLine 3 (opaque black)

        enrlyt = execEC $ do

            layout_title .= "Energy"

            layout_x_axis . laxis_title .= "Time"

            layout_y_axis . laxis_generate .= scaledAxis def (0,20)
            layout_y_axis . laxis_title .= "Joules"

            let (tks,tus,tms) = unzip3 [ let (k,u) = mechanicalEnergy fg qdq in ((t,k),(t,u),(t,k+u)) | (t,qdq) <- tqdqs ]

            plot . liftEC $ do
                plot_lines_title .= "Kinetic"
                plot_lines_values .= [tks]
                plot_lines_style .= solidLine 3 (opaque red)

            plot . liftEC $ do
                plot_lines_title .= "Potential"
                plot_lines_values .= [tus]
                plot_lines_style .= solidLine 3 (opaque blue)

            plot . liftEC $ do
                plot_lines_title .= "Total"
                plot_lines_values .= [tms]
                plot_lines_style .= solidLine 3 (opaque purple)

        knmlyt = execEC $ do

            layout_title .= "Kinematics"

            let (_,qdq) = head tqdqs
                [tht] = listCoordinates $ position qdq
                (x,y) = (l * sin tht, -l * cos tht)

            layout_x_axis . laxis_generate .= scaledAxis def (-2,2)

            layout_y_axis . laxis_generate .= scaledAxis def (-2,2)

            plot . liftEC $ do
                plot_lines_values .= [[(0,0),(x,y)]]
                plot_lines_style .= solidLine 3 (opaque black)

            plot . liftEC $ do
                plot_points_values .= [(0,0)]
                plot_points_style .= hollowCircles 6 4 (opaque black)

            plot . liftEC $ do
                plot_points_values .= [(x,y)]
                plot_points_style .= filledCircles 6 (opaque black)


     in toRenderable . weights (1,1) . wideAbove enrlyt $ tval knmlyt .|. tval phslyt

-- Main --

main :: IO ()
main = do

    lngvn <- runWithSystemRandom $ langevinFlow f sgma qdq0

    let tivl = 2.0
        fps = round $ recip dt

        mly = proc t -> do
            x <- lngvn -< t
            txs <- trajectoryWindow tivl -< (t,x)
            returnA -< pathToRenderable txs

    initGUI

    apst <- newAnimationPost 1.2 (Just fps)

    forkIO $ streamM_ mly (postRenderable apst) ts

    mainGUI