{-# LANGUAGE PatternGuards #-}

-- | Monomials in a countably infinite set of variables x1, x2, x3, ...
module MathObj.Monomial
    ( -- * Type
      T(..)

      -- * Creating monomials
    , mkMonomial
    , constant
    , x

      -- * Utility functions
    , degree
    , pDegree
    , scaleMon

    ) where

import qualified Algebra.Additive     as Additive
import qualified Algebra.Differential as Differential
import qualified Algebra.Field        as Field
import qualified Algebra.Ring         as Ring
import qualified Algebra.ZeroTestable as ZeroTestable

import           Data.List            (intercalate, sort)
import qualified Data.Map             as M
import           Data.Ord             (comparing)

import           NumericPrelude

-- | A monomial is a map from variable indices to integer powers,
--   paired with a (polymorphic) coefficient.  Note that negative
--   integer powers are handled just fine, so monomials form a field.
--
--   Instances are provided for Eq, Ord, ZeroTestable, Additive, Ring,
--   Differential, and Field.  Note that adding two monomials only
--   makes sense if they have matching variables and exponents.  The
--   Differential instance represents partial differentiation with
--   respect to x1.
--
--   The Ord instance for monomials orders them first by permutation
--   degree, then by largest variable index (largest first), then by
--   exponent (largest first).  This may seem a bit odd, but in fact
--   reflects the use of these monomials to implement cycle index
--   series, where this ordering corresponds nicely to generation
--   of integer partitions. To make the library more general we could
--   parameterize monomials by the desired ordering.
data T a = Cons { T a -> a
coeff  :: a
                , T a -> Map Integer Integer
powers :: M.Map Integer Integer
                }

mkMonomial :: a -> [(Integer, Integer)] -> T a
mkMonomial :: a -> [(Integer, Integer)] -> T a
mkMonomial a
a [(Integer, Integer)]
p = a -> Map Integer Integer -> T a
forall a. a -> Map Integer Integer -> T a
Cons a
a ([(Integer, Integer)] -> Map Integer Integer
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [(Integer, Integer)]
p)

negOne :: Ring.C a => a
negOne :: a
negOne = a -> a
forall a. C a => a -> a
negate a
forall a. C a => a
one

instance (ZeroTestable.C a, Ring.C a, Eq a, Show a) => Show (T a) where
  show :: T a -> String
show (Cons a
a Map Integer Integer
pows) | a -> Bool
forall a. C a => a -> Bool
isZero a
a    = String
"0"
                     | Map Integer Integer -> Bool
forall k a. Map k a -> Bool
M.null Map Integer Integer
pows = a -> String
forall a. Show a => a -> String
show a
a
                     | a
a a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
forall a. C a => a
one    = Map Integer Integer -> String
showVars Map Integer Integer
pows
                     | a
a a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
forall a. C a => a
negOne = String
"-" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Map Integer Integer -> String
showVars Map Integer Integer
pows
                     | Bool
otherwise   = a -> String
forall a. Show a => a -> String
show a
a String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
" " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Map Integer Integer -> String
showVars Map Integer Integer
pows

showVars :: M.Map Integer Integer -> String
showVars :: Map Integer Integer -> String
showVars Map Integer Integer
m = String -> [String] -> String
forall a. [a] -> [[a]] -> [a]
intercalate String
" " ([String] -> String) -> [String] -> String
forall a b. (a -> b) -> a -> b
$ ((Integer, Integer) -> [String])
-> [(Integer, Integer)] -> [String]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (Integer, Integer) -> [String]
forall a a. (Eq a, Num a, Show a, Show a) => (a, a) -> [String]
showVar (Map Integer Integer -> [(Integer, Integer)]
forall k a. Map k a -> [(k, a)]
M.toList Map Integer Integer
m)
  where showVar :: (a, a) -> [String]
showVar (a
_,a
0) = []
        showVar (a
v,a
1) = [String
"x" String -> ShowS
forall a. [a] -> [a] -> [a]
++ a -> String
forall a. Show a => a -> String
show a
v]
        showVar (a
v,a
p) = [String
"x" String -> ShowS
forall a. [a] -> [a] -> [a]
++ a -> String
forall a. Show a => a -> String
show a
v String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"^" String -> ShowS
forall a. [a] -> [a] -> [a]
++ a -> String
forall a. Show a => a -> String
show a
p]

-- | The degree of a monomial is the sum of its exponents.
degree :: T a -> Integer
degree :: T a -> Integer
degree (Cons a
_ Map Integer Integer
m) = (Integer -> Integer -> Integer)
-> Integer -> Map Integer Integer -> Integer
forall a b k. (a -> b -> b) -> b -> Map k a -> b
M.foldr Integer -> Integer -> Integer
forall a. C a => a -> a -> a
(+) Integer
0 Map Integer Integer
m

-- | The \"partition degree\" of a monomial is the sum of the products
--   of each variable index with its exponent.  For example, x1^3 x2^2
--   x4^3 has partition degree 1*3 + 2*2 + 4*3 = 19.  The terminology
--   comes from the fact that, for example, we can view x1^3 x2^2 x4^3
--   as corresponding to an integer partition of 19 (namely, 1 + 1 + 1
--   + 2 + 2 + 4 + 4 + 4).
pDegree :: T a -> Integer
pDegree :: T a -> Integer
pDegree (Cons a
_ Map Integer Integer
m) = [Integer] -> Integer
forall a. C a => [a] -> a
sum ([Integer] -> Integer)
-> (Map Integer Integer -> [Integer])
-> Map Integer Integer
-> Integer
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Integer, Integer) -> Integer)
-> [(Integer, Integer)] -> [Integer]
forall a b. (a -> b) -> [a] -> [b]
map ((Integer -> Integer -> Integer) -> (Integer, Integer) -> Integer
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry Integer -> Integer -> Integer
forall a. C a => a -> a -> a
(*)) ([(Integer, Integer)] -> [Integer])
-> (Map Integer Integer -> [(Integer, Integer)])
-> Map Integer Integer
-> [Integer]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map Integer Integer -> [(Integer, Integer)]
forall k a. Map k a -> [(k, a)]
M.assocs (Map Integer Integer -> Integer) -> Map Integer Integer -> Integer
forall a b. (a -> b) -> a -> b
$ Map Integer Integer
m

