-----------------------------------------------------------------------------
-- |
-- Module  :  ForSyDe.Shallow.Utility.PolyArith
-- Copyright   :  (c) ForSyDe Group, KTH 2007-2008
-- License     :  BSD-style (see the file LICENSE)
-- 
-- Maintainer  :  forsyde-dev@ict.kth.se
-- Stability   :  experimental
-- Portability :  portable
--
-- This is the polynomial arithematic library. The arithematic operations include 
-- addition, multiplication, division and power. However, the computation time is 
-- not optimized for multiplication and is O(n2), which could be considered to be 
-- optimized by FFT algorithms later on.
-----------------------------------------------------------------------------
module ForSyDe.Shallow.Utility.PolyArith(
      -- *Polynomial data type
      Poly(..),
      -- *Addition, DmMultiplication, division and power operations
      addPoly, mulPoly, divPoly, powerPoly,
      -- *Some helper functions
      getCoef, scalePoly, addPolyCoef, subPolyCoef, scalePolyCoef
    )
    where 

-- |Polynomial data type.
data Poly a = Poly [a]
            | PolyPair (Poly a, Poly a) deriving (Poly a -> Poly a -> Bool
(Poly a -> Poly a -> Bool)
-> (Poly a -> Poly a -> Bool) -> Eq (Poly a)
forall a. Eq a => Poly a -> Poly a -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Poly a -> Poly a -> Bool
$c/= :: forall a. Eq a => Poly a -> Poly a -> Bool
== :: Poly a -> Poly a -> Bool
$c== :: forall a. Eq a => Poly a -> Poly a -> Bool
Eq)


-- |Multiplication operation of polynomials.
mulPoly :: Num a => Poly a -> Poly a -> Poly a
mulPoly :: Poly a -> Poly a -> Poly a
mulPoly (Poly []) Poly a
_ = [a] -> Poly a
forall a. [a] -> Poly a
Poly []
mulPoly Poly a
_ (Poly []) = [a] -> Poly a
forall a. [a] -> Poly a
Poly []
-- Here is the O(n^2) version of polynomial multiplication
mulPoly (Poly [a]
xs) (Poly [a]
ys) = [a] -> Poly a
forall a. [a] -> Poly a
Poly ([a] -> Poly a) -> [a] -> Poly a
forall a b. (a -> b) -> a -> b
$ (a -> [a] -> [a]) -> [a] -> [a] -> [a]
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (\a
y [a]
zs ->
  let (a
v:[a]
vs) = a -> [a] -> [a]
forall a. Num a => a -> [a] -> [a]
scalePolyCoef a
y [a]
xs in a
v a -> [a] -> [a]
forall a. a -> [a] -> [a]
:[a] -> [a] -> [a]
forall a. Num a => [a] -> [a] -> [a]
addPolyCoef [a]
vs [a]
zs) [] [a]
ys
mulPoly (PolyPair (Poly a
a, Poly a
b)) (PolyPair (Poly a
c, Poly a
d)) =
  (Poly a, Poly a) -> Poly a
forall a. (Poly a, Poly a) -> Poly a
PolyPair (Poly a -> Poly a -> Poly a
forall a. Num a => Poly a -> Poly a -> Poly a
mulPoly Poly a
a Poly a
c, Poly a -> Poly a -> Poly a
forall a. Num a => Poly a -> Poly a -> Poly a
mulPoly Poly a
b Poly a
d)
mulPoly (PolyPair (Poly a
a, Poly a
b)) (Poly [a]
c) =
  (Poly a, Poly a) -> Poly a
forall a. (Poly a, Poly a) -> Poly a
PolyPair (Poly a -> Poly a -> Poly a
forall a. Num a => Poly a -> Poly a -> Poly a
mulPoly Poly a
a ([a] -> Poly a
forall a. [a] -> Poly a
Poly [a]
c), Poly a
b)
mulPoly (Poly [a]
c) (PolyPair (Poly a
a, Poly a
b)) =
  Poly a -> Poly a -> Poly a
forall a. Num a => Poly a -> Poly a -> Poly a
mulPoly ((Poly a, Poly a) -> Poly a
forall a. (Poly a, Poly a) -> Poly a
PolyPair (Poly a
a, Poly a
b)) ([a] -> Poly a
forall a. [a] -> Poly a
Poly [a]
c)

-- |Division operation of polynomials.
divPoly :: Num a => Poly a -> Poly a -> Poly a
divPoly :: Poly a -> Poly a -> Poly a
divPoly (Poly [a]
a) (Poly [a]
b) = (Poly a, Poly a) -> Poly a
forall a. (Poly a, Poly a) -> Poly a
PolyPair ([a] -> Poly a
forall a. [a] -> Poly a
Poly [a]
a,[a] -> Poly a
forall a. [a] -> Poly a
Poly [a]
b)
divPoly (PolyPair (Poly a
a, Poly a
b)) (PolyPair (Poly a
c, Poly a
d)) =
  Poly a -> Poly a -> Poly a
