-- | Operations on integers

module Math.Combinat.Numbers.Integers 
  ( -- * Integer logarithm
    integerLog2
  , ceilingLog2
    -- * Integer square root
  , isSquare
  , integerSquareRoot
  , ceilingSquareRoot
  , integerSquareRoot' 
  , integerSquareRootNewton'
  )
  where

--------------------------------------------------------------------------------

-- import Math.Combinat.Numbers

import Data.List ( group , sort )
import Data.Bits

import System.Random

--------------------------------------------------------------------------------
-- Integer logarithm

-- | Largest integer @k@ such that @2^k@ is smaller or equal to @n@
integerLog2 :: Integer -> Integer
integerLog2 :: Integer -> Integer
integerLog2 Integer
n = Integer -> Integer
forall t p. (Num t, Num p, Bits t) => t -> p
go Integer
n where
  go :: t -> p
go t
0 = -p
1
  go t
k = p
1 p -> p -> p
forall a. Num a => a -> a -> a
+ t -> p
go (t -> Int -> t
forall a. Bits a => a -> Int -> a
shiftR t
k Int
1)

-- | Smallest integer @k@ such that @2^k@ is larger or equal to @n@
ceilingLog2 :: Integer -> Integer
ceilingLog2 :: Integer -> Integer
ceilingLog2 Integer
0 = Integer
0
ceilingLog2 Integer
n = Integer
1 Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer -> Integer
forall t p. (Num t, Num p, Bits t) => t -> p
go (Integer
nInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
-Integer
1) where
  go :: t -> p
go t
0 = -p
1
  go t
k = p
1 p -> p -> p
forall a. Num a => a -> a -> a
+ t -> p
go (t -> Int -> t
forall a. Bits a => a -> Int -> a
shiftR t
k Int
1)
  
--------------------------------------------------------------------------------
-- Integer square root

isSquare :: Integer -> Bool
isSquare :: Integer -> Bool
isSquare Integer
n = 
  if (Integer -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Integer -> Int) -> Integer -> Int
forall a b. (a -> b) -> a -> b
$ Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
mod Integer
n Integer
32) Int -> [Int] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [Int]
rs 
    then (Integer, Integer) -> Integer
forall a b. (a, b) -> b
snd (Integer -> (Integer, Integer)
integerSquareRoot' Integer
n) Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
0
    else Bool
False
  where
    rs :: [Int]
rs = [Int
0,Int
1,Int
4,Int
9,Int
16,Int
17,Int
25] :: [Int]
    
-- | Integer square root (largest integer whose square is smaller or equal to the input)
-- using Newton's method, with a faster (for large numbers) inital guess based on bit shifts.
integerSquareRoot :: Integer -> Integer
integerSquareRoot :: Integer -> Integer
integerSquareRoot = (Integer, Integer) -> Integer
forall a b. (a, b) -> a
fst ((Integer, Integer) -> Integer)
-> (Integer -> (Integer, Integer)) -> Integer -> Integer
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Integer -> (Integer, Integer)
integerSquareRoot'

-- | Smallest integer whose square is larger or equal to the input
ceilingSquareRoot :: Integer -> Integer
ceilingSquareRoot :: Integer -> Integer
ceilingSquareRoot Integer
n = (if Integer
rInteger -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
>Integer
0 then Integer
uInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
+Integer
1 else Integer
u) where (Integer
u,Integer
r) = Integer -> (Integer, Integer)
integerSquareRoot' Integer
n 

-- | We also return the excess residue; that is
--
-- > (a,r) = integerSquareRoot' n
-- 
-- means that
--
-- > a*a + r = n
-- > a*a <= n < (a+1)*(a+1)
integerSquareRoot' :: Integer -> (Integer,Integer)
integerSquareRoot' :: Integer -> (Integer, Integer)
integerSquareRoot' Integer
n
  | Integer
nInteger -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
<Integer
0 = [Char] -> (Integer, Integer)
forall a. HasCallStack => [Char] -> a
error [Char]
"integerSquareRoot: negative input"
  | Integer
nInteger -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
<Integer
2 = (Integer
n,Integer
0)
  | Bool
otherwise = Integer -> (Integer, Integer)
go Integer
firstGuess 
  where
    k :: Integer
k = Integer -> Integer
integerLog2 Integer
n
    firstGuess :: Integer
firstGuess = Integer
2Integer -> Integer -> Integer
forall a b. (Num a, Integral b) => a -> b -> a
^(Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
div (Integer
kInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
+Integer
2) Integer
2) -- !! note that (div (k+1) 2) is NOT enough !!
    go :: Integer -> (Integer, Integer)
go Integer
a = 
      if Integer
m Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
a
        then Integer -> (Integer, Integer)
go Integer
a' 
        else (Integer
a, Integer
r Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
aInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
*(Integer
mInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
-Integer
a))
      where
        (Integer
m,Integer
r) = Integer -> Integer -> (Integer, Integer)
forall a. Integral a => a -> a -> (a, a)
divMod Integer
n Integer
a
        a' :: Integer
a' = Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
div (Integer
m Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
a) Integer
2

-- | Newton's method without an initial guess. For very small numbers (<10^10) it
-- is somewhat faster than the above version.
integerSquareRootNewton' :: Integer -> (Integer,Integer)
integerSquareRootNewton' :: Integer -> (Integer, Integer)
integerSquareRootNewton' Integer
n
  | Integer
nInteger -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
<Integer
0 = [Char] -> (Integer, Integer)
forall a. HasCallStack => [Char] -> a
error [Char]
"integerSquareRootNewton: negative input"
  | Integer
nInteger -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
<Integer
2 = (Integer
n,Integer
0)
  | Bool
otherwise = Integer -> (Integer, Integer)
go (Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
div Integer
n Integer
2) 
  where
    go :: Integer -> (Integer, Integer)
go Integer
a = 
      if Integer
m Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
a
        then Integer -> (Integer, Integer)
go Integer
a' 
        else (Integer
a, Integer
r Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
aInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
*(Integer
mInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
-Integer
a))
      where
        (Integer
m,Integer
r) = Integer -> Integer -> (Integer, Integer)
forall a. Integral a => a -> a -> (a, a)
divMod Integer
n Integer
a
        a' :: Integer
a' = Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
div (Integer
m Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
a) Integer
2

{-
-- brute force test of integer square root
isqrt_test n1 n2 = 
  [ k 
  | k<-[n1..n2] 
  , let (a,r) = integerSquareRoot' k
  , (a*a+r/=k) || (a*a>k) || (a+1)*(a+1)<=k 
  ]
-}

--------------------------------------------------------------------------------