-- |
-- Module:      Math.NumberTheory.MoebiusInversion
-- Copyright:   (c) 2012 Daniel Fischer
-- Licence:     MIT
-- Maintainer:  Daniel Fischer <daniel.is.fischer@googlemail.com>
--
-- Generalised Möbius inversion

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

module Math.NumberTheory.MoebiusInversion
    ( generalInversion
    , totientSum
    ) where

import Control.Monad
import Control.Monad.ST
import Data.Proxy
import qualified Data.Vector.Generic as G
import qualified Data.Vector.Generic.Mutable as MG

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

-- | @totientSum n@ is, for @n > 0@, the sum of @[totient k | k <- [1 .. n]]@,
--   computed via generalised Möbius inversion.
--   See <http://mathworld.wolfram.com/TotientSummatoryFunction.html> for the
--   formula used for @totientSum@.
--
-- >>> import Data.Proxy
-- >>> totientSum (Proxy :: Proxy Data.Vector.Unboxed.Vector) 100 :: Int
-- 3044
-- >>> totientSum (Proxy :: Proxy Data.Vector.Vector) 100 :: Integer
-- 3044
totientSum
    :: (Integral t, G.Vector v t)
    => Proxy v
    -> Word
    -> t
totientSum :: forall t (v :: * -> *).
(Integral t, Vector v t) =>
Proxy v -> Word -> t
totientSum Proxy v
_ Word
0 = t
0
totientSum Proxy v
proxy Word
n = forall t (v :: * -> *).
(Num t, Vector v t) =>
Proxy v -> (Word -> t) -> Word -> t
generalInversion Proxy v
proxy (forall {a}. Integral a => a -> a
triangle forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (Integral a, Num b) => a -> b
fromIntegral) Word
n
  where
    triangle :: a -> a
triangle a
k = (a
k forall a. Num a => a -> a -> a
* (a
k forall a. Num a => a -> a -> a
+ a
1)) forall a. Integral a => a -> a -> a
`quot` a
2

-- | @generalInversion g n@ evaluates the generalised Möbius inversion of @g@
--   at the argument @n@.
--
--   The generalised Möbius inversion implemented here allows an efficient
--   calculation of isolated values of the function @f : N -> Z@ if the function
--   @g@ defined by
--
-- >
-- > g n = sum [f (n `quot` k) | k <- [1 .. n]]
-- >
--
--   can be cheaply computed. By the generalised Möbius inversion formula, then
--
-- >
-- > f n = sum [moebius k * g (n `quot` k) | k <- [1 .. n]]
-- >
--
--   which allows the computation in /O/(n) steps, if the values of the
--   Möbius function are known. A slightly different formula, used here,
--   does not need the values of the Möbius function and allows the
--   computation in /O/(n^0.75) steps, using /O/(n^0.5) memory.
--
--   An example of a pair of such functions where the inversion allows a
--   more efficient computation than the direct approach is
--
-- >
-- > f n = number of reduced proper fractions with denominator <= n
-- >
-- > g n = number of proper fractions with denominator <= n
-- >
--
--   (a /proper fraction/ is a fraction @0 < n/d < 1@). Then @f n@ is the
--   cardinality of the Farey sequence of order @n@ (minus 1 or 2 if 0 and/or
--   1 are included in the Farey sequence), or the sum of the totients of
--   the numbers @2 <= k <= n@. @f n@ is not easily directly computable,
--   but then @g n = n*(n-1)/2@ is very easy to compute, and hence the inversion
--   gives an efficient method of computing @f n@.
--
--   Since the function arguments are used as array indices, the domain of
--   @f@ is restricted to 'Int'.
--
--   The value @f n@ is then computed by @generalInversion g n@. Note that when
--   many values of @f@ are needed, there are far more efficient methods, this
--   method is only appropriate to compute isolated values of @f@.
generalInversion
    :: (Num t, G.Vector v t)
    => Proxy v
    -> (Word -> t)
    -> Word
    -> t
generalInversion :: forall t (v :: * -> *).
(Num t, Vector v t) =>
Proxy v -> (Word -> t) -> Word -> t
generalInversion Proxy v
proxy Word -> t
fun Word
n = case Word
n of
    Word
0 ->forall a. HasCallStack => [Char] -> a
error [Char]
"Möbius inversion only defined on positive domain"
    Word
1 -> Word -> t
fun Word
1
    Word
2 -> Word -> t
fun Word
2 forall a. Num a => a -> a -> a
- Word -> t
fun Word
1
    Word
3 -> Word -> t
fun Word
3 forall a. Num a => a -> a -> a
- t
2forall a. Num a => a -> a -> a
*Word -> t
fun Word
1
    Word
_ -> forall a. (forall s. ST s a) -> a
runST (forall s t (v :: * -> *).
(Num t, Vector v t) =>
Proxy v -> (Int -> t) -> Int -> ST s t
fastInvertST Proxy v
proxy (Word -> t
fun forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Word
intToWord) (Word -> Int
wordToInt Word
n))

