    -- |
    -- 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 (Integer, Num, (*))
    import Data.Bool (Bool)
    import Data.Coerce (coerce)
    import Data.Euclidean (Euclidean, quot, divide)
    import Data.Function (($), (.))
    import Data.Int (Int)
    import Data.List (takeWhile)
    import Data.Map (Map)
    import qualified Data.Map.Strict as M
    import Data.Maybe (isJust)
    import Data.Ord (Ord, (<=))
    import Data.Semiring (Semiring(..))
    import Data.Tuple (fst)
    import Data.Word (Word)
    import Numeric.Natural (Natural)
    import Text.Show (Show)

    -- 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 { forall a b. 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
    $cshowsPrec :: forall a b. (Show a, Show b) => Int -> DirichletSeries a b -> ShowS
    showsPrec :: Int -> DirichletSeries a b -> ShowS
    $cshow :: forall a b. (Show a, Show b) => DirichletSeries a b -> String
    show :: DirichletSeries a b -> String
    $cshowList :: forall a b. (Show a, Show b) => [DirichletSeries a b] -> ShowS
    showList :: [DirichletSeries a b] -> ShowS
    Show)

    fromDistinctAscList :: forall a b. [(a, b)] -> DirichletSeries a b
    fromDistinctAscList :: forall a b. [(a, b)] -> DirichletSeries a b
    fromDistinctAscList = ([(a, b)] -> Map a b) -> [(a, b)] -> DirichletSeries a b
    forall a b. Coercible a b => a -> b
    coerce (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 :: forall a b.
    (Ord a, Num a, Semiring b) =>
    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 :: forall a b.
    (a -> Bool) -> DirichletSeries a b -> DirichletSeries a b
    filter a -> Bool
    predicate = (Map a b -> Map a b) -> DirichletSeries a b -> DirichletSeries a b
    forall a b. Coercible a b => a -> b
    coerce (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 :: forall a b.
    (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)
    forall a b. Coercible a b => a -> b
    coerce (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 :: forall a b.
    (Ord a, Semiring b) =>
    [DirichletSeries a b] -> DirichletSeries a b
    unions = ([Map a b] -> Map a b)
    -> [DirichletSeries a b] -> DirichletSeries a b
    forall a b. Coercible a b => 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 :: forall a b.
    (Ord a, Semiring b) =>
    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
    forall a b. Coercible a b => a -> b
    coerce (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 :: forall a b. DirichletSeries a b -> Int
    size = (Map a b -> Int) -> DirichletSeries a b -> Int
    forall a b. Coercible a b => a -> b
    coerce (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 :: forall a b.
    (Num a, Euclidean a, Ord a, Semiring b) =>
    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 #-}