-- |
-- Module:      Math.NumberTheory.Primes.Factorisation.Utils
-- Copyright:   (c) 2011 Daniel Fischer
-- Licence:     MIT
-- Stability:   Provisional
-- Portability: Non-portable (GHC extensions)
--
-- Some utilities related to factorisation, defined here to avoid import cycles.
{-# LANGUAGE BangPatterns #-}
module Math.NumberTheory.Primes.Factorisation.Utils
( ppTotient
, totientFromCanonical
, carmichaelFromCanonical
, divisorsFromCanonical
, tauFromCanonical
, divisorSumFromCanonical
, sigmaFromCanonical
) where

import Data.Set (Set)
import qualified Data.Set as Set
import Data.Bits
import Data.List

import Math.NumberTheory.Powers.Integer

-- | Totient of a prime power.
ppTotient :: (Integer,Int) -> Integer
ppTotient (p,1) = p-1
ppTotient (p,k) = (p-1)*(integerPower p (k-1))  -- slightly faster than (^) usually

-- | Calculate the totient from the canonical factorisation.
totientFromCanonical :: [(Integer,Int)] -> Integer
totientFromCanonical = product . map ppTotient

-- | Calculate the Carmichael function from the factorisation.
--   Requires that the list of prime factors is strictly ascending.
carmichaelFromCanonical :: [(Integer,Int)] -> Integer
carmichaelFromCanonical = go2
where
go2 ((2,k):ps) = let acc = case k of
1 -> 1
2 -> 2
_ -> 1 `shiftL` (k-2)
in go acc ps
go2 ps = go 1 ps
go !acc ((p,1):pps) = go (lcm acc (p-1)) pps
go acc ((p,k):pps)  = go ((lcm acc (p-1))*integerPower p (k-1)) pps
go acc []           = acc

-- | The set of divisors, efficiently calculated from the canonical factorisation.
divisorsFromCanonical :: [(Integer,Int)] -> Set Integer
divisorsFromCanonical = foldl' step (Set.singleton 1)
where
step st (p,k) = Set.unions (st:[Set.mapMonotonic (*pp) st | pp <- take k (iterate (*p) p) ])

-- | The number of divisors, efficiently calculated from the canonical factorisation.
tauFromCanonical :: [(a,Int)] -> Integer
tauFromCanonical pps = product [fromIntegral k + 1 | (_,k) <- pps]

-- | The sum of all divisors, efficiently calculated from the canonical factorisation.
divisorSumFromCanonical :: [(Integer,Int)] -> Integer
divisorSumFromCanonical = product . map ppDivSum

ppDivSum :: (Integer,Int) -> Integer
ppDivSum (p,1) = p+1
ppDivSum (p,k) = (p^(k+1)-1) `quot` (p-1)

-- | The sum of the powers (with fixed exponent) of all divisors,
--   efficiently calculated from the canonical factorisation.
sigmaFromCanonical :: Int -> [(Integer,Int)] -> Integer
sigmaFromCanonical k = product . map (ppDivPowerSum k)

ppDivPowerSum :: Int -> (Integer,Int) -> Integer
ppDivPowerSum k (p,m) = (p^(k*(m+1)) - 1) `quot` (p^k - 1)