VegaLite.hs

{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}

{- |
Internal Vega-Lite spec model plus aeson encoding, shared by the web plot
backends. Not public. Channel encodings resolve to a 'ResolvedField' carrying
the field name, Vega-Lite type, and inlined column values.
-}
module DataFrame.Display.Internal.VegaLite (
    -- * Spec model
    Mark (..),
    Channel (..),
    FieldType (..),
    ChannelEnc (..),
    Sort (..),
    ScaleSpec (..),
    ScaleType (..),
    defaultScale,
    Transform (..),
    VLSpec (..),
    emptySpec,
    chanEnc,
    channelName,

    -- * Resolving expressions to fields
    ResolvedField (..),
    fieldTypeOf,
    resolveField,
    textField,
    numField,

    -- * Encoding to JSON / HTML
    specToValue,
    inlineRows,
    specHtml,
    rowCountWarning,
) where

import qualified Control.Monad
import Data.Aeson (ToJSON (toJSON), Value (Null), object, (.=))
import qualified Data.Aeson.Key as K
import Data.Aeson.Text (encodeToLazyText)
import qualified Data.List as L
import Data.Maybe (catMaybes, fromMaybe)
import qualified Data.Text as T
import qualified Data.Text.Lazy as TL
import qualified Data.Vector as V
import System.IO (hPutStrLn, stderr)
import Type.Reflection (TypeRep, tyConName, typeRep, typeRepTyCon, pattern App)

import DataFrame.Internal.Column (Column, Columnable, unwrapTypedColumn)
import DataFrame.Internal.DataFrame (DataFrame, getColumn)
import DataFrame.Internal.Expression (Expr (Col))
import DataFrame.Internal.Interpreter (interpret)

import DataFrame.Display.Internal.Common (columnToDoubles, columnToStrings)

-- ---------------------------------------------------------------------------
-- Spec model
-- ---------------------------------------------------------------------------

data Mark = Bar | Line | Point | Area | Boxplot | Arc | Rule | Tick
    deriving (Eq, Show)

markName :: Mark -> T.Text
markName m = case m of
    Bar -> "bar"
    Line -> "line"
    Point -> "point"
    Area -> "area"
    Boxplot -> "boxplot"
    Arc -> "arc"
    Rule -> "rule"
    Tick -> "tick"

data Channel
    = X
    | Y
    | Color
    | Size
    | Shape
    | Column
    | Row
    | Opacity
    | Theta
    | Tooltip
    | Order
    deriving (Eq, Show)

channelName :: Channel -> T.Text
channelName c = case c of
    X -> "x"
    Y -> "y"
    Color -> "color"
    Size -> "size"
    Shape -> "shape"
    Column -> "column"
    Row -> "row"
    Opacity -> "opacity"
    Theta -> "theta"
    Tooltip -> "tooltip"
    Order -> "order"

data FieldType = Quantitative | Nominal | Ordinal | Temporal
    deriving (Eq, Show)

fieldTypeName :: FieldType -> T.Text
fieldTypeName t = case t of
    Quantitative -> "quantitative"
    Nominal -> "nominal"
    Ordinal -> "ordinal"
    Temporal -> "temporal"

{- | A single channel encoding. An empty 'ceField' with a 'ceAggregate' of
@count@ produces a fieldless count aggregation, as Vega-Lite expects.
-}
data ChannelEnc = ChannelEnc
    { ceChannel :: Channel
    , ceField :: T.Text
    , ceType :: FieldType
    , ceAggregate :: Maybe T.Text
    , ceBin :: Bool
    , ceScale :: ScaleSpec
    , ceSort :: Maybe Sort
    }

-- | A bare channel encoding with no aggregation, binning, scale, or sort options.
chanEnc :: Channel -> T.Text -> FieldType -> ChannelEnc
chanEnc ch fld ft = ChannelEnc ch fld ft Nothing False defaultScale Nothing

{- | Domain ordering for a discrete channel, mirroring Vega-Lite's @sort@.
'DataOrder' pins the axis to the order rows appear in the inlined data, so
pre-sorted values (e.g. histogram bins) display in that order.
-}
data Sort = DataOrder
    deriving (Eq, Show)

{- | Per-channel scale configuration, mirroring Vega-Lite's @scale@ object.
'Nothing' fields leave the Vega-Lite default in place; grow this record
(domain, nice, clamp, scheme, ...) as the public vocabulary grows.
-}
data ScaleSpec = ScaleSpec
    { scaleType :: Maybe ScaleType
    , scaleZero :: Maybe Bool
    }
    deriving (Eq, Show)

