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, Columnable)
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

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 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" .= 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]]

-- | 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."