-- |
-- Module      : Data.Rhythm.Binary.BurrowsWheeler
-- Copyright   : (c) Eric Bailey, 2025
--
-- License     : MIT
-- Maintainer  : eric@ericb.me
-- Stability   : experimental
-- Portability : POSIX
--
-- Binary [necklaces](http://combos.org/necklace) via the set of Lyndon words
-- generated by an inverse [Burrows-Wheeler
-- transform](https://en.wikipedia.org/wiki/Burrows%E2%80%93Wheeler_transform).
-- Modified from [Rosetta
-- Code](https://rosettacode.org/wiki/De_Bruijn_sequences#Permutation-based).
module Data.Rhythm.Binary.BurrowsWheeler
  ( necklaces,
    lyndonWords,
    cycleForm,
  )
where

import Data.Functor.Base (ListF (..))
import Data.Functor.Foldable (ana)
import Data.IntMap.Strict ((!))
import Data.IntMap.Strict qualified as IntMap
import Data.List (elemIndices, sortOn)
import Data.List.Extra (snoc)
import Data.Ord (Down (..))

-- | All binary necklaces of a given length.
--
-- >>> necklaces 4
-- [[1,1,1,1],[1,1,1,0],[1,1,0,0],[1,0,1,0],[1,0,0,0],[0,0,0,0]]
necklaces :: Int -> [[Int]]
necklaces :: Int -> [[Int]]
necklaces Int
n =
  ([Int] -> Down [Int]) -> [[Int]] -> [[Int]]
forall b a. Ord b => (a -> b) -> [a] -> [a]
sortOn [Int] -> Down [Int]
forall a. a -> Down a
Down ([[Int]] -> [[Int]]) -> [[Int]] -> [[Int]]
forall a b. (a -> b) -> a -> b
$
    ([Int] -> [Int]) -> [[Int]] -> [[Int]]
forall a b. (a -> b) -> [a] -> [b]
map ([Int] -> [Int]
forall a. [a] -> [a]
reverse ([Int] -> [Int]) -> ([Int] -> [Int]) -> [Int] -> [Int]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> [Int] -> [Int]
forall a. Int -> [a] -> [a]
take Int
n ([Int] -> [Int]) -> ([Int] -> [Int]) -> [Int] -> [Int]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Int] -> [Int]
forall a. HasCallStack => [a] -> [a]
cycle) ([[Int]] -> [[Int]]) -> [[Int]] -> [[Int]]
forall a b. (a -> b) -> a -> b
$
      [Int] -> Int -> [[Int]]
forall a. Ord a => [a] -> Int -> [[a]]
lyndonWords [Int
0, Int
1] Int
n

-- | The set of Lyndon words generated by an inverse Burrows-Wheeler transform.
--
-- >>> lyndonWords [0,1] 4
-- [[0],[0,0,0,1],[0,0,1,1],[0,1],[0,1,1,1],[1]]
lyndonWords :: (Ord a) => [a] -> Int -> [[a]]
lyndonWords :: forall a. Ord a => [a] -> Int -> [[a]]
lyndonWords [a]
s Int
n = (Int -> a) -> [Int] -> [a]
forall a b. (a -> b) -> [a] -> [b]
map ([a]
ref !!) ([Int] -> [a]) -> [[Int]] -> [[a]]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Int] -> [[Int]]
cycleForm [Int]
perm
  where
    ref :: [a]
ref = Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
take ([a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
s Int -> Int -> Int
forall a b. (Num a, Integral b) => a -> b -> a
^ Int
n) ([a] -> [a]
forall a. HasCallStack => [a] -> [a]
cycle [a]
s)
    perm :: [Int]
perm = [a]
s [a] -> (a -> [Int]) -> [Int]
forall a b. [a] -> (a -> [b]) -> [b]
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (a -> [a] -> [Int]
forall a. Eq a => a -> [a] -> [Int]
`elemIndices` [a]
ref)

-- | Represent a permutation in a cycle notation.
--
-- >>> cycleForm [1,5,4,3,2,0]
-- [[1,5,0],[4,2],[3]]
cycleForm :: [Int] -> [[Int]]
cycleForm :: [Int] -> [[Int]]
cycleForm = (IntMap Int -> Base [[Int]] (IntMap Int)) -> IntMap Int -> [[Int]]
forall t a. Corecursive t => (a -> Base t a) -> a -> t
forall a. (a -> Base [[Int]] a) -> a -> [[Int]]
ana IntMap Int -> ListF [Int] (IntMap Int)
IntMap Int -> Base [[Int]] (IntMap Int)
coalgebra (IntMap Int -> [[Int]])
-> ([Int] -> IntMap Int) -> [Int] -> [[Int]]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(Int, Int)] -> IntMap Int
forall a. [(Int, a)] -> IntMap a
IntMap.fromAscList ([(Int, Int)] -> IntMap Int)
-> ([Int] -> [(Int, Int)]) -> [Int] -> IntMap Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Int] -> [Int] -> [(Int, Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0 ..]
  where
    coalgebra :: IntMap Int -> ListF [Int] (IntMap Int)
coalgebra = ListF [Int] (IntMap Int)
-> (((Int, Int), IntMap Int) -> ListF [Int] (IntMap Int))
-> Maybe ((Int, Int), IntMap Int)
-> ListF [Int] (IntMap Int)
forall b a. b -> (a -> b) -> Maybe a -> b
maybe ListF [Int] (IntMap Int)
forall a b. ListF a b
Nil ((Int, Int), IntMap Int) -> ListF [Int] (IntMap Int)
makeCycle (Maybe ((Int, Int), IntMap Int) -> ListF [Int] (IntMap Int))
-> (IntMap Int -> Maybe ((Int, Int), IntMap Int))
-> IntMap Int
-> ListF [Int] (IntMap Int)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IntMap Int -> Maybe ((Int, Int), IntMap Int)
forall a. IntMap a -> Maybe ((Int, a), IntMap a)
IntMap.minViewWithKey
    makeCycle :: ((Int, Int), IntMap Int) -> ListF [Int] (IntMap Int)
makeCycle ((Int
from, Int
to), IntMap Int
perm)
      | Int
from Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
to = [Int] -> IntMap Int -> ListF [Int] (IntMap Int)
forall a b. a -> b -> ListF a b
Cons [Int
from] IntMap Int
perm
      | Bool
otherwise = [Int] -> IntMap Int -> ListF [Int] (IntMap Int)
forall a b. a -> b -> ListF a b
Cons ([Int] -> Int -> [Int]
forall a. [a] -> a -> [a]
snoc [Int]
ciclo Int
from) ((Int -> IntMap Int -> IntMap Int)
-> IntMap Int -> [Int] -> IntMap Int
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr Int -> IntMap Int -> IntMap Int
forall a. Int -> IntMap a -> IntMap a
IntMap.delete IntMap Int
perm [Int]
ciclo)
      where
        ciclo :: [Int]
ciclo = (Int -> Bool) -> [Int] -> [Int]
forall a. (a -> Bool) -> [a] -> [a]
takeWhile (Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
from) ((Int -> Int) -> Int -> [Int]
forall a. (a -> a) -> a -> [a]
iterate (IntMap Int
perm !) Int
to)