{-# LANGUAGE ExplicitNamespaces #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE TemplateHaskell #-}
module DataFrame.Functions where

import DataFrame.Internal.Column
import DataFrame.Internal.DataFrame (DataFrame(..), unsafeGetColumn)
import DataFrame.Internal.Expression (Expr(..), UExpr(..))

import           Control.Monad
import           Data.Function
import qualified Data.List as L
import qualified Data.Map as M
import qualified Data.Text as T
import qualified Data.Vector.Generic as VG
import qualified Data.Vector.Unboxed as VU
import qualified Data.Vector as VB
import           Language.Haskell.TH
import qualified Language.Haskell.TH.Syntax as TH

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

as :: Columnable a => Expr a -> T.Text -> (T.Text, UExpr)
as :: forall a. Columnable a => Expr a -> Text -> (Text, UExpr)
as Expr a
expr Text
name = (Text
name, Expr a -> UExpr
forall a. Columnable a => Expr a -> UExpr
Wrap Expr a
expr)

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 c b a.
(Columnable c, Columnable b, Columnable a) =>
Text -> (c -> b -> a) -> Expr c -> 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 c b a.
(Columnable c, Columnable b, Columnable a) =>
Text -> (c -> b -> a) -> Expr c -> 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 c b a.
(Columnable c, Columnable b, Columnable a) =>
Text -> (c -> b -> a) -> Expr c -> 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 c b a.
(Columnable c, Columnable b, Columnable a) =>
Text -> (c -> b -> a) -> Expr c -> 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 c b a.
(Columnable c, Columnable b, Columnable a) =>
Text -> (c -> b -> a) -> Expr c -> 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 c b a.
(Columnable c, Columnable b, Columnable a) =>
Text -> (c -> b -> a) -> Expr c -> Expr b -> Expr a
BinOp Text
"geq" a -> a -> Bool
forall a. Ord a => a -> a -> Bool
(>=)

count :: Columnable a => Expr a -> Expr Int
count :: forall a. Columnable a => Expr a -> Expr Int
count (Col Text
name) = Text
-> Text
-> (forall (v :: * -> *) b.
    (Vector v b, Columnable b) =>
    v b -> Int)
-> Expr Int
forall a.
Columnable a =>
Text
-> Text
-> (forall (v :: * -> *) b. (Vector v b, Columnable b) => v b -> a)
-> Expr a
GeneralAggregate Text
name Text
"count" v b -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
forall (v :: * -> *) b. (Vector v b, Columnable b) => v b -> Int
VG.length
count Expr a
_ = [Char] -> Expr Int
forall a. HasCallStack => [Char] -> a
error [Char]
"Argument can only be a column reference not an unevaluated expression"

anyValue :: Columnable a => Expr a -> Expr a
anyValue :: forall a. Columnable a => Expr a -> Expr a
anyValue (Col Text
name) = Text
-> Text
-> (forall (v :: * -> *) a1.
    (Vector v a1, Columnable a1) =>
    v a1 -> a1)
-> Expr a
forall a.
Columnable a =>
Text
-> Text
-> (forall (v :: * -> *) a1.
    (Vector v a1, Columnable a1) =>
    v a1 -> a1)
-> Expr a
ReductionAggregate Text
name Text
"anyValue" v a1 -> a1
forall (v :: * -> *) a. Vector v a => v a -> a
forall (v :: * -> *) a1. (Vector v a1, Columnable a1) => v a1 -> a1
VG.head

minimum :: Columnable a => Expr a -> Expr a
minimum :: forall a. Columnable a => Expr a -> Expr a
minimum (Col Text
name) = Text
-> Text
-> (forall (v :: * -> *) a1.
    (Vector v a1, Columnable a1) =>
    v a1 -> a1)
-> Expr a
forall a.
Columnable a =>
Text
-> Text
-> (forall (v :: * -> *) a1.
    (Vector v a1, Columnable a1) =>
    v a1 -> a1)
-> Expr a
ReductionAggregate Text
name Text
"minimum" v a1 -> a1
forall (v :: * -> *) a. (Vector v a, Ord a) => v a -> a
forall (v :: * -> *) a1. (Vector v a1, Columnable a1) => v a1 -> a1
VG.minimum

maximum :: Columnable a => Expr a -> Expr a
maximum :: forall a. Columnable a => Expr a -> Expr a
maximum (Col Text
name) = Text
-> Text
-> (forall (v :: * -> *) a1.
    (Vector v a1, Columnable a1) =>
    v a1 -> a1)
-> Expr a
forall a.
Columnable a =>
Text
-> Text
-> (forall (v :: * -> *) a1.
    (Vector v a1, Columnable a1) =>
    v a1 -> a1)
-> Expr a
ReductionAggregate Text
name Text
"maximum" v a1 -> a1
forall (v :: * -> *) a. (Vector v a, Ord a) => v a -> a
forall (v :: * -> *) a1. (Vector v a1, Columnable a1) => v a1 -> a1
VG.maximum

sum :: (Columnable a, Num a, VU.Unbox a) => Expr a -> Expr a
sum :: forall a. (Columnable a, Num a, Unbox a) => Expr a -> Expr a
sum (Col Text
name) = Text -> Text -> (Vector a -> a) -> Expr a
forall a b.
(Columnable a, Columnable b, Num a, Num b) =>
Text -> Text -> (Vector b -> a) -> Expr a
NumericAggregate Text
name Text
"sum" Vector a -> a
forall (v :: * -> *) a. (Vector v a, Num a) => v a -> a
VG.sum

mean :: (Columnable a, Num a, VU.Unbox a) => Expr a -> Expr Double
mean :: forall a. (Columnable a, Num a, Unbox a) => Expr a -> Expr Double
mean (Col Text
name) = let
        mean' :: v Double -> Double
mean' v Double
samp = let
                (!Double
total, !Int
n) = ((Double, Int) -> Double -> (Double, Int))
-> (Double, Int) -> v Double -> (Double, Int)
forall (v :: * -> *) b a.
Vector v b =>
(a -> b -> a) -> a -> v b -> a
VG.foldl' (\(!Double
total, !Int
n) Double
v -> (Double
total Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
v, Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1))  (Double
0 :: Double, Int
0 :: Int) v Double
samp
            in Double
total Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n
    in Text -> Text -> (Vector Double -> Double) -> Expr Double
forall a b.
(Columnable a, Columnable b, Num a, Num b) =>
Text -> Text -> (Vector b -> a) -> Expr a
NumericAggregate Text
name Text
"mean" Vector Double -> Double
forall {v :: * -> *}. Vector v Double => v Double -> Double
mean'

typeFromString :: String -> Q Type
typeFromString :: [Char] -> Q Type
typeFromString [Char]
s = do
  Maybe Name
maybeType <- [Char] -> Q (Maybe Name)
lookupTypeName [Char]
s
  case Maybe Name
maybeType of
    Just Name
name -> Type -> Q Type
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return (Name -> Type
ConT Name
name)
    Maybe Name
Nothing -> do
      if [Char]
"Maybe " [Char] -> [Char] -> Bool
forall a. Eq a => [a] -> [a] -> Bool
`L.isPrefixOf` [Char]
s
        then do
          let innerType :: [Char]
innerType = Int -> [Char] -> [Char]
forall a. Int -> [a] -> [a]
drop Int
6 [Char]
s
          Type
inner <- [Char] -> Q Type
typeFromString [Char]
innerType
          Type -> Q Type
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return (Type -> Type -> Type
AppT (Name -> Type
ConT ''Maybe) Type
inner)
        else if [Char]
"Either " [Char] -> [Char] -> Bool
forall a. Eq a => [a] -> [a] -> Bool
`L.isPrefixOf` [Char]
s
          then do
            let ([Char]
left: [Char]
right:[[Char]]
_) = [[Char]] -> [[Char]]
forall a. HasCallStack => [a] -> [a]
tail ([Char] -> [[Char]]
words [Char]
s)
            Type
lhs <- [Char] -> Q Type
typeFromString [Char]
left
            Type
