--- Imports ---


-- Goal --

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


--- Program ---


-- Globals --

mnx = -4
mxx = 4
mny = -4
mxy = 4
vr = 2
sps = [ joinReplicated [fromList Normal [x,vr], fromList Normal [y,vr]]
       | x <- tail $ range mnx mxx 10, y <- range mny mxy 10 ]
gn = 10
trns = modulateTransducerGain gn $ buildReplicatedNormalTransducer sps

x0 = -1
y0 = 1
xy0 = [x0,y0]

-- Functions --

rngx = (mnx,mxx,100)
rngy = (mny,mxy,100)
niso = 10
clrs = rgbaGradient (0,0,1,0.6) (1,0,0,0.6) niso

transducerRenderable rs = toRenderable . execEC $ do

    let [x',_,y',_] = listCoordinates $ conditionalObservableDistribution trns rs
        posterior x y = density (conditionalObservableDistribution trns rs) [x,y]
        cntrs = contours rngx rngy niso posterior

    sequence_ $ do

        ((_,cntr),clr) <- zip cntrs clrs

        return . plot . liftEC $ do

            plot_lines_style .= solidLine 3 clr
            plot_lines_values .= cntr

    layout_x_axis . laxis_generate .= scaledAxis def (mnx,mxx)
    layout_y_axis . laxis_generate .= scaledAxis def (mny,mxy)

    plot . liftEC $ do
        plot_points_style .= filledCircles 4 (opaque black)
        plot_points_values .= [(x0, y0)]
        plot_points_title .= "Stimulus"

    plot . liftEC $ do
        plot_points_style .= filledCircles 4 (opaque red)
        plot_points_values .= [(x',y')]
        plot_points_title .= "Estimate"

    plot . liftEC $
        plot_annotation_values .= [(x,y,show r) | (r,[x,_,y,_]) <- zip rs $ listCoordinates <$> sps ]

{-
    plotLeft . liftEC $ do
        plot_lines_style .= solidLine 3 (opaque red)
        plot_lines_values .= [let plts = posterior <$> pltrng in zip pltrng $ (*50) . (/ sum plts) <$> plts ]
        plot_lines_title .= "Posterior Density"
        -}

-- Main --

main = do
    rs <- runWithSystemRandom . generate $ conditionalLatentDistribution trns xy0

    print ("Spike count: " ++ show (sum rs))

    let rnbl = transducerRenderable rs

    void $ renderableToAspectWindow False 800 800 rnbl