Plot.hs

module Goal.Core.Plot
    ( -- * Module Exports
      module Graphics.Rendering.Chart
    , module Data.Colour
    , module Data.Colour.Names
    , module Data.Colour.SRGB.Linear
    , module Graphics.Rendering.Chart.Backend.Cairo
    , module Graphics.Rendering.Chart.Grid
    , module Graphics.Rendering.Chart.Gtk
    , module Graphics.Rendering.Chart.State
    , module Goal.Core.Plot.Contour
    -- * Plots
    -- ** PixMap
    , pixMapPlot
    -- ** Histograms
    , histogramPlot
    , histogramPlot0
    , logHistogramPlot
    , logHistogramPlot0
    -- * Layouts
    -- ** PixMap
    , pixMapLayout
    -- ** Histogram
    , histogramLayout
    , logHistogramLayout
    , histogramLayoutLR
    -- * Util
    , rgbaGradient
    -- * Rendering
    , renderableToAspectWindow
    ) where


--- Imports ---

import Data.List hiding (sum)

import Control.Monad
import Control.Lens.Getter
import Control.Lens.Setter hiding (Identity)

import Data.Default.Class
import Numeric


-- Re-exports --

import Graphics.Rendering.Chart hiding (x0,y0,Point)
import Data.Colour
import Data.Colour.Names
import Data.Colour.SRGB.Linear
import Graphics.Rendering.Chart.Backend.Cairo
import Graphics.Rendering.Chart.State
import Graphics.Rendering.Chart.Grid
import Graphics.Rendering.Chart.Gtk

-- Scientific --

import Goal.Core.Plot.Contour

-- Qualified --

import qualified Graphics.UI.Gtk as G


-- Unqualified --

import Graphics.Rendering.Cairo (liftIO)


--- Util ---

rgbaGradient :: (Double, Double, Double, Double) -> (Double, Double, Double, Double) -> Int
    -> [AlphaColour Double]
-- | Returns an ordered list of colours useful for plotting.
rgbaGradient (rmn,gmn,bmn,amn) (rmx,gmx,bmx,amx) n =
    zipWith (flip withOpacity) [amn,amn + astp .. amx]
    $ zipWith3 rgb [rmn,rmn + rstp .. rmx] [gmn,gmn + gstp .. gmx] [bmn,bmn + bstp .. bmx]
    where rstp = (rmx - rmn) / fromIntegral n
          gstp = (gmx - gmn) / fromIntegral n
          bstp = (bmx - bmn) / fromIntegral n
          astp = (amx - amn) / fromIntegral n


--- Plots ---


-- PixMap --

pixMapPlot :: (Double,Double) -> [[AlphaColour Double]] -> Plot Double Double
-- | Returns a pixmap representation of a matrix style set of doubles. Based on the
-- defaults, the list of colours are assumed to be in (y,x) coordinates, where the
-- origin is at the lower left of the image. If matrix style coordinates are desired,
-- The containing layout should be given a reversed y axis, so that the origin is at the
-- top left of the image.

--  The pair of doubles indicates where corner of the given image should be located.
pixMapPlot (x0,y0) pss = foldr1 joinPlot
    $ concat [ [ boxPlot r c . solidFillStyle $ p | (c,p) <- zip [x0..] ps ] | (r,ps) <- zip [y0..] pss ]
    where boxPlot y x stl = toPlot
              $ plot_fillbetween_style .~ stl
              $ plot_fillbetween_values .~ [(x-0.5,(y-0.5,y+0.5)),(x+0.5,(y-0.5,y+0.5))]
              $ def

pixMapLayout :: Int -> Int -> Layout Double Double -> Layout Double Double
-- | A nice base layout for a pixMap, with a box around the pixmap with one 'pixel'
-- padding, and a reversed y axis for matrix style coordinates.
pixMapLayout rws0 cls0 lyt =
    layout_top_axis_visibility .~ AxisVisibility True True False
    $ layout_left_axis_visibility .~ AxisVisibility True True False
    $ layout_right_axis_visibility .~ AxisVisibility True True False
    $ layout_bottom_axis_visibility .~ AxisVisibility True True False
    $ layout_y_axis . laxis_reverse .~ True
    $ layout_y_axis . laxis_generate .~
        const (makeAxis (const "") ([-1.5,rws+0.5],[-1.5,rws+0.5],[-1.5,rws+0.5]))
    $ layout_x_axis . laxis_generate .~
        const (makeAxis (const "") ([-1.5,cls+0.5],[-1.5,cls+0.5],[-1.5,cls+0.5]))
    $ lyt
    where cls = fromIntegral cls0
          rws = fromIntegral rws0

-- Histogram --

histogramPlot0 :: (Num a,BarsPlotValue a) =>Int -> [[Double]] -> PlotBars Double a -> PlotBars Double a
-- | Generates a histogram plot where the min and max bin value is taken from the data set.
histogramPlot0 n xss plt =
    let mx = maximum $ maximum <$> xss
        mn = minimum $ minimum <$> xss
    in histogramPlot n mn mx xss plt

logHistogramPlot0 :: Int -> [[Double]] -> PlotBars Double Double -> PlotBars Double Double
-- | Generates a histogram plot where the min and max bin value is taken from the data set.
logHistogramPlot0 n xss =
    let mx = maximum $ maximum <$> xss
        mn = minimum $ minimum <$> xss
    in logHistogramPlot n mn mx xss

histogramPlot
    :: (Num a,BarsPlotValue a)
    => Int -- ^ Number of bins
    -> Double -- ^ Min range
    -> Double -- ^ Max range
    -> [[Double]] -- ^ Data set
    -> PlotBars Double a -- ^ Plot
    -> PlotBars Double a -- ^ New Plot
