-- |
-- Module:      Math.NumberTheory.Moduli.SqrtTests
-- Copyright:   (c) 2016 Andrew Lelechenko
-- Licence:     MIT
-- Maintainer:  Andrew Lelechenko <andrew.lelechenko@gmail.com>
--
-- Tests for Math.NumberTheory.Moduli.Sqrt
--

{-# LANGUAGE ViewPatterns    #-}

{-# OPTIONS_GHC -fno-warn-type-defaults #-}
{-# OPTIONS_GHC -Wno-incomplete-uni-patterns #-}

module Math.NumberTheory.Moduli.SqrtTests
  ( testSuite
  ) where

import Test.Tasty
import Test.Tasty.HUnit

import Control.Arrow
import Data.List (sort)
import qualified Data.List.NonEmpty as NE
import Data.Maybe (fromJust, listToMaybe)
import Numeric.Natural

import Math.NumberTheory.Moduli hiding (invertMod)
import Math.NumberTheory.Primes (unPrime, isPrime, Prime)
import Math.NumberTheory.TestUtils

unwrapPP :: (Prime Integer, Power Word) -> (Prime Integer, Word)
unwrapPP (p, Power e) = (p, e `mod` 5)

nubOrd :: Ord a => [a] -> [a]
nubOrd = map NE.head . NE.group . sort

-- | Check that 'sqrtMod' is defined iff a quadratic residue exists.
--   Also check that the result is a solution of input modular equation.
sqrtsModPrimeProperty1 :: AnySign Integer -> Prime Integer -> Bool
sqrtsModPrimeProperty1 (AnySign n) p'@(unPrime -> p) = case sqrtsModPrime n p' of
  []     -> jacobi n p == MinusOne
  rt : _ -> (p == 2 || jacobi n p /= MinusOne) && (rt ^ 2 - n) `rem` p == 0

sqrtsModPrimeProperty2 :: AnySign Integer -> Prime Integer -> Bool
sqrtsModPrimeProperty2 (AnySign n) p'@(unPrime -> p) = all (\rt -> (rt ^ 2 - n) `rem` p == 0) (sqrtsModPrime n p')

sqrtsModPrimeProperty3 :: AnySign Integer -> Prime Integer -> Bool
sqrtsModPrimeProperty3 (AnySign n) p'@(unPrime -> p) = nubOrd rts == sort rts
  where
    rts = map (`mod` p) $ sqrtsModPrime n p'

sqrtsModPrimeProperty4 :: AnySign Integer -> Prime Integer -> Bool
sqrtsModPrimeProperty4 (AnySign n) p'@(unPrime -> p) = all (\rt -> rt >= 0 && rt < p) (sqrtsModPrime n p')

tonelliShanksProperty1 :: Positive Integer -> Prime Integer -> Bool
tonelliShanksProperty1 (Positive n) p'@(unPrime -> p) = p `mod` 4 /= 1 || jacobi n p /= One || rt ^ 2 `mod` p == n `mod` p
  where
    rt : _ = sqrtsModPrime n p'

tonelliShanksProperty2 :: Prime Integer -> Bool
tonelliShanksProperty2 p'@(unPrime -> p) = p `mod` 4 /= 1 || (rt ^ 2 - n) `rem` p == 0
  where
    n : _ = filter (\s -> jacobi s p == One) [2..p-1]
    rt : _ = sqrtsModPrime n p'

tonelliShanksProperty3 :: Prime Integer -> Bool
tonelliShanksProperty3 p'@(unPrime -> p)
  = p `mod` 4 /= 1
  || rt ^ 2 `mod` p == p - 1
  where
    rt : _ = sqrtsModPrime (-1) p'

tonelliShanksSpecialCases :: Assertion
tonelliShanksSpecialCases =
  assertEqual "OEIS A002224" (map Just [6, 32, 219, 439, 1526, 2987, 22193, 11740, 13854, 91168, 326277, 232059, 3230839, 4379725, 11754394, 32020334, 151024619, 345641931, 373671108, 1857111865, 8110112775, 4184367042]) rts
  where
    ps :: [Integer]
    ps = [17, 73, 241, 1009, 2689, 8089, 33049, 53881, 87481, 483289, 515761, 1083289, 3818929, 9257329, 22000801, 48473881, 175244281, 427733329, 898716289, 8114538721, 9176747449, 23616331489]
    rts = map (listToMaybe . sqrtsModPrime 2 . fromJust . isPrime) ps

sqrtsModPrimePowerProperty1 :: AnySign Integer -> (Prime Integer, Power Word) -> Bool
sqrtsModPrimePowerProperty1 (AnySign n) (p'@(unPrime -> p), Power e) = gcd n p > 1
  || all (\rt -> (rt ^ 2 - n) `rem` (p ^ e) == 0) (sqrtsModPrimePower n p' e)

sqrtsModPrimePowerProperty2 :: AnySign Integer -> Power Word -> Bool
sqrtsModPrimePowerProperty2 n e = sqrtsModPrimePowerProperty1 n (fromJust $ isPrime (2 :: Integer), e)

sqrtsModPrimePowerProperty3 :: AnySign Integer -> (Prime Integer, Power Word) -> Bool
sqrtsModPrimePowerProperty3 (AnySign n) (p'@(unPrime -> p), Power e') = nubOrd rts == sort rts
  where
    e = e' `mod` 5
    m = p ^ e
    rts = map (`mod` m) $ sqrtsModPrimePower n p' e

sqrtsModPrimePowerProperty4 :: AnySign Integer -> Power Word -> Bool
sqrtsModPrimePowerProperty4 n e = sqrtsModPrimePowerProperty3 n (fromJust $ isPrime (2 :: Integer), e)

sqrtsModPrimePowerProperty5 :: AnySign Integer -> (Prime Integer, Power Word) -> Bool
sqrtsModPrimePowerProperty5 (AnySign n) (p'@(unPrime -> p), Power e') = all (\rt -> rt >= 0 && rt < m) rts
  where
    e = e' `mod` 5
    m = p ^ e
    rts = sqrtsModPrimePower n p' e

sqrtsModPrimePowerProperty6 :: AnySign Integer -> Power Word -> Bool
sqrtsModPrimePowerProperty6 n e = sqrtsModPrimePowerProperty5 n (fromJust $ isPrime (2 :: Integer), e)

sqrtsModPrimePowerSpecialCase1 :: Assertion
sqrtsModPrimePowerSpecialCase1 =
  assertEqual "should be equal" [0, 2] (sort (sqrtsModPrimePower 16 (fromJust (isPrime (2 :: Integer))) 2))

sqrtsModPrimePowerSpecialCase2 :: Assertion
sqrtsModPrimePowerSpecialCase2 =
  assertEqual "should be equal" [4, 5] (sort (sqrtsModPrimePower 16 (fromJust (isPrime (3 :: Integer))) 2))

sqrtsModPrimePowerSpecialCase3 :: Assertion
sqrtsModPrimePowerSpecialCase3 =
  assertEqual "should be equal" [0, 3, 6] (sort (sqrtsModPrimePower 0 (fromJust (isPrime (3 :: Integer))) 2))

sqrtsModPrimePowerSpecialCase4 :: Assertion
sqrtsModPrimePowerSpecialCase4 =
  assertEqual "should be equal" [0, 9, 18] (sort (sqrtsModPrimePower 0 (fromJust (isPrime (3 :: Integer))) 3))

sqrtsModPrimePowerSpecialCase5 :: Assertion
sqrtsModPrimePowerSpecialCase5 =
  assertEqual "should be equal" [0, 4, 8, 12] (sort (sqrtsModPrimePower 0 (fromJust (isPrime (2 :: Integer))) 4))

sqrtsModPrimePowerSpecialCase6 :: Assertion
sqrtsModPrimePowerSpecialCase6 =
  assertEqual "should be equal" [3, 6, 12, 15, 21, 24] (sort (sqrtsModPrimePower 9 (fromJust (isPrime (3 :: Integer))) 3))

sqrtsModPrimePowerSpecialCase7 :: Assertion
sqrtsModPrimePowerSpecialCase7 =
  assertEqual "should be equal" [2, 6] (sort (sqrtsModPrimePower 4 (fromJust (isPrime (2 :: Integer))) 3))

sqrtsModPrimePowerSpecialCase8 :: Assertion
sqrtsModPrimePowerSpecialCase8 =
  assertEqual "should be equal" [1, 3] (sort (sqrtsModPrimePower 1 (fromJust (isPrime (2 :: Integer))) 2))

sqrtsModPrimePowerSpecialCase9 :: Assertion
sqrtsModPrimePowerSpecialCase9 =
  assertEqual "should be equal" [] (sort (sqrtsModPrimePower (-1) (fromJust (isPrime (2 :: Integer))) 2))

sqrtsModPrimePowerSpecialCase10 :: Assertion
sqrtsModPrimePowerSpecialCase10 =
  assertEqual "should be equal" [2, 6, 10, 14] (sort (sqrtsModPrimePower 4 (fromJust (isPrime (2 :: Integer))) 4))

sqrtsModPrimePowerSpecialCase11 :: Assertion
sqrtsModPrimePowerSpecialCase11 =
  assertEqual "should be equal" [4,12,20,28,36,44,52,60] (sort (sqrtsModPrimePower 16 (fromJust (isPrime (2 :: Integer))) 6))

sqrtsModFactorisationProperty1 :: AnySign Integer -> [(Prime Integer, Power Word)] -> Bool
sqrtsModFactorisationProperty1 (AnySign n) (take 10 . map unwrapPP -> pes'@(map (first unPrime) -> pes))
  = nubOrd ps /= sort ps || all
    (\rt -> all (\(p, e) -> (rt ^ 2 - n) `rem` (p ^ e) == 0) pes)
    (take 1000 $ sqrtsModFactorisation n pes')
  where
    ps = map fst pes

sqrtsModFactorisationProperty2 :: AnySign Integer -> [(Prime Integer, Power Word)] -> Bool
sqrtsModFactorisationProperty2 (AnySign n) (take 10 . map unwrapPP -> pes'@(map (first unPrime) -> pes))
  = nubOrd ps /= sort ps || nubOrd rts == sort rts
  where
    ps = map fst pes
    m = product $ map (uncurry (^)) pes
    rts = map (`mod` m) $ take 1000 $ sqrtsModFactorisation n pes'

sqrtsModFactorisationProperty3 :: AnySign Integer -> [(Prime Integer, Power Word)] -> Bool
sqrtsModFactorisationProperty3 (AnySign n) (take 10 . map unwrapPP -> pes'@(map (first unPrime) -> pes))
  = nubOrd ps /= sort ps || all (\rt -> rt >= 0 && rt < m) rts
  where
    ps = map fst pes
    m = product $ map (uncurry (^)) pes
    rts = take 1000 $ sqrtsModFactorisation n pes'

sqrtsModFactorisationSpecialCase1 :: Assertion
sqrtsModFactorisationSpecialCase1 =
  assertEqual "should be equal" [0]
    (sqrtsModFactorisation 0 $ map (first (fromJust . isPrime)) [(2 :: Integer, 1), (3, 1), (5, 1)])

sqrtsModFactorisationSpecialCase2 :: Assertion
sqrtsModFactorisationSpecialCase2 =
  assertEqual "should be equal" [0]
    (sqrtsModFactorisation 0 $ map (first (fromJust . isPrime)) [(3 :: Integer, 1), (5, 1)])

sqrtsModProperty1 :: AnySign Integer -> Positive Natural -> Bool
sqrtsModProperty1 (AnySign n) (Positive m) = case n `modulo` m of
  SomeMod x -> sort (sqrtsMod sfactors x) == filter (\rt -> rt * rt == x) [minBound .. maxBound]
  InfMod{} -> True

testSuite :: TestTree
testSuite = testGroup "Sqrt"
  [ testGroup "sqrtsModPrime"
    [ testSmallAndQuick "matches jacobi"   sqrtsModPrimeProperty1
    , testSmallAndQuick "is residue"       sqrtsModPrimeProperty2
    , testSmallAndQuick "distinct"         sqrtsModPrimeProperty3
    , testSmallAndQuick "bounded"          sqrtsModPrimeProperty4
    ]
  , testGroup "tonelliShanks"
    [ testSmallAndQuick "generic"          tonelliShanksProperty1
    , testSmallAndQuick "smallest residue" tonelliShanksProperty2
    , testSmallAndQuick "-1"               tonelliShanksProperty3
    , testCase          "OEIS A002224"     tonelliShanksSpecialCases
    ]
  , testGroup "sqrtsModPrimePower"
    [ testSmallAndQuick "generic"        sqrtsModPrimePowerProperty1
    , testSmallAndQuick "_  2 _"         sqrtsModPrimePowerProperty2
    , testSmallAndQuick "distinct"       sqrtsModPrimePowerProperty3
    , testSmallAndQuick "_  2 _"         sqrtsModPrimePowerProperty4
    , testSmallAndQuick "bounded"        sqrtsModPrimePowerProperty5
    , testSmallAndQuick "_  2 _"         sqrtsModPrimePowerProperty6
    , testCase          "16 2 2"         sqrtsModPrimePowerSpecialCase1
    , testCase          "16 3 2"         sqrtsModPrimePowerSpecialCase2
    , testCase          "0  3 2"         sqrtsModPrimePowerSpecialCase3
    , testCase          "0  3 3"         sqrtsModPrimePowerSpecialCase4
    , testCase          "0  2 4"         sqrtsModPrimePowerSpecialCase5
    , testCase          "9  3 3"         sqrtsModPrimePowerSpecialCase6
    , testCase          "4  2 3"         sqrtsModPrimePowerSpecialCase7
    , testCase          "1  2 2"         sqrtsModPrimePowerSpecialCase8
    , testCase          "-1 2 2"         sqrtsModPrimePowerSpecialCase9
    , testCase          "4  2 4"         sqrtsModPrimePowerSpecialCase10
    , testCase          "16 2 6"         sqrtsModPrimePowerSpecialCase11
    ]
  , testGroup "sqrtsModFactorisation"
    [ testSmallAndQuick "generic"                 sqrtsModFactorisationProperty1
    , testSmallAndQuick "distinct"                sqrtsModFactorisationProperty2
    , testSmallAndQuick "bounded"                 sqrtsModFactorisationProperty3
    , testCase          "0 [(2,1), (3,1), (5,1)]" sqrtsModFactorisationSpecialCase1
    , testCase          "0 [(3,1), (5,1)]"        sqrtsModFactorisationSpecialCase2
    ]
  , testGroup "sqrtsMod"
    [ testSmallAndQuick "generic" sqrtsModProperty1
    ]
  ]
