-- |
-- Module:      Math.NumberTheory.ArithmeticFunctions.Mertens
-- Copyright:   (c) 2018 Andrew Lelechenko
-- Licence:     MIT
-- Maintainer:  Andrew Lelechenko <andrew.lelechenko@gmail.com>
--
-- Values of <https://en.wikipedia.org/wiki/Mertens_function Mertens function>.
--

{-# LANGUAGE LambdaCase #-}

module Math.NumberTheory.ArithmeticFunctions.Mertens
  ( mertens
  ) where

import qualified Data.Vector.Unboxed as U

import Math.NumberTheory.Roots
import Math.NumberTheory.ArithmeticFunctions.Moebius
import Math.NumberTheory.Utils.FromIntegral

-- | Compute individual values of Mertens function in O(n^(2/3)) time and space.
--
-- >>> map (mertens . (10 ^)) [0..9]
-- [1,-1,1,2,-23,-48,212,1037,1928,-222]
--
-- The implementation follows Theorem 3.1 from <https://arxiv.org/pdf/1610.08551.pdf Computations of the Mertens function and improved bounds on the Mertens conjecture> by G. Hurst, excluding segmentation of sieves.
mertens :: Word -> Int
mertens :: Word -> Int
mertens Word
0 = Int
0
mertens Word
1 = Int
1
mertens Word
x = (Word -> Moebius) -> (Word -> Int) -> [Word] -> Int
sumMultMoebius Word -> Moebius
lookupMus (\Word
n -> Word -> Int
sfunc (Word
x Word -> Word -> Word
forall a. Integral a => a -> a -> a
`quot` Word
n)) [Word
1 .. Word
x Word -> Word -> Word
forall a. Integral a => a -> a -> a
`quot` Word
u]
  where
    u :: Word
u = (Word -> Word
forall a. Integral a => a -> a
integerSquareRoot Word
x Word -> Word -> Word
forall a. Num a => a -> a -> a
+ Word
1) Word -> Word -> Word
forall a. Ord a => a -> a -> a
`max` (Word -> Word
forall a. Integral a => a -> a
integerCubeRoot Word
x Word -> Word -> Word
forall a b. (Num a, Integral b) => a -> b -> a
^ (Word
2 :: Word) Word -> Word -> Word
forall a. Integral a => a -> a -> a
`quot` Word
2)

    sfunc :: Word -> Int
    sfunc :: Word -> Int
sfunc Word
y
      = Int
1
      Int -> Int -> Int
forall a. Num a => a -> a -> a
- [Int] -> Int
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [ Vector Int -> Int -> Int
forall a. Unbox a => Vector a -> Int -> a
U.unsafeIndex Vector Int
mes (Word -> Int
wordToInt (Word -> Int) -> Word -> Int
forall a b. (a -> b) -> a -> b
$ Word
y Word -> Word -> Word
forall a. Integral a => a -> a -> a
`quot` Word
n) |  Word
n <- [Word
y Word -> Word -> Word
forall a. Integral a => a -> a -> a
`quot` Word
u Word -> Word -> Word
forall a. Num a => a -> a -> a
+ Word
1 .. Word
kappa] ]
      Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Word -> Int
wordToInt Word
kappa Int -> Int -> Int
forall a. Num a => a -> a -> a
* Vector Int -> Int -> Int
forall a. Unbox a => Vector a -> Int -> a
U.unsafeIndex Vector Int
mes (Word -> Int
wordToInt Word
nu)
      Int -> Int -> Int
forall a. Num a => a -> a -> a
- (Word -> Moebius) -> (Word -> Int) -> [Word] -> Int
sumMultMoebius Word -> Moebius
lookupMus (\Word
n -> Word -> Int
wordToInt (Word -> Int) -> Word -> Int
forall a b. (a -> b) -> a -> b
$ Word
y Word -> Word -> Word
forall a. Integral a => a -> a -> a
`quot` Word
n) [Word
1 .. Word
nu]
      where
        nu :: Word
nu = Word -> Word
forall a. Integral a => a -> a
integerSquareRoot Word
y
        kappa :: Word
kappa = Word
y Word -> Word -> Word
forall a. Integral a => a -> a -> a
`quot` (Word
nu Word -> Word -> Word
forall a. Num a => a -> a -> a
+ Word
1)

    -- cacheSize ~ u
    cacheSize :: Word
    cacheSize :: Word
cacheSize = Word
u Word -> Word -> Word
forall a. Ord a => a -> a -> a
`max` (Word
x Word -> Word -> Word
forall a. Integral a => a -> a -> a
`quot` Word
u) Word -> Word -> Word
forall a. Ord a => a -> a -> a
`max` Word -> Word
forall a. Integral a => a -> a
integerSquareRoot Word
x

    -- 1-based index
    mus :: U.Vector Moebius
    mus :: Vector Moebius
mus = Word -> Word -> Vector Moebius
sieveBlockMoebius Word
1 Word
cacheSize

    lookupMus :: Word -> Moebius
    lookupMus :: Word -> Moebius
lookupMus Word
i = Vector Moebius -> Int -> Moebius
forall a. Unbox a => Vector a -> Int -> a
U.unsafeIndex Vector Moebius
mus (Word -> Int
wordToInt (Word
i Word -> Word -> Word
forall a. Num a => a -> a -> a
- Word
1))

    -- 0-based index
    mes :: U.Vector Int
    mes :: Vector Int
mes = (Int -> Moebius -> Int) -> Int -> Vector Moebius -> Vector Int
forall a b.
(Unbox a, Unbox b) =>
(a -> b -> a) -> a -> Vector b -> Vector a
U.scanl' Int -> Moebius -> Int
forall a. Num a => a -> Moebius -> a
go Int
0 Vector Moebius
mus
      where
        go :: a -> Moebius -> a
go a
acc = \case
          Moebius
MoebiusN -> a
acc a -> a -> a
forall a. Num a => a -> a -> a
- a
1
          Moebius
MoebiusZ -> a
acc
          Moebius
MoebiusP -> a
acc a -> a -> a
forall a. Num a => a -> a -> a
+ a
1

-- | Compute sum (map (\x -> runMoebius (mu x) * f x))
sumMultMoebius :: (Word -> Moebius) -> (Word -> Int) -> [Word] -> Int
sumMultMoebius :: (Word -> Moebius) -> (Word -> Int) -> [Word] -> Int
sumMultMoebius Word -> Moebius
mu Word -> Int
f = (Int -> Word -> Int) -> Int -> [Word] -> Int
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl Int -> Word -> Int
go Int
0
  where
    go :: Int -> Word -> Int
go Int
acc Word
i = case Word -> Moebius
mu Word
i of
      Moebius
MoebiusN -> Int
acc Int -> Int -> Int
forall a. Num a => a -> a -> a
- Word -> Int
f Word
i
      Moebius
MoebiusZ -> Int
acc
      Moebius
MoebiusP -> Int
acc Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Word -> Int
f Word
i