-- | A representation of rational numbers as lists of prime powers,
--   allowing efficient representation, multiplication and division of
--   large numbers, especially of the sort occurring in combinatorial
--   computations.
--
--   The module also includes a method for generating factorials in
--   factored form directly, and for generating all divisors of
--   factored integers, or computing Euler's totient (phi) function
--   and the Möbius (mu) function.
module MathObj.FactoredRational
    ( -- * Type
      T

      -- * Utilities
    , factorial
    , eulerPhi
    , divisors
    , mu

    ) where

import qualified Algebra.Additive as Additive
import qualified Algebra.Ring as Ring
import qualified Algebra.Field as Field
import qualified Algebra.IntegralDomain as Integral

import qualified Algebra.ZeroTestable as ZeroTestable
import qualified Algebra.ToRational as ToRational
import qualified Algebra.RealIntegral as RealIntegral
import qualified Algebra.ToInteger as ToInteger

import Data.Numbers.Primes

import qualified Algebra.Absolute as Real
import NumericPrelude

-- Represent rational numbers by their prime factorizations.
-- Perhaps this should use a sparse representation instead, using a Map from
-- primes to powers?  Well, that should be easy enough to change later.

-- | The type of factored rationals.
--
--   Instances are provided for Eq, Ord, Additive, Ring, ZeroTestable,
--   Real, ToRational, Integral, RealIntegral, ToInteger, and Field.
--
--   Note that currently, addition is performed on factored rationals
--   by converting them to normal rationals, performing the addition,
--   and factoring.  This could probably be made more efficient by
--   finding a common denominator, pulling out common factors from the
--   numerators, and performing the addition and factoring only on the
--   relatively prime parts.
data T = FQZero            -- ^ zero
       | FQ Bool [Integer] -- ^ prime exponents with sign bit, True = negative.

-- XXX this ought to be improved.
instance Show T where
  show :: T -> String
show T
FQZero = String
"0"
  show (FQ Bool
True [Integer]
pows)  = String
"(-1)" String -> ShowS
forall a. [a] -> [a] -> [a]
++ [Integer] -> String
showPows [Integer]
pows
  show (FQ Bool
False [])   = String
"1"
  show (FQ Bool
False [Integer]
pows) = [Integer] -> String
showPows [Integer]
pows

showPows :: [Integer] -> String
showPows :: [Integer] -> String
showPows [Integer]
pows = [String] -> String
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([String] -> String) -> [String] -> String
forall a b. (a -> b) -> a -> b
$ (Integer -> Integer -> String)
-> [Integer] -> [Integer] -> [String]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Integer -> Integer -> String
showPow [Integer]
forall int. Integral int => [int]
primes [Integer]
pows
  where showPow :: Integer -> Integer -> String
        showPow :: Integer -> Integer -> String
showPow Integer
_ Integer
0 = String
""
        showPow Integer
p Integer
1 = String
"(" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Integer -> String
forall a. Show a => a -> String
show Integer
p String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
")"
        showPow Integer
p Integer
n = String
"(" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Integer -> String
forall a. Show a => a -> String
show Integer
p String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"^" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Integer -> String
forall a. Show a => a -> String
show Integer
n String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
")"

instance Additive.C T where
  zero :: T
zero = T
FQZero
  T
FQZero + :: T -> T -> T
+ T
a = T
a
  T
a + T
FQZero = T
a
  T
x + T
y = Rational -> T
forall a. C a => Rational -> a
fromRational' (T -> Rational
forall a. C a => a -> Rational
toRational T
x Rational -> Rational -> Rational
forall a. C a => a -> a -> a
+ T -> Rational
forall a. C a => a -> Rational
toRational T
y)

  negate :: T -> T
negate T
FQZero   = T
FQZero
  negate (FQ Bool
s [Integer]
e) = Bool -> [Integer] -> T
FQ (Bool -> Bool
not Bool
s) [Integer]
e

instance Ring.C T where
  T
FQZero * :: T -> T -> T
* T
_ = T
FQZero
  T
