Surface.hs

{-# LANGUAGE OverloadedLabels #-}
{-# LANGUAGE OverloadedStrings #-}

-- | Surface chart combinators.
--
-- A common chart is to present a set of rectangles on the XY plane with colour representing values of the underlying data. This library uses the term /surface/ chart but it is often referred to as a heatmap.
module Chart.Surface
  ( SurfaceData (..),
    SurfaceOptions (..),
    defaultSurfaceOptions,
    SurfaceStyle (..),
    defaultSurfaceStyle,
    mkSurfaceData,
    surfaces,
    surfacef,
    SurfaceLegendOptions (..),
    defaultSurfaceLegendOptions,
    surfaceLegendAxisOptions,
    gridReferenceChart,
    addSurfaceLegend,
  )
where

import Chart.Data
import Chart.Hud
import Chart.Primitive
import Chart.Style
import Data.Bifunctor
import Data.Bool
import Data.Colour
import Data.Data
import Data.Foldable
import Data.FormatN
import Data.Maybe
import GHC.Generics
import NumHask.Space
import Optics.Core
import Prelude

-- $setup
--
-- >>> :set -XOverloadedLabels
-- >>> :set -XOverloadedStrings
-- >>> import Chart
-- >>> import Optics.Core

-- | Options for a Surface chart.
data SurfaceOptions = SurfaceOptions
  { -- | surface style
    soStyle :: SurfaceStyle,
    -- | The grain or granularity of the chart
    soGrain :: Point Int,
    -- | Chart range
    soRange :: Rect Double
  }
  deriving (Eq, Show, Generic, Data)

-- | official style
defaultSurfaceOptions :: SurfaceOptions
defaultSurfaceOptions =
  SurfaceOptions defaultSurfaceStyle (Point 10 10) one

-- | A surface chart is a specialization of a 'RectChart'
--
-- >>> defaultSurfaceStyle
-- SurfaceStyle {surfaceColors = [Colour 0.02 0.73 0.80 1.00,Colour 0.02 0.29 0.48 1.00], surfaceRectStyle = Style {size = 6.0e-2, borderSize = 0.0, color = Colour 0.05 0.05 0.05 1.00, borderColor = Colour 0.00 0.00 0.00 0.00, scaleP = NoScaleP, textAnchor = AnchorMiddle, rotation = Nothing, translate = Nothing, escapeText = EscapeText, frame = Nothing, lineCap = Nothing, lineJoin = Nothing, dasharray = Nothing, dashoffset = Nothing, hsize = 0.6, vsize = 1.1, vshift = -0.25, glyphShape = SquareGlyph}}
--
-- ![surface example](other/surface.svg)
data SurfaceStyle = SurfaceStyle
  { -- | list of colours to interpolate between.
    surfaceColors :: [Colour],
    surfaceRectStyle :: Style
  }
  deriving (Eq, Show, Generic, Data)

-- | The official surface style.
defaultSurfaceStyle :: SurfaceStyle
defaultSurfaceStyle =
  SurfaceStyle (palette <$> [0 .. 1]) (blob dark)

-- | Main surface data elements
data SurfaceData = SurfaceData
  { -- | XY Coordinates of surface.
    surfaceRect :: Rect Double,
    -- | Surface colour.
    surfaceColor :: Colour
  }
  deriving (Eq, Show, Generic, Data)

-- | surface chart without any hud trimmings
surfaces :: Style -> [SurfaceData] -> [Chart]
surfaces rs ps =
  ( \(SurfaceData r c) ->
      Chart
        (rs & set #color c)
        (RectData [r])
  )
    <$> ps

-- | Create surface data from a function on a Point
mkSurfaceData ::
  (Point Double -> Double) ->
  Rect Double ->
  Grid (Rect Double) ->
  [Colour] ->
  ([SurfaceData], Range Double)
mkSurfaceData f r g cs = (zipWith SurfaceData rects (flip mixes cs <$> proj), x)
  where
    ps = gridF f r g
    rects = fst <$> ps
    vs = snd <$> ps
    x = unsafeSpace1 vs :: Range Double
    proj = project x (Range 0 1) <$> vs

-- | Create a surface chart from a function.
surfacef :: (Point Double -> Double) -> SurfaceOptions -> ([Chart], Range Double)
surfacef f cfg =
  first (surfaces (view (#soStyle % #surfaceRectStyle) cfg)) $
    mkSurfaceData
      f
      (view #soRange cfg)
      (view #soGrain cfg)
      (toList $ view (#soStyle % #surfaceColors) cfg)

-- | Legend specialization for a surface chart.
data SurfaceLegendOptions = SurfaceLegendOptions
  { sloAxisOptions :: AxisOptions,
    -- | Width of the legend glyph
    sloWidth :: Double,
    -- | Resolution of the legend glyph
    sloResolution :: Int,
    sloDataRange :: Range Double,
    -- | Placement of the legend versus normalised chart placement
    sloRect :: Rect Double,
    sloSurfaceStyle :: SurfaceStyle
  }
  deriving (Eq, Show, Generic, Data)

-- | 'AxisOptions' for a surface chart legend.
surfaceLegendAxisOptions :: AxisOptions
surfaceLegendAxisOptions =
  AxisOptions
    Nothing
    Nothing
    ( Ticks
        (TickRound (FormatN FSPrec (Just 3) 4 True True) 4 NoTickExtend)
        (Just defaultGlyphTickStyleY)
        (Just (defaultTextTick & set #buffer 0.05))
        Nothing
    )
    PlaceRight

-- | official surface legend options
defaultSurfaceLegendOptions :: SurfaceLegendOptions
defaultSurfaceLegendOptions =
  SurfaceLegendOptions surfaceLegendAxisOptions 0.2 100 one (Rect 0.7 0.9 0 0.5) defaultSurfaceStyle

-- | Chart used as a reference to a surface chart.
gridReferenceChart :: SurfaceLegendOptions -> ChartTree
gridReferenceChart slo =
  named "grid reference" $
    zipWith
      (\r c -> Chart (blob c) (RectData [r]))
      (gridf <$> spaceGrid)
      colorGrid
  where
    spaceGrid = gridSpace (view #sloDataRange slo) (view #sloResolution slo)
    gridf =
      bool
        (\yr -> Ranges (Range 0 (view #sloWidth slo)) yr)
        (\xr -> Ranges xr (Range 0 (view #sloWidth slo)))
        (isHori slo)
    colorGrid =
      (\x -> mixes x (toList $ view (#sloSurfaceStyle % #surfaceColors) slo))
        <$> grid MidPos (Range 0 1) (view #sloResolution slo)

isHori :: SurfaceLegendOptions -> Bool
isHori slo =
  view (#sloAxisOptions % #place) slo == PlaceBottom
    || view (#sloAxisOptions % #place) slo == PlaceTop

-- | Add a surface legend to a main surface chart.
addSurfaceLegend :: SurfaceLegendOptions -> ChartTree -> ChartTree
addSurfaceLegend slo ct = ctBoth
  where
    grc = gridReferenceChart slo
    hoLegend = (mempty :: HudOptions) & set #axes [Priority 1 (view #sloAxisOptions slo)]
    grcLegend = addHud (FixedAspect (view #sloWidth slo)) hoLegend grc
    ctbox = fromMaybe one (view styleBox' ct)
    legbox = projectOnR ctbox one (view #sloRect slo)
    ctBoth = mconcat [projectChartTree legbox grcLegend, ct]