data ScaleType = LogS
    deriving (Eq, Show)

defaultScale :: ScaleSpec
defaultScale = ScaleSpec Nothing Nothing

data Transform
    = -- | Fit @regression(yField) on xField@ (used inside a layer).
      RegressionT T.Text T.Text
    | -- | Kernel-density estimate of a field.
      DensityT T.Text
    deriving (Eq, Show)

{- | A Vega-Lite spec. When 'vlLayers' is non-empty the spec is a layered
container: the top level carries the shared data and each layer carries its own
mark/encoding/transform.
-}
data VLSpec = VLSpec
    { vlMark :: Mark
    , vlEncodings :: [ChannelEnc]
    , vlTransforms :: [Transform]
    , vlTitle :: Maybe T.Text
    , vlWidth :: Int
    , vlHeight :: Int
    , vlLayers :: [VLSpec]
    }

emptySpec :: Mark -> VLSpec
emptySpec m = VLSpec m [] [] Nothing 600 400 []

-- ---------------------------------------------------------------------------
-- Field-type inference from the expression's element type
-- ---------------------------------------------------------------------------

{- | Derive the Vega-Lite field type from a Haskell type: numeric →
'Quantitative', date/time → 'Temporal', everything else → 'Nominal'. @Maybe a@
classifies as its inner type.
-}
fieldTypeOf :: forall a. (Columnable a) => FieldType
fieldTypeOf = classify (typeRep @a)

classify :: forall k (x :: k). TypeRep x -> FieldType
classify tr
    | nm `elem` quantNames = Quantitative
    | nm `elem` temporalNames = Temporal
    | nm == "Maybe" = case tr of
        App _ arg -> classify arg
        _ -> Nominal
    | otherwise = Nominal
  where
    nm = tyConName (typeRepTyCon tr)
    quantNames =
        [ "Int"
        , "Int8"
        , "Int16"
        , "Int32"
        , "Int64"
        , "Word"
        , "Word8"
        , "Word16"
        , "Word32"
        , "Word64"
        , "Integer"
        , "Natural"
        , "Double"
        , "Float"
        , "Scientific"
        ]
    temporalNames =
        ["Day", "UTCTime", "LocalTime", "ZonedTime", "TimeOfDay"]

-- ---------------------------------------------------------------------------
-- Resolving an expression to a field + values
-- ---------------------------------------------------------------------------

data ResolvedField = ResolvedField
    { rfName :: T.Text
    , rfType :: FieldType
    , rfValues :: [Value]
    }

{- | Resolve an expression against a frame. A bare @Col name@ reuses the named
column; any other expression is materialised with the core interpreter and
stored under the given fallback name.
-}
resolveField ::
    forall a. (Columnable a) => DataFrame -> T.Text -> Expr a -> ResolvedField
resolveField df fallbackName expr =
    let ft = fieldTypeOf @a
        (name, col) = case expr of
            Col cname -> (cname, lookupCol cname)
            _ -> (fallbackName, materialiseExpr df expr)
     in ResolvedField name ft (columnToValues ft col)
  where
    lookupCol cname = case getColumn cname df of
        Just c -> c
        Nothing ->
            error $ "DataFrame.Display.Web: column not found: " <> T.unpack cname

materialiseExpr :: (Columnable a) => DataFrame -> Expr a -> Column
materialiseExpr df expr = case interpret df expr of
    Right tc -> unwrapTypedColumn tc
    Left err ->
        error $ "DataFrame.Display.Web: could not evaluate expression: " <> show err

columnToValues :: FieldType -> Column -> [Value]
columnToValues Quantitative col = map toJSON (columnToDoubles col)
columnToValues _ col = map (toJSON :: T.Text -> Value) (columnToStrings col)

-- | A nominal field built directly from text values (for pre-computed data).
textField :: T.Text -> [T.Text] -> ResolvedField
textField name vals = ResolvedField name Nominal (map (toJSON :: T.Text -> Value) vals)

-- | A quantitative field built directly from numeric values (for pre-computed data).
numField :: T.Text -> [Double] -> ResolvedField
numField name vals = ResolvedField name Quantitative (map toJSON vals)

-- ---------------------------------------------------------------------------
-- Encoding to JSON
-- ---------------------------------------------------------------------------

schemaUrl :: T.Text
schemaUrl = "https://vega.github.io/schema/vega-lite/v5.json"

