{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE DataKinds #-}
{-# OPTIONS_GHC -fno-warn-type-defaults #-}

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

import Test.Tasty

import Data.Maybe
import Data.Mod
import Data.Proxy
import Data.Semigroup
import GHC.TypeNats (SomeNat(..), KnownNat, someNatVal)
import Numeric.Natural

import Math.NumberTheory.ArithmeticFunctions (totient)
import Math.NumberTheory.Moduli.Multiplicative
import Math.NumberTheory.Moduli.Singleton
import Math.NumberTheory.Primes
import Math.NumberTheory.TestUtils

nextPrimitiveRoot :: (KnownNat m, UniqueFactorisation a, Integral a) => CyclicGroup a m -> Mod m -> Maybe (PrimitiveRoot m)
nextPrimitiveRoot cg g = listToMaybe $ mapMaybe (isPrimitiveRoot cg) [g..g+100]

nextMultElement :: KnownNat m => Mod m -> Maybe (MultMod m)
nextMultElement g = listToMaybe $ mapMaybe isMultElement [g..g+100]

-- | Ensure 'discreteLogarithm' returns in the appropriate range.
discreteLogRange :: Positive Natural -> Integer -> Integer -> Bool
discreteLogRange (Positive m) a b =
  case someNatVal m of
    SomeNat (_ :: Proxy m) -> (/= Just False) $ do
      cg <- cyclicGroup :: Maybe (CyclicGroup Integer m)
      a' <- nextPrimitiveRoot cg (fromInteger a)
      b' <- nextMultElement (fromInteger b)
      return $ discreteLogarithm cg a' b' < totient m

-- | Check that 'discreteLogarithm' inverts exponentiation.
discreteLogarithmProperty :: Positive Natural -> Integer -> Integer -> Bool
discreteLogarithmProperty (Positive m) a b =
  case someNatVal m of
    SomeNat (_ :: Proxy m) -> (/= Just False) $ do
      cg <- cyclicGroup :: Maybe (CyclicGroup Integer m)
      a' <- nextPrimitiveRoot cg (fromInteger a)
      b' <- nextMultElement (fromInteger b)
      return $ discreteLogarithm cg a' b' `stimes` unPrimitiveRoot a' == b'

-- | Check that 'discreteLogarithm' inverts exponentiation in the other direction.
discreteLogarithmProperty' :: Positive Natural -> Integer -> Natural -> Bool
discreteLogarithmProperty' (Positive m) a k =
  case someNatVal m of
    SomeNat (_ :: Proxy m) -> (/= Just False) $ do
      cg <- cyclicGroup :: Maybe (CyclicGroup Integer m)
      a'' <- nextPrimitiveRoot cg (fromInteger a)
      let a' = unPrimitiveRoot a''
      return $ discreteLogarithm cg a'' (k `stimes` a') == k `mod` totient m

testSuite :: TestTree
testSuite = testGroup "Discrete logarithm"
  [ testSmallAndQuick "output is correct range" discreteLogRange
  , testSmallAndQuick "a^(log_a b) == b"        discreteLogarithmProperty
  , testSmallAndQuick "log_a a^k == k"          discreteLogarithmProperty'
  ]