<commit_before>{-# LANGUAGE
  GADTs
 #-}


module Trace where

import Control.Monad (liftM,liftM2)
import Data.Number.LogFloat

import Primitive
import Base
import Sampler

-- | A random database of stochastic choices made in the program.
data RandomDB a where
    Pure :: a -> RandomDB a
    Node :: Primitive a -> a -> RandomDB a
    -- The arguments to Bind are: trace to this point, trace from this point,
    -- continuation that can be used to generate new trace if the first part changes.
    Bind :: RandomDB b -> RandomDB a -> (b -> Trace a) -> RandomDB a 

-- The three instances below are probably useless.
instance Functor RandomDB where
    fmap f (Pure x) = Pure (f x)
    fmap f d = liftM f d

instance Applicative RandomDB where
    pure = Pure
    (<*>) = liftM2 ($)

instance Monad RandomDB where
    return = pure
    d >>= f = Bind d (f (value d)) (Trace . return . f)

-- | The final value produced by the program.
value :: RandomDB a -> a
value (Pure x) = x
value (Node d x) = x
value (Bind t1 t2 k) = value t2

-- | The number of random choices in the RandomDB.
size :: RandomDB a -> Int
size (Pure _) = 0
size (Node _ _) = 1
size (Bind t1 t2 _) = size t1 + size t2

-- | The product of all densities in the trace.
weight :: RandomDB a -> LogFloat
weight (Pure _) = 1
weight (Node d x) = pdf d x
weight (Bind t1 t2 _) = weight t1 * weight t2

-- | Updates a trace by resampling a randomly selected site.
update :: RandomDB a -> Sampler (RandomDB a)
update (Pure x) = return $ Pure x
update (Node d x) = fmap (Node d) (primitive d)
update (Bind t1 t2 k) = do
  let p = 0.5
  b <- bernoulli p
  if b then 
      do
        t1' <- update t1
        t2' <- runTrace $ k $ value t1'
        return $ Bind t1' t2' k
  else 
      do
        t2' <- update t2
        return $ Bind t1 t2' k

-- | Reuses the random choices from the first computation in the second one.
reuse :: RandomDB a -> RandomDB a -> Trace a
reuse = undefined

newtype Trace a = Trace (Sampler (RandomDB a))
runTrace (Trace d) = d

instance Functor Trace where
    fmap = liftM

instance Applicative Trace where
    pure = return
    (<*>) = liftM2 ($)

instance Monad Trace where
    return = Trace . return . return
    t >>= f = Trace $
        do
          let d = runTrace t
          r1 <- d
          let x = value r1
          r2 <- runTrace (f x)
          return $ Bind r1 r2 f where

instance MonadDist Trace where
    primitive d = Trace $ do
          x <- primitive d
          return $ Node d x
    categorical = error "Can not use Trace on categorical without equality"


mhStep :: RandomDB a -> Sampler (RandomDB a)
mhStep r = do
  r' <- update r
  let ratio = weight r' * fromIntegral (size r') / (weight r * fromIntegral (size r))
  accept <- bernoulli ratio
  return $ if accept then r' else r<commit_msg>Replaced trace with notes on the new design<commit_after>{-# LANGUAGE
  GADTs
 #-}


module Trace where

import Control.Monad (liftM,liftM2)
import Data.Number.LogFloat
import qualified Data.Map.Lazy as Map
import qualified Data.Foldable as Fold

import Primitive
import Base
import Sampler

type DistType = () -- repeat constructors of Primitive
type Label = () -- in the simplest case this is just unit

newtype RandomDB = RandomDB (Map.Map DistType [(Primitive a, a)])

record :: Primitive a -> a -> RandomDB
record = undefined

lookup :: RandomDB -> Primitive a -> Maybe a
lookup = undefined

instance Monoid RandomDB where
  mzero = RandomDB empty
  mplus (RandomDB a) (RandomDB b) = RandomDB $ unionWith (++) a b

instance Monoid LogFloat where
  mempty = 1
  mappend = (*)

weight :: RandomDB -> LogFloat
weight = Fold.foldMap f where
  f xs = Fold.foldMap (uncurry pdf)

update :: MonadDist m => MHKernel m RandomDB
update = undefined


newtype TraceT m a = TraceT {runTraceT :: WriterT RandomDB m a}
                   deriving (Functor, Applicative, Monad, MonadTrans, MonadWriter)

instance MonadDist m => MonadDist (TraceT m) where
  primitive d = do
    x <- lift (primitive d)
    tell (record d x)
    return x


newtype ReuserT m a = ReuserT {runReuserT :: ReaderT RandomDB m a}
                      deriving (Functor, Applicative, Monad, MonadTrans, MonadReader)

instance MonadDist m => MonadDist (ReuserT m) where
  primitive d = do
    rdb <- ask
    case lookup rdb d of Just x -> return x
                         Nohing -> lift (primitive d)
    -- need some sort of update to remove the read value from the database
    -- maybe use StateT instead of ReaderT?
{'fraction_non_alphanumeric': 0.08741103202846975, 'fraction_numerical': 0.008451957295373666, 'mean_word_length': 2.6236905721192585, 'pattern_counts': {'":': 0, '<': 16, '<?xml version=': 0, '>': 33, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '858811', 'n_tokens_mistral': 1468, 'n_tokens_neox': 1413, 'n_words': 624}