{-# LANGUAGE ViewPatterns #-}

module MCMC.Combinators (
                         -- * Target combinators
                         mixTargets
                        -- * Proposal combinators
                        , mixProposals
                        , chooseProposal
                        , mixCondProposals
                        , updateNth
                        , updateBlock
                        , updateFirst
                        , updateSecond
                        -- * Step combinators
                        , mixSteps
                        , chooseStep
                        , cycleStep
                        -- * Action combinators
                        , execute
                        , every
                        ) where

import MCMC.Types
import MCMC.Distributions
import MCMC.SemanticEditors
import Data.Maybe

import Control.Monad

-- Target Mixtures --

-- Assume normalized probabilities
mixDensity :: HasDensity t a => [(t a, Double)] -> Density a
mixDensity catProbs x = foldl1 (+) $ map f catProbs
    where f (t,p) = p*(density t x)

-- | Create a mixture target distribution from a list of 
-- distribution-weight pairs. 
-- 
-- The input distributions can be of any type in the 'HasDensity'
-- class, which includes 'Proposal' and 'Target'.
-- The output is a 'Target' containing a 'Density'.
-- 
-- The input weights are meant to be relative, not
-- required to be normalized.
mixTargets :: HasDensity t a => [(t a, Double)] -> Target a
mixTargets = makeTarget . mixDensity . normalize
       
-- Proposal Mixtures --

-- | Create a mixture proposal distribution from a list of
-- distribution-weight pairs.
-- 
-- The input weights are meant to be relative, not
-- required to be normalized.
mixProposals :: [(Proposal a, Double)] -> Proposal a
mixProposals proposalProbs = 
    let normedPropProbs = normalize proposalProbs
        mixSF g = do
          let (props, probs) = unzip normedPropProbs
              propMap = zip [1..] props
          propKey <- sampleFrom (categoricalNormed (zip [1..] probs)) g
          let proposal = fromMaybe (error "mixProposals: undefined key")
                         (lookup propKey propMap)
          sampleFrom proposal g
    in makeProposal (mixDensity normedPropProbs) mixSF

-- Proposal Choices -- 

-- | Create a uniform mixture of proposals based on the mapping
-- function. The first argument is the total number of proposals
-- in the mixture.
-- 
-- This is a convenience function for cases where there are a large 
-- number of index-based proposals (such as the case where there is a  
-- unique proposal for updating each dimension in a high-dimensional state).
chooseProposal :: Int -> (Int -> Proposal a) -> Proposal a
chooseProposal n f = mixProposals $ [(f i , 1) | i <- [1..n]] -- uniform mix

-- Proposal updaters combinations

-- | Create a mixture of /conditional/ proposal distributions.
-- The result is a mixture of proposal distributions that are all conditioned
-- on the same starting state.
-- 
-- The input weights are meant to be relative, not
-- required to be normalized.
mixCondProposals :: [(a -> Proposal a, Double)] -> a -> Proposal a
mixCondProposals cps a = mixProposals $ map (\(cp,prob) -> (cp a, prob)) cps

getNth :: Int -> [a] -> ([a], a, [a])
getNth n ls = let (front, [e], back) = chopAt n n ls
              in (front, e, back)

-- | Update the nth dimension of a multi-dimensioanl state 
-- represented using a list.
updateNth :: Eq a => Int           -- ^ The dimension index
          -> (a -> Proposal a)     -- ^ A conditional proposal to focus on that dimension
          -> ([a] -> Proposal [a]) -- ^ A conditional proposal that updates only one dimension
updateNth n p x = 
    let den y = let (xBefore, xn, xAfter) = getNth n x
                    (yBefore, yn, yAfter) = getNth n y
                in if length x == length y && xBefore == yBefore
                                           && xAfter == yAfter
                   then density (p xn) yn else 0
        s g = nthM n (\xn -> sampleFrom (p xn) g) x
    in makeProposal den s

-- | Update a contiguous block of dimensions of a multi-dimensional state
-- represented using a list.
updateBlock :: Eq a => Int           -- ^ The start dimension of the block 
            -> Int                   -- ^ The end dimension of the block
            -> ([a] -> Proposal [a]) -- ^ A conditional proposal to focus on the block 
            -> ([a] -> Proposal [a]) -- ^ A conditional proposal that updates only that block
updateBlock n m p x = 
    let den y = let (xBefore, xBlock, xAfter) = chopAt n m x
                    (yBefore, yBlock, yAfter) = chopAt n m y
                in if length x == length y && xBefore == yBefore
                                           && xAfter == yAfter
                   then density (p xBlock) yBlock else 0
        s g = blockM n m (\b -> sampleFrom (p b) g) x
    in makeProposal den s

-- | Update the first element of a multi-dimensional state represented as a tuple.
updateFirst :: Eq b => (a -> Proposal a) -- ^ A conditional proposal to focus on the first element 
            -> ((a,b) -> Proposal (a,b))
updateFirst p (a,b) = makeProposal dens sf
    where dens (x,y) = if y==b then density (p a) x else 0
          sf g = do a' <- sampleFrom (p a) g
                    return (a',b)

-- | Update the second element of a multi-dimensional state represented as a tuple.
updateSecond :: Eq a => (b -> Proposal b) 
             -> ((a,b) -> Proposal (a,b)) -- ^ A conditional proposal to focus on the second element
updateSecond p (a,b) = makeProposal dens sf
    where dens (x,y) = if x==a then density (p b) y else 0
          sf g = do b' <- sampleFrom (p b) g
                    return (a,b')

-- Examples --

-- ex = (second.first.second) not (1,((3,True),2))
-- eg = (cdr.cdr.car) not [False,False,True,False]
-- examp = (nth 4) not [False,False,False,True,False]
-- exampl = (nth 4) (swapWith 4) [1,2,3,5,5]
-- diagex = diag [1,2,3,4]
-- cdrex = cdr (\ls -> map ((+)10) ls) [1,2,3,4,5]
-- cprobs = U.prescanl (+) 0 $ normalize $ U.fromList [0.3, 0.4, 0.6, 0.3, 0.7]
-- gcp = U.maximum $ U.filter ((>=) 0.60) cprobs

-- Kernel Mixtures --

-- | Create a mixture 'Step', representing a categorical distribution over 
-- multiple inference procedures.
-- 
-- The input weights are meant to be relative, not
-- required to be normalized.
mixSteps :: [(Step x, Double)] -> Step x
mixSteps stepProbs = 
    let mixStep g x = do
          let (steps, probs) = unzip stepProbs
              stepMap = zip [1..] steps
          stepKey <- sampleFrom (categorical (zip [1..] probs)) g
          let step = fromMaybe (error "mixSteps: undefined key")
                     (lookup stepKey stepMap)
          step g x
    in mixStep

-- | The analogue of 'chooseProposal', to use a specific 'Step'
-- for each dimension of a multi-dimensional state.
-- 
-- The same 'Kernel' value is used for creating each 'Step' in the mixture.
chooseStep :: Int                      -- ^ The total number of 'Step's in the mixture
           -> (Int -> Target a)        -- ^ The mapping function over targets
           -> (Int -> a -> Proposal a) -- ^ The mapping function over conditional proposals
           -> Kernel x a               -- ^ The transition kernel to use across all 'Step's
           -> Step x                   -- ^ The mixture step
chooseStep n t p k = mixSteps [(k (t i) (p i) , 1) | i <- [1..n]]

-- Kernel Cycles --

-- | Create a cycled transition kernel. This combinator applies the 
-- provided transition kernel to the target, and to the conditional proposals
-- in the order that they appear in the list.
-- Each application generates an intermediate state that is passed to the next 
-- conditional proposal in the list.
cycleStep :: Kernel x a -> Target a -> [a -> Proposal a] -> Step x
cycleStep kernel t cps =
  let steps g = [kernel t cp g | cp <- cps]
      combine comb step = (\iox -> iox >>= comb) . step
      cycled g = foldl combine return (steps g)
  in cycled

-- | Execute the action once, given the current state in the Markov chain.
execute :: Monad m => Action x m a b -> x -> m (Action x m a b)
execute (viewAction -> Action act fin a) x = liftM (makeAction act fin) (act x a)

-- | Execute the action once every @n@ times, where @n@ is the first argument.
every :: Monad m => Int -> Action x m a b -> Action x m (a,Int) b
every n (viewAction -> Action act fin a) = 
  let skip_act x (b,i) = if i == (n-1)
                         then do b' <- act x b
                                 return (b',0)
                         else return (b,i+1)
      skip_fin = fin . fst
  in makeAction skip_act skip_fin (a,0)