module Math.Modulus where

import Control.Monad
import Data.Bits
import qualified Data.Foldable as F

{- |
>>> powMod 2 0 1000000007
1
>>> powMod 0 0 1000000007
1
>>> powMod 2 1000000005 1000000007
500000004
>>> powMod 2 (-1) 1000000007
500000004
>>> powMod 123456789 998244353 998244353
123456789
>>> powMod (-2) 2 1000000007
4
>>> powMod (-2) 3 1000000007
999999999
-}
powMod :: (Integral a) => a -> Int -> a -> a
powMod :: forall a. Integral a => a -> Int -> a -> a
powMod a
x Int
n a
m
  | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0 = a -> a -> Int -> a
forall {a}. (Bits a, Num a) => a -> a -> a -> a
go a
1 a
x Int
n
  | Int
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 = a
1
  | Bool
otherwise = a -> a -> Int -> a
forall {a}. (Bits a, Num a) => a -> a -> a -> a
go a
1 (a -> a -> a
forall a. Integral a => a -> a -> a
recipMod a
x a
m) (-Int
n)
  where
    go :: a -> a -> a -> a
go !a
acc !a
y !a
i
      | a
i a -> a -> a
forall a. Bits a => a -> a -> a
.&. a
1 a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
0 = a -> a -> a -> a
go a
acc (a
y a -> a -> a
forall a. Num a => a -> a -> a
* a
y a -> a -> a
forall a. Integral a => a -> a -> a
`rem` a
m) (a -> Int -> a
forall a. Bits a => a -> Int -> a
unsafeShiftR a
i Int
1)
      | a
i a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
1 = a
acc a -> a -> a
forall a. Num a => a -> a -> a
* a
y a -> a -> a
forall a. Integral a => a -> a -> a
`mod` a
m
      | Bool
