-- |
-- Module:      Math.NumberTheory.ArithmeticFunctions.Standard
-- Copyright:   (c) 2016 Andrew Lelechenko
-- Licence:     MIT
-- Maintainer:  Andrew Lelechenko <andrew.lelechenko@gmail.com>
--
-- Textbook arithmetic functions.
--

{-# LANGUAGE ScopedTypeVariables #-}

module Math.NumberTheory.ArithmeticFunctions.Standard
  ( -- * List divisors
    divisors, divisorsA
  , divisorsList, divisorsListA
  , divisorsSmall, divisorsSmallA
  , divisorsTo, divisorsToA
    -- * Multiplicative functions
  , multiplicative
  , divisorCount, tau, tauA
  , sigma, sigmaA
  , totient, totientA
  , jordan, jordanA
  , ramanujan, ramanujanA
  , moebius, moebiusA, Moebius(..), runMoebius
  , liouville, liouvilleA
    -- * Additive functions
  , additive
  , smallOmega, smallOmegaA
  , bigOmega, bigOmegaA
    -- * Misc
  , carmichael, carmichaelA
  , expMangoldt, expMangoldtA
  , isNFree, isNFreeA, nFrees, nFreesBlock
  ) where

import Data.Coerce
import Data.Euclidean (GcdDomain(divide))
import Data.IntSet (IntSet)
import qualified Data.IntSet as IS
import Data.Maybe
import Data.Set (Set)
import qualified Data.Set as S
import Data.Semigroup

import Math.NumberTheory.ArithmeticFunctions.Class
import Math.NumberTheory.ArithmeticFunctions.Moebius
import Math.NumberTheory.ArithmeticFunctions.NFreedom (nFrees, nFreesBlock)
import Math.NumberTheory.Primes
import Math.NumberTheory.Utils.FromIntegral

import Numeric.Natural

-- | Create a multiplicative function from the function on prime's powers. See examples below.
multiplicative :: Num a => (Prime n -> Word -> a) -> ArithmeticFunction n a
multiplicative :: forall a n.
Num a =>
(Prime n -> Word -> a) -> ArithmeticFunction n a
multiplicative Prime n -> Word -> a
f = forall m n a.
Monoid m =>
(Prime n -> Word -> m) -> (m -> a) -> ArithmeticFunction n a
ArithmeticFunction ((forall a. a -> Product a
Product forall b c a. (b -> c) -> (a -> b) -> a -> c
.) forall b c a. (b -> c) -> (a -> b) -> a -> c
. Prime n -> Word -> a
f) forall a. Product a -> a
getProduct

-- | See 'divisorsA'.
divisors :: (UniqueFactorisation n, Ord n) => n -> Set n
divisors :: forall n. (UniqueFactorisation n, Ord n) => n -> Set n
divisors = forall n a.
UniqueFactorisation n =>
ArithmeticFunction n a -> n -> a
runFunction forall n. (Ord n, Num n) => ArithmeticFunction n (Set n)
divisorsA
{-# SPECIALIZE divisors :: Natural -> Set Natural #-}
{-# SPECIALIZE divisors :: Integer -> Set Integer #-}

-- | The set of all (positive) divisors of an argument.
divisorsA :: (Ord n, Num n) => ArithmeticFunction n (Set n)
divisorsA :: forall n. (Ord n, Num n) => ArithmeticFunction n (Set n)
divisorsA = forall m n a.
Monoid m =>
(Prime n -> Word -> m) -> (m -> a) -> ArithmeticFunction n a
ArithmeticFunction (\Prime n
p -> forall a. Set a -> SetProduct a
SetProduct forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall n. Num n => n -> Word -> Set n
divisorsHelper (forall a. Prime a -> a
unPrime Prime n
p)) (forall a. Ord a => a -> Set a -> Set a
S.insert n
1 forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. SetProduct a -> Set a
getSetProduct)

divisorsHelper :: Num n => n -> Word -> Set n
divisorsHelper :: forall n. Num n => n -> Word -> Set n
divisorsHelper n
_ Word
0 = forall a. Set a
S.empty
divisorsHelper n
p Word
1 = forall a. a -> Set a
S.singleton n
p
divisorsHelper n
p Word
a = forall a. [a] -> Set a
S.fromDistinctAscList forall a b. (a -> b) -> a -> b
$ n
p forall a. a -> [a] -> [a]
: n
p forall a. Num a => a -> a -> a
* n
p forall a. a -> [a] -> [a]
: forall a b. (a -> b) -> [a] -> [b]
map (n
p forall a b. (Num a, Integral b) => a -> b -> a
^) [Int
3 .. Word -> Int
wordToInt Word
a]
{-# INLINE divisorsHelper #-}

-- | See 'divisorsListA'.
divisorsList :: UniqueFactorisation n => n -> [n]
divisorsList :: forall n. UniqueFactorisation n => n -> [n]
divisorsList = forall n a.
UniqueFactorisation n =>
ArithmeticFunction n a -> n -> a
runFunction forall n. Num n => ArithmeticFunction n [n]
divisorsListA

-- | The unsorted list of all (positive) divisors of an argument, produced in lazy fashion.
divisorsListA :: Num n => ArithmeticFunction n [n]
divisorsListA :: forall n. Num n => ArithmeticFunction n [n]
divisorsListA = forall m n a.
Monoid m =>
(Prime n -> Word -> m) -> (m -> a) -> ArithmeticFunction n a
ArithmeticFunction (\Prime n
p -> forall a. [a] -> ListProduct a
ListProduct forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall n. Num n => n -> Word -> [n]
divisorsListHelper (forall a. Prime a -> a
unPrime Prime n
p)) ((n
1 forall a. a -> [a] -> [a]
:) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. ListProduct a -> [a]
getListProduct)

divisorsListHelper :: Num n => n -> Word -> [n]
divisorsListHelper :: forall n. Num n => n -> Word -> [n]
divisorsListHelper n
_ Word
0 = []
divisorsListHelper n
p Word
1 = [n
p]
divisorsListHelper n
p Word
a = n
p forall a. a -> [a] -> [a]
: n
p forall a. Num a => a -> a -> a
* n
p forall a. a -> [a] -> [a]
: forall a b. (a -> b) -> [a] -> [b]
map (n
p forall a b. (Num a, Integral b) => a -> b -> a
^) [Int
3 .. Word -> Int
wordToInt Word
a]
{-# INLINE divisorsListHelper #-}

-- | See 'divisorsSmallA'.
divisorsSmall :: Int -> IntSet
divisorsSmall :: Int -> IntSet
divisorsSmall = forall n a.
UniqueFactorisation n =>
ArithmeticFunction n a -> n -> a
runFunction ArithmeticFunction Int IntSet
divisorsSmallA

-- | Same as 'divisors', but with better performance on cost of type restriction.
divisorsSmallA :: ArithmeticFunction Int IntSet
divisorsSmallA :: ArithmeticFunction Int IntSet
divisorsSmallA = forall m n a.
Monoid m =>
(Prime n -> Word -> m) -> (m -> a) -> ArithmeticFunction n a
ArithmeticFunction (\Prime Int
p -> IntSet -> IntSetProduct
IntSetProduct forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Word -> IntSet
divisorsHelperSmall (forall a. Prime a -> a
unPrime Prime Int
p)) (Int -> IntSet -> IntSet
IS.insert Int
1 forall b c a. (b -> c) -> (a -> b) -> a -> c
. IntSetProduct -> IntSet
getIntSetProduct)

divisorsHelperSmall :: Int -> Word -> IntSet
divisorsHelperSmall :: Int -> Word -> IntSet
divisorsHelperSmall Int
_ Word
0 = IntSet
IS.empty
divisorsHelperSmall Int
p Word
1 = Int -> IntSet
IS.singleton Int
p
divisorsHelperSmall Int
p Word
a = [Int] -> IntSet
IS.fromDistinctAscList forall a b. (a -> b) -> a -> b
$ Int
p forall a. a -> [a] -> [a]
: Int
p forall a. Num a => a -> a -> a
* Int
p forall a. a -> [a] -> [a]
: forall a b. (a -> b) -> [a] -> [b]
map (Int
p forall a b. (Num a, Integral b) => a -> b -> a
^) [Int
3 .. Word -> Int
wordToInt Word
a]
{-# INLINE divisorsHelperSmall #-}

-- | See 'divisorsToA'.
divisorsTo :: (UniqueFactorisation n, Integral n) => n -> n -> Set n
divisorsTo :: forall n. (UniqueFactorisation n, Integral n) => n -> n -> Set n
divisorsTo n
to = forall n a.
UniqueFactorisation n =>
ArithmeticFunction n a -> n -> a
runFunction (forall n.
(UniqueFactorisation n, Integral n) =>
n -> ArithmeticFunction n (Set n)
divisorsToA n
to)

-- | The set of all (positive) divisors up to an inclusive bound.
divisorsToA :: (UniqueFactorisation n, Integral n) => n -> ArithmeticFunction n (Set n)
divisorsToA :: forall n.
(UniqueFactorisation n, Integral n) =>
n -> ArithmeticFunction n (Set n)
divisorsToA n
to = forall m n a.
Monoid m =>
(Prime n -> Word -> m) -> (m -> a) -> ArithmeticFunction n a
ArithmeticFunction forall {a}.
(Ord a, Num a) =>
Prime a -> Word -> BoundedSetProduct a
f BoundedSetProduct n -> Set n
unwrap
  where f :: Prime a -> Word -> BoundedSetProduct a
f Prime a
p Word
k = forall a. (a -> Set a) -> BoundedSetProduct a
BoundedSetProduct (\a
bound -> forall n. (Ord n, Num n) => n -> n -> Word -> Set n
divisorsToHelper a
bound (forall a. Prime a -> a
unPrime Prime a
p) Word
k)
        unwrap :: BoundedSetProduct n -> Set n
unwrap (BoundedSetProduct n -> Set n
res) = if n
1 forall a. Ord a => a -> a -> Bool
<= n
to then forall a. Ord a => a -> Set a -> Set a
S.insert n
1 (n -> Set n
res n
to) else n -> Set n
res n
to

-- | Generate at most @a@ powers of @p@ up to an inclusive bound @b@.
divisorsToHelper :: (Ord n, Num n) => n -> n -> Word -> Set n
divisorsToHelper :: forall n. (Ord n, Num n) => n -> n -> Word -> Set n
divisorsToHelper n
_ n
_ Word
0 = forall a. Set a
S.empty
divisorsToHelper n
b n
p Word
1 = if n
p forall a. Ord a => a -> a -> Bool
<= n
b then forall a. a -> Set a
S.singleton n
p else forall a. Set a
S.empty
divisorsToHelper n
b n
p Word
a = forall a. [a] -> Set a
S.fromDistinctAscList forall a b. (a -> b) -> a -> b
$ forall a. Int -> [a] -> [a]
take (Word -> Int
wordToInt Word
a) forall a b. (a -> b) -> a -> b
$ forall a. (a -> Bool) -> [a] -> [a]
takeWhile (forall a. Ord a => a -> a -> Bool
<=n
b) forall a b. (a -> b) -> a -> b
$ forall a. (a -> a) -> a -> [a]
iterate (n
pforall a. Num a => a -> a -> a
*) n
p
{-# INLINE divisorsToHelper #-}

-- | Synonym for 'tau'.
--
-- >>> map divisorCount [1..10]
-- [1,2,2,3,2,4,2,4,3,4]
divisorCount :: (UniqueFactorisation n, Num a) => n -> a
divisorCount :: forall n a. (UniqueFactorisation n, Num a) => n -> a
divisorCount = forall n a. (UniqueFactorisation n, Num a) => n -> a
tau

-- | See 'tauA'.
tau :: (UniqueFactorisation n, Num a) => n -> a
tau :: forall n a. (UniqueFactorisation n, Num a) => n -> a
tau = forall n a.
UniqueFactorisation n =>
ArithmeticFunction n a -> n -> a
runFunction forall a n. Num a => ArithmeticFunction n a
tauA

-- | The number of (positive) divisors of an argument.
--
-- > tauA = multiplicative (\_ k -> k + 1)
tauA :: Num a => ArithmeticFunction n a
tauA :: forall a n. Num a => ArithmeticFunction n a
tauA = forall a n.
Num a =>
(Prime n -> Word -> a) -> ArithmeticFunction n a
multiplicative forall a b. (a -> b) -> a -> b
$ forall a b. a -> b -> a
const (forall a b. (Integral a, Num b) => a -> b
fromIntegral forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Enum a => a -> a
succ)

-- | See 'sigmaA'.
sigma :: (UniqueFactorisation n, Integral n, Num a, GcdDomain a) => Word -> n -> a
sigma :: forall n a.
(UniqueFactorisation n, Integral n, Num a, GcdDomain a) =>
Word -> n -> a
sigma = forall n a.
UniqueFactorisation n =>
ArithmeticFunction n a -> n -> a
runFunction forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall n a.
(Integral n, Num a, GcdDomain a) =>
Word -> ArithmeticFunction n a
sigmaA
{-# INLINABLE sigma #-}

-- | The sum of the @k@-th powers of (positive) divisors of an argument.
--
-- > sigmaA = multiplicative (\p k -> sum $ map (p ^) [0..k])
-- > sigmaA 0 = tauA
sigmaA :: (Integral n, Num a, GcdDomain a) => Word -> ArithmeticFunction n a
sigmaA :: forall n a.
(Integral n, Num a, GcdDomain a) =>
Word -> ArithmeticFunction n a
sigmaA Word
0 = forall a n. Num a => ArithmeticFunction n a
tauA
sigmaA Word
1 = forall a n.
Num a =>
(Prime n -> Word -> a) -> ArithmeticFunction n a
multiplicative forall a b. (a -> b) -> a -> b
$ forall a. (Num a, GcdDomain a) => a -> Word -> a
sigmaHelper forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (Integral a, Num b) => a -> b
fromIntegral' forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Prime a -> a
unPrime
sigmaA Word
a = forall a n.
Num a =>
(Prime n -> Word -> a) -> ArithmeticFunction n a
multiplicative forall a b. (a -> b) -> a -> b
$ forall a. (Num a, GcdDomain a) => a -> Word -> a
sigmaHelper forall b c a. (b -> c) -> (a -> b) -> a -> c
. (forall a b. (Num a, Integral b) => a -> b -> a
^ Word -> Int
wordToInt Word
a) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (Integral a, Num b) => a -> b
fromIntegral' forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Prime a -> a
unPrime
{-# INLINABLE sigmaA #-}

sigmaHelper :: (Num a, GcdDomain a) => a -> Word -> a
sigmaHelper :: forall a. (Num a, GcdDomain a) => a -> Word -> a
sigmaHelper a
pa Word
1 = a
pa forall a. Num a => a -> a -> a
+ a
1
sigmaHelper a
pa Word
2 = a
pa forall a. Num a => a -> a -> a
* a
pa forall a. Num a => a -> a -> a
+ a
pa forall a. Num a => a -> a -> a
+ a
1
sigmaHelper a
pa Word
k = forall a. HasCallStack => Maybe a -> a
fromJust ((a
pa forall a b. (Num a, Integral b) => a -> b -> a
^ Word -> Int
wordToInt (Word
k forall a. Num a => a -> a -> a
+ Word
1) forall a. Num a => a -> a -> a
- a
1) forall a. GcdDomain a => a -> a -> Maybe a
`divide` (a
pa forall a. Num a => a -> a -> a
- a
1))
{-# INLINE sigmaHelper #-}

-- | See 'totientA'.
totient :: UniqueFactorisation n => n -> n
totient :: forall n. UniqueFactorisation n => n -> n
totient = forall n a.
UniqueFactorisation n =>
ArithmeticFunction n a -> n -> a
runFunction forall n. Num n => ArithmeticFunction n n
totientA
{-# INLINABLE totient #-}

-- | Calculates the totient of a positive number @n@, i.e.
--   the number of @k@ with @1 <= k <= n@ and @'gcd' n k == 1@,
--   in other words, the order of the group of units in @&#8484;/(n)@.
totientA :: Num n => ArithmeticFunction n n
totientA :: forall n. Num n => ArithmeticFunction n n
totientA = forall a n.
Num a =>
(Prime n -> Word -> a) -> ArithmeticFunction n a
multiplicative forall a b. (a -> b) -> a -> b
$ forall n. Num n => n -> Word -> n
jordanHelper forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Prime a -> a
unPrime
{-# INLINABLE totientA #-}

-- | See 'jordanA'.
jordan :: UniqueFactorisation n => Word -> n -> n
jordan :: forall n. UniqueFactorisation n => Word -> n -> n
jordan = forall n a.
UniqueFactorisation n =>
ArithmeticFunction n a -> n -> a
runFunction forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall n. Num n => Word -> ArithmeticFunction n n
jordanA

-- | Calculates the k-th Jordan function of an argument.
--
-- > jordanA 1 = totientA
jordanA :: Num n => Word -> ArithmeticFunction n n
jordanA :: forall n. Num n => Word -> ArithmeticFunction n n
jordanA Word
0 = forall a n.
Num a =>
(Prime n -> Word -> a) -> ArithmeticFunction n a
multiplicative forall a b. (a -> b) -> a -> b
$ \Prime n
_ Word
_ -> n
0
jordanA Word
1 = forall n. Num n => ArithmeticFunction n n
totientA
jordanA Word
a = forall a n.
Num a =>
(Prime n -> Word -> a) -> ArithmeticFunction n a
multiplicative forall a b. (a -> b) -> a -> b
$ forall n. Num n => n -> Word -> n
jordanHelper forall b c a. (b -> c) -> (a -> b) -> a -> c
. (forall a b. (Num a, Integral b) => a -> b -> a
^ Word -> Int
wordToInt Word
a) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Prime a -> a
unPrime

jordanHelper :: Num n => n -> Word -> n
jordanHelper :: forall n. Num n => n -> Word -> n
jordanHelper n
pa Word
1 = n
pa forall a. Num a => a -> a -> a
- n
1
jordanHelper n
pa Word
2 = (n
pa forall a. Num a => a -> a -> a
- n
1) forall a. Num a => a -> a -> a
* n
pa
jordanHelper n
pa Word
k = (n
pa forall a. Num a => a -> a -> a
- n
1) forall a. Num a => a -> a -> a
* n
pa forall a b. (Num a, Integral b) => a -> b -> a
^ Word -> Int
wordToInt (Word
k forall a. Num a => a -> a -> a
- Word
1)
{-# INLINE jordanHelper #-}

-- | See 'ramanujanA'.
ramanujan :: Integer -> Integer
ramanujan :: Integer -> Integer
ramanujan = forall n a.
UniqueFactorisation n =>
ArithmeticFunction n a -> n -> a
runFunction ArithmeticFunction Integer Integer
ramanujanA

-- | Calculates the <https://en.wikipedia.org/wiki/Ramanujan_tau_function Ramanujan tau function>
--   of a positive number @n@, using formulas given <http://www.numbertheory.org/php/tau.html here>
ramanujanA :: ArithmeticFunction Integer Integer
ramanujanA :: ArithmeticFunction Integer Integer
ramanujanA = forall a n.
Num a =>
(Prime n -> Word -> a) -> ArithmeticFunction n a
multiplicative forall a b. (a -> b) -> a -> b
$ Integer -> Word -> Integer
ramanujanHelper forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Prime a -> a
unPrime

ramanujanHelper :: Integer -> Word -> Integer
ramanujanHelper :: Integer -> Word -> Integer
ramanujanHelper Integer
_ Word
0 = Integer
1
ramanujanHelper Integer
2 Word
1 = -Integer
24
ramanujanHelper Integer
p Word
1 = (Integer
65 forall a. Num a => a -> a -> a
* forall a. (Num a, GcdDomain a) => a -> Word -> a
sigmaHelper (Integer
p forall a b. (Num a, Integral b) => a -> b -> a
^ (Int
11 :: Int)) Word
1 forall a. Num a => a -> a -> a
+ Integer
691 forall a. Num a => a -> a -> a
* forall a. (Num a, GcdDomain a) => a -> Word -> a
sigmaHelper (Integer
p forall a b. (Num a, Integral b) => a -> b -> a
^ (Int
5 :: Int)) Word
1 forall a. Num a => a -> a -> a
- Integer
691 forall a. Num a => a -> a -> a
* Integer
252 forall a. Num a => a -> a -> a
* Integer
2 forall a. Num a => a -> a -> a
* forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [forall n a.
(UniqueFactorisation n, Integral n, Num a, GcdDomain a) =>
Word -> n -> a
sigma Word
5 Integer
k forall a. Num a => a -> a -> a
* forall n a.
(UniqueFactorisation n, Integral n, Num a, GcdDomain a) =>
Word -> n -> a
sigma Word
5 (Integer
pforall a. Num a => a -> a -> a
-Integer
k) | Integer
k <- [Integer
1..(Integer
p forall a. Integral a => a -> a -> a
`quot` Integer
2)]]) forall a. Integral a => a -> a -> a
`quot` Integer
756
ramanujanHelper Integer
p Word
k = forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum forall a b. (a -> b) -> a -> b
$ forall a b c d. (a -> b -> c -> d) -> [a] -> [b] -> [c] -> [d]
zipWith3 (\Integer
a Integer
b Integer
c -> Integer
a forall a. Num a => a -> a -> a
* Integer
b forall a. Num a => a -> a -> a
* Integer
c) [Integer]
paPowers [Integer]
tpPowers [Integer]
binomials
  where pa :: Integer
pa = Integer
p forall a b. (Num a, Integral b) => a -> b -> a
^ (Int
11 :: Int)
        tp :: Integer
tp = Integer -> Word -> Integer
ramanujanHelper Integer
p Word
1
        paPowers :: [Integer]
paPowers = forall a. (a -> a) -> a -> [a]
iterate (forall a. Num a => a -> a -> a
* (-Integer
pa)) Integer
1
        binomials :: [Integer]
binomials = forall b a. (b -> a -> b) -> b -> [a] -> [b]
scanl (\Integer
acc Integer
j -> Integer
acc forall a. Num a => a -> a -> a
* (Integer
k' forall a. Num a => a -> a -> a
- Integer
2 forall a. Num a => a -> a -> a
* Integer
j) forall a. Num a => a -> a -> a
* (Integer
k' forall a. Num a => a -> a -> a
- Integer
2 forall a. Num a => a -> a -> a
* Integer
j forall a. Num a => a -> a -> a
- Integer
1) forall a. Integral a => a -> a -> a
`quot` (Integer
k' forall a. Num a => a -> a -> a
- Integer
j) forall a. Integral a => a -> a -> a
`quot` (Integer
j forall a. Num a => a -> a -> a
+ Integer
1)) Integer
1 [Integer
0 .. Integer
k' forall a. Integral a => a -> a -> a
`quot` Integer
2 forall a. Num a => a -> a -> a
- Integer
1]
        k' :: Integer
k' = Word -> Integer
wordToInteger Word
k
        tpPowers :: [Integer]
tpPowers = forall a. [a] -> [a]
reverse forall a b. (a -> b) -> a -> b
$ forall a. Int -> [a] -> [a]
take (forall (t :: * -> *) a. Foldable t => t a -> Int
length [Integer]
binomials) forall a b. (a -> b) -> a -> b
$ forall a. (a -> a) -> a -> [a]
iterate (forall a. Num a => a -> a -> a
* Integer
tpforall a b. (Num a, Integral b) => a -> b -> a
^(Int
2::Int)) (if forall a. Integral a => a -> Bool
even Word
k then Integer
1 else Integer
tp)
{-# INLINE ramanujanHelper #-}

-- | See 'moebiusA'.
moebius :: UniqueFactorisation n => n -> Moebius
moebius :: forall n. UniqueFactorisation n => n -> Moebius
moebius = forall n a.
UniqueFactorisation n =>
ArithmeticFunction n a -> n -> a
runFunction forall n. ArithmeticFunction n Moebius
moebiusA

-- | Calculates the Möbius function of an argument.
moebiusA :: ArithmeticFunction n Moebius
moebiusA :: forall n. ArithmeticFunction n Moebius
moebiusA = forall m n a.
Monoid m =>
(Prime n -> Word -> m) -> (m -> a) -> ArithmeticFunction n a
ArithmeticFunction (forall a b. a -> b -> a
const forall {a}. (Eq a, Num a) => a -> Moebius
f) forall a. a -> a
id
  where
    f :: a -> Moebius
f a
1 = Moebius
MoebiusN
    f a
0 = Moebius
MoebiusP
    f a
_ = Moebius
MoebiusZ

-- | See 'liouvilleA'.
liouville :: (UniqueFactorisation n, Num a) => n -> a
liouville :: forall n a. (UniqueFactorisation n, Num a) => n -> a
liouville = forall n a.
UniqueFactorisation n =>
ArithmeticFunction n a -> n -> a
runFunction forall a n. Num a => ArithmeticFunction n a
liouvilleA

-- | Calculates the Liouville function of an argument.
liouvilleA :: Num a => ArithmeticFunction n a
liouvilleA :: forall a n. Num a => ArithmeticFunction n a
liouvilleA = forall m n a.
Monoid m =>
(Prime n -> Word -> m) -> (m -> a) -> ArithmeticFunction n a
ArithmeticFunction (forall a b. a -> b -> a
const forall a b. (a -> b) -> a -> b
$ Bool -> Xor
Xor forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Integral a => a -> Bool
odd) forall a. Num a => Xor -> a
runXor

-- | See 'carmichaelA'.
carmichael :: (UniqueFactorisation n, Integral n) => n -> n
carmichael :: forall n. (UniqueFactorisation n, Integral n) => n -> n
carmichael = forall n a.
UniqueFactorisation n =>
ArithmeticFunction n a -> n -> a
runFunction forall n. Integral n => ArithmeticFunction n n
carmichaelA
{-# SPECIALIZE carmichael :: Int     -> Int #-}
{-# SPECIALIZE carmichael :: Word    -> Word #-}
{-# SPECIALIZE carmichael :: Integer -> Integer #-}
{-# SPECIALIZE carmichael :: Natural -> Natural #-}

-- | Calculates the Carmichael function for a positive integer, that is,
--   the (smallest) exponent of the group of units in @&#8484;/(n)@.
carmichaelA :: Integral n => ArithmeticFunction n n
carmichaelA :: forall n. Integral n => ArithmeticFunction n n
carmichaelA = forall m n a.
Monoid m =>
(Prime n -> Word -> m) -> (m -> a) -> ArithmeticFunction n a
ArithmeticFunction (\Prime n
p -> forall a. a -> LCM a
LCM forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall {a}. (Eq a, Num a) => a -> Word -> a
f (forall a. Prime a -> a
unPrime Prime n
p)) forall a. LCM a -> a
getLCM
  where
    f :: a -> Word -> a
f a
2 Word
1 = a
1
    f a
2 Word
2 = a
2
    f a
2 Word
k = a
2 forall a b. (Num a, Integral b) => a -> b -> a
^ Word -> Int
wordToInt (Word
k forall a. Num a => a -> a -> a
- Word
2)
    f a
p Word
1 = a
p forall a. Num a => a -> a -> a
- a
1
    f a
p Word
2 = (a
p forall a. Num a => a -> a -> a
- a
1) forall a. Num a => a -> a -> a
* a
p
    f a
p Word
k = (a
p forall a. Num a => a -> a -> a
- a
1) forall a. Num a => a -> a -> a
* a
p forall a b. (Num a, Integral b) => a -> b -> a
^ Word -> Int
wordToInt (Word
k forall a. Num a => a -> a -> a
- Word
1)

-- | Create an additive function from the function on prime's powers. See examples below.
additive :: Num a => (Prime n -> Word -> a) -> ArithmeticFunction n a
additive :: forall a n.
Num a =>
(Prime n -> Word -> a) -> ArithmeticFunction n a
additive Prime n -> Word -> a
f = forall m n a.
Monoid m =>
(Prime n -> Word -> m) -> (m -> a) -> ArithmeticFunction n a
ArithmeticFunction ((forall a. a -> Sum a
Sum forall b c a. (b -> c) -> (a -> b) -> a -> c
.) forall b c a. (b -> c) -> (a -> b) -> a -> c
. Prime n -> Word -> a
f) forall a. Sum a -> a
getSum

-- | See 'smallOmegaA'.
smallOmega :: (UniqueFactorisation n, Num a) => n -> a
smallOmega :: forall n a. (UniqueFactorisation n, Num a) => n -> a
smallOmega = forall n a.
UniqueFactorisation n =>
ArithmeticFunction n a -> n -> a
runFunction forall a n. Num a => ArithmeticFunction n a
smallOmegaA

-- | Number of distinct prime factors.
--
-- > smallOmegaA = additive (\_ _ -> 1)
smallOmegaA :: Num a => ArithmeticFunction n a
smallOmegaA :: forall a n. Num a => ArithmeticFunction n a
smallOmegaA = forall a n.
Num a =>
(Prime n -> Word -> a) -> ArithmeticFunction n a
additive forall a b. (a -> b) -> a -> b
$ forall a b. a -> b -> a
const forall a b. (a -> b) -> a -> b
$ forall a b. a -> b -> a
const a
1

-- | See 'bigOmegaA'.
bigOmega :: UniqueFactorisation n => n -> Word
bigOmega :: forall n. UniqueFactorisation n => n -> Word
bigOmega = forall n a.
UniqueFactorisation n =>
ArithmeticFunction n a -> n -> a
runFunction forall n. ArithmeticFunction n Word
bigOmegaA

-- | Number of prime factors, counted with multiplicity.
--
-- > bigOmegaA = additive (\_ k -> k)
bigOmegaA :: ArithmeticFunction n Word
bigOmegaA :: forall n. ArithmeticFunction n Word
bigOmegaA = forall a n.
Num a =>
(Prime n -> Word -> a) -> ArithmeticFunction n a
additive forall a b. (a -> b) -> a -> b
$ forall a b. a -> b -> a
const forall a. a -> a
id

-- | See 'expMangoldtA'.
expMangoldt :: UniqueFactorisation n => n -> n
expMangoldt :: forall n. UniqueFactorisation n => n -> n
expMangoldt = forall n a.
UniqueFactorisation n =>
ArithmeticFunction n a -> n -> a
runFunction forall n. Num n => ArithmeticFunction n n
expMangoldtA

-- | The exponent of von Mangoldt function. Use @log expMangoldtA@ to recover von Mangoldt function itself.
expMangoldtA :: Num n => ArithmeticFunction n n
expMangoldtA :: forall n. Num n => ArithmeticFunction n n
expMangoldtA = forall m n a.
Monoid m =>
(Prime n -> Word -> m) -> (m -> a) -> ArithmeticFunction n a
ArithmeticFunction (forall a b. a -> b -> a
const forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. a -> Mangoldt a
MangoldtOne forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Prime a -> a
unPrime) forall a. Num a => Mangoldt a -> a
runMangoldt

data Mangoldt a
  = MangoldtZero
  | MangoldtOne a
  | MangoldtMany

runMangoldt :: Num a => Mangoldt a -> a
runMangoldt :: forall a. Num a => Mangoldt a -> a
runMangoldt Mangoldt a
m = case Mangoldt a
m of
  Mangoldt a
MangoldtZero  -> a
1
  MangoldtOne a
a -> a
a
  Mangoldt a
MangoldtMany  -> a
1

instance Semigroup (Mangoldt a) where
  Mangoldt a
MangoldtZero <> :: Mangoldt a -> Mangoldt a -> Mangoldt a
<> Mangoldt a
a = Mangoldt a
a
  Mangoldt a
a <> Mangoldt a
MangoldtZero = Mangoldt a
a
  Mangoldt a
_ <> Mangoldt a
_ = forall a. Mangoldt a
MangoldtMany

instance Monoid (Mangoldt a) where
  mempty :: Mangoldt a
mempty  = forall a. Mangoldt a
MangoldtZero
  mappend :: Mangoldt a -> Mangoldt a -> Mangoldt a
mappend = forall a. Semigroup a => a -> a -> a
(<>)

-- | See 'isNFreeA'.
isNFree :: UniqueFactorisation n => Word -> n -> Bool
isNFree :: forall n. UniqueFactorisation n => Word -> n -> Bool
isNFree Word
n = forall n a.
UniqueFactorisation n =>
ArithmeticFunction n a -> n -> a
runFunction (forall n. Word -> ArithmeticFunction n Bool
isNFreeA Word
n)

-- | Check if an integer is @n@-free. An integer @x@ is @n@-free if in its
-- factorisation into prime factors, no factor has an exponent larger than or
-- equal to @n@.
isNFreeA :: Word -> ArithmeticFunction n Bool
isNFreeA :: forall n. Word -> ArithmeticFunction n Bool
isNFreeA Word
n = forall m n a.
Monoid m =>
(Prime n -> Word -> m) -> (m -> a) -> ArithmeticFunction n a
ArithmeticFunction (\Prime n
_ Word
pow -> Bool -> All
All forall a b. (a -> b) -> a -> b
$ Word
pow forall a. Ord a => a -> a -> Bool
< Word
n) All -> Bool
getAll

newtype LCM a = LCM { forall a. LCM a -> a
getLCM :: a }

instance Integral a => Semigroup (LCM a) where
  <> :: LCM a -> LCM a -> LCM a
(<>) = coerce :: forall a b. Coercible a b => a -> b
coerce (forall a. Integral a => a -> a -> a
lcm :: a -> a -> a)

instance Integral a => Monoid (LCM a) where
  mempty :: LCM a
mempty  = forall a. a -> LCM a
LCM a
1
  mappend :: LCM a -> LCM a -> LCM a
mappend = forall a. Semigroup a => a -> a -> a
(<>)

newtype Xor = Xor { Xor -> Bool
_getXor :: Bool }

runXor :: Num a => Xor -> a
runXor :: forall a. Num a => Xor -> a
runXor Xor
m = case Xor
m of
  Xor Bool
False ->  a
1
  Xor Bool
True  -> -a
1

instance Semigroup Xor where
   <> :: Xor -> Xor -> Xor
(<>) = coerce :: forall a b. Coercible a b => a -> b
coerce (forall a. Eq a => a -> a -> Bool
(/=) :: Bool -> Bool -> Bool)

instance Monoid Xor where
  mempty :: Xor
mempty  = Bool -> Xor
Xor Bool
False
  mappend :: Xor -> Xor -> Xor
mappend = forall a. Semigroup a => a -> a -> a
(<>)

newtype SetProduct a = SetProduct { forall a. SetProduct a -> Set a
getSetProduct :: Set a }

instance (Num a, Ord a) => Semigroup (SetProduct a) where
  SetProduct Set a
s1 <> :: SetProduct a -> SetProduct a -> SetProduct a
<> SetProduct Set a
s2 = forall a. Set a -> SetProduct a
SetProduct forall a b. (a -> b) -> a -> b
$ Set a
s1 forall a. Semigroup a => a -> a -> a
<> Set a
s2 forall a. Semigroup a => a -> a -> a
<> forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (\a
n -> forall a b. (a -> b) -> Set a -> Set b
S.mapMonotonic (forall a. Num a => a -> a -> a
* a
n) Set a
s2) Set a
s1

instance (Num a, Ord a) => Monoid (SetProduct a) where
  mempty :: SetProduct a
mempty  = forall a. Set a -> SetProduct a
SetProduct forall a. Monoid a => a
mempty
  mappend :: SetProduct a -> SetProduct a -> SetProduct a
mappend = forall a. Semigroup a => a -> a -> a
(<>)

newtype ListProduct a = ListProduct { forall a. ListProduct a -> [a]
getListProduct :: [a] }

instance Num a => Semigroup (ListProduct a) where
  ListProduct [a]
s1 <> :: ListProduct a -> ListProduct a -> ListProduct a
<> ListProduct [a]
s2 = forall a. [a] -> ListProduct a
ListProduct forall a b. (a -> b) -> a -> b
$ [a]
s1 forall a. Semigroup a => a -> a -> a
<> [a]
s2 forall a. Semigroup a => a -> a -> a
<> forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (\a
n -> forall a b. (a -> b) -> [a] -> [b]
map (forall a. Num a => a -> a -> a
* a
n) [a]
s2) [a]
s1

instance Num a => Monoid (ListProduct a) where
  mempty :: ListProduct a
mempty  = forall a. [a] -> ListProduct a
ListProduct forall a. Monoid a => a
mempty
  mappend :: ListProduct a -> ListProduct a -> ListProduct a
mappend = forall a. Semigroup a => a -> a -> a
(<>)

-- Represent as a Reader monad
newtype BoundedSetProduct a = BoundedSetProduct { forall a. BoundedSetProduct a -> a -> Set a
_getBoundedSetProduct :: a -> Set a }

takeWhileLE :: Ord a => a -> Set a -> Set a
takeWhileLE :: forall a. Ord a => a -> Set a -> Set a
takeWhileLE a
b Set a
xs = if Bool
m then forall a. Ord a => a -> Set a -> Set a
S.insert a
b Set a
ls else Set a
ls
  where (Set a
ls, Bool
m, Set a
_) = forall a. Ord a => a -> Set a -> (Set a, Bool, Set a)
S.splitMember a
b Set a
xs

instance (Ord a, Num a) => Semigroup (BoundedSetProduct a) where
  BoundedSetProduct a -> Set a
f1 <> :: BoundedSetProduct a -> BoundedSetProduct a -> BoundedSetProduct a
<> BoundedSetProduct a -> Set a
f2 = forall a. (a -> Set a) -> BoundedSetProduct a
BoundedSetProduct a -> Set a
f
    where f :: a -> Set a
f a
b = Set a
s1 forall a. Semigroup a => a -> a -> a
<> Set a
s2 forall a. Semigroup a => a -> a -> a
<> forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (\a
n -> forall a. Ord a => a -> Set a -> Set a
takeWhileLE a
b forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> Set a -> Set b
S.mapMonotonic (forall a. Num a => a -> a -> a
* a
n) Set a
s2) Set a
s1
            where s1 :: Set a
s1 = a -> Set a
f1 a
b
                  s2 :: Set a
s2 = a -> Set a
f2 a
b

instance (Ord a, Num a) => Monoid (BoundedSetProduct a) where
  mempty :: BoundedSetProduct a
mempty = forall a. (a -> Set a) -> BoundedSetProduct a
BoundedSetProduct forall a. Monoid a => a
mempty
  mappend :: BoundedSetProduct a -> BoundedSetProduct a -> BoundedSetProduct a
mappend = forall a. Semigroup a => a -> a -> a
(<>)

newtype IntSetProduct = IntSetProduct { IntSetProduct -> IntSet
getIntSetProduct :: IntSet }

instance Semigroup IntSetProduct where
  IntSetProduct IntSet
s1 <> :: IntSetProduct -> IntSetProduct -> IntSetProduct
<> IntSetProduct IntSet
s2 = IntSet -> IntSetProduct
IntSetProduct forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *). Foldable f => f IntSet -> IntSet
IS.unions forall a b. (a -> b) -> a -> b
$ IntSet
s1 forall a. a -> [a] -> [a]
: IntSet
s2 forall a. a -> [a] -> [a]
: forall a b. (a -> b) -> [a] -> [b]
map (\Int
n -> (Int -> Int) -> IntSet -> IntSet
IS.map (forall a. Num a => a -> a -> a
* Int
n) IntSet
s2) (IntSet -> [Int]
IS.toAscList IntSet
s1)

instance Monoid IntSetProduct where
  mempty :: IntSetProduct
mempty  = IntSet -> IntSetProduct
IntSetProduct forall a. Monoid a => a
mempty
  mappend :: IntSetProduct -> IntSetProduct -> IntSetProduct
mappend = forall a. Semigroup a => a -> a -> a
(<>)