{-# OPTIONS_GHC -fno-warn-name-shadowing #-}

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

      -- * Constructing polynomials
    , fromMonomials
    , lift0
    , lift1
    , lift2
    , x
    , constant

      -- * Operations

    , compose

    , merge

    ) where

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

import qualified MathObj.Monomial as Mon

import qualified Data.Map as M

import NumericPrelude

-- | A polynomial is just a list of monomials, construed as their sum.
--   We maintain the invariant that polynomials are always sorted by
--   the ordering on monomials defined in "MathObj.Monomial": first by
--   partition degree, then by largest variable index (decreasing),
--   then by exponent of the highest-index variable (decreasing).
--   This works out nicely for operations on cycle index series.
--
--   Instances are provided for Additive, Ring, Differential
--   (partial differentiation with respect to x1), and Show.
newtype T a = Cons [Mon.T a]

instance (ZeroTestable.C a, Ring.C a, Ord a, Show a) => Show (T a) where
  show :: T a -> String
show (Cons []) = String
"0"
  show (Cons (T a
m:[T a]
ms)) = T a -> String
forall a. Show a => a -> String
show T a
m String -> ShowS
forall a. [a] -> [a] -> [a]
++ (T a -> String) -> [T a] -> String
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap T a -> String
forall a. (Ord a, C a, C a, Show a) => T a -> String
showMon [T a]
ms
    where showMon :: T a -> String
showMon T a
m | T a -> a
forall a. T a -> a
Mon.coeff T a
m a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
forall a. C a => a
zero = String
" - " String -> ShowS
forall a. [a] -> [a] -> [a]
++ T a -> String
forall a. Show a => a -> String
show (T a -> T a
forall a. C a => a -> a
negate T a
m)
                    | Bool
otherwise          = String
" + " String -> ShowS
forall a. [a] -> [a] -> [a]
++ T a -> String
forall a. Show a => a -> String
show T a
m

