-- |
-- 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 CPP        #-}
{-# LANGUAGE LambdaCase #-}

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

import qualified Data.Vector.Unboxed as U

import Math.NumberTheory.Powers.Cubes
import Math.NumberTheory.Powers.Squares
import Math.NumberTheory.ArithmeticFunctions.Moebius

-- | 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 0 = 0
mertens 1 = 1
mertens x = sumMultMoebius lookupMus (\n -> sfunc (x `quot` n)) [1 .. x `quot` u]
  where
    u = (integerSquareRoot x + 1) `max` ((integerCubeRoot x) ^ (2 :: Word) `quot` 2)

    sfunc :: Word -> Int
    sfunc y
      = 1
      - sum [ U.unsafeIndex mes (fromIntegral $ y `quot` n) |  n <- [y `quot` u + 1 .. kappa] ]
      + fromIntegral kappa * U.unsafeIndex mes (fromIntegral nu)
      - sumMultMoebius lookupMus (\n -> fromIntegral $ y `quot` n) [1 .. nu]
      where
        nu = integerSquareRoot y
        kappa = y `quot` (nu + 1)

    -- cacheSize ~ u
    cacheSize :: Word
    cacheSize = u `max` (x `quot` u) `max` integerSquareRoot x

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

    lookupMus :: Word -> Moebius
    lookupMus i = U.unsafeIndex mus (fromIntegral (i - 1))

    -- 0-based index
    mes :: U.Vector Int
    mes = U.scanl' go 0 mus
      where
        go acc = \case
          MoebiusN -> acc - 1
          MoebiusZ -> acc
          MoebiusP -> acc + 1

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