otherwise = a -> a -> a -> a
go (a
acc a -> a -> a
forall a. Num a => a -> a -> a
* a
y a -> a -> a
forall a. Integral a => a -> a -> a
`rem` a
m) (a
y a -> a -> a
forall a. Num a => a -> a -> a
* a
y a -> a -> a
forall a. Integral a => a -> a -> a
`rem` a
m) (a -> Int -> a
forall a. Bits a => a -> Int -> a
unsafeShiftR (a
i a -> a -> a
forall a. Num a => a -> a -> a
- a
1) Int
1)
{-# INLINE powMod #-}

{- |
>>> recipMod 2 1000000007
500000004
>>> recipMod 10 1000000007
700000005
>>> recipMod 0 1000000007
0

prop> \x m -> not (m >= 1) || mod (x * recipMod x m) m == mod (gcd x m) m
+++ OK, passed 100 tests.
-}
recipMod :: (Integral a) => a -> a -> a
recipMod :: forall a. Integral a => a -> a -> a
recipMod a
x a
m = a -> a -> a -> a -> a
go (a -> a -> a
forall a. Integral a => a -> a -> a
mod a
x a
m) a
m a
1 a
0
  where
    go :: a -> a -> a -> a -> a
go !a
a !a
b !a
u !a
v
      | a
b a -> a -> Bool
forall a. Ord a => a -> a -> Bool
> a
0 = case a
a a -> a -> a
forall a. Integral a => a -> a -> a
`quot` a
b of
          a
q -> a -> a -> a -> a -> a
go a
b (a
a a -> a -> a
forall a. Num a => a -> a -> a
- (a
q a -> a -> a
forall a. Num a => a -> a -> a
* a
b)) a
v (a
u a -> a -> a
forall a. Num a => a -> a -> a
- (a
q a -> a -> a
forall a. Num a => a -> a -> a
* a
v))
      | Bool
otherwise = a
u a -> a -> a
forall a. Integral a => a -> a -> a
`mod` a
m
{-# INLINE recipMod #-}

{- | Extended Euclidean algorithm

@(x, y, g) = extGCD a b (a * x + b * y = g)@

>>> extGCD 3 5
(2,-1,1)
>>> extGCD 4 6
(-1,1,2)

prop> \a b -> not (a > 0 && b > 0) || let (_, _ , g) = extGCD a b in g == gcd a b
+++ OK, passed 100 tests.

prop> \a b -> not (a > 0 && b > 0) || let (x, y, g) = extGCD a b in a * x + b * y == g
+++ OK, passed 100 tests.

prop> \a b -> not (a > 0 && b > 0) || let (x, y, g) = extGCD a b in abs x <= div b g && abs y <= div a g
+++ OK, passed 100 tests.
-}
extGCD :: (Integral a) => a -> a -> (a, a, a)
extGCD :: forall a. Integral a => a -> a -> (a, a, a)
extGCD a
a0 a
b0 = a -> a -> a -> a -> (a, a, a)
go a
a0 a
b0 a
1 a
0
  where
    go :: a -> a -> a -> a -> (a, a, a)
go !a
a !a
b !a
u !a
v
      | a
b a -> a -> Bool
forall a. Ord a => a -> a -> Bool
> a
0 = case a -> a -> a
forall a. Integral a => a -> a -> a
quot a
a a
b of
          a
q -> a -> a -> a -> a -> (a, a, a)
go a
b (a
a a -> a -> a
forall a. Num a => a -> a -> a
- (a
q a -> a -> a
forall a. Num a => a -> a -> a
* a
b)) a
v (a
u a -> a -> a
forall a. Num a => a -> a -> a
- (a
q a -> a -> a
forall a. Num a => a -> a -> a
* a
v))
      | Bool
otherwise = (a
u, a -> a -> a
forall a. Integral a => a -> a -> a
quot (a
a a -> a -> a
forall a. Num a => a -> a -> a
- a
a0 a -> a -> a
forall a. Num a => a -> a -> a
* a
u) a
b0, a
a)
{-# INLINE extGCD #-}

{- | Solve @a * x + b * y == c@

If multiple solutions exists then return the minimum non-negative @x@.

If @a \/= 0, b \/= 0, a * x + b * y == c@
then @a (x - k * b \/ g) + b * (y + k * a \/ g) == c@

>>> linearDiophantine 3 5 1
Just (2,-1)
>>> linearDiophantine 3 5 3
Just (1,0)
>>> linearDiophantine 4 6 2
Just (2,-1)
>>> linearDiophantine 4 6 3
Nothing
prop> \a b c -> maybe True (\(x, y) -> a * x + b * y == c) $ linearDiophantine a b c
+++ OK, passed 100 tests.
-}
linearDiophantine :: (Integral a) => a -> a -> a -> Maybe (a, a)
linearDiophantine :: forall a. Integral a => a -> a -> a -> Maybe (a, a)
linearDiophantine a
_ a
_ a
0 = (a, a) -> Maybe (a, a)
forall a. a -> Maybe a
Just (a
0, a
0)
linearDiophantine a
0 a
0 a
_ = Maybe (a, a)
forall a. Maybe a
Nothing
linearDiophantine a
0 a
b a
c
  | (a
q, a
0) <- a -> a -> (a, a)
forall a. Integral a => a -> a -> (a, a)
divMod a
c a
b = (a, a) -> Maybe (a, a)
forall a. a -> Maybe a
Just (a
0, a
q)
  | Bool
otherwise = Maybe (a, a)
forall a. Maybe a
Nothing
linearDiophantine a
a a
0 a
c
  | (a
q, a
0) <- a -> a -> (a, a)
forall a. Integral a => a -> a -> (a, a)
divMod a
c a
a = (a, a) -> Maybe (a, a)
forall a. a -> Maybe a
Just (a
q, a
0)
  | Bool
otherwise = Maybe (a, a)
forall a. Maybe a
Nothing
linearDiophantine a
a a
b a
c | a
c a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
0 = a -> a -> a -> Maybe (a, a)
forall a. Integral a => a -> a -> a -> Maybe (a, a)
linearDiophantine (-a
a) (-a
b) (-a
c)
linearDiophantine a
a a
b a
c
  | a -> a -> a
forall a. Integral a => a -> a -> a
mod a
c a
g a -> a -> Bool
forall a. Ord a => a -> a -> Bool
> a
0 = Maybe (a, a)
forall a. Maybe a
Nothing
  | Bool
otherwise = (a, a) -> Maybe (a, a)
forall a. a -> Maybe a
Just (a
x, a
y)
  where
    (!a
x0, !a
_, !a
g) = a -> a -> (a, a, a)
forall a. Integral a => a -> a -> (a, a, a)
extGCD (a -> a
forall a. Num a => a -> a
abs a
a) (a -> a
forall a. Num a => a -> a
abs a
b)
    b' :: a
b' = a -> a -> a
forall a. Integral a => a -> a -> a
quot (a -> a
forall a. Num a => a -> a
abs a
b) a
g
    c' :: a
c' = a -> a -> a
forall a. Integral a => a -> a -> a
quot a
c a
g -- c > 0
    x :: a
x = a -> a -> a
forall a. Integral a => a -> a -> a
mod (a -> a
forall a. Num a => a -> a
signum a
a a -> a -> a
forall a. Num a => a -> a -> a
* a
c' a -> a -> a
forall a. Num a => a -> a -> a
* a
x0) a
b'
    y :: a
y = a -> a -> a
forall a. Integral a => a -> a -> a
div (a
c a -> a -> a
forall a. Num a => a -> a -> a
- a
a a -> a -> a
forall a. Num a => a -> a -> a
* a
x) a
b

{- | Chinese Remainder Theorem

@x = r0 (mod m0), x = r1 (mod m1) ==> x (mod (lcm m0 m1))@

>>> crt (10, 20) (10, 30)
Just (10,60)
>>> crt (10, 20) (10, 20)
Just (10,20)
>>> crt (10, 20) (11, 20)
Nothing
-}
crt :: (Integral a) => (a, a) -> (a, a) -> Maybe (a, a)
crt :: forall a. Integral a => (a, a) -> (a, a) -> Maybe (a, a)
crt (!a
r0, !a
m0) (!a
r1, !a
m1)
  | a -> a -> a
forall a. Integral a => a -> a -> a
mod (a
r1 a -> a -> a
forall a. Num a => a -> a -> a
- a
r0) a
g a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
0 = (a, a) -> Maybe (a, a)
forall a. a -> Maybe a
Just (a
r, a
m)
  | Bool
otherwise = Maybe (a, a)
forall a. Maybe a
Nothing
  where
    -- m0 * p + m1 * q == g
    (a
p, a
_q, a
g) = a -> a -> (a, a, a)
forall a. Integral a => a -> a -> (a, a, a)
extGCD a
m0 a
m1
    !m :: a
m = a
m0 a -> a -> a
forall a. Num a => a -> a -> a
* a -> a -> a
forall a. Integral a => a -> a -> a
quot a
m1 a
g
    !r :: a
r = a -> a -> a
forall a. Integral a => a -> a -> a
mod (a
r0 a -> a -> a
forall a. Num a => a -> a -> a
+ a
m0 a -> a -> a
forall a. Num a => a -> a -> a
* a -> a -> a
forall a. Integral a => a -> a -> a
quot (a
r1 a -> a -> a
forall a. Num a => a -> a -> a
- a
r0) a
g a -> a -> a
forall a. Num a => a -> a -> a
* a
p) a
m

{- | * @p0@, @p1@ are prime
   * @p0 /= p1@ or @r0 == r1@

>>> crt' (2,3) (3,5)
8
>>> crt' (3,5) (3,7)
3
>>> crt' (3,7) (3,7)
3
-}
crt' :: (Integral a) => (a, a) -> (a, a) -> a
crt' :: forall a. Integral a => (a, a) -> (a, a) -> a
crt' (a
r0, a
p0) (a
r1, a
p1) =
  a -> a -> a
forall a. Integral a => a -> a -> a
mod (a
r0 a -> a -> a
forall a. Num a => a -> a -> a
+ a
p0 a -> a -> a
forall a. Num a => a -> a -> a
* a -> a -> a
forall a. Integral a => a -> a -> a
recipMod a
p0 a
p1 a -> a -> a
forall a. Integral a => a -> a -> a
`rem` (a
p0 a -> a -> a
forall a. Num a => a -> a -> a
* a
p1) a -> a -> a
forall a. Num a => a -> a -> a
* (a
r1 a -> a -> a
forall a. Num a => a -> a -> a
- a
r0)) (a
p0 a -> a -> a
forall a. Num a => a -> a -> a
* a
p1)

{- |
>>> crts [(20,30),(30,50),(20,70)]
Just (230,1050)
>>> crts []
Just (0,1)
>>> crts [(1, 10), (2, 20), (3, 30)]
Nothing
-}
crts :: (Integral a) => [(a, a)] -> Maybe (a, a)
crts :: forall a. Integral a => [(a, a)] -> Maybe (a, a)
crts [(a, a)]
cs = ((a, a) -> ((a, a) -> Maybe (a, a)) -> (a, a) -> Maybe (a, a))
-> ((a, a) -> Maybe (a, a)) -> [(a, a)] -> (a, a) -> Maybe (a, a)
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (((a, a) -> Maybe (a, a))
-> ((a, a) -> Maybe (a, a)) -> (a, a) -> Maybe (a, a)
forall (m :: * -> *) a b c.
Monad m =>
(a -> m b) -> (b -> m c) -> a -> m c
(>=>) (((a, a) -> Maybe (a, a))
 -> ((a, a) -> Maybe (a, a)) -> (a, a) -> Maybe (a, a))
-> ((a, a) -> (a, a) -> Maybe (a, a))
-> (a, a)
-> ((a, a) -> Maybe (a, a))
-> (a, a)
-> Maybe (a, a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a, a) -> (a, a) -> Maybe (a, a)
forall a. Integral a => (a, a) -> (a, a) -> Maybe (a, a)
crt) (a, a) -> Maybe (a, a)
forall a. a -> Maybe a
forall (m :: * -> *) a. Monad m => a -> m a
return [(a, a)]
cs (a
0, a
1)

{- |
@x (mod m) (x = r[i] (mod m[i]))@

@gcd m[i] m[j] == 1@

/O(n^2)/

>>> garner [(2, 3), (3, 5), (2, 7)] 999
23
>>> garner [(2, 3), (3, 5), (2, 7)] 10
3
-}
garner :: (Integral a) => [(a, a)] -> a -> a
garner :: forall a. Integral a => [(a, a)] -> a -> a
garner [(a, a)]
rms a
m0 = [a] -> [(a, a)] -> a
go [] [(a, a)]
rms
  where
    ms :: [a]
ms = ((a, a) -> a) -> [(a, a)] -> [a]
forall a b. (a -> b) -> [a] -> [b]
map (a, a) -> a
forall a b. (a, b) -> b
snd [(a, a)]
rms
    go :: [a] -> [(a, a)] -> a
go [a]
cs ((a
r, a
m) : [(a, a)]
rest) = [a] -> [(a, a)] -> a
go ([a]
cs [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++ [a
c]) [(a, a)]
rest
      where
        -- a + b * c = r (mod m)
        (a
a, a
b) = ((a, a) -> (a, a) -> (a, a)) -> (a, a) -> [(a, a)] -> (a, a)
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
F.foldl' (a -> (a, a) -> (a, a) -> (a, a)
forall {b}. Integral b => b -> (b, b) -> (b, b) -> (b, b)
step a
m) (a
0, a
1) ([(a, a)] -> (a, a)) -> [(a, a)] -> (a, a)
forall a b. (a -> b) -> a -> b
$ [a] -> [a] -> [(a, a)]
forall a b. [a] -> [b] -> [(a, b)]
zip [a]
cs [a]
ms
        !c :: a
c = a -> a -> a
forall a. Integral a => a -> a -> a
rem (a -> a -> a
forall a. Integral a => a -> a -> a
mod (a
r a -> a -> a
forall a. Num a => a -> a -> a
- a
a) a
m a -> a -> a
forall a. Num a => a -> a -> a
* a -> a -> a
forall a. Integral a => a -> a -> a
recipMod a
b a
m) a
m
    go [a]
cs [] = (a, a) -> a
forall a b. (a, b) -> a
fst ((a, a) -> a) -> ([(a, a)] -> (a, a)) -> [(a, a)] -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((a, a) -> (a, a) -> (a, a)) -> (a, a) -> [(a, a)] -> (a, a)
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
F.foldl' (a -> (a, a) -> (a, a) -> (a, a)
forall {b}. Integral b => b -> (b, b) -> (b, b) -> (b, b)
step a
m0) (a
0, a
1) ([(a, a)] -> a) -> [(a, a)] -> a
forall a b. (a -> b) -> a -> b
$ [a] -> [a] -> [(a, a)]
forall a b. [a] -> [b] -> [(a, b)]
zip [a]
cs [a]
ms

    step :: b -> (b, b) -> (b, b) -> (b, b)
step b
m (!b
s, !b
p) (b
ci, b
mi) = (b
s', b
p')
      where
        !s' :: b
s' = b -> b -> b
forall a. Integral a => a -> a -> a
rem (b
s b -> b -> b
forall a. Num a => a -> a -> a
+ b -> b -> b
forall a. Integral a => a -> a -> a
rem (b
ci b -> b -> b
forall a. Num a => a -> a -> a
* b
p) b
m) b
m
        !p' :: b
p' = b -> b -> b
forall a. Integral a => a -> a -> a
rem (b
p b -> b -> b
forall a. Num a => a -> a -> a
* b
mi) b
m