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

{- |
Runtime schema representation. The Template-Haskell @deriveSchema@ splice
lives in "DataFrame.Internal.Schema.TH" so this module can be used from
packages that do not depend on @template-haskell@.
-}
module DataFrame.Internal.Schema (
    SchemaType (..),
    schemaType,
    Schema (..),
    makeSchema,
    RuntimeSchema (..),
) where

import Data.Kind (Type)
import qualified Data.Map as M
import Data.Maybe (isJust)
import qualified Data.Proxy as P
import qualified Data.Text as T
import Data.Type.Equality (TestEquality (..))
import DataFrame.Internal.Column (Columnable)
import GHC.TypeLits (KnownSymbol, Symbol, symbolVal)
import Type.Reflection (typeRep)

-- | A runtime tag for a column’s element type.
data SchemaType where
    -- | Constructor carrying a 'Proxy' of the element type.
    SType :: (Columnable a, Read a) => P.Proxy a -> SchemaType

{- | Show the underlying element type using 'typeRep'.

==== __Examples__
>>> :set -XTypeApplications
>>> show (schemaType @Bool)
"Bool"
-}
instance Show SchemaType where
    show :: SchemaType -> String
    show :: SchemaType -> String
show (SType (Proxy a
_ :: P.Proxy a)) = TypeRep a -> String
forall a. Show a => a -> String
show (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @a)

{- | Two 'SchemaType's are equal iff their element types are the same.

==== __Examples__
>>> :set -XTypeApplications
>>> schemaType @Int == schemaType @Int
True

>>> schemaType @Int == schemaType @Integer
False
-}
instance Eq SchemaType where
    (==) :: SchemaType -> SchemaType -> Bool
    == :: SchemaType -> SchemaType -> Bool
(==) (SType (Proxy a
_ :: P.Proxy a)) (SType (Proxy a
_ :: P.Proxy b)) =
        Maybe (a :~: a) -> Bool
forall a. Maybe a -> Bool
isJust (TypeRep a -> TypeRep a -> Maybe (a :~: a)
forall a b. TypeRep a -> TypeRep b -> Maybe (a :~: b)
forall {k} (f :: k -> *) (a :: k) (b :: k).
TestEquality f =>
f a -> f b -> Maybe (a :~: b)
testEquality (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @a) (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @b))

{- | Construct a 'SchemaType' for the given @a@.

==== __Examples__
>>> :set -XTypeApplications
>>> schemaType @T.Text == schemaType @T.Text
True

>>> show (schemaType @Double)
"Double"
-}
schemaType :: forall a. (Columnable a, Read a) => SchemaType
schemaType :: forall a. (Columnable a, Read a) => SchemaType
schemaType = Proxy a -> SchemaType
forall a. (Columnable a, Read a) => Proxy a -> SchemaType
SType (forall t. Proxy t
forall {k} (t :: k). Proxy t
P.Proxy @a)

{- | Logical schema of a 'DataFrame': a mapping from column names to their
element types ('SchemaType').
-}
newtype Schema = Schema
    { Schema -> Map Text SchemaType
elements :: M.Map T.Text SchemaType
    {- ^ Mapping from /column name/ to its 'SchemaType'.

    Invariant: keys are unique column names. A missing key means the column
    is not present in the schema.
    -}
    }
    deriving (Int -> Schema -> ShowS
[Schema] -> ShowS
Schema -> String
(Int -> Schema -> ShowS)
-> (Schema -> String) -> ([Schema] -> ShowS) -> Show Schema
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Schema -> ShowS
showsPrec :: Int -> Schema -> ShowS
$cshow :: Schema -> String
show :: Schema -> String
$cshowList :: [Schema] -> ShowS
showList :: [Schema] -> ShowS
Show, Schema -> Schema -> Bool
(Schema -> Schema -> Bool)
-> (Schema -> Schema -> Bool) -> Eq Schema
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Schema -> Schema -> Bool
== :: Schema -> Schema -> Bool
$c/= :: Schema -> Schema -> Bool
/= :: Schema -> Schema -> Bool
Eq)

-- | Construct a 'Schema' from a list of @(columnName, schemaType)@ pairs.
makeSchema :: [(T.Text, SchemaType)] -> Schema
makeSchema :: [(Text, SchemaType)] -> Schema
makeSchema = Map Text SchemaType -> Schema
Schema (Map Text SchemaType -> Schema)
-> ([(Text, SchemaType)] -> Map Text SchemaType)
-> [(Text, SchemaType)]
-> Schema
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(Text, SchemaType)] -> Map Text SchemaType
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList

{- | The runtime 'Schema' behind a type-level schema — names /and/ element
types — so a reader can project to a schema's columns and skip inference for
them in one step.

Every column type must have a 'Read' instance, which 'Columnable' does not
imply; that is what lets the names carry their types across to a reader.

==== __Examples__
>>> :set -XTypeApplications -XDataKinds
>>> elements (runtimeSchema @'[ '("n", Int)])
fromList [("n",Int)]
-}
class RuntimeSchema (cols :: [(Symbol, Type)]) where
    runtimeSchema :: Schema

instance RuntimeSchema '[] where
    runtimeSchema :: Schema
runtimeSchema = [(Text, SchemaType)] -> Schema
makeSchema []

instance
    (KnownSymbol name, Columnable a, Read a, RuntimeSchema rest) =>
    RuntimeSchema ('(name, a) ': rest)
    where
    runtimeSchema :: Schema
runtimeSchema =
        Map Text SchemaType -> Schema
Schema (Map Text SchemaType -> Schema) -> Map Text SchemaType -> Schema
forall a b. (a -> b) -> a -> b
$
            Text -> SchemaType -> Map Text SchemaType -> Map Text SchemaType
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert
                (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
P.Proxy @name)))
                (forall a. (Columnable a, Read a) => SchemaType
schemaType @a)
                (Schema -> Map Text SchemaType
elements (forall (cols :: [(Symbol, *)]). RuntimeSchema cols => Schema
runtimeSchema @rest))