{-# INLINE fromMonomials #-}
fromMonomials :: [Mon.T a] -> T a
fromMonomials :: [T a] -> T a
fromMonomials = [T a] -> T a
forall a. [T a] -> T a
lift0

{-# INLINE lift0 #-}
lift0 :: [Mon.T a] -> T a
lift0 :: [T a] -> T a
lift0 = [T a] -> T a
forall a. [T a] -> T a
Cons

{-# INLINE lift1 #-}
lift1 :: ([Mon.T a] -> [Mon.T a]) -> (T a -> T a)
lift1 :: ([T a] -> [T a]) -> T a -> T a
lift1 [T a] -> [T a]
f (Cons [T a]
xs) = [T a] -> T a
forall a. [T a] -> T a
Cons ([T a] -> [T a]
f [T a]
xs)

{-# INLINE lift2 #-}
lift2 :: ([Mon.T a] -> [Mon.T a] -> [Mon.T a]) -> (T a -> T a -> T a)
lift2 :: ([T a] -> [T a] -> [T a]) -> T a -> T a -> T a
lift2 [T a] -> [T a] -> [T a]
f (Cons [T a]
xs) (Cons [T a]
ys) = [T a] -> T a
forall a. [T a] -> T a
Cons ([T a] -> [T a] -> [T a]
f [T a]
xs [T a]
ys)

-- | Create the polynomial xn for a given n.
x :: (Ring.C a) => Integer -> T a
x :: Integer -> T a
x Integer
n = [T a] -> T a
forall a. [T a] -> T a
fromMonomials [Integer -> T a
forall a. C a => Integer -> T a
Mon.x Integer
n]

-- | Create a constant polynomial.
constant :: a -> T a
constant :: a -> T a
constant a
a = [T a] -> T a
forall a. [T a] -> T a
fromMonomials [a -> T a
forall a. a -> T a
Mon.constant a
a]

-- | Add two polynomials.  We assume that they are already sorted, so
--   that addition works on infinite polynomials.
add :: (Ord a, Additive.C a) => [a] -> [a] -> [a]
add :: [a] -> [a] -> [a]
add [a]
xs [a]
ys = Bool -> (a -> a -> a) -> [a] -> [a] -> [a]
forall a. Ord a => Bool -> (a -> a -> a) -> [a] -> [a] -> [a]
merge Bool
True a -> a -> a
forall a. C a => a -> a -> a
(+) [a]
xs [a]
ys

-- | Merge two sorted lists, with a flag specifying whether to keep
--   singletons, and a combining function for elements that are equal.
merge :: Ord a => Bool -> (a -> a -> a) -> [a] -> [a] -> [a]
merge :: Bool -> (a -> a -> a) -> [a] -> [a] -> [a]
merge Bool
True  a -> a -> a
_ [] [a]
ys = [a]
ys
merge Bool
False a -> a -> a
_ [] [a]
_  = []
merge Bool
True  a -> a -> a
_ [a]
xs [] = [a]
xs
merge Bool
False a -> a -> a
_ [a]
_  [] = []
merge Bool
b a -> a -> a
f xxs :: [a]
xxs@(a
x:[a]
xs) yys :: [a]
yys@(a
y:[a]
ys) | a
x a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
y     = Bool -> ([a] -> [a]) -> ([a] -> [a]) -> [a] -> [a]
forall p. Bool -> p -> p -> p
if' Bool
b (a
xa -> [a] -> [a]
forall a. a -> [a] -> [a]
:) [a] -> [a]
forall a. a -> a
id ([a] -> [a]) -> [a] -> [a]
forall a b. (a -> b) -> a -> b
$ Bool -> (a -> a -> a) -> [a] -> [a] -> [a]
forall a. Ord a => Bool -> (a -> a -> a) -> [a] -> [a] -> [a]
merge Bool
b a -> a -> a
f [a]
xs [a]
yys
                                | a
x a -> a -> Bool
forall a. Ord a => a -> a -> Bool
> a
y     = Bool -> ([a] -> [a]) -> ([a] -> [a]) -> [a] -> [a]
forall p. Bool -> p -> p -> p
if' Bool
b (a
ya -> [a] -> [a]
forall a. a -> [a] -> [a]
:) [a] -> [a]
forall a. a -> a
id ([a] -> [a]) -> [a] -> [a]
forall a b. (a -> b) -> a -> b
$ Bool -> (a -> a -> a) -> [a] -> [a] -> [a]
forall a. Ord a => Bool -> (a -> a -> a) -> [a] -> [a] -> [a]
merge Bool
b a -> a -> a
f [a]
xxs [a]
ys
                                | Bool
otherwise = a -> a -> a
f a
x a
y a -> [a] -> [a]
forall a. a -> [a] -> [a]
: Bool -> (a -> a -> a) -> [a] -> [a] -> [a]
forall a. Ord a => Bool -> (a -> a -> a) -> [a] -> [a] -> [a]
merge Bool
b a -> a -> a
f [a]
xs [a]
ys
  where if' :: Bool -> p -> p -> p
if' Bool
True p
x p
_ = p
x
        if' Bool
False p
_ p
y = p
y

instance (Additive.C a, ZeroTestable.C a) => Additive.C (T a) where
  zero :: T a
zero   = [T a] -> T a
forall a. [T a] -> T a
fromMonomials []
  negate :: T a -> T a
negate = ([T a] -> [T a]) -> T a -> T a
forall a. ([T a] -> [T a]) -> T a -> T a
lift1 (([T a] -> [T a]) -> T a -> T a) -> ([T a] -> [T a]) -> T a -> T a
forall a b. (a -> b) -> a -> b
$ (T a -> T a) -> [T a] -> [T a]
forall a b. (a -> b) -> [a] -> [b]
map T a -> T a
forall a. C a => a -> a
negate
  + :: T a -> T a -> T a
(+)    = ([T a] -> [T a] -> [T a]) -> T a -> T a -> T a
forall a. ([T a] -> [T a] -> [T a]) -> T a -> T a -> T a
lift2 [T a] -> [T a] -> [T a]
forall a. (Ord a, C a) => [a] -> [a] -> [a]
add

-- | Multiply two (sorted) polynomials.
mul :: (Ring.C a, Ord a) => [a] -> [a] -> [a]
mul :: [a] -> [a] -> [a]
mul [] [a]
_ = []
mul [a]
_ [] = []
mul (a
x:[a]
xs) (a
y:[a]
ys) = a
xa -> a -> a
forall a. C a => a -> a -> a
*a
y a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a] -> [a] -> [a]
forall a. (Ord a, C a) => [a] -> [a] -> [a]
add ((a -> a) -> [a] -> [a]
forall a b. (a -> b) -> [a] -> [b]
map (a
xa -> a -> a
forall a. C a => a -> a -> a
*) [a]
ys) ([a] -> [a] -> [a]
forall a. (C a, Ord a) => [a] -> [a] -> [a]
mul [a]
xs (a
ya -> [a] -> [a]
forall a. a -> [a] -> [a]
:[a]
ys))

instance (Ring.C a, ZeroTestable.C a) => Ring.C (T a) where
  fromInteger :: Integer -> T a
fromInteger Integer
n = [T a] -> T a
forall a. [T a] -> T a
fromMonomials [Integer -> T a
forall a. C a => Integer -> a
fromInteger Integer
n]
  * :: T a -> T a -> T a
(*) = ([T a] -> [T a] -> [T a]) -> T a -> T a -> T a
forall a. ([T a] -> [T a] -> [T a]) -> T a -> T a -> T a
lift2 [T a] -> [T a] -> [T a]
forall a. (C a, Ord a) => [a] -> [a] -> [a]
mul

-- Partial differentiation with respect to x1.
instance (ZeroTestable.C a, Ring.C a) => Differential.C (T a) where
  differentiate :: T a -> T a
differentiate = ([T a] -> [T a]) -> T a -> T a
forall a. ([T a] -> [T a]) -> T a -> T a
lift1 (([T a] -> [T a]) -> T a -> T a) -> ([T a] -> [T a]) -> T a -> T a
forall a b. (a -> b) -> a -> b
$ (T a -> Bool) -> [T a] -> [T a]
forall a. (a -> Bool) -> [a] -> [a]
filter (Bool -> Bool
not (Bool -> Bool) -> (T a -> Bool) -> T a -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. T a -> Bool
forall a. C a => a -> Bool
isZero) ([T a] -> [T a]) -> ([T a] -> [T a]) -> [T a] -> [T a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (T a -> T a) -> [T a] -> [T a]
forall a b. (a -> b) -> [a] -> [b]
map T a -> T a
forall a. C a => a -> a
Differential.differentiate

-- | Plethyistic substitution: F o G = F(G(x1,x2,x3...),
--   G(x2,x4,x6...), G(x3,x6,x9...), ...)  See Bergeron, Labelle, and
--   Leroux, \"Combinatorial Species and Tree-Like Structures\",
--   p. 43.
compose :: (Ring.C a, ZeroTestable.C a) => T a -> T a -> T a
compose :: T a -> T a -> T a
compose (Cons []) T a
_ = [T a] -> T a
forall a. [T a] -> T a
Cons []
compose (Cons (T a
x:[T a]
_)) (Cons []) = [T a] -> T a
forall a. [T a] -> T a
Cons [T a
x]
compose (Cons [T a]
xs) yys :: T a
yys@(Cons (T a
y:[T a]
_))
  | T a -> Integer
forall a. T a -> Integer
Mon.degree T a
y Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
0 Bool -> Bool -> Bool
&& (Bool -> Bool
not (Bool -> Bool) -> (T a -> Bool) -> T a -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Bool
forall a. C a => a -> Bool
isZero (a -> Bool) -> (T a -> a) -> T a -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. T a -> a
forall a. T a -> a
Mon.coeff (T a -> Bool) -> T a -> Bool
forall a b. (a -> b) -> a -> b
$ T a
y)
    = String -> T a
forall a. HasCallStack => String -> a
error (String -> T a) -> String -> T a
forall a b. (a -> b) -> a -> b
$ String
"MultiVarPolynomial.compose: inner series must not have a constant term."
  | Bool
otherwise = [T a] -> T a -> T a
forall a. (C a, C a) => [T a] -> T a -> T a
comp [T a]
xs T a
yys

-- | We need to be careful to make sure this is suitably
--   lazy. For example, this works for finite polynomials:
--
-- > comp ms p = sum . map (substMon p) $ ms
--
--   but not for infinite ones!
--
--   This is accomplished by calling a recursive helper function
--   taking as an extra argument a running sum containing only terms
--   with partition degree greater than or equal to the most recently
--   processed monomial.  Plethyistically substituting a polynomial
--   (with no constant term) into a monomial of partition degree d
--   produces a polynomial with all terms of partition degree >= d, so
--   when we encounter a monomial with partition degree d, we know we
--   are done with all terms in the running sum of lesser partition
--   degree.
--
--   Precondition: the second argument has no constant term.
comp :: (Ring.C a, ZeroTestable.C a) => [Mon.T a] -> T a -> T a
comp :: [T a] -> T a -> T a
comp [T a]
ms T a
p = T a -> [T a] -> T a
comp' T a
forall a. C a => a
zero [T a]
ms
  where -- comp' :: T a -> [Mon.T a] -> T a
        comp' :: T a -> [T a] -> T a
comp' T a
part []     = T a
part
        comp' T a
part (T a
m:[T a]
ms) = ([T a] -> [T a] -> [T a]) -> T a -> T a -> T a
forall a. ([T a] -> [T a] -> [T a]) -> T a -> T a -> T a
lift2 [T a] -> [T a] -> [T a]
forall a. [a] -> [a] -> [a]
(++) T a
done (T a -> T a) -> T a -> T a
forall a b. (a -> b) -> a -> b
$ T a -> [T a] -> T a
comp' (T a
rest T a -> T a -> T a
forall a. C a => a -> a -> a
+ T a -> T a -> T a
forall a. (C a, C a) => T a -> T a -> T a
substMon T a
p T a
m) [T a]
ms
          where (T a
done,T a
rest) = (T a -> Bool) -> T a -> (T a, T a)
forall a. (T a -> Bool) -> T a -> (T a, T a)
splitPoly ((Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< T a -> Integer
forall a. T a -> Integer
Mon.pDegree T a
m) (Integer -> Bool) -> (T a -> Integer) -> T a -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. T a -> Integer
forall a. T a -> Integer
Mon.pDegree) T a
part

-- | Plethyistic substitution of a polynomial into a monomial.
substMon :: (ZeroTestable.C a, Ring.C a) => T a -> Mon.T a -> T a
substMon :: T a -> T a -> T a
substMon T a
poly T a
m
  = (a -> T a
forall a. a -> T a
constant (T a -> a
forall a. T a -> a
Mon.coeff T a
m) T a -> T a -> T a
forall a. C a => a -> a -> a
*)
  (T a -> T a)
-> (Map Integer Integer -> T a) -> Map Integer Integer -> T a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Integer -> Integer -> T a -> T a)
-> T a -> Map Integer Integer -> T a
forall k a b. (k -> a -> b -> b) -> b -> Map k a -> b
M.foldrWithKey (\Integer
sub Integer
pow -> T a -> T a -> T a
forall a. C a => a -> a -> a
(*) (Integer -> T a -> T a
forall a. Integer -> T a -> T a
scalePoly Integer
sub T a
poly T a -> Integer -> T a
forall a. C a => a -> Integer -> a
^Integer
pow)) T a
forall a. C a => a
one
  (Map Integer Integer -> T a) -> Map Integer Integer -> T a
forall a b. (a -> b) -> a -> b
$ T a -> Map Integer Integer
forall a. T a -> Map Integer Integer
Mon.powers T a
m

-- | @scalePoly n Z@ changes Z(x_1, x_2, x_3, ...) into Z(x_n, x_2n, x_3n, ...)
scalePoly :: Integer -> T a -> T a
scalePoly :: Integer -> T a -> T a
scalePoly Integer
n = ([T a] -> [T a]) -> T a -> T a
forall a. ([T a] -> [T a]) -> T a -> T a
lift1 (([T a] -> [T a]) -> T a -> T a) -> ([T a] -> [T a]) -> T a -> T a
forall a b. (a -> b) -> a -> b
$ (T a -> T a) -> [T a] -> [T a]
forall a b. (a -> b) -> [a] -> [b]
map (Integer -> T a -> T a
forall a. Integer -> T a -> T a
Mon.scaleMon Integer
n)

-- | Split a polynomial into two pieces based on a predicate.
splitPoly :: (Mon.T a -> Bool) -> T a -> (T a, T a)
splitPoly :: (T a -> Bool) -> T a -> (T a, T a)
splitPoly T a -> Bool
p (Cons [T a]
xs) = ([T a] -> T a
forall a. [T a] -> T a
Cons [T a]
ys, [T a] -> T a
forall a. [T a] -> T a
Cons [T a]
zs)
  where ([T a]
ys, [T a]
zs) = (T a -> Bool) -> [T a] -> ([T a], [T a])
forall a. (a -> Bool) -> [a] -> ([a], [a])
span T a -> Bool
p [T a]
xs