{-# LANGUAGE DataKinds #-}
{-# LANGUAGE ExplicitNamespaces #-}

-- |
-- Module      : Data.Rhythm.Random
-- Description : Correlated random numbers
-- Copyright   : (c) Eric Bailey, 2025
--
-- License     : MIT
-- Maintainer  : eric@ericb.me
-- Stability   : stable
-- Portability : POSIX
--
-- Generate random numbers with specified correlation.
module Data.Rhythm.Random where

import Control.Monad (foldM, (>=>))
import Control.Monad.IO.Class (MonadIO, liftIO)
import Control.Monad.Loops (unfoldrM)
import Data.Bool (bool)
import Data.Finite (Finite, finite, getFinite)
import Data.Functor ((<&>))
import Data.Maybe (fromJust)
import Data.Proxy (Proxy (..))
import Data.Vector.Sized (Vector)
import Data.Vector.Sized qualified as VS
import GHC.TypeLits (KnownNat, natVal, type (+))
import GHC.TypeNats (SomeNat (..), someNatVal)
import System.Random (randomIO)

-- | See 'randomFinites'.
randomNumbers :: (MonadIO m) => Integer -> Integer -> Integer -> Integer -> m [Integer]
randomNumbers :: forall (m :: * -> *).
MonadIO m =>
Integer -> Integer -> Integer -> Integer -> m [Integer]
randomNumbers Integer
maxNumber Integer
start Integer
correlation Integer
n =
  case (Natural -> SomeNat
someNatVal (Integer -> Natural
forall a b. (Integral a, Num b) => a -> b
fromIntegral Integer
maxNumber Natural -> Natural -> Natural
forall a. Num a => a -> a -> a
+ Natural
1), Natural -> SomeNat
someNatVal (Integer -> Natural
forall a b. (Integral a, Num b) => a -> b
fromIntegral Integer
n Natural -> Natural -> Natural
forall a. Num a => a -> a -> a
- Natural
1)) of
    (SomeNat (Proxy n
_ :: Proxy x), SomeNat (Proxy n
_ :: Proxy y)) ->
      (Finite n -> Integer) -> [Finite n] -> [Integer]
forall a b. (a -> b) -> [a] -> [b]
map Finite n -> Integer
forall (n :: Natural). Finite n -> Integer
getFinite ([Finite n] -> [Integer])
-> (Vector (1 + n) (Finite n) -> [Finite n])
-> Vector (1 + n) (Finite n)
-> [Integer]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector (1 + n) (Finite n) -> [Finite n]
forall (n :: Natural) a. Vector n a -> [a]
VS.toList
        (Vector (1 + n) (Finite n) -> [Integer])
-> m (Vector (1 + n) (Finite n)) -> m [Integer]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (x :: Natural) (y :: Natural) (m :: * -> *).
(KnownNat x, KnownNat y, MonadIO m) =>
Finite x -> Finite x -> m (Vector (1 + y) (Finite x))
randomFinites @x @y (Integer -> Finite n
forall (n :: Natural). KnownNat n => Integer -> Finite n
finite Integer
start) (Integer -> Finite n
forall (n :: Natural). KnownNat n => Integer -> Finite n
finite Integer
correlation)

-- | Generate a vector of random numbers with specified correlation.
randomFinites ::
  forall x y m.
  (KnownNat x, KnownNat y, MonadIO m) =>
  Finite x ->
  Finite x ->
  m (Vector (1 + y) (Finite x))
randomFinites :: forall (x :: Natural) (y :: Natural) (m :: * -> *).
(KnownNat x, KnownNat y, MonadIO m) =>
Finite x -> Finite x -> m (Vector (1 + y) (Finite x))
randomFinites Finite x
startingNumber Finite x
correlation =
  Finite x -> Vector y (Finite x) -> Vector (1 + y) (Finite x)
forall (n :: Natural) a. a -> Vector n a -> Vector (1 + n) a
VS.cons Finite x
startingNumber (Vector y (Finite x) -> Vector (1 + y) (Finite x))
-> ([Integer] -> Vector y (Finite x))
-> [Integer]
-> Vector (1 + y) (Finite x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Integer -> Finite x) -> Vector y Integer -> Vector y (Finite x)
forall a b (n :: Natural). (a -> b) -> Vector n a -> Vector n b
VS.map Integer -> Finite x
forall (n :: Natural). KnownNat n => Integer -> Finite n
finite (Vector y Integer -> Vector y (Finite x))
-> ([Integer] -> Vector y Integer)
-> [Integer]
-> Vector y (Finite x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Maybe (Vector y Integer) -> Vector y Integer
forall a. HasCallStack => Maybe a -> a
fromJust (Maybe (Vector y Integer) -> Vector y Integer)
-> ([Integer] -> Maybe (Vector y Integer))
-> [Integer]
-> Vector y Integer
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Integer] -> Maybe (Vector y Integer)
forall (n :: Natural) a. KnownNat n => [a] -> Maybe (Vector n a)
VS.fromListN
    ([Integer] -> Vector (1 + y) (Finite x))
-> m [Integer] -> m (Vector (1 + y) (Finite x))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((Integer, Integer) -> m (Maybe (Integer, (Integer, Integer))))
-> (Integer, Integer) -> m [Integer]
forall (m :: * -> *) a b.
Monad m =>
(a -> m (Maybe (b, a))) -> a -> m [b]
unfoldrM (Integer, Integer) -> m (Maybe (Integer, (Integer, Integer)))
forall {a}.
(Eq a, Num a) =>
(a, Integer) -> m (Maybe (Integer, (a, Integer)))
go (Proxy y -> Integer
forall (n :: Natural) (proxy :: Natural -> *).
KnownNat n =>
proxy n -> Integer
natVal (forall (t :: Natural). Proxy t
forall {k} (t :: k). Proxy t
Proxy @y), Finite x -> Integer
forall (n :: Natural). Finite n -> Integer
getFinite Finite x
startingNumber)
  where
    go :: (a, Integer) -> m (Maybe (Integer, (a, Integer)))
go (a
0, Integer
_) = Maybe (Integer, (a, Integer)) -> m (Maybe (Integer, (a, Integer)))
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe (Integer, (a, Integer))
forall a. Maybe a
Nothing
    go (a
n, Integer
prev) =
      Integer -> m Integer
applyCorrelation Integer
prev m Integer
-> (Integer -> Maybe (Integer, (a, Integer)))
-> m (Maybe (Integer, (a, Integer)))
forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> \Integer
next ->
        (Integer, (a, Integer)) -> Maybe (Integer, (a, Integer))
forall a. a -> Maybe a
Just (Integer
next, (a
n a -> a -> a
forall a. Num a => a -> a -> a
- a
1, Integer
next))

    applyCorrelation :: Integer -> m Integer
applyCorrelation =
      (Integer -> [Integer] -> m Integer)
-> [Integer] -> Integer -> m Integer
forall a b c. (a -> b -> c) -> b -> a -> c
flip ((Integer -> Integer -> m Integer)
-> Integer -> [Integer] -> m Integer
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM Integer -> Integer -> m Integer
forall {f :: * -> *} {a} {a}.
(Integral a, Integral a, MonadIO f) =>
a -> a -> f a
doDecrement) [Integer
m, Integer
m Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
1 .. Integer
m Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
c Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
1]
        (Integer -> m Integer)
-> (Integer -> m Integer) -> Integer -> m Integer
forall (m :: * -> *) a b c.
Monad m =>
(a -> m b) -> (b -> m c) -> a -> m c
>=> (Integer -> [Integer] -> m Integer)
-> [Integer] -> Integer -> m Integer
forall a b c. (a -> b -> c) -> b -> a -> c
flip ((Integer -> Integer -> m Integer)
-> Integer -> [Integer] -> m Integer
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM Integer -> Integer -> m Integer
forall {f :: * -> *} {p}. MonadIO f => Integer -> p -> f Integer
doIncrement) [Integer
1 .. Integer
c]

    doDecrement :: a -> a -> f a
doDecrement a
prev a
divisor =
      a -> a -> Bool -> a
forall a. a -> a -> Bool -> a
bool a
prev (a -> a -> a
forall a. Ord a => a -> a -> a
max a
0 (a
prev a -> a -> a
forall a. Num a => a -> a -> a
- a
1))
        (Bool -> a) -> (Double -> Bool) -> Double -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< a -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral a
prev Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ a -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral a
divisor)
        (Double -> a) -> f Double -> f a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO Double -> f Double
forall a. IO a -> f a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (forall a (m :: * -> *). (Random a, MonadIO m) => m a
randomIO @Double)

    doIncrement :: Integer -> p -> f Integer
doIncrement Integer
prev p
_ =
      Integer -> Integer -> Bool -> Integer
forall a. a -> a -> Bool -> a
bool Integer
prev (Integer -> Integer -> Integer
forall a. Ord a => a -> a -> a
min Integer
m (Integer
prev Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
1))
        (Bool -> Integer) -> f Bool -> f Integer
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO Bool -> f Bool
forall a. IO a -> f a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO IO Bool
forall a (m :: * -> *). (Random a, MonadIO m) => m a
randomIO

    m :: Integer
m = Proxy x -> Integer
forall (n :: Natural) (proxy :: Natural -> *).
KnownNat n =>
proxy n -> Integer
natVal (forall (t :: Natural). Proxy t
forall {k} (t :: k). Proxy t
Proxy @x) Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
1
    c :: Integer
c = Finite x -> Integer
forall (n :: Natural). Finite n -> Integer
getFinite Finite x
correlation