--- Imports ---


-- Goal --

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

import qualified Data.Vector.Storable as C


--- Globals ---


nsmps = 10
tru = chart Standard $ fromList (MultivariateNormal 2) [0,0.5,1,0.5,0,1]

rng = (-4,4,400)
niso = 10

axprms = LinearAxisParams (show . round) 5 5

vectorToPair xs = (xs C.! 0, xs C.! 1)
pairToVector (x,y) = C.fromList [x,y]

--- Main ---


main :: IO ()
main = do

    smps <- runWithSystemRandom . replicateM nsmps $ generate tru

    let mlenrm = chart Standard $ mle (MultivariateNormal 2) smps
        --efnrm = chart Natural $ mle (MultivariateNormal 2) smps

        truf x y = density tru $ pairToVector (x,y)
        mlef x y = density mlenrm $ pairToVector (x,y)
        --eff x y = density efnrm $ pairToVector (x,y)

        trucntrs = contours rng rng niso truf
        mlecntrs = contours rng rng niso mlef
        --efcntrs = contours rng rng niso eff

        truclrs = rgbaGradient (1,0,0,0.5) (1,0,0,1) niso
        mleclrs = rgbaGradient  (0,0,1,0.5) (0,0,1,1) niso
        --efclrs = rgbaGradient (0,1,0,0.5) (0,1,0,1) niso
        bls = True : repeat False

        rnbl = toRenderable . execEC $ do

            --layout_title .= ("Multivariate Normal" ++ "; KLD: " ++ showFFloat (Just 3) (klDivergence mlenrm tru) "")

            layout_x_axis . laxis_generate .= scaledAxis axprms (-4,4)
            layout_x_axis . laxis_override .= axisGridHide
            layout_x_axis . laxis_title .= "x"
            layout_y_axis . laxis_generate .= scaledAxis axprms (-4,4)
            layout_y_axis . laxis_override .= axisGridHide
            layout_y_axis . laxis_title .= "y"

            sequence_ $ do

                ((_,cntr),clr,bl) <- zip3 trucntrs truclrs bls

                return . plot . liftEC $ do

                    --when bl $ plot_lines_title .= "True"
                    plot_lines_style .= solidLine 3 clr
                    plot_lines_values .= cntr

            sequence_ $ do

                ((_,cntr),clr,bl) <- zip3 mlecntrs mleclrs bls

                return . plot . liftEC $ do

                    --when bl $ plot_lines_title .= "Standard MLE"
                    plot_lines_style .= solidLine 3 clr
                    plot_lines_values .= cntr

            plot . liftEC $ do
                --plot_points_title .= "Samples"
                plot_points_values .= map vectorToPair smps
                plot_points_style .= filledCircles 4 (opaque black)

{-
            sequence $ do

                ((_,cntr),clr,bl) <- zip3 efcntrs efclrs bls

                return . plot . liftEC $ do

                    when bl $ plot_lines_title .= "Exponential Family MLE"
                    plot_lines_style .= solidLine 3 clr
                    plot_lines_values .= cntr
                    -}

    --renderableToAspectWindow False 800 600 rnbl
    void $ renderableToFile (FileOptions (250,250) PDF) "multivariate.pdf" rnbl