{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE ExplicitNamespaces #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE TemplateHaskell #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}

-- |
-- Module      : Data.Rhythm.DeBruijn
-- Description : De Bruijn sequences
-- Copyright   : (c) Eric Bailey, 2024-2025
--
-- License     : MIT
-- Maintainer  : eric@ericb.me
-- Stability   : stable
-- Portability : POSIX
--
-- Constructing de Bruijn sequences using [a greedy
-- algorithm](http://debruijnsequence.org/db/greedy).
--
-- = References
--   - De Bruijn Sequence and Universal Cycle Constructions
--     (2024). http://debruijnsequence.org.
module Data.Rhythm.DeBruijn
  ( -- * Ergonomic
    deBruijnSequence,

    -- * Safe
    deBruijnSequence',

    -- * Determining symbols
    DeBruijnState (..),
    initialDeBruijnState,
  )
where

import Control.Lens (makeLenses, uses, (%=))
import Control.Monad (when)
import Control.Monad.State (evalState)
import Data.FastDigits (undigits)
import Data.Finite (Finite, getFinite)
import Data.Proxy (Proxy (..))
import Data.Set (Set)
import Data.Set qualified as Set
import Data.Vector.Sized (Vector, (//))
import Data.Vector.Sized qualified as VS
import GHC.TypeNats (KnownNat, Nat, SomeNat (..), natVal, someNatVal, type (^))

-- | Generator for a de Bruijn sequence of order @n@ and alphabet size @k@.
data DeBruijnState (k :: Nat) (n :: Nat) = DeBruijnState
  { forall (k :: Natural) (n :: Natural).
DeBruijnState k n -> Vector n (Finite k)
_seed :: Vector n (Finite k),
    forall (k :: Natural) (n :: Natural).
DeBruijnState k n -> Set Integer
_seen :: Set Integer
  }

makeLenses ''DeBruijnState

-- | Seed with \(0^n\).
initialDeBruijnState :: (KnownNat k, KnownNat n) => DeBruijnState k n
initialDeBruijnState :: forall (k :: Natural) (n :: Natural).
(KnownNat k, KnownNat n) =>
DeBruijnState k n
initialDeBruijnState = Vector n (Finite k) -> Set Integer -> DeBruijnState k n
forall (k :: Natural) (n :: Natural).
Vector n (Finite k) -> Set Integer -> DeBruijnState k n
DeBruijnState (Finite k -> Vector n (Finite k)
forall (n :: Natural) a. KnownNat n => a -> Vector n a
VS.replicate Finite k
0) Set Integer
forall a. Set a
Set.empty

-- | Generate the largest de Bruijn sequence of a given order.
--
-- >>> concatMap show $ deBruijnSequence 2 4
-- "1111011001010000"
--
-- >>> concatMap show $ deBruijnSequence 4 3
-- "3332331330322321320312311310302301300222122021121020120011101000"
deBruijnSequence :: Integer -> Integer -> [Integer]
deBruijnSequence :: Integer -> Integer -> [Integer]
deBruijnSequence Integer
k Integer
n =
  case (Natural -> SomeNat
someNatVal (Integer -> Natural
forall a. Num a => Integer -> a
fromInteger Integer
k), Natural -> SomeNat
someNatVal (Integer -> Natural
forall a. Num a => Integer -> a
fromInteger Integer
n)) of
    (SomeNat (Proxy n
_ :: Proxy k), SomeNat (Proxy n
_ :: Proxy n)) ->
      Finite n -> Integer
forall (n :: Natural). Finite n -> Integer
getFinite (Finite n -> Integer) -> [Finite n] -> [Integer]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Vector (n ^ n) (Finite n) -> [Finite n]
forall (n :: Natural) a. Vector n a -> [a]
VS.toList (forall (k :: Natural) (n :: Natural).
(KnownNat k, KnownNat n) =>
Vector (k ^ n) (Finite k)
deBruijnSequence' @k @n)

-- | See 'deBruijnSequence'.
deBruijnSequence' :: forall k n. (KnownNat k, KnownNat n) => Vector (k ^ n) (Finite k)
deBruijnSequence' :: forall (k :: Natural) (n :: Natural).
(KnownNat k, KnownNat n) =>
Vector (k ^ n) (Finite k)
deBruijnSequence' =
  State (DeBruijnState k n) (Vector Vector (k ^ n) (Finite k))
-> DeBruijnState k n -> Vector Vector (k ^ n) (Finite k)
forall s a. State s a -> s -> a
evalState ((Finite (k ^ n) -> StateT (DeBruijnState k n) Identity (Finite k))
-> State (DeBruijnState k n) (Vector Vector (k ^ n) (Finite k))
forall (n :: Natural) (m :: * -> *) a.
(KnownNat n, Monad m) =>
(Finite n -> m a) -> m (Vector n a)
VS.generateM (StateT (DeBruijnState k n) Identity (Finite k)
-> Finite (k ^ n) -> StateT (DeBruijnState k n) Identity (Finite k)
forall a b. a -> b -> a
const (Finite k -> StateT (DeBruijnState k n) Identity (Finite k)
forall {m :: * -> *} {k :: Natural}.
(KnownNat k, MonadState (DeBruijnState k n) m) =>
Finite k -> m (Finite k)
nextSymbol Finite k
forall a. Bounded a => a
maxBound))) DeBruijnState k n
forall (k :: Natural) (n :: Natural).
(KnownNat k, KnownNat n) =>
DeBruijnState k n
initialDeBruijnState
  where
    nextSymbol :: Finite k -> m (Finite k)
nextSymbol Finite k
bit =
      Finite k -> m Bool
forall {m :: * -> *} {k :: Natural} {n :: Natural}.
(KnownNat k, KnownNat n, MonadState (DeBruijnState k n) m) =>
Finite k -> m Bool
try Finite k
bit m Bool -> (Bool -> m (Finite k)) -> m (Finite k)
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
        Bool
True -> Finite k
bit Finite k -> m () -> m (Finite k)
forall a b. a -> m b -> m a
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ ((Vector n (Finite k) -> Identity (Vector n (Finite k)))
-> DeBruijnState k n -> Identity (DeBruijnState k n)
forall (k :: Natural) (n :: Natural) (k :: Natural) (n :: Natural)
       (f :: * -> *).
Functor f =>
(Vector n (Finite k) -> f (Vector n (Finite k)))
-> DeBruijnState k n -> f (DeBruijnState k n)
seed ((Vector n (Finite k) -> Identity (Vector n (Finite k)))
 -> DeBruijnState k n -> Identity (DeBruijnState k n))
-> (Vector n (Finite k) -> Vector n (Finite k)) -> m ()
forall s (m :: * -> *) a b.
MonadState s m =>
ASetter s s a b -> (a -> b) -> m ()
%= forall (n :: Natural) a.
KnownNat n =>
a -> Vector n a -> Vector n a
pushEnd @n Finite k
bit)
        Bool
False -> Finite k -> m (Finite k)
nextSymbol (Finite k
bit Finite k -> Finite k -> Finite k
forall a. Num a => a -> a -> a
- Finite k
1)

    try :: Finite k -> m Bool
try Finite k
bit =
      do
        (Vector n (Finite k) -> Identity (Vector n (Finite k)))
-> DeBruijnState k n -> Identity (DeBruijnState k n)
forall (k :: Natural) (n :: Natural) (k :: Natural) (n :: Natural)
       (f :: * -> *).
Functor f =>
(Vector n (Finite k) -> f (Vector n (Finite k)))
-> DeBruijnState k n -> f (DeBruijnState k n)
seed ((Vector n (Finite k) -> Identity (Vector n (Finite k)))
 -> DeBruijnState k n -> Identity (DeBruijnState k n))
-> (Vector n (Finite k) -> Vector n (Finite k)) -> m ()
forall s (m :: * -> *) a b.
MonadState s m =>
ASetter s s a b -> (a -> b) -> m ()
%= (Vector n (Finite k)
-> [(Finite n, Finite k)] -> Vector n (Finite k)
forall (m :: Natural) a.
Vector m a -> [(Finite m, a)] -> Vector m a
// [(Finite n
forall a. Bounded a => a
maxBound, Finite k
bit)])
        Integer
num <- LensLike' (Const Integer) (DeBruijnState k n) (Vector n (Finite k))
-> (Vector n (Finite k) -> Integer) -> m Integer
forall s (m :: * -> *) r a.
MonadState s m =>
LensLike' (Const r) s a -> (a -> r) -> m r
uses LensLike' (Const Integer) (DeBruijnState k n) (Vector n (Finite k))
forall (k :: Natural) (n :: Natural) (k :: Natural) (n :: Natural)
       (f :: * -> *).
Functor f =>
(Vector n (Finite k) -> f (Vector n (Finite k)))
-> DeBruijnState k n -> f (DeBruijnState k n)
seed (Natural -> [Finite k] -> Integer
forall a b. (Integral a, Integral b) => a -> [b] -> Integer
undigits Natural
k ([Finite k] -> Integer)
-> (Vector n (Finite k) -> [Finite k])
-> Vector n (Finite k)
-> Integer
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector n (Finite k) -> [Finite k]
forall (n :: Natural) a. Vector n a -> [a]
VS.toList)
        Bool
isNew <- LensLike' (Const Bool) (DeBruijnState k n) (Set Integer)
-> (Set Integer -> Bool) -> m Bool
forall s (m :: * -> *) r a.
MonadState s m =>
LensLike' (Const r) s a -> (a -> r) -> m r
uses LensLike' (Const Bool) (DeBruijnState k n) (Set Integer)
forall (k :: Natural) (n :: Natural) (f :: * -> *).
Functor f =>
(Set Integer -> f (Set Integer))
-> DeBruijnState k n -> f (DeBruijnState k n)
seen (Integer -> Set Integer -> Bool
forall a. Ord a => a -> Set a -> Bool
Set.notMember Integer
num)
        Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
isNew (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$
          (Set Integer -> Identity (Set Integer))
-> DeBruijnState k n -> Identity (DeBruijnState k n)
forall (k :: Natural) (n :: Natural) (f :: * -> *).
Functor f =>
(Set Integer -> f (Set Integer))
-> DeBruijnState k n -> f (DeBruijnState k n)
seen ((Set Integer -> Identity (Set Integer))
 -> DeBruijnState k n -> Identity (DeBruijnState k n))
-> (Set Integer -> Set Integer) -> m ()
forall s (m :: * -> *) a b.
MonadState s m =>
ASetter s s a b -> (a -> b) -> m ()
%= Integer -> Set Integer -> Set Integer
forall a. Ord a => a -> Set a -> Set a
Set.insert Integer
num
        Bool -> m Bool
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
isNew

    k :: Natural
k = Proxy k -> Natural
forall (n :: Natural) (proxy :: Natural -> *).
KnownNat n =>
proxy n -> Natural
natVal (forall (t :: Natural). Proxy t
forall {k} (t :: k). Proxy t
Proxy @k)

-- | Push an item from the end of a vector.
--
-- Conceptually:
--
-- @
-- 'pushEnd' x v = 'VS.snoc' x ('VS.tail' v)
-- @
pushEnd :: (KnownNat n) => a -> VS.Vector n a -> VS.Vector n a
pushEnd :: forall (n :: Natural) a.
KnownNat n =>
a -> Vector n a -> Vector n a
pushEnd a
x Vector n a
v = (Finite n -> a) -> Vector n a
forall (n :: Natural) a.
KnownNat n =>
(Finite n -> a) -> Vector n a
VS.generate ((Finite n -> a) -> Vector n a) -> (Finite n -> a) -> Vector n a
forall a b. (a -> b) -> a -> b
$ \Finite n
i ->
  if Finite n
i Finite n -> Finite n -> Bool
forall a. Eq a => a -> a -> Bool
== Finite n
forall a. Bounded a => a
maxBound Finite n -> Finite n -> Finite n
forall a. Num a => a -> a -> a
- Finite n
1
    then a
x
    else Vector n a -> Finite n -> a
forall (n :: Natural) a. Vector n a -> Finite n -> a
VS.index Vector n a
v (Finite n
i Finite n -> Finite n -> Finite n
forall a. Num a => a -> a -> a
+ Finite n
1)