{-# LANGUAGE ExplicitNamespaces #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StrictData #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
module DataFrame.Internal.Expression where

import qualified Data.Map as M
import Data.Type.Equality (type (:~:)(Refl), TestEquality (testEquality))
import Data.Data (Typeable)
import DataFrame.Internal.Column
import DataFrame.Internal.DataFrame
import DataFrame.Internal.Types
import qualified Data.Text as T
import qualified Data.Vector as V
import qualified Data.Vector.Unboxed as VU
import Type.Reflection (typeRep)
import DataFrame.Errors (DataFrameException(ColumnNotFoundException))
import Control.Exception (throw)
import Data.Maybe (fromMaybe)

data Expr a where
    Col :: Columnable a => T.Text -> Expr a 
    Lit :: Columnable a => a -> Expr a
    Apply :: (Columnable a, Columnable b) => T.Text -> (b -> a) -> Expr b -> Expr a
    BinOp :: (Columnable c, Columnable b, Columnable a) => T.Text -> (c -> b -> a) -> Expr c -> Expr b -> Expr a

interpret :: forall a b . (Columnable a) => DataFrame -> Expr a -> TypedColumn a
interpret :: forall a b. Columnable a => DataFrame -> Expr a -> TypedColumn a
interpret DataFrame
df (Lit a
value) = Column -> TypedColumn a
forall a. Columnable a => Column -> TypedColumn a
TColumn (Column -> TypedColumn a) -> Column -> TypedColumn a
forall a b. (a -> b) -> a -> b
$ Vector a -> Column
forall a.
(Columnable a, ColumnifyRep (KindOf a) a) =>
Vector a -> Column
toColumn' (Vector a -> Column) -> Vector a -> Column
forall a b. (a -> b) -> a -> b
$ Int -> a -> Vector a
forall a. Int -> a -> Vector a
V.replicate ((Int, Int) -> Int
forall a b. (a, b) -> a
fst ((Int, Int) -> Int) -> (Int, Int) -> Int
forall a b. (a -> b) -> a -> b
$ DataFrame -> (Int, Int)
dataframeDimensions DataFrame
df) a
value
interpret DataFrame
df (Col Text
name) = case Text -> DataFrame -> Maybe Column
getColumn Text
name DataFrame
df of
    Maybe Column
Nothing -> DataFrameException -> TypedColumn a
forall a e. Exception e => e -> a
throw (DataFrameException -> TypedColumn a)
-> DataFrameException -> TypedColumn a
forall a b. (a -> b) -> a -> b
$ Text -> Text -> [Text] -> DataFrameException
ColumnNotFoundException Text
name Text
"" (((Text, Int) -> Text) -> [(Text, Int)] -> [Text]
forall a b. (a -> b) -> [a] -> [b]
map (Text, Int) -> Text
forall a b. (a, b) -> a
fst ([(Text, Int)] -> [Text]) -> [(Text, Int)] -> [Text]
forall a b. (a -> b) -> a -> b
$ Map Text Int -> [(Text, Int)]
forall k a. Map k a -> [(k, a)]
M.toList (Map Text Int -> [(Text, Int)]) -> Map Text Int -> [(Text, Int)]
forall a b. (a -> b) -> a -> b
$ DataFrame -> Map Text Int
columnIndices DataFrame
df)
    Just Column
col -> Column -> TypedColumn a
forall a. Columnable a => Column -> TypedColumn a
TColumn Column
col
interpret DataFrame
df (Apply Text
_ (b -> a
f :: c -> d) Expr b
value) = let
        (TColumn Column
value') = forall a b. Columnable a => DataFrame -> Expr a -> TypedColumn a
interpret @c DataFrame
df Expr b
value
    in Column -> TypedColumn a
forall a. Columnable a => Column -> TypedColumn a
TColumn (Column -> TypedColumn a) -> Column -> TypedColumn a
forall a b. (a -> b) -> a -> b
$ Column -> Maybe Column -> Column
forall a. a -> Maybe a -> a
fromMaybe ([Char] -> Column
forall a. HasCallStack => [Char] -> a
error [Char]
"transform returned nothing") ((b -> a) -> Column -> Maybe Column
forall b c.
(Columnable b, Columnable c, UnboxIf c, Typeable b, Typeable c) =>
(b -> c) -> Column -> Maybe Column
transform b -> a
f Column
value')
interpret DataFrame
df (BinOp Text
_ (c -> b -> a
f :: c -> d -> e) Expr c
left Expr b
right) = let
        (TColumn Column
left') = forall a b. Columnable a => DataFrame -> Expr a -> TypedColumn a
interpret @c DataFrame
df Expr c
left
        (TColumn Column
right') = forall a b. Columnable a => DataFrame -> Expr a -> TypedColumn a
interpret @d DataFrame
df Expr b
right
    in Column -> TypedColumn a
forall a. Columnable a => Column -> TypedColumn a
TColumn (Column -> TypedColumn a) -> Column -> TypedColumn a
forall a b. (a -> b) -> a -> b
$ (c -> b -> a) -> Column -> Column -> Column
forall a b c.
(Columnable a, Columnable b, Columnable c) =>
(a -> b -> c) -> Column -> Column -> Column
zipWithColumns c -> b -> a
f Column
left' Column
right'

instance (Num a, Columnable a) => Num (Expr a) where
    (+) :: Expr a -> Expr a -> Expr a
    + :: Expr a -> Expr a -> Expr a
(+) = Text -> (a -> a -> a) -> Expr a -> Expr a -> Expr a
forall b b a.
(Columnable b, Columnable b, Columnable a) =>
Text -> (b -> b -> a) -> Expr b -> Expr b -> Expr a
BinOp Text
"add" a -> a -> a
forall a. Num a => a -> a -> a
(+)

    (*) :: Expr a -> Expr a -> Expr a
    * :: Expr a -> Expr a -> Expr a
(*) = Text -> (a -> a -> a) -> Expr a -> Expr a -> Expr a
forall b b a.
(Columnable b, Columnable b, Columnable a) =>
Text -> (b -> b -> a) -> Expr b -> Expr b -> Expr a
BinOp Text
"mult" a -> a -> a
forall a. Num a => a -> a -> a
(*)

    fromInteger :: Integer -> Expr a
    fromInteger :: Integer -> Expr a
fromInteger = a -> Expr a
forall a. Columnable a => a -> Expr a
Lit (a -> Expr a) -> (Integer -> a) -> Integer -> Expr a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Integer -> a
forall a. Num a => Integer -> a
fromInteger 

    negate :: Expr a -> Expr a
    negate :: Expr a -> Expr a
negate = Text -> (a -> a) -> Expr a -> Expr a
forall a b.
(Columnable a, Columnable b) =>
Text -> (b -> a) -> Expr b -> Expr a
Apply Text
"negate" a -> a
forall a. Num a => a -> a
negate

    abs :: Num a => Expr a -> Expr a
    abs :: Num a => Expr a -> Expr a
abs = Text -> (a -> a) -> Expr a -> Expr a
forall a b.
(Columnable a, Columnable b) =>
Text -> (b -> a) -> Expr b -> Expr a
Apply Text
"abs" a -> a
forall a. Num a => a -> a
abs

    signum :: Num a => Expr a -> Expr a
    signum :: Num a => Expr a -> Expr a
signum = Text -> (a -> a) -> Expr a -> Expr a
forall a b.
(Columnable a, Columnable b) =>
Text -> (b -> a) -> Expr b -> Expr a
Apply Text
"signum" a -> a
forall a. Num a => a -> a
signum

instance (Fractional a, Columnable a) => Fractional (Expr a) where
    fromRational :: (Fractional a, Columnable a) => Rational -> Expr a
    fromRational :: (Fractional a, Columnable a) => Rational -> Expr a
fromRational = a -> Expr a
forall a. Columnable a => a -> Expr a
Lit (a -> Expr a) -> (Rational -> a) -> Rational -> Expr a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Rational -> a
forall a. Fractional a => Rational -> a
fromRational

    (/) :: (Fractional a, Columnable a) => Expr a -> Expr a -> Expr a
    / :: (Fractional a, Columnable a) => Expr a -> Expr a -> Expr a
(/) = Text -> (a -> a -> a) -> Expr a -> Expr a -> Expr a
forall b b a.
(Columnable b, Columnable b, Columnable a) =>
Text -> (b -> b -> a) -> Expr b -> Expr b -> Expr a
BinOp Text
"divide" a -> a -> a
forall a. Fractional a => a -> a -> a
(/)

instance (Floating a, Columnable a) => Floating (Expr a) where
    pi :: (Floating a, Columnable a) => Expr a
    pi :: (Floating a, Columnable a) => Expr a
pi = a -> Expr a
forall a. Columnable a => a -> Expr a
Lit a
forall a. Floating a => a
pi
    exp :: (Floating a, Columnable a) => Expr a -> Expr a
    exp :: (Floating a, Columnable a) => Expr a -> Expr a
exp = Text -> (a -> a) -> Expr a -> Expr a
forall a b.
(Columnable a, Columnable b) =>
Text -> (b -> a) -> Expr b -> Expr a
Apply Text
"exp" a -> a
forall a. Floating a => a -> a
exp
    log :: (Floating a, Columnable a) => Expr a -> Expr a
    log :: (Floating a, Columnable a) => Expr a -> Expr a
log = Text -> (a -> a) -> Expr a -> Expr a
forall a b.
(Columnable a, Columnable b) =>
Text -> (b -> a) -> Expr b -> Expr a
Apply Text
"log" a -> a
forall a. Floating a => a -> a
log
    sin :: (Floating a, Columnable a) => Expr a -> Expr a
    sin :: (Floating a, Columnable a) => Expr a -> Expr a
sin = Text -> (a -> a) -> Expr a -> Expr a
forall a b.
(Columnable a, Columnable b) =>
Text -> (b -> a) -> Expr b -> Expr a
Apply Text
"sin" a -> a
forall a. Floating a => a -> a
sin
    cos :: (Floating a, Columnable a) => Expr a -> Expr a
    cos :: (Floating a, Columnable a) => Expr a -> Expr a
cos = Text -> (a -> a) -> Expr a -> Expr a
forall a b.
(Columnable a, Columnable b) =>
Text -> (b -> a) -> Expr b -> Expr a
Apply Text
"cos" a -> a
forall a. Floating a => a -> a
cos
    asin :: (Floating a, Columnable a) => Expr a -> Expr a
    asin :: (Floating a, Columnable a) => Expr a -> Expr a
asin = Text -> (a -> a) -> Expr a -> Expr a
forall a b.
(Columnable a, Columnable b) =>
Text -> (b -> a) -> Expr b -> Expr a
Apply Text
"asin" a -> a
forall a. Floating a => a -> a
asin
    acos :: (Floating a, Columnable a) => Expr a -> Expr a
    acos :: (Floating a, Columnable a) => Expr a -> Expr a
acos = Text -> (a -> a) -> Expr a -> Expr a
forall a b.
(Columnable a, Columnable b) =>
Text -> (b -> a) -> Expr b -> Expr a
Apply Text
"acos" a -> a
forall a. Floating a => a -> a
acos 
    atan :: (Floating a, Columnable a) => Expr a -> Expr a
    atan :: (Floating a, Columnable a) => Expr a -> Expr a
atan = Text -> (a -> a) -> Expr a -> Expr a
forall a b.
(Columnable a, Columnable b) =>
Text -> (b -> a) -> Expr b -> Expr a
Apply Text
"atan" a -> a
forall a. Floating a => a -> a
atan
    sinh :: (Floating a, Columnable a) => Expr a -> Expr a
    sinh :: (Floating a, Columnable a) => Expr a -> Expr a
sinh = Text -> (a -> a) -> Expr a -> Expr a
forall a b.
(Columnable a, Columnable b) =>
Text -> (b -> a) -> Expr b -> Expr a
Apply Text
"sinh" a -> a
forall a. Floating a => a -> a
sinh
    cosh :: (Floating a, Columnable a) => Expr a -> Expr a
    cosh :: (Floating a, Columnable a) => Expr a -> Expr a
cosh = Text -> (a -> a) -> Expr a -> Expr a
forall a b.
(Columnable a, Columnable b) =>
Text -> (b -> a) -> Expr b -> Expr a
Apply Text
"cosh" a -> a
forall a. Floating a => a -> a
cosh
    asinh :: (Floating a, Columnable a) => Expr a -> Expr a
    asinh :: (Floating a, Columnable a) => Expr a -> Expr a
asinh = Text -> (a -> a) -> Expr a -> Expr a
forall a b.
(Columnable a, Columnable b) =>
Text -> (b -> a) -> Expr b -> Expr a
Apply Text
"asinh" a -> a
forall a. Floating a => a -> a
sinh
    acosh :: (Floating a, Columnable a) => Expr a -> Expr a
    acosh :: (Floating a, Columnable a) => Expr a -> Expr a
acosh = Text -> (a -> a) -> Expr a -> Expr a
forall a b.
(Columnable a, Columnable b) =>
Text -> (b -> a) -> Expr b -> Expr a
Apply Text
"acosh" a -> a
forall a. Floating a => a -> a
acosh
    atanh :: (Floating a, Columnable a) => Expr a -> Expr a
    atanh :: (Floating a, Columnable a) => Expr a -> Expr a
atanh = Text -> (a -> a) -> Expr a -> Expr a
forall a b.
(Columnable a, Columnable b) =>
Text -> (b -> a) -> Expr b -> Expr a
Apply Text
"atanh" a -> a
forall a. Floating a => a -> a
atanh


instance (Show a) => Show (Expr a) where
    show :: forall a . Show a => Expr a -> String
    show :: forall a. Show a => Expr a -> [Char]
show (Col Text
name) = [Char]
"col@" [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ TypeRep a -> [Char]
forall a. Show a => a -> [Char]
show (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @a) [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ [Char]
"(" [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ Text -> [Char]
T.unpack Text
name [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ [Char]
")"
    show (Lit a
value) = a -> [Char]
forall a. Show a => a -> [Char]
show a
value
    show (Apply Text
name b -> a
f Expr b
value) = Text -> [Char]
T.unpack Text
name [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ [Char]
"(" [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ Expr b -> [Char]
forall a. Show a => a -> [Char]
show Expr b
value [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ [Char]
")"
    show (BinOp Text
name c -> b -> a
f Expr c
a Expr b
b) = Text -> [Char]
T.unpack Text
name [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ [Char]
"(" [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ Expr c -> [Char]
forall a. Show a => a -> [Char]
show Expr c
a [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ [Char]
", " [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ Expr b -> [Char]
forall a. Show a => a -> [Char]
show Expr b
b [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ [Char]
")" 

col :: Columnable a => T.Text -> Expr a
col :: forall a. Columnable a => Text -> Expr a
col = Text -> Expr a
forall a. Columnable a => Text -> Expr a
Col

lit :: Columnable a => a -> Expr a
lit :: forall a. Columnable a => a -> Expr a
lit = a -> Expr a
forall a. Columnable a => a -> Expr a
Lit

lift :: (Columnable a, Columnable b) => (a -> b) -> Expr a -> Expr b
lift :: forall a b.
(Columnable a, Columnable b) =>
(a -> b) -> Expr a -> Expr b
lift = Text -> (a -> b) -> Expr a -> Expr b
forall a b.
(Columnable a, Columnable b) =>
Text -> (b -> a) -> Expr b -> Expr a
Apply Text
"udf"

lift2 :: (Columnable c, Columnable b, Columnable a) => (c -> b -> a) -> Expr c -> Expr b -> Expr a 
lift2 :: forall c b a.
(Columnable c, Columnable b, Columnable a) =>
(c -> b -> a) -> Expr c -> Expr b -> Expr a
lift2 = Text -> (c -> b -> a) -> Expr c -> Expr b -> Expr a
forall b b a.
(Columnable b, Columnable b, Columnable a) =>
Text -> (b -> b -> a) -> Expr b -> Expr b -> Expr a
BinOp Text
"udf"

eq :: (Columnable a, Eq a) => Expr a -> Expr a -> Expr Bool
eq :: forall a. (Columnable a, Eq a) => Expr a -> Expr a -> Expr Bool
eq = Text -> (a -> a -> Bool) -> Expr a -> Expr a -> Expr Bool
forall b b a.
(Columnable b, Columnable b, Columnable a) =>
Text -> (b -> b -> a) -> Expr b -> Expr b -> Expr a
BinOp Text
"eq" a -> a -> Bool
forall a. Eq a => a -> a -> Bool
(==)

lt :: (Columnable a, Ord a) => Expr a -> Expr a -> Expr Bool
lt :: forall a. (Columnable a, Ord a) => Expr a -> Expr a -> Expr Bool
lt = Text -> (a -> a -> Bool) -> Expr a -> Expr a -> Expr Bool
forall b b a.
(Columnable b, Columnable b, Columnable a) =>
Text -> (b -> b -> a) -> Expr b -> Expr b -> Expr a
BinOp Text
"lt" a -> a -> Bool
forall a. Ord a => a -> a -> Bool
(<)

gt :: (Columnable a, Ord a) => Expr a -> Expr a -> Expr Bool
gt :: forall a. (Columnable a, Ord a) => Expr a -> Expr a -> Expr Bool
gt = Text -> (a -> a -> Bool) -> Expr a -> Expr a -> Expr Bool
forall b b a.
(Columnable b, Columnable b, Columnable a) =>
Text -> (b -> b -> a) -> Expr b -> Expr b -> Expr a
BinOp Text
"gt" a -> a -> Bool
forall a. Ord a => a -> a -> Bool
(<)

leq :: (Columnable a, Ord a, Eq a) => Expr a -> Expr a -> Expr Bool
leq :: forall a.
(Columnable a, Ord a, Eq a) =>
Expr a -> Expr a -> Expr Bool
leq = Text -> (a -> a -> Bool) -> Expr a -> Expr a -> Expr Bool
forall b b a.
(Columnable b, Columnable b, Columnable a) =>
Text -> (b -> b -> a) -> Expr b -> Expr b -> Expr a
BinOp Text
"leq" a -> a -> Bool
forall a. Ord a => a -> a -> Bool
(<=)

geq :: (Columnable a, Ord a, Eq a) => Expr a -> Expr a -> Expr Bool
geq :: forall a.
(Columnable a, Ord a, Eq a) =>
Expr a -> Expr a -> Expr Bool
geq = Text -> (a -> a -> Bool) -> Expr a -> Expr a -> Expr Bool
forall b b a.
(Columnable b, Columnable b, Columnable a) =>
Text -> (b -> b -> a) -> Expr b -> Expr b -> Expr a
BinOp Text
"geq" a -> a -> Bool
forall a. Ord a => a -> a -> Bool
(<=)