-- |
-- Module      : Data.Rhythm.Markov
-- Copyright   : (c) Eric Bailey, 2025
--
-- License     : MIT
-- Maintainer  : eric@ericb.me
-- Stability   : experimental
-- Portability : POSIX
--
-- Generating random numbers using a Markov chain.
module Data.Rhythm.Markov
  ( TransitionMatrix (..),
    transitionMatrix,
    markovGen,
  )
where

import Control.Monad.IO.Class (MonadIO)
import Data.Functor ((<&>))
import Data.Vector (Vector)
import Data.Vector qualified as V
import System.Random (randomIO)
import Text.Trifecta (Parser, count, decimal, double, newline)

data TransitionMatrix = TransitionMatrix
  { TransitionMatrix -> Int
size :: Int,
    TransitionMatrix -> Vector (Vector Double)
unTransitionMatrix :: Vector (Vector Double)
  }
  deriving (TransitionMatrix -> TransitionMatrix -> Bool
(TransitionMatrix -> TransitionMatrix -> Bool)
-> (TransitionMatrix -> TransitionMatrix -> Bool)
-> Eq TransitionMatrix
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: TransitionMatrix -> TransitionMatrix -> Bool
== :: TransitionMatrix -> TransitionMatrix -> Bool
$c/= :: TransitionMatrix -> TransitionMatrix -> Bool
/= :: TransitionMatrix -> TransitionMatrix -> Bool
Eq, Int -> TransitionMatrix -> ShowS
[TransitionMatrix] -> ShowS
TransitionMatrix -> String
(Int -> TransitionMatrix -> ShowS)
-> (TransitionMatrix -> String)
-> ([TransitionMatrix] -> ShowS)
-> Show TransitionMatrix
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> TransitionMatrix -> ShowS
showsPrec :: Int -> TransitionMatrix -> ShowS
$cshow :: TransitionMatrix -> String
show :: TransitionMatrix -> String
$cshowList :: [TransitionMatrix] -> ShowS
showList :: [TransitionMatrix] -> ShowS
Show)

transitionMatrix :: Parser (Vector (Vector Double))
transitionMatrix :: Parser (Vector (Vector Double))
transitionMatrix =
  do
    Int
n <- Integer -> Int
forall a. Num a => Integer -> a
fromInteger (Integer -> Int) -> Parser Integer -> Parser Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Parser Integer
forall (m :: * -> *). TokenParsing m => m Integer
decimal Parser Int -> Parser Char -> Parser Int
forall a b. Parser a -> Parser b -> Parser a
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f a
<* Parser Char
forall (m :: * -> *). CharParsing m => m Char
newline
    [Vector Double] -> Vector (Vector Double)
forall a. [a] -> Vector a
V.fromList ([Vector Double] -> Vector (Vector Double))
-> Parser [Vector Double] -> Parser (Vector (Vector Double))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Parser (Vector Double) -> Parser [Vector Double]
forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
count Int
n ([Double] -> Vector Double
forall a. [a] -> Vector a
V.fromList ([Double] -> Vector Double)
-> Parser [Double] -> Parser (Vector Double)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Parser Double -> Parser [Double]
forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
count Int
n Parser Double
forall (m :: * -> *). TokenParsing m => m Double
double)

-- | Given how many numbers to create, the initial state, and a transition
-- matrix, generate random numbers using a Markov chain.
markovGen :: (MonadIO m) => Int -> Int -> Vector (Vector Double) -> m (Vector Int)
markovGen :: forall (m :: * -> *).
MonadIO m =>
Int -> Int -> Vector (Vector Double) -> m (Vector Int)
markovGen Int
0 Int
_ Vector (Vector Double)
_ = Vector Int -> m (Vector Int)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Vector Int
forall a. Vector a
V.empty
markovGen Int
n Int
s Vector (Vector Double)
m = Int -> Vector Int -> Vector Int
forall a. a -> Vector a -> Vector a
V.cons Int
s (Vector Int -> Vector Int) -> m (Vector Int) -> m (Vector Int)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((Int, Vector Double) -> m (Maybe (Int, (Int, Vector Double))))
-> (Int, Vector Double) -> m (Vector Int)
forall (m :: * -> *) b a.
Monad m =>
(b -> m (Maybe (a, b))) -> b -> m (Vector a)
V.unfoldrM (Int, Vector Double) -> m (Maybe (Int, (Int, Vector Double)))
forall {f :: * -> *} {a}.
(Random a, MonadIO f, Ord a, Num a) =>
(Int, Vector a) -> f (Maybe (Int, (Int, Vector Double)))
go (Int
1, Vector (Vector Double)
m Vector (Vector Double) -> Int -> Vector Double
forall a. Vector a -> Int -> a
V.! Int
s)
  where
    go :: (Int, Vector a) -> f (Maybe (Int, (Int, Vector Double)))
go (Int
k, Vector a
row)
      | Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
n =
          f a
forall a (m :: * -> *). (Random a, MonadIO m) => m a
randomIO f a
-> (a -> Maybe (Int, (Int, Vector Double)))
-> f (Maybe (Int, (Int, Vector Double)))
forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> \a
p ->
            (a -> Bool) -> Vector a -> Maybe Int
forall a. (a -> Bool) -> Vector a -> Maybe Int
V.findIndex (a -> a -> Bool
forall a. Ord a => a -> a -> Bool
> a
p) ((a -> a -> a) -> Vector a -> Vector a
forall a. (a -> a -> a) -> Vector a -> Vector a
V.scanl1 a -> a -> a
forall a. Num a => a -> a -> a
(+) Vector a
row) Maybe Int
-> (Int -> (Int, (Int, Vector Double)))
-> Maybe (Int, (Int, Vector Double))
forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> \Int
i ->
              (Int
i, (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1, Vector (Vector Double)
m Vector (Vector Double) -> Int -> Vector Double
forall a. Vector a -> Int -> a
V.! Int
i))
      | Bool
otherwise = Maybe (Int, (Int, Vector Double))
-> f (Maybe (Int, (Int, Vector Double)))
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe (Int, (Int, Vector Double))
forall a. Maybe a
Nothing