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

{- |
Module      : DataFrame.Typed.Generic
License     : MIT

Generic-based opt-in for record-to-schema derivation. Mirrors the Template
Haskell splice in "DataFrame.Typed.TH" but builds the schema type from a
@GHC.Generics.Generic@ instance instead of @reify@.

Use it like this:

@
data Order = Order
  { orderId :: Int64
  , region  :: Text
  , amount  :: Double
  } deriving (Show, Eq, Generic)

type OrderSchema = SchemaOf Order

instance HasSchema Order OrderSchema where
  toColumns   = genericToColumns
  fromColumns = genericFromColumns
@

Field names are translated with the @CamelCase -> snake_case@ rule
(matching 'DataFrame.Typed.TH.camelToSnake'); use 'SchemaOfRaw' if you
want the schema to keep the record selector names verbatim — in that
case you cannot use 'genericToColumns' \/ 'genericFromColumns' and must
either hand-roll the instance or use the TH splice with a custom name
transform.
-}
module DataFrame.Typed.Generic (
    -- * Type-level schema derivation
    NameCase (..),
    SchemaOf,
    SchemaOfRaw,
    RepToSchema,
    CamelToSnake,

    -- * Value-level default methods
    genericToColumns,
    genericFromColumns,
    GHasColumns,
) where

import Data.Kind (Type)
import Data.Proxy (Proxy (..))
import qualified Data.Text as T
import qualified Data.Vector as VB
import GHC.Generics (
    C,
    D,
    Generic (..),
    K1 (..),
    M1 (..),
    Meta (..),
    S,
    type (:*:) (..),
 )
import GHC.TypeLits (
    CharToNat,
    ConsSymbol,
    KnownSymbol,
    NatToChar,
    Symbol,
    UnconsSymbol,
    symbolVal,
    type (+),
 )

import Data.Type.Bool (If, type (&&))
import Data.Type.Ord (type (<=?))

import qualified DataFrame.Internal.Column as C
import qualified DataFrame.Internal.DataFrame as D
import DataFrame.Typed.Record (requireColumn)
import DataFrame.Typed.Schema (Append)
import DataFrame.Typed.Util (camelToSnake)

{- | Field-name policy applied to record selectors when computing
'RepToSchema'.

* 'SnakeCase' — translate @camelCaseField@ to @\"camel_case_field\"@.
* 'IdentityCase' — keep the selector name verbatim.
-}
data NameCase = SnakeCase | IdentityCase

{- | The schema type @'[ '(name, ty), ...]@ derived from the 'Rep' of a
record type, with the given 'NameCase' applied to each field name.
-}
type family RepToSchema (nc :: NameCase) (r :: Type -> Type) :: [(Symbol, Type)] where
    RepToSchema nc (M1 D _ f) = RepToSchema nc f
    RepToSchema nc (M1 C _ f) = RepToSchema nc f
    RepToSchema nc (a :*: b) = Append (RepToSchema nc a) (RepToSchema nc b)
    RepToSchema nc (M1 S ('MetaSel ('Just name) _ _ _) (K1 _ a)) =
        '(TransformName nc name, a) ': '[]

type family TransformName (nc :: NameCase) (name :: Symbol) :: Symbol where
    TransformName 'SnakeCase s = CamelToSnake s
    TransformName 'IdentityCase s = s

-- | Type-level camelCase -> snake_case. Matches 'camelToSnake' at the value level.
type family CamelToSnake (s :: Symbol) :: Symbol where
    CamelToSnake s = SnakeStart (UnconsSymbol s)

type family SnakeStart (mu :: Maybe (Char, Symbol)) :: Symbol where
    SnakeStart 'Nothing = ""
    SnakeStart ('Just '(c, r)) =
        ConsSymbol (ToLowerChar c) (SnakeRest (UnconsSymbol r))

type family SnakeRest (mu :: Maybe (Char, Symbol)) :: Symbol where
    SnakeRest 'Nothing = ""
    SnakeRest ('Just '(c, r)) =
        SnakeStep (IsUpperChar c) c (SnakeRest (UnconsSymbol r))

