-- | 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 FQZero = "0"
  show (FQ True pows)  = "(-1)" ++ showPows pows
  show (FQ False [])   = "1"
  show (FQ False pows) = showPows pows

showPows :: [Integer] -> String
showPows pows = concat $ zipWith showPow primes pows
  where showPow :: Integer -> Integer -> String
        showPow _ 0 = ""
        showPow p 1 = "(" ++ show p ++ ")"
        showPow p n = "(" ++ show p ++ "^" ++ show n ++ ")"

instance Additive.C T where
  zero = FQZero
  FQZero + a = a
  a + FQZero = a
  x + y = fromRational' (toRational x + toRational y)

  negate FQZero   = FQZero
  negate (FQ s e) = FQ (not s) e

instance Ring.C T where
  FQZero * _ = FQZero
  _ * FQZero = FQZero
  (FQ s1 e1) * (FQ s2 e2) = FQ (s1 /= s2) (zipWithExt 0 0 (+) e1 e2)

  fromInteger 0 = FQZero
  fromInteger n | n < 0     = FQ True (factor (negate n))
                | otherwise = FQ False (factor n)

  _ ^ 0 = one
  FQZero   ^ _ = FQZero
  (FQ s e) ^ n = FQ s (map (*n) 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 da db f = zipWithExt'
  where zipWithExt' []     bs     = zipWith f (repeat da) bs
        zipWithExt' as     []     = zipWith f as (repeat db)
        zipWithExt' (a:as) (b:bs) = f a b : zipWithExt' as 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 m = factor' m primes
  where
    factor' 1 _ = []
    factor' n (p:ps) = let (k,n') = logRem n p
                       in  k : factor' n' ps
    factor' _ [] = error "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 = logRem' 0
  where logRem' k n p | n `mod` p == 0 = logRem' (k+1) (n `div` p) p
                      | otherwise = (k,n)

instance ZeroTestable.C T where
  isZero FQZero = True
  isZero _      = False

instance Eq T where
  FQZero == FQZero   = True
  FQZero == (FQ _ _) = False
  (FQ _ _) == FQZero = False
  (FQ s1 e1) == (FQ s2 e2) = s1 == s2 && dropZeros e1 == dropZeros e2
    where dropZeros = reverse . dropWhile (==0) . reverse

instance Ord T where
  compare FQZero FQZero = EQ
  compare FQZero (FQ False _) = LT
  compare FQZero (FQ True  _) = GT
  compare (FQ False _) FQZero = GT
  compare (FQ True  _) FQZero = LT
  compare (FQ False _) (FQ True  _)   = GT
  compare (FQ True  _) (FQ False _)   = LT
  compare fq1 fq2 = compare (toRational fq1) (toRational fq2)

instance Real.C T where
  abs FQZero   = FQZero
  abs (FQ _ e) = FQ False e
  signum = Real.signumOrd

instance ToRational.C T where
  toRational FQZero   = 0
  toRational (FQ s e) = (if s then negate else id)
                      . product
                      . zipWith (^-) (map (% 1) primes)
                      $ e

instance Integral.C T where
  divMod a b =
    if isZero b
      then (undefined,a)
      else (a/b,zero)

instance RealIntegral.C T
  -- default definition is fine

instance ToInteger.C T where
  toInteger FQZero = 0
  toInteger (FQ s e) | any (<0) e = error "non-integer in FactoredRational.toInteger"
                     | otherwise  = (if s then negate else id)
                                  . product
                                  . zipWith (^) primes
                                  $ e

instance Field.C T where
  recip FQZero = error "division by zero"
  recip (FQ s e) = FQ s (map negate e)

-- | Efficiently compute n! directly as a factored rational.
factorial :: Integer -> T
factorial 0 = one
factorial 1 = one
factorial n = FQ False (takeWhile (>0) . map (factorialFactors n) $ primes)

-- | @factorialFactors n p@ computes the power of prime p in the
--   factorization of n!.
factorialFactors :: Integer -> Integer -> Integer
factorialFactors n p = sum
                     . takeWhile (>0)
                     . map (n `div`)
                     $ iterate (*p) 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 FQZero = 1
eulerPhi (FQ _ pows) = product $ zipWith phiP primes pows
  where phiP _ 0 = 1
        phiP p a = p^(a-1) * (p-1)

-- | List of the divisors of n.  Only makes sense for integers.
divisors :: T -> [T]
divisors FQZero = [one]
divisors (FQ b pows) = map (FQ b) $ mapM (enumFromTo 0) 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 FQZero      = error "FactoredRational.mu: zero argument"
mu (FQ True _) = error "FactoredRational.mu: negative argument"
mu (FQ _ [])   = 1   -- mu(1) = 1
mu (FQ _ pows) | all (`elem` [0,1]) pows  =  (-1)^(sum pows)
               | otherwise                =  0