{-# LANGUAGE GADTs, MultiParamTypeClasses, KindSignatures, FlexibleInstances,
  ViewPatterns #-}

module MCMC.Kernels ( 
                     -- * Making the random walk
                     walk
                     -- * Transition kernels
                     -- ** Metropolis-Hastings
                    , metropolisHastings
                    , vizMH
                    , printMH
                    -- ** Simulated Annealing
                    , Temp
                    , CoolingSchedule
                    , StateSA
                    , simulatedAnnealing
                    , vizSA
                    , printSA
                    -- ** Gibbs
                    , gibbs
                    ) where

import MCMC.Actions
import MCMC.Combinators
import MCMC.Distributions
import MCMC.Types
import Text.Printf

-- | Execute a random walk and create a Markov chain.
walk :: Step x          -- ^ The stepping style (based on the transition kernel) 
     -> x               -- ^ The starting state
     -> Int             -- ^ The number of steps to take
     -> Rand            -- ^ A PRNG
     -> Action x IO a b -- ^ An action to take at each step in the walk
     -> IO b            -- ^ The action-dependent output at the end of the walk
walk _ _ 0 _ (viewAction -> Action _ f a) = f a
walk step x n r action = do 
  x' <- step r x
  execute action x' >>= walk step x' (n-1) r

-- Metropolis Hastings --

metropolisHastings :: Kernel a a
metropolisHastings t c_p = 
    let mhStep g xi = do
          u <- sampleFrom (uniform 0 1) g
          xstar <- sampleFrom (c_p xi) g
          let accept = min 1 (numer / denom)
              numer = density t xstar * density (c_p xstar) xi
              denom = density t xi * density (c_p xi) xstar
          return $ if u < accept then xstar else xi
    in mhStep

vizMH :: PrintF Double Double
vizMH = id

-- Visualizes only the first dimension
vizMHFirstDim :: PrintF [Double] Double
vizMHFirstDim = map head

printMH :: PrintF [Double] [String]
printMH lls = let p d = printf "%0.3f" d :: String
              in map (map p) lls

-- Simulated Annealing --

type Temp = Double
-- | This is the tempering function used in the simulated annealing process.
type CoolingSchedule = Temp -> Temp
type StateSA a = (a, Temp, CoolingSchedule)

simulatedAnnealing :: Kernel (StateSA a) a
simulatedAnnealing t c_p  = 
    let saStep g (xi,temp,cool) = do
          u <- sampleFrom (uniform 0 1) g
          xstar <- sampleFrom (c_p xi) g
          let accept = min 1 (numer / denom)
              numer = (*) (density (c_p xstar) xi) $ (**) (density t xstar) (1 / temp)
              denom = (*) (density (c_p xi) xstar) $ (**) (density t xi) (1 / temp)
              new_temp = cool temp
          return $ if u < accept then (xstar,new_temp,cool) else (xi,new_temp,cool)
    in saStep

tripleFirst :: (a, b, c) -> a
tripleFirst (a,_,_) = a

myFilter :: [[Double]] -> [[Double]]
myFilter = filter (\x -> x < (repeat 15) && x > (repeat $ -5))

vizSA :: PrintF (StateSA Double) Double
vizSA = map tripleFirst

-- Visualizes only the first dimension
vizSAFirstDim :: PrintF (StateSA [Double]) Double
vizSAFirstDim = vizMHFirstDim . myFilter . map tripleFirst

-- Print the samples without the temp/cooling schedule
printSA :: PrintF (StateSA [Double]) [String]
printSA = printMH . myFilter . map tripleFirst

-- Gibbs --

-- | The full conditional proposals must be specified to this transition kernel.
gibbs :: Target a -> [a -> Proposal a] -> Step a
gibbs = cycleStep alwaysAccept
    where alwaysAccept _ q g x = sampleFrom (q x) g

-- MCMC-EM --

mhEM :: Kernel theta theta
mhEM t c_p = undefined