content: {-
  Copyright   : (c) Jean-Christophe Mincke, 2016-2017
  
  This Source Code Form is subject to the terms of the Mozilla Public
  License, v. 2.0. If a copy of the MPL was not distributed with this
  file, You can obtain one at http://mozilla.org/MPL/2.0/.
  -}
  
  
  {-# LANGUAGE DataKinds #-}
  {-# LANGUAGE DeriveAnyClass #-}
  {-# LANGUAGE DeriveFunctor #-}
  {-# LANGUAGE DeriveGeneric #-}
  {-# LANGUAGE FlexibleContexts #-}
  {-# LANGUAGE FlexibleInstances #-}
  {-# LANGUAGE GeneralizedNewtypeDeriving #-}
  {-# LANGUAGE MultiParamTypeClasses #-}
  {-# LANGUAGE RankNTypes #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  {-# LANGUAGE TypeFamilies #-}
  {-# LANGUAGE UndecidableInstances #-}
  
  module Control.Distributed.Blast.Syntax
  (
    -- * Specialized syntax primitives.
    rmap
    , rflatmap
    , rfilter
    , (<**>)
    , (<$$>)
    , lfold
    , lfold'
    , rfold
    , rfold'
    , rjoin
    , count
    , rKeyedJoin
  
    -- * Types.
    , KeyedVal (..)
    , Range (..)
  
    -- * Miscellanous functions.
    , rangeToList
  
    -- * Class.
    , Joinable (..)
  )
  where
  
  --import Debug.Trace
  import            Control.Monad hiding (join)
  import            Data.Foldable
  import            Data.Hashable
  import qualified  Data.HashMap.Lazy as M
  import qualified  Data.List as L
  import            Data.Maybe (catMaybes)
  import            Data.Proxy
  import qualified  Data.Vector as Vc
  import qualified  Data.Serialize as S
  
  import            GHC.Generics (Generic)
  
  
  import            Control.Distributed.Blast.Types
  
  class Joinable a b where
    join :: a -> b -> Maybe (a, b)
  
  
  
  -- | Maps a closure over a remote collection.
  rmap :: (Monad m, Builder m e, Traversable t) =>
          Fun e a b -> e 'Remote (t a) -> Computation m e 'Remote (t b)
  rmap fm e  = do
    cs <- mkRemoteClosure fm
    rapply' cs e
    where
    mkRemoteClosure (Pure f) = do
      ue <- lconst ()
      return $ ExpClosure ue (\() a -> mapM f a)
    mkRemoteClosure (Closure ce f) = return $ ExpClosure ce (\c a -> mapM (f c) a)
  
  -- | Maps a closure over a remote collection and concatenates the results.
  rflatmap :: (Monad m, Foldable t, Builder m e, Monoid b) =>
              Fun e a b -> e 'Remote (t a) -> Computation m e 'Remote b
  rflatmap fp e = do
    cs <- mkRemoteClosure fp
    rapply' cs e
    where
    mkRemoteClosure (Pure f) = do
      ue <- lconst ()
      return $ ExpClosure ue (\() a -> foldMap f a)
    mkRemoteClosure (Closure ce f) = return $ ExpClosure ce (\c a -> foldMap (f c) a)
  
  -- | Applies a filter on a remote collection.
  rfilter :: (Monad m, Applicative f, Foldable t, Monoid (f a),
              Monoid (IO (f a)), Builder m e) =>
    Fun e a Bool -> e 'Remote (t a) -> Computation m e 'Remote (f a)
  rfilter p e = do
    cs <- mkRemoteClosure p
    rapply' cs e
    where
    mkRemoteClosure (Pure f) = do
      ue <- lconst ()
      return $ ExpClosure ue (\() ta -> do
          r <- foldMap (\a -> do
              b <- f a
              return $ if b then pure a else mempty) ta
          return r)
    mkRemoteClosure (Closure ce f) = return $ ExpClosure ce (\c ta -> do
          r <- foldMap (\a -> do
                b <- f c a
                return $ if b then pure a else mempty) ta
          return r)
  
  -- | Pseudo applicative syntax for local values.
  (<**>) :: (Monad m, Builder m e)
            =>  Computation m e 'Local (a -> b)
            -> e 'Local a -> Computation m e 'Local b
  f <**> a = do
    cs <- f
    lapply cs a
  
  
  -- | Pseudo applicative syntax for local values.
  (<$$>) :: (Monad m, Builder m e)
            => (a -> b) -> e 'Local a -> Computation m e 'Local b
  f <$$> e = lconst f <**> e
  
  
  -- | Local fold.
  lfold ::  (Monad m, Foldable t, Builder m e)
            => e 'Local (b -> a -> b)
            -> e 'Local b
            -> e 'Local (t a)
            -> Computation m e 'Local b
  lfold f zero a = do
    f' <- foldl <$$> f <**> zero
    lapply f' a
  
  
  -- | Local fold.
  lfold' :: (Monad m, Foldable t, Builder m e) =>
    (b -> a -> b)
    -> e 'Local b
    -> e 'Local (t a)
    -> Computation m e 'Local b
  lfold' f zero a = do
    f' <- lconst f
    lfold f' zero a
  
  -- | Counts the number of elements in a collection.
  count ::  (Monad m, Foldable t, Builder m e)
            => e 'Local (t a) -> Computation m e 'Local Int
  count e = do
    zero <- lconst (0::Int)
    f <- lconst (\b _ -> b+1)
    lfold f zero e
  
  -- | Remote fold. Returns a value of type '[r]' which is guaranteed to be "Unchunkable".
  rfold ::  (Builder m e, Traversable t, Applicative t, S.Serialize r, Monad m, ChunkableFreeVar r)
            => FoldFun e a r -> e 'Local r -> e 'Remote (t a) -> Computation m e 'Remote [r]
  rfold fp zero e = do
    cs <- mkRemoteClosure fp
    rapply' cs e
    where
    mkRemoteClosure (FoldPure f) = do
        cv <- (\z -> ((), z)) <$$> zero
        return $ ExpClosure cv (\((), z) a -> do
                  r <- foldM f z a
                  return [r])
    mkRemoteClosure (FoldClosure ce f) = do
        cv <- (\c z -> (c, z)) <$$> ce <**> zero
        return $ ExpClosure cv (\(c,z) a -> do
                  r <- foldM  (f c) z a
                  return [r])
  
  
  -- | Remote fold followed by a local aggregation.
  -- Correct if and only if the folding function is both associative and commutative.
  rfold' :: (Monad m, Applicative t, Traversable t, S.Serialize r, Builder m e, ChunkableFreeVar r) =>
    FoldFun e a r
    -> ([r] -> b)
    -> e 'Local r
    -> e 'Remote (t a)
    ->  Computation m e 'Local b
  rfold' f aggregator zero a = do
    rs <- rfold f zero a
    ars <- collect' unChunk rs
    aggregator <$$> ars
  
  
  instance Joinable a b where
    join a b = Just (a, b)
  
  
  fromList' :: (Applicative t, Foldable t, Monoid (t a)) => [a] -> t a
  fromList' l = foldMap pure l
  
  rjoin :: (Monad m, Applicative t, Foldable t, Foldable t1, Foldable t2,
            Monoid (t (a1, a2)), S.Serialize (t1 a1),
            Builder m e,
            Joinable a1 a2,
            ChunkableFreeVar (t1 a1),
            UnChunkable (t1 a1) (t1 a1) ) =>
    e 'Remote (t1 a1) -> e 'Remote (t2 a2) -> Computation m e 'Remote (t (a1, a2))
  rjoin a b = do
    a' <- collect' unChunk a
    let cs = ExpClosure a' (\av bv -> return $ fromList' $ catMaybes [join x y | x <- toList av, y <- toList bv])
    rapply' cs b
  
  data KeyedVal k v = KeyedVal k v
    deriving (Generic, S.Serialize, Show)
  
  instance (Eq k) => Joinable (KeyedVal k a) (KeyedVal k b) where
    join (KeyedVal k1 a) (KeyedVal k2 b) | k1 == k2 = Just (KeyedVal k1 a, KeyedVal k2 b)
    join (KeyedVal _ _) (KeyedVal _ _) = Nothing
  
  
  data OptiT t k v = OptiT (t (KeyedVal k v))
    deriving (Generic)
  
  instance (Show (t (KeyedVal k v))) => Show (OptiT t k v) where
    show (OptiT x) = show x
  
  instance (S.Serialize (t (KeyedVal k v))) => S.Serialize (OptiT t k v)
  
  instance  {-# OVERLAPPING #-} (Hashable k) => Chunkable [KeyedVal k v] [KeyedVal k v] where
    chunk nbBuckets l =
      Vc.reverse $ Vc.generate nbBuckets (\i -> buckets M.! i)
      where
      buckets = L.foldl proc M.empty l
      proc bucket' kv@(KeyedVal k _) = let
        i = hash k `mod` nbBuckets
        in M.insertWith (++) i [kv] bucket'
  
  
  
  instance (Applicative t, Foldable t, Monoid (t (KeyedVal k v))
           , Chunkable (t (KeyedVal k v)) (t (KeyedVal k v))
            ) =>
    ChunkableFreeVar (OptiT t k v) where
    chunk' n (OptiT tkvs) = fmap OptiT $ chunk n tkvs
  
  
  -- | Optimized remote join operation between 2 collections of (key, value) pairs.
  -- Only works if the remote value "e 'Remote (t2 (KeyedVal k b))" is chunked according to "Chunkable (t (KeyedVal k v)) (t (KeyedVal k v))"
  rKeyedJoin
    :: (Eq k, Monad m, Applicative t, Applicative t1,
        Foldable t, Foldable t1, Foldable t2,
        Monoid (t (KeyedVal k (a, b))), Monoid (t1 (KeyedVal k a)),
        UnChunkable (t1 (KeyedVal k a)) (t1 (KeyedVal k a)), Chunkable (t1 (KeyedVal k a)) (t1 (KeyedVal k a)),
        Builder m e, S.Serialize (t1 (KeyedVal k a))) =>
       Proxy t
       -> e 'Remote (t1 (KeyedVal k a))
       -> e 'Remote (t2 (KeyedVal k b))
       -> Computation m e 'Remote (t (KeyedVal k (a, b)))
  rKeyedJoin _ a b = do
    a' <- collect' unChunk a
    ja <- OptiT <$$> a'
    let cs = ExpClosure ja (\(OptiT av) bv -> return $ fromList' $ catMaybes [doJoin x y | x <- toList av, y <- toList bv])
    rapply' cs b
    where
    doJoin (KeyedVal k1 a') (KeyedVal k2 b') | k1 == k2 = Just $ KeyedVal k1 (a', b')
    doJoin (KeyedVal _ _) (KeyedVal _ _) = Nothing
  
  
  
  
  -- | A Range defined by two integer: [a, b[
  data Range = Range Int Int
    deriving (Eq, Show, Generic, S.Serialize)
  
  -- | Transforms a "Range" into a list.
  --
  -- @
  -- rangeToList (Range 1 4) == [1,2,3]
  -- @
  rangeToList :: Range -> [Int]
  rangeToList (Range a b) = [a .. (b-1)]
  
  
  instance Chunkable Range Range where
    chunk nbBuckets (Range minV maxV) =
      Vc.fromList $ L.reverse $ go [] minV nbBuckets
      where
      delta = (maxV - minV) `div` nbBuckets
      go ranges current 1 = (Range current maxV):ranges
      go ranges current n | current >= maxV = go (Range current current : ranges) current (n - 1)
      go ranges current n =
        go (Range current end' : ranges) end' (n - 1)
        where
        end' = if end > maxV then maxV else end
        end = current + delta
  
  
  
  
  
  
  
  
  
  
  
  
  
  
  
  
  
  