specToValue :: [ResolvedField] -> VLSpec -> Value
specToValue fields spec
    | null (vlLayers spec) =
        object $
            commonPairs ++ unitPairs spec
    | otherwise =
        object $
            commonPairs ++ [("layer", toJSON (map (object . unitPairs) (vlLayers spec)))]
  where
    commonPairs =
        catMaybes
            [ Just ("$schema" .= schemaUrl)
            , fmap ("title" .=) (vlTitle spec)
            , Just ("width" .= vlWidth spec)
            , Just ("height" .= vlHeight spec)
            , Just ("data" .= object ["values" .= inlineRows fields])
            ]

-- | The mark/encoding/transform pairs of a unit spec (no data — that is shared).
unitPairs :: VLSpec -> [(K.Key, Value)]
unitPairs spec =
    catMaybes
        [ if null (vlTransforms spec)
            then Nothing
            else Just ("transform" .= map transformValue (vlTransforms spec))
        , Just ("mark" .= markValue (vlMark spec))
        , Just ("encoding" .= encodingValue (vlEncodings spec))
        ]

markValue :: Mark -> Value
markValue m = object ["type" .= markName m, "tooltip" .= True]

encodingValue :: [ChannelEnc] -> Value
encodingValue encs =
    object [K.fromText (channelName (ceChannel e)) .= channelValue e | e <- encs]

channelValue :: ChannelEnc -> Value
channelValue e =
    object $
        catMaybes
            [ if T.null (ceField e) then Nothing else Just ("field" .= ceField e)
            , Just ("type" .= fieldTypeName (ceType e))
            , fmap ("aggregate" .=) (ceAggregate e)
            , if ceBin e then Just ("bin" .= True) else Nothing
            , fmap ("scale" .=) (scaleValue (ceScale e))
            , ("sort" .=) . sortValue <$> ceSort e
            ]

-- | Render a sort spec to its Vega-Lite JSON value.
sortValue :: Sort -> Value
sortValue DataOrder = Null

-- | Render a scale spec, or 'Nothing' when every field is at its default.
scaleValue :: ScaleSpec -> Maybe Value
scaleValue s
    | null pairs = Nothing
    | otherwise = Just (object pairs)
  where
    pairs =
        catMaybes
            [ fmap (("type" .=) . scaleTypeName) (scaleType s)
            , fmap ("zero" .=) (scaleZero s)
            ]

scaleTypeName :: ScaleType -> T.Text
scaleTypeName LogS = "log"

transformValue :: Transform -> Value
transformValue (RegressionT yField xField) =
    object ["regression" .= yField, "on" .= xField]
transformValue (DensityT field) =
    object ["density" .= field]

-- | Build the @data.values@ array of row objects, deduplicating fields by name.
inlineRows :: [ResolvedField] -> Value
inlineRows fields =
    let uniq = L.nubBy (\a b -> rfName a == rfName b) fields
        vecs = [(rfName f, V.fromList (rfValues f)) | f <- uniq]
        n = maximum (0 : map (V.length . snd) vecs)
        row i = object [K.fromText nm .= fromMaybe Null (vs V.!? i) | (nm, vs) <- vecs]
     in toJSON [row i | i <- [0 .. n - 1]]

-- ---------------------------------------------------------------------------
-- HTML embedding (vega-embed via CDN)
-- ---------------------------------------------------------------------------

{- | Render a spec to a self-contained HTML snippet that loads vega/vega-lite/
vega-embed from a CDN and embeds the chart. Data is inlined, so the snippet
renders correctly even from a @file://@ URL.
-}
specHtml :: T.Text -> [ResolvedField] -> VLSpec -> T.Text
specHtml chartId fields spec =
    let specJson = TL.toStrict (encodeToLazyText (specToValue fields spec))
     in T.concat
            [ "<div id=\""
            , chartId
            , "\"></div>\n"
            , "<script src=\"https://cdn.jsdelivr.net/npm/vega@5\"></script>\n"
            , "<script src=\"https://cdn.jsdelivr.net/npm/vega-lite@5\"></script>\n"
            , "<script src=\"https://cdn.jsdelivr.net/npm/vega-embed@6\"></script>\n"
            , "<script>vegaEmbed('#"
            , chartId
            , "', "
            , specJson
            , ");</script>\n"
            ]

{- | Warn on stderr when a large number of rows is being inlined, which bloats
the spec and can slow the browser.
-}
rowCountWarning :: [ResolvedField] -> IO ()
rowCountWarning fields = do
    let n = maximum (0 : map (length . rfValues) fields)
    Control.Monad.when (n > 5000) $
        hPutStrLn stderr $
            "DataFrame.Display.Web: inlining "
                ++ show n
                ++ " rows into the plot spec; consider filtering or aggregating first."