rhs <- [Char] -> Q Type
typeFromString [Char]
right
            Type -> Q Type
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return (Type -> Type -> Type
AppT (Type -> Type -> Type
AppT (Name -> Type
ConT ''Either) Type
lhs) Type
rhs)
          else [Char] -> Q Type
forall a. [Char] -> Q a
forall (m :: * -> *) a. MonadFail m => [Char] -> m a
fail ([Char] -> Q Type) -> [Char] -> Q Type
forall a b. (a -> b) -> a -> b
$ [Char]
"Unsupported type: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
s

declareColumns :: DataFrame -> DecsQ
declareColumns :: DataFrame -> DecsQ
declareColumns DataFrame
df = let
        names :: [Text]
names = (((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])
-> (DataFrame -> [(Text, Int)]) -> DataFrame -> [Text]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Text, Int) -> (Text, Int) -> Ordering)
-> [(Text, Int)] -> [(Text, Int)]
forall a. (a -> a -> Ordering) -> [a] -> [a]
L.sortBy (Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
compare (Int -> Int -> Ordering)
-> ((Text, Int) -> Int) -> (Text, Int) -> (Text, Int) -> Ordering
forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` (Text, Int) -> Int
forall a b. (a, b) -> b
snd)([(Text, Int)] -> [(Text, Int)])
-> (DataFrame -> [(Text, Int)]) -> DataFrame -> [(Text, Int)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map Text Int -> [(Text, Int)]
forall k a. Map k a -> [(k, a)]
M.toList (Map Text Int -> [(Text, Int)])
-> (DataFrame -> Map Text Int) -> DataFrame -> [(Text, Int)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. DataFrame -> Map Text Int
columnIndices) DataFrame
df
        types :: [[Char]]
types = (Text -> [Char]) -> [Text] -> [[Char]]
forall a b. (a -> b) -> [a] -> [b]
map (Column -> [Char]
columnTypeString (Column -> [Char]) -> (Text -> Column) -> Text -> [Char]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Text -> DataFrame -> Column
`unsafeGetColumn` DataFrame
df)) [Text]
names
        specs :: [(Text, [Char])]
specs = [Text] -> [[Char]] -> [(Text, [Char])]
forall a b. [a] -> [b] -> [(a, b)]
zip [Text]
names [[Char]]
types
    in ([[Dec]] -> [Dec]) -> Q [[Dec]] -> DecsQ
forall a b. (a -> b) -> Q a -> Q b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [[Dec]] -> [Dec]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat (Q [[Dec]] -> DecsQ) -> Q [[Dec]] -> DecsQ
forall a b. (a -> b) -> a -> b
$ [(Text, [Char])] -> ((Text, [Char]) -> DecsQ) -> Q [[Dec]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [(Text, [Char])]
specs (((Text, [Char]) -> DecsQ) -> Q [[Dec]])
-> ((Text, [Char]) -> DecsQ) -> Q [[Dec]]
forall a b. (a -> b) -> a -> b
$ \(Text
nm, [Char]
tyStr) -> do
        Type
ty  <- [Char] -> Q Type
typeFromString [Char]
tyStr
        let n :: Name
n  = [Char] -> Name
mkName (Text -> [Char]
T.unpack Text
nm)
        Dec
sig <- Name -> Q Type -> Q Dec
forall (m :: * -> *). Quote m => Name -> m Type -> m Dec
sigD Name
n [t| Expr $(Type -> Q Type
forall a. a -> Q a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Type
ty) |]
        Dec
val <- Q Pat -> Q Body -> [Q Dec] -> Q Dec
forall (m :: * -> *).
Quote m =>
m Pat -> m Body -> [m Dec] -> m Dec
valD (Name -> Q Pat
forall (m :: * -> *). Quote m => Name -> m Pat
varP Name
n) (Q Exp -> Q Body
forall (m :: * -> *). Quote m => m Exp -> m Body
normalB [| col $(Text -> Q Exp
forall t (m :: * -> *). (Lift t, Quote m) => t -> m Exp
forall (m :: * -> *). Quote m => Text -> m Exp
TH.lift Text
nm) |]) []
        [Dec] -> DecsQ
forall a. a -> Q a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [Dec
sig, Dec
val]