-- | Creates a histogram out of a data set. The data set is a list of list of values, where
-- each sublist is a collection of data along an axis. Under and overflow is put into the
-- first and last bin, respectively. The bars are centered at the mid point between each
-- pair of bins.
histogramPlot n mn mx xss plt =
    let bns = range mn mx (n+1)
        vls = transpose $ toHistogram (tail $ take n bns) . sort <$> xss
        stp = (head (tail bns) - head bns) / 2
        bns' = (+ stp) <$> take n bns
     in plot_bars_alignment .~ BarsCentered $ plot_bars_values .~ zip bns' vls $ plt
    where toHistogram _ [] = repeat 0
          toHistogram [] xs = [genericLength xs]
          toHistogram (bn:bns') xs =
              let (hds,xs') = span (< bn) xs
              in genericLength hds : toHistogram bns' xs'

range _ _ 0 = []
range mn mx 1 = [(mn + mx) / 2]
range mn mx n =
    [ x * mx + (1 - x) * mn | x <- (/ (fromIntegral n - 1)) . fromIntegral <$> [0 .. n-1] ]

logHistogramPlot
    :: Int -- ^ Number of bins
    -> Double -- ^ Min range
    -> Double -- ^ Max range
    -> [[Double]] -- ^ Data set
    -> PlotBars Double Double -- ^ Plot
    -> PlotBars Double Double -- ^ New Plot
logHistogramPlot n mn mx xss plt =
    let bplt = histogramPlot n mn mx xss plt
        vls = modVals <$> bplt ^. plot_bars_values
        --cl = fromIntegral . ceiling . maximum . concat $ snd <$> vls
    in plot_bars_values .~ vls $ bplt
    where lbs = 10
          modVals (x,ys) = (x, logBase lbs . (+1) <$> ys)

histogramLayout :: BarsPlotValue a => PlotBars Double a -> Layout Double a -> Layout Double a
-- | The base layout for a histogram.
histogramLayout pbrs lyt =
    let bns = fst <$> pbrs ^. plot_bars_values
        stp = (head (tail bns) - head bns) / 2
        bns' = (head bns - stp) : map (+stp) bns
        rng = abs $ maximum bns'
        labelFun x
            | rng >= 1000 = showEFloat (Just 2) x ""
            | rng <= 0.01 = showEFloat (Just 2) x ""
            | otherwise = reverse . dropWhile (== '.') . dropWhile (== '0') . reverse $ showFFloat (Just 2) x ""
    in layout_x_axis . laxis_generate .~ const (makeAxis labelFun (bns',[],bns'))
       $ layout_plots %~ (plotBars pbrs:)
       $ lyt

logHistogramLayout :: PlotBars Double Double -> Layout Double Double -> Layout Double Double
-- | Base layout for a log-histogram.
logHistogramLayout pbrs lyt =
    let vls = pbrs ^. plot_bars_values
        bns = fst <$> vls
        stp = (head (tail bns) - head bns) / 2
        bns' = (head bns - stp) : map (+stp) bns
        cl = fromIntegral . ceiling . maximum . concat $ snd <$> vls
        rng = abs $ maximum bns'
        xLabelFun x
            | rng >= 1000 = showEFloat (Just 2) x ""
            | rng <= 0.01 = showEFloat (Just 2) x ""
            | otherwise = reverse . dropWhile (== '.') . dropWhile (== '0') . reverse $ showFFloat (Just 2) x ""
    in layout_plots %~ (plotBars pbrs:)
        $ layout_y_axis . laxis_generate .~ const (makeAxis yLabelFun ([0..cl],[],[0..cl]))
        $ layout_x_axis . laxis_generate .~ const (makeAxis xLabelFun (bns',[],bns'))
        $ lyt
    where lbs = 10
          yLabelFun 0 = show 0
          yLabelFun x = show (round lbs) ++ "e" ++ show (round x) ++ "-1"

histogramLayoutLR :: (BarsPlotValue a,PlotValue b) => PlotBars Double a -> LayoutLR Double a b -> LayoutLR Double a b
-- | The base layout for a histogramLR.
histogramLayoutLR pbrs lyt =
    let bns = fst <$> pbrs ^. plot_bars_values
        stp = (head (tail bns) - head bns) / 2
        bns' = (head bns - stp) : map (+stp) bns
        rng = abs $ maximum bns'
        labelFun x
            | rng >= 1000 = showEFloat (Just 2) x ""
            | rng <= 0.01 = showEFloat (Just 2) x ""
            | otherwise = reverse . dropWhile (== '.') . dropWhile (== '0') . reverse $ showFFloat (Just 2) x ""
    in layoutlr_x_axis . laxis_generate .~ const (makeAxis labelFun (bns',[],bns'))
       $ layoutlr_plots %~ (Left (plotBars pbrs):)
       $ lyt

--- IO ---


renderableToAspectWindow
    :: Bool -- ^ Display Full Screen
    -> Int -- ^ Image width
    -> Int -- ^ Image height
    -> Renderable a -- ^ The Renderable
    -> IO () -- ^ Renders the renderable to the screen

-- | Displays a renderable in a GTK aspect window.
renderableToAspectWindow fs wdth hght rnbl = do

    G.initGUI

    win <- G.windowNew
    afrm <- G.aspectFrameNew 0.5 0.5 . Just $ realToFrac (fromIntegral wdth / fromIntegral hght)
    da <- G.drawingAreaNew

    G.set afrm [ G.containerChild G.:= da ]
    G.set win [ G.containerChild G.:= afrm ]
    when fs $ G.windowFullscreen win
    G.onDestroy win G.mainQuit

    (da `G.on` G.exposeEvent) . liftIO $ do

        updateCanvas rnbl da
        return True

    G.widgetShowAll win
    G.mainGUI