-- | Internal module whose primary export is 'uniformPartition'.  Use
-- 'RandomCycle.Vector' instead.
module RandomCycle.Vector.Partition where

import Data.Bits
import qualified Data.Vector.Generic as GV
import GHC.Natural (Natural)
import System.Random.Stateful

{- UTILITIES -}

---- | Internal. Find the first index where the bit is flipped, shifting the
-- bits as you go and returning the final shifted bit vector. The degenerate
-- case @bs == 0@ returns the otherwise unreachable point @(0, 0)@ to guarantee
-- termination, but note that case is nonsensical in 'partitionFromBits' and
-- handled explicitly there.
commonSubseqBits :: Natural -> (Natural, Int)
commonSubseqBits 0 = (0, 0)
commonSubseqBits bs = until done (\(bs', i) -> (bs' `shiftR` 1, i + 1)) (bs, 0)
  where
    done = if bs `testBit` 0 then not . (`testBit` 0) . fst else (`testBit` 0) . fst

-- | Partition a vector 'v' according to groupings provided by the bits 'bs'.
-- If the first set bit in 'bs' is at a position larger than the last index of
-- 'v', this returns @[v]@. More generally, bits set at positions after the
-- last index of 'v' do not contribute to the grouping. @bs == 0@ always
-- results in @[v]@.
--
-- See 'RandomCycle.List.partitionFromBits' for other examples.
--
-- ==== __Examples__
--
-- >>> import qualified Data.Vector as V
-- >>> partitionFromBits 5 (V.fromList [0..2::Int])
-- [[0],[1],[2]]
-- >>> partitionFromBits 13 (V.fromList [0..2::Int])
-- [[0],[1],[2]]
-- >>> partitionFromBits 4 (V.fromList [0..2::Int])
-- [[0,1],[2]]
-- >>> partitionFromBits 8 (V.fromList [0..2::Int])
-- [[0,1,2]]
partitionFromBits :: (GV.Vector v a) => Natural -> v a -> [v a]
partitionFromBits _ v | GV.null v = []
partitionFromBits 0 v = [v]
partitionFromBits bs v =
  let (bs', idx) = commonSubseqBits bs
      (v1, v2) = GV.splitAt idx v
   in v1 : partitionFromBits bs' v2

{- RANDOM -}

-- | Draw a random partition of the input vector @xs@ from the uniform
-- distribution on partitions. This proceeds by randomizing the placement of
-- each breakpoint, in other words by walking a random path in a perfect binary
-- tree. /O(n)/ for a vector length /n/.
--
-- This function preserves the order of the input list.
uniformPartition :: (GV.Vector v a, StatefulGen g m) => v a -> g -> m [v a]
uniformPartition xs g = do
  let d = GV.length xs
  bs <- uniformRM (0, 2 ^ d - 1) g
  pure $ partitionFromBits bs xs