module Opaleye.SQLite.Internal.Aggregate where

import           Control.Applicative (Applicative, pure, (<*>))

import qualified Data.Profunctor as P
import qualified Data.Profunctor.Product as PP

import qualified Opaleye.SQLite.Internal.PackMap as PM
import qualified Opaleye.SQLite.Internal.PrimQuery as PQ
import qualified Opaleye.SQLite.Internal.Tag as T
import qualified Opaleye.SQLite.Internal.Column as C

import qualified Opaleye.SQLite.Internal.HaskellDB.PrimQuery as HPQ

{-|
An 'Aggregator' takes a collection of rows of type @a@, groups
them, and transforms each group into a single row of type @b@. This
corresponds to aggregators using @GROUP BY@ in SQL.

An 'Aggregator' corresponds closely to a 'Control.Foldl.Fold' from the
@foldl@ package.  Whereas an 'Aggregator' @a@ @b@ takes each group of
type @a@ to a single row of type @b@, a 'Control.Foldl.Fold' @a@ @b@
takes a list of @a@ and returns a single row of type @b@.
-}
newtype Aggregator a b = Aggregator
                         (PM.PackMap (Maybe HPQ.AggrOp, HPQ.PrimExpr) HPQ.PrimExpr
                                     a b)

makeAggr' :: Maybe HPQ.AggrOp -> Aggregator (C.Column a) (C.Column b)
makeAggr' :: Maybe AggrOp -> Aggregator (Column a) (Column b)
makeAggr' Maybe AggrOp
m = PackMap (Maybe AggrOp, PrimExpr) PrimExpr (Column a) (Column b)
-> Aggregator (Column a) (Column b)
forall a b.
PackMap (Maybe AggrOp, PrimExpr) PrimExpr a b -> Aggregator a b
Aggregator ((forall (f :: * -> *).
 Applicative f =>
 ((Maybe AggrOp, PrimExpr) -> f PrimExpr)
 -> Column a -> f (Column b))
-> PackMap (Maybe AggrOp, PrimExpr) PrimExpr (Column a) (Column b)
forall a b s t.
(forall (f :: * -> *). Applicative f => (a -> f b) -> s -> f t)
-> PackMap a b s t
PM.PackMap
                          (\(Maybe AggrOp, PrimExpr) -> f PrimExpr
f (C.Column e) -> (PrimExpr -> Column b) -> f PrimExpr -> f (Column b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap PrimExpr -> Column b
forall a. PrimExpr -> Column a
C.Column ((Maybe AggrOp, PrimExpr) -> f PrimExpr
f (Maybe AggrOp
m, PrimExpr
e))))

makeAggr :: HPQ.AggrOp -> Aggregator (C.Column a) (C.Column b)
makeAggr :: AggrOp -> Aggregator (Column a) (Column b)
makeAggr = Maybe AggrOp -> Aggregator (Column a) (Column b)
forall a b. Maybe AggrOp -> Aggregator (Column a) (Column b)
makeAggr' (Maybe AggrOp -> Aggregator (Column a) (Column b))
-> (AggrOp -> Maybe AggrOp)
-> AggrOp
-> Aggregator (Column a) (Column b)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AggrOp -> Maybe AggrOp
forall a. a -> Maybe a
Just

runAggregator :: Applicative f => Aggregator a b
              -> ((Maybe HPQ.AggrOp, HPQ.PrimExpr) -> f HPQ.PrimExpr) -> a -> f b
runAggregator :: Aggregator a b
-> ((Maybe AggrOp, PrimExpr) -> f PrimExpr) -> a -> f b
runAggregator (Aggregator PackMap (Maybe AggrOp, PrimExpr) PrimExpr a b
a) = PackMap (Maybe AggrOp, PrimExpr) PrimExpr a b
-> ((Maybe AggrOp, PrimExpr) -> f PrimExpr) -> a -> f b
forall (f :: * -> *) a b s t.
Applicative f =>
PackMap a b s t -> (a -> f b) -> s -> f t
PM.traversePM PackMap (Maybe AggrOp, PrimExpr) PrimExpr a b
a

aggregateU :: Aggregator a b
           -> (a, PQ.PrimQuery, T.Tag) -> (b, PQ.PrimQuery, T.Tag)
aggregateU :: Aggregator a b -> (a, PrimQuery, Tag) -> (b, PrimQuery, Tag)
aggregateU Aggregator a b
agg (a
c0, PrimQuery
primQ, Tag
t0) = (b
c1, PrimQuery
primQ', Tag -> Tag
T.next Tag
t0)
  where (b
c1, [(Symbol, (Maybe AggrOp, PrimExpr))]
projPEs) =
          PM [(Symbol, (Maybe AggrOp, PrimExpr))] b
-> (b, [(Symbol, (Maybe AggrOp, PrimExpr))])
forall a r. PM [a] r -> (r, [a])
PM.run (Aggregator a b
-> ((Maybe AggrOp, PrimExpr)
    -> StateT
         ([(Symbol, (Maybe AggrOp, PrimExpr))], Int) Identity PrimExpr)
-> a
-> PM [(Symbol, (Maybe AggrOp, PrimExpr))] b
forall (f :: * -> *) a b.
Applicative f =>
Aggregator a b
-> ((Maybe AggrOp, PrimExpr) -> f PrimExpr) -> a -> f b
runAggregator Aggregator a b
agg (Tag
-> (Maybe AggrOp, PrimExpr)
-> StateT
     ([(Symbol, (Maybe AggrOp, PrimExpr))], Int) Identity PrimExpr
extractAggregateFields Tag
t0) a
c0)

        primQ' :: PrimQuery
primQ' = [(Symbol, (Maybe AggrOp, PrimExpr))] -> PrimQuery -> PrimQuery
PQ.Aggregate [(Symbol, (Maybe AggrOp, PrimExpr))]
projPEs PrimQuery
primQ

extractAggregateFields :: T.Tag -> (Maybe HPQ.AggrOp, HPQ.PrimExpr)
      -> PM.PM [(HPQ.Symbol, (Maybe HPQ.AggrOp, HPQ.PrimExpr))] HPQ.PrimExpr
extractAggregateFields :: Tag
-> (Maybe AggrOp, PrimExpr)
-> StateT
     ([(Symbol, (Maybe AggrOp, PrimExpr))], Int) Identity PrimExpr
extractAggregateFields = String
-> Tag
-> (Maybe AggrOp, PrimExpr)
-> StateT
     ([(Symbol, (Maybe AggrOp, PrimExpr))], Int) Identity PrimExpr
forall primExpr.
String -> Tag -> primExpr -> PM [(Symbol, primExpr)] PrimExpr
PM.extractAttr String
"result"

-- { Boilerplate instances

instance Functor (Aggregator a) where
  fmap :: (a -> b) -> Aggregator a a -> Aggregator a b
fmap a -> b
f (Aggregator PackMap (Maybe AggrOp, PrimExpr) PrimExpr a a
g) = PackMap (Maybe AggrOp, PrimExpr) PrimExpr a b -> Aggregator a b
forall a b.
PackMap (Maybe AggrOp, PrimExpr) PrimExpr a b -> Aggregator a b
Aggregator ((a -> b)
-> PackMap (Maybe AggrOp, PrimExpr) PrimExpr a a
-> PackMap (Maybe AggrOp, PrimExpr) PrimExpr a b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f PackMap (Maybe AggrOp, PrimExpr) PrimExpr a a
g)

instance Applicative (Aggregator a) where
  pure :: a -> Aggregator a a
pure = PackMap (Maybe AggrOp, PrimExpr) PrimExpr a a -> Aggregator a a
forall a b.
PackMap (Maybe AggrOp, PrimExpr) PrimExpr a b -> Aggregator a b
Aggregator (PackMap (Maybe AggrOp, PrimExpr) PrimExpr a a -> Aggregator a a)
-> (a -> PackMap (Maybe AggrOp, PrimExpr) PrimExpr a a)
-> a
-> Aggregator a a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> PackMap (Maybe AggrOp, PrimExpr) PrimExpr a a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
  Aggregator PackMap (Maybe AggrOp, PrimExpr) PrimExpr a (a -> b)
f <*> :: Aggregator a (a -> b) -> Aggregator a a -> Aggregator a b
<*> Aggregator PackMap (Maybe AggrOp, PrimExpr) PrimExpr a a
x = PackMap (Maybe AggrOp, PrimExpr) PrimExpr a b -> Aggregator a b
forall a b.
PackMap (Maybe AggrOp, PrimExpr) PrimExpr a b -> Aggregator a b
Aggregator (PackMap (Maybe AggrOp, PrimExpr) PrimExpr a (a -> b)
f PackMap (Maybe AggrOp, PrimExpr) PrimExpr a (a -> b)
-> PackMap (Maybe AggrOp, PrimExpr) PrimExpr a a
-> PackMap (Maybe AggrOp, PrimExpr) PrimExpr a b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> PackMap (Maybe AggrOp, PrimExpr) PrimExpr a a
x)

instance P.Profunctor Aggregator where
  dimap :: (a -> b) -> (c -> d) -> Aggregator b c -> Aggregator a d
dimap a -> b
f c -> d
g (Aggregator PackMap (Maybe AggrOp, PrimExpr) PrimExpr b c
q) = PackMap (Maybe AggrOp, PrimExpr) PrimExpr a d -> Aggregator a d
forall a b.
PackMap (Maybe AggrOp, PrimExpr) PrimExpr a b -> Aggregator a b
Aggregator ((a -> b)
-> (c -> d)
-> PackMap (Maybe AggrOp, PrimExpr) PrimExpr b c
-> PackMap (Maybe AggrOp, PrimExpr) PrimExpr a d
forall (p :: * -> * -> *) a b c d.
Profunctor p =>
(a -> b) -> (c -> d) -> p b c -> p a d
P.dimap a -> b
f c -> d
g PackMap (Maybe AggrOp, PrimExpr) PrimExpr b c
q)

instance PP.ProductProfunctor Aggregator where
  empty :: Aggregator () ()
empty = Aggregator () ()
forall (p :: * -> * -> *). Applicative (p ()) => p () ()
PP.defaultEmpty
  ***! :: Aggregator a b -> Aggregator a' b' -> Aggregator (a, a') (b, b')
(***!) = Aggregator a b -> Aggregator a' b' -> Aggregator (a, a') (b, b')
forall (p :: * -> * -> *) a a' b b'.
(Applicative (p (a, a')), Profunctor p) =>
p a b -> p a' b' -> p (a, a') (b, b')
PP.defaultProfunctorProduct

instance PP.SumProfunctor Aggregator where
  Aggregator PackMap (Maybe AggrOp, PrimExpr) PrimExpr a b
x1 +++! :: Aggregator a b
-> Aggregator a' b' -> Aggregator (Either a a') (Either b b')
+++! Aggregator PackMap (Maybe AggrOp, PrimExpr) PrimExpr a' b'
x2 = PackMap
  (Maybe AggrOp, PrimExpr) PrimExpr (Either a a') (Either b b')
-> Aggregator (Either a a') (Either b b')
forall a b.
PackMap (Maybe AggrOp, PrimExpr) PrimExpr a b -> Aggregator a b
Aggregator (PackMap (Maybe AggrOp, PrimExpr) PrimExpr a b
x1 PackMap (Maybe AggrOp, PrimExpr) PrimExpr a b
-> PackMap (Maybe AggrOp, PrimExpr) PrimExpr a' b'
-> PackMap
     (Maybe AggrOp, PrimExpr) PrimExpr (Either a a') (Either b b')
forall (p :: * -> * -> *) a b a' b'.
SumProfunctor p =>
p a b -> p a' b' -> p (Either a a') (Either b b')
PP.+++! PackMap (Maybe AggrOp, PrimExpr) PrimExpr a' b'
x2)

-- }