_ * T
FQZero = T
FQZero
  (FQ Bool
s1 [Integer]
e1) * (FQ Bool
s2 [Integer]
e2) = Bool -> [Integer] -> T
FQ (Bool
s1 Bool -> Bool -> Bool
forall a. Eq a => a -> a -> Bool
/= Bool
s2) (Integer
-> Integer
-> (Integer -> Integer -> Integer)
-> [Integer]
-> [Integer]
-> [Integer]
forall a b c. a -> b -> (a -> b -> c) -> [a] -> [b] -> [c]
zipWithExt Integer
0 Integer
0 Integer -> Integer -> Integer
forall a. C a => a -> a -> a
(+) [Integer]
e1 [Integer]
e2)

  fromInteger :: Integer -> T
fromInteger Integer
0 = T
FQZero
  fromInteger Integer
n | Integer
n Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
0     = Bool -> [Integer] -> T
FQ Bool
True (Integer -> [Integer]
factor (Integer -> Integer
forall a. C a => a -> a
negate Integer
n))
                | Bool
otherwise = Bool -> [Integer] -> T
FQ Bool
False (Integer -> [Integer]
factor Integer
n)

  T
_ ^ :: T -> Integer -> T
^ Integer
0 = T
forall a. C a => a
one
  T
FQZero   ^ Integer
_ = T
FQZero
  (FQ Bool
s [Integer]
e) ^ Integer
n = Bool -> [Integer] -> T
FQ Bool
s ((Integer -> Integer) -> [Integer] -> [Integer]
forall a b. (a -> b) -> [a] -> [b]
map (Integer -> Integer -> Integer
forall a. C a => a -> a -> a
*Integer
n) [Integer]
e)

-- | Zip two lists together with a combining function, using default
--   values to extend the lists if one is shorter than the other.
zipWithExt :: a -> b -> (a -> b -> c) -> [a] -> [b] -> [c]
zipWithExt :: a -> b -> (a -> b -> c) -> [a] -> [b] -> [c]
zipWithExt a
da b
db a -> b -> c
f = [a] -> [b] -> [c]
zipWithExt'
  where zipWithExt' :: [a] -> [b] -> [c]
zipWithExt' []     [b]
bs     = (a -> b -> c) -> [a] -> [b] -> [c]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith a -> b -> c
f (a -> [a]
forall a. a -> [a]
repeat a
da) [b]
bs
        zipWithExt' [a]
as     []     = (a -> b -> c) -> [a] -> [b] -> [c]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith a -> b -> c
f [a]
as (b -> [b]
forall a. a -> [a]
repeat b
db)
        zipWithExt' (a
a:[a]
as) (b
b:[b]
bs) = a -> b -> c
f a
a b
b c -> [c] -> [c]
forall a. a -> [a] -> [a]
: [a] -> [b] -> [c]
zipWithExt' [a]
as [b]
bs

-- | A simple factoring method.
--
--   We should probably just depend on another module with some
--   dedicated, efficient factoring code written by someone really
--   smart, but this simple method works OK for now.
--
--   Precondition: argument is positive.
factor :: Integer -> [Integer]
factor :: Integer -> [Integer]
factor Integer
m = Integer -> [Integer] -> [Integer]
factor' Integer
m [Integer]
forall int. Integral int => [int]
primes
  where
    factor' :: Integer -> [Integer] -> [Integer]
factor' Integer
1 [Integer]
_ = []
    factor' Integer