forall a. Num a => Poly a -> Poly a -> Poly a
mulPoly ((Poly a, Poly a) -> Poly a
forall a. (Poly a, Poly a) -> Poly a
PolyPair (Poly a
a, Poly a
b)) ((Poly a, Poly a) -> Poly a
forall a. (Poly a, Poly a) -> Poly a
PolyPair (Poly a
d, Poly a
c))
divPoly (PolyPair (Poly a
a, Poly a
b)) (Poly [a]
c) =
  (Poly a, Poly a) -> Poly a
forall a. (Poly a, Poly a) -> Poly a
PolyPair (Poly a
a, Poly a -> Poly a -> Poly a
forall a. Num a => Poly a -> Poly a -> Poly a
mulPoly Poly a
b ([a] -> Poly a
forall a. [a] -> Poly a
Poly [a]
c))
divPoly (Poly [a]
c) (PolyPair (Poly a
a, Poly a
b)) =
  (Poly a, Poly a) -> Poly a
forall a. (Poly a, Poly a) -> Poly a
PolyPair (Poly a -> Poly a -> Poly a
forall a. Num a => Poly a -> Poly a -> Poly a
mulPoly Poly a
b ([a] -> Poly a
forall a. [a] -> Poly a
Poly [a]
c), Poly a
a)

-- |Addition operations of polynomials.
addPoly :: (Num a, Eq a) => Poly a -> Poly a -> Poly a
addPoly :: Poly a -> Poly a -> Poly a
addPoly (Poly [a]
a) (Poly [a]
b) = [a] -> Poly a
forall a. [a] -> Poly a
Poly ([a] -> Poly a) -> [a] -> Poly a
forall a b. (a -> b) -> a -> b
$ [a] -> [a] -> [a]
forall a. Num a => [a] -> [a] -> [a]
addPolyCoef [a]
a [a]
b
addPoly (PolyPair (Poly a
a, Poly a
b)) (PolyPair (Poly a
c, Poly a
d)) =
    if Poly a
bPoly a -> Poly a -> Bool
forall a. Eq a => a -> a -> Bool
==Poly a
d then  -- simplifyPolyPair $
      (Poly a, Poly a) -> Poly a
forall a. (Poly a, Poly a) -> Poly a
PolyPair (Poly a -> Poly a -> Poly a
forall a. (Num a, Eq a) => Poly a -> Poly a -> Poly a
addPoly Poly a
a Poly a
c, Poly a
d)
    else  -- simplifyPolyPair $
      (Poly a, Poly a) -> Poly a
forall a. (Poly a, Poly a) -> Poly a
PolyPair (Poly a
dividedPoly, Poly a
divisorPoly)
  where
    divisorPoly :: Poly a
divisorPoly = if Poly a
b Poly a -> Poly a -> Bool
forall a. Eq a => a -> a -> Bool
==Poly a
d then Poly a
b else Poly a -> Poly a -> Poly a
forall a. Num a => Poly a -> Poly a -> Poly a
mulPoly Poly a
b Poly a
d
    dividedPoly :: Poly a
dividedPoly = if Poly a
b Poly a -> Poly a -> Bool
forall a. Eq a => a -> a -> Bool
== Poly a
d then Poly a -> Poly a -> Poly a
forall a. (Num a, Eq a) => Poly a -> Poly a -> Poly a
addPoly Poly a
a Poly a
c
      else Poly a -> Poly a -> Poly a
forall a. (Num a, Eq a) => Poly a -> Poly a -> Poly a
addPoly (Poly a -> Poly a -> Poly a
forall a. Num a => Poly a -> Poly a -> Poly a
mulPoly Poly a
a Poly a
d) (Poly a -> Poly a -> Poly a
forall a. Num a => Poly a -> Poly a -> Poly a
mulPoly Poly a
b Poly a
c)
addPoly (Poly [a]
a) (PolyPair (Poly a
c, Poly a
d) ) =
    Poly a -> Poly a -> Poly a
forall a. (Num a, Eq a) => Poly a -> Poly a -> Poly a
addPoly ((Poly a, Poly a) -> Poly a
forall a. (Poly a, Poly a) -> Poly a
PolyPair (Poly a
multiPolyHelper, Poly a
d)) ((Poly a, Poly a) -> Poly a
forall a. (Poly a, Poly a) -> Poly a
PolyPair (Poly a
c,Poly a
d) )
  where
    multiPolyHelper :: Poly a
multiPolyHelper = Poly a -> Poly a -> Poly a
forall a. Num a => Poly a -> Poly a -> Poly a
mulPoly ([a] -> Poly a
forall a. [a] -> Poly a
Poly [a]
a) Poly a
d
addPoly  abPoly :: Poly a
abPoly@(PolyPair (Poly a, Poly a)
_) cPoly :: Poly a
cPoly@(Poly [a]
_) = Poly a -> Poly a -> Poly a
forall a. (Num a, Eq a) => Poly a -> Poly a -> Poly a
addPoly Poly a
cPoly Poly a
abPoly
 
