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

module MCMC.Distributions (
                           -- * Methods over distributions
                           HasDensity(..)
                          , productDensity
                          , sampleFrom
                          , fromProposal
                          -- * Standard distributions
                          , uniform
                          , mvUniform
                          , diag
                          , normal
                          , mvNormal
                          , categorical
                          , normalize
                          , categoricalNormed
                          , beta
                          , bern
                          , poisson
                          ) where

import qualified System.Random.MWC as MWC
import qualified System.Random.MWC.Distributions as MWC.D
import Control.Monad
import qualified Data.Packed.Matrix as M
import qualified Numeric.LinearAlgebra.Algorithms as LA
import qualified Numeric.Container as C
import qualified Language.Hakaru.Distribution as HD
import qualified Language.Hakaru.Types as HT
import Data.Maybe
import Data.Ord
import Data.List as L

import MCMC.Types
import MCMC.SemanticEditors

-- | The class of distributions for which there exists a density method.
class HasDensity d a where
    -- | A method that provides the probability density at a point in the distribution.
    density :: d a -> Density a

-- | Compute the product density of the input distributions.
-- An example use is in constructing a target distribution
-- whose density can be expressed as a product of probability densities.
productDensity :: HasDensity d a => [d a] -> Density a
productDensity ds x = product $ map (\d -> density d x) ds

-- | Get the probability density at a point for any target distribution.
instance HasDensity Target a where
    density (viewTarget -> Target d) = d

-- | Get the probability density at a point for any proposal distribution.
instance HasDensity Proposal a where
    density (viewProposal -> Proposal d _) = d

-- | This function can be used to call the sampling method of any proposal distribution.
sampleFrom :: Proposal a -> Sample a
sampleFrom (viewProposal -> Proposal _ s) = s

{-| 
  Convenience function for constructing target distributions from predefined
  or custom proposal distributions. One use case is in testing that the samplers
  in the library correctly simulate standard distributions.
 -}
fromProposal :: Proposal a -> Target a
fromProposal = makeTarget . density

-- Uniform -- 

-- Univariate

-- | Univariate uniform distribution over real numbers. The parameters should 
-- not be equal. 
uniform :: (MWC.Variate a, Real a) => a -> a -> Proposal a
uniform a b
    | b < a = makeUniform b a
    | a < b = makeUniform a b
    | otherwise = error "Wrong parameters for Uniform distribution"

unif1D :: Real a => a -> a -> a -> Double
unif1D a b x
    | x < a = 0
    | x > b = 0
    | otherwise = 1 / realToFrac (b - a)

makeUniform :: (MWC.Variate a, Real a) => a -> a -> Proposal a
makeUniform a b = makeProposal (unif1D a b) (MWC.uniformR (a,b))

-- Multivariate

-- | Multivariate uniform distribution over real numbers.
mvUniform :: (MWC.Variate a, Real a) => [a] -> [a] -> Proposal [a]
mvUniform a b
    | b < a = makeMVUniform b a
    | a < b = makeMVUniform a b
    | otherwise = error "Wrong parameters for multi-variate Uniform distribution"

makeMVUniform :: (MWC.Variate a, Real a) => [a] -> [a] -> Proposal [a]
makeMVUniform a b = 
    let tupF f (p,q,r) = f p q r
        uniD x = product . map (tupF unif1D) $ zip3 a b x
        uniSF g = mapM (flip MWC.uniformR g) $ zip a b
    in makeProposal uniD uniSF
                                         
-- Normal --

-- Univariate

-- | Univariate Gaussian distribution over real numbers.
normal :: Double -> Double -> Proposal Double
normal mean cov = makeProposal dens sf
    where hakaruNormal = HD.normal mean (sqrt cov)
          dens x = exp $ HT.logDensity hakaruNormal (HT.Lebesgue x)
          sf g = liftM HT.fromLebesgue $ HT.distSample hakaruNormal g

-- Multivariate

type CovMatrix = M.Matrix Double
type Mu a = M.Matrix a

mu :: M.Element a => [a] -> Mu a
mu mean = M.fromLists [mean]

