# L-9 MCS 507 Mon 11 Sep 2023 : laguerre.jl
# To run with redirected output, type 
# julia laguerre.jl > output.txt
# at the command prompt.

using Printf  # for formatted printing of numbers
"""
This simple program illustrates the method of Laguerre
to approximate roots of polynomials in one variable.

To evaluate polynomials, we use evalpoly(x, p),
where x is the value for the input and p[1], p[2], ...
are the coefficients of the polynomial p,
in ascending order of the power of x.  The evalpoly(x, p) 
returns p[1] + p[2]*x + p[3]*x^2 + ... + p[d+1]*x^d,
where d is the degree of p, d = size(p, 1) - 1.
"""

"""
The function diffpoly returns the coefficients of the derivative
of the polynomial with coefficients in c. \\
The zero polynomial is represented as an array of one zero.

EXAMPLE : \\
    c = Array([Complex{Float64}(i,i) for i = 1:3]) \\
    diffpoly(c)
"""
function diffpoly(c::Array{Complex{Float64},1})
    sz = size(c, 1)
    if sz < 2
        result = Array([Complex{Float64}(0)])
    else
        result = Array([Complex{Float64}(0) for _ = 1:sz-1])
        for i=2:sz
            result[i-1] = (i-1)*c[i]
        end
   end
   return result
end

"""
The function laguerre applies the method of Laguerre.

ON ENTRY : (p, d1p, d2p, z0, dxtol, pxtol, maxit, verbose)         \\
  p        are the coefficients of a polynomial in one variable    \\
  d1p      are the coefficients of the first derivative of p       \\
  d2p      are the coefficients of the second derivative of p      \\
  z0       is an approximation for the root                        \\
  dxtol    is the tolerance on the forward error                   \\
  pxtol    is the tolerance on the backward error                  \\
  maxit    is the maximum number of iterations                     \\
  verbose  is the verbose flag, if true, writes one line each step
                                                                 
ON RETURN : (root, absdx, abspx, nbrit, fail)                      \\
  root     is an approximation for the root                        \\
  absdx    is the estimated forward error                          \\
  abspx    is the estimated backward error                         \\
  nbrit    is the number of iterations                             \\
  fail     is true if tolerances not reached,                      \\
           false otherwise.
"""
function laguerre(p::Array{Complex{Float64},1},
                  d1p::Array{Complex{Float64},1},
                  d2p::Array{Complex{Float64},1},
                  z0::Complex{Float64},
                  dxtol::Float64=1.0e-8,
                  pxtol::Float64=1.0e-8,
                  maxit::Int64=20,
                  verbose::Bool=true)
    root = z0; dx = 1; pval = 1
    degm1 = size(p, 1) - 2 # degree of p minus one
    deg = degm1 + 1
    if verbose
        title = "       real(root)             imag(root)"
        println("step : $title          |dx|     |p(x)|")
        stri = @sprintf("%3d", 0)
        strx = @sprintf("%.16e  %.16e", real(root), imag(root))
        println("$stri  : $strx")
    end
    pval = evalpoly(root, p)
    for i=1:maxit
        if abs(pval) < pxtol
            if verbose
                stri = string(i-1)
                println("succeeded after $stri step(s)")
            end
            return (root, abs(dx), abs(pval), i, false)
        end
        d1val = evalpoly(root, d1p)
        d2val = evalpoly(root, d2p)
        Lroot = d1val/pval
        Mroot = Lroot^2 - d2val/pval
        valsqrt = sqrt(degm1*(deg*Mroot - Lroot^2))
        yplus = Lroot + valsqrt
        yminus = Lroot - valsqrt
        if abs(yplus) > abs(yminus)
            dx = deg/yplus
        else
            dx = deg/yminus
        end
        root = root - dx
        pval = evalpoly(root, p)
        if verbose
            stri = @sprintf("%3d", i)
            strx = @sprintf("%.16e  %.16e", real(root), imag(root))
            strdx = @sprintf("%.2e", abs(dx))
            strpx = @sprintf("%.2e", abs(pval))
            println("$stri  : $strx  $strdx  $strpx")
        end
        if abs(dx) < dxtol
            if verbose
                stri = string(i)
                println("succeeded after $stri step(s)")
            end
            return (root, abs(dx), abs(pval), i, false)
        end
    end
    strN = string(maxit)
    if verbose
        println("failed requirements after $strN step(s)")
    end
    return (root, abs(dx), abs(pval), maxit, true)
end

"""
Returns the index of the root in roots that is within tol of point.
"""
function rank(roots::Array{Complex{Float64},1},
              point::Complex{Float64},
              tol::Float64=1.0e-4)
    for i=1:size(roots, 1)
        if(abs(roots[i] - point) < tol)
            return i
        end
    end
    return 0
end

"""
The function runlaguerre demonstrates the method of Laguerre
on the polynomial (x-1)*(x-2) = x^2 - 3*x + 2.
"""
function runlaguerre()
    c1 = Complex{Float64}(2)
    c2 = Complex{Float64}(-3)
    c3 = Complex{Float64}(1)
    p = [c1, c2, c3]
    d1p = diffpoly(p)
    d2p = diffpoly(d1p)
    roots = [Complex{Float64}(1.0), Complex{Float64}(2.0)]
    z0 = Complex{Float64}(rand(),rand())
    rnd = rand()
    if(rnd > 0.5)
        z0 = 3*z0
    end
    println(z0)
    println("running on x^2 - 3*x + 2 ...")
    result = laguerre(p,d1p,d2p,z0)
    srnk = string(rank(roots, result[1]))
    println("the rank : $srnk")
end

function main()
    p = [Complex{Float64}(-1.0), Complex{Float64}(0.0), Complex{Float64}(0.0),
         Complex{Float64}(0.0),  Complex{Float64}(0.0), Complex{Float64}(0.0),
         Complex{Float64}(0.0), Complex{Float64}(0.0), Complex{Float64}(1.0)]
    d1p = diffpoly(p)
    d2p = diffpoly(d1p)
    sq2 = sqrt(2.0)/2.0
    roots = [complex(1), complex(sq2, sq2), complex(0, 1), 
             complex(-sq2, sq2), complex(-1), complex(-sq2, -sq2),
             complex(0, -1), complex(sq2, -sq2)]
    dim = 10001
    left = -1.1
    right = 1.1
    dz = (right - left)/(dim-1)
    mat = zeros(Int8, dim, dim)
    for i=1:dim
        for j=1:dim
            z0 = complex(left + (i-1)*dz, left + (j-1)*dz)
            if(rank(roots, z0, 0.01) != 0)
                mat[i,j] = size(roots, 1)
            else
                result = laguerre(p, d1p, d2p, z0, 1.0e-8, 1.0e-8, 20, false)
                mat[i,j] = rank(roots, result[1]) - 1
            end
        end
    end
    println(dim)
    for i=1:dim
        print(mat[i,1])
        for j=2:dim
            print(",", mat[i,j])
        end
        println()
    end
end

# runlaguerre()
main()