module Roots where                
import Data.Complex
import Data.List(genericLength)

roots :: RealFloat a => a -> Int -> [Complex a] -> [Complex a]
roots eps count as =
      --
      -- List of complex roots of a polynomial
      -- a0 + a1*x + a2*x^2...
      -- represented by the list as=[a0,a1,a2...]
      -- where
      --     eps is a desired accuracy
      --     count is a maximum count of iterations allowed
      -- Require: list 'as' must have at least two elements
      --     and the last element must not be zero 
      roots' eps count as []
      where
          roots' eps count as xs 
              | length as <= 2  = x:xs
              | otherwise       = 
                  roots' eps count (deflate x bs [last as]) (x:xs)
              where
                  x  = laguerre eps count as 0
                  bs = drop 1 (reverse (drop 1 as))
                  deflate z bs cs
                      | bs == []   = cs
                      | otherwise  = 
                          deflate z (tail bs) (((head bs)+z*(head cs)):cs)


laguerre :: RealFloat a => a -> Int -> [Complex a] -> Complex a -> Complex a       
laguerre eps count as x
      --
      -- One of the roots of the polynomial 'as',
      -- where
      --    eps is a desired accuracy
      --    count is a maximum count of iterations allowed
      --    x is initial guess of the root
      -- This method is due to Laguerre.
      --
      | count <= 0               = x
      | magnitude (x - x') < eps = x'
      | otherwise                = laguerre eps (count - 1) as x'
      where
          x'     = laguerre2 eps as as' as'' x
          as'    = polynomial_derivative as
          as''   = polynomial_derivative as' 
          laguerre2 eps as as' as'' x
              -- One iteration step
              | magnitude b < eps           = x
              | magnitude gp < magnitude gm = 
                  if gm == 0 then x - 1 else x - n/gm
              | otherwise                   = 
                  if gp == 0 then x - 1 else x - n/gp
              where
                  gp    = g + delta
                  gm    = g - delta
                  g     = d/b
                  delta = sqrt ((n-1)*(n*h - g2))
                  h     = g2 - f/b
                  b     = polynomial_value as x
                  d     = polynomial_value as' x
                  f     = polynomial_value as'' x
                  g2    = g^2
                  n     = genericLength as

polynomial_value :: Num a => [a] -> a -> a
polynomial_value as x =
      --
      -- Value of polynomial a0 + a1 x  + a2 x^2 ...
      -- evaluated for 'x',
      -- where 'as' is a list [a0,a1,a2...]
      --
      foldr (u x) 0 as
      where
          u x a b = a + b*x

polynomial_derivative :: Num a => [a] -> [a]
polynomial_derivative as
      --
      -- List of coefficients for derivative of polynomial
      -- a0 + a1 x + a2 x^2 ...
      --
      | as == []  = []
      | otherwise = deriv 1 (drop 1 as) []
      where
          deriv n bs cs
             | bs == []   = reverse2 cs
             | otherwise  = deriv (n+1) (tail bs) ((n*(head bs)):cs) 
          reverse2 cs
              | cs == []  = []
              | otherwise = reverse cs



-----------------------------------------------------------------------------
--
-- Copyright:
--
--      (C) 1998 Numeric Quest Inc., All rights reserved
--
-- Email:
--
--      jans@numeric-quest.com
--
-- License:
--
--      GNU General Public License, GPL
-- 
-----------------------------------------------------------------------------