{-# LANGUAGE TypeFamilies #-}

--- Imports ---


-- Goal --

import Goal.Core

import Goal.Geometry
import Goal.Probability


--- Program ---


-- Globals --

vr = 1
mn = -4
mx = 4
nkrns = 10
mus = range mn mx nkrns
sps = [ fromList Normal [mu,vr] | mu <- mus]

gn1 = 2
gn2 = 4

trns1 = modulateTransducerGain gn1 $ buildNormalTransducer sps
trns2 = modulateTransducerGain gn2 $ buildNormalTransducer sps

x0 = 0

stps = 2000
pltrng = range mn mx stps
laxprms = LinearAxisParams (show . round) 2 2
iaxprms = LinearAxisParams show 3 3
xaxprms = LinearAxisParams (show . round) 5 5


-- Functions --


-- Main --

main = do

    rs1 <- runWithSystemRandom . generate $ conditionalLatentDistribution trns1 x0
    rs2 <- runWithSystemRandom . generate $ conditionalLatentDistribution trns2 x0

    let tclyt = execEC $ do

            layout_y_axis . laxis_generate .= scaledAxis laxprms (0,1.5)
            layout_x_axis . laxis_generate .= autoScaledAxis xaxprms
            --layout_y_axis . laxis_title .= "Activation"
            layout_y_axis . laxis_override .= axisGridHide

            --layout_x_axis . laxis_title .= "Stimulus"
            layout_x_axis . laxis_override .= axisGridHide

            plot . liftEC $ do
                --plot_lines_title .= "Tuning Curves"
                plot_lines_style .= solidLine 2 (opaque black)
                plot_lines_values .= ( zip pltrng <$> transpose
                    (listCoordinates . (gn1 />) . potentialMapping <$> conditionalLatentDistributions trns1 pltrng) )

    let rsplytfun trns rs = execEC $ do

            let posterior = conditionalObservableDistribution trns rs
                scl = 10

            --layoutlr_title .= ("μ=" ++ showFFloat (Just 3) mu "" ++ "; σ=" ++ showFFloat (Just 3) sd "")

            layoutlr_left_axis . laxis_generate .= scaledAxis laxprms (0,2)
            --layoutlr_left_axis . laxis_title .= "Probability Density"
            layoutlr_left_axis . laxis_override .= axisGridHide

            layoutlr_right_axis . laxis_generate .= scaledIntAxis iaxprms (0,round scl)
            --layoutlr_right_axis . laxis_title .= "Response Count"
            layoutlr_right_axis . laxis_override .= axisGridHide

            --layoutlr_x_axis . laxis_title .= "Stimulus"
            layoutlr_margin .= 10

            layoutlr_x_axis . laxis_override .= axisGridHide
            layoutlr_x_axis . laxis_generate .= autoScaledAxis xaxprms

            layoutlr_plots
                .= [ Left $ vlinePlot "" (solidLine 2 $ opaque black) x0 ]

            plotRight . liftEC $ do
                plot_points_style .= filledCircles 3 (opaque black)
                plot_points_values .= zip mus rs
                --plot_points_title .= "Response"

            plotLeft . liftEC $ do
                plot_lines_style .= solidLine 2 (opaque red)
                plot_lines_values .= [zip pltrng $ density posterior <$> pltrng]
                --plot_lines_title .= "Posterior Density"

    let rsplyt1 = rsplytfun trns1 rs1
        rsplyt2 = rsplytfun trns2 rs2
        rsplyt3 = rsplytfun trns2 (zipWith (+) rs1 rs2)

    let rnbl = toRenderable . weights (1,1)
            $ tval (StackedLayouts [StackedLayout tclyt, StackedLayoutLR rsplyt2] True)
                .|. tval (StackedLayouts [StackedLayoutLR rsplyt1, StackedLayoutLR rsplyt3] True)

    void $ renderableToAspectWindow False 1200 800 rnbl
    --void $ renderableToFile (FileOptions (600,300) PDF) "population-code.pdf" rnbl