{-# LANGUAGE ExplicitNamespaces #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE ViewPatterns #-}
{-# LANGUAGE PatternSynonyms #-}

module Data.DataFrame.Internal.Function where

import qualified Data.Text as T
import qualified Data.Vector as V

import Data.DataFrame.Internal.Types
import Data.Typeable ( Typeable, type (:~:)(Refl) )
import Data.Type.Equality (TestEquality(testEquality))
import Type.Reflection (typeRep, typeOf)

-- A GADT to wrap functions so we can have hetegeneous lists of functions.
data Function where
    F1 :: forall a b . (Columnable a, Columnable b) => (a -> b) -> Function
    F2 :: forall a b c . (Columnable a, Columnable b, Columnable c) => (a -> b -> c) -> Function
    F3 :: forall a b c d . (Columnable a, Columnable b, Columnable c, Columnable d) => (a -> b -> c -> d) -> Function
    F4 :: forall a b c d e . (Columnable a, Columnable b, Columnable c, Columnable d, Columnable e) => (a -> b -> c -> d -> e) -> Function
    Cond :: forall a . (Columnable a) => (a -> Bool) -> Function
    ICond :: forall a . (Columnable a) => (Int -> a -> Bool) -> Function

-- Helper class to do the actual wrapping
class WrapFunction a where
    wrapFunction :: a -> Function

-- Instance for 1-argument functions
instance (Columnable a, Columnable b) => WrapFunction (a -> b) where
    wrapFunction :: (Columnable a, Columnable b) => (a -> b) -> Function
    wrapFunction :: (Columnable a, Columnable b) => (a -> b) -> Function
wrapFunction = (a -> b) -> Function
forall a b. (Columnable a, Columnable b) => (a -> b) -> Function
F1

-- Instance for 2-argument functions
instance {-# INCOHERENT #-} (Columnable a, Columnable b, Columnable c) => WrapFunction (a -> b -> c) where
    wrapFunction :: (Columnable a, Columnable b, Columnable c) => (a -> b -> c) -> Function
    wrapFunction :: (Columnable a, Columnable b, Columnable c) =>
(a -> b -> c) -> Function
wrapFunction = (a -> b -> c) -> Function
forall a b c.
(Columnable a, Columnable b, Columnable c) =>
(a -> b -> c) -> Function
F2

-- Instance for 3-argument functions
instance {-# INCOHERENT #-} (Columnable a, Columnable b, Columnable c, Columnable d) => WrapFunction (a -> b -> c -> d) where
    wrapFunction :: (Columnable a, Columnable b, Columnable c, Columnable d) => (a -> b -> c -> d) -> Function
    wrapFunction :: (Columnable a, Columnable b, Columnable c, Columnable d) =>
(a -> b -> c -> d) -> Function
wrapFunction = (a -> b -> c -> d) -> Function
forall a b c d.
(Columnable a, Columnable b, Columnable c, Columnable d) =>
(a -> b -> c -> d) -> Function
F3

instance {-# INCOHERENT #-} (Columnable a, Columnable b, Columnable c, Columnable d, Columnable e) => WrapFunction (a -> b -> c -> d -> e) where
    wrapFunction :: (Columnable a, Columnable b, Columnable c, Columnable d, Columnable e) => (a -> b -> c -> d -> e) -> Function
    wrapFunction :: (Columnable a, Columnable b, Columnable c, Columnable d,
 Columnable e) =>
(a -> b -> c -> d -> e) -> Function
wrapFunction = (a -> b -> c -> d -> e) -> Function
forall a b c d e.
(Columnable a, Columnable b, Columnable c, Columnable d,
 Columnable e) =>
(a -> b -> c -> d -> e) -> Function
F4

-- The main function that wraps arbitrary functions
func :: forall fn . WrapFunction fn => fn -> Function
func :: forall fn. WrapFunction fn => fn -> Function
func = fn -> Function
forall fn. WrapFunction fn => fn -> Function
wrapFunction

pattern Empty :: V.Vector a
pattern $mEmpty :: forall {r} {a}. Vector a -> ((# #) -> r) -> ((# #) -> r) -> r
$bEmpty :: forall a. Vector a
Empty <- (V.null -> True) where Empty = Vector a
forall a. Vector a
V.empty 

uncons :: V.Vector a -> Maybe (a, V.Vector a)
uncons :: forall a. Vector a -> Maybe (a, Vector a)
uncons Vector a
Empty = Maybe (a, Vector a)
forall a. Maybe a
Nothing
uncons Vector a
v     = (a, Vector a) -> Maybe (a, Vector a)
forall a. a -> Maybe a
Just (Vector a -> a
forall a. Vector a -> a
V.unsafeHead Vector a
v, Vector a -> Vector a
forall a. Vector a -> Vector a
V.unsafeTail Vector a
v)

pattern (:<|)  :: a -> V.Vector a -> V.Vector a
pattern x $m:<| :: forall {r} {a}.
Vector a -> (a -> Vector a -> r) -> ((# #) -> r) -> r
:<| xs <- (uncons -> Just (x, xs))

funcApply :: forall c . (Columnable c) => V.Vector RowValue -> Function ->  c
funcApply :: forall c. Columnable c => Vector RowValue -> Function -> c
funcApply Vector RowValue
Empty Function
_ = [Char] -> c
forall a. HasCallStack => [Char] -> a
error [Char]
"Empty args"
funcApply (Value (a
x :: a') :<| Vector RowValue
Empty) (F1 (a -> b
f :: (a -> b))) = case 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 @a) of
        Just a :~: a
Refl -> case TypeRep b -> TypeRep c -> Maybe (b :~: c)
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 (b -> TypeRep b
forall a. Typeable a => a -> TypeRep a
typeOf (a -> b
f a
a
x)) (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @c) of
            Just b :~: c
Refl -> a -> b
f a
a
x
            Maybe (b :~: c)
Nothing -> [Char] -> c
forall a. HasCallStack => [Char] -> a
error [Char]
"Result type mismatch"
        Maybe (a :~: a)
Nothing -> [Char] -> c
forall a. HasCallStack => [Char] -> a
error [Char]
"Arg type mismatch"
funcApply (Value (a
x :: a') :<| Vector RowValue
xs) (F2 (a -> b -> c
f :: (a -> b))) = case 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 (a -> TypeRep a
forall a. Typeable a => a -> TypeRep a
typeOf a
x) (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @a) of
        Just a :~: a
Refl -> Vector RowValue -> Function -> c
forall c. Columnable c => Vector RowValue -> Function -> c
funcApply Vector RowValue
xs ((b -> c) -> Function
forall a b. (Columnable a, Columnable b) => (a -> b) -> Function
F1 (a -> b -> c
f a
a
x))
        Maybe (a :~: a)
Nothing -> [Char] -> c
forall a. HasCallStack => [Char] -> a
error [Char]
"Arg type mismatch"
funcApply (Value (a
x :: a') :<| Vector RowValue
xs) (F3 (a -> b -> c -> d
f :: (a -> b))) = case 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 (a -> TypeRep a
forall a. Typeable a => a -> TypeRep a
typeOf a
x) (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @a) of
        Just a :~: a
Refl -> Vector RowValue -> Function -> c
forall c. Columnable c => Vector RowValue -> Function -> c
funcApply Vector RowValue
xs ((b -> c -> d) -> Function
forall a b c.
(Columnable a, Columnable b, Columnable c) =>
(a -> b -> c) -> Function
F2 (a -> b -> c -> d
f a
a
x))
        Maybe (a :~: a)
Nothing -> [Char] -> c
forall a. HasCallStack => [Char] -> a
error [Char]
"Arg type mismatch"
funcApply (Value (a
x :: a') :<| Vector RowValue
xs) (F4 (a -> b -> c -> d -> e
f :: (a -> b))) = case 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 (a -> TypeRep a
forall a. Typeable a => a -> TypeRep a
typeOf a
x) (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @a) of
        Just a :~: a
Refl -> Vector RowValue -> Function -> c
forall c. Columnable c => Vector RowValue -> Function -> c
funcApply Vector RowValue
xs ((b -> c -> d -> e) -> Function
forall a b c d.
(Columnable a, Columnable b, Columnable c, Columnable d) =>
(a -> b -> c -> d) -> Function
F3 (a -> b -> c -> d -> e
f a
a
x))
        Maybe (a :~: a)
Nothing -> [Char] -> c
forall a. HasCallStack => [Char] -> a
error [Char]
"Arg type mismatch"