fastInvertST
    :: forall s t v.
       (Num t, G.Vector v t)
    => Proxy v
    -> (Int -> t)
    -> Int
    -> ST s t
fastInvertST :: forall s t (v :: * -> *).
(Num t, Vector v t) =>
Proxy v -> (Int -> t) -> Int -> ST s t
fastInvertST Proxy v
_ Int -> t
fun Int
n = do
    let !k0 :: Int
k0 = forall {a}. Integral a => a -> a
integerSquareRoot (Int
n forall a. Integral a => a -> a -> a
`quot` Int
2)
        !mk0 :: Int
mk0 = Int
n forall a. Integral a => a -> a -> a
`quot` (Int
2forall a. Num a => a -> a -> a
*Int
k0forall a. Num a => a -> a -> a
+Int
1)
        kmax :: a -> a -> a
kmax a
a a
m = (a
a forall a. Integral a => a -> a -> a
`quot` a
m forall a. Num a => a -> a -> a
- a
1) forall a. Integral a => a -> a -> a
`quot` a
2

    Mutable v s t
small <- forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
Int -> m (v (PrimState m) a)
MG.unsafeNew (Int
mk0 forall a. Num a => a -> a -> a
+ Int
1) :: ST s (G.Mutable v s t)
    forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
MG.unsafeWrite Mutable v s t
small Int
0 t
0
    forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
MG.unsafeWrite Mutable v s t
small Int
1 forall a b. (a -> b) -> a -> b
$! Int -> t
fun Int
1
    forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
mk0 forall a. Ord a => a -> a -> Bool
>= Int
2) forall a b. (a -> b) -> a -> b
$
        forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
MG.unsafeWrite Mutable v s t
small Int
2 forall a b. (a -> b) -> a -> b
$! (Int -> t
fun Int
2 forall a. Num a => a -> a -> a
- Int -> t
fun Int
1)

    let calcit :: Int -> Int -> Int -> ST s (Int, Int)
        calcit :: Int -> Int -> Int -> ST s (Int, Int)
calcit Int
switch Int
change Int
i
            | Int
mk0 forall a. Ord a => a -> a -> Bool
< Int
i   = forall (m :: * -> *) a. Monad m => a -> m a
return (Int
switch,Int
change)
            | Int
i forall a. Eq a => a -> a -> Bool
== Int
change = Int -> Int -> Int -> ST s (Int, Int)
calcit (Int
switchforall a. Num a => a -> a -> a
+Int
1) (Int
change forall a. Num a => a -> a -> a
+ Int
4forall a. Num a => a -> a -> a
*Int
switchforall a. Num a => a -> a -> a
+Int
6) Int
i
            | Bool
otherwise = do
                let mloop :: t -> Int -> Int -> ST s (Int, Int)
mloop !t
acc Int
k !Int
m
                        | Int
k forall a. Ord a => a -> a -> Bool
< Int
switch    = t -> Int -> ST s (Int, Int)
kloop t
acc Int
k
                        | Bool
otherwise     = do
                            t
val <- forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
MG.unsafeRead Mutable v s t
small Int
m
                            let nxtk :: Int
nxtk = forall a. Integral a => a -> a -> a
kmax Int
i (Int
mforall a. Num a => a -> a -> a
+Int
1)
                            t -> Int -> Int -> ST s (Int, Int)
mloop (t
acc forall a. Num a => a -> a -> a
- forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int
kforall a. Num a => a -> a -> a
-Int
nxtk)forall a. Num a => a -> a -> a
*t
val) Int
nxtk (Int
mforall a. Num a => a -> a -> a
+Int
1)
                    kloop :: t -> Int -> ST s (Int, Int)
kloop !t
acc Int
k
                        | Int
k forall a. Eq a => a -> a -> Bool
== Int
0    = do
                            forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
MG.unsafeWrite Mutable v s t
small Int
i forall a b. (a -> b) -> a -> b
$! t
acc
                            Int -> Int -> Int -> ST s (Int, Int)
calcit Int
switch Int
change (Int
iforall a. Num a => a -> a -> a
+Int
1)
                        | Bool
otherwise = do
                            t
val <- forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
MG.unsafeRead Mutable v s t
small (Int
i forall a. Integral a => a -> a -> a
`quot` (Int
2forall a. Num a => a -> a -> a
*Int
kforall a. Num a => a -> a -> a
+Int
1))
                            t -> Int -> ST s (Int, Int)
kloop (t
accforall a. Num a => a -> a -> a
-t
val) (Int
kforall a. Num a => a -> a -> a
-Int
1)
                t -> Int -> Int -> ST s (Int, Int)