-- |Power operation of polynomials.
powerPoly :: Num a => Poly a -> Int -> Poly a
powerPoly :: Poly a -> Int -> Poly a
powerPoly Poly a
p Int
n = Poly a -> Poly a -> Int -> Poly a
forall a. Num a => Poly a -> Poly a -> Int -> Poly a
powerX' ([a] -> Poly a
forall a. [a] -> Poly a
Poly [a
1]) Poly a
p Int
n
  where
    powerX' :: Num a => Poly a -> Poly a -> Int -> Poly a
    powerX' :: Poly a -> Poly a -> Int -> Poly a
powerX' Poly a
p' Poly a
_ Int
0 = Poly a
p'
    powerX' Poly a
p' Poly a
p Int
n = Poly a -> Poly a -> Int -> Poly a
forall a. Num a => Poly a -> Poly a -> Int -> Poly a
powerX' (Poly a -> Poly a -> Poly a
forall a. Num a => Poly a -> Poly a -> Poly a
mulPoly Poly a
p' Poly a
p) Poly a
p (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1)

-- |Some helper functions below.

-- |To get the coefficients of the polynomial.
getCoef :: Num a => Poly a -> ([a],[a])
getCoef :: Poly a -> ([a], [a])
getCoef (Poly [a]
xs) = ([a]
xs,[a
1])
getCoef (PolyPair (Poly [a]
xs,Poly [a]
ys)) = ([a]
xs,[a]
ys)
getCoef Poly a
_ = [Char] -> ([a], [a])
forall a. HasCallStack => [Char] -> a
error [Char]
"getCoef: Nested fractions found"

scalePoly :: (Num a) => a -> Poly a -> Poly a
scalePoly :: a -> Poly a -> Poly a
scalePoly a
s Poly a
p = Poly a -> Poly a -> Poly a
forall a. Num a => Poly a -> Poly a -> Poly a
mulPoly ([a] -> Poly a
forall a. [a] -> Poly a
Poly [a
s]) Poly a
p

addPolyCoef :: Num a => [a] -> [a] -> [a]
addPolyCoef :: [a] -> [a] -> [a]
addPolyCoef = (a, a) -> (a -> a -> a) -> [a] -> [a] -> [a]
forall a b c. (a, b) -> (a -> b -> c) -> [a] -> [b] -> [c]
zipWithExt (a
0,a
0) a -> a -> a
forall a. Num a => a -> a -> a
(+)
subPolyCoef :: RealFloat a => [a] -> [a] -> [a]
subPolyCoef :: [a] -> [a] -> [a]
subPolyCoef = (a, a) -> (a -> a -> a) -> [a] -> [a] -> [a]
forall a b c. (a, b) -> (a -> b -> c) -> [a] -> [b] -> [c]
zipWithExt (a
0,a
0) (-)

scalePolyCoef :: (Num a) => a -> [a] -> [a]
scalePolyCoef :: a -> [a] -> [a]
scalePolyCoef a
s [a]
p = (a -> a) -> [a] -> [a]
forall a b. (a -> b) -> [a] -> [b]
map (a
sa -> a -> a
forall a. Num a => a -> a -> a
*) [a]
p

-- |Extended version of 'zipWith', which will add zeros to the shorter list.
zipWithExt :: (a,b) -> (a -> b -> c) -> [a] -> [b] -> [c]
zipWithExt :: (a, b) -> (a -> b -> c) -> [a] -> [b] -> [c]
zipWithExt (a, b)
_ a -> b -> c
_ [] [] = []
zipWithExt (a
x0,b
y0) a -> b -> c
f (a
x:[a]
xs) [] = a -> b -> c
f a
x b
y0 c -> [c] -> [c]
forall a. a -> [a] -> [a]
: ((a, b) -> (a -> b -> c) -> [a] -> [b] -> [c]
forall a b c. (a, b) -> (a -> b -> c) -> [a] -> [b] -> [c]
zipWithExt (a
x0,b
y0) a -> b -> c
f [a]
xs [])
zipWithExt (a
x0,b
y0) a -> b -> c
f [] (b
y:[b]
ys)  = a -> b -> c
f a
x0 b
y c -> [c] -> [c]
forall a. a -> [a] -> [a]
: ((a, b) -> (a -> b -> c) -> [a] -> [b] -> [c]
forall a b c. (a, b) -> (a -> b -> c) -> [a] -> [b] -> [c]
zipWithExt (a
x0,b
y0) a -> b -> c
f [] [b]
ys)
zipWithExt (a
x0,b
y0) a -> b -> c
f (a
x:[a]
xs) (b
y:[b]
ys)  = a -> b -> c
f a
x b
y c -> [c] -> [c]
forall a. a -> [a] -> [a]
: ((a, b) -> (a -> b -> c) -> [a] -> [b] -> [c]
forall a b c. (a, b) -> (a -> b -> c) -> [a] -> [b] -> [c]
zipWithExt (a
x0,b
y0) a -> b -> c
f [a]
xs [b]
ys)