    -- |
    -- Module:      Math.NumberTheory.ArithmeticFunctions.NFreedom
    -- Copyright:   (c) 2018 Alexandre Rodrigues Baldé
    -- Licence:     MIT
    -- Maintainer:  Alexandre Rodrigues Baldé <alexandrer_b@outlook.com>
    --
    -- N-free number generation.
    --

    {-# LANGUAGE FlexibleContexts    #-}
    {-# LANGUAGE ScopedTypeVariables #-}

    module Math.NumberTheory.ArithmeticFunctions.NFreedom
      ( nFrees
      , nFreesBlock
      , sieveBlockNFree
      ) where

    import Control.Monad                         (forM_)
    import Control.Monad.ST                      (runST)
    import Data.Bits (Bits)
    import Data.List                             (scanl')
    import qualified Data.Vector.Unboxed         as U
    import qualified Data.Vector.Unboxed.Mutable as MU

    import Math.NumberTheory.Roots
    import Math.NumberTheory.Primes
    import Math.NumberTheory.Utils.FromIntegral

    -- | Evaluate the `Math.NumberTheory.ArithmeticFunctions.isNFree` function over a block.
    -- Value at @0@, if zero falls into block, is undefined.
    --
    -- This function should __**not**__ be used with a negative lower bound.
    -- If it is, the result is undefined.
    -- Furthermore, do not:
    --
    -- * use a block length greater than @maxBound :: Int@.
    -- * use a power that is either of @0, 1@.
    --
    -- None of these preconditions are checked, and if any occurs, the result is
    -- undefined, __if__ the function terminates.
    --
    -- >>> sieveBlockNFree 2 1 10
    -- [True,True,True,False,True,True,True,False,False,True]
    sieveBlockNFree
      :: forall a. (Integral a, Enum (Prime a), Bits a, UniqueFactorisation a)
      => Word
      -- ^ Power whose @n@-freedom will be checked.
      -> a
      -- ^ Lower index of the block.
      -> Word
      -- ^ Length of the block.
      -> U.Vector Bool
      -- ^ Vector of flags, where @True@ at index @i@ means the @i@-th element of
      -- the block is @n@-free.
    sieveBlockNFree :: forall a.
    (Integral a, Enum (Prime a), Bits a, UniqueFactorisation a) =>
    Word -> a -> Word -> Vector Bool
    sieveBlockNFree Word
    _ a
    _ Word
    0 = Vector Bool
    forall a. Unbox a => Vector a
    U.empty
    sieveBlockNFree Word
    n a
    lowIndex Word
    len'
      = (forall s. ST s (Vector Bool)) -> Vector Bool
    forall a. (forall s. ST s a) -> a
    runST ((forall s. ST s (Vector Bool)) -> Vector Bool)
    -> (forall s. ST s (Vector Bool)) -> Vector Bool
    forall a b. (a -> b) -> a -> b
    $ do
        MVector s Bool
    as <- Int -> Bool -> ST s (MVector (PrimState (ST s)) Bool)
    forall (m :: * -> *) a.
    (PrimMonad m, Unbox a) =>
    Int -> a -> m (MVector (PrimState m) a)
    MU.replicate (Word -> Int
    wordToInt Word
    len') Bool
    True
        [a] -> (a -> ST s ()) -> ST s ()
    forall (t :: * -> *) (m :: * -> *) a b.
    (Foldable t, Monad m) =>
    t a -> (a -> m b) -> m ()
    forM_ [a]
    ps ((a -> ST s ()) -> ST s ()) -> (a -> ST s ()) -> ST s ()
    forall a b. (a -> b) -> a -> b
    $ \a
    p -> do
          let pPow :: a
              pPow :: a
    pPow = a
    p a -> Word -> a
    forall a b. (Num a, Integral b) => a -> b -> a
    ^ Word
    n
              offset :: a
              offset :: a
    offset = a -> a
    forall a. Num a => a -> a
    negate a
    lowIndex a -> a -> a
    forall a. Integral a => a -> a -> a
    `mod` a
    pPow
              -- The second argument in @Data.Vector.Unboxed.Mutable.write@ is an
              -- @Int@, so to avoid segmentation faults or out-of-bounds errors,
              -- the enumeration's higher bound must always be less than
              -- @maxBound :: Int@.
              -- As mentioned above, this is not checked when using this function
              -- by itself, but is carefully managed when this function is called
              -- by @nFrees@, see the comments in it.
              indices :: [a]
              indices :: [a]
    indices = [a
    offset, a
    offset a -> a -> a
    forall a. Num a => a -> a -> a
    + a
    pPow .. a
    len a -> a -> a
    forall a. Num a => a -> a -> a
    - a
    1]
          [a] -> (a -> ST s ()) -> ST s ()
    forall (t :: * -> *) (m :: * -> *) a b.
    (Foldable t, Monad m) =>
    t a -> (a -> m b) -> m ()
    forM_ [a]
    indices ((a -> ST s ()) -> ST s ()) -> (a -> ST s ()) -> ST s ()
    forall a b. (a -> b) -> a -> b
    $ \a
    ix ->
              MVector (PrimState (ST s)) Bool -> Int -> Bool -> ST s ()
    forall (m :: * -> *) a.
    (PrimMonad m, Unbox a) =>
    MVector (PrimState m) a -> Int -> a -> m ()
    MU.write MVector s Bool
    MVector (PrimState (ST s)) Bool
    as (a -> Int
    forall a b. (Integral a, Num b) => a -> b
    fromIntegral a
    ix) Bool
    False
        MVector (PrimState (ST s)) Bool -> ST s (Vector Bool)
    forall a (m :: * -> *).
    (Unbox a, PrimMonad m) =>
    MVector (PrimState m) a -> m (Vector a)
    U.freeze MVector s Bool
    MVector (PrimState (ST s)) Bool
    as

      where
        len :: a
        len :: a
    len = Word -> a
    forall a b. (Integral a, Num b) => a -> b
    fromIntegral Word
    len'

        highIndex :: a
        highIndex :: a
    highIndex = a
    lowIndex a -> a -> a
    forall a. Num a => a -> a -> a
    + a
    len a -> a -> a
    forall a. Num a => a -> a -> a
    - a
    1

        ps :: [a]
        ps :: [a]
    ps = if a
    highIndex a -> a -> Bool
    forall a. Ord a => a -> a -> Bool
    < a
    4 then [] else (Prime a -> a) -> [Prime a] -> [a]
    forall a b. (a -> b) -> [a] -> [b]
    map Prime a -> a
    forall a. Prime a -> a
    unPrime [a -> Prime a
    forall a.
    (Bits a, Integral a, UniqueFactorisation a) =>
    a -> Prime a
    nextPrime a
    2 .. a -> Prime a
    forall a.
    (Bits a, Integral a, UniqueFactorisation a) =>
    a -> Prime a
    precPrime (a -> a
    forall a. Integral a => a -> a
    integerSquareRoot a
    highIndex)]

    -- | For a given nonnegative integer power @n@, generate all @n@-free
    -- numbers in ascending order, starting at @1@.
    --
    -- When @n@ is @0@ or @1@, the resulting list is @[1]@.
    nFrees
        :: forall a. (Integral a, Bits a, UniqueFactorisation a, Enum (Prime a))
        => Word
        -- ^ Power @n@ to be used to generate @n@-free numbers.
        -> [a]
        -- ^ Generated infinite list of @n@-free numbers.
    nFrees :: forall a.
    (Integral a, Bits a, UniqueFactorisation a, Enum (Prime a)) =>
    Word -> [a]
    nFrees Word
    0 = [a
    1]
    nFrees Word
    1 = [a
    1]
    nFrees Word
    n = ((a, Word) -> [a]) -> [(a, Word)] -> [a]
    forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
    concatMap ((a -> Word -> [a]) -> (a, Word) -> [a]
    forall a b c. (a -> b -> c) -> (a, b) -> c
    uncurry (Word -> a -> Word -> [a]
    forall a.
    (Integral a, Bits a, UniqueFactorisation a, Enum (Prime a)) =>
    Word -> a -> Word -> [a]
    nFreesBlock Word
    n)) ([(a, Word)] -> [a]) -> [(a, Word)] -> [a]
    forall a b. (a -> b) -> a -> b
    $ [a] -> [Word] -> [(a, Word)]
    forall a b. [a] -> [b] -> [(a, b)]
    zip [a]
    bounds [Word]
    strides
      where
        -- The 56th element of @iterate (2 *) 256@ is @2^64 :: Word == 0@, so to
        -- avoid overflow only the first 55 elements of this list are used.
        -- After those, since @maxBound :: Int@ is the largest a vector can be,
        -- this value is just repeated. This means after a few dozen iterations,
        -- the sieve will stop increasing in size.
        strides :: [Word]
        strides :: [Word]
    strides = Int -> [Word] -> [Word]
    forall a. Int -> [a] -> [a]
    take Int
    55 ((Word -> Word) -> Word -> [Word]
    forall a. (a -> a) -> a -> [a]
    iterate (Word
    2 Word -> Word -> Word
    forall a. Num a => a -> a -> a
    *) Word
    256) [Word] -> [Word] -> [Word]
    forall a. [a] -> [a] -> [a]
    ++ Word -> [Word]
    forall a. a -> [a]
    repeat (Int -> Word
    intToWord (Int
    forall a. Bounded a => a
    maxBound :: Int))

        -- Infinite list of lower bounds at which @sieveBlockNFree@ will be
        -- applied. This has type @Integral a => a@ instead of @Word@ because
        -- unlike the sizes of the sieve that eventually stop increasing (see
        -- above comment), the lower bound at which @sieveBlockNFree@ is called does not.
        bounds :: [a]
        bounds :: [a]
    bounds = (a -> a -> a) -> a -> [a] -> [a]
    forall b a. (b -> a -> b) -> b -> [a] -> [b]
    scanl' a -> a -> a
    forall a. Num a => a -> a -> a
    (+) a
    1 ([a] -> [a]) -> [a] -> [a]
    forall a b. (a -> b) -> a -> b
    $ (Word -> a) -> [Word] -> [a]
    forall a b. (a -> b) -> [a] -> [b]
    map Word -> a
    forall a b. (Integral a, Num b) => a -> b
    fromIntegral [Word]
    strides

    -- | Generate @n@-free numbers in a block starting at a certain value.
    -- The length of the list is determined by the value passed in as the third
    -- argument. It will be lesser than or equal to this value.
    --
    -- This function should not be used with a negative lower bound. If it is,
    -- the result is undefined.
    --
    -- The block length cannot exceed @maxBound :: Int@, this precondition is not
    -- checked.
    --
    -- As with @nFrees@, passing @n = 0, 1@ results in an empty list.
    nFreesBlock
        :: forall a . (Integral a, Bits a, UniqueFactorisation a, Enum (Prime a))
        => Word
        -- ^ Power @n@ to be used to generate @n@-free numbers.
        -> a
        -- ^ Starting number in the block.
        -> Word
        -- ^ Maximum length of the block to be generated.
        -> [a]
        -- ^ Generated list of @n@-free numbers.
    nFreesBlock :: forall a.
    (Integral a, Bits a, UniqueFactorisation a, Enum (Prime a)) =>
    Word -> a -> Word -> [a]
    nFreesBlock Word
    0 a
    lo Word
    _ = a -> [a]
    forall a. Integral a => a -> [a]
    help a
    lo
    nFreesBlock Word
    1 a
    lo Word
    _ = a -> [a]
    forall a. Integral a => a -> [a]
    help a
    lo
    nFreesBlock Word
    n a
    lowIndex Word
    len =
        let -- When indexing the array of flags @bs@, the index has to be an
            -- @Int@. As such, it's necessary to cast @strd@ twice.
            -- * Once, immediately below, to create the range of values whose
            -- @n@-freedom will be tested. Since @nFrees@ has return type
            -- @[a]@, this cannot be avoided as @strides@ has type @[Word]@.
            len' :: Int
            len' :: Int
    len' = Word -> Int
    wordToInt Word
    len
            -- * Twice, immediately below, to create the range of indices with
            -- which to query @bs@.
            len'' :: a
            len'' :: a
    len'' = Word -> a
    forall a b. (Integral a, Num b) => a -> b
    fromIntegral Word
    len
            bs :: Vector Bool
    bs  = Word -> a -> Word -> Vector Bool
    forall a.
    (Integral a, Enum (Prime a), Bits a, UniqueFactorisation a) =>
    Word -> a -> Word -> Vector Bool
    sieveBlockNFree Word
    n a
    lowIndex Word
    len
        in ((Int, a) -> a) -> [(Int, a)] -> [a]
    forall a b. (a -> b) -> [a] -> [b]
    map (Int, a) -> a
    forall a b. (a, b) -> b
    snd ([(Int, a)] -> [a]) -> ([a] -> [(Int, a)]) -> [a] -> [a]
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    .
           ((Int, a) -> Bool) -> [(Int, a)] -> [(Int, a)]
    forall a. (a -> Bool) -> [a] -> [a]
    filter ((Vector Bool
    bs Vector Bool -> Int -> Bool
    forall a. Unbox a => Vector a -> Int -> a
    U.!) (Int -> Bool) -> ((Int, a) -> Int) -> (Int, a) -> Bool
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (Int, a) -> Int
    forall a b. (a, b) -> a
    fst) ([(Int, a)] -> [(Int, a)])
    -> ([a] -> [(Int, a)]) -> [a] -> [(Int, a)]
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    .
           [Int] -> [a] -> [(Int, a)]
    forall a b. [a] -> [b] -> [(a, b)]
    zip [Int
    0 .. Int
    len' Int -> Int -> Int
    forall a. Num a => a -> a -> a
    - Int
    1] ([a] -> [a]) -> [a] -> [a]
    forall a b. (a -> b) -> a -> b
    $ [a
    lowIndex .. a
    lowIndex a -> a -> a
    forall a. Num a => a -> a -> a
    + a
    len'']
    {-# INLINE nFreesBlock #-}

    help :: Integral a => a -> [a]
    help :: forall a. Integral a => a -> [a]
    help a
    1 = [a
    1]
    help a
    _ = []
    {-# INLINE help #-}