{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}

module DataFrame.Typed.Access (
    -- * Typed column access
    columnAsVector,
    columnAsList,
) where

import Control.Exception (throw)
import Data.Proxy (Proxy (..))
import qualified Data.Text as T
import qualified Data.Vector as V
import GHC.TypeLits (KnownSymbol, symbolVal)

import DataFrame.Internal.Column (Columnable)
import DataFrame.Internal.Expression (Expr (Col))
import qualified DataFrame.Operations.Core as D
import DataFrame.Typed.Schema (AssertPresent, Lookup)
import DataFrame.Typed.Types (TypedDataFrame (..))

{- | Retrieve a column as a boxed 'Vector', with the type determined by
the schema. The column must exist (enforced at compile time).
-}
columnAsVector ::
    forall name cols a.
    ( KnownSymbol name
    , a ~ Lookup name cols
    , Columnable a
    , AssertPresent name cols
    ) =>
    TypedDataFrame cols -> V.Vector a
columnAsVector :: forall (name :: Symbol) (cols :: [*]) a.
(KnownSymbol name, a ~ Lookup name cols, Columnable a,
 AssertPresent name cols) =>
TypedDataFrame cols -> Vector a
columnAsVector (TDF DataFrame
df) =
    (DataFrameException -> Vector a)
-> (Vector a -> Vector a)
-> Either DataFrameException (Vector a)
-> Vector a
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either DataFrameException -> Vector a
forall a e. Exception e => e -> a
throw Vector a -> Vector a
forall a. a -> a
id (Either DataFrameException (Vector a) -> Vector a)
-> Either DataFrameException (Vector a) -> Vector a
forall a b. (a -> b) -> a -> b
$ Expr a -> DataFrame -> Either DataFrameException (Vector a)
forall a.
Columnable a =>
Expr a -> DataFrame -> Either DataFrameException (Vector a)
D.columnAsVector (forall a. Columnable a => Text -> Expr a
Col @a Text
colName) DataFrame
df
  where
    colName :: Text
colName = String -> Text
T.pack (Proxy name -> String
forall (n :: Symbol) (proxy :: Symbol -> *).
KnownSymbol n =>
proxy n -> String
symbolVal (forall {k} (t :: k). Proxy t
forall (t :: Symbol). Proxy t
Proxy @name))

-- | Retrieve a column as a list, with the type determined by the schema.
columnAsList ::
    forall name cols a.
    ( KnownSymbol name
    , a ~ Lookup name cols
    , Columnable a
    , AssertPresent name cols
    ) =>
    TypedDataFrame cols -> [a]
columnAsList :: forall (name :: Symbol) (cols :: [*]) a.
(KnownSymbol name, a ~ Lookup name cols, Columnable a,
 AssertPresent name cols) =>
TypedDataFrame cols -> [a]
columnAsList (TDF DataFrame
df) =
    Expr a -> DataFrame -> [a]
forall a. Columnable a => Expr a -> DataFrame -> [a]
D.columnAsList (forall a. Columnable a => Text -> Expr a
Col @a Text
colName) DataFrame
df
  where
    colName :: Text
colName = String -> Text
T.pack (Proxy name -> String
forall (n :: Symbol) (proxy :: Symbol -> *).
KnownSymbol n =>
proxy n -> String
symbolVal (forall {k} (t :: k). Proxy t
forall (t :: Symbol). Proxy t
Proxy @name))