{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE KindSignatures #-}

{- |
The schema-checked, typed-expression analogue of "DataFrame.Display.Web.Chart".
Charts are built over a 'TypedDataFrame' using typed expressions ('TExpr',
@#col@), so column references are checked against the schema at compile time
and the field type follows from the column's Haskell type.

@
{\-\# LANGUAGE OverloadedLabels \#-\}
import qualified DataFrame.Display.Web.Chart.Typed as Plt

-- one-liner
Plt.scatter #age #fare tdf

-- composable grammar
Plt.showChart
  ( Plt.chart tdf
      |> Plt.mark Plt.Boxplot
      |> Plt.enc Plt.X #region
      |> Plt.enc Plt.Y #value
      |> Plt.enc Plt.Color #grp
      |> Plt.facet #grp
  )
@

Every combinator delegates to "DataFrame.Display.Web.Chart" by unwrapping the
typed expression ('unTExpr') and frame ('unTDF'); the @cols@ phantom is erased
at the boundary.
-}
module DataFrame.Display.Web.Chart.Typed (
    -- * Chart
    Chart,
    chart,

    -- * Re-exported spec vocabulary
    Mark (..),
    Channel (..),
    FieldType (..),
    Agg (..),

    -- * Building blocks
    mark,
    enc,
    encAs,
    aggregateOn,
    binX,
    binY,
    facet,
    row,
    column,
    layer,
    regression,
    density,
    logScale,
    includeZero,
    title,
    size,

    -- * Rendering
    toVegaSpec,
    toHtml,
    showChart,

    -- * One-shot convenience plots
    scatter,
    bar,
    histogram,
    line,
    pie,
    box,
) where

import Data.Aeson (Value)
import Data.Kind (Type)
import qualified Data.Text as T

import DataFrame.Display.Internal.Common (Agg (..))
import DataFrame.Display.Internal.VegaLite (
    Channel (..),
    FieldType (..),
    Mark (..),
 )
import qualified DataFrame.Display.Web.Chart as C
import DataFrame.Internal.Column (Columnable)
import DataFrame.Typed.Types (TExpr (..), TypedDataFrame (..))
import GHC.TypeLits (Symbol)

-- | A typed chart: a phantom-@cols@ wrapper over the untyped builder.
newtype Chart (cols :: [(Symbol, Type)]) = Chart C.Chart

-- | Start a typed chart from a typed frame.
chart :: TypedDataFrame cols -> Chart cols
chart :: forall (cols :: [(Symbol, *)]). TypedDataFrame cols -> Chart cols
chart TypedDataFrame cols
tdf = Chart -> Chart cols
forall (cols :: [(Symbol, *)]). Chart -> Chart cols
Chart (DataFrame -> Chart
C.chart (TypedDataFrame cols -> DataFrame
forall (cols :: [(Symbol, *)]). TypedDataFrame cols -> DataFrame
unTDF TypedDataFrame cols
tdf))

-- | Set the mark type.
mark :: Mark -> Chart cols -> Chart cols
mark :: forall (cols :: [(Symbol, *)]). Mark -> Chart cols -> Chart cols
mark Mark
m (Chart Chart
c) = Chart -> Chart cols
forall (cols :: [(Symbol, *)]). Chart -> Chart cols
Chart (Mark -> Chart -> Chart
C.mark Mark
m Chart
c)

-- | Encode a typed expression on a channel (field type inferred from its type).
enc :: (Columnable a) => Channel -> TExpr cols a -> Chart cols -> Chart cols
enc :: forall a (cols :: [(Symbol, *)]).
Columnable a =>
Channel -> TExpr cols a -> Chart cols -> Chart cols
enc Channel
ch TExpr cols a
te (Chart Chart
c) = Chart -> Chart cols
forall (cols :: [(Symbol, *)]). Chart -> Chart cols
Chart (Channel -> Expr a -> Chart -> Chart
forall a. Columnable a => Channel -> Expr a -> Chart -> Chart
C.enc Channel
ch (TExpr cols a -> Expr a
forall (cols :: [(Symbol, *)]) a. TExpr cols a -> Expr a
unTExpr TExpr cols a
te) Chart
c)

-- | Like 'enc' but force the Vega-Lite field type.
encAs ::
    (Columnable a) =>
    Channel -> TExpr cols a -> FieldType -> Chart cols -> Chart cols
encAs :: forall a (cols :: [(Symbol, *)]).
Columnable a =>
Channel -> TExpr cols a -> FieldType -> Chart cols -> Chart cols
encAs Channel
ch TExpr cols a
te FieldType
ft (Chart Chart
c) = Chart -> Chart cols
forall (cols :: [(Symbol, *)]). Chart -> Chart cols
Chart (Channel -> Expr a -> FieldType -> Chart -> Chart
forall a.
Columnable a =>
Channel -> Expr a -> FieldType -> Chart -> Chart
C.encAs Channel
ch (TExpr cols a -> Expr a
forall (cols :: [(Symbol, *)]) a. TExpr cols a -> Expr a
unTExpr TExpr cols a
te) FieldType
ft Chart
c)

-- | Apply a declarative aggregation to a channel.
aggregateOn :: Channel -> Agg -> Chart cols -> Chart cols
aggregateOn :: forall (cols :: [(Symbol, *)]).
Channel -> Agg -> Chart cols -> Chart cols
aggregateOn Channel
ch Agg
a (Chart Chart
c) = Chart -> Chart cols
forall (cols :: [(Symbol, *)]). Chart -> Chart cols
Chart (Channel -> Agg -> Chart -> Chart
C.aggregateOn Channel
ch Agg
a Chart
c)

-- | Bin the X (resp. Y) channel.
binX, binY :: Chart cols -> Chart cols
binX :: forall (cols :: [(Symbol, *)]). Chart cols -> Chart cols
binX (Chart Chart
c) = Chart -> Chart cols
forall (cols :: [(Symbol, *)]). Chart -> Chart cols
Chart (Chart -> Chart
C.binX Chart
c)
binY :: forall (cols :: [(Symbol, *)]). Chart cols -> Chart cols
binY (Chart Chart
c) = Chart -> Chart cols
forall (cols :: [(Symbol, *)]). Chart -> Chart cols
Chart (Chart -> Chart
C.binY Chart
c)

-- | Put a channel on a log scale.
logScale :: Channel -> Chart cols -> Chart cols
logScale :: forall (cols :: [(Symbol, *)]). Channel -> Chart cols -> Chart cols
logScale Channel
ch (Chart Chart
c) = Chart -> Chart cols
forall (cols :: [(Symbol, *)]). Chart -> Chart cols
Chart (Channel -> Chart -> Chart
C.logScale Channel
ch Chart
c)

-- | Anchor (@True@) or release (@False@) a channel's scale at zero.
includeZero :: Channel -> Bool -> Chart cols -> Chart cols
includeZero :: forall (cols :: [(Symbol, *)]).
Channel -> Bool -> Chart cols -> Chart cols
includeZero Channel
ch Bool
b (Chart Chart
c) = Chart -> Chart cols
forall (cols :: [(Symbol, *)]). Chart -> Chart cols
Chart (Channel -> Bool -> Chart -> Chart
C.includeZero Channel
ch Bool
b Chart
c)

-- | Facet into small multiples by a column (alias for 'column').
facet :: (Columnable a) => TExpr cols a -> Chart cols -> Chart cols
facet :: forall a (cols :: [(Symbol, *)]).
Columnable a =>
TExpr cols a -> Chart cols -> Chart cols
facet = TExpr cols a -> Chart cols -> Chart cols
forall a (cols :: [(Symbol, *)]).
Columnable a =>
TExpr cols a -> Chart cols -> Chart cols
column

-- | Facet across columns / down rows.
column, row :: (Columnable a) => TExpr cols a -> Chart cols -> Chart cols
column :: forall a (cols :: [(Symbol, *)]).
Columnable a =>
TExpr cols a -> Chart cols -> Chart cols
column TExpr cols a
te (Chart Chart
c) = Chart -> Chart cols
forall (cols :: [(Symbol, *)]). Chart -> Chart cols
Chart (Expr a -> Chart -> Chart
forall a. Columnable a => Expr a -> Chart -> Chart
C.column (TExpr cols a -> Expr a
forall (cols :: [(Symbol, *)]) a. TExpr cols a -> Expr a
unTExpr TExpr cols a
te) Chart
c)
row :: forall a (cols :: [(Symbol, *)]).
Columnable a =>
TExpr cols a -> Chart cols -> Chart cols
row TExpr cols a
te (Chart Chart
c) = Chart -> Chart cols
forall (cols :: [(Symbol, *)]). Chart -> Chart cols
Chart (Expr a -> Chart -> Chart
forall a. Columnable a => Expr a -> Chart -> Chart
C.row (TExpr cols a -> Expr a
forall (cols :: [(Symbol, *)]) a. TExpr cols a -> Expr a
unTExpr TExpr cols a
te) Chart
c)

-- | Set the chart title.
title :: T.Text -> Chart cols -> Chart cols
title :: forall (cols :: [(Symbol, *)]). Text -> Chart cols -> Chart cols
title Text
t (Chart Chart
c) = Chart -> Chart cols
forall (cols :: [(Symbol, *)]). Chart -> Chart cols
Chart (Text -> Chart -> Chart
C.title Text
t Chart
c)

-- | Set the chart size in pixels.
size :: Int -> Int -> Chart cols -> Chart cols
size :: forall (cols :: [(Symbol, *)]).
Int -> Int -> Chart cols -> Chart cols
size Int
w Int
h (Chart Chart
c) = Chart -> Chart cols
forall (cols :: [(Symbol, *)]). Chart -> Chart cols
Chart (Int -> Int -> Chart -> Chart
C.size Int
w Int
h Chart
c)

-- | Overlay several charts sharing data into a single layered chart.
layer :: [Chart cols] -> Chart cols
layer :: forall (cols :: [(Symbol, *)]). [Chart cols] -> Chart cols
layer [Chart cols]
cs = Chart -> Chart cols
forall (cols :: [(Symbol, *)]). Chart -> Chart cols
Chart ([Chart] -> Chart
C.layer [Chart
c | Chart Chart
c <- [Chart cols]
cs])

-- | Add a least-squares regression line fitting @y@ on @x@.
regression ::
    (Columnable a, Columnable b) =>
    TExpr cols a -> TExpr cols b -> Chart cols -> Chart cols
regression :: forall a b (cols :: [(Symbol, *)]).
(Columnable a, Columnable b) =>
TExpr cols a -> TExpr cols b -> Chart cols -> Chart cols
regression TExpr cols a
xE TExpr cols b
yE (Chart Chart
c) = Chart -> Chart cols
forall (cols :: [(Symbol, *)]). Chart -> Chart cols
Chart (Expr a -> Expr b -> Chart -> Chart
forall a b.
(Columnable a, Columnable b) =>
Expr a -> Expr b -> Chart -> Chart
C.regression (TExpr cols a -> Expr a
forall (cols :: [(Symbol, *)]) a. TExpr cols a -> Expr a
unTExpr TExpr cols a
xE) (TExpr cols b -> Expr b
forall (cols :: [(Symbol, *)]) a. TExpr cols a -> Expr a
unTExpr TExpr cols b
yE) Chart
c)

-- | Kernel-density estimate of an expression, drawn as an area.
density :: (Columnable a) => TExpr cols a -> Chart cols -> Chart cols
density :: forall a (cols :: [(Symbol, *)]).
Columnable a =>
TExpr cols a -> Chart cols -> Chart cols
density TExpr cols a
e (Chart Chart
c) = Chart -> Chart cols
forall (cols :: [(Symbol, *)]). Chart -> Chart cols
Chart (Expr a -> Chart -> Chart
forall a. Columnable a => Expr a -> Chart -> Chart
C.density (TExpr cols a -> Expr a
forall (cols :: [(Symbol, *)]) a. TExpr cols a -> Expr a
unTExpr TExpr cols a
e) Chart
c)

-- | The Vega-Lite spec as an aeson 'Value' (escape hatch for advanced use / hvega).
toVegaSpec :: Chart cols -> Value
toVegaSpec :: forall (cols :: [(Symbol, *)]). Chart cols -> Value
toVegaSpec (Chart Chart
c) = Chart -> Value
C.toVegaSpec Chart
c

-- | A self-contained HTML snippet embedding the chart.
toHtml :: Chart cols -> String
toHtml :: forall (cols :: [(Symbol, *)]). Chart cols -> String
toHtml (Chart Chart
c) = Chart -> String
C.toHtml Chart
c

-- | Render the chart to a temp file and open it in the default browser.
showChart :: Chart cols -> IO ()
showChart :: forall (cols :: [(Symbol, *)]). Chart cols -> IO ()
showChart (Chart Chart
c) = Chart -> IO ()
C.showChart Chart
c

-- | Scatter plot of two typed expressions.
scatter ::
    (Columnable a, Columnable b) =>
    TExpr cols a -> TExpr cols b -> TypedDataFrame cols -> IO ()
scatter :: forall a b (cols :: [(Symbol, *)]).
(Columnable a, Columnable b) =>
TExpr cols a -> TExpr cols b -> TypedDataFrame cols -> IO ()
scatter TExpr cols a
xE TExpr cols b
yE TypedDataFrame cols
tdf = Expr a -> Expr b -> DataFrame -> IO ()
forall a b.
(Columnable a, Columnable b) =>
Expr a -> Expr b -> DataFrame -> IO ()
C.scatter (TExpr cols a -> Expr a
forall (cols :: [(Symbol, *)]) a. TExpr cols a -> Expr a
unTExpr TExpr cols a
xE) (TExpr cols b -> Expr b
forall (cols :: [(Symbol, *)]) a. TExpr cols a -> Expr a
unTExpr TExpr cols b
yE) (TypedDataFrame cols -> DataFrame
forall (cols :: [(Symbol, *)]). TypedDataFrame cols -> DataFrame
unTDF TypedDataFrame cols
tdf)

-- | Count of rows per category, as bars.
bar :: (Columnable a) => TExpr cols a -> TypedDataFrame cols -> IO ()
bar :: forall a (cols :: [(Symbol, *)]).
Columnable a =>
TExpr cols a -> TypedDataFrame cols -> IO ()
bar TExpr cols a
xE TypedDataFrame cols
tdf = Expr a -> DataFrame -> IO ()
forall a. Columnable a => Expr a -> DataFrame -> IO ()
C.bar (TExpr cols a -> Expr a
forall (cols :: [(Symbol, *)]) a. TExpr cols a -> Expr a
unTExpr TExpr cols a
xE) (TypedDataFrame cols -> DataFrame
forall (cols :: [(Symbol, *)]). TypedDataFrame cols -> DataFrame
unTDF TypedDataFrame cols
tdf)

-- | Histogram of a numeric expression.
histogram :: (Columnable a) => TExpr cols a -> TypedDataFrame cols -> IO ()
histogram :: forall a (cols :: [(Symbol, *)]).
Columnable a =>
TExpr cols a -> TypedDataFrame cols -> IO ()
histogram TExpr cols a
xE TypedDataFrame cols
tdf = Expr a -> DataFrame -> IO ()
forall a. Columnable a => Expr a -> DataFrame -> IO ()
C.histogram (TExpr cols a -> Expr a
forall (cols :: [(Symbol, *)]) a. TExpr cols a -> Expr a
unTExpr TExpr cols a
xE) (TypedDataFrame cols -> DataFrame
forall (cols :: [(Symbol, *)]). TypedDataFrame cols -> DataFrame
unTDF TypedDataFrame cols
tdf)

-- | Line chart of @y@ over @x@.
line ::
    (Columnable a, Columnable b) =>
    TExpr cols a -> TExpr cols b -> TypedDataFrame cols -> IO ()
line :: forall a b (cols :: [(Symbol, *)]).
(Columnable a, Columnable b) =>
TExpr cols a -> TExpr cols b -> TypedDataFrame cols -> IO ()
line TExpr cols a
xE TExpr cols b
yE TypedDataFrame cols
tdf = Expr a -> Expr b -> DataFrame -> IO ()
forall a b.
(Columnable a, Columnable b) =>
Expr a -> Expr b -> DataFrame -> IO ()
C.line (TExpr cols a -> Expr a
forall (cols :: [(Symbol, *)]) a. TExpr cols a -> Expr a
unTExpr TExpr cols a
xE) (TExpr cols b -> Expr b
forall (cols :: [(Symbol, *)]) a. TExpr cols a -> Expr a
unTExpr TExpr cols b
yE) (TypedDataFrame cols -> DataFrame
forall (cols :: [(Symbol, *)]). TypedDataFrame cols -> DataFrame
unTDF TypedDataFrame cols
tdf)

-- | Pie chart counting rows per category.
pie :: (Columnable a) => TExpr cols a -> TypedDataFrame cols -> IO ()
pie :: forall a (cols :: [(Symbol, *)]).
Columnable a =>
TExpr cols a -> TypedDataFrame cols -> IO ()
pie TExpr cols a
cE TypedDataFrame cols
tdf = Expr a -> DataFrame -> IO ()
forall a. Columnable a => Expr a -> DataFrame -> IO ()
C.pie (TExpr cols a -> Expr a
forall (cols :: [(Symbol, *)]) a. TExpr cols a -> Expr a
unTExpr TExpr cols a
cE) (TypedDataFrame cols -> DataFrame
forall (cols :: [(Symbol, *)]). TypedDataFrame cols -> DataFrame
unTDF TypedDataFrame cols
tdf)

-- | Box-and-whisker plot of a numeric expression.
box :: (Columnable a) => TExpr cols a -> TypedDataFrame cols -> IO ()
box :: forall a (cols :: [(Symbol, *)]).
Columnable a =>
TExpr cols a -> TypedDataFrame cols -> IO ()
box TExpr cols a
yE TypedDataFrame cols
tdf = Expr a -> DataFrame -> IO ()
forall a. Columnable a => Expr a -> DataFrame -> IO ()
C.box (TExpr cols a -> Expr a
forall (cols :: [(Symbol, *)]) a. TExpr cols a -> Expr a
unTExpr TExpr cols a
yE) (TypedDataFrame cols -> DataFrame
forall (cols :: [(Symbol, *)]). TypedDataFrame cols -> DataFrame
unTDF TypedDataFrame cols
tdf)