{-#LANGUAGE DeriveDataTypeable, TemplateHaskell #-}

-- | Basic combinators fo building enumerations
-- most users will want to use the type class 
-- based combinators in "Test.Feat.Class" instead. 

module Test.Feat.Enumerate(
  
  Part,
  Index,
  Enumerate(..),
  
  -- ** Combinators for building enumerations
  module Control.Applicative,
  module Data.Monoid,
  pay,
  
  -- ** Memoisation
  mem,
  mempay,
    
  -- *** Polymorphic memoisation
  module Data.Typeable,
  Tag(Source),
  tag,
  tagShare,
  optimise  
  ) where

-- testing-feat
import Control.Monad.TagShare(Sharing, runSharing, share)
import Test.Feat.Internals.Tag(Tag(Source))
-- base
import Control.Applicative
import Control.Monad
import Data.Monoid
import Data.Typeable
import Language.Haskell.TH
-- data-memocombinators
import Data.MemoCombinators



type Part = Int
type Index = Integer

-- | A functional enumeration of type t is a partition of
-- t into finite numbered sets called Parts. The number that
-- identifies each part is called the cost of the values in 
-- that part.
data Enumerate a = Enumerate
   {  
   -- | Computes the cardinality of a given part.
   card      ::  Part -> Index,
   -- | Selects a value from the enumeration
   -- For @select e p i@, @i@ should be less than @card e p@
   select    ::  Part -> Index -> a,
   -- | A self-optimising function. 
   optimal   ::  Sharing Tag (Enumerate a)
   } deriving Typeable     

-- | Only use fmap with bijective functions (e.g. data constructors)
instance Functor Enumerate where 
  fmap f cf = cf
    {select    = \p n -> f (select cf p n)
    , optimal  = liftM (fmap f) (optimal cf) }

-- | mappend = union
instance Monoid (Enumerate a) where
  mempty      = let e = Enumerate  (\p -> 0) 
                                   (\p i -> error "select: empty")
                                   (return e) in e
  mappend     = union

-- | Disjoint union
union :: Enumerate a -> Enumerate a -> Enumerate a
union a b  =  infinite part (liftM2 union (optimal a) (optimal b)) where
  part p   =  finUnion (finite a p) (finite b p)

-- | <*> corresponds to product (as with lists)
instance Applicative Enumerate where
  pure     = singleton
  f <*> a  = fmap (uncurry ($)) (cartesian f a)

-- | The product of two enumerations
cartesian :: Enumerate a -> Enumerate b -> Enumerate (a,b)
cartesian a b = infinite (\p -> foldl1 finUnion
  [finCart (finite a x) (finite b (p-x))| x <- [0..p]])
    (liftM2 cartesian (optimal a) (optimal b))

-- | The definition of @pure@ for the applicaive instance. 
singleton :: a -> Enumerate a
singleton a = let e = Enumerate car sel (return e) in e 
  where  car p    = if p == 0 then 1 else 0
         sel 0 0  = a
         sel _ _  = 
           error "select: index out of bounds"

-- | Increases the cost of all values in an enumeration by one.
pay :: Enumerate a -> Enumerate a
pay sel = Enumerate
    {  card      = \p -> if p <= 0 then 0 else card sel (p-1)
    ,  select    = \p -> select sel (p-1)
    ,  optimal   = liftM pay (optimal sel)
    }

-------------------------------------------------------
-- Memoisation

mem :: Enumerate a -> Enumerate a
mem sel = sel
    { card      = bits (card sel)
    , optimal   = fmap mem (optimal sel)
    }

-- | A conventient combination of memoisation and guarded recursion.
mempay :: Enumerate a -> Enumerate a
mempay = mem . pay
           

-------------------------------------------------------
-- Polymorphic memoisation
tag :: Q Exp -- :: Tag
tag = location >>= makeTag where
   makeTag Loc{  loc_package  = p,    
                 loc_module   = m,    
                 loc_start    = (r,c) }
       = [|Source p m r c|]

tagShare :: Typeable a => Tag -> Enumerate a -> Enumerate a
tagShare t e = e{optimal = share t (optimal e)}

optimise :: Enumerate a -> Enumerate a
optimise e = let e' = runSharing (optimal e) in
  e'{optimal = return e'}   

           
--------------------------------------------------------
-- Operations on finite sets
data Finite a = Finite {fCard :: Index, fSelect :: Index -> a}

finite :: Enumerate a -> Part -> Finite a
finite e p = Finite (card e p) (select e p) 

infinite :: (Part -> Finite a) -> Sharing Tag (Enumerate a) -> Enumerate a
infinite f m = Enumerate (fCard . f) (fSelect . f) m

finUnion :: Finite a -> Finite a -> Finite a
finUnion f1 f2 = Finite car sel where
  car = fCard f1 + fCard f2
  sel i = if i < fCard f1
    then fSelect f1 i
    else fSelect f2 (i-fCard f1)  

finCart :: Finite a -> Finite b -> Finite (a,b)
finCart f1 f2 = Finite car sel where
  car = fCard f1 * fCard f2
  sel i = let (q, r) = (i `quotRem` fCard f2) 
    in (fSelect f1 q, fSelect f2 r)