-- | Create a constant monomial.
constant :: a -> T a
constant :: a -> T a
constant a
a = a -> Map Integer Integer -> T a
forall a. a -> Map Integer Integer -> T a
Cons a
a Map Integer Integer
forall k a. Map k a
M.empty

-- | Create the monomial xn for a given n.
x :: (Ring.C a) => Integer -> T a
x :: Integer -> T a
x Integer
n = a -> Map Integer Integer -> T a
forall a. a -> Map Integer Integer -> T a
Cons a
forall a. C a => a
Ring.one (Integer -> Integer -> Map Integer Integer
forall k a. k -> a -> Map k a
M.singleton Integer
n Integer
1)

-- | Scale all the variable subscripts by a constant.  Useful for
--   operations like plethyistic substitution or Mobius inversion.
scaleMon :: Integer -> T a -> T a
scaleMon :: Integer -> T a -> T a
scaleMon Integer
n (Cons a
a Map Integer Integer
m) = a -> Map Integer Integer -> T a
forall a. a -> Map Integer Integer -> T a
Cons a
a ((Integer -> Integer) -> Map Integer Integer -> Map Integer Integer
forall k2 k1 a. Ord k2 => (k1 -> k2) -> Map k1 a -> Map k2 a
M.mapKeys (Integer
nInteger -> Integer -> Integer
forall a. C a => a -> a -> a
*) Map Integer Integer
m)

instance Eq (T a) where
  (Cons a
_ Map Integer Integer
m1) == :: T a -> T a -> Bool
== (Cons a
_ Map Integer Integer
m2) = Map Integer Integer
m1 Map Integer Integer -> Map Integer Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Map Integer Integer
m2

instance Ord (T a) where
  compare :: T a -> T a -> Ordering
compare T a
m1 T a
m2
    | Integer
