{-# LANGUAGE TypeOperators #-}

--- Imports ---


-- Goal --

import Pendulum

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

-- Qualified --

import qualified System.Directory as D


--- Program ---


-- Globals --

nbns = 20
mnbn = -2
mxbn = 6
npths = 1000
iaxprms = LinearAxisParams show 5 5


-- Functions --

generatePaths
    :: Function Mixture Mixture :#: NeuralNetwork (Replicated Poisson) (Replicated Bernoulli) (Replicated Poisson)
    -> Rand (ST s) [[(Partials, Mixture, Mixture) :#: (PhaseSpace Pendulum, Replicated Poisson, Replicated Poisson)]]
generatePaths nnp = do
    let randomPath = randomState >>= generatePath nstpssml nnp
    replicateM npths randomPath

isLatentPathBounded
    :: [(Partials, Mixture, Mixture) :#: (PhaseSpace Pendulum, Replicated Poisson, Replicated Poisson)]
    -> Bool
isLatentPathBounded xnzs =
    let (xs,_,_) = unzip3 $ splitTriple <$> xnzs
     in and [ q > -pi && q < pi | [q,_] <- listCoordinates <$> xs ]

isLogLikelihoodPathFinite
    :: [(Partials, Mixture, Mixture) :#: (PhaseSpace Pendulum, Replicated Poisson, Replicated Poisson)]
    -> Bool
isLogLikelihoodPathFinite xnzs =
    let lzs = snd . unzip $ beliefNegativeLogLikelihoods trns <$> xnzs
     in and $ not . isInfinite <$> lzs

isLogLikelihoodPlotted
    :: (Partials, Mixture, Mixture) :#: (PhaseSpace Pendulum, Replicated Poisson, Replicated Poisson)
    -> Bool
isLogLikelihoodPlotted xnz =
    let lz = snd $ beliefNegativeLogLikelihoods trns xnz
     in lz > mnbn && lz < mxbn


-- Main --

main :: IO ()
main = do

    bl <- D.doesFileExist flnm
    c0s <- if bl
              then read <$> readFile flnm
              else error "Script requires a 'ppc-dynamics' file"

    let nnp = fromList nn c0s

    xnzss <- runWithSystemRandom $ generatePaths nnp

    let bxnzss = filter isLatentPathBounded xnzss
        fbxnzss = filter isLogLikelihoodPathFinite bxnzss
        (pxnzs,pxnzs') = partition isLogLikelihoodPlotted $ concat fbxnzss
        (lns,lzs) = unzip $ beliefNegativeLogLikelihoods trns <$> pxnzs
        (lns',lzs') = unzip $ beliefNegativeLogLikelihoods trns <$> concat fbxnzss
        (_,ns,zs) = unzip3 $ splitTriple <$> pxnzs

    putStrLn "Percent in Bounds:"
    print $ 100 * genericLength bxnzss / genericLength xnzss
    putStrLn "Percent of Finite Rate:"
    print $ 100 * genericLength fbxnzss / genericLength bxnzss
    putStrLn "Percent within Histogram Bounds:"
    print $ 100 * genericLength pxnzs / genericLength (concat fbxnzss)
    putStrLn "Max out of Histogram Bounds:"
    print $ maximum . snd . unzip $ beliefNegativeLogLikelihoods trns <$> pxnzs'
    putStrLn "Average Observation Spike Count:"
    print . mean $ sum . listCoordinates <$> ns
    putStrLn "Average Belief Rate:"
    print . mean $ sum . listCoordinates <$> zs
    putStrLn "Average Information Gain:"
    print $ mean lns' - mean lzs'

    let plt =
            plot_bars_item_styles .~ [(FillStyleSolid $ opaque black,Nothing),(FillStyleSolid $ opaque red,Nothing)]
            $ histogramPlot nbns mnbn mxbn [lns,lzs] def

        lyt =
            layout_plots .~ [plotBars plt]
            -- $ layout_x_axis . laxis_title .~ "-Log-Likelihood"
            $ layout_y_axis . laxis_generate .~ autoScaledIntAxis iaxprms
            $ layout_x_axis . laxis_override .~ axisGridHide
            $ layout_y_axis . laxis_override .~ axisGridHide
            $ histogramLayout plt def

    --void $ renderableToAspectWindow False 1200 800 $ toRenderable (lyt :: Layout Double Int)
    void $ renderableToFile (FileOptions (600,200) PDF) "histogram.pdf" $ toRenderable (lyt :: Layout Double Int)