VegaLite.hs

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

-- | Internal Vega-Lite spec model and aeson encoding for the web backends.
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 Control.Exception (throw)
import DataFrame.Internal.Column (
    Column (MergedColumn),
    Columnable,
    bitmapTestBit,
    columnBitmap,
    materializeMerged,
 )
import DataFrame.Internal.DataFrame (DataFrame, dataframeDimensions, getColumn)
import DataFrame.Internal.Expression (Expr (Col))
import DataFrame.Internal.Interpreter (Ctx (..), eval, materialize)

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

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.
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: 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@.
data Sort = DataOrder
    deriving (Eq, Show)

-- | Per-channel scale configuration, mirroring Vega-Lite's @scale@ object.
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, layered when 'vlLayers' is non-empty.
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 []

-- | Derive the Vega-Lite field type from a Haskell 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"]

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

-- | Resolve an expression against a frame, reusing a bare column directly.
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 :: forall a. (Columnable a) => DataFrame -> Expr a -> Column
materialiseExpr df expr = either throw id $ do
    v <- eval (FlatCtx df) expr
    pure $ materialize @a (fst (dataframeDimensions df)) v

-- | A column's values as JSON, with null slots as 'Null'.
columnToValues :: FieldType -> Column -> [Value]
columnToValues ft col = case col of
    MergedColumn _ _ -> columnToValues ft (materializeMerged col)
    _ -> maskNulls (raw ft)
  where
    raw Quantitative = map toJSON (columnToDoubles col)
    raw _ = map (toJSON :: T.Text -> Value) (columnToStrings col)
    maskNulls vs = case columnBitmap col of
        Nothing -> vs
        Just bm -> zipWith (\i v -> if bitmapTestBit bm i then v else Null) [0 ..] vs

-- | 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 pre-computed numeric values.
numField :: T.Text -> [Double] -> ResolvedField
numField name vals = ResolvedField name Quantitative (map toJSON vals)

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 (data 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" .= fieldRef (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" .= fieldRef yField, "on" .= fieldRef xField]
transformValue (DensityT field) =
    object ["density" .= fieldRef field]

-- | A column as a Vega-Lite field; escapes @.@, @[@, @]@ (else nested).
fieldRef :: T.Text -> T.Text
fieldRef = T.concatMap esc
  where
    esc c
        | c `elem` ['.', '[', ']', '\\'] = T.pack ['\\', c]
        | otherwise = T.singleton c

-- | 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]]

-- | Render a spec to a self-contained HTML snippet with inlined data.
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.
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."