content: -- |
  -- Module:      Math.NumberTheory.Primes.Counting.Impl
  -- Copyright:   (c) 2011 Daniel Fischer
  -- Licence:     MIT
  -- Maintainer:  Daniel Fischer <daniel.is.fischer@googlemail.com>
  --
  -- Number of primes not exceeding @n@, @&#960;(n)@, and @n@-th prime.
  --
  {-# LANGUAGE BangPatterns        #-}
  {-# LANGUAGE FlexibleContexts    #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  
  {-# OPTIONS_GHC -fspec-constr-count=24 #-}
  {-# OPTIONS_GHC -Wno-incomplete-uni-patterns #-}
  
  module Math.NumberTheory.Primes.Counting.Impl
      ( primeCount
      , primeCountMaxArg
      , nthPrime
      ) where
  
  import Math.NumberTheory.Primes.Sieve.Eratosthenes
      (PrimeSieve(..), primeList, primeSieve, psieveFrom, sieveTo, sieveBits, sieveRange)
  import Math.NumberTheory.Primes.Sieve.Indexing (toPrim, idxPr)
  import Math.NumberTheory.Primes.Counting.Approximate (nthPrimeApprox, approxPrimeCount)
  import Math.NumberTheory.Primes.Types
  import Math.NumberTheory.Roots
  import Math.NumberTheory.Utils.FromIntegral
  
  import Control.Monad.ST
  import Data.Bit
  import Data.Bits
  import Data.Int
  import qualified Data.Vector.Unboxed as U
  import qualified Data.Vector.Unboxed.Mutable as MU
  
  -- | Maximal allowed argument of 'primeCount'. Currently 8e18.
  primeCountMaxArg :: Integer
  primeCountMaxArg = 8000000000000000000
  
  -- | @'primeCount' n == &#960;(n)@ is the number of (positive) primes not exceeding @n@.
  --
  --   For efficiency, the calculations are done on 64-bit signed integers, therefore @n@ must
  --   not exceed 'primeCountMaxArg'.
  --
  --   Requires @/O/(n^0.5)@ space, the time complexity is roughly @/O/(n^0.7)@.
  --   For small bounds, @'primeCount' n@ simply counts the primes not exceeding @n@,
  --   for bounds from @30000@ on, Meissel's algorithm is used in the improved form due to
  --   D.H. Lehmer, cf.
  --   <http://en.wikipedia.org/wiki/Prime_counting_function#Algorithms_for_evaluating_.CF.80.28x.29>.
  primeCount :: Integer -> Integer
  primeCount n
      | n > primeCountMaxArg = error $ "primeCount: can't handle bound " ++ show n
      | n < 2     = 0
      | n < 1000  = intToInteger . length . takeWhile (<= n) . map unPrime . primeList . primeSieve $ max 242 n
      | n < 30000 = runST $ do
          ba <- sieveTo n
          let (s, e) = (0, MU.length ba - 1)
          ct <- countFromTo s e ba
          return (intToInteger $ ct+3)
      | otherwise =
          let !ub = cop $ fromInteger n
              !sr = integerSquareRoot ub
              !cr = nxtEnd $ integerCubeRoot ub + 15
              nxtEnd k = k - (k `rem` 30) + 31
              !phn1 = calc ub cr
              !cs = cr+6
              !pdf = sieveCount ub cs sr
          in phn1 - pdf
  
  -- | @'nthPrime' n@ calculates the @n@-th prime. Numbering of primes is
  --   @1@-based, so @'nthPrime' 1 == 2@.
  --
  --   Requires @/O/((n*log n)^0.5)@ space, the time complexity is roughly @/O/((n*log n)^0.7@.
  --   The argument must be strictly positive.
  nthPrime :: Int -> Prime Integer
  nthPrime 1 = Prime 2
  nthPrime 2 = Prime 3
  nthPrime 3 = Prime 5
  nthPrime 4 = Prime 7
  nthPrime 5 = Prime 11
  nthPrime 6 = Prime 13
  nthPrime n
      | n < 1
      = error "Prime indexing starts at 1"
      | n < 200000
      = Prime $ countToNth (n - 3) [primeSieve (p0 + p0 `quot` 32 + 37)]
      | p0 > toInteger (maxBound :: Int)
      = error $ "nthPrime: index " ++ show n ++ " is too large to handle"
      | miss > 0
      = Prime $ tooLow  n (fromInteger p0) miss
      | otherwise
      = Prime $ tooHigh n (fromInteger p0) (negate miss)
        where
          p0 = nthPrimeApprox (toInteger n)
          miss = n - fromInteger (primeCount p0)
  
  --------------------------------------------------------------------------------
  --                                The Works                                   --
  --------------------------------------------------------------------------------
  
  -- TODO: do something better in case we guess too high.
  -- Not too pressing, since I think a) nthPrimeApprox always underestimates
  -- in the range we can handle, and b) it's always "goodEnough"
  
  tooLow :: Int -> Int -> Int -> Integer
  tooLow n p0 shortage
    | p1 > toInteger (maxBound :: Int)
    = error $ "nthPrime: index " ++ show n ++ " is too large to handle"
    | goodEnough
    = lowSieve p0 shortage
    | c1 < n
    = lowSieve (fromInteger p1) (n-c1)
    | otherwise
    = lowSieve p0 shortage   -- a third count wouldn't make it faster, I think
    where
      gap = truncate (log (intToDouble p0 :: Double))
      est = toInteger shortage * gap
      p1  = toInteger p0 + est
      goodEnough = 3*est*est*est < 2*p1*p1    -- a second counting would be more work than sieving
      c1  = fromInteger (primeCount p1)
  
  tooHigh :: Int -> Int -> Int -> Integer
  tooHigh n p0 surplus
    | c < n
    = lowSieve b (n-c)
    | otherwise
    = tooHigh n b (c-n)
    where
      gap = truncate (log (intToDouble p0 :: Double))
      b = p0 - (surplus * gap * 11) `quot` 10
      c = fromInteger (primeCount (toInteger b))
  
  lowSieve :: Int -> Int -> Integer
  lowSieve a miss = countToNth (miss+rep) psieves
        where
          strt = a + 1 + (a .&. 1)
          psieves@(PS vO ba:_) = psieveFrom (toInteger strt)
          rep | o0 < 0    = 0
              | otherwise = sum [1 | i <- [0 .. r2], unBit (ba `U.unsafeIndex` i)]
                where
                  o0 = toInteger strt - vO - 9   -- (strt - 2) - v0 - 7
                  r0 = fromInteger o0 `rem` 30
                  r1 = r0 `quot` 3
                  r2 = min 7 (if r1 > 5 then r1-1 else r1)
  
  -- highSieve :: Integer -> Integer -> Integer -> Integer
  -- highSieve a surp gap = error "Oh shit"
  
  sieveCount :: Int64 -> Int64 -> Int64 -> Integer
  sieveCount ub cr sr = runST (sieveCountST ub cr sr)
  
  sieveCountST :: forall s. Int64 -> Int64 -> Int64 -> ST s Integer
  sieveCountST ub cr sr = do
      let psieves = psieveFrom (int64ToInteger cr)
          pisr = approxPrimeCount sr
          picr = approxPrimeCount cr
          diff = pisr - picr
          size = int64ToInt (diff + diff `quot` 50) + 30
      store <- MU.unsafeNew size :: ST s (MU.MVector s Int64)
      let feed :: Int64 -> Int -> Int -> U.Vector Bit -> [PrimeSieve] -> ST s Integer
          feed voff !wi !ri uar sves
            | ri == sieveBits = case sves of
                                  (PS vO ba : more) -> feed (fromInteger vO) wi 0 ba more
                                  _ -> error "prime stream ended prematurely"
            | pval > sr   = do
                stu <- U.unsafeThaw uar
                eat 0 0 voff (wi-1) ri stu sves
            | unBit (uar `U.unsafeIndex` ri) = do
                MU.unsafeWrite store wi (ub `quot` pval)
                feed voff (wi+1) (ri+1) uar sves
            | otherwise = feed voff wi (ri+1) uar sves
              where
                pval = voff + toPrim ri
          eat :: Integer -> Integer -> Int64 -> Int -> Int -> MU.MVector s Bit -> [PrimeSieve] -> ST s Integer
          eat !acc !btw voff !wi !si stu sves
              | si == sieveBits =
                  case sves of
                    [] -> error "Premature end of prime stream"
                    (PS vO ba : more) -> do
                        nstu <- U.unsafeThaw ba
                        eat acc btw (fromInteger vO) wi 0 nstu more
              | wi < 0    = return acc
              | otherwise = do
                  qb <- MU.unsafeRead store wi
                  let dist = qb - voff - 7
                  if dist < intToInt64 sieveRange
                    then do
                        let (b,j) = idxPr (dist+7)
                            !li = (b `shiftL` 3) .|. j
                        new <- if li < si then return 0 else countFromTo si li stu
                        let nbtw = btw + intToInteger new + 1
                        eat (acc+nbtw) nbtw voff (wi-1) (li+1) stu sves
                    else do
                        let (cpl,fds) = dist `quotRem` intToInt64 sieveRange
                            (b,j) = idxPr (fds+7)
                            !li = (b `shiftL` 3) .|. j
                            ctLoop !lac 0 (PS vO ba : more) = do
                                nstu <- U.unsafeThaw ba
                                new <- countFromTo 0 li nstu
                                let nbtw = btw + lac + 1 + intToInteger new
                                eat (acc+nbtw) nbtw (integerToInt64 vO) (wi-1) (li+1) nstu more
                            ctLoop lac s (ps : more) = do
                                let !new = countAll ps
                                ctLoop (lac + intToInteger new) (s-1) more
                            ctLoop _ _ [] = error "Primes ended"
                        new <- countFromTo si (sieveBits-1) stu
                        ctLoop (intToInteger new) (cpl-1) sves
      case psieves of
        (PS vO ba : more) -> feed (fromInteger vO) 0 0 ba more
        _ -> error "No primes sieved"
  
  calc :: Int64 -> Int64 -> Integer
  calc lim plim = runST (calcST lim plim)
  
  calcST :: forall s. Int64 -> Int64 -> ST s Integer
  calcST lim plim = do
      !parr <- sieveTo (int64ToInteger plim)
      let (plo, phi) = (0, MU.length parr - 1)
      !pct <- countFromTo plo phi parr
      !ar1 <- MU.unsafeNew end
      MU.unsafeWrite ar1 0 lim
      MU.unsafeWrite ar1 1 1
      !ar2 <- MU.unsafeNew end
      let go :: Int -> Int -> MU.MVector s Int64 -> MU.MVector s Int64 -> ST s Integer
          go cap pix old new
              | pix == 2  =   coll cap old
              | otherwise = do
                  Bit isp <- MU.unsafeRead parr pix
                  if isp
                      then do
                          let !n = fromInteger (toPrim pix)
                          !ncap <- treat cap n old new
                          go ncap (pix-1) new old
                      else go cap (pix-1) old new
          coll :: Int -> MU.MVector s Int64 -> ST s Integer
          coll stop ar =
              let cgo !acc i
                      | i < stop  = do
                          !k <- MU.unsafeRead ar i
                          !v <- MU.unsafeRead ar (i+1)
                          cgo (acc + int64ToInteger v*cp6 k) (i+2)
                      | otherwise = return (acc+intToInteger pct+2)
              in cgo 0 0
      go 2 start ar1 ar2
    where
      (bt,ri) = idxPr plim
      !start = 8*bt + ri
      !size = int64ToInt $ integerSquareRoot lim `quot` 4
      !end = 2*size
  
  treat :: Int -> Int64 -> MU.MVector s Int64 -> MU.MVector s Int64 -> ST s Int
  treat end n old new = do
      qi0 <- locate n 0 (end `quot` 2 - 1) old
      let collect stop !acc ix
              | ix < end  = do
                  !k <- MU.unsafeRead old ix
                  if k < stop
                      then do
                          v <- MU.unsafeRead old (ix+1)
                          collect stop (acc-v) (ix+2)
                      else return (acc,ix)
              | otherwise = return (acc,ix)
          goTreat !wi !ci qi
              | qi < end  = do
                  !key <- MU.unsafeRead old qi
                  !val <- MU.unsafeRead old (qi+1)
                  let !q0 = key `quot` n
                      !r0 = int64ToInt (q0 `rem` 30030)
                      !nkey = q0 - int8ToInt64 (cpDfAr `U.unsafeIndex` r0)
                      nk0 = q0 + int8ToInt64 (cpGpAr `U.unsafeIndex` (r0+1) + 1)
                      !nlim = n*nk0
                  (wi1,ci1) <- copyTo end nkey old ci new wi
                  ckey <- MU.unsafeRead old ci1
                  (!acc, !ci2) <- if ckey == nkey
                                    then do
                                      !ov <- MU.unsafeRead old (ci1+1)
                                      return (ov-val,ci1+2)
                                    else return (-val,ci1)
                  (!tot, !nqi) <- collect nlim acc (qi+2)
                  MU.unsafeWrite new wi1 nkey
                  MU.unsafeWrite new (wi1+1) tot
                  goTreat (wi1+2) ci2 nqi
              | otherwise = copyRem end old ci new wi
      goTreat 0 0 qi0
  
  --------------------------------------------------------------------------------
  --                               Auxiliaries                                  --
  --------------------------------------------------------------------------------
  
  locate :: Int64 -> Int -> Int -> MU.MVector s Int64 -> ST s Int
  locate p low high arr = do
      let go lo hi
            | lo < hi     = do
              let !md = (lo+hi) `quot` 2
              v <- MU.unsafeRead arr (2*md)
              case compare p v of
                  LT -> go lo md
                  EQ -> return (2*md)
                  GT -> go (md+1) hi
            | otherwise   = return (2*lo)
      go low high
  
  {-# INLINE copyTo #-}
  copyTo :: Int -> Int64 -> MU.MVector s Int64 -> Int
         -> MU.MVector s Int64 -> Int -> ST s (Int,Int)
  copyTo end lim old oi new ni = do
      let go ri wi
              | ri < end  = do
                  ok <- MU.unsafeRead old ri
                  if ok < lim
                      then do
                          !ov <- MU.unsafeRead old (ri+1)
                          MU.unsafeWrite new wi ok
                          MU.unsafeWrite new (wi+1) ov
                          go (ri+2) (wi+2)
                      else return (wi,ri)
              | otherwise = return (wi,ri)
      go oi ni
  
  {-# INLINE copyRem #-}
  copyRem :: Int -> MU.MVector s Int64 -> Int -> MU.MVector s Int64 -> Int -> ST s Int
  copyRem end old oi new ni = do
    let len = end - oi
    MU.copy (MU.slice ni len new) (MU.slice oi len old)
    pure $ ni + len
  
  {-# INLINE cp6 #-}
  cp6 :: Int64 -> Integer
  cp6 k =
    case k `quotRem` 30030 of
      (q,r) -> 5760*int64ToInteger q +
                  int16ToInteger (cpCtAr `U.unsafeIndex` int64ToInt r)
  
  cop :: Int64 -> Int64
  cop m = m - int8ToInt64 (cpDfAr `U.unsafeIndex` int64ToInt (m `rem` 30030))
  
  
  --------------------------------------------------------------------------------
  --                           Ugly helper arrays                               --
  --------------------------------------------------------------------------------
  
  cpCtAr :: U.Vector Int16
  cpCtAr = runST $ do
      ar <- MU.replicate 30030 1
      let zilch s i
              | i < 30030 = MU.unsafeWrite ar i 0 >> zilch s (i+s)
              | otherwise = return ()
          accumulate ct i
              | i < 30030 = do
                  v <- MU.unsafeRead ar i
                  let !ct' = ct+v
                  MU.unsafeWrite ar i ct'
                  accumulate ct' (i+1)
              | otherwise = return ar
      zilch 2 0
      zilch 6 3
      zilch 10 5
      zilch 14 7
      zilch 22 11
      zilch 26 13
      vec <- accumulate 1 2
      U.unsafeFreeze vec
  
  cpDfAr :: U.Vector Int8
  cpDfAr = runST $ do
      ar <- MU.replicate 30030 0
      let note s i
              | i < 30029 = MU.unsafeWrite ar i 1 >> note s (i+s)
              | otherwise = return ()
          accumulate d i
              | i < 30029 = do
                  v <- MU.unsafeRead ar i
                  if v == 0
                      then accumulate 2 (i+2)
                      else do MU.unsafeWrite ar i d
                              accumulate (d+1) (i+1)
              | otherwise = return ar
      note 2 0
      note 6 3
      note 10 5
      note 14 7
      note 22 11
      note 26 13
      vec <- accumulate 2 3
      U.unsafeFreeze vec
  
  cpGpAr :: U.Vector Int8
  cpGpAr = runST $ do
      ar <- MU.replicate 30031 0
      MU.unsafeWrite ar 30030 1
      let note s i
              | i < 30029 = MU.unsafeWrite ar i 1 >> note s (i+s)
              | otherwise = return ()
          accumulate d i
              | i < 1     = return ar
              | otherwise = do
                  v <- MU.unsafeRead ar i
                  if v == 0
                      then accumulate 2 (i-2)
                      else do MU.unsafeWrite ar i d
                              accumulate (d+1) (i-1)
      note 2 0
      note 6 3
      note 10 5
      note 14 7
      note 22 11
      note 26 13
      vec <- accumulate 2 30027
      U.unsafeFreeze vec
  
  -------------------------------------------------------------------------------
  -- Prime counting
  
  -- find the n-th set bit in a list of PrimeSieves,
  -- aka find the (n+3)-rd prime
  countToNth :: Int -> [PrimeSieve] -> Integer
  countToNth !_ [] = error "countToNth: Prime stream ended prematurely"
  countToNth !n (PS v0 bs : more) = case nthBitIndex (Bit True) n bs of
    Just i -> v0 + toPrim i
    Nothing -> countToNth (n - countBits bs) more
  
  -- count all set bits in a chunk, do it wordwise for speed.
  countAll :: PrimeSieve -> Int
  countAll (PS _ bs) = countBits bs
  
  -- count set bits between two indices (inclusive)
  -- start and end must both be valid indices and start <= end
  countFromTo :: Int -> Int -> MU.MVector s Bit -> ST s Int
  countFromTo start end =
    fmap countBits . U.unsafeFreeze . MU.slice start (end - start + 1)
  