module Math.Modulus.Log where

import qualified Data.IntMap.Strict as IM

import Math.Modulus (powMod)

{- |
Baby-step Giant-step

@a^x = b (mod p)@ p is prime

/O(sqrt P * log P)/

>>> logMod 3 27 998244353
Just 3
>>> logMod 3 123456789 998244353
Just 772453214
>>> logMod 1 2 1000000007
Nothing
-}
logMod :: Int -> Int -> Int -> Maybe Int
logMod :: Key -> Key -> Key -> Maybe Key
logMod Key
a Key
b Key
p = Key -> Key -> Maybe Key
go Key
0 Key
b
  where
    !sqrtP :: Key
sqrtP = Double -> Key
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
ceiling (Double -> Key) -> (Double -> Double) -> Double -> Key
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Floating a => a -> a
sqrt @Double (Double -> Key) -> Double -> Key
forall a b. (a -> b) -> a -> b
$ Key -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Key
p
    !g :: Key
g = Key -> Key -> Key -> Key
forall a. Integral a => a -> Key -> a -> a
powMod Key
a (-Key
sqrtP) Key
p
    babyStep :: Key -> Key
babyStep Key
x = Key
a Key -> Key -> Key
forall a. Num a => a -> a -> a
* Key
x Key -> Key -> Key
forall a. Integral a => a -> a -> a
`rem` Key
p
    giantStep :: Key -> Key
giantStep Key
x = Key
g Key -> Key -> Key
forall a. Num a => a -> a -> a
* Key
x Key -> Key -> Key
forall a. Integral a => a -> a -> a
`rem` Key
p

    table :: IM.IntMap Int
    !table :: IntMap Key
table = [(Key, Key)] -> IntMap Key
forall a. [(Key, a)] -> IntMap a
IM.fromList ([(Key, Key)] -> IntMap Key) -> [(Key, Key)] -> IntMap Key
forall a b. (a -> b) -> a -> b
$ [Key] -> [Key] -> [(Key, Key)]
forall a b. [a] -> [b] -> [(a, b)]
zip ((Key -> Key) -> Key -> [Key]
forall a. (a -> a) -> a -> [a]
iterate Key -> Key
babyStep Key
1) [Key
0 .. Key
sqrtP]

    go :: Key -> Key -> Maybe Key
go !Key
i !Key
x
      | Key
i Key -> Key -> Bool
forall a. Ord a => a -> a -> Bool
< Key
sqrtP = case Key -> IntMap Key -> Maybe Key
forall a. Key -> IntMap a -> Maybe a
IM.lookup Key
x IntMap Key
table of
          Just Key
j -> Key -> Maybe Key
forall a. a -> Maybe a
Just (Key -> Maybe Key) -> Key -> Maybe Key
forall a b. (a -> b) -> a -> b
$! Key
i Key -> Key -> Key
forall a. Num a => a -> a -> a
* Key
sqrtP Key -> Key -> Key
forall a. Num a => a -> a -> a
+ Key
j
          Maybe Key
Nothing -> Key -> Key -> Maybe Key
go (Key
i Key -> Key -> Key
forall a. Num a => a -> a -> a
+ Key
1) (Key -> Maybe Key) -> Key -> Maybe Key
forall a b. (a -> b) -> a -> b
$ Key -> Key
giantStep Key
x
      | Bool
otherwise = Maybe Key
forall a. Maybe a
Nothing