{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE TemplateHaskell #-}

-- |
-- Module      : Data.Rhythm.Binary
-- Copyright   : (c) Eric Bailey, 2024-2025
--
-- License     : MIT
-- Maintainer  : eric@ericb.me
-- Stability   : experimental
-- Portability : POSIX
--
-- De Bruijn sequences, and conversion between binary strings and lists of
-- intervals.
module Data.Rhythm.Binary
  ( binaryToIntervals,
    deBruijnSequence,
    intervalsToBinary,
    RSW.necklaces,
    RSW.necklaces',
    necklacesAllowed,
    necklacesPopCount,
  )
where

import Control.Lens (makeLenses, uses, (%=))
import Control.Monad (when)
import Control.Monad.State (State, evalState)
import Data.Bits (Bits (testBit), popCount)
import Data.FastDigits (undigits)
import Data.IntSet qualified as IntSet
import Data.List (unfoldr)
import Data.List.Extra (snoc, splitOn)
import Data.Rhythm.Binary.RuskeySavageWang qualified as RSW
import Data.Set (Set)
import Data.Set qualified as Set
import Data.Tree (flatten)

-- | Convert a binary string to a list of intervals.
--
-- >>> binaryToIntervals "1010010001001000"
-- [2,3,4,3,4]
binaryToIntervals :: String -> [Int]
binaryToIntervals :: String -> [Int]
binaryToIntervals =
  (String -> Int) -> [String] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map ((Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int -> Int) -> (String -> Int) -> String -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length) ([String] -> [Int]) -> (String -> [String]) -> String -> [Int]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [String] -> [String]
forall a. HasCallStack => [a] -> [a]
tail ([String] -> [String])
-> (String -> [String]) -> String -> [String]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> String -> [String]
forall a. (HasCallStack, Eq a) => [a] -> [a] -> [[a]]
splitOn String
"1"

data DeBruijnState = DeBruijnState
  { DeBruijnState -> [Int]
_seed :: [Int],
    DeBruijnState -> Set Integer
_seen :: Set Integer
  }

makeLenses ''DeBruijnState

mkDeBruijnState :: Int -> DeBruijnState
mkDeBruijnState :: Int -> DeBruijnState
mkDeBruijnState Int
n = [Int] -> Set Integer -> DeBruijnState
DeBruijnState (Int -> Int -> [Int]
forall a. Int -> a -> [a]
replicate Int
n Int
0) Set Integer
forall a. Set a
Set.empty

-- | Generate the largest de Bruijn sequence of a given order.
--
-- Based on http://debruijnsequence.org/db/greedy.
--
-- >>> deBruijnSequence 4
-- [1,1,1,1,0,1,1,0,0,1,0,1,0,0,0,0]
deBruijnSequence :: Int -> [Int]
deBruijnSequence :: Int -> [Int]
deBruijnSequence Int
n = State DeBruijnState [Int] -> DeBruijnState -> [Int]
forall s a. State s a -> s -> a
evalState ([Int] -> State DeBruijnState [Int]
outer []) (Int -> DeBruijnState
mkDeBruijnState Int
n)
  where
    outer :: [Int] -> State DeBruijnState [Int]
outer [Int]
neck =
      do
        ([Int] -> Identity [Int])
-> DeBruijnState -> Identity DeBruijnState
Lens' DeBruijnState [Int]
seed (([Int] -> Identity [Int])
 -> DeBruijnState -> Identity DeBruijnState)
-> ([Int] -> [Int]) -> StateT DeBruijnState Identity ()
forall s (m :: * -> *) a b.
MonadState s m =>
ASetter s s a b -> (a -> b) -> m ()
%= (([Int] -> Int -> [Int]) -> Int -> [Int] -> [Int]
forall a b c. (a -> b -> c) -> b -> a -> c
flip [Int] -> Int -> [Int]
forall a. [a] -> a -> [a]
snoc Int
0 ([Int] -> [Int]) -> ([Int] -> [Int]) -> [Int] -> [Int]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Int] -> [Int]
forall a. HasCallStack => [a] -> [a]
tail)
        Int -> [Int] -> StateT DeBruijnState Identity (Maybe [Int])
inner Int
1 [Int]
neck StateT DeBruijnState Identity (Maybe [Int])
-> (Maybe [Int] -> State DeBruijnState [Int])
-> State DeBruijnState [Int]
forall a b.
StateT DeBruijnState Identity a
-> (a -> StateT DeBruijnState Identity b)
-> StateT DeBruijnState Identity b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= State DeBruijnState [Int]
-> ([Int] -> State DeBruijnState [Int])
-> Maybe [Int]
-> State DeBruijnState [Int]
forall b a. b -> (a -> b) -> Maybe a -> b
maybe ([Int] -> State DeBruijnState [Int]
forall a. a -> StateT DeBruijnState Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([Int] -> [Int]
forall a. [a] -> [a]
reverse [Int]
neck)) [Int] -> State DeBruijnState [Int]
outer

    inner :: Int -> [Int] -> StateT DeBruijnState Identity (Maybe [Int])