n (Integer
p:[Integer]
ps) = let (Integer
k,Integer
n') = Integer -> Integer -> (Integer, Integer)
logRem Integer
n Integer
p
                       in  Integer
k Integer -> [Integer] -> [Integer]
forall a. a -> [a] -> [a]
: Integer -> [Integer] -> [Integer]
factor' Integer
n' [Integer]
ps
    factor' Integer
_ [] = String -> [Integer]
forall a. HasCallStack => String -> a
error String
"Oh noes, Euclid was wrong!"

-- | @logRem n p@ computes (k,n'), where k is the highest power of p
--   that divides n, and n' = n `div` p^k.
logRem :: Integer -> Integer -> (Integer, Integer)
logRem :: Integer -> Integer -> (Integer, Integer)
logRem = Integer -> Integer -> Integer -> (Integer, Integer)
forall b a. (Eq b, C b, Num b, Num a, C a) => a -> b -> b -> (a, b)
logRem' Integer
0
  where logRem' :: a -> b -> b -> (a, b)
logRem' a
k b
n b
p | b
n b -> b -> b
forall a. C a => a -> a -> a
`mod` b
p b -> b -> Bool
forall a. Eq a => a -> a -> Bool
== b
0 = a -> b -> b -> (a, b)
logRem' (a
ka -> a -> a
forall a. C a => a -> a -> a
+a
1) (b
n b -> b -> b
forall a. C a => a -> a -> a
`div` b
p) b
p
                      | Bool
otherwise = (a
k,b
n)

instance ZeroTestable.C T where
  isZero :: T -> Bool
isZero T
FQZero = Bool
True
  isZero T
_      = Bool
False

instance Eq T where
  T
FQZero == :: T -> T -> Bool
== T
FQZero   = Bool
True
  T
FQZero == (FQ Bool
_ [Integer]
_) = Bool
False
  (FQ Bool
_ [Integer]
_) == T
FQZero = Bool
False
  (FQ Bool
s1 [Integer]
e1) == (FQ Bool
s2 [Integer]
e2) = Bool
s1 Bool -> Bool -> Bool
forall a. Eq a => a -> a -> Bool
== Bool
s2 Bool -> Bool -> Bool
&& [Integer] -> [Integer]
dropZeros [Integer]
e1 [Integer] -> [Integer] -> Bool
forall a. Eq a => a -> a -> Bool
== [Integer] -> [Integer]
dropZeros [Integer]
e2
    where dropZeros :: [Integer] -> [Integer]
dropZeros = [Integer] -> [Integer]
forall a. [a] -> [a]
reverse ([Integer] -> [Integer])
-> ([Integer] -> [Integer]) -> [Integer] -> [Integer]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Integer -> Bool) -> [Integer] -> [Integer]
forall a. (a -> Bool) -> [a] -> [a]
dropWhile (Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
==Integer
0) ([Integer] -> [Integer])
-> ([Integer] -> [Integer]) -> [Integer] -> [Integer]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Integer] -> [Integer]
forall a. [a] -> [a]
reverse

instance Ord T where
  compare :: T -> T -> Ordering
compare T
FQZero T
FQZero = Ordering
EQ
  compare T
FQZero (FQ Bool
False [Integer]
_) = Ordering
LT
  compare T
FQZero (FQ Bool
True  [Integer]
_) = Ordering
GT
  compare (FQ Bool
False [Integer]
_) T
FQZero = Ordering
GT
  compare (FQ Bool
True  [Integer]
_) T
FQZero = Ordering
LT
  compare (FQ Bool
False [Integer]
_) (FQ Bool
True  [Integer]
_)   = Ordering
GT
  compare (FQ Bool
True  [Integer]
_) (FQ Bool
False [Integer]
_)   = Ordering
LT
  compare T
fq1 T
fq2 = Rational -> Rational -> Ordering
forall a. Ord a => a -> a -> Ordering
compare (T -> Rational
forall a. C a => a -> Rational
toRational T
fq1) (T -> Rational
forall a. C a => a -> Rational
toRational T
fq2)

instance Real.C T where
  abs :: T -> T
abs T
FQZero   = T
FQZero
  abs (FQ Bool
_ [Integer]
e) = Bool -> [Integer] -> T
FQ Bool
False [Integer]
e
  signum :: T -> T
signum = T -> T
forall a. (C a, Ord a) => a -> a
Real.signumOrd

instance ToRational.C T where
  toRational :: T -> Rational
toRational T
FQZero   = Rational
0
  toRational (FQ Bool
s [Integer]
e) = (if Bool
s then Rational -> Rational
forall a. C a => a -> a
negate else Rational -> Rational
forall a. a -> a
id)
                      (Rational -> Rational)
