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

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

    -- * Numeric vector extraction
    columnAsIntVector,
    columnAsDoubleVector,
    columnAsFloatVector,
    columnAsUnboxedVector,

    -- * Matrix extraction
    toDoubleMatrix,
    toFloatMatrix,
    toIntMatrix,
) where

import Control.Exception (throw)
import Data.Proxy (Proxy (..))
import qualified Data.Text as T
import qualified Data.Vector as V
import qualified Data.Vector.Unboxed as VU
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 (
    AllColumnsReal,
    AssertPresent,
    AssertRealColumn,
    SafeLookup,
 )
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 ~ SafeLookup name cols
    , Columnable a
    , AssertPresent name cols
    ) =>
    TypedDataFrame cols -> V.Vector a
columnAsVector :: forall (name :: Symbol) (cols :: [(Symbol, *)]) a.
(KnownSymbol name, a ~ SafeLookup 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 ~ SafeLookup name cols
    , Columnable a
    , AssertPresent name cols
    ) =>
    TypedDataFrame cols -> [a]
columnAsList :: forall (name :: Symbol) (cols :: [(Symbol, *)]) a.
(KnownSymbol name, a ~ SafeLookup 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))

{- | Retrieve a column coerced to an unboxed 'Int' vector, named by type
application. The column must exist and be numeric — both are compile-time
checks via 'SafeLookup', so this is total (no 'Either', no runtime throw).
-}
columnAsIntVector ::
    forall name cols a.
    ( KnownSymbol name
    , a ~ SafeLookup name cols
    , Columnable a
    , AssertRealColumn "columnAsIntVector" name a
    , Real a
    , VU.Unbox a
    , AssertPresent name cols
    ) =>
    TypedDataFrame cols -> VU.Vector Int
columnAsIntVector :: forall (name :: Symbol) (cols :: [(Symbol, *)]) a.
(KnownSymbol name, a ~ SafeLookup name cols, Columnable a,
 AssertRealColumn "columnAsIntVector" name a, Real a, Unbox a,
 AssertPresent name cols) =>
