-- |
-- Module:      Math.NumberTheory.Moduli.Equations
-- Copyright:   (c) 2018 Andrew Lelechenko
-- Licence:     MIT
-- Maintainer:  Andrew Lelechenko <andrew.lelechenko@gmail.com>
--
-- Polynomial modular equations.
--

{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE ViewPatterns        #-}

module Math.NumberTheory.Moduli.Equations
  ( solveLinear
  , solveQuadratic
  ) where

import Data.Constraint
import Data.Maybe
import Data.Mod
import GHC.Integer.GMP.Internals
import GHC.TypeNats (KnownNat, natVal)

import Math.NumberTheory.Moduli.Chinese
import Math.NumberTheory.Moduli.Singleton
import Math.NumberTheory.Moduli.Sqrt
import Math.NumberTheory.Primes
import Math.NumberTheory.Utils (recipMod)

-------------------------------------------------------------------------------
-- Linear equations

-- | Find all solutions of ax + b ≡ 0 (mod m).
--
-- >>> :set -XDataKinds
-- >>> solveLinear (6 :: Mod 10) 4 -- solving 6x + 4 ≡ 0 (mod 10)
-- [(1 `modulo` 10),(6 `modulo` 10)]
solveLinear
  :: KnownNat m
  => Mod m   -- ^ a
  -> Mod m   -- ^ b
  -> [Mod m] -- ^ list of x
solveLinear :: forall (m :: Nat). KnownNat m => Mod m -> Mod m -> [Mod m]
solveLinear Mod m
a Mod m
b = forall a b. (a -> b) -> [a] -> [b]
map forall a. Num a => Integer -> a
fromInteger forall a b. (a -> b) -> a -> b
$ Integer -> Integer -> Integer -> [Integer]
solveLinear' (forall a. Integral a => a -> Integer
toInteger (forall (n :: Nat) (proxy :: Nat -> *). KnownNat n => proxy n -> Nat
natVal Mod m
a)) (forall a. Integral a => a -> Integer
toInteger (forall (m :: Nat). Mod m -> Nat
unMod Mod m
a)) (forall a. Integral a => a -> Integer
toInteger (forall (m :: Nat). Mod m -> Nat
unMod Mod m
b))

solveLinear' :: Integer -> Integer -> Integer -> [Integer]
solveLinear' :: Integer -> Integer -> Integer -> [Integer]
solveLinear' Integer
m Integer
a Integer
b = case Integer -> Integer -> Integer -> Maybe Integer
solveLinearCoprime Integer
m' (Integer
a forall a. Integral a => a -> a -> a
`quot` Integer
d) (Integer
b forall a. Integral a => a -> a -> a
`quot` Integer
d) of
  Maybe Integer
Nothing -> []
  Just Integer
x  -> forall a b. (a -> b) -> [a] -> [b]
map (\Integer
i -> Integer
x forall a. Num a => a -> a -> a
+ Integer
m' forall a. Num a => a -> a -> a
* Integer
i) [Integer
0 .. Integer
d forall a. Num a => a -> a -> a
- Integer
1]
  where
    d :: Integer
d = Integer
m forall a. Integral a => a -> a -> a
`gcd` Integer
a forall a. Integral a => a -> a -> a
`gcd` Integer
b
    m' :: Integer
m' = Integer
m forall a. Integral a => a -> a -> a
`quot` Integer
d

solveLinearCoprime :: Integer -> Integer -> Integer -> Maybe Integer
solveLinearCoprime :: Integer -> Integer -> Integer -> Maybe Integer
solveLinearCoprime Integer
1 Integer
_ Integer
_ = forall a. a -> Maybe a
Just Integer
0
solveLinearCoprime Integer
m Integer
a Integer
b = (\Integer
a1 -> forall a. Num a => a -> a
negate Integer
b forall a. Num a => a -> a -> a
* Integer
a1 forall a. Integral a => a -> a -> a
`mod` Integer
m) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Integer -> Integer -> Maybe Integer
recipMod Integer
a Integer
m

-------------------------------------------------------------------------------
-- Quadratic equations