inner Int
i [Int]
neck =
      Int -> State DeBruijnState Bool
visit Int
i State DeBruijnState Bool
-> (Bool -> StateT DeBruijnState Identity (Maybe [Int]))
-> StateT DeBruijnState Identity (Maybe [Int])
forall a b.
StateT DeBruijnState Identity a
-> (a -> StateT DeBruijnState Identity b)
-> StateT DeBruijnState Identity b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
        Bool
True -> Maybe [Int] -> StateT DeBruijnState Identity (Maybe [Int])
forall a. a -> StateT DeBruijnState Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([Int] -> Maybe [Int]
forall a. a -> Maybe a
Just (Int
i Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: [Int]
neck))
        Bool
False | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0 -> Int -> [Int] -> StateT DeBruijnState Identity (Maybe [Int])
inner (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) [Int]
neck
        Bool
False -> Maybe [Int] -> StateT DeBruijnState Identity (Maybe [Int])
forall a. a -> StateT DeBruijnState Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe [Int]
forall a. Maybe a
Nothing

visit :: Int -> State DeBruijnState Bool
visit :: Int -> State DeBruijnState Bool
visit Int
v =
  do
    ([Int] -> Identity [Int])
-> DeBruijnState -> Identity DeBruijnState
Lens' DeBruijnState [Int]
seed (([Int] -> Identity [Int])
 -> DeBruijnState -> Identity DeBruijnState)
-> ([Int] -> [Int]) -> StateT DeBruijnState Identity ()
forall s (m :: * -> *) a b.
MonadState s m =>
ASetter s s a b -> (a -> b) -> m ()
%= (([Int] -> Int -> [Int]) -> Int -> [Int] -> [Int]
forall a b c. (a -> b -> c) -> b -> a -> c
flip [Int] -> Int -> [Int]
forall a. [a] -> a -> [a]
snoc Int
v ([Int] -> [Int]) -> ([Int] -> [Int]) -> [Int] -> [Int]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Int] -> [Int]
forall a. HasCallStack => [a] -> [a]
init)
    Integer
num <- LensLike' (Const Integer) DeBruijnState [Int]
-> ([Int] -> Integer) -> StateT DeBruijnState Identity Integer
forall s (m :: * -> *) r a.
MonadState s m =>
LensLike' (Const r) s a -> (a -> r) -> m r
uses LensLike' (Const Integer) DeBruijnState [Int]
Lens' DeBruijnState [Int]
seed (forall a b. (Integral a, Integral b) => a -> [b] -> Integer
undigits @Int Int
2)
    Bool
isNew <- LensLike' (Const Bool) DeBruijnState (Set Integer)
-> (Set Integer -> Bool) -> State DeBruijnState Bool
forall s (m :: * -> *) r a.
MonadState s m =>
LensLike' (Const r) s a -> (a -> r) -> m r
uses LensLike' (Const Bool) DeBruijnState (Set Integer)
Lens' DeBruijnState (Set Integer)
seen (Integer -> Set Integer -> Bool
forall a. Ord a => a -> Set a -> Bool
Set.notMember Integer
num)
    Bool
-> StateT DeBruijnState Identity ()
-> StateT DeBruijnState Identity ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
isNew (StateT DeBruijnState Identity ()
 -> StateT DeBruijnState Identity ())
-> StateT DeBruijnState Identity ()
-> StateT DeBruijnState Identity ()
forall a b. (a -> b) -> a -> b
$
      (Set Integer -> Identity (Set Integer))
-> DeBruijnState -> Identity DeBruijnState
Lens' DeBruijnState (Set Integer)
seen ((Set Integer -> Identity (Set Integer))
 -> DeBruijnState -> Identity DeBruijnState)
-> (Set Integer -> Set Integer) -> StateT DeBruijnState Identity ()
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 -> State DeBruijnState Bool
forall a. a -> StateT DeBruijnState Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
isNew

-- | Convert a list of intervals to a binary string.
--
-- >>> intervalsToBinary [2,3,4,3,4]
-- "1010010001001000"
intervalsToBinary :: [Int] -> String
intervalsToBinary :: [Int] -> String
intervalsToBinary =
  (Int -> String) -> [Int] -> String
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap ((Char
'1' :) (String -> String) -> (Int -> String) -> Int -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int -> Char -> String
forall a. Int -> a -> [a]
`replicate` Char
'0') (Int -> String) -> (Int -> Int) -> Int -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Int -> Int
forall a. Num a => a -> a -> a
subtract Int
1)

necklacesAllowed :: [Int] -> Int -> [[Int]]
necklacesAllowed :: [Int] -> Int -> [[Int]]
necklacesAllowed [Int]
allowed Int
n =
  Int -> [Integer] -> [[Int]]