type family SnakeStep (up :: Bool) (c :: Char) (rest :: Symbol) :: Symbol where
    SnakeStep 'True c rest = ConsSymbol '_' (ConsSymbol (ToLowerChar c) rest)
    SnakeStep 'False c rest = ConsSymbol c rest

type family IsUpperChar (c :: Char) :: Bool where
    IsUpperChar c =
        (CharToNat 'A' <=? CharToNat c) && (CharToNat c <=? CharToNat 'Z')

type family ToLowerChar (c :: Char) :: Char where
    ToLowerChar c = If (IsUpperChar c) (NatToChar (CharToNat c + 32)) c

-- | Snake_case schema derived from @a@'s 'Generic' representation.
type SchemaOf a = RepToSchema 'SnakeCase (Rep a)

-- | Identity-cased schema derived from @a@'s 'Generic' representation.
type SchemaOfRaw a = RepToSchema 'IdentityCase (Rep a)

{- | Walks the 'Rep' tree of a record, producing or consuming a list of
named columns. Used by 'genericToColumns' \/ 'genericFromColumns'.
-}
class GHasColumns (r :: Type -> Type) where
    gToColumns :: [r p] -> [(T.Text, C.Column)]
    gFromColumns :: D.DataFrame -> Either T.Text [r p]

instance (GHasColumns f) => GHasColumns (M1 D meta f) where
    gToColumns :: forall p. [M1 D meta f p] -> [(Text, Column)]
gToColumns [M1 D meta f p]
rs = [f p] -> [(Text, Column)]
forall p. [f p] -> [(Text, Column)]
forall (r :: * -> *) p. GHasColumns r => [r p] -> [(Text, Column)]
gToColumns ((M1 D meta f p -> f p) -> [M1 D meta f p] -> [f p]
forall a b. (a -> b) -> [a] -> [b]
map M1 D meta f p -> f p
forall k i (c :: Meta) (f :: k -> *) (p :: k). M1 i c f p -> f p
unM1 [M1 D meta f p]
rs)
    gFromColumns :: forall p. DataFrame -> Either Text [M1 D meta f p]
gFromColumns DataFrame
df = ([f p] -> [M1 D meta f p])
-> Either Text [f p] -> Either Text [M1 D meta f p]
forall a b. (a -> b) -> Either Text a -> Either Text b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((f p -> M1 D meta f p) -> [f p] -> [M1 D meta f p]
forall a b. (a -> b) -> [a] -> [b]
map f p -> M1 D meta f p
forall k i (c :: Meta) (f :: k -> *) (p :: k). f p -> M1 i c f p
M1) (DataFrame -> Either Text [f p]
forall p. DataFrame -> Either Text [f p]
forall (r :: * -> *) p.
GHasColumns r =>
DataFrame -> Either Text [r p]
gFromColumns DataFrame
df)

instance (GHasColumns f) => GHasColumns (M1 C meta f) where
    gToColumns :: forall p. [M1 C meta f p] -> [(Text, Column)]
gToColumns [M1 C meta f p]
rs = [f p] -> [(Text, Column)]
forall p. [f p] -> [(Text, Column)]
forall (r :: * -> *) p. GHasColumns r => [r p] -> [(Text, Column)]
gToColumns ((M1 C meta f p -> f p) -> [M1 C meta f p] -> [f p]
forall a b. (a -> b) -> [a] -> [b]
map M1 C meta f p -> f p
forall k i (c :: Meta) (f :: k -> *) (p :: k). M1 i c f p -> f p
unM1 [M1 C meta f p]
rs)
    gFromColumns :: forall p. DataFrame -> Either Text [M1 C meta f p]
