{-# LANGUAGE GADTs, TypeFamilies, TypeOperators, CPP, FlexibleContexts, FlexibleInstances, ScopedTypeVariables, MultiParamTypeClasses, UndecidableInstances #-}
{-# OPTIONS_GHC -fenable-rewrite-rules #-}
----------------------------------------------------------------------
-- |
-- Module      :  Control.Monad.Representable.Trie
-- Copyright   :  (c) Edward Kmett 2011
--                (c) Conal Elliott 2008
-- License     :  BSD3
-- 
-- Maintainer  :  ekmett@gmail.com
-- Stability   :  experimental
-- 
----------------------------------------------------------------------

module Control.Monad.Reader.Trie
  ( 
  -- * Representations of polynomial functors
    HasTrie(..)
  -- * A Trie-based Reader monad transformer
  , ReaderTrieT(..)
  -- * Memoizing functions
  , mup, memo, memo2, memo3
  , inTrie, inTrie2, inTrie3
  -- * Workarounds for current GHC limitations
  , (:=)(..)
  , trie, untrie
  , coerceKey, uncoerceKey
  ) where

import Control.Applicative
import Control.Arrow (first,(&&&))
import Control.Comonad
import Control.Monad.Trans.Class
import Control.Monad.Trans.Reader
import Control.Monad.Trans.Identity
import Control.Monad.IO.Class
import Control.Monad.Reader.Class
import Control.Monad.Representable
import Control.Monad.Writer.Class as Writer
import Data.Bits
import Data.Distributive
import Data.Int
import Data.Function (on)
import Data.Functor.Bind
import Data.Functor.Identity
import Data.Foldable
import Data.Key
import Data.Monoid
import Data.Traversable
import Data.Semigroup
import Data.Semigroup.Foldable
import Data.Semigroup.Traversable
import Data.Word
import Prelude hiding (lookup)

data a := b where Refl :: a := a

-- class (TraversableWithKey1 (Trie a), Representable (Trie a), Key (Trie a) ~ a) => HasTrie a where
class (TraversableWithKey1 (Trie a), Representable (Trie a)) => HasTrie a where
  type Trie a :: * -> *
  -- | Ideally we would have the constraint @Key (Trie a) ~ a@ as a class constraint. 
  -- We are forced to approximate this using an explicit equality witness until GHC implements this feature.
  keyRefl :: a := Key (Trie a)

coerceKey :: HasTrie a => a -> Key (Trie a)
coerceKey = go keyRefl where
  go :: HasTrie a => (a := Key (Trie a)) -> a -> Key (Trie a)
  go Refl = id

uncoerceKey :: HasTrie a => Key (Trie a) -> a
uncoerceKey = go keyRefl where
  go :: HasTrie a => (a := Key (Trie a)) -> Key (Trie a) -> a
  go Refl = id

instance HasTrie () where
  type Trie () = Identity
  keyRefl = Refl

-- Matt Hellige's notation for @argument f . result g@.
-- <http://matt.immute.net/content/pointless-fun>
(~>) :: (a' -> a) -> (b -> b') -> (a -> b) -> a' -> b'
g ~> f = (f .) . (. g)

untrie :: HasTrie t => Trie t a -> t -> a
untrie = go keyRefl where
  go :: HasTrie t => (t := Key (Trie t)) -> Trie t a -> t -> a
  go Refl = index

trie :: HasTrie t => (t -> a) -> Trie t a
trie = go keyRefl where
  go :: HasTrie t => (t := Key (Trie t)) -> (t -> a) -> Trie t a
  go Refl = tabulate

memo :: HasTrie t => (t -> a) -> t -> a
memo = untrie . trie

-- | Lift a memoizer to work with one more argument.
mup :: HasTrie t => (b -> c) -> (t -> b) -> t -> c
mup mem f = memo (mem . f)

-- | Memoize a binary function, on its first argument and then on its
-- second.  Take care to exploit any partial evaluation.
memo2 :: (HasTrie s, HasTrie t) => (s -> t -> a) -> s -> t -> a
memo2 = mup memo

-- | Memoize a ternary function on successive arguments.  Take care to
-- exploit any partial evaluation.
memo3 :: (HasTrie r, HasTrie s, HasTrie t) => (r -> s -> t -> a) -> r -> s -> t -> a
memo3 = mup memo2

-- | Apply a unary function inside of a tabulate
inTrie 
  :: (HasTrie a, HasTrie c) 
  => ((a -> b) -> c -> d)
  -> Trie a b -> Trie c d
inTrie = untrie ~> trie

-- | Apply a binary function inside of a tabulate
inTrie2 
  :: (HasTrie a, HasTrie c, HasTrie e) 
  => ((a -> b) -> (c -> d) -> e -> f)
  -> Trie a b -> Trie c d -> Trie e f
inTrie2 = untrie ~> inTrie

-- | Apply a ternary function inside of a tabulate
inTrie3 
  :: (HasTrie a, HasTrie c, HasTrie e, HasTrie g) 
  => ((a -> b) -> (c -> d) -> (e -> f) -> g -> h)
  -> Trie a b -> Trie c d -> Trie e f -> Trie g h
inTrie3 = untrie ~> inTrie2

instance (HasTrie a, HasTrie b) => HasTrie (a,b) where
  type Trie (a,b) = RepT (Trie a) (Trie b)
  keyRefl = go keyRefl keyRefl where
    go :: (a := Key (Trie a)) -> (b := Key (Trie b)) -> (a, b) := Key (Trie (a,b))
    go Refl Refl = Refl

type instance Key (ReaderTrieT a m) = (a, Key m)

newtype ReaderTrieT a m b = ReaderTrieT { runReaderTrieT :: Trie a (m b) } 

instance (HasTrie a, Functor m) => Functor (ReaderTrieT a m) where
  fmap f = ReaderTrieT . fmap (fmap f) . runReaderTrieT

instance (HasTrie a, Apply m) => Apply (ReaderTrieT a m) where
  ReaderTrieT ff <.> ReaderTrieT fa = ReaderTrieT ((<.>) <$> ff <.> fa)

instance (HasTrie a, Applicative m) => Applicative (ReaderTrieT a m) where
  pure = ReaderTrieT . pure . pure 
  ReaderTrieT ff <*> ReaderTrieT fa = ReaderTrieT ((<*>) <$> ff <*> fa)

instance (HasTrie a, Bind m) => Bind (ReaderTrieT a m) where
  ReaderTrieT fm >>- f = ReaderTrieT $ tabulate (\a -> index fm a >>- flip index a . runReaderTrieT . f)

instance (HasTrie a, Monad m) => Monad (ReaderTrieT a m) where
  return = ReaderTrieT . pure . return
  ReaderTrieT fm >>= f = ReaderTrieT $ tabulate (\a -> index fm a >>= flip index a . runReaderTrieT . f)

instance (HasTrie a, Monad m) => MonadReader a (ReaderTrieT a m) where 
  ask = ReaderTrieT (trie return)
  local f (ReaderTrieT fm) = ReaderTrieT (tabulate (index fm . coerceKey . f . uncoerceKey))

instance HasTrie a => MonadTrans (ReaderTrieT a) where
  lift = ReaderTrieT . pure 

instance (HasTrie a, Distributive m) => Distributive (ReaderTrieT a m) where
  distribute = ReaderTrieT . fmap distribute . collect runReaderTrieT

instance (HasTrie a, Keyed m) => Keyed (ReaderTrieT a m) where
  mapWithKey f = ReaderTrieT . mapWithKey (\k -> mapWithKey (f . (,) (uncoerceKey k))) . runReaderTrieT

instance (HasTrie a, Index m) => Index (ReaderTrieT a m) where
  index = uncurry . fmap index . untrie . runReaderTrieT

instance (HasTrie a, Lookup (Trie a), Lookup m) => Lookup (ReaderTrieT a m) where
  lookup (k,k') (ReaderTrieT fm) = lookup (coerceKey k) fm >>= lookup k'

instance (HasTrie a, Representable m) => Representable (ReaderTrieT a m) where
  tabulate = ReaderTrieT . trie . fmap tabulate . curry
  
instance (HasTrie a, Foldable m) => Foldable (ReaderTrieT a m) where
  foldMap f = foldMap (foldMap f) . runReaderTrieT

instance (HasTrie a, Foldable1 m) => Foldable1 (ReaderTrieT a m) where
  foldMap1 f = foldMap1 (foldMap1 f) . runReaderTrieT

instance (HasTrie a, FoldableWithKey m) => FoldableWithKey (ReaderTrieT a m) where
  foldMapWithKey f = foldMapWithKey (\k -> foldMapWithKey (f . (,) (uncoerceKey k))) . runReaderTrieT

instance (HasTrie a, FoldableWithKey1 m) => FoldableWithKey1 (ReaderTrieT a m) where
  foldMapWithKey1 f = foldMapWithKey1 (\k -> foldMapWithKey1 (f . (,) (uncoerceKey k))) . runReaderTrieT 

instance (HasTrie a, Traversable m) => Traversable (ReaderTrieT a m) where
  traverse f = fmap ReaderTrieT . traverse (traverse f) . runReaderTrieT

instance (HasTrie a, Traversable1 m) => Traversable1 (ReaderTrieT a m) where
  traverse1 f = fmap ReaderTrieT . traverse1 (traverse1 f) . runReaderTrieT

instance (HasTrie a, TraversableWithKey m) => TraversableWithKey (ReaderTrieT a m) where
  traverseWithKey f = fmap ReaderTrieT . traverseWithKey (\k -> traverseWithKey (f . (,) (uncoerceKey k))) . runReaderTrieT

instance (HasTrie a, TraversableWithKey1 m) => TraversableWithKey1 (ReaderTrieT a m) where
  traverseWithKey1 f = fmap ReaderTrieT . traverseWithKey1 (\k -> traverseWithKey1 (f . (,) (uncoerceKey k))) . runReaderTrieT

instance (HasTrie a, Representable m, Semigroup a, Semigroup (Key m)) => Extend (ReaderTrieT a m) where
  extend = extendRep
  duplicate = duplicateRep

instance (HasTrie a, Representable m, Semigroup a, Semigroup (Key m), Monoid a, Monoid (Key m)) => Comonad (ReaderTrieT a m) where
  extract = extractRep

instance (HasTrie a, MonadIO m) => MonadIO (ReaderTrieT a m) where
  liftIO = lift . liftIO 

instance (HasTrie a, MonadWriter w m) => MonadWriter w (ReaderTrieT a m) where
  tell = lift . tell
  listen = ReaderTrieT . tabulate . fmap Writer.listen . index . runReaderTrieT
  pass = ReaderTrieT . tabulate . fmap Writer.pass . index . runReaderTrieT