{-# LANGUAGE DataKinds #-}
{-# LANGUAGE ExplicitNamespaces #-}
{-# LANGUAGE GADTs #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}

-- |
-- Module      : Data.Rhythm.Binary.FoldSequences
-- Description : Binary fold sequences
-- Copyright   : (c) Eric Bailey, 2025
--
-- License     : MIT
-- Maintainer  : eric@ericb.me
-- Stability   : stable
-- Portability : POSIX
--
-- Binary fold sequences.
module Data.Rhythm.Binary.FoldSequences
  ( -- * Ergonomic
    foldSequence,

    -- * Safe
    foldSequence',
  )
where

import Data.Bits (countTrailingZeros, shiftL, (.&.))
import Data.Bool (bool)
import Data.Finite (Finite, finite, getFinite)
import Data.Proxy (Proxy (..))
import Data.Vector.Sized (Vector)
import Data.Vector.Sized qualified as VS
import GHC.TypeNats (KnownNat, SomeNat (..), natVal, someNatVal, type (^))

-- | See 'foldSequence''.
--
-- >>> foldSequence 7 2 3
-- [1,1,0,1,1,0,0]
foldSequence :: Int -> Integer -> Integer -> [Integer]
foldSequence :: Int -> Integer -> Integer -> [Integer]
foldSequence Int
n Integer
m Integer
f =
  case (Natural -> SomeNat
someNatVal (Int -> Natural
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n), Natural -> SomeNat
someNatVal (Integer -> Natural
forall a b. (Integral a, Num b) => a -> b
fromIntegral Integer
m)) of
    (SomeNat (Proxy n
_ :: Proxy n), SomeNat (Proxy n
_ :: Proxy m)) ->
      (Finite 2 -> [Integer] -> [Integer])
-> [Integer] -> Vector Vector n (Finite 2) -> [Integer]
forall a b. (a -> b -> b) -> b -> Vector Vector n a -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr ((:) (Integer -> [Integer] -> [Integer])
-> (Finite 2 -> Integer) -> Finite 2 -> [Integer] -> [Integer]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Finite 2 -> Integer
forall (n :: Natural). Finite n -> Integer
getFinite) [] (Vector Vector n (Finite 2) -> [Integer])
-> Vector Vector n (Finite 2) -> [Integer]
forall a b. (a -> b) -> a -> b
$
        forall (n :: Natural) (m :: Natural).
(KnownNat n, KnownNat m) =>
Finite (2 ^ m) -> Vector n (Finite 2)
foldSequence' @n @m (Integer -> Finite (2 ^ n)
forall (n :: Natural). KnownNat n => Integer -> Finite n
finite Integer
f)

-- | Generate a fold sequence of @n@ terms with @m@ bits from function number
-- @f@ \(\in \{0,\dotsc,2^m-1\}\).
--
-- >>> map getFinite $ VS.toList $ foldSequence' @7 @2 3
-- [1,1,0,1,1,0,0]
foldSequence' :: forall n m. (KnownNat n, KnownNat m) => Finite (2 ^ m) -> Vector n (Finite 2)
foldSequence' :: forall (n :: Natural) (m :: Natural).
(KnownNat n, KnownNat m) =>
Finite (2 ^ m) -> Vector n (Finite 2)
foldSequence' Finite (2 ^ m)
f = (Int -> Finite 2) -> Vector n Int -> Vector n (Finite 2)
forall a b (n :: Natural). (a -> b) -> Vector n a -> Vector n b
VS.map Int -> Finite 2
go (forall (n :: Natural) a. (KnownNat n, Num a) => a -> Vector n a
VS.enumFromN @n @Int Int
1)
  where
    go :: Int -> Finite 2
go Int
i = Finite 2 -> Finite 2 -> Bool -> Finite 2
forall a. a -> a -> Bool -> a
bool Finite 2
x (Finite 2
1 Finite 2 -> Finite 2 -> Finite 2
forall a. Num a => a -> a -> a
- Finite 2
x) ((Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
b Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
4 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
1)
      where
        j :: Int
j = Int
i Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. (-Int
i)
        (Int
a, Int
b) = (Int -> Int
forall b. FiniteBits b => b -> Int
countTrailingZeros Int
j Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
m, Int
i Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
j))
        x :: Finite 2
x = Integer -> Finite 2
forall (n :: Natural). KnownNat n => Integer -> Finite n
finite (Int -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Bool -> Int
forall a. Enum a => a -> Int
fromEnum (Integer
0 Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== (Integer
g Integer -> Integer -> Integer
forall a. Bits a => a -> a -> a
.&. (Integer
1 Integer -> Int -> Integer
forall a. Bits a => a -> Int -> a
`shiftL` Int
a)))))
    m :: Int
m = Natural -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Proxy m -> Natural
forall (n :: Natural) (proxy :: Natural -> *).
KnownNat n =>
proxy n -> Natural
natVal (forall (t :: Natural). Proxy t
forall {k} (t :: k). Proxy t
Proxy @m))
    g :: Integer
g = Integer
pow2m Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- (Finite (2 ^ m) -> Integer
forall (n :: Natural). Finite n -> Integer
getFinite Finite (2 ^ m)
f Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`mod` Integer
pow2m) Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
1
    pow2m :: Integer
pow2m = Integer
2 Integer -> Int -> Integer
forall a b. (Num a, Integral b) => a -> b -> a
^ Int
m