{-# LANGUAGE CPP #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE ViewPatterns #-}
module Math.NumberTheory.DirichletCharacters
(
OrZero, pattern Zero, pattern NonZero
, orZeroToNum
, DirichletCharacter
, indexToChar
, indicesToChars
, characterNumber
, allChars
, fromTable
, eval
, evalGeneral
, evalAll
, principalChar
, isPrincipal
, orderChar
, RealCharacter
, isRealCharacter
, getRealChar
, toRealFunction
, jacobiCharacter
, PrimitiveCharacter
, isPrimitive
, getPrimitiveChar
, induced
, makePrimitive
, WithNat(..)
, RootOfUnity(..)
, toRootOfUnity
, toComplex
, validChar
) where
#if !MIN_VERSION_base(4,12,0)
import Control.Applicative (liftA2)
#endif
import Data.Bits (Bits(..))
import Data.Foldable (for_)
import Data.Functor.Identity (Identity(..))
import Data.Kind
import Data.List (mapAccumL, foldl', sort, find, unfoldr)
import Data.Maybe (mapMaybe, fromJust, fromMaybe)
#if MIN_VERSION_base(4,12,0)
import Data.Monoid (Ap(..))
#endif
import Data.Proxy (Proxy(..))
import Data.Ratio ((%), numerator, denominator)
import Data.Semigroup (Semigroup(..),Product(..))
import qualified Data.Vector as V
import qualified Data.Vector.Mutable as MV
import Data.Vector (Vector, (!))
import GHC.TypeNats (Nat, SomeNat(..), natVal, someNatVal)
import Numeric.Natural (Natural)
import Math.NumberTheory.ArithmeticFunctions (totient)
import Math.NumberTheory.Moduli.Chinese
import Math.NumberTheory.Moduli.Class (KnownNat, Mod, getVal)
import Math.NumberTheory.Moduli.Internal (isPrimitiveRoot', discreteLogarithmPP)
import Math.NumberTheory.Moduli.Multiplicative (MultMod(..), isMultElement)
import Math.NumberTheory.Moduli.Singleton (Some(..), cyclicGroupFromFactors)
import Math.NumberTheory.Powers.Modular (powMod)
import Math.NumberTheory.Primes (Prime(..), UniqueFactorisation, factorise, nextPrime)
import Math.NumberTheory.RootsOfUnity
import Math.NumberTheory.Utils.FromIntegral (wordToInt)
import Math.NumberTheory.Utils
newtype DirichletCharacter (n :: Nat) = Generated [DirichletFactor]
data DirichletFactor = OddPrime { _getPrime :: Prime Natural
, _getPower :: Word
, _getGenerator :: Natural
, _getValue :: RootOfUnity
}
| TwoPower { _getPower2 :: Int
, _getFirstValue :: RootOfUnity
, _getSecondValue :: RootOfUnity
}
| Two
instance Eq (DirichletCharacter n) where
Generated a == Generated b = a == b
instance Eq DirichletFactor where
TwoPower _ x1 x2 == TwoPower _ y1 y2 = x1 == y1 && x2 == y2
OddPrime _ _ _ x == OddPrime _ _ _ y = x == y
Two == Two = True
_ == _ = False
generator :: (Integral a, UniqueFactorisation a) => Prime a -> Word -> a
generator p k
| k == 1 = modP
| otherwise = if powMod modP (p'-1) (p'*p') == 1 then modP + p' else modP
where p' = unPrime p
modP = case cyclicGroupFromFactors [(p,k)] of
Just (Some cg) -> head $ filter (isPrimitiveRoot' cg) [2..p'-1]
_ -> error "illegal"
lambda :: Integer -> Int -> Integer
lambda x e = ((powMod x (2*modulus) largeMod - 1) `shiftR` (e+1)) .&. (modulus - 1)
where modulus = bit (e-2)
largeMod = bit (2*e - 1)
eval :: DirichletCharacter n -> MultMod n -> RootOfUnity
eval (Generated ds) m = foldMap (evalFactor m') ds
where m' = getVal $ multElement m
evalFactor :: Integer -> DirichletFactor -> RootOfUnity
evalFactor m =
\case
OddPrime (toInteger . unPrime -> p) k (toInteger -> a) b ->
discreteLogarithmPP p k a (m `rem` p^k) `stimes` b
TwoPower k s b -> (if testBit m 1 then s else mempty)
<> lambda (thingy k m) k `stimes` b
Two -> mempty
thingy :: (Bits p, Num p) => Int -> p -> p
thingy k m = if testBit m 1
then bit k - m'
else m'
where m' = m .&. (bit k - 1)
evalGeneral :: KnownNat n => DirichletCharacter n -> Mod n -> OrZero RootOfUnity
evalGeneral chi t = case isMultElement t of
Nothing -> Zero
Just x -> NonZero $ eval chi x
principalChar :: KnownNat n => DirichletCharacter n
principalChar = minBound
mulChars :: DirichletCharacter n -> DirichletCharacter n -> DirichletCharacter n
mulChars (Generated x) (Generated y) = Generated (zipWith combine x y)
where combine :: DirichletFactor -> DirichletFactor -> DirichletFactor
combine Two Two = Two
combine (OddPrime p k g n) (OddPrime _ _ _ m) =
OddPrime p k g (n <> m)
combine (TwoPower k a n) (TwoPower _ b m) =
TwoPower k (a <> b) (n <> m)
combine _ _ = error "internal error: malformed DirichletCharacter"
instance Semigroup (DirichletCharacter n) where
(<>) = mulChars
stimes = stimesChar
instance KnownNat n => Monoid (DirichletCharacter n) where
mempty = principalChar
mappend = (<>)
stimesChar :: Integral a => a -> DirichletCharacter n -> DirichletCharacter n
stimesChar s (Generated xs) = Generated (map mult xs)
where mult :: DirichletFactor -> DirichletFactor
mult (OddPrime p k g n) = OddPrime p k g (s `stimes` n)
mult (TwoPower k a b) = TwoPower k (s `stimes` a) (s `stimes` b)
mult Two = Two
instance KnownNat n => Enum (DirichletCharacter n) where
toEnum = indexToChar . fromIntegral
fromEnum = fromIntegral . characterNumber
succ x = makeChar x (characterNumber x + 1)
pred x = makeChar x (characterNumber x - 1)
enumFromTo x y = bulkMakeChars x [fromEnum x..fromEnum y]
enumFrom x = bulkMakeChars x [fromEnum x..]
enumFromThenTo x y z = bulkMakeChars x [fromEnum x, fromEnum y..fromEnum z]
enumFromThen x y = bulkMakeChars x [fromEnum x, fromEnum y..]
instance KnownNat n => Bounded (DirichletCharacter n) where
minBound = indexToChar 0
maxBound = indexToChar (totient n - 1)
where n = natVal (Proxy :: Proxy n)
characterNumber :: DirichletCharacter n -> Integer
characterNumber (Generated y) = foldl' go 0 y
where go x (OddPrime p k _ a) = x * m + numerator (fromRootOfUnity a * fromIntegral m)
where p' = fromIntegral (unPrime p)
m = p'^(k-1)*(p'-1)
go x (TwoPower k a b) = x' * 2 + numerator (fromRootOfUnity a * 2)
where m = bit (k-2) :: Integer
x' = x `shiftL` (k-2) + numerator (fromRootOfUnity b * fromIntegral m)
go x Two = x
indexToChar :: forall n. KnownNat n => Natural -> DirichletCharacter n
indexToChar = runIdentity . indicesToChars . Identity
indicesToChars :: forall n f. (KnownNat n, Functor f) => f Natural -> f (DirichletCharacter n)
indicesToChars = fmap (Generated . unroll t . (`mod` m))
where n = natVal (Proxy :: Proxy n)
(Product m, t) = mkTemplate n
allChars :: forall n. KnownNat n => [DirichletCharacter n]
allChars = indicesToChars [0..m-1]
where m = totient $ natVal (Proxy :: Proxy n)
makeChar :: Integral a => DirichletCharacter n -> a -> DirichletCharacter n
makeChar x = runIdentity . bulkMakeChars x . Identity
bulkMakeChars :: (Integral a, Functor f) => DirichletCharacter n -> f a -> f (DirichletCharacter n)
bulkMakeChars x = fmap (Generated . unroll t . (`mod` m) . fromIntegral)
where (Product m, t) = templateFromCharacter x
data Template = OddTemplate { _getPrime' :: Prime Natural
, _getPower' :: Word
, _getGenerator' :: !Natural
, _getModulus' :: !Natural
}
| TwoPTemplate { _getPower2' :: Int
, _getModulus' :: !Natural
}
| TwoTemplate
templateFromCharacter :: DirichletCharacter n -> (Product Natural, [Template])
templateFromCharacter (Generated t) = traverse go t
where go (OddPrime p k g _) = (Product m, OddTemplate p k g m)
where p' = unPrime p
m = p'^(k-1)*(p'-1)
go (TwoPower k _ _) = (Product (2*m), TwoPTemplate k m)
where m = bit (k-2)
go Two = (Product 1, TwoTemplate)
mkTemplate :: Natural -> (Product Natural, [Template])
mkTemplate = go . sort . factorise
where go :: [(Prime Natural, Word)] -> (Product Natural, [Template])
go ((unPrime -> 2, 1): xs) = (Product 1, [TwoTemplate]) <> traverse odds xs
go ((unPrime -> 2, wordToInt -> k): xs) = (Product (2*m), [TwoPTemplate k m]) <> traverse odds xs
where m = bit (k-2)
go xs = traverse odds xs
odds :: (Prime Natural, Word) -> (Product Natural, Template)
odds (p, k) = (Product m, OddTemplate p k (generator p k) m)
where p' = unPrime p
m = p'^(k-1)*(p'-1)
unroll :: [Template] -> Natural -> [DirichletFactor]
unroll t m = snd (mapAccumL func m t)
where func :: Natural -> Template -> (Natural, DirichletFactor)
func a (OddTemplate p k g n) = (a1, OddPrime p k g (toRootOfUnity $ (toInteger a2) % (toInteger n)))
where (a1,a2) = quotRem a n
func a (TwoPTemplate k n) = (b1, TwoPower k (toRootOfUnity $ (toInteger a2) % 2) (toRootOfUnity $ (toInteger b2) % (toInteger n)))
where (a1,a2) = quotRem a 2
(b1,b2) = quotRem a1 n
func a TwoTemplate = (a, Two)
isPrincipal :: DirichletCharacter n -> Bool
isPrincipal chi = characterNumber chi == 0
induced :: forall n d. (KnownNat d, KnownNat n) => DirichletCharacter d -> Maybe (DirichletCharacter n)
induced (Generated start) = if n `rem` d == 0
then Just (Generated (combine (snd $ mkTemplate n) start))
else Nothing
where n = natVal (Proxy :: Proxy n)
d = natVal (Proxy :: Proxy d)
combine :: [Template] -> [DirichletFactor] -> [DirichletFactor]
combine [] _ = []
combine ts [] = map newFactor ts
combine (t:xs) (y:ys) = case (t,y) of
(TwoTemplate, Two) -> Two: combine xs ys
(TwoTemplate, _) -> Two: combine xs (y:ys)
(TwoPTemplate k _, Two) -> TwoPower k mempty mempty: combine xs ys
(TwoPTemplate k _, TwoPower _ a b) -> TwoPower k a b: combine xs ys
(TwoPTemplate k _, _) -> TwoPower k mempty mempty: combine xs (y:ys)
(OddTemplate p k _ _, OddPrime q _ g a) | p == q -> OddPrime p k g a: combine xs ys
(OddTemplate p k g _, OddPrime q _ _ _) | p < q -> OddPrime p k g mempty: combine xs (y:ys)
_ -> error "internal error in induced: please report this as a bug"
newFactor :: Template -> DirichletFactor
newFactor TwoTemplate = Two
newFactor (TwoPTemplate k _) = TwoPower k mempty mempty
newFactor (OddTemplate p k g _) = OddPrime p k g mempty
jacobiCharacter :: forall n. KnownNat n => Maybe (RealCharacter n)
jacobiCharacter = if odd n
then Just $ RealChar $ Generated $ map go $ snd $ mkTemplate n
else Nothing
where n = natVal (Proxy :: Proxy n)
go :: Template -> DirichletFactor
go (OddTemplate p k g _) = OddPrime p k g $ toRootOfUnity ((toInteger k) % 2)
go _ = error "internal error in jacobiCharacter: please report this as a bug"
newtype RealCharacter n = RealChar {
getRealChar :: DirichletCharacter n
}
deriving Eq
isRealCharacter :: DirichletCharacter n -> Maybe (RealCharacter n)
isRealCharacter t@(Generated xs) = if all real xs then Just (RealChar t) else Nothing
where real :: DirichletFactor -> Bool
real (OddPrime _ _ _ a) = a <> a == mempty
real (TwoPower _ _ b) = b <> b == mempty
real Two = True
toRealFunction :: KnownNat n => RealCharacter n -> Mod n -> Int
toRealFunction (RealChar chi) m = case evalGeneral chi m of
Zero -> 0
NonZero t | t == mempty -> 1
NonZero t | t == RootOfUnity (1 % 2) -> -1
_ -> error "internal error in toRealFunction: please report this as a bug"
validChar :: forall n. KnownNat n => DirichletCharacter n -> Bool
validChar (Generated xs) = correctDecomposition && all correctPrimitiveRoot xs && all validValued xs
where correctDecomposition = sort (factorise n) == map getPP xs
getPP (TwoPower k _ _) = (two, fromIntegral k)
getPP (OddPrime p k _ _) = (p, k)
getPP Two = (two,1)
correctPrimitiveRoot (OddPrime p k g _) = g == generator p k
correctPrimitiveRoot _ = True
validValued (TwoPower k a b) = a <> a == mempty && (bit (k-2) :: Integer) `stimes` b == mempty
validValued (OddPrime (unPrime -> p) k _ a) = (p^(k-1)*(p-1)) `stimes` a == mempty
validValued Two = True
n = natVal (Proxy :: Proxy n)
two = nextPrime 2
orderChar :: DirichletCharacter n -> Integer
orderChar (Generated xs) = foldl' lcm 1 $ map orderFactor xs
where orderFactor (TwoPower _ (RootOfUnity a) (RootOfUnity b)) = denominator a `lcm` denominator b
orderFactor (OddPrime _ _ _ (RootOfUnity a)) = denominator a
orderFactor Two = 1
isPrimitive :: DirichletCharacter n -> Maybe (PrimitiveCharacter n)
isPrimitive t@(Generated xs) = if all primitive xs then Just (PrimitiveCharacter t) else Nothing
where primitive :: DirichletFactor -> Bool
primitive Two = False
primitive (OddPrime _ 1 _ a) = a /= mempty
primitive (OddPrime (unPrime -> p) k _ a) = (p^(k-2)*(p-1)) `stimes` a /= mempty
primitive (TwoPower 2 a _) = a /= mempty
primitive (TwoPower k _ b) = (bit (k-3) :: Integer) `stimes` b /= mempty
newtype PrimitiveCharacter n = PrimitiveCharacter {
getPrimitiveChar :: DirichletCharacter n
}
deriving Eq
data WithNat (a :: Nat -> Type) where
WithNat :: KnownNat m => a m -> WithNat a
makePrimitive :: DirichletCharacter n -> WithNat PrimitiveCharacter
makePrimitive (Generated xs) =
case someNatVal (product mods) of
SomeNat (Proxy :: Proxy m) -> WithNat (PrimitiveCharacter (Generated ys) :: PrimitiveCharacter m)
where (mods,ys) = unzip (mapMaybe prim xs)
prim :: DirichletFactor -> Maybe (Natural, DirichletFactor)
prim Two = Nothing
prim (OddPrime p' k g a) = case find works options of
Nothing -> error "invalid character"
Just (0,_) -> Nothing
Just (i,_) -> Just (p^i, OddPrime p' i g a)
where options = (0,1): [(i,p^(i-1)*(p-1)) | i <- [1..k]]
works (_,phi) = phi `stimes` a == mempty
p = unPrime p'
prim (TwoPower k a b) = case find worksb options of
Nothing -> error "invalid character"
Just (2,_) | a == mempty -> Nothing
Just (i,_) -> Just (bit i :: Natural, TwoPower i a b)
where options = [(i, bit (i-2) :: Natural) | i <- [2..k]]
worksb (_,phi) = phi `stimes` b == mempty
#if !MIN_VERSION_base(4,12,0)
newtype Ap f a = Ap { getAp :: f a }
deriving (Eq, Functor, Applicative, Monad)
instance (Applicative f, Semigroup a) => Semigroup (Ap f a) where
(<>) = liftA2 (<>)
instance (Applicative f, Semigroup a, Monoid a) => Monoid (Ap f a) where
mempty = pure mempty
mappend = (<>)
#endif
type OrZero a = Ap Maybe a
pattern Zero :: OrZero a
pattern Zero = Ap Nothing
pattern NonZero :: a -> OrZero a
pattern NonZero x = Ap (Just x)
{-# COMPLETE Zero, NonZero #-}
orZeroToNum :: Num a => (b -> a) -> OrZero b -> a
orZeroToNum _ Zero = 0
orZeroToNum f (NonZero x) = f x
evalAll :: forall n. KnownNat n => DirichletCharacter n -> Vector (OrZero RootOfUnity)
evalAll (Generated xs) = V.generate (fromIntegral n) func
where n = natVal (Proxy :: Proxy n)
vectors = map mkVector xs
func :: Int -> OrZero RootOfUnity
func m = foldMap go vectors
where go :: (Int, Vector (OrZero RootOfUnity)) -> OrZero RootOfUnity
go (modulus,v) = v ! (m `mod` modulus)
mkVector :: DirichletFactor -> (Int, Vector (OrZero RootOfUnity))
mkVector Two = (2, V.fromList [Zero, mempty])
mkVector (OddPrime p k (fromIntegral -> g) a) = (modulus, w)
where
p' = unPrime p
modulus = fromIntegral (p'^k) :: Int
w = V.create $ do
v <- MV.replicate modulus Zero
let powers = iterateMaybe go (1,mempty)
go (m,x) = if m' > 1
then Just (m', x<>a)
else Nothing
where m' = m*g `mod` modulus
for_ powers $ \(m,x) -> MV.unsafeWrite v m (NonZero x)
return v
mkVector (TwoPower k a b) = (modulus, w)
where
modulus = bit k
w = V.generate modulus f
f m
| even m = Zero
| otherwise = NonZero ((if testBit m 1 then a else mempty) <> lambda (toInteger m'') k `stimes` b)
where m'' = thingy k m
iterateMaybe :: (a -> Maybe a) -> a -> [a]
iterateMaybe f x = unfoldr (fmap (\t -> (t, f t))) (Just x)
fromTable :: forall n. KnownNat n => Vector (OrZero RootOfUnity) -> Maybe (DirichletCharacter n)
fromTable v = if length v == fromIntegral n
then Generated <$> traverse makeFactor tmpl >>= check
else Nothing
where n = natVal (Proxy :: Proxy n)
n' = fromIntegral n :: Integer
tmpl = snd (mkTemplate n)
check :: DirichletCharacter n -> Maybe (DirichletCharacter n)
check chi = if evalAll chi == v then Just chi else Nothing
makeFactor :: Template -> Maybe DirichletFactor
makeFactor TwoTemplate = Just Two
makeFactor (TwoPTemplate k _) = TwoPower k <$> getValue (-1,bit k) <*> getValue (exp4 k, bit k)
makeFactor (OddTemplate p k g _) = OddPrime p k g <$> getValue (toInteger g, toInteger (unPrime p)^k)
getValue :: (Integer,Integer) -> Maybe RootOfUnity
getValue (g,m) = getAp (v ! fromInteger (fromJust (chinese (g,m) (1,n' `quot` m)) `mod` n'))
exp4terms :: [Rational]
exp4terms = [4^k % product [1..k] | k <- [0..]]
exp4 :: Int -> Integer
exp4 n = (`mod` bit n) $ sum $ map (`mod` bit n) $ map (\q -> numerator q * fromMaybe (error "error in exp4") (recipMod (denominator q) (bit n))) $ take n $ exp4terms