RSW.nodesToNecklaces Int
n ([Integer] -> [[Int]]) -> [Integer] -> [[Int]]
forall a b. (a -> b) -> a -> b
$
    (Integer -> Bool) -> [Integer] -> [Integer]
forall a. (a -> Bool) -> [a] -> [a]
filter ((Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
> Integer
0) (Integer -> Bool) -> (Integer -> Bool) -> Integer -> Bool
forall (f :: * -> *). Applicative f => f Bool -> f Bool -> f Bool
<&&> Integer -> Bool
isAllowed) ([Integer] -> [Integer]) -> [Integer] -> [Integer]
forall a b. (a -> b) -> a -> b
$
      Tree Integer -> [Integer]
forall a. Tree a -> [a]
flatten (Int -> Tree Integer
forall a. (Integral a, Bits a) => Int -> Tree a
RSW.necklaces' Int
n)
  where
    isAllowed :: Integer -> Bool
isAllowed = (Int -> Bool) -> [Int] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (Int -> IntSet -> Bool
`IntSet.member` [Int] -> IntSet
IntSet.fromList [Int]
allowed) ([Int] -> Bool) -> (Integer -> [Int]) -> Integer -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Integer -> [Int]
forall a. (Integral a, Bits a) => Int -> a -> [Int]
countParts Int
n

-- | All binary necklaces with a given number of ones of a given length.
--
-- >>> necklacesPopCount 3 6
-- [[1,1,1,0,0,0],[1,1,0,1,0,0],[1,0,1,1,0,0],[1,0,1,0,1,0]]
necklacesPopCount :: Int -> Int -> [[Int]]
necklacesPopCount :: Int -> Int -> [[Int]]
necklacesPopCount !Int
m !Int
n =
  Int -> [Integer] -> [[Int]]
RSW.nodesToNecklaces Int
n ([Integer] -> [[Int]]) -> [Integer] -> [[Int]]
forall a b. (a -> b) -> a -> b
$
    (Integer -> Bool) -> [Integer] -> [Integer]
forall a. (a -> Bool) -> [a] -> [a]
filter ((Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
m) (Int -> Bool) -> (Integer -> Int) -> Integer -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Integer -> Int
forall a. Bits a => a -> Int
popCount) ([Integer] -> [Integer]) -> [Integer] -> [Integer]
forall a b. (a -> b) -> a -> b
$
      Tree Integer -> [Integer]
forall a. Tree a -> [a]
flatten (Int -> Tree Integer
forall a. (Integral a, Bits a) => Int -> Tree a
RSW.necklaces' Int
n)

-- | Count the parts in the n-digit little-endian binary representation of x.
--
-- A part is the length of a substring 10* composing the necklace.
-- For example the necklace 10100 has parts of size 2 and 3.
--
-- >>> countParts 5 5
-- [2,3]
countParts :: (Integral a, Bits a) => Int -> a -> [Int]
countParts :: forall a. (Integral a, Bits a) => Int -> a -> [Int]
countParts Int
n a
x = ((Int, Int) -> Maybe (Int, (Int, Int))) -> (Int, Int) -> [Int]
forall b a. (b -> Maybe (a, b)) -> b -> [a]
unfoldr (Int, Int) -> Maybe (Int, (Int, Int))
forall {b} {b}.
(Num b, Num b, Eq b) =>
(Int, b) -> Maybe (b, (Int, b))
go (Int
0, Int
0)
  where
    go :: (Int, b) -> Maybe (b, (Int, b))
go (Int
i, b
0)
      | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
n = Maybe (b, (Int, b))
forall a. Maybe a
Nothing
      | a -> Int -> Bool
forall a. Bits a => a -> Int -> Bool
testBit a
x Int
i = (Int, b) -> Maybe (b, (Int, b))
go (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1, b
1)
      | Bool
otherwise = (Int, b) -> Maybe (b, (Int, b))
go (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1, b
0)
    go (Int
i, b
len)
      | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
n = (b, (Int, b)) -> Maybe (b, (Int, b))
forall a. a -> Maybe a
Just (b
len, (Int
i, b
0))
      | a -> Int -> Bool
forall a. Bits a => a -> Int -> Bool
testBit a
x Int
i = (b, (Int, b)) -> Maybe (b, (Int, b))
forall a. a -> Maybe a
Just (b
len, (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1, b
1))
      | Bool
otherwise = (Int, b) -> Maybe (b, (Int, b))
go (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1, b
len b -> b -> b
forall a. Num a => a -> a -> a
+ b
1)

(<&&>) :: (Applicative f) => f Bool -> f Bool -> f Bool
<&&> :: forall (f :: * -> *). Applicative f => f Bool -> f Bool -> f Bool
(<&&>) = (Bool -> Bool -> Bool) -> f Bool -> f Bool -> f Bool
forall a b c. (a -> b -> c) -> f a -> f b -> f c
forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 Bool -> Bool -> Bool
(&&)

infixr 3 <&&> -- same as (&&)