using LinearAlgebra
using Printf
using Plots
using LaTeXStrings

function UNM(f;
        x::Union{Nothing, Real}=nothing,
        eps::Real=1e-6,
        finf::Real=1e+8,
        MaxFeval::Integer=30,
        plt::Union{Plots.Plot, Nothing}=nothing,
        plotatend::Bool=true,
        Plotg::Integer=0
    )::Tuple{Real, String}

    # Plotg
    # 1 = the function value / gap are plotted
    # 2 = the function and the second-order model are plotted
    # 3 = the function, the first-order model and the second-order model are
    #     plotted (the first-order model has no role in the algorithm, but
    #     this shows how much better the second-order model is than the
    #     first-order one)

    resolution = 1000

    Interactive = false

    # reading and checking input- - - - - - - - - - - - - - - - - - - - - - - -
    # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -

    (fStar, _, _, rangeplt) = f(nothing)

    if x == nothing
        x = rangeplt[1]
    end

    if finf ≤ 0
        throw(DomainError(finf, "finf must be in > 0"))
    end

    if MaxFeval < 2
        throw(ArgumentError("At least two function computations are required"))
        return (NaN, "empty")
    end

    # initializations - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -

    feval = 0
    status = "optimal"
    fbest = Inf

    if Plotg == 1
        gap = []
    end

    if Plotg == 1 && plt == nothing
        plt = plot(yscale = :log,
                   xlims=(0, 35),
                   ylims=(1e-15, Inf),
                   guidefontsize=16)
    elseif Plotg == 2 && plt == nothing
        plt = plot(legend = false)
    elseif Plotg == 3 && plt == nothing
        plt = plot(legend = false)
    end


    println("Univariate Newton's Method")

    if fStar > -Inf
        println("feval\trel gap\t\tx\t\tf(x)\t\tf'(x)\t\tf''(x)")
    else
        println("feval\tfbest\t\tx\t\tf'(x)\t\tf'(x)\t\tf''(x)")
    end

    # main loop - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -

    while true
        # compute f(x), f'(x), f''(x) - - - - - - - - - - - - - - - - - - - - -

        (fx, f1x, f2x) = f(x)

        feval += 1

        if fx < fbest
            fbest = fx
        end

        if abs(f2x) ≤ 1e-16
            status = "stopped"
            @printf("numerical issue: f''(x) = %1.4e\n", f2x)
            break
        end

        # main logic- - - - - - - - - - - - - - - - - - - - - - - - - - - - - -

        xn = x - f1x/f2x

        # output statistics - - - - - - - - - - - - - - - - - - - - - - - - - -

        if fStar > -Inf
            gapk = (fbest - fStar) / max(abs(fStar), 1)

            if Plotg == 1
                push!(gap, gapk)
            end
        else
            gapk = fbest
        end

        @printf("%4d\t%1.4e\t%1.8e\t%1.4e\t%1.4e\t%1.4e\n", feval, gapk, x, fx, f1x, f2x)

        if Plotg > 1
            xm = min(x, xn) - abs(x - xn) / 5
            xp = max(x, xn) + abs(x - xn) / 5

            xx = range(xm, xp, resolution)
            yy = map(v -> v[1], f.(xx))

            for e in plt.series_list
                e[:linealpha] = 0
            end

            old_ylims = ylims(plt)
            ylims!(plt, (minimum(yy), max(yy[1], yy[end])))

            xlims!(plt, (xm, xp))

            plot!(plt, xx, yy)

            if Plotg == 3
                # first-order model is
                # f( y ) = f( x ) + f'( x )( y - x )
                #        = [ f( x ) - f'( x ) x ]
                #        + f'( x ) y
                b = f1x
                c = fx - f1x*x

                yy = b .* xx .+ c
                plot!(plt, xx, yy)
            end
            # second-order model is
            # f( y ) = f( x ) + f'( x )( y - x ) + f''( x )( y - x )^2 / 2
            #        = [ f( x ) - f'( x ) x + f''( x ) x^2 / 2 ]
            #        + [ f'( x ) - f''( x ) x ] y
            #        + f''( x )y^2 / 2
            a = f2x/2
            b = f1x - f2x*x
            c = fx - f1x*x + f2x*x^2 / 2

            yy = @. a * xx^2 + b * xx + c
            plot!(plt, xx, yy)

            new_xticks = []
            if x < xn
                new_xticks = ([xm, x, xn, xp],
                    [@sprintf("%1.1g", xm), L"x^k", L"x^{k+1}", @sprintf("%1.1g", xp)])
            else
                new_xticks = ([xm, xn, x, xp],
                    [@sprintf("%1.1g", xm), L"x^{k+1}", L"x^k", @sprintf("%1.1g", xp)])
            end
            xticks!(plt, new_xticks)
        end

        # stopping criteria - - - - - - - - - - - - - - - - - - - - - - - - - -

        if abs(f1x) ≤ eps
            break
        end

        if feval > MaxFeval
            status = "stopped"
            break
        end

        if abs(fx) > finf
            status = "error"
            break
        end

        if Plotg ≠ 0 && Interactive
            IJulia.clear_output(wait=true)
            display(plt)
            sleep(0.1)
            readline()
        end

        # iterate - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -

        x = xn
    end
    
    if Plotg == 1
        plot!(plt,
            gap,
            linewidth=2)
    end

    if plotatend && Plotg ≠ 0
        display(plt)
    end

    (x, status)
end