content: -- |
  -- Module:       Math.NumberTheory.DirichletCharactersTests
  -- Copyright:    (c) 2018 Bhavik Mehta
  -- License:      MIT
  -- Maintainer:   Andrew Lelechenko <andrew.lelechenko@gmail.com>
  --
  -- Tests for Math.NumberTheory.DirichletCharacters
  --
  
  {-# LANGUAGE GADTs #-}
  {-# LANGUAGE Rank2Types #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  {-# LANGUAGE TypeApplications #-}
  {-# LANGUAGE ViewPatterns #-}
  
  module Math.NumberTheory.DirichletCharactersTests where
  
  import Test.Tasty
  
  import Data.Complex
  import Data.List (genericLength)
  import Data.Maybe (isJust, mapMaybe)
  import Data.Proxy
  import Data.Semigroup
  import qualified Data.Vector as V
  import Numeric.Natural
  
  import GHC.TypeNats (SomeNat(..), someNatVal, KnownNat, natVal, sameNat)
  import Data.Type.Equality
  
  import Math.NumberTheory.ArithmeticFunctions (totient, divisorsList)
  import Math.NumberTheory.DirichletCharacters
  import qualified Math.NumberTheory.Moduli.Sqrt as J
  import Math.NumberTheory.Moduli.Class (SomeMod(..), modulo)
  import Math.NumberTheory.TestUtils (testSmallAndQuick, Positive(..))
  
  -- | This tests property 6 from https://en.wikipedia.org/wiki/Dirichlet_character#Axiomatic_definition
  dirCharOrder :: forall n. KnownNat n => DirichletCharacter n -> Bool
  dirCharOrder chi = isPrincipal (totient n `stimes` chi)
    where n = natVal @n Proxy
  
  -- | Tests wikipedia's property 3 (note 1,2,5 are essentially enforced by the type system).
  testMultiplicative :: KnownNat n => DirichletCharacter n -> Natural -> Natural -> Bool
  testMultiplicative chi (fromIntegral -> a) (fromIntegral -> b) = chiAB == chiAchiB
    where chi' = evalGeneral chi
          chiAB = chi' (a*b)
          chiAchiB = (<>) <$> chi' a <*> chi' b
  
  -- | Test property 4 from wikipedia
  testAtOne :: KnownNat n => DirichletCharacter n -> Bool
  testAtOne chi = eval chi mempty == mempty
  
  dirCharProperty :: (forall n. KnownNat n => DirichletCharacter n -> a) -> Positive Natural -> Natural -> a
  dirCharProperty test (Positive n) i =
    case someNatVal n of
      SomeNat (Proxy :: Proxy n) -> test chi
        where chi = indexToChar @n (i `mod` totient n)
  
  realCharProperty :: (forall n. KnownNat n => RealCharacter n -> a) -> Positive Natural -> Int -> a
  realCharProperty test (Positive n) i =
    case someNatVal n of
      SomeNat (Proxy :: Proxy n) -> test chi
        where chi = chars !! (i `mod` length chars)
              chars = mapMaybe isRealCharacter [principalChar @n .. maxBound]
  
  -- | There should be totient(n) characters
  countCharacters :: Positive Natural -> Bool
  countCharacters (Positive n) =
    case someNatVal n of
      SomeNat (Proxy :: Proxy n) ->
        genericLength (allChars @n) == totient n
  
  -- | The principal character should be 1 if gcd k n is 1 and 0 otherwise
  principalCase :: Positive Natural -> Positive Integer -> Bool
  principalCase (Positive n) (Positive k) =
    case k `modulo` n of
      SomeMod a -> evalGeneral chi a == if gcd k (fromIntegral n) > 1
                                           then Zero
                                           else mempty
        where chi = principalChar
      InfMod{} -> False
  
  -- | Test the orthogonality relations https://en.wikipedia.org/wiki/Dirichlet_character#Character_orthogonality
  orthogonality1 :: forall n. KnownNat n => DirichletCharacter n -> Bool
  orthogonality1 chi = magnitude (total - correct) < (1e-13 :: Double)
    where n = natVal @n Proxy
          total = sum [orZeroToNum toComplex (evalGeneral chi a) | a <- [0 .. maxBound]]
          correct = if isPrincipal chi
                       then fromIntegral $ totient n
                       else 0
  
  orthogonality2 :: Positive Natural -> Integer -> Bool
  orthogonality2 (Positive n) a =
    case a `modulo` n of
      SomeMod a' -> magnitude (total - correct) < (1e-13 :: Double)
        where total = sum [orZeroToNum toComplex (evalGeneral chi a') | chi <- allChars]
              correct = if a' == 1
                           then fromIntegral $ totient n
                           else 0
      InfMod {} -> False
  
  -- | Manually confirm isRealCharacter is correct (in both directions)
  realityCheck :: KnownNat n => DirichletCharacter n -> Bool
  realityCheck chi = isJust (isRealCharacter chi) == isReal'
    where isReal' = and [real (evalGeneral chi t) | t <- [minBound..maxBound]]
          real Zero = True
          real (NonZero t) = t <> t == mempty
  
  -- | Check real character evaluation matches normal evaluation
  realEvalCheck :: KnownNat n => RealCharacter n -> Int -> Bool
  realEvalCheck chi i' = fromIntegral (toRealFunction chi i) == (orZeroToNum toComplex (evalGeneral (getRealChar chi) i) :: Complex Double)
    where i = fromIntegral i'
  
  -- | The jacobi character agrees with the jacobi symbol
  jacobiCheck :: Positive Natural -> Bool
  jacobiCheck (Positive n) =
    case someNatVal (2*n+1) of
      SomeNat (Proxy :: Proxy n) ->
        case jacobiCharacter @n of
          Just chi -> and [toRealFunction chi (fromIntegral j) == J.symbolToNum (J.jacobi j (2*n+1)) | j <- [0..2*n]]
          _ -> False
  
  -- | Bulk evaluation agrees with pointwise evaluation
  evalAllCheck :: forall n. KnownNat n => DirichletCharacter n -> Bool
  evalAllCheck chi = V.generate (fromIntegral $ natVal @n Proxy) (evalGeneral chi . fromIntegral) == evalAll chi
  
  -- | Induced characters agree with the original character.
  -- (Except for when d=1, where chi(0) = 1, which is true for no other d)
  inducedCheck :: forall d. KnownNat d => DirichletCharacter d -> Positive Natural -> Bool
  inducedCheck chi (Positive k) =
    case someNatVal (d*k) of
      SomeNat (Proxy :: Proxy n) ->
        case induced @n chi of
          Just chi2 -> and (V.izipWith matchedValue (V.concat (replicate (fromIntegral k) (evalAll chi))) (evalAll chi2))
          Nothing -> False
    where d = natVal @d Proxy
          matchedValue i x1 x2 = if gcd (fromIntegral i) (d*k) > 1
                                    then x2 == Zero
                                    else x2 == x1
  
  -- | Primitive checker is correct (in both directions)
  primitiveCheck :: forall n. KnownNat n => DirichletCharacter n -> Bool
  primitiveCheck chi = isJust (isPrimitive chi) == isPrimitive'
    where isPrimitive' = all (\m -> m == n || testModulus m) (divisorsList n)
          n = fromIntegral (natVal @n Proxy) :: Int
          table = evalAll chi
          testModulus d = not $ null [a | a <- [1..n-1], gcd a n == 1, a `mod` d == 1 `mod` d, table V.! a /= mempty]
  
  -- | Ensure that makePrimitive gives primitive characters
  makePrimitiveCheck :: DirichletCharacter n -> Bool
  makePrimitiveCheck chi = case makePrimitive chi of
                              WithNat chi' -> isJust (isPrimitive (getPrimitiveChar chi'))
  
  -- | sameNat also ensures the two new moduli are the same
  makePrimitiveIdem :: DirichletCharacter n -> Bool
  makePrimitiveIdem chi = case makePrimitive chi of
                            WithNat (chi' :: PrimitiveCharacter n') ->
                              case makePrimitive (getPrimitiveChar chi') of
                                WithNat (chi'' :: PrimitiveCharacter n'') ->
                                  case sameNat (Proxy :: Proxy n') (Proxy :: Proxy n'') of
                                    Just Refl -> chi' == chi''
                                    Nothing -> False
  
  orderCheck :: DirichletCharacter n -> Bool
  orderCheck chi = isPrincipal (n `stimes` chi) && and [not (isPrincipal (i `stimes` chi)) | i <- [1..n-1]]
    where n = orderChar chi
  
  fromTableCheck :: forall n. KnownNat n => DirichletCharacter n -> Bool
  fromTableCheck chi = isJust (fromTable @n (evalAll chi))
  
  -- A bunch of functions making sure that every function which can produce a character (in
  -- particular by fiddling internal representation) produces a valid character
  indexToCharValid :: KnownNat n => DirichletCharacter n -> Bool
  indexToCharValid = validChar
  
  principalCharValid :: Positive Natural -> Bool
  principalCharValid (Positive n) =
    case someNatVal n of
      SomeNat (Proxy :: Proxy n) -> validChar (principalChar @n)
  
  mulCharsValid :: KnownNat n => DirichletCharacter n -> DirichletCharacter n -> Bool
  mulCharsValid chi1 chi2 = validChar (chi1 <> chi2)
  
  mulCharsValid' :: Positive Natural -> Natural -> Natural -> Bool
  mulCharsValid' (Positive n) i j =
    case someNatVal n of
      SomeNat (Proxy :: Proxy n) ->
        mulCharsValid (indexToChar @n (i `mod` totient n)) (indexToChar @n (j `mod` totient n))
  
  stimesCharValid :: KnownNat n => DirichletCharacter n -> Int -> Bool
  stimesCharValid chi n = validChar (n `stimes` chi)
  
  succValid :: KnownNat n => DirichletCharacter n -> Bool
  succValid = validChar . succ
  
  inducedValid :: forall d. KnownNat d => DirichletCharacter d -> Positive Natural -> Bool
  inducedValid chi (Positive k) = case someNatVal (natVal @d Proxy * k) of
    SomeNat (Proxy :: Proxy n) -> maybe False validChar (induced @n chi)
  
  jacobiValid :: Positive Natural -> Bool
  jacobiValid (Positive n) =
    case someNatVal (2*n+1) of
      SomeNat (Proxy :: Proxy n) ->
        case jacobiCharacter @n of
          Just chi -> validChar (getRealChar chi)
          _ -> False
  
  makePrimitiveValid :: DirichletCharacter n -> Bool
  makePrimitiveValid chi = case makePrimitive chi of
                              WithNat chi' -> validChar (getPrimitiveChar chi')
  
  testSuite :: TestTree
  testSuite = testGroup "DirichletCharacters"
    [ testSmallAndQuick "Dirichlet characters divide the right order" (dirCharProperty dirCharOrder)
    , testSmallAndQuick "Dirichlet characters are multiplicative" (dirCharProperty testMultiplicative)
    , testSmallAndQuick "Dirichlet characters are 1 at 1" (dirCharProperty testAtOne)
    , testSmallAndQuick "Right number of Dirichlet characters" countCharacters
    , testSmallAndQuick "Principal character behaves as expected" principalCase
    , testSmallAndQuick "Orthogonality relation 1" (dirCharProperty orthogonality1)
    , testSmallAndQuick "Orthogonality relation 2" orthogonality2
    , testSmallAndQuick "Real character checking is correct" (dirCharProperty realityCheck)
    , testSmallAndQuick "Real character evaluation is accurate" (realCharProperty realEvalCheck)
    , testSmallAndQuick "Jacobi character matches symbol" jacobiCheck
    , testSmallAndQuick "Bulk evaluation matches pointwise" (dirCharProperty evalAllCheck)
    , testSmallAndQuick "Induced character is correct" (dirCharProperty inducedCheck)
    , testSmallAndQuick "Primitive character checking is correct" (dirCharProperty primitiveCheck)
    , testSmallAndQuick "makePrimitive produces primitive character" (dirCharProperty makePrimitiveCheck)
    , testSmallAndQuick "makePrimitive is idempotent" (dirCharProperty makePrimitiveIdem)
    , testSmallAndQuick "Calculates correct order" (dirCharProperty orderCheck)
    , testSmallAndQuick "Can construct from table" (dirCharProperty fromTableCheck)
    , testGroup "Creates valid characters"
      [ testSmallAndQuick "indexToChar" (dirCharProperty indexToCharValid)
      , testSmallAndQuick "principalChar" principalCharValid
      , testSmallAndQuick "mulChars" mulCharsValid'
      , testSmallAndQuick "stimesChar" (dirCharProperty stimesCharValid)
      , testSmallAndQuick "succ" (dirCharProperty succValid)
      , testSmallAndQuick "induced" (dirCharProperty inducedValid)
      , testSmallAndQuick "jacobi" jacobiValid
      , testSmallAndQuick "makePrimitive" (dirCharProperty makePrimitiveValid)
      ]
    ]
  