-> ([Integer] -> Rational) -> [Integer] -> Rational
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Rational] -> Rational
forall a. C a => [a] -> a
product
                      ([Rational] -> Rational)
-> ([Integer] -> [Rational]) -> [Integer] -> Rational
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Rational -> Integer -> Rational)
-> [Rational] -> [Integer] -> [Rational]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Rational -> Integer -> Rational
forall a. C a => a -> Integer -> a
(^-) ((Integer -> Rational) -> [Integer] -> [Rational]
forall a b. (a -> b) -> [a] -> [b]
map (Integer -> Integer -> Rational
forall a. C a => a -> a -> T a
% Integer
1) [Integer]
forall int. Integral int => [int]
primes)
                      ([Integer] -> Rational) -> [Integer] -> Rational
forall a b. (a -> b) -> a -> b
$ [Integer]
e

instance Integral.C T where
  divMod :: T -> T -> (T, T)
divMod T
a T
b =
    if T -> Bool
forall a. C a => a -> Bool
isZero T
b
      then (T
forall a. HasCallStack => a
undefined,T
a)
      else (T
aT -> T -> T
forall a. C a => a -> a -> a
/T
b,T
forall a. C a => a
zero)

instance RealIntegral.C T
  -- default definition is fine

instance ToInteger.C T where
  toInteger :: T -> Integer
toInteger T
FQZero = Integer
0
  toInteger (FQ Bool
s [Integer]
e) | (Integer -> Bool) -> [Integer] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
<Integer
0) [Integer]
e = String -> Integer
forall a. HasCallStack => String -> a
error String
"non-integer in FactoredRational.toInteger"
                     | Bool
otherwise  = (if Bool
s then Integer -> Integer
forall a. C a => a -> a
negate else Integer -> Integer
forall a. a -> a
id)
                                  (Integer -> Integer)
-> ([Integer] -> Integer) -> [Integer] -> Integer
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Integer] -> Integer
forall a. C a => [a] -> a
product
                                  ([Integer] -> Integer)
-> ([Integer] -> [Integer]) -> [Integer] -> Integer
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Integer -> Integer -> Integer)
-> [Integer] -> [Integer] -> [Integer]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Integer -> Integer -> Integer
forall a. C a => a -> Integer -> a
(^) [Integer]
forall int. Integral int => [int]
primes
                                  ([Integer] -> Integer) -> [Integer] -> Integer
forall a b. (a -> b) -> a -> b
$ [Integer]
e

instance Field.C T where
  recip :: T -> T
recip T
FQZero = String -> T
forall a. HasCallStack => String -> a
error String
"division by zero"
  recip (FQ Bool
s [Integer]
e) = Bool -> [Integer] -> T
FQ Bool
s ((Integer -> Integer) -> [Integer] -> [Integer]
forall a b. (a -> b) -> [a] -> [b]
map Integer -> Integer
forall a. C a => a -> a
negate [Integer]
e)

-- | Efficiently compute n! directly as a factored rational.
factorial :: Integer -> T
factorial :: Integer -> T
factorial Integer
0 = T
forall a. C a => a
one
factorial Integer
1 = T
forall a. C a => a
one
factorial Integer
n = Bool -> [Integer] -> T
FQ Bool
False ((Integer -> Bool) -> [Integer] -> [Integer]
forall a. (a -> Bool) -> [a] -> [a]
takeWhile (Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
>Integer
0) ([Integer] -> [Integer])
-> ([Integer] -> [Integer]) -> [Integer] -> [Integer]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Integer -> Integer) -> [Integer] -> [Integer]
forall a b. (a -> b) -> [a] -> [b]
map (Integer -> Integer -> Integer
factorialFactors Integer
n) ([Integer] -> [Integer]) -> [Integer] -> [Integer]
forall a b. (a -> b) -> a -> b
$ [Integer]
forall int. Integral int => [int]
primes)

-- | @factorialFactors n p@ computes the power of prime p in the
--   factorization of n!.
factorialFactors :: Integer -> Integer -> Integer
factorialFactors :: Integer -> Integer -> Integer
factorialFactors Integer
n Integer
p = [Integer] -> Integer
forall a. C a => [a] -> a
sum
                     ([Integer] -> Integer)
