    {-# LANGUAGE TemplateHaskell, FlexibleInstances #-}

    module Graphics.Rendering.Chart.Plot.Histogram
      ( -- * Histograms
        PlotHist (..)
      , histToPlot
      , defaultPlotHist
      , defaultFloatPlotHist
      , defaultNormedPlotHist
        -- * Accessors
      , plot_hist_title
      , plot_hist_bins
      , plot_hist_values
      , plot_hist_no_zeros
      , plot_hist_range
      , plot_hist_drop_lines
      , plot_hist_line_style
      , plot_hist_fill_style
      , plot_hist_norm_func
      ) where

    import Control.Monad (when)
    import Data.Monoid
    import Data.Maybe (fromMaybe)
    import qualified Data.Foldable as F
    import qualified Data.Vector as V

    import Control.Lens
    import Graphics.Rendering.Chart.Plot.Types
    import Graphics.Rendering.Chart.Geometry
    import Graphics.Rendering.Chart.Drawing
    import Data.Default.Class

    import Data.Colour (opaque)
    import Data.Colour.Names (blue)
    import Data.Colour.SRGB (sRGB)

    import Numeric.Histogram

    data PlotHist x y = PlotHist
        { -- | Plot title
          _plot_hist_title                :: String

          -- | Number of bins
        , _plot_hist_bins                 :: Int

          -- | Values to histogram
        , _plot_hist_values               :: [x]

          -- | Don't attempt to plot bins with zero counts. Useful when
          -- the y-axis is logarithmically scaled.
        , _plot_hist_no_zeros             :: Bool

          -- | Override the range of the histogram. If @Nothing@ the
          -- range of @_plot_hist_values@ is used.
          --
          -- Note that any normalization is always computed over the full
          -- data set, including samples not falling in the histogram range.
        , _plot_hist_range                :: Maybe (x,x)

          -- | Plot vertical lines between bins
        , _plot_hist_drop_lines           :: Bool

          -- | Fill style of the bins
        , _plot_hist_fill_style           :: FillStyle

          -- | Line style of the bin outlines
        , _plot_hist_line_style           :: LineStyle

          -- | Normalization function
        , _plot_hist_norm_func            :: Double -> Int -> y
        }

    instance Default (PlotHist x Int) where
        def = defaultPlotHist

    -- | The default style is an unnormalized histogram of 20 bins.
    defaultPlotHist :: PlotHist x Int
    defaultPlotHist = PlotHist { _plot_hist_bins        = 20
                               , _plot_hist_title       = ""
                               , _plot_hist_values      = []
                               , _plot_hist_no_zeros    = False
                               , _plot_hist_range       = Nothing
                               , _plot_hist_drop_lines  = False
                               , _plot_hist_line_style  = defaultLineStyle
                               , _plot_hist_fill_style  = defaultFillStyle
                               , _plot_hist_norm_func   = const id
                               }

    -- | @defaultPlotHist@ but with real counts
    defaultFloatPlotHist :: PlotHist x Double
    defaultFloatPlotHist = defaultPlotHist { _plot_hist_norm_func = const realToFrac }

    -- | @defaultPlotHist@ but normalized such that the integral of the
    -- histogram is one.
    defaultNormedPlotHist :: PlotHist x Double
    defaultNormedPlotHist = defaultPlotHist { _plot_hist_norm_func = \n y->realToFrac y / n }

    defaultFillStyle :: FillStyle
    defaultFillStyle = solidFillStyle (opaque $ sRGB 0.5 0.5 1.0)

    defaultLineStyle :: LineStyle
    defaultLineStyle = (solidLine 1 $ opaque blue)
         { _line_cap  = LineCapButt
         , _line_join = LineJoinMiter
         }

    -- | Convert a @PlotHist@ to a @Plot@
    --
    -- N.B. In principle this should be Chart's @ToPlot@ class but unfortunately
    -- this does not allow us to set bounds on the x and y axis types, hence
    -- the need for this function.
    histToPlot :: (RealFrac x, Num y, Ord y) => PlotHist x y -> Plot x y
    histToPlot p = Plot {
            _plot_render      = renderPlotHist p,
            _plot_legend      = [(_plot_hist_title p, renderPlotLegendHist p)],
            _plot_all_points  = unzip
                                $ concatMap (\((x1,x2), y)->[(x1,y), (x2,y)])
                                $ histToBins p
        }

    buildHistPath :: (RealFrac x, Num y)
                  => PointMapFn x y -> [((x,x), y)] -> Path
    buildHistPath _ [] = End
    buildHistPath pmap bins = MoveTo (pt xb 0) (go bins)
        where go [((x1,x2),y)]      = LineTo (pt x1 y)
                                    $ LineTo (pt x2 y)
                                    $ LineTo (pt x2 0)
                                    $ End
              go (((x1,x2),y):rest) = LineTo (pt x1 y)
                                    $ LineTo (pt x2 y)
                                    $ go rest
              go []                 = End
              ((xb,_),_) = head bins
              pt x y = pmap (LValue x, LValue y)

    renderPlotHist :: (RealFrac x, Num y, Ord y)
                   => PlotHist x y -> PointMapFn x y -> ChartBackend ()
    renderPlotHist p pmap
        | null bins = return ()
        | otherwise = do
            withFillStyle (_plot_hist_fill_style p) $
                alignFillPath (buildHistPath pmap bins) >>= fillPath
            withLineStyle (_plot_hist_line_style p) $ do
                when (_plot_hist_drop_lines p) $
                    alignStrokePath dropLinesPath >>= strokePath
                alignStrokePath (buildHistPath pmap bins) >>= strokePath
        where bins = histToBins p
              pt x y = pmap (LValue x, LValue y)
              dropLinesPath = F.foldMap (\((x1,_), y)->moveTo (pt x1 0)
                                                    <> lineTo (pt x1 y)
                                        ) $ tail bins

    renderPlotLegendHist :: PlotHist x y -> Rect -> ChartBackend ()
    renderPlotLegendHist p (Rect p1 p2) =
        withLineStyle (_plot_hist_line_style p) $
            let y = (p_y p1 + p_y p2) / 2
            in strokePath $ moveTo' (p_x p1) y <> lineTo' (p_x p2) y

    histToBins :: (RealFrac x, Num y, Ord y) => PlotHist x y -> [((x,x), y)]
    histToBins hist =
        filter_zeros $ zip bounds $ counts
        where n = _plot_hist_bins hist
              (a,b) = realHistRange hist
              dx = realToFrac (b-a) / realToFrac n
              bounds = binBounds a b n
              values = V.fromList (_plot_hist_values hist)
              filter_zeros | _plot_hist_no_zeros hist  = filter (\(_,c)->c > 0)
                           | otherwise                 = id
              norm = dx * realToFrac (V.length values)
              normalize = _plot_hist_norm_func hist norm
              counts = V.toList $ V.map (normalize . snd)
                       $ histWithBins (V.fromList bounds) (zip (repeat 1) $ V.toList values)

    -- TODO: Determine more aesthetically pleasing range
    realHistRange :: (RealFrac x) => PlotHist x y -> (x,x)
    realHistRange hist = fromMaybe range $ _plot_hist_range hist
        where values = V.fromList (_plot_hist_values hist)
              range = if V.null values
                        then (0,0)
                        else (V.minimum values, V.maximum values)

    $( makeLenses ''PlotHist )