{-
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 AllowAmbiguousTypes #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE UndecidableInstances #-}


module Control.Distributed.Blast.Types
(
  Computation
  , LocalComputation
  , RemoteComputation
  , Kind (..)
  , Partition
  , Chunkable (..)
  , UnChunkable (..)
  , ChunkableFreeVar (..)
  , ChunkFun
  , UnChunkFun
  , Fun (..)
  , FoldFun (..)
  , ExpClosure (..)
  , Indexable (..)
  , Builder (..)
  , Syntax (..)
  , GenericInfoMap
  , GenericInfo (..)
  , rapply'
  , rconst'
  , rconst
  , rconstIO
  , lconst
  , lconstIO
  , collect'
  , collect
  , lapply
  , rapply
  , refCount
  , generateReferenceMap
  , build
  , JobDesc (..)
  , Config (..)
  , defaultConfig
  , fun
  , closure
  , foldFun
  , foldClosure
  , funIO
  , closureIO
  , foldFunIO
  , foldClosureIO
  , partitionSize
  , getPart

)
where

--import Debug.Trace
import qualified  Control.Lens as Lens ()
import            Control.Monad.Operational
import qualified  Data.List as L
import qualified  Data.Map as M
import qualified  Data.Set as S
import qualified  Data.Serialize as S
import qualified  Data.Vector as Vc


data GenericInfo i = GenericInfo {
  giRefs :: S.Set Int -- set of parents, that is, nodes that reference this node
  , giInfo :: i
  }
    deriving Show

-- $(Lens.makeLenses ''GenericInfo)

type GenericInfoMap i = M.Map Int (GenericInfo i)

-- | Generic type describing a computation.
type Computation m e (k::Kind) a =
    Control.Monad.Operational.ProgramT (Syntax m) m (e k a)

-- | A computation that evaluates as a local value.
type LocalComputation a = forall e m. (Monad m, Builder m e) => Computation m e 'Local a


-- | A computation that evaluates as a remote value.
type RemoteComputation a = forall e m. (Monad m, Builder m e) => Computation m e 'Remote a

-- | Kind of computation.
data Kind = Remote | Local

-- | Represents the partitioning of a remote value.
type Partition a = Vc.Vector a

partitionSize :: Partition a -> Int
partitionSize v = Vc.length v

getPart :: Int -> Partition a -> Maybe a
getPart i p | i <= partitionSize p - 1 = Just $ p Vc.! i
getPart _ _ = Nothing

-- | Values that can be partitionned.
class Chunkable a b | a -> b, b -> a where
  -- | Given a value "a", chunk it into 'n' parts.
  chunk :: Int -> a -> Partition b

-- | Values that can be reconstructed from a list of parts.
class UnChunkable b a | b -> a, b -> a where
  -- | Given a list of parts, reconstruct a value.
  unChunk :: [b] -> a

-- | Values that can be reconstructed from a list of parts.
-- This applies to local values that are captured by a closure.
-- Helps optimize the implementation of remote relational operators or more generally, remote dyadic operators.
class ChunkableFreeVar a where
  -- | Given a list of parts, reconstruct a value.
  chunk' :: Int -> a -> Partition a
  chunk' n a = Vc.generate n (const a)


data Fun e a b =
  Pure (a -> IO b)
  |forall c . (S.Serialize c, ChunkableFreeVar c) => Closure (e 'Local c) (c -> a -> IO b)

data FoldFun e a r =
  FoldPure (r -> a -> IO r)
  |forall c . (S.Serialize c,ChunkableFreeVar c) => FoldClosure (e 'Local c) (c -> r -> a -> IO r)

data ExpClosure e a b =
  forall c . (S.Serialize c, ChunkableFreeVar c) => ExpClosure (e 'Local c) (c -> a -> IO b)


class Indexable e where
  getIndex :: e (k::Kind) a -> Int

type ChunkFun a b = Int -> a -> Partition b
type UnChunkFun b a = [b] -> a

class (Indexable e) => Builder m e where
  makeRApply :: Int -> ExpClosure e a b -> e 'Remote a -> m (e 'Remote b)
  makeRConst :: (S.Serialize b) => Int -> ChunkFun a b -> IO a -> m (e 'Remote b)
  makeLConst :: Int -> IO a -> m (e 'Local a)
  makeCollect :: (S.Serialize b) => Int -> UnChunkFun b a -> e 'Remote b -> m (e 'Local a)
  makeLApply :: Int -> e 'Local (a -> b) -> e 'Local a -> m (e 'Local b)

data Syntax m e where
  StxRApply :: (Builder m e) => ExpClosure e a b -> e 'Remote a -> Syntax m (e 'Remote b)
  StxRConst :: (Builder m e, S.Serialize b) => ChunkFun a b -> IO a -> Syntax m (e 'Remote b)
  StxLConst :: (Builder m e) => IO a -> Syntax m (e 'Local a)
  StxCollect :: (Builder m e, S.Serialize b) => UnChunkFun b a -> e 'Remote b -> Syntax m (e 'Local a)
  StxLApply :: (Builder m e) => e 'Local (a -> b) -> e 'Local a -> Syntax m (e 'Local b)

-- | Applies a ExpClosure to remote value.
rapply' :: (Builder m e)
  => ExpClosure e a b
  -> e 'Remote a
  -> Computation m e 'Remote b
rapply' f a = singleton (StxRApply f a)

-- | Creates a remote value, passing a specific chunk function.
rconst' :: (S.Serialize b) =>
  ChunkFun a b -> IO a -> RemoteComputation b
rconst' f a = singleton (StxRConst f a)

-- | Creates a remote value.
rconst :: (S.Serialize b, Chunkable a b) => a -> RemoteComputation b
rconst a = rconst' chunk (return a)

-- | Creates a remote value.
rconstIO :: (S.Serialize b, Chunkable a b) => IO a -> RemoteComputation b
rconstIO a = rconst' chunk a

-- | Creates a local value.
lconst :: a -> LocalComputation a
lconst a = singleton (StxLConst (return a))

-- | Creates a local value.
lconstIO :: IO a -> LocalComputation a
lconstIO a = singleton (StxLConst a)

-- | Creates a local value from a remote value, passing a specific chunk function.
collect' :: (S.Serialize b, Builder m e) =>
  UnChunkFun b a ->  e 'Remote b -> Computation m e 'Local a
collect' f a = singleton (StxCollect f a)

-- | Creates a local value from a remote value, passing a specific chunk function.
collect :: (S.Serialize b, Builder m e, UnChunkable b a) =>
  e 'Remote b -> Computation m e 'Local a
collect a = collect' unChunk a

-- | Applies a function to a local value.
lapply :: (Builder m e) =>
  e 'Local (a -> b) -> e 'Local a -> Computation m e 'Local b
lapply f a = singleton (StxLApply f a)


-- | Applies a closure to remote value.
rapply :: (Monad m, Builder m e) =>
        Fun e a b -> e 'Remote a -> Computation m e 'Remote b
rapply fm e  = do
  cs <- mkRemoteClosure fm
  rapply' cs e
  where
  mkRemoteClosure (Pure f) = do
    ue <- lconst ()
    return $ ExpClosure ue (\() a -> f a)
  mkRemoteClosure (Closure ce f) = return $ ExpClosure ce (\c a -> f c a)



refCount :: Int -> GenericInfoMap i -> Int
refCount n m =
  case M.lookup n m of
    Just (GenericInfo refs _) -> S.size refs
    Nothing -> error ("Ref count not found for node: " ++ show n)

addUnitInfo :: Int -> GenericInfoMap () -> GenericInfoMap ()
addUnitInfo n refMap =
  case M.lookup n refMap of
    Just _ -> error $  ("Node " ++ show n ++ " already exists")
    Nothing -> M.insert n (GenericInfo S.empty ()) refMap


reference :: Int -> Int -> GenericInfoMap i -> GenericInfoMap i
reference parent child refMap = do
  case M.lookup child refMap of
    Just inf@(GenericInfo old _) -> M.insert child (inf {giRefs = S.insert parent old}) refMap
    Nothing -> error $  ("Node " ++ show child ++ " is referenced before being visited")

generateReferenceMap ::forall a m e. (Builder m e, Monad m) =>  Int -> GenericInfoMap () -> ProgramT (Syntax m) m (e 'Local a) -> m (GenericInfoMap (), Int)
generateReferenceMap counter refMap p = do
    pv <- viewT p
    eval pv
    where
    eval :: (Builder m e, Monad m) => ProgramViewT (Syntax m) m (e 'Local a) -> m (GenericInfoMap(), Int)
    eval (StxRApply cs@(ExpClosure ce _) a :>>=  is) = do
      e <- makeRApply counter cs a
      let refMap' = addUnitInfo counter refMap
      let refMap'' = reference counter (getIndex ce) refMap'
      let refMap''' = reference counter (getIndex a) refMap''
      generateReferenceMap (counter+1) refMap''' (is e)
    eval (StxRConst f a :>>=  is) = do
      e <- makeRConst counter f a
      let refMap' = addUnitInfo counter refMap
      generateReferenceMap (counter+1) refMap' (is e)
    eval (StxLConst a :>>=  is) = do
      e <- makeLConst counter a
      let refMap' = addUnitInfo counter refMap
      generateReferenceMap (counter+1) refMap' (is e)
    eval (StxCollect f a :>>=  is) = do
      e <- makeCollect counter f a
      let refMap' = addUnitInfo counter refMap
      let refMap'' = reference counter (getIndex a) refMap'
      generateReferenceMap (counter+1) refMap'' (is e)
    eval (StxLApply f a :>>=  is) = do
      e <- makeLApply counter f a
      let refMap' = addUnitInfo counter refMap
      let refMap'' = reference counter (getIndex f) refMap'
      let refMap''' = reference counter (getIndex a) refMap''
      generateReferenceMap (counter+1) refMap''' (is e)
    eval (Return _) = return (refMap, counter)



build ::forall a m e. (Builder m e, Monad m) => GenericInfoMap () -> Int -> ProgramT (Syntax m) m (e 'Local a)  -> m (e 'Local a)
build refMap counter p = do
    pv <- viewT p
    eval pv
    where
    eval :: (Builder m e, Monad m) => ProgramViewT (Syntax m) m (e 'Local a) -> m (e 'Local a)
    eval (StxRApply cs@(ExpClosure _ _) a :>>=  is) = do
      e <- makeRApply counter cs a
      build refMap (counter+1) (is e)
    eval (StxRConst chunkFun a :>>=  is) = do
      e <- makeRConst counter chunkFun a
      build refMap (counter+1) (is e)
    eval (StxLConst a :>>=  is) = do
      e <- makeLConst counter a
      build refMap (counter+1) (is e)
    eval (StxCollect f a :>>=  is) = do
      e <- makeCollect counter f a
      build refMap (counter+1) (is e)
    eval (StxLApply f a :>>=  is) = do
      e <- makeLApply counter f a
      build refMap (counter+1) (is e)
    eval (Return a) = return a


-- | Definition of a recursive job.
data JobDesc a b = MkJobDesc {
    -- | The initial value passed to the computation generator.
    seed :: a
  -- | The computation generator.
  , computationGen :: a -> (forall e m. (Monad m, Builder m e) => Computation m e 'Local (a, b))
  -- | An action that is executed after each iteration.
  , reportingAction :: a -> b -> IO a
  -- | Predicate that determines whether or not to continue the computation (False to continue, True to exit)
  , shouldStop  :: a -> a -> b -> Bool
  }


data Config = MkConfig
  {
    slaveAvailability :: Float    -- ^ Probability of slave failure. Used in testing.
    , statefullSlaves :: Bool       -- ^ True turns on the statefull slave mode. Slaves are stateless if False.
  }

-- | Default configuration
-- @
-- defaultConfig = MkConfig False 1.0 True
-- @
defaultConfig :: Config
defaultConfig = MkConfig 1.0 True

-- instances

instance {-# OVERLAPPABLE #-} ChunkableFreeVar a
instance {-# OVERLAPPABLE #-} (ChunkableFreeVar a , ChunkableFreeVar b) => ChunkableFreeVar (a,b) where
  chunk' n (a, c) =
    Vc.zip pb pd
    where
    pb = chunk' n a
    pd = chunk' n c

instance ChunkableFreeVar ()


instance {-# OVERLAPPABLE #-} Chunkable [a] [a] where
  chunk nbBuckets l =
    Vc.reverse $ Vc.fromList $ go [] nbBuckets l
    where
    go acc 1 ls = ls:acc
    go acc n ls = go (L.take nbPerBucket ls : acc) (n-1) (L.drop nbPerBucket ls)
    len = L.length l
    nbPerBucket = len `div` nbBuckets

instance {-# OVERLAPPABLE #-} UnChunkable [a] [a] where
  unChunk l = L.concat l



-- | Creates a closure from a pure function.
fun :: (a -> b) -> Fun e a b
fun f = Pure (return . f)

-- | Creates a closure from a pure function and a local value.
closure :: (S.Serialize c, ChunkableFreeVar c) => e 'Local c -> (c -> a -> b) -> Fun e a b
closure ce f = Closure ce (\c a -> return $ f c a)


-- | Creates a folding closure from a pure function.
foldFun :: (r -> a -> r) -> FoldFun e a r
foldFun f = FoldPure (\r a -> return $ f r a)

-- | Creates a folding closure from a pure function and a local value.
foldClosure :: (S.Serialize c, ChunkableFreeVar c) => e 'Local c -> (c -> r -> a -> r) -> FoldFun e a r
foldClosure ce f = FoldClosure ce (\c r a -> return $ f c r a)

-- | Creates a closure from a impure function.
funIO :: (a -> IO b) -> Fun k a b
funIO f = Pure f

-- | Creates a closure from a impure function and a local value.
closureIO :: (S.Serialize c, ChunkableFreeVar c) => e 'Local c -> (c -> a -> IO b) -> Fun e a b
closureIO ce f = Closure ce f

-- | Creates a folding closure from a impure function.
foldFunIO :: (r -> a -> IO r) -> FoldFun e a r
foldFunIO f = FoldPure f

-- | Creates a folding closure from a impure function and a local value.
foldClosureIO :: (S.Serialize c, ChunkableFreeVar c) => e 'Local c -> (c -> r -> a -> IO r) -> FoldFun e a r
foldClosureIO ce f = FoldClosure ce f