{-# LANGUAGE ExplicitNamespaces #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TupleSections #-}

module DataFrame.Monad where

import DataFrame (DataFrame)
import qualified DataFrame as D
import DataFrame.Internal.Column (Columnable)
import DataFrame.Internal.Expression (Expr (..))

import qualified Data.Text as T
import System.Random

-- A re-implementation of the state monad.
-- `mtl` might be too heavy a dependency just to get
-- a single monad instance.
newtype FrameM a = FrameM {forall a. FrameM a -> DataFrame -> (DataFrame, a)
runFrameM_ :: DataFrame -> (DataFrame, a)}

instance Functor FrameM where
    fmap :: (a -> b) -> FrameM a -> FrameM b
    fmap :: forall a b. (a -> b) -> FrameM a -> FrameM b
fmap a -> b
f (FrameM DataFrame -> (DataFrame, a)
g) = (DataFrame -> (DataFrame, b)) -> FrameM b
forall a. (DataFrame -> (DataFrame, a)) -> FrameM a
FrameM ((DataFrame -> (DataFrame, b)) -> FrameM b)
-> (DataFrame -> (DataFrame, b)) -> FrameM b
forall a b. (a -> b) -> a -> b
$ \DataFrame
df ->
        let (DataFrame
df', a
x) = DataFrame -> (DataFrame, a)
g DataFrame
df
         in (DataFrame
df', a -> b
f a
x)

instance Applicative FrameM where
    pure :: forall a. a -> FrameM a
pure a
x = (DataFrame -> (DataFrame, a)) -> FrameM a
forall a. (DataFrame -> (DataFrame, a)) -> FrameM a
FrameM (,a
x)
    (<*>) :: FrameM (a -> b) -> FrameM a -> FrameM b
    FrameM DataFrame -> (DataFrame, a -> b)
ff <*> :: forall a b. FrameM (a -> b) -> FrameM a -> FrameM b
<*> FrameM DataFrame -> (DataFrame, a)
fx = (DataFrame -> (DataFrame, b)) -> FrameM b
forall a. (DataFrame -> (DataFrame, a)) -> FrameM a
FrameM ((DataFrame -> (DataFrame, b)) -> FrameM b)
-> (DataFrame -> (DataFrame, b)) -> FrameM b
forall a b. (a -> b) -> a -> b
$ \DataFrame
df ->
        let (DataFrame
df1, a -> b
f) = DataFrame -> (DataFrame, a -> b)
ff DataFrame
df
            (DataFrame
df2, a
x) = DataFrame -> (DataFrame, a)
fx DataFrame
df1
         in (DataFrame
df2, a -> b
f a
x)

instance Monad FrameM where
    (>>=) :: FrameM a -> (a -> FrameM b) -> FrameM b
    FrameM DataFrame -> (DataFrame, a)
g >>= :: forall a b. FrameM a -> (a -> FrameM b) -> FrameM b
>>= a -> FrameM b
f = (DataFrame -> (DataFrame, b)) -> FrameM b
forall a. (DataFrame -> (DataFrame, a)) -> FrameM a
FrameM ((DataFrame -> (DataFrame, b)) -> FrameM b)
-> (DataFrame -> (DataFrame, b)) -> FrameM b
forall a b. (a -> b) -> a -> b
$ \DataFrame
df ->
        let (DataFrame
df1, a
x) = DataFrame -> (DataFrame, a)
g DataFrame
df
            FrameM DataFrame -> (DataFrame, b)
h = a -> FrameM b
f a
x
         in DataFrame -> (DataFrame, b)
h DataFrame
df1

modifyM :: (DataFrame -> DataFrame) -> FrameM ()
modifyM :: (DataFrame -> DataFrame) -> FrameM ()
modifyM DataFrame -> DataFrame
f = (DataFrame -> (DataFrame, ())) -> FrameM ()
forall a. (DataFrame -> (DataFrame, a)) -> FrameM a
FrameM ((DataFrame -> (DataFrame, ())) -> FrameM ())
-> (DataFrame -> (DataFrame, ())) -> FrameM ()
forall a b. (a -> b) -> a -> b
$ \DataFrame
df -> (DataFrame -> DataFrame
f DataFrame
df, ())

inspectM :: (DataFrame -> b) -> FrameM b
inspectM :: forall b. (DataFrame -> b) -> FrameM b
inspectM DataFrame -> b
f = (DataFrame -> (DataFrame, b)) -> FrameM b
forall a. (DataFrame -> (DataFrame, a)) -> FrameM a
FrameM ((DataFrame -> (DataFrame, b)) -> FrameM b)
-> (DataFrame -> (DataFrame, b)) -> FrameM b
forall a b. (a -> b) -> a -> b
$ \DataFrame
df -> (DataFrame
df, DataFrame -> b
f DataFrame
df)

deriveM :: (Columnable a) => T.Text -> Expr a -> FrameM (Expr a)
deriveM :: forall a. Columnable a => Text -> Expr a -> FrameM (Expr a)
deriveM Text
name Expr a
expr = (DataFrame -> (DataFrame, Expr a)) -> FrameM (Expr a)
forall a. (DataFrame -> (DataFrame, a)) -> FrameM a
FrameM ((DataFrame -> (DataFrame, Expr a)) -> FrameM (Expr a))
-> (DataFrame -> (DataFrame, Expr a)) -> FrameM (Expr a)
forall a b. (a -> b) -> a -> b
$ \DataFrame
df ->
    let df' :: DataFrame
df' = Text -> Expr a -> DataFrame -> DataFrame
forall a. Columnable a => Text -> Expr a -> DataFrame -> DataFrame
D.derive Text
name Expr a
expr DataFrame
df
     in (DataFrame
df', Text -> Expr a
forall a. Columnable a => Text -> Expr a
Col Text
name)

renameM :: (Columnable a) => Expr a -> T.Text -> FrameM (Expr a)
renameM :: forall a. Columnable a => Expr a -> Text -> FrameM (Expr a)
renameM (Col Text
oldName) Text
newName = (DataFrame -> (DataFrame, Expr a)) -> FrameM (Expr a)
forall a. (DataFrame -> (DataFrame, a)) -> FrameM a
FrameM ((DataFrame -> (DataFrame, Expr a)) -> FrameM (Expr a))
-> (DataFrame -> (DataFrame, Expr a)) -> FrameM (Expr a)
forall a b. (a -> b) -> a -> b
$ \DataFrame
df ->
    let df' :: DataFrame
df' = Text -> Text -> DataFrame -> DataFrame
D.rename Text
oldName Text
newName DataFrame
df
     in (DataFrame
df', Text -> Expr a
forall a. Columnable a => Text -> Expr a
Col Text
newName)
renameM Expr a
expr Text
newName = Text -> Expr a -> FrameM (Expr a)
forall a. Columnable a => Text -> Expr a -> FrameM (Expr a)
deriveM Text
newName Expr a
expr

filterWhereM :: Expr Bool -> FrameM ()
filterWhereM :: Expr Bool -> FrameM ()
filterWhereM Expr Bool
p = (DataFrame -> DataFrame) -> FrameM ()
modifyM (Expr Bool -> DataFrame -> DataFrame
D.filterWhere Expr Bool
p)

sampleM :: (RandomGen g) => g -> Double -> FrameM ()
sampleM :: forall g. RandomGen g => g -> Double -> FrameM ()
sampleM g
pureGen Double
p = (DataFrame -> DataFrame) -> FrameM ()
modifyM (g -> Double -> DataFrame -> DataFrame
forall g. RandomGen g => g -> Double -> DataFrame -> DataFrame
D.sample g
pureGen Double
p)

takeM :: Int -> FrameM ()
takeM :: Int -> FrameM ()
takeM Int
n = (DataFrame -> DataFrame) -> FrameM ()
modifyM (Int -> DataFrame -> DataFrame
D.take Int
n)

filterJustM :: (Columnable a) => Expr (Maybe a) -> FrameM (Expr a)
filterJustM :: forall a. Columnable a => Expr (Maybe a) -> FrameM (Expr a)
filterJustM (Col Text
name) = (DataFrame -> (DataFrame, Expr a)) -> FrameM (Expr a)
forall a. (DataFrame -> (DataFrame, a)) -> FrameM a
FrameM ((DataFrame -> (DataFrame, Expr a)) -> FrameM (Expr a))
-> (DataFrame -> (DataFrame, Expr a)) -> FrameM (Expr a)
forall a b. (a -> b) -> a -> b
$ \DataFrame
df ->
    let df' :: DataFrame
df' = Text -> DataFrame -> DataFrame
D.filterJust Text
name DataFrame
df
     in (DataFrame
df', Text -> Expr a
forall a. Columnable a => Text -> Expr a
Col Text
name)
filterJustM Expr (Maybe a)
expr =
    [Char] -> FrameM (Expr a)
forall a. HasCallStack => [Char] -> a
error ([Char] -> FrameM (Expr a)) -> [Char] -> FrameM (Expr a)
forall a b. (a -> b) -> a -> b
$ [Char]
"Cannot filter on compound expression: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Expr (Maybe a) -> [Char]
forall a. Show a => a -> [Char]
show Expr (Maybe a)
expr

imputeM :: (Columnable a) => Expr (Maybe a) -> a -> FrameM (Expr a)
imputeM :: forall a. Columnable a => Expr (Maybe a) -> a -> FrameM (Expr a)
imputeM expr :: Expr (Maybe a)
expr@(Col Text
name) a
value = (DataFrame -> (DataFrame, Expr a)) -> FrameM (Expr a)
forall a. (DataFrame -> (DataFrame, a)) -> FrameM a
FrameM ((DataFrame -> (DataFrame, Expr a)) -> FrameM (Expr a))
-> (DataFrame -> (DataFrame, Expr a)) -> FrameM (Expr a)
forall a b. (a -> b) -> a -> b
$ \DataFrame
df ->
    let df' :: DataFrame
df' = Expr (Maybe a) -> a -> DataFrame -> DataFrame
forall b.
Columnable b =>
Expr (Maybe b) -> b -> DataFrame -> DataFrame
D.impute Expr (Maybe a)
expr a
value DataFrame
df
     in (DataFrame
df', Text -> Expr a
forall a. Columnable a => Text -> Expr a
Col Text
name)
imputeM Expr (Maybe a)
expr a
_ = [Char] -> FrameM (Expr a)
forall a. HasCallStack => [Char] -> a
error ([Char] -> FrameM (Expr a)) -> [Char] -> FrameM (Expr a)
forall a b. (a -> b) -> a -> b
$ [Char]
"Cannot impute on compound expression: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Expr (Maybe a) -> [Char]
forall a. Show a => a -> [Char]
show Expr (Maybe a)
expr

runFrameM :: DataFrame -> FrameM a -> (a, DataFrame)
runFrameM :: forall a. DataFrame -> FrameM a -> (a, DataFrame)
runFrameM DataFrame
df (FrameM DataFrame -> (DataFrame, a)
action) =
    let (DataFrame
df', a
a) = DataFrame -> (DataFrame, a)
action DataFrame
df
     in (a
a, DataFrame
df')

evalFrameM :: DataFrame -> FrameM a -> a
evalFrameM :: forall a. DataFrame -> FrameM a -> a
evalFrameM DataFrame
df FrameM a
m = (a, DataFrame) -> a
forall a b. (a, b) -> a
fst (DataFrame -> FrameM a -> (a, DataFrame)
forall a. DataFrame -> FrameM a -> (a, DataFrame)
runFrameM DataFrame
df FrameM a
m)

execFrameM :: DataFrame -> FrameM a -> DataFrame
execFrameM :: forall a. DataFrame -> FrameM a -> DataFrame
execFrameM DataFrame
df FrameM a
m = (a, DataFrame) -> DataFrame
forall a b. (a, b) -> b
snd (DataFrame -> FrameM a -> (a, DataFrame)
forall a. DataFrame -> FrameM a -> (a, DataFrame)
runFrameM DataFrame
df FrameM a
m)