{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE OverloadedStrings #-}

module DataFrame.Operators where

import Data.Function ((&))
import qualified Data.Text as T
import DataFrame.Internal.Column (Columnable)
import DataFrame.Internal.Expression (
    BinaryOp (
        MkBinaryOp,
        binaryCommutative,
        binaryFn,
        binaryName,
        binaryPrecedence,
        binarySymbol
    ),
    Expr (Binary, Col, If, Lit),
    NamedExpr,
    UExpr (UExpr),
 )

infix 8 .^^
infix 4 .==, .<, .<=, .>=, .>, ./=
infixr 3 .&&
infixr 2 .||
infixr 0 .=

(|>) :: a -> (a -> b) -> b
|> :: forall a b. a -> (a -> b) -> b
(|>) = a -> (a -> b) -> b
forall a b. a -> (a -> b) -> b
(&)

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

name :: (Show a) => Expr a -> T.Text
name :: forall a. Show a => Expr a -> Text
name (Col Text
n) = Text
n
name Expr a
other =
    [Char] -> Text
forall a. HasCallStack => [Char] -> a
error ([Char] -> Text) -> [Char] -> Text
forall a b. (a -> b) -> a -> b
$
        [Char]
"You must call `name` on a column reference. Not the expression: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Expr a -> [Char]
forall a. Show a => a -> [Char]
show Expr a
other

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

ifThenElse :: (Columnable a) => Expr Bool -> Expr a -> Expr a -> Expr a
ifThenElse :: forall a. Columnable a => Expr Bool -> Expr a -> Expr a -> Expr a
ifThenElse = Expr Bool -> Expr a -> Expr a -> Expr a
forall a. Columnable a => Expr Bool -> Expr a -> Expr a -> Expr a
If

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

(.=) :: (Columnable a) => T.Text -> Expr a -> NamedExpr
.= :: forall a. Columnable a => Text -> Expr a -> NamedExpr
(.=) = (Expr a -> Text -> NamedExpr) -> Text -> Expr a -> NamedExpr
forall a b c. (a -> b -> c) -> b -> a -> c
flip Expr a -> Text -> NamedExpr
forall a. Columnable a => Expr a -> Text -> NamedExpr
as

(.==) :: (Columnable a, Eq a) => Expr a -> Expr a -> Expr Bool
.== :: forall a. (Columnable a, Eq a) => Expr a -> Expr a -> Expr Bool
(.==) =
    BinaryOp a a Bool -> Expr a -> Expr a -> Expr Bool
forall c b a.
(Columnable c, Columnable b, Columnable a) =>
BinaryOp c b a -> Expr c -> Expr b -> Expr a
Binary
        ( MkBinaryOp
            { binaryFn :: a -> a -> Bool
binaryFn = a -> a -> Bool
forall a. Eq a => a -> a -> Bool
(==)
            , binaryName :: Text
binaryName = Text
"eq"
            , binarySymbol :: Maybe Text
binarySymbol = Text -> Maybe Text
forall a. a -> Maybe a
Just Text
"=="
            , binaryCommutative :: Bool
binaryCommutative = Bool
True
            , binaryPrecedence :: Int
binaryPrecedence = Int
4
            }
        )

(./=) :: (Columnable a, Eq a) => Expr a -> Expr a -> Expr Bool
./= :: forall a. (Columnable a, Eq a) => Expr a -> Expr a -> Expr Bool
(./=) =
    BinaryOp a a Bool -> Expr a -> Expr a -> Expr Bool
forall c b a.
(Columnable c, Columnable b, Columnable a) =>
BinaryOp c b a -> Expr c -> Expr b -> Expr a
Binary
        ( MkBinaryOp
            { binaryFn :: a -> a -> Bool
binaryFn = a -> a -> Bool
forall a. Eq a => a -> a -> Bool
(/=)
            , binaryName :: Text
binaryName = Text
"neq"
            , binarySymbol :: Maybe Text
binarySymbol = Text -> Maybe Text
forall a. a -> Maybe a
Just Text
"/="
            , binaryCommutative :: Bool
binaryCommutative = Bool
True
            , binaryPrecedence :: Int
binaryPrecedence = Int
4
            }
        )

(.<) :: (Columnable a, Ord a) => Expr a -> Expr a -> Expr Bool
.< :: forall a. (Columnable a, Ord a) => Expr a -> Expr a -> Expr Bool
(.<) =
    BinaryOp a a Bool -> Expr a -> Expr a -> Expr Bool
forall c b a.
(Columnable c, Columnable b, Columnable a) =>
BinaryOp c b a -> Expr c -> Expr b -> Expr a
Binary
        ( MkBinaryOp
            { binaryFn :: a -> a -> Bool
binaryFn = a -> a -> Bool
forall a. Ord a => a -> a -> Bool
(<)
            , binaryName :: Text
binaryName = Text
"lt"
            , binarySymbol :: Maybe Text
binarySymbol = Text -> Maybe Text
forall a. a -> Maybe a
Just Text
"<"
            , binaryCommutative :: Bool
binaryCommutative = Bool
False
            , binaryPrecedence :: Int
binaryPrecedence = Int
4
            }
        )

(.>) :: (Columnable a, Ord a) => Expr a -> Expr a -> Expr Bool
.> :: forall a. (Columnable a, Ord a) => Expr a -> Expr a -> Expr Bool
(.>) =
    BinaryOp a a Bool -> Expr a -> Expr a -> Expr Bool
forall c b a.
(Columnable c, Columnable b, Columnable a) =>
BinaryOp c b a -> Expr c -> Expr b -> Expr a
Binary
        ( MkBinaryOp
            { binaryFn :: a -> a -> Bool
binaryFn = a -> a -> Bool
forall a. Ord a => a -> a -> Bool
(>)
            , binaryName :: Text
binaryName = Text
"gt"
            , binarySymbol :: Maybe Text
binarySymbol = Text -> Maybe Text
forall a. a -> Maybe a
Just Text
">"
            , binaryCommutative :: Bool
binaryCommutative = Bool
False
            , binaryPrecedence :: Int
binaryPrecedence = Int
4
            }
        )

(.<=) :: (Columnable a, Ord a, Eq a) => Expr a -> Expr a -> Expr Bool
.<= :: forall a.
(Columnable a, Ord a, Eq a) =>
Expr a -> Expr a -> Expr Bool
(.<=) =
    BinaryOp a a Bool -> Expr a -> Expr a -> Expr Bool
forall c b a.
(Columnable c, Columnable b, Columnable a) =>
BinaryOp c b a -> Expr c -> Expr b -> Expr a
Binary
        ( MkBinaryOp
            { binaryFn :: a -> a -> Bool
binaryFn = a -> a -> Bool
forall a. Ord a => a -> a -> Bool
(<=)
            , binaryName :: Text
binaryName = Text
"leq"
            , binarySymbol :: Maybe Text
binarySymbol = Text -> Maybe Text
forall a. a -> Maybe a
Just Text
"<="
            , binaryCommutative :: Bool
binaryCommutative = Bool
False
            , binaryPrecedence :: Int
binaryPrecedence = Int
4
            }
        )

(.>=) :: (Columnable a, Ord a, Eq a) => Expr a -> Expr a -> Expr Bool
.>= :: forall a.
(Columnable a, Ord a, Eq a) =>
Expr a -> Expr a -> Expr Bool
(.>=) =
    BinaryOp a a Bool -> Expr a -> Expr a -> Expr Bool
forall c b a.
(Columnable c, Columnable b, Columnable a) =>
BinaryOp c b a -> Expr c -> Expr b -> Expr a
Binary
        ( MkBinaryOp
            { binaryFn :: a -> a -> Bool
binaryFn = a -> a -> Bool
forall a. Ord a => a -> a -> Bool
(>=)
            , binaryName :: Text
binaryName = Text
"geq"
            , binarySymbol :: Maybe Text
binarySymbol = Text -> Maybe Text
forall a. a -> Maybe a
Just Text
">="
            , binaryCommutative :: Bool
binaryCommutative = Bool
False
            , binaryPrecedence :: Int
binaryPrecedence = Int
4
            }
        )

(.&&) :: Expr Bool -> Expr Bool -> Expr Bool
.&& :: Expr Bool -> Expr Bool -> Expr Bool
(.&&) =
    BinaryOp Bool Bool Bool -> Expr Bool -> Expr Bool -> Expr Bool
forall c b a.
(Columnable c, Columnable b, Columnable a) =>
BinaryOp c b a -> Expr c -> Expr b -> Expr a
Binary
        ( MkBinaryOp
            { binaryFn :: Bool -> Bool -> Bool
binaryFn = Bool -> Bool -> Bool
(&&)
            , binaryName :: Text
binaryName = Text
"and"
            , binarySymbol :: Maybe Text
binarySymbol = Text -> Maybe Text
forall a. a -> Maybe a
Just Text
"&&"
            , binaryCommutative :: Bool
binaryCommutative = Bool
True
            , binaryPrecedence :: Int
binaryPrecedence = Int
3
            }
        )

(.||) :: Expr Bool -> Expr Bool -> Expr Bool
.|| :: Expr Bool -> Expr Bool -> Expr Bool
(.||) =
    BinaryOp Bool Bool Bool -> Expr Bool -> Expr Bool -> Expr Bool
forall c b a.
(Columnable c, Columnable b, Columnable a) =>
BinaryOp c b a -> Expr c -> Expr b -> Expr a
Binary
        ( MkBinaryOp
            { binaryFn :: Bool -> Bool -> Bool
binaryFn = Bool -> Bool -> Bool
(||)
            , binaryName :: Text
binaryName = Text
"or"
            , binarySymbol :: Maybe Text
binarySymbol = Text -> Maybe Text
forall a. a -> Maybe a
Just Text
"||"
            , binaryCommutative :: Bool
binaryCommutative = Bool
True
            , binaryPrecedence :: Int
binaryPrecedence = Int
2
            }
        )

(.^^) :: (Columnable a, Num a) => Expr a -> Int -> Expr a
.^^ :: forall a. (Columnable a, Num a) => Expr a -> Int -> Expr a
(.^^) Expr a
expr Int
i =
    BinaryOp a Int a -> Expr a -> Expr Int -> Expr a
forall c b a.
(Columnable c, Columnable b, Columnable a) =>
BinaryOp c b a -> Expr c -> Expr b -> Expr a
Binary
        ( MkBinaryOp
            { binaryFn :: a -> Int -> a
binaryFn = a -> Int -> a
forall a b. (Num a, Integral b) => a -> b -> a
(^)
            , binaryName :: Text
binaryName = Text
"pow"
            , binarySymbol :: Maybe Text
binarySymbol = Text -> Maybe Text
forall a. a -> Maybe a
Just Text
"^"
            , binaryCommutative :: Bool
binaryCommutative = Bool
False
            , binaryPrecedence :: Int
binaryPrecedence = Int
8
            }
        )
        Expr a
expr
        (Int -> Expr Int
forall a. Columnable a => a -> Expr a
Lit Int
i)