mloop (Int -> t
fun Int
i forall a. Num a => a -> a -> a
- Int -> t
fun (Int
i forall a. Integral a => a -> a -> a
`quot` Int
2)) ((Int
iforall a. Num a => a -> a -> a
-Int
1) forall a. Integral a => a -> a -> a
`quot` Int
2) Int
1

    (Int
sw, Int
ch) <- Int -> Int -> Int -> ST s (Int, Int)
calcit Int
1 Int
8 Int
3
    Mutable v s t
large <- forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
Int -> m (v (PrimState m) a)
MG.unsafeNew Int
k0 :: ST s (G.Mutable v s t)

    let calcbig :: Int -> Int -> Int -> ST s (G.Mutable v s t)
        calcbig :: Int -> Int -> Int -> ST s (Mutable v s t)
calcbig Int
switch Int
change Int
j
            | Int
j forall a. Eq a => a -> a -> Bool
== Int
0    = forall (m :: * -> *) a. Monad m => a -> m a
return Mutable v s t
large
            | (Int
2forall a. Num a => a -> a -> a
*Int
jforall a. Num a => a -> a -> a
-Int
1)forall a. Num a => a -> a -> a
*Int
change forall a. Ord a => a -> a -> Bool
<= Int
n   = Int -> Int -> Int -> ST s (Mutable v s t)
calcbig (Int
switchforall a. Num a => a -> a -> a
+Int
1) (Int
change forall a. Num a => a -> a -> a
+ Int
4forall a. Num a => a -> a -> a
*Int
switchforall a. Num a => a -> a -> a
+Int
6) Int
j
            | Bool
otherwise = do
                let i :: Int
i = Int
n forall a. Integral a => a -> a -> a
`quot` (Int
2forall a. Num a => a -> a -> a
*Int
jforall a. Num a => a -> a -> a
-Int
1)
                    mloop :: t -> Int -> Int -> ST s (Mutable v s t)
mloop !t
acc Int
k Int
m
                        | Int
k forall a. Ord a => a -> a -> Bool
< Int
switch    = t -> Int -> ST s (Mutable v s t)
kloop t
acc Int
k
                        | Bool
otherwise     = do
                            t
val <- forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
MG.unsafeRead Mutable v s t
small Int
m
                            let nxtk :: Int
nxtk = forall a. Integral a => a -> a -> a
kmax Int
i (Int
mforall a. Num a => a -> a -> a
+Int
1)
                            t -> Int -> Int -> ST s (Mutable v s t)
mloop (t
acc forall a. Num a => a -> a -> a
- forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int
kforall a. Num a => a -> a -> a
-Int
nxtk)forall a. Num a => a -> a -> a
*t
val) Int
nxtk (Int
mforall a. Num a => a -> a -> a
+Int
1)
                    kloop :: t -> Int -> ST s (Mutable v s t)
kloop !t
acc Int
k
                        | Int
k forall a. Eq a => a -> a -> Bool
== Int
0    = do
                            forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
MG.unsafeWrite Mutable v s t
large (Int
jforall a. Num a => a -> a -> a
-Int
1) forall a b. (a -> b) -> a -> b
$! t
acc
                            Int -> Int -> Int -> ST s (Mutable v s t)
calcbig Int
switch Int
change (Int
jforall a. Num a => a -> a -> a
-Int
1)
                        | Bool
otherwise = do
                            let m :: Int
m = Int
i forall a. Integral a => a -> a -> a
`quot` (Int
2forall a. Num a => a -> a -> a
*Int
kforall a. Num a => a -> a -> a
+Int
1)
                            t
val <- if Int
m forall a. Ord a => a -> a -> Bool
<= Int
mk0
                                     then forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
MG.unsafeRead Mutable v s t
small Int
m
                                     else forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
MG.unsafeRead Mutable v s t
large (Int
kforall a. Num a => a -> a -> a
*(Int
2forall a. Num a => a -> a -> a
*Int
jforall a. Num a => a -> a -> a
-Int
1)forall a. Num a => a -> a -> a
+Int
jforall a. Num a => a -> a -> a
-Int
1)
                            t -> Int -> ST s (Mutable v s t)
kloop (t
accforall a. Num a => a -> a -> a
-t
val) (Int
kforall a. Num a => a -> a -> a
-Int
1)
                t -> Int -> Int -> ST s (Mutable v s t)
mloop (Int -> t
fun Int
i forall a. Num a => a -> a -> a
- Int -> t
fun (Int
i forall a. Integral a => a -> a -> a
`quot` Int
2)) ((Int
iforall a. Num a => a -> a -> a
-Int
1) forall a. Integral a => a -> a -> a
`quot` Int
2) Int
1

    Mutable v s t
mvec <- Int -> Int -> Int -> ST s (Mutable v s t)
calcbig Int
sw Int
ch Int
k0
    forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
MG.unsafeRead Mutable v s t
mvec Int
0