{-# LANGUAGE DataKinds #-}
{-# LANGUAGE UnicodeSyntax #-}

-- |
-- Module      : Data.Rhythm.Binary.Necklaces
-- Description : Binary necklaces
-- Copyright   : (c) Eric Bailey, 2024-2025
--
-- License     : MIT
-- Maintainer  : eric@ericb.me
-- Stability   : stable
-- Portability : POSIX
--
-- Binary [necklaces](http://combos.org/necklace), internally encoded as
-- numbers.
--
-- = References
--   - Frank Ruskey, Carla Savage, Terry Min Yih Wang, Generating necklaces,
--     Journal of Algorithms, Volume 13, Issue 3, 1992, Pages 414-430, ISSN
--     0196-6774, https://doi.org/10.1016/0196-6774(92)90047-G.
module Data.Rhythm.Binary.Necklaces
  ( -- * Ergonomic
    necklaces,
    necklacesAllowed,
    necklacesPopCount,
    necklacesPopCountAllowed,

    -- * Efficient
    necklaces',
  )
where

import Control.Bool ((<&&>))
import Data.Bits (Bits (complementBit, rotateL, shiftL), popCount)
import Data.IntSet qualified as IntSet
import Data.List.NonEmpty qualified as NE
import Data.Rhythm.Internal (countParts, nodesToNecklaces)
import Data.Tree (Tree (..), flatten, unfoldTree)

-- $setup
-- >>> import Data.Finite (getFinite)

-- | 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 -> [Integer] -> [[Int]]
nodesToNecklaces Int
n ([Integer] -> [[Int]]) -> [Integer] -> [[Int]]
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
necklaces' Int
n)

-- | All binary necklaces of a given length with allowed parts.
--
-- >>> necklacesAllowed [1,2,3] 5
-- [[1,1,1,1,1],[1,1,1,1,0],[1,1,1,0,0],[1,1,0,1,0],[1,0,1,0,0]]
necklacesAllowed :: [Int] -> Int -> [[Int]]
necklacesAllowed :: [Int] -> Int -> [[Int]]
necklacesAllowed [Int]
allowed Int
n =
  Int -> [Integer] -> [[Int]]
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
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]]
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
necklaces' Int
n)

-- | All binary necklaces with a given number of ones of a given length with
-- allowed parts.
--
-- >>> necklacesPopCountAllowed 3 [1,2,3] 5
-- [[1,1,1,0,0],[1,1,0,1,0]]
necklacesPopCountAllowed :: Int -> [Int] -> Int -> [[Int]]
necklacesPopCountAllowed :: Int -> [Int] -> Int -> [[Int]]
necklacesPopCountAllowed Int
m [Int]
allowed Int
n =
  Int -> [Integer] -> [[Int]]
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
<&&> ((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 -> 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
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 of a given length, encoded as numbers.
--
-- \[
--   \begin{align*}
--     \sigma(x_1 ... x_n) &= x_2 ... x_n x_1 \\
--     \tau(x_1 ... x_{n-1}) &= x_1 ... x_{n-1}\overline{x_n}
--   \end{align*}
-- \]
--
-- Generate the tree of binary necklaces of length \(n\), starting with
-- \(x = 0^n\) as root, where children of \(x\) are the necklaces of the form
-- \(\tau\sigma^j(x)\) for \(1 \le j \le n -1\).
--
-- >>> flatten (necklaces' 4)
-- [0,1,3,7,15,5]
necklaces' :: (Integral a, Bits a) => Int -> Tree a
necklaces' :: forall a. (Integral a, Bits a) => Int -> Tree a
necklaces' Int
0 = a -> [Tree a] -> Tree a
forall a. a -> [Tree a] -> Tree a
Node a
0 []
necklaces' !Int
n = a -> [Tree a] -> Tree a
forall a. a -> [Tree a] -> Tree a
Node a
0 [(a -> (a, [a])) -> a -> Tree a
forall b a. (b -> (a, [b])) -> b -> Tree a
unfoldTree a -> (a, [a])
search a
1]
  where
    -- rotation
    σ :: a -> a
σ a
necklace = a -> Int -> a
forall a. Bits a => a -> Int -> a
rotateL a
necklace Int
1 a -> a -> a
forall a. Integral a => a -> a -> a
`mod` a
m
    σʲ :: a -> [a]
σʲ = Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
take (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) ([a] -> [a]) -> (a -> [a]) -> a -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. NonEmpty a -> [a]
forall a. NonEmpty a -> [a]
NE.tail (NonEmpty a -> [a]) -> (a -> NonEmpty a) -> a -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> a) -> a -> NonEmpty a
forall a. (a -> a) -> a -> NonEmpty a
NE.iterate a -> a
σ

    τ :: a -> a
τ = (a -> Int -> a) -> Int -> a -> a
forall a b c. (a -> b -> c) -> b -> a -> c
flip a -> Int -> a
forall a. Bits a => a -> Int -> a
complementBit Int
0

    -- build the tree of n-ary binary necklaces with a given root
    search :: a -> (a, [a])
search !a
necklace
      | a
necklace a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
m = (a
necklace, [])
      | Bool
otherwise = (a
necklace, (a -> Bool) -> [a] -> [a]
forall a. (a -> Bool) -> [a] -> [a]
takeWhile a -> Bool
isNecklace ((a -> a) -> [a] -> [a]
forall a b. (a -> b) -> [a] -> [b]
map a -> a
τ (a -> [a]
σʲ a
necklace)))

    -- a necklace is the lexicographically smallest rotation
    isNecklace :: a -> Bool
isNecklace a
necklace =
      a
necklace a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
m
        Bool -> Bool -> Bool
|| (a -> Bool) -> [a] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (a
necklace <=) (a -> [a]
σʲ a
necklace)

    -- 1ⁿ
    m :: a
m = a -> Int -> a
forall a. Bits a => a -> Int -> a
shiftL a
1 Int
n a -> a -> a
forall a. Num a => a -> a -> a
- a
1