using LinearAlgebra # for the norm(x⃗) function = ‖x⃗‖ = √x⃗ᵀx⃗

function forces(θ⃗, q⃗)
    x⃗₀ = [-1, 0]  # q₀ is fixed at (-1,0)
    q₀ = 1
    F⃗ = map(1:length(θ⃗)) do i
        x⃗ᵢ = [cos(θ⃗[i]), sin(θ⃗[i])]  # position of charge i on unit circle

        # inverse-square force law: F⃗ = qᵢqⱼ r̂ᵢⱼ / rᵢⱼ² = qᵢqⱼ r⃗ᵢⱼ / rᵢⱼ³
        
        # force from x⃗₀:
        F⃗ᵢ = (x⃗ᵢ - x⃗₀) * (q⃗[i] * q₀ / norm(x⃗ᵢ - x⃗₀)^3)

        for j = 1:length(θ⃗)  # add forces from other charges
            if i ≠ j
                x⃗ⱼ = [cos(θ⃗[j]), sin(θ⃗[j])]
                F⃗ᵢ += (x⃗ᵢ - x⃗ⱼ) * (q⃗[i] * q⃗[j] / norm(x⃗ᵢ - x⃗ⱼ)^3)
            end
        end
        
        # construct force component in tangential direction
        F⃗ᵢ ⋅ [-x⃗ᵢ[2], x⃗ᵢ[1]]  # dot product with tangent vector [-sin, cos]
    end
    return F⃗
end

forces([1,2], [1,1])

using ForwardDiff

ForwardDiff.jacobian(θ⃗ -> forces(θ⃗, [1,1]), [1,2])

θ⃗ = [1,2] # initial "guess"
q⃗ = [1,1]

for iteration = 1:8
    @show θ⃗ .* (180/π)
    @show f⃗ = forces(θ⃗, q⃗)
    J = ForwardDiff.jacobian(θ⃗ -> forces(θ⃗, q⃗), θ⃗)
    θ⃗ = θ⃗ - J \ f⃗
    println() # print a blank line
end

# show the final values:
@show θ⃗ .* (180/π)
@show f⃗ = forces(θ⃗, q⃗);

using PyPlot

function plotcharges(θ⃗)
    gca().add_patch(plt.Circle((0,0), 1.0, clip_on=false, fill=false))
    plot(-1, 0, "ko")
    text(-1 + 0.05, 0, "q₀")
    axis("equal")
    for (i,θ) in enumerate(θ⃗)
        plot(cos(θ), sin(θ), "o")
        text(cos(θ)+0.1, sin(θ), "q" * ('₀' + i))
    end
    xlabel(L"x")
    ylabel(L"y")
end

plotcharges(θ⃗)
title(L"equilibrium distribution for $N=3$ equal charges")

using Interact

fig = figure()
@manipulate for iterations = slider(0:16, value=0, label="Newton iterations"),
                q₁ in slider([0.1,1,2,4,8], value=1, label="q₁"),
                q₂ in slider([0.1,1,2,4,8], value=1, label="q₂"),
                N in slider(2:10, value=3, label="Number of particles")
    n = N-1
    θ⃗ = [1:n;] ./ n
    q⃗ = ones(n)
    q⃗[1] = q₁
    if n > 1
        q⃗[2] = q₂
    end
    for iteration = 1:iterations
        f⃗ = forces(θ⃗, q⃗)
        J = ForwardDiff.jacobian(θ⃗ -> forces(θ⃗, q⃗), θ⃗)
        θ⃗ = θ⃗ - J \ f⃗
    end
    f⃗ = forces(θ⃗, q⃗)
    withfig(fig) do
        plotcharges(θ⃗)
        title("$N particles, forces = $(round.(f⃗, sigdigits=3))")
    end
end