gFromColumns DataFrame
df = ([f p] -> [M1 C meta f p])
-> Either Text [f p] -> Either Text [M1 C meta f p]
forall a b. (a -> b) -> Either Text a -> Either Text b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((f p -> M1 C meta f p) -> [f p] -> [M1 C meta f p]
forall a b. (a -> b) -> [a] -> [b]
map f p -> M1 C meta f p
forall k i (c :: Meta) (f :: k -> *) (p :: k). f p -> M1 i c f p
M1) (DataFrame -> Either Text [f p]
forall p. DataFrame -> Either Text [f p]
forall (r :: * -> *) p.
GHasColumns r =>
DataFrame -> Either Text [r p]
gFromColumns DataFrame
df)

instance (GHasColumns a, GHasColumns b) => GHasColumns (a :*: b) where
    gToColumns :: forall p. [(:*:) a b p] -> [(Text, Column)]
gToColumns [(:*:) a b p]
rs =
        [a p] -> [(Text, Column)]
forall p. [a p] -> [(Text, Column)]
forall (r :: * -> *) p. GHasColumns r => [r p] -> [(Text, Column)]
gToColumns (((:*:) a b p -> a p) -> [(:*:) a b p] -> [a p]
forall a b. (a -> b) -> [a] -> [b]
map (\(a p
x :*: b p
_) -> a p
x) [(:*:) a b p]
rs)
            [(Text, Column)] -> [(Text, Column)] -> [(Text, Column)]
forall a. [a] -> [a] -> [a]
++ [b p] -> [(Text, Column)]
forall p. [b p] -> [(Text, Column)]
forall (r :: * -> *) p. GHasColumns r => [r p] -> [(Text, Column)]
gToColumns (((:*:) a b p -> b p) -> [(:*:) a b p] -> [b p]
forall a b. (a -> b) -> [a] -> [b]
map (\(a p
_ :*: b p
y) -> b p
y) [(:*:) a b p]
rs)
    gFromColumns :: forall p. DataFrame -> Either Text [(:*:) a b p]
gFromColumns DataFrame
df = do
        [a p]
as <- DataFrame -> Either Text [a p]
forall p. DataFrame -> Either Text [a p]
forall (r :: * -> *) p.
GHasColumns r =>
DataFrame -> Either Text [r p]
gFromColumns DataFrame
df
        [b p]
bs <- DataFrame -> Either Text [b p]
forall p. DataFrame -> Either Text [b p]
forall (r :: * -> *) p.
GHasColumns r =>
DataFrame -> Either Text [r p]
gFromColumns DataFrame
df
        [(:*:) a b p] -> Either Text [(:*:) a b p]
forall a. a -> Either Text a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((a p -> b p -> (:*:) a b p) -> [a p] -> [b p] -> [(:*:) a b p]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith a p -> b p -> (:*:) a b p
forall k (f :: k -> *) (g :: k -> *) (p :: k).
f p -> g p -> (:*:) f g p
(:*:) [a p]
as [b p]
bs)

instance
    (KnownSymbol name, C.Columnable a) =>
    GHasColumns
        ( M1
            S
            ('MetaSel ('Just name) su ss ds)
            (K1 i a)
        )
    where
    gToColumns :: forall p.
[M1 S ('MetaSel ('Just name) su ss ds) (K1 i a) p]
-> [(Text, Column)]
gToColumns [M1 S ('MetaSel ('Just name) su ss ds) (K1 i a) p]
rs =
        let colName :: Text
colName = String -> Text
T.pack (String -> String
camelToSnake (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)))
            vals :: [a]
vals = (M1 S ('MetaSel ('Just name) su ss ds) (K1 i a) p -> a)
-> [M1 S ('MetaSel ('Just name) su ss ds) (K1 i a) p] -> [a]
forall a b. (a -> b) -> [a] -> [b]
map (K1 i a p -> a
forall k i c (p :: k). K1 i c p -> c
unK1 (K1 i a p -> a)
-> (M1 S ('MetaSel ('Just name) su ss ds) (K1 i a) p -> K1 i a p)
-> M1 S ('MetaSel ('Just name) su ss ds) (K1 i a) p
-> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. M1 S ('MetaSel ('Just name) su ss ds) (K1 i a) p -> K1 i a p
forall k i (c :: Meta) (f :: k -> *) (p :: k). M1 i c f p -> f p
unM1) [M1 S ('MetaSel ('Just name) su ss ds) (K1 i a) p]
rs
         in [(Text
colName, [a] -> Column
forall a.
(Columnable a, ColumnifyRep (KindOf a) a) =>
[a] -> Column
C.fromList [a]
vals)]
    gFromColumns :: forall p.
