    -- |
    -- Module:      Math.NumberTheory.Roots.Squares
    -- Copyright:   (c) 2011 Daniel Fischer, 2016-2020 Andrew Lelechenko
    -- Licence:     MIT
    -- Maintainer:  Daniel Fischer <daniel.is.fischer@googlemail.com>
    --
    -- Functions dealing with squares. Efficient calculation of integer square roots
    -- and efficient testing for squareness.

    {-# LANGUAGE BangPatterns #-}
    {-# LANGUAGE MagicHash    #-}

    module Math.NumberTheory.Roots.Squares
        ( -- * Square root calculation
          integerSquareRoot
        , integerSquareRoot'
        , integerSquareRootRem
        , integerSquareRootRem'
        , exactSquareRoot
          -- * Tests for squares
        , isSquare
        , isSquare'
        , isPossibleSquare
        ) where

    import Data.Bits (finiteBitSize, (.&.))
    import Data.Int (Int64)
    import Data.Word (Word64)
    import GHC.Exts (Ptr(..))
    import Numeric.Natural (Natural)

    import Math.NumberTheory.Roots.Squares.Internal
    import Math.NumberTheory.Utils.BitMask (indexBitSet)

    -- | For a non-negative input \( n \)
    --   calculate its integer square root \( \lfloor \sqrt{n} \rfloor \).
    --   Throw an error on negative input.
    --
    -- >>> integerSquareRoot 99
    -- 9
    -- >>> integerSquareRoot 100
    -- 10
    -- >>> integerSquareRoot 101
    -- 10
    {-# SPECIALISE integerSquareRoot :: Int -> Int #-}
    {-# SPECIALISE integerSquareRoot :: Word -> Word #-}
    {-# SPECIALISE integerSquareRoot :: Int64 -> Int64 #-}
    {-# SPECIALISE integerSquareRoot :: Word64 -> Word64 #-}
    {-# SPECIALISE integerSquareRoot :: Integer -> Integer #-}
    {-# SPECIALISE integerSquareRoot :: Natural -> Natural #-}
    integerSquareRoot :: Integral a => a -> a
    integerSquareRoot :: forall a. Integral a => a -> a
    integerSquareRoot a
    n
      | a
    n a -> a -> Bool
    forall a. Ord a => a -> a -> Bool
    < a
    0       = [Char] -> a
    forall a. HasCallStack => [Char] -> a
    error [Char]
    "integerSquareRoot: negative argument"
      | Bool
    otherwise   = a -> a
    forall a. Integral a => a -> a
    integerSquareRoot' a
    n

    -- | Calculate the integer square root of a non-negative number @n@,
    --   that is, the largest integer @r@ with @r*r <= n@.
    --   The precondition @n >= 0@ is not checked.
    {-# RULES
    "integerSquareRoot'/Int"     integerSquareRoot' = isqrtInt'
    "integerSquareRoot'/Word"    integerSquareRoot' = isqrtWord
    "integerSquareRoot'/Int64"   integerSquareRoot' = isqrtInt64'
    "integerSquareRoot'/Word64"  integerSquareRoot' = isqrtWord64
    "integerSquareRoot'/Integer" integerSquareRoot' = isqrtInteger
    "integerSquareRoot'/Natural" integerSquareRoot' = fromInteger . isqrtInteger . toInteger
      #-}
    {-# INLINE [1] integerSquareRoot' #-}
    integerSquareRoot' :: Integral a => a -> a
    integerSquareRoot' :: forall a. Integral a => a -> a
    integerSquareRoot' = a -> a
    forall a. Integral a => a -> a
    isqrtA

    -- | For a non-negative input \( n \)
    --   calculate its integer square root \( r = \lfloor \sqrt{n} \rfloor \)
    --   and remainder \( s = n - r^2 \).
    --   Throw an error on negative input.
    --
    -- >>> integerSquareRootRem 99
    -- (9,18)
    -- >>> integerSquareRootRem 100
    -- (10,0)
    -- >>> integerSquareRootRem 101
    -- (10,1)
    {-# SPECIALISE integerSquareRootRem :: Int -> (Int, Int) #-}
    {-# SPECIALISE integerSquareRootRem :: Word -> (Word, Word) #-}
    {-# SPECIALISE integerSquareRootRem :: Integer -> (Integer, Integer) #-}
    {-# SPECIALISE integerSquareRootRem :: Natural -> (Natural, Natural) #-}
    integerSquareRootRem :: Integral a => a -> (a, a)
    integerSquareRootRem :: forall a. Integral a => a -> (a, a)
    integerSquareRootRem a
    n
      | a
    n a -> a -> Bool
    forall a. Ord a => a -> a -> Bool
    < a
    0       = [Char] -> (a, a)
    forall a. HasCallStack => [Char] -> a
    error [Char]
    "integerSquareRootRem: negative argument"
      | Bool
    otherwise   = a -> (a, a)
    forall a. Integral a => a -> (a, a)
    integerSquareRootRem' a
    n

    -- | Calculate the integer square root of a non-negative number as well as
    --   the difference of that number with the square of that root, that is if
    --   @(s,r) = integerSquareRootRem' n@ then @s^2 <= n == s^2+r < (s+1)^2@.
    --   The precondition @n >= 0@ is not checked.
    {-# RULES
    "integerSquareRootRem'/Integer" integerSquareRootRem' = karatsubaSqrt
      #-}
    {-# INLINE [1] integerSquareRootRem' #-}
    integerSquareRootRem' :: Integral a => a -> (a, a)
    integerSquareRootRem' :: forall a. Integral a => a -> (a, a)
    integerSquareRootRem' a
    n = (a
    s, a
    n a -> a -> a
    forall a. Num a => a -> a -> a
    - a
    s a -> a -> a
    forall a. Num a => a -> a -> a
    * a
    s)
      where
        s :: a
    s = a -> a
    forall a. Integral a => a -> a
    integerSquareRoot' a
    n

    -- | Calculate the exact integer square root if it exists,
    -- otherwise return 'Nothing'.
    --
    -- >>> map exactSquareRoot [-100, 99, 100, 101]
    -- [Nothing,Nothing,Just 10,Nothing]
    {-# SPECIALISE exactSquareRoot :: Int -> Maybe Int #-}
    {-# SPECIALISE exactSquareRoot :: Word -> Maybe Word #-}
    {-# SPECIALISE exactSquareRoot :: Integer -> Maybe Integer #-}
    {-# SPECIALISE exactSquareRoot :: Natural -> Maybe Natural #-}
    exactSquareRoot :: Integral a => a -> Maybe a
    exactSquareRoot :: forall a. Integral a => a -> Maybe a
    exactSquareRoot a
    n
      | a
    n a -> a -> Bool
    forall a. Ord a => a -> a -> Bool
    >= a
    0
      , a -> Bool
    forall a. Integral a => a -> Bool
    isPossibleSquare a
    n
      , (a
    r, a
    0) <- a -> (a, a)
    forall a. Integral a => a -> (a, a)
    integerSquareRootRem' a
    n = a -> Maybe a
    forall a. a -> Maybe a
    Just a
    r
      | Bool
    otherwise                         = Maybe a
    forall a. Maybe a
    Nothing

    -- | Test whether the argument is a perfect square.
    --
    -- >>> map isSquare [-100, 99, 100, 101]
    -- [False,False,True,False]
    {-# SPECIALISE isSquare :: Int -> Bool #-}
    {-# SPECIALISE isSquare :: Word -> Bool #-}
    {-# SPECIALISE isSquare :: Integer -> Bool #-}
    {-# SPECIALISE isSquare :: Natural -> Bool #-}
    isSquare :: Integral a => a -> Bool
    isSquare :: forall a. Integral a => a -> Bool
    isSquare a
    n = a
    n a -> a -> Bool
    forall a. Ord a => a -> a -> Bool
    >= a
    0 Bool -> Bool -> Bool
    && a -> Bool
    forall a. Integral a => a -> Bool
    isSquare' a
    n

    -- | Test whether the input (a non-negative number) @n@ is a square.
    --   The same as 'isSquare', but without the negativity test,
    --   so marginally faster.
    --
    --   The precondition @n >= 0@ is not tested, passing negative
    --   arguments may cause any kind of havoc.
    {-# SPECIALISE isSquare' :: Int -> Bool #-}
    {-# SPECIALISE isSquare' :: Word -> Bool #-}
    {-# SPECIALISE isSquare' :: Integer -> Bool #-}
    {-# SPECIALISE isSquare' :: Natural -> Bool #-}
    isSquare' :: Integral a => a -> Bool
    isSquare' :: forall a. Integral a => a -> Bool
    isSquare' a
    n
        | a -> Bool
    forall a. Integral a => a -> Bool
    isPossibleSquare a
    n
        , (a
    _, a
    0) <- a -> (a, a)
    forall a. Integral a => a -> (a, a)
    integerSquareRootRem' a
    n = Bool
    True
        | Bool
    otherwise                         = Bool
    False

    -- | Test whether a non-negative number may be a square.
    --   Non-negativity is not checked, passing negative arguments may
    --   cause any kind of havoc.
    --
    --   First the remainder modulo 256 is checked (that can be calculated
    --   easily without division and eliminates about 82% of all numbers).
    --   After that, the remainders modulo 9, 25, 7, 11 and 13 are tested
    --   to eliminate altogether about 99.436% of all numbers.
    {-# SPECIALISE isPossibleSquare :: Int -> Bool #-}
    {-# SPECIALISE isPossibleSquare :: Word -> Bool #-}
    {-# SPECIALISE isPossibleSquare :: Integer -> Bool #-}
    {-# SPECIALISE isPossibleSquare :: Natural -> Bool #-}
    isPossibleSquare :: Integral a => a -> Bool
    isPossibleSquare :: forall a. Integral a => a -> Bool
    isPossibleSquare a
    n'
      =  Ptr Word -> Int -> Bool
    indexBitSet Ptr Word
    mask256 (Integer -> Int
    forall a. Num a => Integer -> a
    fromInteger (Integer
    n Integer -> Integer -> Integer
    forall a. Bits a => a -> a -> a
    .&. Integer
    255))
      Bool -> Bool -> Bool
    && Ptr Word -> Int -> Bool
    indexBitSet Ptr Word
    mask693 (Integer -> Int
    forall a. Num a => Integer -> a
    fromInteger (Integer
    n Integer -> Integer -> Integer
    forall a. Integral a => a -> a -> a
    `rem` Integer
    693))
      Bool -> Bool -> Bool
    && Ptr Word -> Int -> Bool
    indexBitSet Ptr Word
    mask325 (Integer -> Int
    forall a. Num a => Integer -> a
    fromInteger (Integer
    n Integer -> Integer -> Integer
    forall a. Integral a => a -> a -> a
    `rem` Integer
    325))
      where
        n :: Integer
    n = a -> Integer
    forall a. Integral a => a -> Integer
    toInteger a
    n'

    -----------------------------------------------------------------------------
    -- Generated by 'Math.NumberTheory.Utils.BitMask.vectorToAddrLiteral'

    mask256 :: Ptr Word
    mask256 :: Ptr Word
    mask256 = Addr# -> Ptr Word
    forall a. Addr# -> Ptr a
    Ptr Addr#
    "\DC3\STX\ETX\STX\DC2\STX\STX\STX\DC3\STX\STX\STX\DC2\STX\STX\STX\DC2\STX\ETX\STX\DC2\STX\STX\STX\DC2\STX\STX\STX\DC2\STX\STX\STX"#

    mask693 :: Ptr Word
    mask693 :: Ptr Word
    mask693 = Addr# -> Ptr Word
    forall a. Addr# -> Ptr a
    Ptr Addr#
    "\DC3\STXA\STX0\NUL\STX\EOTI\NUL\STX\t\CAN\NUL\NULB\164\NUL\DC1\EOT\b\STX\NUL@P\128@\NUL\STX\t\128 \SOH\DLE\NUL\SOH\130$\NUL\128\DC4(\NUL\NUL\SOH\DC2\NUL\f\STX\DC4\SOH\NUL \b\NUL\"\NUL\128\EOT`\144\NUL\b\129\NULE\DC2\DLE@\STX\EOT\NUL\129\NUL\t\b\EOT\SOH\194\128\NUL\DLE\EOT\NUL\DLE\NUL\NUL"#

    mask325 :: Ptr Word
    mask325 :: Ptr Word
    mask325 = Addr# -> Ptr Word
    forall a. Addr# -> Ptr a
    Ptr Addr#
    "\DC3B\SOH&\144\NUL\n!%\140\STXH0\SOH\DC4BJ\b\ENQ\144@\STX(\132\148\DLE\n \131\EOTP\f)!\DC4@\STX\EM\160\DLE\DC2"#

    -- -- Make an array indicating whether a remainder is a square remainder.
    -- sqRemArray :: Int -> V.Vector Bool
    -- sqRemArray md = runST $ do
    --   ar <- MV.replicate md False
    --   let !stop = (md `quot` 2) + 1
    --       fill k
    --         | k < stop  = MV.unsafeWrite ar ((k*k) `rem` md) True >> fill (k+1)
    --         | otherwise = return ()
    --   MV.unsafeWrite ar 0 True
    --   MV.unsafeWrite ar 1 True
    --   fill 2
    --   V.unsafeFreeze ar

    -- sr256 :: V.Vector Bool
    -- sr256 = sqRemArray 256

    -- sr693 :: V.Vector Bool
    -- sr693 = sqRemArray 693

    -- sr325 :: V.Vector Bool
    -- sr325 = sqRemArray 325

    -----------------------------------------------------------------------------
    -- Specialisations for Int, Word, and Integer

    -- For @n <= 2^64@, the result of
    --
    -- > truncate (sqrt $ fromIntegral n)
    --
    -- is never too small and never more than one too large.
    -- The multiplication doesn't overflow for 32 or 64 bit Ints.
    isqrtInt' :: Int -> Int
    isqrtInt' :: Int -> Int
    isqrtInt' Int
    n
        | Int
    n Int -> Int -> Bool
    forall a. Ord a => a -> a -> Bool
    < Int
    rInt -> Int -> Int
    forall a. Num a => a -> a -> a
    *Int
    r   = Int
    rInt -> Int -> Int
    forall a. Num a => a -> a -> a
    -Int
    1
        | Bool
    otherwise = Int
    r
          where
            !r :: Int
    r = (Double -> Int
    forall b. Integral b => Double -> b
    forall a b. (RealFrac a, Integral b) => a -> b
    truncate :: Double -> Int) (Double -> Int) -> (Double -> Double) -> Double -> Int
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Double -> Double
    forall a. Floating a => a -> a
    sqrt (Double -> Int) -> Double -> Int
    forall a b. (a -> b) -> a -> b
    $ Int -> Double
    forall a b. (Integral a, Num b) => a -> b
    fromIntegral Int
    n

    isqrtInt64' :: Int64 -> Int64
    isqrtInt64' :: Int64 -> Int64
    isqrtInt64' Int64
    n
        | Int64
    n Int64 -> Int64 -> Bool
    forall a. Ord a => a -> a -> Bool
    < Int64
    rInt64 -> Int64 -> Int64
    forall a. Num a => a -> a -> a
    *Int64
    r   = Int64
    rInt64 -> Int64 -> Int64
    forall a. Num a => a -> a -> a
    -Int64
    1
        | Bool
    otherwise = Int64
    r
          where
            !r :: Int64
    r = (Double -> Int64
    forall b. Integral b => Double -> b
    forall a b. (RealFrac a, Integral b) => a -> b
    truncate :: Double -> Int64) (Double -> Int64) -> (Double -> Double) -> Double -> Int64
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Double -> Double
    forall a. Floating a => a -> a
    sqrt (Double -> Int64) -> Double -> Int64
    forall a b. (a -> b) -> a -> b
    $ Int64 -> Double
    forall a b. (Integral a, Num b) => a -> b
    fromIntegral Int64
    n

    -- Same for Word.
    isqrtWord :: Word -> Word
    isqrtWord :: Word -> Word
    isqrtWord Word
    n
        | Word
    n Word -> Word -> Bool
    forall a. Ord a => a -> a -> Bool
    < (Word
    rWord -> Word -> Word
    forall a. Num a => a -> a -> a
    *Word
    r)
          -- Double interprets values near maxBound as 2^64, we don't have that problem for 32 bits
          Bool -> Bool -> Bool
    || Word -> Int
    forall b. FiniteBits b => b -> Int
    finiteBitSize (Word
    0 :: Word) Int -> Int -> Bool
    forall a. Eq a => a -> a -> Bool
    == Int
    64 Bool -> Bool -> Bool
    && Word
    r Word -> Word -> Bool
    forall a. Eq a => a -> a -> Bool
    == Word
    4294967296
                    = Word
    rWord -> Word -> Word
    forall a. Num a => a -> a -> a
    -Word
    1
        | Bool
    otherwise = Word
    r
          where
            !r :: Word
    r = (Int -> Word
    forall a b. (Integral a, Num b) => a -> b
    fromIntegral :: Int -> Word) (Int -> Word) -> (Double -> Int) -> Double -> Word
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (Double -> Int
    forall b. Integral b => Double -> b
    forall a b. (RealFrac a, Integral b) => a -> b
    truncate :: Double -> Int) (Double -> Int) -> (Double -> Double) -> Double -> Int
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Double -> Double
    forall a. Floating a => a -> a
    sqrt (Double -> Word) -> Double -> Word
    forall a b. (a -> b) -> a -> b
    $ Word -> Double
    forall a b. (Integral a, Num b) => a -> b
    fromIntegral Word
    n

    isqrtWord64 :: Word64 -> Word64
    isqrtWord64 :: Word64 -> Word64
    isqrtWord64 Word64
    n
        | Word64
    n Word64 -> Word64 -> Bool
    forall a. Ord a => a -> a -> Bool
    < (Word64
    rWord64 -> Word64 -> Word64
    forall a. Num a => a -> a -> a
    *Word64
    r)
          -- Double interprets values near maxBound as 2^64
          Bool -> Bool -> Bool
    || Word64
    r Word64 -> Word64 -> Bool
    forall a. Eq a => a -> a -> Bool
    == Word64
    4294967296
                    = Word64
    rWord64 -> Word64 -> Word64
    forall a. Num a => a -> a -> a
    -Word64
    1
        | Bool
    otherwise = Word64
    r
          where
            !r :: Word64
    r = (Int64 -> Word64
    forall a b. (Integral a, Num b) => a -> b
    fromIntegral :: Int64 -> Word64) (Int64 -> Word64) -> (Double -> Int64) -> Double -> Word64
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (Double -> Int64
    forall b. Integral b => Double -> b
    forall a b. (RealFrac a, Integral b) => a -> b
    truncate :: Double -> Int64) (Double -> Int64) -> (Double -> Double) -> Double -> Int64
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Double -> Double
    forall a. Floating a => a -> a
    sqrt (Double -> Word64) -> Double -> Word64
    forall a b. (a -> b) -> a -> b
    $ Word64 -> Double
    forall a b. (Integral a, Num b) => a -> b
    fromIntegral Word64
    n

    {-# INLINE isqrtInteger #-}
    isqrtInteger :: Integer -> Integer
    isqrtInteger :: Integer -> Integer
    isqrtInteger = (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)
    karatsubaSqrt