    -- |
    -- 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 a. Num a => [a] -> a
    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 b a. (b -> a -> b) -> b -> [a] -> b
    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