module Geometry.Algorithms.Sampling (
    -- * Poisson-Disc sampling
      poissonDisc

    -- * Other distributions
    , uniformlyDistributedPoints
    , gaussianDistributedPoints
    , rejection
) where



import           Control.Applicative
import           Control.Monad.Primitive
import           Data.Function
import           Data.Vector                     (Vector)
import qualified Data.Vector                     as V
import           System.Random.MWC
import           System.Random.MWC.Distributions

import Geometry
import Geometry.Algorithms.Sampling.PoissonDisc


-- $setup
-- >>> import qualified System.Random.MWC         as MWC
-- >>> import           Graphics.Rendering.Cairo
-- >>> import           Draw
-- >>> import           Numerics.Functions
-- >>> import           Control.Monad.ST



-- | Generate uniformly distributed points.
--
-- <<docs/haddock/Geometry/Algorithms/Sampling/uniform.svg>>
--
-- === __(image code)__
-- >>> :{
-- haddockRender "Geometry/Algorithms/Sampling/uniform.svg" 300 300 $ \_ -> do
--     let numPoints = 1000
--         bb = shrinkBoundingBox 30 [zero, Vec2 300 300]
--         points = runST $ do
--             gen <- MWC.create
--             uniformlyDistributedPoints gen bb numPoints
--     V.iforM_ points $ \i p -> do
--         setColor (mma i)
--         sketch (Circle p 2)
--         fill
-- :}
-- Generated file: size 263KB, crc32: 0x29e19f0e
uniformlyDistributedPoints
    :: (PrimMonad m, HasBoundingBox boundingBox)
    => Gen (PrimState m) -- ^ RNG from mwc-random. 'create' yields the default (static) RNG.
    -> boundingBox       -- ^ Region to generate points in
    -> Int               -- ^ Number of points
    -> m (Vector Vec2)
uniformlyDistributedPoints :: forall (m :: * -> *) boundingBox.
(PrimMonad m, HasBoundingBox boundingBox) =>
Gen (PrimState m) -> boundingBox -> Int -> m (Vector Vec2)
uniformlyDistributedPoints Gen (PrimState m)
gen boundingBox
bb Int
count = Int -> m Vec2 -> m (Vector Vec2)
forall (m :: * -> *) a. Monad m => Int -> m a -> m (Vector a)
V.replicateM Int
count m Vec2
randomPoint
  where
    BoundingBox (Vec2 Double
xMin Double
yMin) (Vec2 Double
xMax Double
yMax) = boundingBox -> BoundingBox
forall a. HasBoundingBox a => a -> BoundingBox
boundingBox boundingBox
bb
    randomPoint :: m Vec2
randomPoint = (Double -> Double -> Vec2) -> m Double -> m Double -> m Vec2
forall a b c. (a -> b -> c) -> m a -> m b -> m c
forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 Double -> Double -> Vec2
Vec2 ((Double, Double) -> Gen (PrimState m) -> m Double
forall a (m :: * -> *).
(Variate a, PrimMonad m) =>
(a, a) -> Gen (PrimState m) -> m a
forall (m :: * -> *).
PrimMonad m =>
(Double, Double) -> Gen (PrimState m) -> m Double
uniformR (Double
xMin, Double
xMax) Gen (PrimState m)
gen) ((Double, Double) -> Gen (PrimState m) -> m Double
forall a (m :: * -> *).
(Variate a, PrimMonad m) =>
(a, a) -> Gen (PrimState m) -> m a
forall (m :: * -> *).
PrimMonad m =>
(Double, Double) -> Gen (PrimState m) -> m Double
uniformR (Double
yMin, Double
yMax) Gen (PrimState m)
gen)

-- | Generate Gaussian/normal distributed points.
--
-- Note: This is a rejection algorithm which discards samples outside of the
-- 'BoundingBox'. If you choose the covariance much larger than the height or
-- width, performance will deteriorate as more and more points are rejected.
--
-- <<docs/haddock/Geometry/Algorithms/Sampling/gaussian.svg>>
--
-- === __(image code)__
-- >>> :{
-- haddockRender "Geometry/Algorithms/Sampling/gaussian.svg" 300 300 $ \_ -> do
--     let numPoints = 1000
--         bb = shrinkBoundingBox 30 [zero, Vec2 300 300]
--         points = runST $ do
--             gen <- MWC.create
--             gaussianDistributedPoints gen bb (30 *. mempty) numPoints
--     V.iforM_ points $ \i p -> do
--         setColor (mma i)
--         sketch (Circle p 2)
--         fill
-- :}
-- Generated file: size 267KB, crc32: 0xdffc4138
gaussianDistributedPoints
    :: (PrimMonad m, HasBoundingBox boundingBox)
    => Gen (PrimState m) -- ^ RNG from mwc-random. 'create' yields the default (static) RNG.
    -> boundingBox       -- ^ Determines width and height. The center of this is the mean \(\mathbf\mu\).
    -> Mat2              -- ^ Covariance matrix \(\mathbf\Sigma\).
    -> Int               -- ^ Number of points.
    -> m (Vector Vec2)
gaussianDistributedPoints :: forall (m :: * -> *) boundingBox.
(PrimMonad m, HasBoundingBox boundingBox) =>
Gen (PrimState m) -> boundingBox -> Mat2 -> Int -> m (Vector Vec2)
gaussianDistributedPoints Gen (PrimState m)
gen boundingBox
container Mat2
covariance Int
count = Int -> m Vec2 -> m (Vector Vec2)
forall (m :: * -> *) a. Monad m => Int -> m a -> m (Vector a)
V.replicateM Int
count m Vec2
randomPoint
  where
    bb :: BoundingBox
