{-# LANGUAGE GeneralizedNewtypeDeriving, FlexibleInstances #-}

module SMC where

import PPL ( Sample(sample), Infer(score, observe, logscore), Model, runModel, sampleWeighted, WeightedT )
import Randomness ( Random, Weighted, Weight, parallel )
import Control.Monad ( ap, replicateM, zipWithM )
import Data.Either ( partitionEithers )
import Particle ( ParticleT(..), reweight )
import Control.Monad.Trans.Class (lift)
import Data.Monoid (Product(..))
import Control.Monad.Trans.Writer (WriterT, tell, runWriterT)
import Numeric.Log (Log(..))
import Distribution ( Dist (samplefrom), logpdf, Uniform (Uniform) )
import Statistics.Resampling (Bootstrap(resamples))
import Control.Monad.ST (ST)
import Control.Parallel (pseq)
import Data.List ( partition, sortOn, groupBy )
import Debug.Trace ( trace )
import Control.Monad.RWS (gets)


type Address = Int

instance (Monoid w, Sample m) => Sample (ParticleT l w m) where
  sample = Compute . fmap return . sample

instance Sample m => Infer (ParticleT l (Product Weight) m) where
  score w = Done (Product $ Exp $ log w) ()
  logscore w = Done (Product w) ()
  observe d x = Done (Product $ logpdf d x) ()

-- SMC with seperate score and resample statements
type SMCModel = ParticleT () (Product Weight) Random
type AddressSMC = ParticleT Address (Product Weight) Random

avg :: [Weighted Double] -> Double
avg xs = let total = sum (map snd xs) in
         exp $ ln $ sum [Exp (log a) * w / total | (a, w) <- xs]

stats :: Functor f => f [Weighted (Double, Double) ] -> f Double
stats = fmap (avg . map (\((a,b),c) -> (a,c)))

weight :: ParticleT l (Product Weight) m a -> Weight
weight (Done w _) = getProduct w
weight (Resample l w _) = getProduct w
weight (Compute _) = error "Compute has no weight"

-- A stochastic monad is one that can perform some sort of stochastic compute step
class Monad m => Stochastic m  where
  compute :: m a -> Random (m a)

instance Stochastic (ParticleT l (Product Weight) Random) where
  -- run a computation until it reaches a resample or is done
  compute (Compute ma) = ma >>= compute
  compute p = return p

class Stochastic m => Resampleable m where
  resample :: m ()

instance Resampleable SMCModel where
  resample = Resample () mempty $ return ()

class Infer m => SMC m where
  resampleSMC :: Int -> [m a] -> Random [Weighted a]
  normalisingConstant :: Int -> [m a] -> Random Weight

class Infer m => SampleWeighted m where
  doResample :: [m a] -> Random [m a]

-- SMC where every resample is explicit
newtype ExplicitResample a = ExplicitResample { unexplicit :: SMCModel a } deriving (Functor, Applicative, Monad, Sample, Resampleable, Infer, SMC)

-- SMC where every score triggers a resample
newtype AlwaysResample a = AlwaysResample { unalways :: SMCModel a } deriving (Functor, Applicative, Monad, Sample, Resampleable, SMC)

-- SMC where weight zero particles are resampled
newtype AliveFilter a = AliveFilter { unalive :: SMCModel a } deriving (Functor, Applicative, Monad, Sample, Resampleable, SMC)

-- SMC where subsets of particles are resampled based on addresses
newtype AlphaSMC a = AlphaSMC { unalpha :: AddressSMC a } deriving (Functor, Applicative, Monad, Sample, Infer)

-- SMC where subsets of particles are resampled based on addresses
newtype Adaptive a = Adaptive { unadapt :: SMCModel a } deriving (Functor, Applicative, Monad, Sample, Resampleable, Show)

-- SMC with Startified Resampling
newtype Stratified a = Stratified { unstrat :: SMCModel a } deriving (Functor, Applicative, Monad, Sample, Resampleable)

-- SMC with Subset Resampling
newtype Labelled a = Labelled { unlab :: AddressSMC a } deriving (Functor, Applicative, Monad, Sample, Infer)

-- these all have implicit resamples:
instance Infer AliveFilter where
  score w = AliveFilter (score w >> resample)
  observe d x = AliveFilter (observe d x >> resample)

instance Infer Adaptive where
  score w = Adaptive (score w >> resample)
  observe d x = Adaptive (observe d x >> resample)

instance Infer AlwaysResample where
  score w = AlwaysResample (score w >> resample)
  observe d x = AlwaysResample (observe d x >> resample)

instance Infer Stratified where
  score w = Stratified (score w >> resample)
  observe d x = Stratified (observe d x >> resample)

-- Generalized Newtype Deriving isn't quite enough to do these 
instance Stochastic ExplicitResample where
  compute (ExplicitResample a) = ExplicitResample <$> compute a

instance Stochastic Labelled where
  compute (Labelled a) = Labelled <$> compute a

instance Stochastic AlwaysResample where
  compute (AlwaysResample a) = AlwaysResample <$> compute a

instance Stochastic Adaptive where
  compute (Adaptive a) = Adaptive <$> compute a

instance Stochastic Stratified where
  compute (Stratified a) = Stratified <$> compute a

resampleLab :: [Labelled a] -> Random [Labelled a]
resampleLab ms = do
    -- do all the stochastic compuations so everything
    -- is either done or at a resample statement 
    ms' <- parallel $ map (compute . unlab) ms
    -- check if all particles are done...
    if all isDone ms'
      -- ... in which case return 
      then return $ map Labelled ms'
      -- otherwise we do resampling
      else do
        -- we need to prune the resample we dealt with and reset weights
        resampled <- labResample $ map getweight ms'
        -- recursive call to continue inference
        resampleLab $ map Labelled resampled

labResample :: [Weighted (AddressSMC a)] -> Random [AddressSMC a]
labResample ms = let grouped = groupBy (\(x,_) (y,_) -> address x == address y) ms in 
    fmap concat $ parallel $ map stratResample grouped

resampleStrat :: [Stratified a] -> Random [Stratified a]
resampleStrat ms = do
    -- do all the stochastic compuations so everything
    -- is either done or at a resample statement 
    ms' <- parallel $ map (compute . unstrat) ms
    -- check if all particles are done...
    if all isDone ms'
      -- ... in which case return 
      then return $ map Stratified ms'
      -- otherwise we do resampling
      else do
        -- we need to prune the resample we dealt with and reset weights
        resampled <- stratResample $ map getweight ms'
        -- recursive call to continue inference
        resampleStrat $ map Stratified resampled

getSamples :: [Weight] -> [Weighted a] -> [a]
getSamples [] _ = []
getSamples (r:rs) values = let (y:ys) = dropWhile (\(x,w) -> w < r) values
                           in fst y : getSamples rs (y:ys)

residual :: [Weighted a] -> ([a], [Weighted a])
residual [] = ([], [])
residual ((a, w) : rest) = let (xs, wxs) = residual rest
                           in let w' = floor w
                           in (replicate w' a ++ xs, (a, w - fromIntegral w') : wxs)


