{-# LANGUAGE FlexibleContexts,TypeOperators #-}

--- Imports ---


-- Scientific --

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

--- Program ---


-- Globals --

res = 200
niso = 10


-- Functions --

divergenceLayout :: (ExponentialFamily m, Transition c Mixture m, Transition c Natural m)
    => (Double, Double) -> AlphaColour Double -> c -> m -> Layout Double Double
divergenceLayout (mn,mx) clr c m = execEC $ do

    let f x y = relativeEntropy (chart c $ fromList m [x]) (chart c $ fromList m [y])
        cntrs = contours (mn,mx,res) (mn,mx,res) niso f
        x0 = (mx + mn) / 2
        y0 = x0
        str0 = "0.0"
        hgh = 0.95 * mx + 0.05 * mn
        lw = 0.05 * mx + 0.95 * mn
        x1 = hgh
        y1 = lw
        str1 = showFFloat (Just 1) (f x1 y1) ""
        x2 = lw
        y2 = hgh
        str2 = showFFloat (Just 1) (f x2 y2) ""

    plot . liftEC $ do
        plot_lines_style .= solidLine 2 clr
        plot_lines_values .= [[ (x,x) | x <- range mn mx 3 ]]

    sequence_ $ do

        (_,cntr) <- cntrs

        return . plot . liftEC $ do

            plot_lines_style .= solidLine 3 clr
            plot_lines_values .= cntr

    plot . liftEC $ do
        plot_points_values .= [(x0,y0),(x1,y1),(x2,y2)]
        plot_points_style .= filledCircles 9 (opaque white)

    plot . liftEC $ do
        plot_annotation_values .= [(x0,y0,str0),(x1,y1,str1),(x2,y2,str2)]
        plot_annotation_style . font_weight .= FontWeightBold


-- Main --

main :: IO ()
main = do

    let [blyt0,blyt1,plyt0,plyt1] =
            [ toRenderable $ divergenceLayout (0.02,0.98) (opaque blue) Mixture Bernoulli
            , toRenderable $ divergenceLayout (-5,5) (opaque red) Natural Bernoulli
            , toRenderable $ divergenceLayout (0.1,4) (opaque blue) Mixture Poisson
            , toRenderable $ divergenceLayout (-2,2) (opaque red) Natural Poisson ]

    let bgrd = tval blyt0 ./. tval blyt1
        pgrd = tval plyt0 ./. tval plyt1

    let rnbl = gridToRenderable . weights (1,1) $ bgrd .|. pgrd
    --void $ renderableToFile (FileOptions (500,500) PDF) "divergence.pdf" grd
    void $ renderableToAspectWindow False 1000 1000 rnbl