    -- |
    -- 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 fourth powers. Efficient calculation of integer fourth
    -- roots and efficient testing for being a square's square.

    {-# LANGUAGE CPP       #-}
    {-# LANGUAGE MagicHash #-}
    {- HLINT ignore "Use fewer imports" -}

    module Math.NumberTheory.Roots.Fourth
        ( integerFourthRoot
        , integerFourthRoot'
        , exactFourthRoot
        , isFourthPower
        , isFourthPower'
        , isPossibleFourthPower
        ) where

    import Data.Bits (finiteBitSize, (.&.))
    import GHC.Exts (Int#, Ptr(..), int2Double#, double2Int#, isTrue#, sqrtDouble#, (<#))
    import Numeric.Natural (Natural)

    #ifdef MIN_VERSION_integer_gmp
    import GHC.Exts (uncheckedIShiftRA#, (*#), (-#))
    import GHC.Integer.GMP.Internals (Integer(..), shiftLInteger, shiftRInteger, sizeofBigNat#)
    import GHC.Integer.Logarithms (integerLog2#)
    #define IS S#
    #define IP Jp#
    #define bigNatSize sizeofBigNat
    #else
    import GHC.Exts (uncheckedShiftRL#, minusWord#, timesWord#)
    import GHC.Num.BigNat (bigNatSize#)
    import GHC.Num.Integer (Integer(..), integerLog2#, integerShiftR#, integerShiftL#)
    #endif

    import Math.NumberTheory.Utils.BitMask (indexBitSet)

    -- | Calculate the integer fourth root of a nonnegative number,
    --   that is, the largest integer @r@ with @r^4 <= n@.
    --   Throws an error on negaitve input.
    {-# SPECIALISE integerFourthRoot :: Int -> Int #-}
    {-# SPECIALISE integerFourthRoot :: Word -> Word #-}
    {-# SPECIALISE integerFourthRoot :: Integer -> Integer #-}
    {-# SPECIALISE integerFourthRoot :: Natural -> Natural #-}
    integerFourthRoot :: Integral a => a -> a
    integerFourthRoot :: forall a. Integral a => a -> a
    integerFourthRoot 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]
    "integerFourthRoot: negative argument"
        | Bool
    otherwise = a -> a
    forall a. Integral a => a -> a
    integerFourthRoot' a
    n

    -- | Calculate the integer fourth root of a nonnegative number,
    --   that is, the largest integer @r@ with @r^4 <= n@.
    --   The condition is /not/ checked.
    {-# RULES
    "integerFourthRoot'/Int"     integerFourthRoot' = biSqrtInt
    "integerFourthRoot'/Word"    integerFourthRoot' = biSqrtWord
    "integerFourthRoot'/Integer" integerFourthRoot' = biSqrtIgr
      #-}
    {-# INLINE [1] integerFourthRoot' #-}
    integerFourthRoot' :: Integral a => a -> a
    integerFourthRoot' :: forall a. Integral a => a -> a
    integerFourthRoot' a
    0 = a
    0
    integerFourthRoot' a
    n = a -> a -> a
    forall a. Integral a => a -> a -> a
    newton4 a
    n (a -> a
    forall a. Integral a => a -> a
    approxBiSqrt a
    n)

    -- | Returns @Nothing@ if @n@ is not a fourth power,
    --   @Just r@ if @n == r^4@ and @r >= 0@.
    {-# SPECIALISE exactFourthRoot :: Int -> Maybe Int #-}
    {-# SPECIALISE exactFourthRoot :: Word -> Maybe Word #-}
    {-# SPECIALISE exactFourthRoot :: Integer -> Maybe Integer #-}
    {-# SPECIALISE exactFourthRoot :: Natural -> Maybe Natural #-}
    exactFourthRoot :: Integral a => a -> Maybe a
    exactFourthRoot :: forall a. Integral a => a -> Maybe a
    exactFourthRoot a
    0 = a -> Maybe a
    forall a. a -> Maybe a
    Just a
    0
    exactFourthRoot a
    n
        | a
    n a -> a -> Bool
    forall a. Ord a => a -> a -> Bool
    < a
    0     = Maybe a
    forall a. Maybe a
    Nothing
        | a -> Bool
    forall a. Integral a => a -> Bool
    isPossibleFourthPower a
    n Bool -> Bool -> Bool
    && a
    r2a -> a -> a
    forall a. Num a => a -> a -> a
    *a
    r2 a -> a -> Bool
    forall a. Eq a => a -> a -> Bool
    == a
    n = a -> Maybe a
    forall a. a -> Maybe a
    Just a
    r
        | Bool
    otherwise = Maybe a
    forall a. Maybe a
    Nothing
          where
            r :: a
    r = a -> a
    forall a. Integral a => a -> a
    integerFourthRoot' a
    n
            r2 :: a
    r2 = a
    ra -> a -> a
    forall a. Num a => a -> a -> a
    *a
    r

    -- | Test whether an integer is a fourth power.
    --   First nonnegativity is checked, then the unchecked
    --   test is called.
    {-# SPECIALISE isFourthPower :: Int -> Bool #-}
    {-# SPECIALISE isFourthPower :: Word -> Bool #-}
    {-# SPECIALISE isFourthPower :: Integer -> Bool #-}
    {-# SPECIALISE isFourthPower :: Natural -> Bool #-}
    isFourthPower :: Integral a => a -> Bool
    isFourthPower :: forall a. Integral a => a -> Bool
    isFourthPower a
    0 = Bool
    True
    isFourthPower 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
    isFourthPower' a
    n

    -- | Test whether a nonnegative number is a fourth power.
    --   The condition is /not/ checked. If a number passes the
    --   'isPossibleFourthPower' test, its integer fourth root
    --   is calculated.
    {-# SPECIALISE isFourthPower' :: Int -> Bool #-}
    {-# SPECIALISE isFourthPower' :: Word -> Bool #-}
    {-# SPECIALISE isFourthPower' :: Integer -> Bool #-}
    {-# SPECIALISE isFourthPower' :: Natural -> Bool #-}
    isFourthPower' :: Integral a => a -> Bool
    isFourthPower' :: forall a. Integral a => a -> Bool
    isFourthPower' a
    n = a -> Bool
    forall a. Integral a => a -> Bool
    isPossibleFourthPower a
    n Bool -> Bool -> Bool
    && a
    r2a -> a -> a
    forall a. Num a => a -> a -> a
    *a
    r2 a -> a -> Bool
    forall a. Eq a => a -> a -> Bool
    == a
    n
      where
        r :: a
    r = a -> a
    forall a. Integral a => a -> a
    integerFourthRoot' a
    n
        r2 :: a
    r2 = a
    ra -> a -> a
    forall a. Num a => a -> a -> a
    *a
    r

    -- | Test whether a nonnegative number is a possible fourth power.
    --   The condition is /not/ checked.
    --   This eliminates about 99.958% of numbers.
    {-# SPECIALISE isPossibleFourthPower :: Int -> Bool #-}
    {-# SPECIALISE isPossibleFourthPower :: Word -> Bool #-}
    {-# SPECIALISE isPossibleFourthPower :: Integer -> Bool #-}
    {-# SPECIALISE isPossibleFourthPower :: Natural -> Bool #-}
    isPossibleFourthPower :: Integral a => a -> Bool
    isPossibleFourthPower :: forall a. Integral a => a -> Bool
    isPossibleFourthPower 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
    mask425 (Integer -> Int
    forall a. Num a => Integer -> a
    fromInteger (Integer
    n Integer -> Integer -> Integer
    forall a. Integral a => a -> a -> a
    `rem` Integer
    425))
      Bool -> Bool -> Bool
    && Ptr Word -> Int -> Bool
    indexBitSet Ptr Word
    mask377 (Integer -> Int
    forall a. Num a => Integer -> a
    fromInteger (Integer
    n Integer -> Integer -> Integer
    forall a. Integral a => a -> a -> a
    `rem` Integer
    377))
      where
        n :: Integer
    n = a -> Integer
    forall a. Integral a => a -> Integer
    toInteger a
    n'

    {-# SPECIALISE newton4 :: Integer -> Integer -> Integer #-}
    newton4 :: Integral a => a -> a -> a
    newton4 :: forall a. Integral a => a -> a -> a
    newton4 a
    n a
    a = a -> a
    go (a -> a
    step a
    a)
          where
            step :: a -> a
    step a
    k = (a
    3a -> a -> a
    forall a. Num a => a -> a -> a
    *a
    k a -> a -> a
    forall a. Num a => a -> a -> a
    + a
    n a -> a -> a
    forall a. Integral a => a -> a -> a
    `quot` (a
    ka -> a -> a
    forall a. Num a => a -> a -> a
    *a
    ka -> a -> a
    forall a. Num a => a -> a -> a
    *a
    k)) a -> a -> a
    forall a. Integral a => a -> a -> a
    `quot` a
    4
            go :: a -> a
    go a
    k
                | a
    m a -> a -> Bool
    forall a. Ord a => a -> a -> Bool
    < a
    k     = a -> a
    go a
    m
                | Bool
    otherwise = a
    k
                  where
                    m :: a
    m = a -> a
    step a
    k

    {-# SPECIALISE approxBiSqrt :: Integer -> Integer #-}
    approxBiSqrt :: Integral a => a -> a
    approxBiSqrt :: forall a. Integral a => a -> a
    approxBiSqrt = Integer -> a
    forall a. Num a => Integer -> a
    fromInteger (Integer -> a) -> (a -> Integer) -> a -> a
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Integer -> Integer
    appBiSqrt (Integer -> Integer) -> (a -> Integer) -> a -> Integer
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . a -> Integer
    forall a b. (Integral a, Num b) => a -> b
    fromIntegral

    -- Find a fairly good approximation to the fourth root.
    -- About 48 bits should be correct for large Integers.
    appBiSqrt :: Integer -> Integer
    appBiSqrt :: Integer -> Integer
    appBiSqrt (IS Int#
    i#) = Int# -> Integer
    IS (Double# -> Int#
    double2Int# (Double# -> Double#
    sqrtDouble# (Double# -> Double#
    sqrtDouble# (Int# -> Double#
    int2Double# Int#
    i#))))
    appBiSqrt n :: Integer
    n@(IP ByteArray#
    bn#)
        | Int# -> Bool
    isTrue# (ByteArray# -> Int#
    bigNatSize# ByteArray#
    bn# Int# -> Int# -> Int#
    <# Int#
    thresh#) =
              Double -> Integer
    forall b. Integral b => Double -> b
    forall a b. (RealFrac a, Integral b) => a -> b
    floor (Double -> Double
    forall a. Floating a => a -> a
    sqrt (Double -> Double) -> (Double -> Double) -> Double -> Double
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Double -> Double
    forall a. Floating a => a -> a
    sqrt (Double -> Double) -> Double -> Double
    forall a b. (a -> b) -> a -> b
    $ Integer -> Double
    forall a. Num a => Integer -> a
    fromInteger Integer
    n :: Double)
        | Bool
    otherwise = case Integer -> Word#
    integerLog2# Integer
    n of
    #ifdef MIN_VERSION_integer_gmp
                        l# -> case uncheckedIShiftRA# l# 2# -# 47# of
                                h# -> case shiftRInteger n (4# *# h#) of
                                        m -> case floor (sqrt $ sqrt $ fromInteger m :: Double) of
                                                r -> shiftLInteger r h#
    #else
                        Word#
    l# -> case Word# -> Int# -> Word#
    uncheckedShiftRL# Word#
    l# Int#
    2# Word# -> Word# -> Word#
    `minusWord#` Word#
    47## of
                                Word#
    h# -> case Integer -> Word# -> Integer
    integerShiftR# Integer
    n (Word#
    4## Word# -> Word# -> Word#
    `timesWord#` Word#
    h#) of
                                        Integer
    m -> case Double -> Integer
    forall b. Integral b => Double -> b
    forall a b. (RealFrac a, Integral b) => a -> b
    floor (Double -> Double
    forall a. Floating a => a -> a
    sqrt (Double -> Double) -> Double -> Double
    forall a b. (a -> b) -> a -> b
    $ Double -> Double
    forall a. Floating a => a -> a
    sqrt (Double -> Double) -> Double -> Double
    forall a b. (a -> b) -> a -> b
    $ Integer -> Double
    forall a. Num a => Integer -> a
    fromInteger Integer
    m :: Double) of
                                                Integer
    r -> Integer -> Word# -> Integer
    integerShiftL# Integer
    r Word#
    h#
    #endif
        where
            -- threshold for shifting vs. direct fromInteger
            -- we shift when we expect more than 256 bits
            thresh# :: Int#
            thresh# :: Int#
    thresh# = if Word -> Int
    forall b. FiniteBits b => b -> Int
    finiteBitSize (Word
    0 :: Word) Int -> Int -> Bool
    forall a. Eq a => a -> a -> Bool
    == Int
    64 then Int#
    5# else Int#
    9#

    -- There's already a check for negative in integerFourthRoot,
    -- but integerFourthRoot' is exported directly too.
    appBiSqrt Integer
    _ = [Char] -> Integer
    forall a. HasCallStack => [Char] -> a
    error [Char]
    "integerFourthRoot': negative argument"

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

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

    mask425 :: Ptr Word
    mask425 :: Ptr Word
    mask425 = Addr# -> Ptr Word
    forall a. Addr# -> Ptr a
    Ptr Addr#
    "\ETX\NUL!\NUL\NUL\NUL\f\NUL\NUL\NULB\NUL \EOT\NUL\NUL\NUL\SOH\NUL\NUL@\b\NUL\132\NUL\SOH\NUL \STX\NUL\NUL\b\SOH\128\DLE\NUL\NUL\NUL\EOT\NUL\NUL\NUL!\NUL\DLE\STX\128\NUL\128\NUL\NUL\NUL \NUL"#

    mask377 :: Ptr Word
    mask377 :: Ptr Word
    mask377 = Addr# -> Ptr Word
    forall a. Addr# -> Ptr a
    Ptr Addr#
    "\ETX\NUL\SOH \NUL\NUL0\NUL\STXD\130@\NUL\b \NUL\NUL\b\EOT\SOH \ACK\NUL\NUL@\DLE\NUL\NUL\NUL\NUL\NUL\SOH!\NUL\NUL@\NUL\NUL\NUL\n@\NUL\b\NUL\NUL\DLE \NUL"#

    -- biSqRes256 :: V.Vector Bool
    -- biSqRes256 = runST $ do
    --     ar <- MV.replicate 256 False
    --     let note 257 = return ()
    --         note i = MV.unsafeWrite ar i True >> note (i+16)
    --     MV.unsafeWrite ar 0 True
    --     MV.unsafeWrite ar 16 True
    --     note 1
    --     V.unsafeFreeze ar

    -- biSqRes425 :: V.Vector Bool
    -- biSqRes425 = runST $ do
    --     ar <- MV.replicate 425 False
    --     let note 154 = return ()
    --         note i = MV.unsafeWrite ar ((i*i*i*i) `rem` 425) True >> note (i+1)
    --     note 0
    --     V.unsafeFreeze ar

    -- biSqRes377 :: V.Vector Bool
    -- biSqRes377 = runST $ do
    --     ar <- MV.replicate 377 False
    --     let note 144 = return ()
    --         note i = MV.unsafeWrite ar ((i*i*i*i) `rem` 377) True >> note (i+1)
    --     note 0
    --     V.unsafeFreeze ar

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

    biSqRootIntLimit :: Int
    biSqRootIntLimit :: Int
    biSqRootIntLimit = if Word -> Int
    forall b. FiniteBits b => b -> Int
    finiteBitSize (Word
    0 :: Word) Int -> Int -> Bool
    forall a. Eq a => a -> a -> Bool
    == Int
    64 then Int
    55108 else Int
    215

    biSqrtInt :: Int -> Int
    biSqrtInt :: Int -> Int
    biSqrtInt Int
    0 = Int
    0
    biSqrtInt Int
    n
        | Int
    r Int -> Int -> Bool
    forall a. Ord a => a -> a -> Bool
    > Int
    biSqRootIntLimit = Int
    biSqRootIntLimit
        | Int
    n Int -> Int -> Bool
    forall a. Ord a => a -> a -> Bool
    < Int
    r4    = Int
    rInt -> Int -> Int
    forall a. Num a => a -> a -> a
    -Int
    1
        | Bool
    otherwise = Int
    r
          where
            x :: Double
            x :: Double
    x = Int -> Double
    forall a b. (Integral a, Num b) => a -> b
    fromIntegral Int
    n
            -- timed faster than x**0.25, never too small
            r :: Int
    r = Double -> Int
    forall b. Integral b => Double -> b
    forall a b. (RealFrac a, Integral b) => a -> b
    truncate (Double -> Double
    forall a. Floating a => a -> a
    sqrt (Double -> Double
    forall a. Floating a => a -> a
    sqrt Double
    x))
            r2 :: Int
    r2 = Int
    rInt -> Int -> Int
    forall a. Num a => a -> a -> a
    *Int
    r
            r4 :: Int
    r4 = Int
    r2Int -> Int -> Int
    forall a. Num a => a -> a -> a
    *Int
    r2

    biSqRootWordLimit :: Word
    biSqRootWordLimit :: Word
    biSqRootWordLimit = if Word -> Int
    forall b. FiniteBits b => b -> Int
    finiteBitSize (Word
    0 :: Word) Int -> Int -> Bool
    forall a. Eq a => a -> a -> Bool
    == Int
    64 then Word
    65535 else Word
    255

    biSqrtWord :: Word -> Word
    biSqrtWord :: Word -> Word
    biSqrtWord Word
    0 = Word
    0
    biSqrtWord Word
    n
        | Word
    r Word -> Word -> Bool
    forall a. Ord a => a -> a -> Bool
    > Word
    biSqRootWordLimit = Word
    biSqRootWordLimit
        | Word
    n Word -> Word -> Bool
    forall a. Ord a => a -> a -> Bool
    < Word
    r4    = Word
    rWord -> Word -> Word
    forall a. Num a => a -> a -> a
    -Word
    1
        | Bool
    otherwise = Word
    r
          where
            x :: Double
            x :: Double
    x = Word -> Double
    forall a b. (Integral a, Num b) => a -> b
    fromIntegral Word
    n
            r :: Word
    r = Double -> Word
    forall b. Integral b => Double -> b
    forall a b. (RealFrac a, Integral b) => a -> b
    truncate (Double -> Double
    forall a. Floating a => a -> a
    sqrt (Double -> Double
    forall a. Floating a => a -> a
    sqrt Double
    x))
            r2 :: Word
    r2 = Word
    rWord -> Word -> Word
    forall a. Num a => a -> a -> a
    *Word
    r
            r4 :: Word
    r4 = Word
    r2Word -> Word -> Word
    forall a. Num a => a -> a -> a
    *Word
    r2

    biSqrtIgr :: Integer -> Integer
    biSqrtIgr :: Integer -> Integer
    biSqrtIgr Integer
    0 = Integer
    0
    biSqrtIgr Integer
    n = Integer -> Integer -> Integer
    forall a. Integral a => a -> a -> a
    newton4 Integer
    n (Integer -> Integer
    forall a. Integral a => a -> a
    approxBiSqrt Integer
    n)