{-| 
  Convenience function to create a diagonal matrix from a list representing
  the diagonal. Useful for creating a diagonal covariance matrix for the
  multivariate Gaussian distribution.
 -}
diag :: [Double] -> [[Double]]
diag d = [(nth i) (swapWith e) (replicate (length d) 0) | (i,e) <- zip [1..] d]

-- | Multivariate Gaussian distribution over real numbers.
mvNormal :: [Double] -> [[Double]] -> Proposal [Double]
mvNormal mean cov =
    let covMat = M.fromLists cov
        (muMat, n) = (mu mean, length mean)
    in makeProposal (mvNormalDensity muMat covMat) (mvNormalSF muMat covMat n)

mvNormalDensity :: Mu Double -> CovMatrix -> Density [Double]
mvNormalDensity m cov x = c * exp (-d / 2)
    where (covInv, (lndet, sign)) = LA.invlndet cov
          c1 = (2*pi) ^^ (length x)
          c = 1 / (sqrt $ sign * (exp lndet) * c1)
          xm = C.sub (M.fromLists [x]) m
          prod = xm C.<> covInv C.<> (M.trans xm)
          d = (M.@@>) prod (0,0)

mvNormalSF :: Mu Double -> CovMatrix -> Int -> Sample [Double]
mvNormalSF m cov n g = do
      z <- replicateM n (MWC.D.standard g)
      let zt = M.trans $ M.fromLists [z]
          a = LA.chol cov
      return . head . M.toLists $ C.add m $ C.trans $ a C.<> zt

-- Categorical --

-- | Categorical distribution over instances of the Eq typeclass.
-- The input argument is a list of category-proportion pairs. 
-- 
-- The input proportions represent relative weights and are not 
-- required to be normalized.

-- Look at 'MCMC.Combinators.mixProposals' and 'MCMC.Combinators.mixSteps' for 
-- examples of using categorical over elements outside of the 'Eq' typeclass.
categorical :: Eq a => [(a, Double)] -> Proposal a
categorical = categoricalNormed . normalize

-- | Normalize the weights in a list of category-weight pairs.
normalize :: [(a, Double)] -> [(a, Double)]
normalize catProbs = map norm catProbs
    where norm = second (flip (/) s)
          s = foldl (+) 0 $ (snd.unzip) catProbs

-- | Assume that the weights are already normalized. This is useful
-- as an optimized version of @categorical@.
categoricalNormed :: Eq a => [(a,Double)] -> Proposal a
categoricalNormed catProbs =
    -- CHECK: Should fromMaybe default to "error" instead of 0?
    let dens a = snd $ fromMaybe (a,0) $ find ((==)a.fst) catProbs
    in makeProposal dens (categoricalSF catProbs)

categoricalSF :: [(a, Double)] -> Sample a
categoricalSF catProbs g = do
  u <- sampleFrom (uniform 0 1) g
  let (cats, probs) = unzip catProbs
      catsCDF = zip cats $ init $ scanl (+) 0 probs
  return $ fst $ maximumBy (comparing snd) $ filter ((u>=).snd) catsCDF

-- Beta --

-- | Beta distribution over real numbers. Requires non-negative arguments.
beta :: Double -> Double -> Proposal Double
beta a b = makeProposal dens betaSF
    where hakaruBeta = HD.beta a b
          dens x = exp $ HT.logDensity hakaruBeta (HT.Lebesgue x)
          betaSF g = liftM HT.fromLebesgue $ HT.distSample hakaruBeta g

-- Bern --

-- | Bernoulli distribution
bern :: Double -> Proposal Bool
bern p = makeProposal dens bernSF
    where hakaruBern = HD.bern p
          dens x = exp $ HT.logDensity hakaruBern (HT.Discrete x)
          bernSF g = liftM HT.fromDiscrete $ HT.distSample hakaruBern g

-- Poisson --

-- | Univariate Poisson distribution. Requires non-negative argument.
poisson :: Double -> Proposal Int
poisson lambda = makeProposal (exp . HT.logDensity d . HT.Discrete)
                   (fmap HT.fromDiscrete . HT.distSample d)
  where d = HD.poisson lambda