{-# LANGUAGE GADTs #-}
{-# LANGUAGE FunctionalDependencies #-}

-- | Combinators for generating a list of dependencies between calculations, to
-- be executed later in DependencyRunner.
module Byg.DependencyGenerator
  ( DepGenM
  , DepGenM'
  , TokenableTo(..)
  , toTupleToken
  , evalDepGenM
  , inject
  , onToken
  , onTupleToken
  , runFunctionIO
  , runFunctionIO_
  , mapDepGenM
  , mapDepGenM_
  , forDepGenM
  , forDepGenM_
  , filterDepGenM
  , filterDepGenM'
  , zipDepGenM
  , untupleFstDepGenM
  , untupleSndDepGenM
  , untupleDepGenM
  , unzipFstDepGenM
  , unzipSndDepGenM
  , unzipDepGenM
  ) where

import Byg.Types.Token (Token(..))
import Byg.Types.Functions (IsFunctionIO(..))
import Byg.Types.Dependency (Action(..), F(..), Dependency, makeDependency)

import Type.Reflection (Typeable, TypeRep, typeRep)
import Control.Monad.State (MonadState, State, runState, put, get)
import Control.Monad.Writer (MonadWriter, WriterT, runWriterT, tell)

newtype DepGenM a = DepGenM { unDepGenM :: WriterT [Dependency] (State Int) a }
  deriving (Functor, Applicative, Monad, MonadState Int, MonadWriter [Dependency])

type DepGenM' a = DepGenM (Token a)

runDepGenM :: Int -> DepGenM a -> ((a, [Dependency]), Int)
runDepGenM top m = runState (runWriterT (unDepGenM m)) top

evalDepGenM :: DepGenM () -> [Dependency]
evalDepGenM = snd . fst . runDepGenM 0

tellDep :: Dependency -> DepGenM ()
tellDep dep = tell [dep]

newToken :: (Typeable a, Show a) => DepGenM (Token a)
newToken = do
  top <- get
  let top' = top + 1
      target = Token top'
  put top'
  pure target

genDependencyM :: (Typeable a, Show a) => (Token a -> DepGenM Dependency) -> DepGenM (Token a)
genDependencyM f = do
  target <- newToken
  result <- f target
  tellDep result
  pure target

genDependency :: (Typeable a, Show a) => (Token a -> Dependency) -> DepGenM (Token a)
genDependency f = genDependencyM (pure . f)

inject :: (Show a, Typeable a) => a -> DepGenM (Token a)
inject x = genDependency (makeDependency NoToken (Inject x))

onToken :: (TokenableTo a t, Show a, Typeable a, Show b, Typeable b) => (a -> b) -> t -> DepGenM (Token b)
onToken f input = do
  input' <- toToken input
  genDependency (makeDependency input' (Function (F f)))

onTupleToken :: (TokenableTo a t1, Show a, Typeable a, TokenableTo b t2, Show b, Typeable b, Show r, Typeable r) => (a -> b -> r) -> t1 -> t2 -> DepGenM (Token r)
onTupleToken f input1 input2 = do
  tup <- toTupleToken input1 input2
  genDependency (makeDependency tup (Function (F (uncurry f))))

runFunctionIO :: IsFunctionIO f a b => f -> Token a -> DepGenM (Token b)
runFunctionIO f input = genDependency (makeDependency input (FunctionIO f))

runFunctionIO_ :: IsFunctionIO f a () => f -> Token a -> DepGenM ()
runFunctionIO_ f input = tellDep $ makeDependency input (FunctionIO f) NoToken

class (Show t, Typeable t) => TokenableTo t s | s -> t where
  toToken :: s -> DepGenM (Token t)
  tokenTypeRep :: s -> TypeRep t

instance (Show a, Typeable a) => TokenableTo a (Token a) where
  toToken = pure
  tokenTypeRep _ = typeRep

instance (Show a, Typeable a) => TokenableTo [a] [Token a] where
  toToken = pure . ListToken
  tokenTypeRep _ = typeRep

instance (Show a, Typeable a) => TokenableTo a (DepGenM (Token a)) where
  toToken = id
  tokenTypeRep _ = typeRep

instance (Show a, Typeable a) => TokenableTo [a] [DepGenM (Token a)] where
  toToken = fmap ListToken . sequence
  tokenTypeRep _ = typeRep

toTupleToken :: (TokenableTo ta a, TokenableTo tb b) => a -> b -> DepGenM (Token (ta, tb))
toTupleToken a b = TupleToken <$> toToken a <*> toToken b

mapDepGenM :: (TokenableTo [a] v, Typeable a, Show a, Typeable b, Show b) => (Token a -> DepGenM (Token b)) -> v -> DepGenM (Token [b])
mapDepGenM f input = do
  input' <- toToken input
  genDependencyM $ \target -> do
    top <- get
    let (((innerInp, innerOutp), subDeps), top') = runDepGenM top $ do
          inp <- newToken
          outp <- f inp
          pure (inp, outp)
    put top'
    pure (makeDependency input' (MapComp subDeps innerInp innerOutp) target)

mapDepGenM_ :: (TokenableTo [a] v, Typeable a, Show a) => (Token a -> DepGenM ()) -> v -> DepGenM ()
mapDepGenM_ f input = do
  _ <- mapDepGenM (\x -> f x >> pure NoToken) input
  pure ()

forDepGenM :: (TokenableTo [a] v, Typeable a, Show a, Typeable b, Show b) => v -> (Token a -> DepGenM (Token b)) -> DepGenM (Token [b])
forDepGenM = flip mapDepGenM

forDepGenM_ :: (TokenableTo [a] v, Typeable a, Show a) => v -> (Token a -> DepGenM ()) -> DepGenM ()
forDepGenM_ = flip mapDepGenM_

filterDepGenM' :: (TokenableTo [Bool] v, TokenableTo [a] u, Typeable a, Show a) => v -> u -> DepGenM (Token [a])
filterDepGenM' mask input = do
  tup <- toTupleToken input mask
  genDependency (makeDependency tup FilterComp)

filterDepGenM :: (TokenableTo [a] v, Typeable a, Show a) => (Token a -> DepGenM (Token Bool)) -> v -> DepGenM (Token [a])
filterDepGenM f input = do
  mask <- mapDepGenM f input
  filterDepGenM' mask input

zipDepGenM :: (TokenableTo [a] v, TokenableTo [b] u, Typeable a, Show a, Typeable b, Show b) => v -> u -> DepGenM (Token [(a, b)])
zipDepGenM a b = do
  a' <- toToken a
  b' <- toToken b
  pure $ ZipToken a' b'

untupleFstDepGenM :: (TokenableTo (a, b) t, Typeable a, Show a, Typeable b, Show b) => t -> DepGenM (Token a)
untupleFstDepGenM t = do
  t' <- toToken t
  case t' of
    TupleToken a _ ->
      pure a
    Token _ ->
      genDependency (makeDependency t' UntupleFst)

untupleSndDepGenM :: (TokenableTo (a, b) t, Typeable a, Show a, Typeable b, Show b) => t -> DepGenM (Token b)
untupleSndDepGenM t = do
  t' <- toToken t
  case t' of
    TupleToken _ b ->
      pure b
    Token _ ->
      genDependency (makeDependency t' UntupleSnd)

untupleDepGenM :: (TokenableTo (a, b) t, Typeable a, Show a, Typeable b, Show b) => t -> DepGenM (Token a, Token b)
untupleDepGenM t = do
  t' <- toToken t
  a <- untupleFstDepGenM t'
  b <- untupleSndDepGenM t'
  pure (a, b)

unzipFstDepGenM :: (TokenableTo [(a, b)] t, Typeable a, Show a, Typeable b, Show b) => t -> DepGenM (Token [a])
unzipFstDepGenM t = do
  t' <- toToken t
  case t' of
    ZipToken a _ ->
      pure a
    _ ->
      genDependency (makeDependency t' UnzipFst)

unzipSndDepGenM :: (TokenableTo [(a, b)] t, Typeable a, Show a, Typeable b, Show b) => t -> DepGenM (Token [b])
unzipSndDepGenM t = do
  t' <- toToken t
  case t' of
    ZipToken _ b ->
      pure b
    _ ->
      genDependency (makeDependency t' UnzipSnd)

unzipDepGenM :: (TokenableTo [(a, b)] t, Typeable a, Show a, Typeable b, Show b) => t -> DepGenM (Token [a], Token [b])
unzipDepGenM t = do
  t' <- toToken t
  a <- unzipFstDepGenM t'
  b <- unzipSndDepGenM t'
  pure (a, b)