# Pkg.add.(["Interact", "PyPlot", "ForwardDiff"]) # uncomment this line to install packages
using Interact, PyPlot

fig = figure()
xs = linspace(-5,5,1000)
@manipulate for step in slider(1:20, value=1), start in slider(-4:0.1:4, value=-0.1)
    withfig(fig) do
        x = start
        local xprev, f, f′
        for i = 1:step
            xprev = x
            f = 2cos(x) - x + x^2/10
            f′ = -2sin(x) - 1 + 2x/10
            x = x - f/f′
        end
        plot(xs, 0*xs, "k-")
        plot(xs, 2cos.(xs) - xs + xs.^2/10, "b-")
        newf = 2cos(x) - x + x^2/10
        plot([xprev, x], [f, 0.0], "ro")
        plot(x, newf, "bo")
        plot(xs, f + f′ * (xs - xprev), "r--")
        plot([x, x], [0, newf], "k--")
        xlim(minimum(xs), maximum(xs))
        ylim(-5,4)
        title("Newton step $step: f($x) = $newf")
    end
end

fₑₓ(x) = [ sin(x[1]*x[2]) - 0.5
           x[1]^2 - x[2]^2 ]

# manual Jacobian
Jₑₓ(x) = [ x[2]*cos(x[1]*x[2])  x[1]*cos(x[1]*x[2])
         2x[1]               -2x[2]               ]

fₑₓ([2,3])

Jₑₓ([2,3]) # manual Jacobian

using ForwardDiff

ForwardDiff.jacobian(fₑₓ, [2,3]) # automatic Jacobian

ForwardDiff.jacobian(fₑₓ, [2,3]) - Jₑₓ([2,3])

figure(figsize=(10,4))

subplot(1,2,1)
X = linspace(-3,3,250)
pcolor(X', X, [-log(0.01 + norm(fₑₓ([x,y]))) for y in X, x in X], cmap="gray")
xlabel(L"x_1")
ylabel(L"x_2")
title(L"-\log [ 0.01 + \Vert f_\mathrm{ex}(x) \Vert]")
colorbar()

subplot(1,2,2)
x = linspace(0,3,200)
plot(x, sin.(x.^2) .- 0.5, "r-")
plot(x, 0*x, "k-")
plot([sqrt(pi/6), sqrt(5pi/6), sqrt(13pi/6)], [0,0,0], "ro", markersize=8)
text(0.65, -0.23, L"\sqrt{\pi/6}")
title(L"\sin(x^2) - \frac{1}{2}")
xlabel(L"x = x_1 = x_2")
grid()

sqrt(π/6)

x = [0.5, 0.8]
for i = 1:10
    x = x - Jₑₓ(x) \ fₑₓ(x)  # Newton step
    println("$i Newton steps: x = $x, f(x) = ", fₑₓ(x))
end

A = [ -1   0   0   1   0   0
       0   0   0  -1   1   0
       0   0   0   0  -1   1
       0   0   1   0   0  -1
       0   1  -1   0   0   0
       1  -1   0   0   0   0
       0  -1   0   0   0   1
       0   0   0  -1   0   1 ]

Yₖ = 1
αₖ = 0.5
Ỹₖ(d) = Yₖ / (1 + αₖ * d^2)

f(v) = A' * (diagm(Ỹₖ.(A*v)) * (A*v)) - [1,-1,0,0,0,0]
J(v) = ForwardDiff.jacobian(f, v)

newtonstep(v) = v - pinv(J(v)) * f(v)

function newton(v, nsteps)
    for i = 1:nsteps
        v = newtonstep(v)
    end
    return v
end
newton(v, nsteps::AbstractVector) = map(n -> newton(v,n), nsteps)

M₀ = A' * diagm(Ỹₖ.(zeros(8))) * A # matrix AᵀYA for zero voltage drops (i.e. linear problem)
v₀ = pinv(M₀) * [1,-1,0,0,0,0]  # initial guess = solution to linear problem
vsteps = newton(v₀, 0:10)
plot([vsteps[i][j] for i=1:length(vsteps), j=1:6], "o-")
xlabel("Newton steps")
ylabel("node voltages")
legend([L"v_1", L"v_2", L"v_3", L"v_4", L"v_5", L"v_6"], loc="upper right")
title(L"Newton steps for $\alpha = 0.5$")

xlabel("Newton steps")
ylabel(L"\Vert f(v) \Vert")
semilogy(norm.(f.(vsteps)), "bo-")
title(L"Newton convergence for $\alpha = 0.5$")

Ỹₖ′(d) = -2αₖ*d*Yₖ / (1 + αₖ * d^2)^2
Kₖ(d) = Ỹₖ(d) + Ỹₖ′(d)*d
J_manual(v) = A' * diagm(Kₖ.(A*v)) * A

v = [0.1,1.2,3.4,5.6,-0.3,0.7] # a "random" vector
ForwardDiff.jacobian(f, v)

J_manual(v)

ForwardDiff.jacobian(f, v) - J_manual(v)