# P-1 MCS 572 due Fri 27 Sep 2024 :  muller.jl

# A basic implementation of Muller's method in Julia.

using Printf
"""
    function muller(f::Vector{Complex{Float64}},
                    x0::Complex{Float64},
                    x1::Complex{Float64},
                    x2::Complex{Float64},
                    dxtol::Float64=1.0e-8,
                    fxtol::Float64=1.0e-8,
                    N::Int64=10;
                    verbose::Bool=true)

applies Muller's method to approximate a root of a polynomial f.

ON ENTRY :

    f          coefficients of a polynomial
    x0,x1,x2   are three distinct start points
    dxtol      tolerance on the forward error
    fxtol      tolerance on the backward error
    N          maximum number of iterations
    verbose    to print results

ON RETURN : (x, absdx, absfx, nbrit, fail)

    x          approximation for a root
    absdx      estimated forward error
    absfx      estimated backward error
    nbrit      number of iterations
    fail       true if tolerances not reached,
               false otherwise.

EXAMPLE:
    f = rand(6) + rand(6)*complex(0,1)
    x0 = rand(Complex{Float64},1)[1]
    x1 = rand(Complex{Float64},1)[1]
    x2 = rand(Complex{Float64},1)[1]
    (root, err, res, nit, fail) = muller(f, x0, x1, x2)
"""
function muller(f::Vector{Complex{Float64}},
                x0::Complex{Float64},
                x1::Complex{Float64},
                x2::Complex{Float64},
                dxtol::Float64=1.0e-8,
                fxtol::Float64=1.0e-8,
                N::Int64=10;
                verbose::Bool=true)
    if verbose
        println("running Muller's method...")
        titlep1 = "        real(root)            imag(root)      "
        titlep2 = "    |dx|     |f(x)|"
        println("step : $titlep1 $titlep2")
    end
    fx0 = evalpoly(x0, f)
    fx1 = evalpoly(x1, f)
    fx2 = evalpoly(x2, f)
    (root, proot, dx) = (x2, fx2, 1)
    froot = proot
    for i = 1:N 
        # the quadric q passes through (x0, fx0), (x1, fx1), (x2, fx2)
        (h0, h1) = (x1 - x0, x2 - x1)
        (d0, d1) = ((fx1 - fx0)/h0, (fx2 - fx1)/h1)
        # q = a*x^2 + b*x + c
        a = (d1 - d0)/(h1 + h0)
        b = a*h1 + d1
        c = fx2
        # apply the quadratic formula
        rad = sqrt(b*b - 4*a*c)
        # the closest root has the largest denominator
        if(abs(b + rad) > abs(b - rad))
           den = b + rad
        else
           den = b - rad
        end
        dx = -2*c/den
        root = x2 + dx
        froot = evalpoly(root, f)
        if verbose
            stri = @sprintf("%3d", i)
            strreal = @sprintf("%+.16e", real(root))
            strimag = @sprintf("%+.16e", imag(root))
            strdx = @sprintf("%.2e", abs(dx))
            strfx = @sprintf("%.2e", abs(froot))
            println("$stri  : $strreal  $strimag  $strdx  $strfx")
        end
        if((abs(dx) < dxtol) | (abs(proot) < fxtol))
            stri = string(i)
            if verbose
                println("succeeded after $stri steps")
            end
            return (root, abs(dx), abs(froot), i, false)
        end
        (x0, fx0) = (x1, fx1)
        (x1, fx1) = (x2, fx2)
        (x2, fx2) = (root, froot)
    end
    strN = string(N)
    if verbose
        println("failed requirements after $strN steps")
    end
    return (root, abs(dx), abs(froot), N, true)
end

"""
    function test(d::Int, verbose=true)

tests a basic implementation of Muller's method,
on a polynomial with random complex coefficients of degree d.

If not verbose, then only the roots at the end are printed,
otherwise each intermediate value is shown.
"""
function test(d::Int, verbose=true)
    f = rand(d) + rand(d)*complex(0,1)
    x0 = rand(Complex{Float64},1)[1]
    x1 = x0 + 1.0e-4*rand(Complex{Float64},1)[1]
    x2 = x0 + 1.0e-4*rand(Complex{Float64},1)[1]
    result = muller(f, x0, x1, x2; verbose)
    println(" approximation  : ", result[1])
    println("  forward error : ", result[2])
    println(" backward error : ", result[3])
    println("number of steps : ", result[4])
    println("         failed : ", result[5])
end

# test(6)  # test polynomial of degree 6

"""
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

""" 
    function main()

prints the matrix for the basins of attraction of Muller's method.
"""
function main()
    c = [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)]
    p = Vector{Complex{Float64}}(c)
    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{Float64}(left + (i-1)*dz, left + (j-1)*dz)
            if(rank(roots, z0, 0.01) != 0)
                mat[i,j] = size(roots, 1)
            else
                z1 = z0 + 1.0e-4*rand(Complex{Float64},1)[1]
                z2 = z0 + 1.0e-4*rand(Complex{Float64},1)[1]
                # result = muller(p, z0, z1, z2, 1.0e-8, 1.0e-8, 20, false)
                result = muller(p, z0, z1, z2; verbose=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

main()