-- Goal --

import Pendulum

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

-- Qualified --

import qualified System.Directory as D


--- Program ---


-- Globals --

xaxprms = LinearAxisParams (show . round) 5 5
yaxprms = LinearAxisParams (show . round) 5 5

-- Functions --

coordinateLayout n rng (xs,ns,zs) = execEC $ do

    layout_x_axis . laxis_override .= axisGridHide
    layout_y_axis . laxis_generate .= scaledAxis yaxprms rng
    layout_y_axis . laxis_override .= axisGridHide
    layout_x_axis . laxis_generate .= autoScaledAxis xaxprms

    plot . liftEC $ do

        plot_lines_values .= [ zip ts $ coordinate n <$> xs ]
        plot_lines_style .= solidLine 2 (opaque black)

    plot . liftEC $ do

        plot_points_style .= filledCircles 2 (opaque black)
        plot_points_values .=
            zip ts (coordinate (2*n) . potentialMapping <$> (harmoniumTranspose trns >$> ns))

    plot . liftEC $ do

        plot_lines_style .= solidLine 2 (opaque red)
        plot_lines_values .=
            [zip ts $ coordinate (2*n) . potentialMapping <$> (harmoniumTranspose trns >$> zs)]


-- Main --

main :: IO ()
main = do

    bl <- D.doesFileExist flnm
    c0s <- if bl
              then read <$> readFile flnm
              else error "Script requires a 'ppc-dynamics' file"

    let nnp = fromList nn c0s
        qdq0 = fromList (Bundle pndl) [1,0]

    xnzs <- runWithSystemRandom $ generatePath nstpssml nnp qdq0

    let (xs,ns,zs) = unzip3 $ splitTriple <$> xnzs
        (lns,lzs) = unzip $ beliefNegativeLogLikelihoods trns <$> xnzs

        gnlyt = execEC $ do

            layout_x_axis . laxis_override .= axisGridHide
            layout_y_axis . laxis_generate .= scaledAxis xaxprms (0,100)
            layout_y_axis . laxis_override .= axisGridHide
            layout_x_axis . laxis_generate .= autoScaledAxis xaxprms

            plot . liftEC $ do
                plot_points_values .= zip ts (sum . listCoordinates <$> ns)
                plot_points_style .= filledCircles 2 (opaque black)

            plot . liftEC $ do
                plot_lines_values .= [zip ts $ sum . listCoordinates <$> zs]
                plot_lines_style .= solidLine 2 (opaque red)

{-
            plot . liftEC $ do
                plot_lines_values .= [zip ts $ sum . listCoordinates <$> zipWith (<->) zs ns]
                plot_lines_style .= solidLine 2 (opaque blue)
                -}

        qlyt = coordinateLayout 0 (-2,2) (xs,ns,zs)
        dqlyt = coordinateLayout 1 (-4,4) (xs,ns,zs)

        rnbl = toRenderable $ StackedLayouts [StackedLayout qlyt, StackedLayout dqlyt, StackedLayout gnlyt] True

    putStrLn "Average Observation Likelihood:"
    print $ mean lns
    putStrLn "Average Belief Likelihood:"
    print $ mean lzs
    putStrLn "Average Observation Spike Count:"
    print . mean $ sum . listCoordinates <$> ns
    putStrLn "Average Belief Rate:"
    print . mean $ sum . listCoordinates <$> zs

    --void $ renderableToAspectWindow False 1200 1200 rnbl
    void $ renderableToFile (FileOptions (600,400) PDF) "simulation.pdf" rnbl