DataFrame
-> Either Text [M1 S ('MetaSel ('Just name) su ss ds) (K1 i a) p]
gFromColumns DataFrame
df = do
        let colName :: Text
colName = String -> Text
T.pack (String -> String
camelToSnake (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)))
        Vector a
v <- forall a.
Columnable a =>
Text -> DataFrame -> Either Text (Vector a)
requireColumn @a Text
colName DataFrame
df
        [M1 S ('MetaSel ('Just name) su ss ds) (K1 i a) p]
-> Either Text [M1 S ('MetaSel ('Just name) su ss ds) (K1 i a) p]
forall a. a -> Either Text a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((a -> M1 S ('MetaSel ('Just name) su ss ds) (K1 i a) p)
-> [a] -> [M1 S ('MetaSel ('Just name) su ss ds) (K1 i a) p]
forall a b. (a -> b) -> [a] -> [b]
map (K1 i a p -> M1 S ('MetaSel ('Just name) su ss ds) (K1 i a) p
forall k i (c :: Meta) (f :: k -> *) (p :: k). f p -> M1 i c f p
M1 (K1 i a p -> M1 S ('MetaSel ('Just name) su ss ds) (K1 i a) p)
-> (a -> K1 i a p)
-> a
-> M1 S ('MetaSel ('Just name) su ss ds) (K1 i a) p
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> K1 i a p
forall k i c (p :: k). c -> K1 i c p
K1) (Vector a -> [a]
forall a. Vector a -> [a]
VB.toList Vector a
v))

{- | Default implementation of 'DataFrame.Typed.Record.toColumns' for any
@Generic@ record. Field names are translated with @camelCase -> snake_case@.

@
instance HasSchema Order (SchemaOf Order) where
  toColumns   = genericToColumns
  fromColumns = genericFromColumns
@
-}
genericToColumns ::
    forall a. (Generic a, GHasColumns (Rep a)) => [a] -> [(T.Text, C.Column)]
genericToColumns :: forall a.
(Generic a, GHasColumns (Rep a)) =>
[a] -> [(Text, Column)]
genericToColumns = [Rep a Any] -> [(Text, Column)]
forall p. [Rep a p] -> [(Text, Column)]
forall (r :: * -> *) p. GHasColumns r => [r p] -> [(Text, Column)]
gToColumns ([Rep a Any] -> [(Text, Column)])
-> ([a] -> [Rep a Any]) -> [a] -> [(Text, Column)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> Rep a Any) -> [a] -> [Rep a Any]
forall a b. (a -> b) -> [a] -> [b]
map a -> Rep a Any
forall x. a -> Rep a x
forall a x. Generic a => a -> Rep a x
from

{- | Default implementation of 'DataFrame.Typed.Record.fromColumns' for any
@Generic@ record.
-}
genericFromColumns ::
    forall a. (Generic a, GHasColumns (Rep a)) => D.DataFrame -> Either T.Text [a]
genericFromColumns :: forall a.
(Generic a, GHasColumns (Rep a)) =>
DataFrame -> Either Text [a]
genericFromColumns DataFrame
df = ([Rep a Any] -> [a]) -> Either Text [Rep a Any] -> Either Text [a]
forall a b. (a -> b) -> Either Text a -> Either Text b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((Rep a Any -> a) -> [Rep a Any] -> [a]
forall a b. (a -> b) -> [a] -> [b]
map Rep a Any -> a
forall a x. Generic a => Rep a x -> a
forall x. Rep a x -> a
to) (DataFrame -> Either Text [Rep a Any]
forall p. DataFrame -> Either Text [Rep a p]
forall (r :: * -> *) p.
GHasColumns r =>
DataFrame -> Either Text [r p]
gFromColumns DataFrame
df)