-- | Find all solutions of ax² + bx + c ≡ 0 (mod m).
--
-- >>> :set -XDataKinds
-- >>> solveQuadratic sfactors (1 :: Mod 32) 0 (-17) -- solving x² - 17 ≡ 0 (mod 32)
-- [(9 `modulo` 32),(25 `modulo` 32),(7 `modulo` 32),(23 `modulo` 32)]
solveQuadratic
  :: SFactors Integer m
  -> Mod m   -- ^ a
  -> Mod m   -- ^ b
  -> Mod m   -- ^ c
  -> [Mod m] -- ^ list of x
solveQuadratic :: forall (m :: Nat).
SFactors Integer m -> Mod m -> Mod m -> Mod m -> [Mod m]
solveQuadratic SFactors Integer m
sm Mod m
a Mod m
b Mod m
c = case forall a (m :: Nat).
Integral a =>
SFactors a m -> (() :: Constraint) :- KnownNat m
proofFromSFactors SFactors Integer m
sm of
  Sub Dict (KnownNat m)
(() :: Constraint) => Dict (KnownNat m)
Dict ->
    forall a b. (a -> b) -> [a] -> [b]
map forall a. Num a => Integer -> a
fromInteger
    forall a b. (a -> b) -> a -> b
$ forall a b. (a, b) -> a
fst
    forall a b. (a -> b) -> a -> b
$ [([Integer], Integer)] -> ([Integer], Integer)
combine
    forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map (\(Prime Integer
p, Word
n) -> (Integer -> Integer -> Integer -> Prime Integer -> Word -> [Integer]
solveQuadraticPrimePower Integer
a' Integer
b' Integer
c' Prime Integer
p Word
n, forall a. Prime a -> a
unPrime Prime Integer
p forall a b. (Num a, Integral b) => a -> b -> a
^ Word
n))
    forall a b. (a -> b) -> a -> b
$ forall a (m :: Nat). SFactors a m -> [(Prime a, Word)]
unSFactors SFactors Integer m
sm
  where
    a' :: Integer
a' = forall a. Integral a => a -> Integer
toInteger forall a b. (a -> b) -> a -> b
$ forall (m :: Nat). Mod m -> Nat
unMod Mod m
a
    b' :: Integer
b' = forall a. Integral a => a -> Integer
toInteger forall a b. (a -> b) -> a -> b
$ forall (m :: Nat). Mod m -> Nat
unMod Mod m
b
    c' :: Integer
c' = forall a. Integral a => a -> Integer
toInteger forall a b. (a -> b) -> a -> b
$ forall (m :: Nat). Mod m -> Nat
unMod Mod m
c

    combine :: [([Integer], Integer)] -> ([Integer], Integer)
    combine :: [([Integer], Integer)] -> ([Integer], Integer)
combine = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl
      (\([Integer]
xs, Integer
xm) ([Integer]
ys, Integer
ym) -> ([ forall a b. (a, b) -> a
fst forall a b. (a -> b) -> a -> b
$ forall a. HasCallStack => Maybe a -> a
fromJust forall a b. (a -> b) -> a -> b
$ forall a.
(Eq a, Ring a, Euclidean a) =>
(a, a) -> (a, a) -> Maybe (a, a)
chinese (Integer
x, Integer
xm) (Integer
y, Integer
ym) | Integer
x <- [Integer]
xs, Integer
y <- [Integer]
ys ], Integer
xm forall a. Num a => a -> a -> a
* Integer
ym))
      ([Integer
0], Integer
1)

solveQuadraticPrimePower
  :: Integer
  -> Integer
  -> Integer
  -> Prime Integer
  -> Word
  -> [Integer]
solveQuadraticPrimePower :: Integer -> Integer -> Integer -> Prime Integer -> Word -> [Integer]
solveQuadraticPrimePower Integer
a Integer
b Integer
c Prime Integer
p = Word -> [Integer]
go
  where
    go :: Word -> [Integer]
    go :: Word -> [Integer]
go Word
0 = [Integer
0]
    go Word
1 = Integer -> Integer -> Integer -> Prime Integer -> [Integer]
solveQuadraticPrime Integer
a Integer
b Integer
c Prime Integer
p
    go Word
k = forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (Word -> Integer -> [Integer]
liftRoot Word
k) (Word -> [Integer]
go (Word
k forall a. Num a => a -> a -> a
- Word
1))

    -- Hensel lifting
    -- https://en.wikipedia.org/wiki/Hensel%27s_lemma#Hensel_lifting
    liftRoot :: Word -> Integer -> [Integer]
    liftRoot :: Word -> Integer -> [Integer]
liftRoot Word
k Integer
r = case Integer -> Integer -> Maybe Integer
recipMod (Integer
2 forall a. Num a => a -> a -> a
* Integer
a forall a. Num a => a -> a -> a
* Integer
r forall a. Num a => a -> a -> a
+ Integer
b) Integer
pk of
      Maybe Integer
Nothing -> case Integer
fr of
        Integer
0 -> forall a b. (a -> b) -> [a] -> [b]
map (\Integer
i -> Integer
r forall a. Num a => a -> a -> a
+ Integer
pk forall a. Integral a => a -> a -> a
`quot` Integer
p' forall a. Num a => a -> a -> a
* Integer
i) [Integer
0 .. Integer
p' forall a. Num a => a -> a -> a
- Integer
1]
        Integer
_ -> []
      Just Integer
invDeriv -> [(Integer
r forall a. Num a => a -> a -> a
- Integer
fr forall a. Num a => a -> a -> a
* Integer
invDeriv) forall a. Integral a => a -> a -> a
`mod` Integer
pk]
      where
        pk :: Integer
pk = Integer
p' forall a b. (Num a, Integral b) => a -> b -> a
^ Word
k
        fr :: Integer
fr = (Integer
a forall a. Num a => a -> a -> a
* Integer
r forall a. Num a => a -> a -> a
* Integer
r forall a. Num a => a -> a -> a
+ Integer
b forall a. Num a => a -> a -> a
* Integer
r forall a. Num a => a -> a -> a
+ Integer
c) forall a. Integral a => a -> a -> a
`mod` Integer
pk

    p' :: Integer
    p' :: Integer
p' = forall a. Prime a -> a
unPrime Prime Integer
p

solveQuadraticPrime
  :: Integer
  -> Integer
  -> Integer
  -> Prime Integer
  -> [Integer]
solveQuadraticPrime :: Integer -> Integer -> Integer -> Prime Integer -> [Integer]
solveQuadraticPrime Integer
a Integer
b Integer
c (forall a. Prime a -> a
unPrime -> Integer
2 :: Integer)
  = case (forall a. Integral a => a -> Bool
even Integer
c, forall a. Integral a => a -> Bool
even (Integer
a forall a. Num a => a -> a -> a
+ Integer
b)) of
    (Bool
True, Bool
True) -> [Integer
0, Integer
1]
    (Bool
True, Bool
_)    -> [Integer
0]
    (Bool
_, Bool
False)   -> [Integer
1]
    (Bool, Bool)
_            -> []
solveQuadraticPrime Integer
a Integer
b Integer
c Prime Integer
p
  | Integer
a forall a. Integral a => a -> a -> a
`rem` Integer
p' forall a. Eq a => a -> a -> Bool
== Integer
0
  = Integer -> Integer -> Integer -> [Integer]
solveLinear' Integer
p' Integer
b Integer
c
  | Bool
otherwise
  = forall a b. (a -> b) -> [a] -> [b]
map (\Integer
n -> (Integer
n forall a. Num a => a -> a -> a
- Integer
b) forall a. Num a => a -> a -> a
* Integer -> Integer -> Integer
recipModInteger (Integer
2 forall a. Num a => a -> a -> a
* Integer
a) Integer
p' forall a. Integral a => a -> a -> a
`mod` Integer
p')
  forall a b. (a -> b) -> a -> b
$ Integer -> Prime Integer -> [Integer]
sqrtsModPrime (Integer
b forall a. Num a => a -> a -> a
* Integer
b forall a. Num a => a -> a -> a
- Integer
4 forall a. Num a => a -> a -> a
* Integer
a forall a. Num a => a -> a -> a
* Integer
c) Prime Integer
p
    where
      p' :: Integer
      p' :: Integer
p' = forall a. Prime a -> a
unPrime Prime Integer
p