-- |
-- Module:      Math.NumberTheory.Utils.DirichletSeries
-- Copyright:   (c) 2018 Andrew Lelechenko
-- Licence:     MIT
-- Maintainer:  Andrew Lelechenko <andrew.lelechenko@gmail.com>
--
-- An abstract representation of a Dirichlet series over a semiring.
--

{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE TypeApplications      #-}

module Math.NumberTheory.Utils.DirichletSeries
  ( DirichletSeries
  , fromDistinctAscList
  , lookup
  , filter
  , partition
  , unions
  , union
  , size
  , timesAndCrop
  ) where

import Prelude hiding (filter, last, rem, quot, snd, lookup)
import Data.Coerce
import Data.Euclidean
import Data.Map (Map)
import qualified Data.Map.Strict as M
import Data.Maybe
import Data.Semiring (Semiring(..))
import Numeric.Natural

-- Sparse Dirichlet series are represented by an ascending list of pairs.
-- For instance, [(a, b), (c, d)] stands for 1 + b/a^s + d/c^s.
-- Note that the representation still may include a term (1, b), so
-- [(1, b), (c, d)] means (1 + b) + d/c^s.
newtype DirichletSeries a b = DirichletSeries { DirichletSeries a b -> Map a b
_unDirichletSeries :: Map a b }
  deriving (Int -> DirichletSeries a b -> ShowS
[DirichletSeries a b] -> ShowS
DirichletSeries a b -> String
(Int -> DirichletSeries a b -> ShowS)
-> (DirichletSeries a b -> String)
-> ([DirichletSeries a b] -> ShowS)
-> Show (DirichletSeries a b)
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
forall a b. (Show a, Show b) => Int -> DirichletSeries a b -> ShowS
forall a b. (Show a, Show b) => [DirichletSeries a b] -> ShowS
forall a b. (Show a, Show b) => DirichletSeries a b -> String
showList :: [DirichletSeries a b] -> ShowS
$cshowList :: forall a b. (Show a, Show b) => [DirichletSeries a b] -> ShowS
show :: DirichletSeries a b -> String
$cshow :: forall a b. (Show a, Show b) => DirichletSeries a b -> String
showsPrec :: Int -> DirichletSeries a b -> ShowS
$cshowsPrec :: forall a b. (Show a, Show b) => Int -> DirichletSeries a b -> ShowS
Show)

fromDistinctAscList :: forall a b. [(a, b)] -> DirichletSeries a b
fromDistinctAscList :: [(a, b)] -> DirichletSeries a b
fromDistinctAscList = ([(a, b)] -> Map a b) -> [(a, b)] -> DirichletSeries a b
coerce ([(a, b)] -> Map a b
forall k a. [(k, a)] -> Map k a
M.fromDistinctAscList @a @b)

lookup :: (Ord a, Num a, Semiring b) => a -> DirichletSeries a b -> b
lookup :: a -> DirichletSeries a b -> b
lookup a
1 (DirichletSeries Map a b
m) = b -> a -> Map a b -> b
forall k a. Ord k => a -> k -> Map k a -> a
M.findWithDefault b
forall a. Semiring a => a
zero a
1 Map a b
m b -> b -> b
forall a. Semiring a => a -> a -> a
`plus` b
forall a. Semiring a => a
one
lookup a
a (DirichletSeries Map a b
m) = b -> a -> Map a b -> b
forall k a. Ord k => a -> k -> Map k a -> a
M.findWithDefault b
forall a. Semiring a => a
zero a
a Map a b
m

filter :: forall a b. (a -> Bool) -> DirichletSeries a b -> DirichletSeries a b
filter :: (a -> Bool) -> DirichletSeries a b -> DirichletSeries a b
filter a -> Bool
predicate = (Map a b -> Map a b) -> DirichletSeries a b -> DirichletSeries a b
coerce ((a -> b -> Bool) -> Map a b -> Map a b
forall k a. (k -> a -> Bool) -> Map k a -> Map k a
M.filterWithKey @a @b (\a
k b
_ -> a -> Bool
predicate a
k))

partition :: forall a b. (a -> Bool) -> DirichletSeries a b -> (DirichletSeries a b, DirichletSeries a b)
partition :: (a -> Bool)
-> DirichletSeries a b
-> (DirichletSeries a b, DirichletSeries a b)
partition a -> Bool
predicate = (Map a b -> (Map a b, Map a b))
-> DirichletSeries a b
-> (DirichletSeries a b, DirichletSeries a b)
coerce ((a -> b -> Bool) -> Map a b -> (Map a b, Map a b)
forall k a. (k -> a -> Bool) -> Map k a -> (Map k a, Map k a)
M.partitionWithKey @a @b (\a
k b
_ -> a -> Bool
predicate a
k))

unions :: forall a b. (Ord a, Semiring b) => [DirichletSeries a b] -> DirichletSeries a b
unions :: [DirichletSeries a b] -> DirichletSeries a b
unions = ([Map a b] -> Map a b)
-> [DirichletSeries a b] -> DirichletSeries a b
coerce ((b -> b -> b) -> [Map a b] -> Map a b
forall (f :: * -> *) k a.
(Foldable f, Ord k) =>
(a -> a -> a) -> f (Map k a) -> Map k a
M.unionsWith b -> b -> b
forall a. Semiring a => a -> a -> a
plus :: [Map a b] -> Map a b)

union :: forall a b. (Ord a, Semiring b) => DirichletSeries a b -> DirichletSeries a b -> DirichletSeries a b
union :: DirichletSeries a b -> DirichletSeries a b -> DirichletSeries a b
union = (Map a b -> Map a b -> Map a b)
-> DirichletSeries a b
-> DirichletSeries a b
-> DirichletSeries a b
coerce ((b -> b -> b) -> Map a b -> Map a b -> Map a b
forall k a. Ord k => (a -> a -> a) -> Map k a -> Map k a -> Map k a
M.unionWith @a @b b -> b -> b
forall a. Semiring a => a -> a -> a
plus)

size :: forall a b. DirichletSeries a b -> Int
size :: DirichletSeries a b -> Int
size = (Map a b -> Int) -> DirichletSeries a b -> Int
coerce (Map a b -> Int
forall k a. Map k a -> Int
M.size @a @b)

-- | Let as = sum_i k_i/a_i^s and bs = sum_i l_i/b_i^s be Dirichlet series,
-- and all a_i and b_i are divisors of n. Return Dirichlet series cs,
-- which contains all terms as * bs = sum_i m_i/c_i^s such that c_i divides n.
timesAndCrop
  :: (Num a, Euclidean a, Ord a, Semiring b)
  => a
  -> DirichletSeries a b
  -> DirichletSeries a b
  -> DirichletSeries a b
timesAndCrop :: a
-> DirichletSeries a b
-> DirichletSeries a b
-> DirichletSeries a b
timesAndCrop a
n (DirichletSeries Map a b
as) (DirichletSeries Map a b
bs)
  = Map a b -> DirichletSeries a b
forall a b. Map a b -> DirichletSeries a b
DirichletSeries
  (Map a b -> DirichletSeries a b) -> Map a b -> DirichletSeries a b
forall a b. (a -> b) -> a -> b
$ (b -> b -> b) -> Map a b -> Map a b -> Map a b
forall k a. Ord k => (a -> a -> a) -> Map k a -> Map k a -> Map k a
M.unionWith b -> b -> b
forall a. Semiring a => a -> a -> a
plus ((b -> b -> b) -> Map a b -> Map a b -> Map a b
forall k a. Ord k => (a -> a -> a) -> Map k a -> Map k a -> Map k a
M.unionWith b -> b -> b
forall a. Semiring a => a -> a -> a
plus Map a b
as Map a b
bs)
  (Map a b -> Map a b) -> Map a b -> Map a b
forall a b. (a -> b) -> a -> b
$ (b -> b -> b) -> [(a, b)] -> Map a b
forall k a. Ord k => (a -> a -> a) -> [(k, a)] -> Map k a
M.fromListWith b -> b -> b
forall a. Semiring a => a -> a -> a
plus
  [ (a
a a -> a -> a
forall a. Num a => a -> a -> a
* a
b, b
fa b -> b -> b
forall a. Semiring a => a -> a -> a
`times` b
fb)
  | (a
b, b
fb) <- Map a b -> [(a, b)]
forall k a. Map k a -> [(k, a)]
M.assocs Map a b
bs
  , let nb :: a
nb = a
n a -> a -> a
forall a. Euclidean a => a -> a -> a
`quot` a
b
  , (a
a, b
fa) <- ((a, b) -> Bool) -> [(a, b)] -> [(a, b)]
forall a. (a -> Bool) -> [a] -> [a]
takeWhile ((a -> a -> Bool
forall a. Ord a => a -> a -> Bool
<= a
nb) (a -> Bool) -> ((a, b) -> a) -> (a, b) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a, b) -> a
forall a b. (a, b) -> a
fst) (Map a b -> [(a, b)]
forall k a. Map k a -> [(k, a)]
M.assocs Map a b
as)
  , Maybe a -> Bool
forall a. Maybe a -> Bool
isJust (a
nb a -> a -> Maybe a
forall a. GcdDomain a => a -> a -> Maybe a
`divide` a
a)
  ]
{-# SPECIALISE timesAndCrop :: Semiring b => Int -> DirichletSeries Int b -> DirichletSeries Int b -> DirichletSeries Int b #-}
{-# SPECIALISE timesAndCrop :: Semiring b => Word -> DirichletSeries Word b -> DirichletSeries Word b -> DirichletSeries Word b #-}
{-# SPECIALISE timesAndCrop :: Semiring b => Integer -> DirichletSeries Integer b -> DirichletSeries Integer b -> DirichletSeries Integer b #-}
{-# SPECIALISE timesAndCrop :: Semiring b => Natural -> DirichletSeries Natural b -> DirichletSeries Natural b -> DirichletSeries Natural b #-}