module MCMC.Tests where

import MCMC.Types
import MCMC.Distributions
import MCMC.Kernels
import MCMC.Actions
import MCMC.Combinators
import qualified System.Random.MWC as MWC

-- Bimodal distribution from section 3.1 of
-- "An Introduction to MCMC for Machine Learning" by C. Andrieu et al.
exampleTarget :: Target Double
exampleTarget = 
    makeTarget $ \x -> 0.3 * exp (-0.2*x*x) + 0.7 * exp (-0.2 * (x-10)**2)

gaussianProposal :: Double -> Proposal Double
gaussianProposal x = normal x 10000

exampleMH :: Step Double
exampleMH = metropolisHastings exampleTarget gaussianProposal

mhTest :: IO ()
mhTest = do
  g <- MWC.createSystemRandom
  let a = batchViz vizMH 50
      e = every 100 a
  walk exampleMH 0 (10^6) g e

exampleSA :: Step (StateSA Double)
exampleSA = simulatedAnnealing exampleTarget gaussianProposal

saTest :: IO ()
saTest = do
  g <- MWC.createSystemRandom
  let coolSch = (*) (1 - 1e-3) :: Temp -> Temp
      x0 = (0, 1, coolSch)
      a = batchViz vizSA 50
      e = every 100 a
  walk exampleSA x0 (10^6) g e

gMix :: Target [Double]
gMix = let g1 = mvNormal [0,0] (diag [1,1])
           g2 = mvNormal [5,5] (diag [2,2])
       in mixTargets $ zip [g1, g2] [0.3, 0.7] 

prop1 :: [Double] -> Proposal [Double]
prop1 x = updateNth 1 (\y -> normal y 1) x

prop2 :: [Double] -> Proposal [Double]
prop2 x = updateNth 2 (\y -> normal y 1) x

mhMix :: Step [Double]
mhMix = let mh1 = metropolisHastings gMix prop1
            mh2 = metropolisHastings gMix prop2
        in mixSteps $ zip [mh1, mh2] [0.7, 0.3] 

mixTest :: IO ()
mixTest = do
  g <- MWC.createSystemRandom
  let a = batchPrint printMH 50
  walk mhMix [0,0] (10^6) g a

mhCycle :: Step [Double]
mhCycle = cycleStep metropolisHastings gMix [prop1, prop2]

cycleTest :: IO ()
cycleTest = do
  g <- MWC.createSystemRandom
  let a = batchPrint printMH 50
  walk mhCycle [0,0] (10^6) g a

blockMH :: Step [Double]
blockMH = let target = fromProposal $ mvNormal [0,1,4,7] (diag [2,2,2,2])
              mh4D = metropolisHastings target
              mh1 = mh4D $ updateBlock 1 2 (\y -> mvNormal y (diag [1,1]))
              mh2 = mh4D $ updateBlock 3 3 (\y -> mvNormal y [[1]])
              mh3 = mh4D $ updateNth 4 (\y -> normal y 1)
          in mixSteps $ zip [mh1, mh2, mh3] [0.5, 0.4, 0.7] 

blockTest :: IO ()
blockTest = do
  g <- MWC.createSystemRandom
  let a = batchPrint printMH 50
  walk blockMH [0,0,0,0] (10^6) g a

-- main :: IO ()
-- main = blockTest

gTest :: IO ()
gTest = do 
  print $ density (prop1 [1,2]) [2,4]
  
gaussMix :: Proposal Double
gaussMix = mixProposals $ zip [normal 0 1, normal 5 20, normal 10 0.1] [1..]

betaMix :: Proposal Double
betaMix = mixProposals $ zip [beta 1 1, beta 2 2, beta 3 3] [1..]

mixOfMixes :: Proposal Double
mixOfMixes = mixProposals [(gaussMix, 3), (betaMix, 1)]

lol :: [[Double]] -> Proposal [[Double]]
lol = updateNth 4 (updateNth 2 (\y -> normal y 2))

betaUpdate :: [Double] -> Proposal [Double]
betaUpdate = updateNth 1 (\x -> beta x 3)

gaussUpdate :: [Double] -> Proposal [Double]
gaussUpdate = updateBlock 1 2 (\y -> mvNormal y (diag [1,1]))

mhMix2 :: Step [Double]
mhMix2 = let mh1 = metropolisHastings gMix betaUpdate
             mh2 = metropolisHastings gMix gaussUpdate
         in mixSteps $ zip [mh1, mh2] [0.7, 0.3]

bimodal :: Target [Double]
bimodal = makeTarget dens
    where dens [x,y] = 0.3 * exp (-0.2*x*x) + 0.7 * exp (-0.2 * (y-10)**2)

mhSampler :: Step [Double]
mhSampler = metropolisHastings bimodal betaUpdate

toggle :: (Double, (Bool,String)) -> Proposal (Double, (Bool,String))
toggle = updateSecond (updateFirst (\ b -> bern $ if b then 0 else 1))