d1 Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
d2   = Ordering
LT
    | Integer
d1 Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
> Integer
d2   = Ordering
GT
    | Bool
otherwise = (T a -> [Rev (Integer, Integer)]) -> T a -> T a -> Ordering
forall a b. Ord a => (b -> a) -> b -> b -> Ordering
comparing T a -> [Rev (Integer, Integer)]
forall a. T a -> [Rev (Integer, Integer)]
q T a
m1 T a
m2
    where d1 :: Integer
d1 = T a -> Integer
forall a. T a -> Integer
pDegree T a
m1
          d2 :: Integer
d2 = T a -> Integer
forall a. T a -> Integer
pDegree T a
m2
          q :: T a -> [Rev (Integer, Integer)]
q  = ((Integer, Integer) -> Rev (Integer, Integer))
-> [(Integer, Integer)] -> [Rev (Integer, Integer)]
forall a b. (a -> b) -> [a] -> [b]
map (Integer, Integer) -> Rev (Integer, Integer)
forall a. a -> Rev a
Rev ([(Integer, Integer)] -> [Rev (Integer, Integer)])
-> (T a -> [(Integer, Integer)]) -> T a -> [Rev (Integer, Integer)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(Integer, Integer)] -> [(Integer, Integer)]
forall a. [a] -> [a]
reverse ([(Integer, Integer)] -> [(Integer, Integer)])
-> (T a -> [(Integer, Integer)]) -> T a -> [(Integer, Integer)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(Integer, Integer)] -> [(Integer, Integer)]
forall a. Ord a => [a] -> [a]
sort ([(Integer, Integer)] -> [(Integer, Integer)])
-> (T a -> [(Integer, Integer)]) -> T a -> [(Integer, Integer)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map Integer Integer -> [(Integer, Integer)]
forall k a. Map k a -> [(k, a)]
M.assocs (Map Integer Integer -> [(Integer, Integer)])
-> (T a -> Map Integer Integer) -> T a -> [(Integer, Integer)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. T a -> Map Integer Integer
forall a. T a -> Map Integer Integer
powers

newtype Rev a = Rev a
  deriving Rev a -> Rev a -> Bool
(Rev a -> Rev a -> Bool) -> (Rev a -> Rev a -> Bool) -> Eq (Rev a)
forall a. Eq a => Rev a -> Rev a -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Rev a -> Rev a -> Bool
$c/= :: forall a. Eq a => Rev a -> Rev a -> Bool
== :: Rev a -> Rev a -> Bool
$c== :: forall a. Eq a => Rev a -> Rev a -> Bool
Eq
instance Ord a => Ord (Rev a) where
  compare :: Rev a -> Rev a -> Ordering
compare (Rev a
a) (Rev a
b) = a -> a -> Ordering
forall a. Ord a => a -> a -> Ordering
compare a
b a
a

instance (ZeroTestable.C a) => ZeroTestable.C (T a) where
  isZero :: T a -> Bool
isZero (Cons a
a Map Integer Integer
_) = a -> Bool
forall a. C a => a -> Bool
isZero a
a

instance (Additive.C a, ZeroTestable.C a) => Additive.C (T a) where
  zero :: T a
zero = a -> Map Integer Integer -> T a
forall a. a -> Map Integer Integer -> T a
Cons a
forall a. C a => a
zero Map Integer Integer
forall k a. Map k a
M.empty
  negate :: T a -> T a
negate (Cons a
a Map Integer Integer
m) = a -> Map Integer Integer -> T a
forall a. a -> Map Integer Integer -> T a
Cons (a -> a
forall a. C a => a -> a
negate a
a) Map Integer Integer
m

  -- precondition: m1 == m2
  (Cons a
a1 Map Integer Integer
m1) + :: T a -> T a -> T a
+ (Cons a
a2 Map Integer Integer
_m2) | a -> Bool
forall a. C a => a -> Bool
isZero a
s  = a -> Map Integer Integer -> T a
forall a. a -> Map Integer Integer -> T a
Cons a
s Map Integer Integer
forall k a. Map k a
M.empty
                               | Bool
otherwise = a -> Map Integer Integer -> T a
forall a. a -> Map Integer Integer -> T a
Cons a
s Map Integer Integer
m1
                               where s :: a
s = a
a1 a -> a -> a
forall a. C a => a -> a -> a
+ a
a2

instance (Ring.C a, ZeroTestable.C a) => Ring.C (T a) where
  fromInteger :: Integer -> T a
fromInteger Integer
n = a -> Map Integer Integer -> T a
forall a. a -> Map Integer Integer -> T a
Cons (Integer -> a
forall a. C a => Integer -> a
fromInteger Integer
n) Map Integer Integer
forall k a. Map k a
M.empty
  (Cons a
a1 Map Integer Integer
m1) * :: T a -> T a -> T a
* (Cons a
a2 Map Integer Integer
m2) = a -> Map Integer Integer -> T a
forall a. a -> Map Integer Integer -> T a
Cons (a
a1a -> a -> a
forall a. C a => a -> a -> a
*a
a2)
                                     ((Integer -> Integer -> Bool)
-> Map Integer Integer -> Map Integer Integer
forall k a. (k -> a -> Bool) -> Map k a -> Map k a
M.filterWithKey (\Integer
_ Integer
p -> Bool -> Bool
not (Integer -> Bool
forall a. C a => a -> Bool
isZero Integer
p)) (Map Integer Integer -> Map Integer Integer)
-> Map Integer Integer -> Map Integer Integer
forall a b. (a -> b) -> a -> b
$
                                        (Integer -> Integer -> Integer)
-> Map Integer Integer
-> Map Integer Integer
-> Map Integer Integer
forall k a. Ord k => (a -> a -> a) -> Map k a -> Map k a -> Map k a
M.unionWith Integer -> Integer -> Integer
forall a. C a => a -> a -> a
(+) Map Integer Integer
m1 Map Integer Integer
m2
                                     )

-- Partial differentiation with respect to x1.
instance (ZeroTestable.C a, Ring.C a) => Differential.C (T a) where
  differentiate :: T a -> T a
differentiate (Cons a
a Map Integer Integer
m)
    | Just Integer
p <- Integer -> Map Integer Integer -> Maybe Integer
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Integer
1 Map Integer Integer
m = a -> Map Integer Integer -> T a
forall a. a -> Map Integer Integer -> T a
Cons (a
aa -> a -> a
forall a. C a => a -> a -> a
*Integer -> a
forall a. C a => Integer -> a
fromInteger Integer
p) ((Integer -> Maybe Integer)
-> Integer -> Map Integer Integer -> Map Integer Integer
forall k a. Ord k => (a -> Maybe a) -> k -> Map k a -> Map k a
M.update Integer -> Maybe Integer
forall a. (Eq a, Num a, C a) => a -> Maybe a
powerPred Integer
1 Map Integer Integer
m)
    | Bool
otherwise              = a -> Map Integer Integer -> T a
forall a. a -> Map Integer Integer -> T a
Cons a
forall a. C a => a
zero Map Integer Integer
forall k a. Map k a
M.empty
    where
      powerPred :: a -> Maybe a
powerPred a
1 = Maybe a
forall a. Maybe a
Nothing
      powerPred a
p = a -> Maybe a
forall a. a -> Maybe a
Just (a
pa -> a -> a
forall a. C a => a -> a -> a
-a
1)

instance (ZeroTestable.C a, Field.C a, Eq a) => Field.C (T a) where
  recip :: T a -> T a
recip (Cons a
a Map Integer Integer
pows) = if a -> Bool
forall a. C a => a -> Bool
isZero a
a
    then String -> T a
forall a. HasCallStack => String -> a
error String
"Monomial.recip: division by zero"
    else a -> Map Integer Integer -> T a
forall a. a -> Map Integer Integer -> T a
Cons (a -> a
forall a. C a => a -> a
recip a
a) ((Integer -> Integer) -> Map Integer Integer -> Map Integer Integer
forall a b k. (a -> b) -> Map k a -> Map k b
M.map Integer -> Integer
forall a. C a => a -> a
negate Map Integer Integer
pows)