{-# LANGUAGE ParallelListComp, ViewPatterns, FlexibleContexts #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}
module Math.Polynomial
    ( Endianness(..)
    , Poly, poly, polyDegree
    , polyCoeffs, polyIsZero, polyIsOne
    , zero, one, constPoly, x
    , scalePoly, negatePoly
    , composePoly
    , addPoly, sumPolys, multPoly, powPoly
    , quotRemPoly, quotPoly, remPoly
    , evalPoly, evalPolyDeriv, evalPolyDerivs
    , contractPoly
    , monicPoly
    , gcdPoly, separateRoots
    , polyDeriv, polyDerivs, polyIntegral
    ) where

import Math.Polynomial.Type
import Math.Polynomial.Pretty ({- instance -})

import Math.Polynomial.VectorSpace (one, x) -- to re-export
import qualified Math.Polynomial.VectorSpace as VS
import Data.VectorSpace.WrappedNum

-- |Given some constant 'k', construct the polynomial whose value is 
-- constantly 'k'.
constPoly :: (Num a, Eq a) => a -> Poly a
constPoly :: a -> Poly a
constPoly x :: a
x = Poly (WrappedNum a) -> Poly a
forall a. Poly (WrappedNum a) -> Poly a
unwrapPoly (WrappedNum a -> Poly (WrappedNum a)
forall a. (Eq a, AdditiveGroup a) => a -> Poly a
VS.constPoly (a -> WrappedNum a
forall a. a -> WrappedNum a
WrapNum a
x))

-- |Given some scalar 's' and a polynomial 'f', computes the polynomial 'g'
-- such that:
-- 
-- > evalPoly g x = s * evalPoly f x
scalePoly :: (Num a, Eq a) => a -> Poly a -> Poly a
scalePoly :: a -> Poly a -> Poly a
scalePoly x :: a
x f :: Poly a
f = Poly (WrappedNum a) -> Poly a
forall a. Poly (WrappedNum a) -> Poly a
unwrapPoly (Scalar (WrappedNum a) -> Poly (WrappedNum a) -> Poly (WrappedNum a)
forall a.
(Eq a, VectorSpace a, AdditiveGroup (Scalar a), Eq (Scalar a)) =>
Scalar a -> Poly a -> Poly a
VS.scalePoly (a -> WrappedNum a
forall a. a -> WrappedNum a
WrapNum a
x) (Poly a -> Poly (WrappedNum a)
forall a. Poly a -> Poly (WrappedNum a)
wrapPoly Poly a
f))

-- |Given some polynomial 'f', computes the polynomial 'g' such that:
-- 
-- > evalPoly g x = negate (evalPoly f x)
negatePoly :: (Num a, Eq a) => Poly a -> Poly a
negatePoly :: Poly a -> Poly a
negatePoly f :: Poly a
f = Poly (WrappedNum a) -> Poly a
forall a. Poly (WrappedNum a) -> Poly a
unwrapPoly (Poly (WrappedNum a) -> Poly (WrappedNum a)
forall a. (AdditiveGroup a, Eq a) => Poly a -> Poly a
VS.negatePoly (Poly a -> Poly (WrappedNum a)
forall a. Poly a -> Poly (WrappedNum a)
wrapPoly Poly a
f))

-- |Given polynomials 'f' and 'g', computes the polynomial 'h' such that:
-- 
-- > evalPoly h x = evalPoly f x + evalPoly g x
addPoly :: (Num a, Eq a) => Poly a -> Poly a -> Poly a
addPoly :: Poly a -> Poly a -> Poly a
addPoly p :: Poly a
p q :: Poly a
q = Poly (WrappedNum a) -> Poly a
forall a. Poly (WrappedNum a) -> Poly a
unwrapPoly (Poly (WrappedNum a) -> Poly (WrappedNum a) -> Poly (WrappedNum a)
forall a. (AdditiveGroup a, Eq a) => Poly a -> Poly a -> Poly a
VS.addPoly (Poly a -> Poly (WrappedNum a)
forall a. Poly a -> Poly (WrappedNum a)
wrapPoly Poly a
p) (Poly a -> Poly (WrappedNum a)
forall a. Poly a -> Poly (WrappedNum a)
wrapPoly Poly a
q))

{-# RULES
  "sum Poly"    forall ps. foldl addPoly zero ps = sumPolys ps
  #-}
sumPolys :: (Num a, Eq a) => [Poly a] -> Poly a
sumPolys :: [Poly a] -> Poly a
sumPolys ps :: [Poly a]
ps = Poly (WrappedNum a) -> Poly a
forall a. Poly (WrappedNum a) -> Poly a
unwrapPoly ([Poly (WrappedNum a)] -> Poly (WrappedNum a)
forall a. (AdditiveGroup a, Eq a) => [Poly a] -> Poly a
VS.sumPolys ((Poly a -> Poly (WrappedNum a))
-> [Poly a] -> [Poly (WrappedNum a)]
forall a b. (a -> b) -> [a] -> [b]
map Poly a -> Poly (WrappedNum a)
forall a. Poly a -> Poly (WrappedNum a)
wrapPoly [Poly a]
ps))

-- |Given polynomials 'f' and 'g', computes the polynomial 'h' such that:
-- 
-- > evalPoly h x = evalPoly f x * evalPoly g x
multPoly :: (Num a, Eq a) => Poly a -> Poly a -> Poly a
multPoly :: Poly a -> Poly a -> Poly a
multPoly p :: Poly a
p q :: Poly a
q = Poly (WrappedNum a) -> Poly a
forall a. Poly (WrappedNum a) -> Poly a
unwrapPoly ((WrappedNum a -> WrappedNum a -> WrappedNum a)
-> Poly (WrappedNum a)
-> Poly (WrappedNum a)
-> Poly (WrappedNum a)
forall a.
(AdditiveGroup a, Eq a) =>
(a -> a -> a) -> Poly a -> Poly a -> Poly a
VS.multPolyWith WrappedNum a -> WrappedNum a -> WrappedNum a
forall a. Num a => a -> a -> a
(*) (Poly a -> Poly (WrappedNum a)
forall a. Poly a -> Poly (WrappedNum a)
wrapPoly Poly a
p) (Poly a -> Poly (WrappedNum a)
forall a. Poly a -> Poly (WrappedNum a)
wrapPoly Poly a
q))

-- |Given a polynomial 'f' and exponent 'n', computes the polynomial 'g'
-- such that:
-- 
-- > evalPoly g x = evalPoly f x ^ n
powPoly :: (Num a, Eq a, Integral b) => Poly a -> b -> Poly a
powPoly :: Poly a -> b -> Poly a
powPoly p :: Poly a
p n :: b
n = Poly (WrappedNum a) -> Poly a
forall a. Poly (WrappedNum a) -> Poly a
unwrapPoly (WrappedNum a
-> (WrappedNum a -> WrappedNum a -> WrappedNum a)
-> Poly (WrappedNum a)
-> b
-> Poly (WrappedNum a)
forall a b.
(AdditiveGroup a, Eq a, Integral b) =>
a -> (a -> a -> a) -> Poly a -> b -> Poly a
VS.powPolyWith 1 WrappedNum a -> WrappedNum a -> WrappedNum a
forall a. Num a => a -> a -> a
(*) (Poly a -> Poly (WrappedNum a)
forall a. Poly a -> Poly (WrappedNum a)
wrapPoly Poly a
p) b
n)

-- |Given polynomials @a@ and @b@, with @b@ not 'zero', computes polynomials
-- @q@ and @r@ such that:
-- 
-- > addPoly (multPoly q b) r == a
quotRemPoly :: (Fractional a, Eq a) => Poly a -> Poly a -> (Poly a, Poly a)
quotRemPoly :: Poly a -> Poly a -> (Poly a, Poly a)
quotRemPoly u :: Poly a
u v :: Poly a
v = (Poly (WrappedNum a) -> Poly a
forall a. Poly (WrappedNum a) -> Poly a
unwrapPoly Poly (WrappedNum a)
q, Poly (WrappedNum a) -> Poly a
forall a. Poly (WrappedNum a) -> Poly a
unwrapPoly Poly (WrappedNum a)
r)
    where
        ~(q :: Poly (WrappedNum a)
q, r :: Poly (WrappedNum a)
r) = (WrappedNum a -> WrappedNum a -> WrappedNum a)
-> (WrappedNum a -> WrappedNum a -> WrappedNum a)
-> Poly (WrappedNum a)
-> Poly (WrappedNum a)
-> (Poly (WrappedNum a), Poly (WrappedNum a))
forall a.
(AdditiveGroup a, Eq a) =>
(a -> a -> a)
-> (a -> a -> a) -> Poly a -> Poly a -> (Poly a, Poly a)
VS.quotRemPolyWith WrappedNum a -> WrappedNum a -> WrappedNum a
forall a. Num a => a -> a -> a
(*) WrappedNum a -> WrappedNum a -> WrappedNum a
forall a. Fractional a => a -> a -> a
(/) (Poly a -> Poly (WrappedNum a)
forall a. Poly a -> Poly (WrappedNum a)
wrapPoly Poly a
u) (Poly a -> Poly (WrappedNum a)
forall a. Poly a -> Poly (WrappedNum a)
wrapPoly Poly a
v)

quotPoly :: (Fractional a, Eq a) => Poly a -> Poly a -> Poly a
quotPoly :: Poly a -> Poly a -> Poly a
quotPoly u :: Poly a
u v :: Poly a
v = Poly (WrappedNum a) -> Poly a
forall a. Poly (WrappedNum a) -> Poly a
unwrapPoly ((WrappedNum a -> WrappedNum a -> WrappedNum a)
-> (WrappedNum a -> WrappedNum a -> WrappedNum a)
-> Poly (WrappedNum a)
-> Poly (WrappedNum a)
-> Poly (WrappedNum a)
forall a.
(AdditiveGroup a, Eq a) =>
(a -> a -> a) -> (a -> a -> a) -> Poly a -> Poly a -> Poly a
VS.quotPolyWith WrappedNum a -> WrappedNum a -> WrappedNum a
forall a. Num a => a -> a -> a
(*) WrappedNum a -> WrappedNum a -> WrappedNum a
forall a. Fractional a => a -> a -> a
(/) (Poly a -> Poly (WrappedNum a)
forall a. Poly a -> Poly (WrappedNum a)
wrapPoly Poly a
u) (Poly a -> Poly (WrappedNum a)
forall a. Poly a -> Poly (WrappedNum a)
wrapPoly Poly a
v))

remPoly :: (Fractional a,  Eq a) => Poly a -> Poly a -> Poly a
remPoly :: Poly a -> Poly a -> Poly a
remPoly u :: Poly a
u v :: Poly a
v = Poly (WrappedNum a) -> Poly a
forall a. Poly (WrappedNum a) -> Poly a
unwrapPoly ((WrappedNum a -> WrappedNum a -> WrappedNum a)
-> (WrappedNum a -> WrappedNum a -> WrappedNum a)
-> Poly (WrappedNum a)
-> Poly (WrappedNum a)
-> Poly (WrappedNum a)
forall a.
(AdditiveGroup a, Eq a) =>
(a -> a -> a) -> (a -> a -> a) -> Poly a -> Poly a -> Poly a
VS.remPolyWith WrappedNum a -> WrappedNum a -> WrappedNum a
forall a. Num a => a -> a -> a
(*) WrappedNum a -> WrappedNum a -> WrappedNum a
forall a. Fractional a => a -> a -> a
(/) (Poly a -> Poly (WrappedNum a)
forall a. Poly a -> Poly (WrappedNum a)
wrapPoly Poly a
u) (Poly a -> Poly (WrappedNum a)
forall a. Poly a -> Poly (WrappedNum a)
wrapPoly Poly a
v))

-- |@composePoly f g@ constructs the polynomial 'h' such that:
-- 
-- > evalPoly h = evalPoly f . evalPoly g
-- 
-- This is a very expensive operation and, in general, returns a polynomial 
-- that is quite a bit more expensive to evaluate than @f@ and @g@ together
-- (because it is of a much higher order than either).  Unless your 
-- polynomials are quite small or you are quite certain you need the
-- coefficients of the composed polynomial, it is recommended that you 
-- simply evaluate @f@ and @g@ and explicitly compose the resulting 
-- functions.  This will usually be much more efficient.
composePoly :: (Num a, Eq a) => Poly a -> Poly a -> Poly a
composePoly :: Poly a -> Poly a -> Poly a
composePoly p :: Poly a
p q :: Poly a
q = Poly (WrappedNum a) -> Poly a
forall a. Poly (WrappedNum a) -> Poly a
unwrapPoly ((WrappedNum a -> WrappedNum a -> WrappedNum a)
-> Poly (WrappedNum a)
-> Poly (WrappedNum a)
-> Poly (WrappedNum a)
forall a.
(AdditiveGroup a, Eq a) =>
(a -> a -> a) -> Poly a -> Poly a -> Poly a
VS.composePolyWith WrappedNum a -> WrappedNum a -> WrappedNum a
forall a. Num a => a -> a -> a
(*) (Poly a -> Poly (WrappedNum a)
forall a. Poly a -> Poly (WrappedNum a)
wrapPoly Poly a
p) (Poly a -> Poly (WrappedNum a)
forall a. Poly a -> Poly (WrappedNum a)
wrapPoly Poly a
q))

-- |Evaluate a polynomial at a point or, equivalently, convert a polynomial
-- to the function it represents.  For example, @evalPoly 'x' = 'id'@ and 
-- @evalPoly ('constPoly' k) = 'const' k.@
evalPoly :: (Num a, Eq a) => Poly a -> a -> a
evalPoly :: Poly a -> a -> a
evalPoly f :: Poly a
f x :: a
x = WrappedNum a -> a
forall a. WrappedNum a -> a
unwrapNum (Poly (WrappedNum a) -> Scalar (WrappedNum a) -> WrappedNum a
forall a.
(VectorSpace a, Eq a, AdditiveGroup (Scalar a), Eq (Scalar a)) =>
Poly a -> Scalar a -> a
VS.evalPoly (Poly a -> Poly (WrappedNum a)
forall a. Poly a -> Poly (WrappedNum a)
wrapPoly Poly a
f) (a -> WrappedNum a
forall a. a -> WrappedNum a
WrapNum a
x))

-- |Evaluate a polynomial and its derivative (respectively) at a point.
evalPolyDeriv :: (Num a, Eq a) => Poly a -> a -> (a,a)
evalPolyDeriv :: Poly a -> a -> (a, a)
evalPolyDeriv f :: Poly a
f x :: a
x = (WrappedNum a -> a
forall a. WrappedNum a -> a
unwrapNum WrappedNum a
y, WrappedNum a -> a
forall a. WrappedNum a -> a
unwrapNum WrappedNum a
y')
    where
        ~(y :: WrappedNum a
y, y' :: WrappedNum a
y') = Poly (WrappedNum a)
-> Scalar (WrappedNum a) -> (WrappedNum a, WrappedNum a)
forall a. (VectorSpace a, Eq a) => Poly a -> Scalar a -> (a, a)
VS.evalPolyDeriv (Poly a -> Poly (WrappedNum a)
forall a. Poly a -> Poly (WrappedNum a)
wrapPoly Poly a
f) (a -> WrappedNum a
forall a. a -> WrappedNum a
WrapNum a
x)

-- |Evaluate a polynomial and all of its nonzero derivatives at a point.  
-- This is roughly equivalent to:
-- 
-- > evalPolyDerivs p x = map (`evalPoly` x) (takeWhile (not . polyIsZero) (iterate polyDeriv p))
evalPolyDerivs :: (Num a, Eq a) => Poly a -> a -> [a]
evalPolyDerivs :: Poly a -> a -> [a]
evalPolyDerivs f :: Poly a
f x :: a
x = (WrappedNum a -> a) -> [WrappedNum a] -> [a]
forall a b. (a -> b) -> [a] -> [b]
map WrappedNum a -> a
forall a. WrappedNum a -> a
unwrapNum (Poly (WrappedNum a) -> Scalar (WrappedNum a) -> [WrappedNum a]
forall a.
(VectorSpace a, Eq a, Num (Scalar a)) =>
Poly a -> Scalar a -> [a]
VS.evalPolyDerivs (Poly a -> Poly (WrappedNum a)
forall a. Poly a -> Poly (WrappedNum a)
wrapPoly Poly a
f) (a -> WrappedNum a
forall a. a -> WrappedNum a
WrapNum a
x))

-- |\"Contract\" a polynomial by attempting to divide out a root.
--
-- @contractPoly p a@ returns @(q,r)@ such that @q*(x-a) + r == p@
contractPoly :: (Num a, Eq a) => Poly a -> a -> (Poly a, a)
contractPoly :: Poly a -> a -> (Poly a, a)
contractPoly p :: Poly a
p a :: a
a = (Poly (WrappedNum a) -> Poly a
forall a. Poly (WrappedNum a) -> Poly a
unwrapPoly Poly (WrappedNum a)
q, WrappedNum a -> a
forall a. WrappedNum a -> a
unwrapNum WrappedNum a
r)
    where
        (q :: Poly (WrappedNum a)
q, r :: WrappedNum a
r) = Poly (WrappedNum a)
-> Scalar (WrappedNum a) -> (Poly (WrappedNum a), WrappedNum a)
forall a.
(VectorSpace a, Eq a) =>
Poly a -> Scalar a -> (Poly a, a)
VS.contractPoly (Poly a -> Poly (WrappedNum a)
forall a. Poly a -> Poly (WrappedNum a)
wrapPoly Poly a
p) (a -> WrappedNum a
forall a. a -> WrappedNum a
WrapNum a
a)

-- |@gcdPoly a b@ computes the highest order monic polynomial that is a 
-- divisor of both @a@ and @b@.  If both @a@ and @b@ are 'zero', the 
-- result is undefined.
gcdPoly :: (Fractional a, Eq a) => Poly a -> Poly a -> Poly a
gcdPoly :: Poly a -> Poly a -> Poly a
gcdPoly a :: Poly a
a b :: Poly a
b = Poly (WrappedNum a) -> Poly a
forall a. Poly (WrappedNum a) -> Poly a
unwrapPoly (WrappedNum a
-> (WrappedNum a -> WrappedNum a -> WrappedNum a)
-> (WrappedNum a -> WrappedNum a -> WrappedNum a)
-> Poly (WrappedNum a)
-> Poly (WrappedNum a)
-> Poly (WrappedNum a)
forall a.
(AdditiveGroup a, Eq a) =>
a -> (a -> a -> a) -> (a -> a -> a) -> Poly a -> Poly a -> Poly a
VS.gcdPolyWith 1 WrappedNum a -> WrappedNum a -> WrappedNum a
forall a. Num a => a -> a -> a
(*) WrappedNum a -> WrappedNum a -> WrappedNum a
forall a. Fractional a => a -> a -> a
(/) (Poly a -> Poly (WrappedNum a)
forall a. Poly a -> Poly (WrappedNum a)
wrapPoly Poly a
a) (Poly a -> Poly (WrappedNum a)
forall a. Poly a -> Poly (WrappedNum a)
wrapPoly Poly a
b))

-- |Normalize a polynomial so that its highest-order coefficient is 1
monicPoly :: (Fractional a, Eq a) => Poly a -> Poly a
monicPoly :: Poly a -> Poly a
monicPoly p :: Poly a
p = Poly (WrappedNum a) -> Poly a
forall a. Poly (WrappedNum a) -> Poly a
unwrapPoly (WrappedNum a
-> (WrappedNum a -> WrappedNum a -> WrappedNum a)
-> Poly (WrappedNum a)
-> Poly (WrappedNum a)
forall a.
(AdditiveGroup a, Eq a) =>
a -> (a -> a -> a) -> Poly a -> Poly a
VS.monicPolyWith 1 WrappedNum a -> WrappedNum a -> WrappedNum a
forall a. Fractional a => a -> a -> a
(/) (Poly a -> Poly (WrappedNum a)
forall a. Poly a -> Poly (WrappedNum a)
wrapPoly Poly a
p))

-- |Compute the derivative of a polynomial.
polyDeriv :: (Num a, Eq a) => Poly a -> Poly a
polyDeriv :: Poly a -> Poly a
polyDeriv p :: Poly a
p = Poly (WrappedNum a) -> Poly a
forall a. Poly (WrappedNum a) -> Poly a
unwrapPoly (Poly (WrappedNum a) -> Poly (WrappedNum a)
forall a. (VectorSpace a, Eq a, Num (Scalar a)) => Poly a -> Poly a
VS.polyDeriv (Poly a -> Poly (WrappedNum a)
forall a. Poly a -> Poly (WrappedNum a)
wrapPoly Poly a
p))

-- |Compute all nonzero derivatives of a polynomial, starting with its 
-- \"zero'th derivative\", the original polynomial itself.
polyDerivs :: (Num a, Eq a) => Poly a -> [Poly a]
polyDerivs :: Poly a -> [Poly a]
polyDerivs p :: Poly a
p = (Poly (WrappedNum a) -> Poly a)
-> [Poly (WrappedNum a)] -> [Poly a]
forall a b. (a -> b) -> [a] -> [b]
map Poly (WrappedNum a) -> Poly a
forall a. Poly (WrappedNum a) -> Poly a
unwrapPoly (Poly (WrappedNum a) -> [Poly (WrappedNum a)]
forall a.
(VectorSpace a, Eq a, Num (Scalar a)) =>
Poly a -> [Poly a]
VS.polyDerivs (Poly a -> Poly (WrappedNum a)
forall a. Poly a -> Poly (WrappedNum a)
wrapPoly Poly a
p))

-- |Compute the definite integral (from 0 to x) of a polynomial.
polyIntegral :: (Fractional a, Eq a) => Poly a -> Poly a
polyIntegral :: Poly a -> Poly a
polyIntegral p :: Poly a
p = Poly (WrappedNum a) -> Poly a
forall a. Poly (WrappedNum a) -> Poly a
unwrapPoly (Poly (WrappedNum a) -> Poly (WrappedNum a)
forall a.
(VectorSpace a, Eq a, Fractional (Scalar a)) =>
Poly a -> Poly a
VS.polyIntegral (Poly a -> Poly (WrappedNum a)
forall a. Poly a -> Poly (WrappedNum a)
wrapPoly Poly a
p))

-- |Separate a nonzero polynomial into a set of factors none of which have
-- multiple roots, and the product of which is the original polynomial.
-- Note that if division is not exact, it may fail to separate roots.  
-- Rational coefficients is a good idea.
--
-- Useful when applicable as a way to simplify root-finding problems.
separateRoots :: (Fractional a, Eq a) => Poly a -> [Poly a]
separateRoots :: Poly a -> [Poly a]
separateRoots p :: Poly a
p
    | Poly a -> Bool
forall a. (Num a, Eq a) => Poly a -> Bool
polyIsZero Poly a
q  = [Char] -> [Poly a]
forall a. HasCallStack => [Char] -> a
error "separateRoots: zero polynomial"
    | Poly a -> Bool
forall a. (Num a, Eq a) => Poly a -> Bool
polyIsOne Poly a
q   = [Poly a
p]
    | Bool
otherwise     = Poly a
p Poly a -> Poly a -> Poly a
forall a. (Fractional a, Eq a) => Poly a -> Poly a -> Poly a
`quotPoly` Poly a
q Poly a -> [Poly a] -> [Poly a]
forall a. a -> [a] -> [a]
: Poly a -> [Poly a]
forall a. (Fractional a, Eq a) => Poly a -> [Poly a]
separateRoots Poly a
q
    where
        q :: Poly a
q = Poly a -> Poly a -> Poly a
forall a. (Fractional a, Eq a) => Poly a -> Poly a -> Poly a
gcdPoly Poly a
p (Poly a -> Poly a
forall a. (Num a, Eq a) => Poly a -> Poly a
polyDeriv Poly a
p)