-> ([Integer] -> [Integer]) -> [Integer] -> Integer
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Integer -> Bool) -> [Integer] -> [Integer]
forall a. (a -> Bool) -> [a] -> [a]
takeWhile (Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
>Integer
0)
                     ([Integer] -> [Integer])
-> ([Integer] -> [Integer]) -> [Integer] -> [Integer]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Integer -> Integer) -> [Integer] -> [Integer]
forall a b. (a -> b) -> [a] -> [b]
map (Integer
n Integer -> Integer -> Integer
forall a. C a => a -> a -> a
`div`)
                     ([Integer] -> Integer) -> [Integer] -> Integer
forall a b. (a -> b) -> a -> b
$ (Integer -> Integer) -> Integer -> [Integer]
forall a. (a -> a) -> a -> [a]
iterate (Integer -> Integer -> Integer
forall a. C a => a -> a -> a
*Integer
p) Integer
p

-- | Compute Euler's totient function (@eulerPhi n@ is the number of
--   integers less than and relatively prime to n).  Only makes sense
--   for (nonnegative) integers.
eulerPhi :: T -> Integer
eulerPhi :: T -> Integer
eulerPhi T
FQZero = Integer
1
eulerPhi (FQ Bool
_ [Integer]
pows) = [Integer] -> Integer
forall a. C a => [a] -> a
product ([Integer] -> Integer) -> [Integer] -> Integer
forall a b. (a -> b) -> a -> b
$ (Integer -> Integer -> Integer)
-> [Integer] -> [Integer] -> [Integer]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Integer -> Integer -> Integer
forall p. (Num p, C p) => p -> Integer -> p
phiP [Integer]
forall int. Integral int => [int]
primes [Integer]
pows
  where phiP :: p -> Integer -> p
phiP p
_ Integer
0 = p
1
        phiP p
p Integer
a = p
pp -> Integer -> p
forall a. C a => a -> Integer -> a
^(Integer
aInteger -> Integer -> Integer
forall a. C a => a -> a -> a
-Integer
1) p -> p -> p
forall a. C a => a -> a -> a
* (p
pp -> p -> p
forall a. C a => a -> a -> a
-p
1)

-- | List of the divisors of n.  Only makes sense for integers.
divisors :: T -> [T]
divisors :: T -> [T]
divisors T
FQZero = [T
forall a. C a => a
one]
divisors (FQ Bool
b [Integer]
pows) = ([Integer] -> T) -> [[Integer]] -> [T]
forall a b. (a -> b) -> [a] -> [b]
map (Bool -> [Integer] -> T
FQ Bool
b) ([[Integer]] -> [T]) -> [[Integer]] -> [T]
forall a b. (a -> b) -> a -> b
$ (Integer -> [Integer]) -> [Integer] -> [[Integer]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (Integer -> Integer -> [Integer]
forall a. Enum a => a -> a -> [a]
enumFromTo Integer
0) [Integer]
pows

-- | Möbius (mu) function of a positive integer: mu(n) = 0 if one or
--   more repeated prime factors, 1 if n = 1, and (-1)^k if n is a
--   product of k distinct primes.
mu :: T -> Integer
mu :: T -> Integer
mu T
FQZero      = String -> Integer
forall a. HasCallStack => String -> a
error String
"FactoredRational.mu: zero argument"
mu (FQ Bool
True [Integer]
_) = String -> Integer
forall a. HasCallStack => String -> a
error String
"FactoredRational.mu: negative argument"
mu (FQ Bool
_ [])   = Integer
1   -- mu(1) = 1
mu (FQ Bool
_ [Integer]
pows) | (Integer -> Bool) -> [Integer] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (Integer -> [Integer] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [Integer
0,Integer
1]) [Integer]
pows  =  (-Integer
1)Integer -> Integer -> Integer
forall a. C a => a -> Integer -> a
^([Integer] -> Integer
forall a. C a => [a] -> a
sum [Integer]
pows)
               | Bool
otherwise                =  Integer
0