TypedDataFrame cols -> Vector Int
columnAsIntVector (TDF DataFrame
df) = (DataFrameException -> Vector Int)
-> (Vector Int -> Vector Int)
-> Either DataFrameException (Vector Int)
-> Vector Int
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either DataFrameException -> Vector Int
forall a e. Exception e => e -> a
throw Vector Int -> Vector Int
forall a. a -> a
id (Expr a -> DataFrame -> Either DataFrameException (Vector Int)
forall a.
(Columnable a, Num a) =>
Expr a -> DataFrame -> Either DataFrameException (Vector Int)
D.columnAsIntVector (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 coerced to an unboxed 'Double' vector. See 'columnAsIntVector'.
columnAsDoubleVector ::
    forall name cols a.
    ( KnownSymbol name
    , a ~ SafeLookup name cols
    , Columnable a
    , AssertRealColumn "columnAsDoubleVector" name a
    , Real a
    , VU.Unbox a
    , AssertPresent name cols
    ) =>
    TypedDataFrame cols -> VU.Vector Double
columnAsDoubleVector :: forall (name :: Symbol) (cols :: [(Symbol, *)]) a.
(KnownSymbol name, a ~ SafeLookup name cols, Columnable a,
 AssertRealColumn "columnAsDoubleVector" name a, Real a, Unbox a,
 AssertPresent name cols) =>
TypedDataFrame cols -> Vector Double
columnAsDoubleVector (TDF DataFrame
df) =
    (DataFrameException -> Vector Double)
-> (Vector Double -> Vector Double)
-> Either DataFrameException (Vector Double)
-> Vector Double
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either DataFrameException -> Vector Double
forall a e. Exception e => e -> a
throw Vector Double -> Vector Double
forall a. a -> a
id (Expr a -> DataFrame -> Either DataFrameException (Vector Double)
forall a.
(Columnable a, Num a) =>
Expr a -> DataFrame -> Either DataFrameException (Vector Double)
D.columnAsDoubleVector (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 coerced to an unboxed 'Float' vector. See 'columnAsIntVector'.
columnAsFloatVector ::
    forall name cols a.
    ( KnownSymbol name
    , a ~ SafeLookup name cols
    , Columnable a
    , AssertRealColumn "columnAsFloatVector" name a
    , Real a
    , VU.Unbox a
    , AssertPresent name cols
    ) =>
    TypedDataFrame cols -> VU.Vector Float
columnAsFloatVector :: forall (name :: Symbol) (cols :: [(Symbol, *)]) a.
(KnownSymbol name, a ~ SafeLookup name cols, Columnable a,
 AssertRealColumn "columnAsFloatVector" name a, Real a, Unbox a,
 AssertPresent name cols) =>
TypedDataFrame cols -> Vector Float
columnAsFloatVector (TDF DataFrame
df) =
    (DataFrameException -> Vector Float)
-> (Vector Float -> Vector Float)
-> Either DataFrameException (Vector Float)
-> Vector Float
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either DataFrameException -> Vector Float
forall a e. Exception e => e -> a
throw Vector Float -> Vector Float
forall a. a -> a
id (Expr a -> DataFrame -> Either DataFrameException (Vector Float)
forall a.
(Columnable a, Num a) =>
Expr a -> DataFrame -> Either DataFrameException (Vector Float)
D.columnAsFloatVector (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 an unboxed vector of its own element type. The column
must exist and be unboxable — both compile-time checks, so this is total.
-}
columnAsUnboxedVector ::
    forall name cols a.
    ( KnownSymbol name
    , a ~ SafeLookup name cols
    , Columnable a
    , VU.Unbox a
    , AssertPresent name cols
    ) =>
    TypedDataFrame cols -> VU.Vector a
columnAsUnboxedVector :: forall (name :: Symbol) (cols :: [(Symbol, *)]) a.
(KnownSymbol name, a ~ SafeLookup name cols, Columnable a, Unbox a,
 AssertPresent name cols) =>
TypedDataFrame cols -> Vector a
columnAsUnboxedVector (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 (Expr a -> DataFrame -> Either DataFrameException (Vector a)
forall a.
(Columnable a, Unbox a) =>
Expr a -> DataFrame -> Either DataFrameException (Vector a)
D.columnAsUnboxedVector (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))

{- | Convert every column to 'Double' and transpose into a row-major matrix.
Total: 'AllColumnsReal' proves at compile time that every column is numeric and
unboxed, so the conversion cannot fail.
-}
toDoubleMatrix ::
    (AllColumnsReal "toDoubleMatrix" cols) =>
    TypedDataFrame cols -> V.Vector (VU.Vector Double)
toDoubleMatrix :: forall (cols :: [(Symbol, *)]).
AllColumnsReal "toDoubleMatrix" cols =>
TypedDataFrame cols -> Vector (Vector Double)
toDoubleMatrix (TDF DataFrame
df) = (DataFrameException -> Vector (Vector Double))
-> (Vector (Vector Double) -> Vector (Vector Double))
-> Either DataFrameException (Vector (Vector Double))
-> Vector (Vector Double)
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either DataFrameException -> Vector (Vector Double)
forall a e. Exception e => e -> a
throw Vector (Vector Double) -> Vector (Vector Double)
forall a. a -> a
id (DataFrame -> Either DataFrameException (Vector (Vector Double))
D.toDoubleMatrix DataFrame
df)

-- | Convert every column to 'Float' and transpose into a row-major matrix. See 'toDoubleMatrix'.
toFloatMatrix ::
    (AllColumnsReal "toFloatMatrix" cols) =>
    TypedDataFrame cols -> V.Vector (VU.Vector Float)
toFloatMatrix :: forall (cols :: [(Symbol, *)]).
AllColumnsReal "toFloatMatrix" cols =>
TypedDataFrame cols -> Vector (Vector Float)
toFloatMatrix (TDF DataFrame
df) = (DataFrameException -> Vector (Vector Float))
-> (Vector (Vector Float) -> Vector (Vector Float))
-> Either DataFrameException (Vector (Vector Float))
-> Vector (Vector Float)
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either DataFrameException -> Vector (Vector Float)
forall a e. Exception e => e -> a
throw Vector (Vector Float) -> Vector (Vector Float)
forall a. a -> a
id (DataFrame -> Either DataFrameException (Vector (Vector Float))
D.toFloatMatrix DataFrame
df)

-- | Convert every column to 'Int' and transpose into a row-major matrix. See 'toDoubleMatrix'.
toIntMatrix ::
    (AllColumnsReal "toIntMatrix" cols) =>
    TypedDataFrame cols -> V.Vector (VU.Vector Int)
toIntMatrix :: forall (cols :: [(Symbol, *)]).
AllColumnsReal "toIntMatrix" cols =>
TypedDataFrame cols -> Vector (Vector Int)
toIntMatrix (TDF DataFrame
df) = (DataFrameException -> Vector (Vector Int))
-> (Vector (Vector Int) -> Vector (Vector Int))
-> Either DataFrameException (Vector (Vector Int))
-> Vector (Vector Int)
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either DataFrameException -> Vector (Vector Int)
forall a e. Exception e => e -> a
throw Vector (Vector Int) -> Vector (Vector Int)
forall a. a -> a
id (DataFrame -> Either DataFrameException (Vector (Vector Int))
D.toIntMatrix DataFrame
df)