{-# LANGUAGE DataKinds #-}
{-# LANGUAGE ExplicitNamespaces #-}

-- |
-- Module      : Data.Rhythm.Internal
-- Description : Here be dragons.
-- Copyright   : (c) Eric Bailey, 2025
--
-- License     : MIT
-- Maintainer  : eric@ericb.me
-- Stability   : experimental
-- Portability : POSIX
module Data.Rhythm.Internal
  ( -- * Digits
    countParts,
    padUpTo,

    -- * Parsers
    binaryDigit,

    -- * Necklaces
    nodesToNecklaces,

    -- * Vectors
    cycleVector,
  )
where

import Control.Applicative ((<|>))
import Data.Bits (Bits, testBit)
import Data.FastDigits (digits)
import Data.Finite (Finite, modulo)
import Data.List (sortOn, unfoldr)
import Data.Ord (Down (..))
import Data.Vector.Sized (Vector)
import Data.Vector.Sized qualified as VS
import GHC.TypeNats (KnownNat)
import Text.Trifecta (Parser, char, (<?>))

-- $setup
-- >>> import Text.Trifecta (parseString)

-- | 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 sizes \(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)

-- | Right pad a list with zeros up to a given length.
--
-- >>> padUpTo 5 [0,1]
-- [0,1,0,0,0]
padUpTo :: (Num a) => Int -> [a] -> [a]
padUpTo :: forall a. Num a => Int -> [a] -> [a]
padUpTo !Int
n [] = Int -> a -> [a]
forall a. Int -> a -> [a]
replicate Int
n a
0
padUpTo !Int
n (a
x : [a]
xs) = a
x a -> [a] -> [a]
forall a. a -> [a] -> [a]
: Int -> [a] -> [a]
forall a. Num a => Int -> [a] -> [a]
padUpTo (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) [a]
xs

-- | Parse a binary digit, i.e., @0@ or @1@.
--
-- >>> parseString binaryDigit mempty "0"
-- Success (finite 0)
--
-- >>> parseString binaryDigit mempty "1"
-- Success (finite 1)
--
-- >>> parseString binaryDigit mempty "?"
-- Failure (ErrInfo {_errDoc = (interactive):1:1: error: expected: one, zero
-- 1 | ?<EOF>
--   | ^      , _errDeltas = [Columns 0 0]})
binaryDigit :: Parser (Finite 2)
binaryDigit :: Parser (Finite 2)
binaryDigit = Parser (Finite 2)
zero Parser (Finite 2) -> Parser (Finite 2) -> Parser (Finite 2)
forall a. Parser a -> Parser a -> Parser a
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> Parser (Finite 2)
one
  where
    zero :: Parser (Finite 2)
zero = (Finite 2
0 Finite 2 -> Parser Char -> Parser (Finite 2)
forall a b. a -> Parser b -> Parser a
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ Char -> Parser Char
forall (m :: * -> *). CharParsing m => Char -> m Char
char Char
'0') Parser (Finite 2) -> String -> Parser (Finite 2)
forall a. Parser a -> String -> Parser a
forall (m :: * -> *) a. Parsing m => m a -> String -> m a
<?> String
"zero"
    one :: Parser (Finite 2)
one = (Finite 2
1 Finite 2 -> Parser Char -> Parser (Finite 2)
forall a b. a -> Parser b -> Parser a
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ Char -> Parser Char
forall (m :: * -> *). CharParsing m => Char -> m Char
char Char
'1') Parser (Finite 2) -> String -> Parser (Finite 2)
forall a. Parser a -> String -> Parser a
forall (m :: * -> *) a. Parsing m => m a -> String -> m a
<?> String
"one"

-- | Convert a list of nodes to binary necklaces of a given length.
--
-- >>> nodesToNecklaces 4 [3,5]
-- [[1,1,0,0],[1,0,1,0]]
nodesToNecklaces :: Int -> [Integer] -> [[Int]]
nodesToNecklaces :: Int -> [Integer] -> [[Int]]
nodesToNecklaces !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]])
-> ([Integer] -> [[Int]]) -> [Integer] -> [[Int]]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Integer -> [Int]) -> [Integer] -> [[Int]]
forall a b. (a -> b) -> [a] -> [b]
map (Int -> [Int] -> [Int]
forall a. Num a => Int -> [a] -> [a]
padUpTo Int
n ([Int] -> [Int]) -> (Integer -> [Int]) -> Integer -> [Int]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Integer -> [Int]
digits Int
2)

-- | Cycle a vector of length @n@ to produce a vector of length @m@.
--
-- Conceptually @'take' m '.' 'cycle'@, but for 'Vector's.
cycleVector :: forall m n a. (KnownNat m, KnownNat n) => Vector n a -> Vector m a
cycleVector :: forall (m :: Nat) (n :: Nat) a.
(KnownNat m, KnownNat n) =>
Vector n a -> Vector m a
cycleVector Vector n a
v = (Integer -> (a, Integer)) -> Integer -> Vector m a
forall (n :: Nat) a b.
KnownNat n =>
(b -> (a, b)) -> b -> Vector n a
VS.unfoldrN Integer -> (a, Integer)
go Integer
0
  where
    go :: Integer -> (a, Integer)
go Integer
i = (Vector n a
v Vector n a -> Finite n -> a
forall (n :: Nat) a. Vector n a -> Finite n -> a
`VS.index` Integer -> Finite n
forall (n :: Nat). KnownNat n => Integer -> Finite n
modulo Integer
i, Integer
i Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
1)