module Math.Prime where

import qualified Data.IntMap as IM
import qualified Data.List as L
import qualified Data.List.NonEmpty as NE

smallPrimes :: (Integral i) => [i]
smallPrimes :: forall i. Integral i => [i]
smallPrimes = i
2 i -> [i] -> [i]
forall a. a -> [a] -> [a]
: [i
n | i
n <- [i
3, i
5 .. i
46337], (i -> Bool) -> [i] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all ((i -> i -> Bool
forall a. Ord a => a -> a -> Bool
> i
0) (i -> Bool) -> (i -> i) -> i -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. i -> i -> i
forall a. Integral a => a -> a -> a
rem i
n) ([i] -> Bool) -> [i] -> Bool
forall a b. (a -> b) -> a -> b
$ (i -> Bool) -> [i] -> [i]
forall a. (a -> Bool) -> [a] -> [a]
takeWhile (\i
x -> i
x i -> i -> i
forall a. Num a => a -> a -> a
* i
x i -> i -> Bool
forall a. Ord a => a -> a -> Bool
<= i
n) [i]
forall i. Integral i => [i]
smallPrimes]
{-# SPECIALIZE smallPrimes :: [Int] #-}

{- |
>>> primeFactors 60
[2,2,3,5]
>>> primeFactors 0
[]
>>> primeFactors 1
[]
>>> primeFactors 2
[2]
>>> primeFactors 2147483647
[2147483647]
>>> primeFactors 999999999989
[999999999989]
>>> primeFactors 999999999997
[5507,181587071]
-}
primeFactors :: (Integral i) => i -> [i]
primeFactors :: forall i. Integral i => i -> [i]
primeFactors i
n | i
n i -> i -> Bool
forall a. Ord a => a -> a -> Bool
< i
2 = []
primeFactors i
n0 = i -> [i] -> [i]
forall {t}. Integral t => t -> [t] -> [t]
go i
n0 [i]
forall i. Integral i => [i]
smallPrimes
  where
    go :: t -> [t] -> [t]
go !t
n pps :: [t]
pps@(t
p : [t]
ps)
      | t
n t -> t -> Bool
forall a. Ord a => a -> a -> Bool
< t
p t -> t -> t
forall a. Num a => a -> a -> a
* t
p = [t
n]
      | t
r t -> t -> Bool
forall a. Ord a => a -> a -> Bool
> t
0 = t -> [t] -> [t]
go t
n [t]
ps
      | Bool
otherwise = t
p t -> [t] -> [t]
forall a. a -> [a] -> [a]
: t -> [t] -> [t]
go t
q [t]
pps
      where
        (t
q, t
r) = t -> t -> (t, t)
forall a. Integral a => a -> a -> (a, a)
quotRem t
n t
p
    go t
n [] = t -> [t] -> [t]
go t
n [t
46339, t
46341 ..]
{-# SPECIALIZE primeFactors :: Int -> [Int] #-}

{- |
>>> isPrime 4649
True
>>> isPrime 0
False
>>> isPrime 1
False
>>> isPrime 2
True
>>> isPrime 2147483647
True
>>> isPrime 999999999989
True
>>> isPrime 999999999997
False
-}
isPrime :: (Integral i) => i -> Bool
isPrime :: forall i. Integral i => i -> Bool
isPrime i
n = [i
n] [i] -> [i] -> Bool
forall a. Eq a => a -> a -> Bool
== i -> [i]
forall i. Integral i => i -> [i]
primeFactors i
n
{-# SPECIALIZE isPrime :: Int -> Bool #-}

{- |
prop> \n -> not (n >= 0) || totient n == length [x|x<-[1..n], gcd n x == 1]
+++ OK, passed 100 tests.

>>> totient 0
0
>>> totient 1
1
-}
totient :: Int -> Int
totient :: Int -> Int
totient Int
n = Int
n Int -> Int -> Int
forall a. Integral a => a -> a -> a
`quot` [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product [Int]
ps Int -> Int -> Int
forall a. Num a => a -> a -> a
* [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product ((Int -> Int) -> [Int] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (Int -> Int -> Int
forall a. Num a => a -> a -> a
subtract Int
1) [Int]
ps)
  where
    ps :: [Int]
ps = (NonEmpty Int -> Int) -> [NonEmpty Int] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map NonEmpty Int -> Int
forall a. NonEmpty a -> a
NE.head ([NonEmpty Int] -> [Int])
-> ([Int] -> [NonEmpty Int]) -> [Int] -> [Int]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Int] -> [NonEmpty Int]
forall (f :: * -> *) a. (Foldable f, Eq a) => f a -> [NonEmpty a]
NE.group ([Int] -> [Int]) -> [Int] -> [Int]
forall a b. (a -> b) -> a -> b
$ Int -> [Int]
forall i. Integral i => i -> [i]
primeFactors Int
n

{- |
>>> divisors 12
[1,2,3,4,6,12]
>>> divisors 0
[1]
>>> divisors 1
[1]
>>> length (divisors 735134400)
1344
-}
divisors :: Int -> [Int]
divisors :: Int -> [Int]
divisors Int
n = [Int] -> [Int]
forall a. Ord a => [a] -> [a]
L.sort ([Int] -> [Int]) -> ([Int] -> [Int]) -> [Int] -> [Int]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([Int] -> Int) -> [[Int]] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product ([[Int]] -> [Int]) -> ([Int] -> [[Int]]) -> [Int] -> [Int]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([Int] -> [Int]) -> [[Int]] -> [[Int]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM ((Int -> Int -> Int) -> Int -> [Int] -> [Int]
forall b a. (b -> a -> b) -> b -> [a] -> [b]
scanl Int -> Int -> Int
forall a. Num a => a -> a -> a
(*) Int
1) ([[Int]] -> [[Int]]) -> ([Int] -> [[Int]]) -> [Int] -> [[Int]]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Int] -> [[Int]]
forall a. Eq a => [a] -> [[a]]
L.group ([Int] -> [Int]) -> [Int] -> [Int]
forall a b. (a -> b) -> a -> b
$ Int -> [Int]
forall i. Integral i => i -> [i]
primeFactors Int
n

{- |
>>> moebius 1
1
>>> moebius 2
-1
>>> moebius 3
-1
>>> moebius (2 * 2)
0
>>> moebius (2 * 3)
1
>>> moebius (2 * 2 * 3)
0
-}
moebius :: Int -> Int
moebius :: Int -> Int
moebius Int
n = [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product ([Int] -> Int) -> ([Int] -> [Int]) -> [Int] -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([Int] -> Int) -> [[Int]] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map [Int] -> Int
forall {a} {a}. Num a => [a] -> a
f ([[Int]] -> [Int]) -> ([Int] -> [[Int]]) -> [Int] -> [Int]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Int] -> [[Int]]
forall a. Eq a => [a] -> [[a]]
L.group ([Int] -> Int) -> [Int] -> Int
forall a b. (a -> b) -> a -> b
$ Int -> [Int]
forall i. Integral i => i -> [i]
primeFactors Int
n
  where
    f :: [a] -> a
f [a
_] = -a
1
    f [a]
_ = a
0

{- |
>>> moebiusInversion 12 (length.divisors)
fromList [(1,1),(2,1),(3,1),(4,1),(6,1),(12,1)]
>>> moebiusInversion 999999999997 (length.divisors)
fromList [(1,1),(5507,1),(181587071,1),(999999999997,1)]
>>> IM.size $ moebiusInversion 735134400 (length.divisors)
1344
-}
moebiusInversion :: (Num a) => Int -> (Int -> a) -> IM.IntMap a
moebiusInversion :: forall a. Num a => Int -> (Int -> a) -> IntMap a
moebiusInversion Int
n Int -> a
f = (IntMap a -> Int -> IntMap a) -> IntMap a -> [Int] -> IntMap a
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
L.foldl' IntMap a -> Int -> IntMap a
forall {b}. Num b => IntMap b -> Int -> IntMap b
step ([(Int, a)] -> IntMap a
forall a. [(Int, a)] -> IntMap a
IM.fromAscList [(Int
d, Int -> a
f Int
d) | Int
d <- [Int]
ds]) [Int]
ds
  where
    !ds :: [Int]
ds = Int -> [Int]
divisors Int
n
    step :: IntMap b -> Int -> IntMap b
step IntMap b
g Int
d =
      let !gd :: b
gd = IntMap b
g IntMap b -> Int -> b
forall a. IntMap a -> Int -> a
IM.! Int
d
       in (Int -> b -> b) -> IntMap b -> IntMap b
forall a b. (Int -> a -> b) -> IntMap a -> IntMap b
IM.mapWithKey
            ( \Int
k b
v ->
                if Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
d Bool -> Bool -> Bool
&& Int -> Int -> Int
forall a. Integral a => a -> a -> a
rem Int
k Int
d Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 then b
v b -> b -> b
forall a. Num a => a -> a -> a
- b
gd else b
v
            )
            IntMap b
g
{-# INLINE moebiusInversion #-}