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

--- Imports ---


-- Goal --

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


--- Globals ---

nsmps = 20

-- Bernoulli --

(mnB,mxB) = (0,1)
bnsB = 2
truB = chart Standard $ fromList Bernoulli [0.7]
toDoubleB = coordinate 0 . sufficientStatistic Bernoulli
rngB = [False,True]

-- Categorical --

(mnC,mxC) = (0,4)
bnsC = 5
toDoubleC = fromIntegral
truC = chart Standard $ fromList (Categorical [0,1,2,3,4]) [0.1,0.4,0.1,0.2]
rngC = [0..4]

-- Poisson --

(mnP,mxP) = (0,20)
bnsP = 20
toDoubleP = fromIntegral
truP = chart Standard $ fromList Poisson [5]
rngP = [0..20]

-- Normal --

(mnN,mxN) = (-3,7)
bnsN = 20
toDoubleN = id
truN = chart Standard $ fromList Normal [2,0.7]
rngN = [-3,-2.99..7]

-- Layout --

generateLayout :: ( Show m, Transition Standard Mixture m, Transition Standard Natural m
    , MaximumLikelihood Standard m, AbsolutelyContinuous Standard m, Generative Standard m , ExponentialFamily m )
    => Int -> Double -> Double -> (Sample m -> Double) -> [Sample m] -> Standard :#: m -> IO (LayoutLR Double Int Double)
generateLayout nb mn mx toDouble rng p = do

    let m = manifold p
        lineFun1 p' = zip (toDouble <$> rng) $ density p' <$> rng
        lineFun2 p' = zip (toDouble <$> rng) $ density p' <$> rng

    smps <- runWithSystemRandom . replicateM nsmps $ generate p

    let mle1 = chart Standard $ mle m smps
    let hstplt = histogramPlot nb mn mx [toDouble <$> smps] . execEC $ do
            plot_bars_titles .= ["Samples"]
            plot_bars_item_styles .= [(solidFillStyle $ opaque blue, Nothing)]

    return . histogramLayoutLR hstplt . execEC $ do

        layoutlr_title .= (show (manifold p) ++ "; KLD: " ++ take 5 (showFFloat (Just 3) (klDivergence mle1 p) ""))
        layoutlr_left_axis . laxis_title .= "Sample Count"
        layoutlr_right_axis . laxis_title .= "Probability Mass"
        layoutlr_x_axis . laxis_title .= "Value"

        plotRight . liftEC $ do
            plot_lines_style .= dashedLine 3 [2,1] (opaque black)
            plot_lines_title .= "True"
            plot_lines_values .= [lineFun1 p]

        plotRight . liftEC $ do
            plot_lines_style .= dashedLine 3 [10,5] (opaque red)
            plot_lines_title .= "Standard MLE"
            plot_lines_values .= [ lineFun1 mle1 ]

        plotRight . liftEC $ do
            plot_lines_style .= dashedLine 3 [7,3] (opaque purple)
            plot_lines_title .= "Exponential Family MLE"
            plot_lines_values .= [ lineFun2 . chart Natural $ mle m smps ]

main = do

    lytB <- tval <$> generateLayout bnsB mnB mxB toDoubleB rngB truB
    lytC <- tval <$> generateLayout bnsC mnC mxC toDoubleC rngC truC
    lytP <- tval <$> generateLayout bnsP mnP mxP toDoubleP rngP truP
    lytN <- tval <$> generateLayout bnsN mnN mxN toDoubleN rngN truN

    let grd1 = lytB .|. lytC
        grd2 = lytP .|. lytN

    renderableToAspectWindow False 800 600 . toRenderable . weights (1,1) $ grd1 ./. grd2