-- |
-- Module:      Math.NumberTheory.Logarithms
-- Copyright:   (c) 2011 Daniel Fischer
-- Licence:     MIT
-- Maintainer:  Daniel Fischer <daniel.is.fischer@googlemail.com>
-- Stability:   Provisional
-- Portability: Non-portable (GHC extensions)
--
-- Integer Logarithms. For efficiency, the internal representation of 'Integer's
-- from integer-gmp is used.
--
{-# LANGUAGE CPP #-}
{-# LANGUAGE MagicHash #-}
module Math.NumberTheory.Logarithms
    ( -- * Integer logarithms with input checks
      integerLogBase
    , integerLog2
    , integerLog10

    , naturalLogBase
    , naturalLog2
    , naturalLog10

    , intLog2
    , wordLog2

      -- * Integer logarithms without input checks
      --
      -- | These functions are total, however, don't rely on the values with out-of-domain arguments.
    , integerLogBase'
    , integerLog2'
    , integerLog10'

    , intLog2'
    , wordLog2'
    ) where

import GHC.Exts

import Data.Bits
import Data.Array.Unboxed
import Numeric.Natural

import GHC.Integer.Logarithms.Compat
#if Base48 && defined(MIN_VERSION_integer_gmp)
import GHC.Integer.GMP.Internals (Integer (..))
import GHC.Natural
#endif

#if CheckBounds
import Data.Array.IArray (IArray, (!))
#else
import Data.Array.Base (unsafeAt)
#endif

-- | Calculate the integer logarithm for an arbitrary base.
--   The base must be greater than 1, the second argument, the number
--   whose logarithm is sought, must be positive, otherwise an error is thrown.
--   If @base == 2@, the specialised version is called, which is more
--   efficient than the general algorithm.
--
--   Satisfies:
--
-- > base ^ integerLogBase base m <= m < base ^ (integerLogBase base m + 1)
--
-- for @base > 1@ and @m > 0@.
integerLogBase :: Integer -> Integer -> Int
integerLogBase b n
  | n < 1       = error "Math.NumberTheory.Logarithms.integerLogBase: argument must be positive."
  | n < b       = 0
  | b == 2      = integerLog2' n
  | b < 2       = error "Math.NumberTheory.Logarithms.integerLogBase: base must be greater than one."
  | otherwise   = integerLogBase' b n

-- | Calculate the integer logarithm of an 'Integer' to base 2.
--   The argument must be positive, otherwise an error is thrown.
integerLog2 :: Integer -> Int
integerLog2 n
  | n < 1       = error "Math.NumberTheory.Logarithms.integerLog2: argument must be positive"
  | otherwise   = I# (integerLog2# n)

-- | Cacluate the integer logarithm for an arbitrary base.
--   The base must be greater than 1, the second argument, the number
--   whose logarithm is sought, must be positive, otherwise an error is thrown.
--   If @base == 2@, the specialised version is called, which is more
--   efficient than the general algorithm.
--
--   Satisfies:
--
-- > base ^ integerLogBase base m <= m < base ^ (integerLogBase base m + 1)
--
-- for @base > 1@ and @m > 0@.
naturalLogBase :: Natural -> Natural -> Int
naturalLogBase b n
  | n < 1       = error "Math.NumberTheory.Logarithms.naturalLogBase: argument must be positive."
  | n < b       = 0
  | b == 2      = naturalLog2' n
  | b < 2       = error "Math.NumberTheory.Logarithms.naturalLogBase: base must be greater than one."
  | otherwise   = naturalLogBase' b n

-- | Calculate the natural logarithm of an 'Natural' to base 2.
--   The argument must be non-zero, otherwise an error is thrown.
naturalLog2 :: Natural -> Int
naturalLog2 n
  | n < 1       = error "Math.NumberTheory.Logarithms.naturalLog2: argument must be non-zero"
  | otherwise   = I# (naturalLog2# n)

-- | Calculate the integer logarithm of an 'Int' to base 2.
--   The argument must be positive, otherwise an error is thrown.
intLog2 :: Int -> Int
intLog2 (I# i#)
  | isTrue# (i# <# 1#)  = error "Math.NumberTheory.Logarithms.intLog2: argument must be positive"
  | otherwise           = I# (wordLog2# (int2Word# i#))

-- | Calculate the integer logarithm of a 'Word' to base 2.
--   The argument must be positive, otherwise an error is thrown.
wordLog2 :: Word -> Int
wordLog2 (W# w#)
  | isTrue# (w# `eqWord#` 0##)  = error "Math.NumberTheory.Logarithms.wordLog2: argument must not be 0."
  | otherwise                   = I# (wordLog2# w#)

-- | Same as 'integerLog2', but without checks, saves a little time when
--   called often for known good input.
integerLog2' :: Integer -> Int
integerLog2' n = I# (integerLog2# n)

-- | Same as 'naturalLog2', but without checks, saves a little time when
--   called often for known good input.
naturalLog2' :: Natural -> Int
naturalLog2' n = I# (naturalLog2# n)

-- | Same as 'intLog2', but without checks, saves a little time when
--   called often for known good input.
intLog2' :: Int -> Int
intLog2' (I# i#) = I# (wordLog2# (int2Word# i#))

-- | Same as 'wordLog2', but without checks, saves a little time when
--   called often for known good input.
wordLog2' :: Word -> Int
wordLog2' (W# w#) = I# (wordLog2# w#)

-- | Calculate the integer logarithm of an 'Integer' to base 10.
--   The argument must be positive, otherwise an error is thrown.
integerLog10 :: Integer -> Int
integerLog10 n
  | n < 1     = error "Math.NumberTheory.Logarithms.integerLog10: argument must be positive"
  | otherwise = integerLog10' n

-- | Calculate the integer logarithm of an 'Integer' to base 10.
--   The argument must be not zero, otherwise an error is thrown.
naturalLog10 :: Natural -> Int
naturalLog10 n
  | n < 1     = error "Math.NumberTheory.Logarithms.naturalaLog10: argument must be non-zero"
  | otherwise = naturalLog10' n

-- | Same as 'integerLog10', but without a check for a positive
--   argument. Saves a little time when called often for known good
--   input.
integerLog10' :: Integer -> Int
integerLog10' n
  | n < 10      = 0
  | n < 100     = 1
  | otherwise   = ex + integerLog10' (n `quot` 10 ^ ex)
    where
      ln = I# (integerLog2# n)
      -- u/v is a good approximation of log 2/log 10
      u  = 1936274
      v  = 6432163
      -- so ex is a good approximation to integerLogBase 10 n
      ex = fromInteger ((u * fromIntegral ln) `quot` v)

-- | Same as 'naturalLog10', but without a check for a positive
--   argument. Saves a little time when called often for known good
--   input.
naturalLog10' :: Natural -> Int
naturalLog10' n
  | n < 10      = 0
  | n < 100     = 1
  | otherwise   = ex + naturalLog10' (n `quot` 10 ^ ex)
    where
      ln = I# (naturalLog2# n)
      -- u/v is a good approximation of log 2/log 10
      u  = 1936274
      v  = 6432163
      -- so ex is a good approximation to naturalLogBase 10 n
      ex = fromInteger ((u * fromIntegral ln) `quot` v)

-- | Same as 'integerLogBase', but without checks, saves a little time when
--   called often for known good input.
integerLogBase' :: Integer -> Integer -> Int
integerLogBase' b n
  | n < b       = 0
  | ln-lb < lb  = 1     -- overflow safe version of ln < 2*lb, implies n < b*b
  | b < 33      = let bi = fromInteger b
                      ix = 2*bi-4
                      -- u/v is a good approximation of log 2/log b
                      u  = logArr `unsafeAt` ix
                      v  = logArr `unsafeAt` (ix+1)
                      -- hence ex is a rather good approximation of integerLogBase b n
                      -- most of the time, it will already be exact
                      ex = fromInteger ((fromIntegral u * fromIntegral ln) `quot` fromIntegral v)
                  in case u of
                      1 -> ln `quot` v      -- a power of 2, easy
                      _ -> ex + integerLogBase' b (n `quot` b ^ ex)
  | otherwise   = let -- shift b so that 16 <= bi < 32
                      bi = fromInteger (b `shiftR` (lb-4))
                      -- we choose an approximation of log 2 / log (bi+1) to
                      -- be sure we underestimate
                      ix = 2*bi-2
                      -- u/w is a reasonably good approximation to log 2/log b
                      -- it is too small, but not by much, so the recursive call
                      -- should most of the time be caught by one of the first
                      -- two guards unless n is huge, but then it'd still be
                      -- a call with a much smaller second argument.
                      u  = fromIntegral $ logArr `unsafeAt` ix
                      v  = fromIntegral $ logArr `unsafeAt` (ix+1)
                      w  = v + u*fromIntegral (lb-4)
                      ex = fromInteger ((u * fromIntegral ln) `quot` w)
                  in ex + integerLogBase' b (n `quot` b ^ ex)
    where
      lb = integerLog2' b
      ln = integerLog2' n

-- | Same as 'naturalLogBase', but without checks, saves a little time when
--   called often for known good input.
naturalLogBase' :: Natural -> Natural -> Int
naturalLogBase' b n
    | n < b       = 0
  | ln-lb < lb  = 1     -- overflow safe version of ln < 2*lb, implies n < b*b
  | b < 33      = let bi = fromIntegral b
                      ix = 2*bi-4
                      -- u/v is a good approximation of log 2/log b
                      u  = logArr `unsafeAt` ix
                      v  = logArr `unsafeAt` (ix+1)
                      -- hence ex is a rather good approximation of integerLogBase b n
                      -- most of the time, it will already be exact
                      ex = fromNatural ((fromIntegral u * fromIntegral ln) `quot` fromIntegral v)
                  in case u of
                      1 -> ln `quot` v      -- a power of 2, easy
                      _ -> ex + naturalLogBase' b (n `quot` b ^ ex)
  | otherwise   = let -- shift b so that 16 <= bi < 32
                      bi = fromNatural (b `shiftR` (lb-4))
                      -- we choose an approximation of log 2 / log (bi+1) to
                      -- be sure we underestimate
                      ix = 2*bi-2
                      -- u/w is a reasonably good approximation to log 2/log b
                      -- it is too small, but not by much, so the recursive call
                      -- should most of the time be caught by one of the first
                      -- two guards unless n is huge, but then it'd still be
                      -- a call with a much smaller second argument.
                      u  = fromIntegral $ logArr `unsafeAt` ix
                      v  = fromIntegral $ logArr `unsafeAt` (ix+1)
                      w  = v + u*fromIntegral (lb-4)
                      ex = fromNatural ((u * fromIntegral ln) `quot` w)
                  in ex + naturalLogBase' b (n `quot` b ^ ex)
    where
      lb = naturalLog2' b
      ln = naturalLog2' n

-- Lookup table for logarithms of 2 <= k <= 32
-- In each row "x , y", x/y is a good rational approximation of log 2  / log k.
-- For the powers of 2, it is exact, otherwise x/y < log 2/log k, since we don't
-- want to overestimate integerLogBase b n = floor $ (log 2/log b)*logBase 2 n.
logArr :: UArray Int Int
logArr = listArray (0, 61)
          [ 1 , 1,
            190537 , 301994,
            1 , 2,
            1936274 , 4495889,
            190537 , 492531,
            91313 , 256348,
            1 , 3,
            190537 , 603988,
            1936274 , 6432163,
            1686227 , 5833387,
            190537 , 683068,
            5458 , 20197,
            91313 , 347661,
            416263 , 1626294,
            1 , 4,
            32631 , 133378,
            190537 , 794525,
            163451 , 694328,
            1936274 , 8368437,
            1454590 , 6389021,
            1686227 , 7519614,
            785355 , 3552602,
            190537 , 873605,
            968137 , 4495889,
            5458 , 25655,
            190537 , 905982,
            91313 , 438974,
            390321 , 1896172,
            416263 , 2042557,
            709397 , 3514492,
            1 , 5
          ]

-------------------------------------------------------------------------------
-- Unsafe
-------------------------------------------------------------------------------

#if CheckBounds
unsafeAt :: (IArray a e, Ix i) => a i e -> i -> e
unsafeAt = (!)`
#endif

-------------------------------------------------------------------------------
-- Natural helpers
-------------------------------------------------------------------------------

fromNatural :: Num a => Natural -> a
fromNatural = fromIntegral

naturalLog2# :: Natural -> Int#
#if Base48 && defined(MIN_VERSION_integer_gmp)
naturalLog2# (NatS# b) = wordLog2# b
naturalLog2# (NatJ# n) = integerLog2# (Jp# n)
#else
naturalLog2# n = integerLog2# (toInteger n)
#endif

#if __GLASGOW_HASKELL__ < 707
-- The times they are a-changing. The types of primops too :(
isTrue# :: Bool -> Bool
isTrue# = id
#endif