stratResample :: Functor m => [Weighted (ParticleT l (Product Weight) m a)] -> Random [ParticleT l (Product Weight) m a]
stratResample xs = do
    -- get sum of weights
    let weight = sum $ map snd xs
    let n = length xs
    let avgweight = weight / Exp (log $ fromIntegral n)
    -- residuals
    let (guaranteed, rest) = residual [(a, w / avgweight) | (a,w) <- xs]
    let n' = n - length guaranteed
    -- rest
    let weighted = let (as, ws) = unzip rest in zip as (scanl1 (+) ws)
    let avgrest = snd (last weighted) / Exp (log $ fromIntegral n')

    -- list of n uniform[i-1/n,i/n] random values 
    let between a b = do
          r <- samplefrom (Uniform (fromIntegral a) (fromIntegral b))
          return $ avgrest * Exp (log r)

    rs <- zipWithM between [0..n'-1] [1..n']
    return $ map (reweight $ Product avgweight) (getSamples rs weighted ++ guaranteed)
{-
stratResample :: [Weighted (SMCModel a)] -> Random [SMCModel a]
stratResample xs = do
    -- [(a1, w1), (a2, w2), ...] becomes [(a1, w1), (a2, w1+w2), ...]
    let weighted = let (as, ws) = unzip xs in zip as (scanl1 (+) ws)
    -- get sum of weights
    let weight = snd $ last weighted
    let n = length weighted
    let avgweight = weight / Exp (log $ fromIntegral n)

    -- list of n uniform[i-1/n,i/n] random values 
    let between a b = do
          r <- samplefrom (Uniform (fromIntegral a) (fromIntegral b))
          return $ avgweight * Exp (log r)
    rs <- zipWithM between [0..n-1] [1..n]

    -- list of k weighted samples 
    let result = getSamples rs weighted
    return $ map (reweight $ Product avgweight) result
-}

class Infer m => Addressable m where
  resample' :: Address -> m ()

  score' :: Address -> Double -> m ()
  score' a w = score w >> resample' a

  observe' :: Dist d v => Address -> d -> v -> m ()
  observe' a d v = observe d v >> resample' a
  
-- subset SMC specific functions
instance Addressable Labelled where
  resample' a = Labelled $ Resample a mempty $ return ()

-- alpha-SMC specific functions
instance Addressable AlphaSMC where
  resample' a = AlphaSMC $ Resample a mempty $ return ()

-- get address from a Dome or Resample
address :: Bounded l => ParticleT l w m a -> l
address (Resample a _ _) = a
address (Done _ _) = maxBound
address (Compute _ ) = error "compute has no address"

-- helper function used by alpha-SMC
resampleMinSubset :: [AddressSMC a] -> Random [AddressSMC a]
resampleMinSubset ps = do
  let mini = minimum $ map address ps
  let (minima, rest) = partition ((== mini) . address) ps
  let w = sum $ map weight minima
  let n = length minima
  minima' <- sampleWeighted n $ map getweight minima
  let normalized = map (reweight $ Product (w / Exp (log $ fromIntegral n))) minima'
  return (normalized ++ rest)

resampleAlpha :: Int -> [AlphaSMC a] -> Random [AlphaSMC a]
resampleAlpha n ms = do
    -- do all the stochastic compuations so everything
    -- is either done or at a resample statement 
    ms' <- parallel $ map (compute . unalpha) ms
    -- check if all particles are done...
    if all isDone ms'
      -- ... in which case return 
      then return $ map AlphaSMC ms'
      -- otherwise we do resampling
      else do
        -- we need to prune the resample we dealt with and reset weights
        resampled <- resampleMinSubset ms'
        -- recursive call to continue inference
        resampleAlpha n $ map AlphaSMC resampled

resampleESS :: Int -> ([Adaptive a], Weight) -> Random ([Adaptive a], Weight)
resampleESS n (ms, w) = do
    -- do all the stochastic compuations so everything
    -- is either done or at a resample statement 
    ms' <- parallel $ map (compute . unadapt) ms
    let w1 = sum (map weight ms')
    let w2 = sum (map ((^2) . weight) ms')
    let ess = w1 * w1 / w2
    let w' = w * (w1 / Exp (log $ fromIntegral n))

    -- check if all particles are done...
    if all isDone ms'
      -- ... in which case return 
      then return (map Adaptive ms', w')
      -- otherwise we do resampling
      else do
        if ess < Exp (log (1 * fromIntegral n))
        then do
          resampled <- sampleWeighted n (map getweight ms')
          resampleESS n (map Adaptive resampled, w')
        else
          resampleESS n (map (Adaptive . ignoreResample) ms', w)

instance SMC AlphaSMC where
  resampleSMC n ms = do
    ms' <- resampleAlpha n ms
    return (map (\(AlphaSMC (Done w a)) -> (a, getProduct w)) ms')
  normalisingConstant n ms = do
    ms' <- resampleAlpha n ms
    let w = sum (map (\(AlphaSMC (Done w a)) -> getProduct w) ms')
    return $ w / Exp (log $ fromIntegral n)

instance SMC Labelled where
  resampleSMC n ms = do
    ms' <- resampleLab ms
    return (map (\(Labelled (Done w a)) -> (a, getProduct w)) ms')
  normalisingConstant n ms = do
    ms' <- resampleLab ms
    return $ getProduct (sum $ map (\(Labelled (Done w a)) -> w) ms') / Exp (log $ fromIntegral $ length ms')

instance SMC Adaptive where
  resampleSMC n ms = do
    (ms', w) <- resampleESS n (ms, Exp 0)
    return (map (\(Adaptive (Done w a)) -> (a, getProduct w)) ms')
  normalisingConstant n ms = snd <$> resampleESS n (ms, Exp 0)

instance SMC Stratified where
  resampleSMC n ms = do
    ms' <- resampleStrat ms
    return (map (\(Stratified (Done w a)) -> (a, getProduct w)) ms')
  normalisingConstant n ms = do
    ms' <- resampleStrat ms
    return $ getProduct (sum $ map (\(Stratified (Done w a)) -> w) ms') / Exp (log $ fromIntegral $ length ms')

ignoreResample :: ParticleT l (Product Weight) Random a -> ParticleT l (Product Weight) Random a
ignoreResample (Resample _ w a) = reweight w a
ignoreResample p@(Done _ _) = p
ignoreResample (Compute _) = error "unexpected compute"



instance Stochastic AliveFilter where
  -- keep trying until we get a nonzero weight
  compute p@(AliveFilter (Compute ma)) = do
    a <- compute (Compute ma)
    if weight a == Exp (log 0) -- note that log 0 = -Infinity
    then compute p
    else return $ AliveFilter a
  compute p = return p

instance Stochastic AlphaSMC where
  -- keep trying until we get a nonzero weight
  compute p@(AlphaSMC (Compute ma)) = do
    a <- compute (Compute ma)
    if weight a == Exp (log 0) -- note that log 0 = -Infinity
    then compute p
    else return $ AlphaSMC a
  compute p = return p

-- get weights, clear internal weights, remove resample
getweight :: ParticleT l (Product Weight) m a -> Weighted (ParticleT l (Product Weight) m a)
getweight (Done w a) = (Done mempty a, getProduct w)
getweight (Resample l w a) = (a, getProduct w)
getweight (Compute a) = error "unexpected compuation"

-- check if particle is done
isDone :: ParticleT l w m a -> Bool
isDone (Done _ _) = True
isDone _ = False

smcmodel :: Int -> ([SMCModel a], Weight) -> Random ([SMCModel a], Weight)
smcmodel n (ms, w) = do
  -- do all the stochastic compuations so everything
    -- is either done or at a resample statement 
    ms' <- parallel $ map compute ms
    let w' = w * sum (map weight ms') / Exp (log $ fromIntegral n)
    -- check if all particles are done...
    if all isDone ms'
      -- ... in which case return 
      then return (ms', w')
      -- otherwise we do resampling
      else do
        -- we need to prune the resample we dealt with and reset weights
        resampled <- sampleWeighted n (map getweight ms')
        -- recursive call to continue inference
        smcmodel n (resampled, w')

instance SMC SMCModel where
  resampleSMC n ms = do
    (ms', w) <- smcmodel n (ms, Exp 0)
    return (map (\(Done w a) -> (a, getProduct w)) ms')
  normalisingConstant n ms = snd <$> smcmodel n (ms, Exp 0)


inferSMC :: SMC m => Int -> m a -> Random [Weighted a]
inferSMC n m = resampleSMC n (replicate n m)

inferConstant :: SMC m => Int -> m a -> Random Weight
inferConstant n m = normalisingConstant n (replicate n m)


