{-# LANGUAGE RecordWildCards     #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE ParallelListComp    #-}

module Data.Analysis.Fulcrum.Plot (doPlot, doLogPlot) where

import Control.Lens((.~))
import Data.Colour (opaque, Colour)
import Data.Colour.Names
import Data.Default.Class (def)
import Data.Map.Strict ((!), member)
import Data.Maybe (fromJust)
import Data.Monoid ((<>))
import Graphics.Rendering.Chart
import Math.Statistics (stddev, average)

import Data.Analysis.Fulcrum.Analysis

colors :: [Colour Double]
colors = cycle [ blue, red, green, orange, purple, yellow ]

plotLayout :: PlotValue b => Plan r f -> [r] -> Layout Int b
plotLayout plan@(Plan { .. }) rows
  = layout_background .~ solidFillStyle (opaque white)
  $ layout_x_axis.laxis_generate .~
      autoIndexAxis (map absShowS axisNames)
  $ layout_y_axis.laxis_title .~ planFCName
  $ layout_x_axis.laxis_title .~ slName planAxis
  $ layout_title .~ planName
  $ setLayoutForeground (opaque black)
  $ def
  where axisNames = getAxisValues plan rows

plotLines :: (RealFloat f, PlotValue g) =>
  Bool -> (Result f) -> String -> [(Int, AbstractTuple)]
    -> AbstractTuple -> Colour Double -> (f -> g) -> [PlotLines Int g]
plotLines True vals axisTitle axisNames ln c conv =
  [ plot_lines_style .~ lineStyle 4 c
  $ plot_lines_values .~
   [ [ (idx, conv (vals ! (ax, ln)))
     | (idx, ax) <- axisNames
     , (ax, ln) `member` vals
     , not (isNaN (vals ! (ax, ln))) ] ]
  $ plot_lines_title .~ axisTitle
  $ def
  ]
plotLines False _ _ _ _ _ _ = []

plotDots :: (RealFloat f, PlotValue g) => Result f -> [(Int, AbstractTuple)] -> AbstractTuple
         -> Colour Double -> (f -> g) -> PlotPoints Int g
plotDots vals axisNames ln color conv
  = plot_points_style .~ filledCircles 5 (opaque color)
  $ plot_points_values .~
    [ (idx, conv (vals ! (ax, ln)))
    | (idx, ax) <- axisNames
    , (ax, ln) `member` vals
    , not (isNaN (vals ! (ax, ln))) ]
  $ def

plotStds :: (RealFloat f) => Result f -> Result f ->
  [(Int, AbstractTuple)] -> AbstractTuple -> Colour Double -> PlotErrBars Int f
plotStds avgs stds axisNames ln color
  = plot_errbars_line_style .~ lineStyle 1 color
  $ plot_errbars_values .~
    [ symErrPoint idx (avgs ! (ax, ln)) 0 (stds ! (ax, ln))
    | (idx, ax) <- axisNames
    , (ax, ln) `member` avgs
    , not (isNaN (stds ! (ax, ln)))
    , not (isNaN (avgs ! (ax, ln))) ]
  $ def

-- | Given a 'Plan' and a list of rows in which the focus variable is
--   numeric and plottable, produce a log-scale plot of the relevant
--   values with error bars. The 'doPlot' function will rerun your plan
--   taking the 'average' of your focus variables, so custom 'planMerge'
--   values will be replaced.
doLogPlot :: Plan r Double -> [r] -> Renderable ()
doLogPlot plan@(Plan { .. }) rows =
  toRenderable $ layout_plots .~ map toPlot (concat plots) <>
                                 map toPlot dots
               $ plotLayout plan rows
  where plots = [ plotLines planLinAxis avgs "" axisNums ln c LogValue
                | ln <- lineNames
                | c <- colors
                ]
        dots  = [ plotDots avgs axisNums ln c LogValue
                | ln <- lineNames
                | c <- colors
                ]
        lineNames = getLineValues plan rows
        axisNames = getAxisValues plan rows
        runPlan m = runPlanToMap (plan { planMerge = m })
        avgs      = runPlan average rows
        axisNums  = zip [(0::Int)..] axisNames

-- | Given a 'Plan' and a list of rows in which the focus variable is
--   numeric and plottable, produce a (linear-scale) plot of the relevant
--   values with error bars. The 'doPlot' function will rerun your plan
--   taking the 'average' and the 'stddev' of your focus variables, so
--   custom 'planMerge' values will be replaced.
doPlot :: (RealFloat f, PlotValue f) => Plan r f -> [r] -> Renderable ()
doPlot plan@(Plan { .. }) rows =
  toRenderable $ layout_plots .~ map toPlot bars <>
                                 map toPlot (concat plots) <>
                                 map toPlot dots
               $ plotLayout plan rows
  where plots = [ plotLines planLinAxis avgs "" axisNums ln c id
                | ln <- lineNames
                | c  <- colors
                ]
        dots  = [ plotDots avgs axisNums ln c id
                | ln <- lineNames
                | c <- colors
                ]
        bars = [ plotStds avgs stds axisNums ln c
               | ln <- lineNames
               | c <- colors
               ]
        lineNames = getLineValues plan rows
        axisNames = getAxisValues plan rows
        runPlan m = runPlanToMap (plan { planMerge = m })
        avgs      = runPlan average rows
        stds      = runPlan stddev rows
        axisNums  = zip [(0::Int)..] axisNames

lineStyle :: Double -> Colour Double -> LineStyle
lineStyle n c
  = line_width .~ n
  $ line_color .~ opaque c
  $ def