bb = boundingBox -> BoundingBox
forall a. HasBoundingBox a => a -> BoundingBox
boundingBox boundingBox
container
    center :: Vec2
center = BoundingBox -> Vec2
forall a. HasBoundingBox a => a -> Vec2
boundingBoxCenter BoundingBox
bb

    randomPoint :: m Vec2
randomPoint = do
        let t :: Transformation
t = Mat2 -> Vec2 -> Transformation
Transformation Mat2
covariance Vec2
center
        Vec2
vec <- Transformation -> (Double -> Vec2) -> Double -> Vec2
forall geo. Transform geo => Transformation -> geo -> geo
transform Transformation
t ((Double -> Vec2) -> Double -> Vec2)
-> (Double -> Double -> Vec2) -> Double -> Double -> Vec2
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Double -> Double -> Vec2
Vec2 (Double -> Double -> Vec2) -> m Double -> m (Double -> Vec2)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Gen (PrimState m) -> m Double
forall g (m :: * -> *). StatefulGen g m => g -> m Double
standard Gen (PrimState m)
gen m (Double -> Vec2) -> m Double -> m Vec2
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Gen (PrimState m) -> m Double
forall g (m :: * -> *). StatefulGen g m => g -> m Double
standard Gen (PrimState m)
gen
        if Vec2
vec Vec2 -> BoundingBox -> Bool
forall thing bigObject.
(HasBoundingBox thing, HasBoundingBox bigObject) =>
thing -> bigObject -> Bool
`insideBoundingBox` BoundingBox
bb
            then Vec2 -> m Vec2
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Vec2
vec
            else m Vec2
randomPoint

-- | Sample a distribution via rejection sampling: pick a random coordinate, check
-- whether the distribution’s value exceeds that value (accept the sample) or
-- retry.
--
-- The distribution does not need to be normalized, but its values should be
-- reasonably close to 1 for part of the domain, or the algorithm will take a long
-- time to sample the points.
--
-- <<docs/haddock/Geometry/Algorithms/Sampling/rejection.svg>>
--
-- === __(image code)__
-- >>> :{
-- haddockRender "Geometry/Algorithms/Sampling/rejection.svg" 300 300 $ \_ -> do
--     let numPoints = 1000
--         bb = shrinkBoundingBox 30 [zero, Vec2 300 300]
--         distribution p = let r = norm (p -. Vec2 150 150) in gaussianFalloff 75 20 r
--         points = runST $ do
--             gen <- MWC.create
--             rejection gen bb distribution numPoints
--     V.iforM_ points $ \i p -> do
--         setColor (mma i)
--         sketch (Circle p 2)
--         fill
-- :}
-- Generated file: size 263KB, crc32: 0x38511a1
rejection
    :: (PrimMonad m, HasBoundingBox boundingBox)
    => Gen (PrimState m) -- ^ RNG from mwc-random. 'create' yields the default (static) RNG.
    -> boundingBox       -- ^ Area to sample in
    -> (Vec2 -> Double)  -- ^ Distribution \(w(\mathbf p) \in [0\ldots 1]\).
    -> Int               -- ^ Number of points
    -> m (Vector Vec2)
rejection :: forall (m :: * -> *) boundingBox.
(PrimMonad m, HasBoundingBox boundingBox) =>
Gen (PrimState m)
-> boundingBox -> (Vec2 -> Double) -> Int -> m (Vector Vec2)
rejection Gen (PrimState m)
gen boundingBox
bb Vec2 -> Double
distribution Int
numPoints = Int -> (Int -> m Vec2) -> m (Vector Vec2)
forall (m :: * -> *) a.
Monad m =>
Int -> (Int -> m a) -> m (Vector a)
V.generateM Int
numPoints ((Int -> m Vec2) -> m (Vector Vec2))
-> (Int -> m Vec2) -> m (Vector Vec2)
forall a b. (a -> b) -> a -> b
$ \Int
_ -> m Vec2
sampleSinglePoint
  where
    sampleSinglePoint :: m Vec2
sampleSinglePoint = (m Vec2 -> m Vec2) -> m Vec2
forall a. (a -> a) -> a
fix ((m Vec2 -> m Vec2) -> m Vec2) -> (m Vec2 -> m Vec2) -> m Vec2
forall a b. (a -> b) -> a -> b
$ \m Vec2
loop -> do
        Vec2
p <- (Vec2, Vec2) -> Gen (PrimState m) -> m Vec2
forall a g (m :: * -> *).
(UniformRange a, StatefulGen g m) =>
(a, a) -> g -> m a
forall g (m :: * -> *).
StatefulGen g m =>
(Vec2, Vec2) -> g -> m Vec2
uniformRM (Vec2
bbMin, Vec2
bbMax) Gen (PrimState m)
gen
        let pVal :: Double
pVal = Vec2 -> Double
distribution Vec2
p
        Double
threshold <- (Double, Double) -> Gen (PrimState m) -> m Double
forall a g (m :: * -> *).
(UniformRange a, StatefulGen g m) =>
(a, a) -> g -> m a
forall g (m :: * -> *).
StatefulGen g m =>
(Double, Double) -> g -> m Double
uniformRM (Double
0,Double
1) Gen (PrimState m)
gen
        if Double
threshold Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
pVal
            then Vec2 -> m Vec2
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Vec2
p
            else m Vec2
loop

    BoundingBox Vec2
bbMin Vec2
bbMax = boundingBox -> BoundingBox
forall a. HasBoundingBox a => a -> BoundingBox
boundingBox boundingBox
bb