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

--- Imports ---


-- Goal --

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


--- Globals ---

f x = exp . sin $ 2 * x
nsmps = 20
mnx = -3
mxx = 3
xs = range mnx mxx nsmps

-- Neural Network --

m = Poisson
n = Replicated Bernoulli 20
o = MeanNormal 1

nn = NeuralNetwork m n o

-- Training --

eps = 0.05
nepchs = 10000

-- Plot --

nplts = 100
pltrng = range mnx mxx nplts

-- Layout --

main = do

    smps <- runWithSystemRandom $ mapM (noisyFunction (chart Standard $ fromList Normal [0,0.1]) f) xs
    let xps = sufficientStatistic o <$> xs
        tps = [ fromList Poisson [smp] | smp <- smps ]

    cs0 <- runWithSystemRandom . replicateM (dimension nn) . generate . chart Standard $ fromList Normal [0,0.1]
    let nnp0 = fromList nn cs0

    let gradient nnp = meanSquaredBackpropagation nnp xps tps
        nnps = vanillaGradientDescent eps gradient nnp0
        nnp1 = nnps !! nepchs

        fhat x = coordinate 0 $ nnp1 >.> sufficientStatistic o x

    let lyt1 = execEC $ do

            layout_title .= "Regression"

            plot . liftEC $ do

                plot_lines_title .= "True"
                plot_lines_style .= solidLine 3 (opaque black)
                plot_lines_values .= [zip pltrng (f <$> pltrng)]

            plot . liftEC $ do

                plot_points_title .= "Samples"
                plot_points_style .=  filledCircles 4 (opaque black)
                plot_points_values .= zip xs smps

            plot . liftEC $ do

                plot_lines_title .= "MLP"
                plot_lines_style .= solidLine 3 (opaque red)
                plot_lines_values .= [zip pltrng (fhat <$> pltrng)]

    let (mp,mtx1,np,mtx2) = splitNeuralNetwork nnp1
    let lyt2 = coordinateLogHistogram 10 "Network Weights" ["B1","I1","B2","I2"]
            [coordinates mp, coordinates mtx1, coordinates np, coordinates mtx2]

    renderableToAspectWindow False 800 800 . toRenderable . weights (1,1) $ tval lyt2 ./. tval lyt1

{-
    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 ]

    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
    -}