#!/usr/bin/env python

import sys
import math
import matplotlib.pyplot as plt
import matplotlib.cm as cm

K = 9.0e9
PHI_0 = 1.e6
D_PHI = 1.e2
DS = 0.01
R_MAX = 50

def field(x, y, q, xq, yq):
    Ex = 0.
    Ey = 0.
    
    for i in range(len(q)):

        dx = x - xq[i]
        dy = y - yq[i]
        dr2 = dx*dx + dy*dy
        dr3i = K*q[i]/(dr2*math.sqrt(dr2));

        Ex += dx*dr3i
        Ey += dy*dr3i

    return Ex,Ey

def potential(x, y, q, xq, yq):
    phi = 0.
    for i in range(len(q)):

        dx = x - xq[i]
        dy = y - yq[i]
        dr = math.sqrt(dx*dx + dy*dy)

        phi += K*q[i]/dr

    return phi

def compute_equipotential(xi, yi, q, xq, yq, direction):

    xstart = xi
    ystart = yi
    phistart = potential(xi, yi, q, xq, yq)
    dprev2 = 0.
    dcurr2 = 0.
    dmin2 = 0.01				# require minimum inside dmin
    steps = 0

    xplot = [xi]
    yplot = [yi]

    while xi*xi+yi*yi < R_MAX*R_MAX and steps < 10000:

        Ex,Ey = field(xi, yi, q, xq, yq)
        E = math.sqrt(Ex**2 + Ey**2)
        ds = DS

        dx = -direction*ds*Ey/E			# first order
        dy = direction*ds*Ex/E

        if 1:
            xx = xi + dx
            yy = yi + dy
            Exx,Eyy = field(xx, yy, q, xq, yq)
            Exmean = 0.5*(Ex+Exx)
            Eymean = 0.5*(Ey+Eyy)
            E = math.sqrt(Exmean**2 + Eymean**2)

            dx = -direction*ds*Eymean/E		# second order
            dy = direction*ds*Exmean/E

        xi += dx
        yi += dy

        if 1:
            dphi = potential(xi, yi, q, xq, yq) - phistart
            if dphi != 0:
                Ex,Ey = field(xi, yi, q, xq, yq)
                E2 = Ex**2 + Ey**2
                xi += dphi*Ex/E2		# corrector
                yi += dphi*Ey/E2

        # Stop at the first minimum in |x-xstart|.

        dnext2 = (xi-xstart)**2 + (yi-ystart)**2
        if dcurr2 < dmin2 and dcurr2 < dprev2 and dnext2 > dcurr2: break

        dprev2 = dcurr2
        dcurr2 = dnext2
        steps += 1

        xplot.append(xi)
        yplot.append(yi)

    return xplot,yplot

def plot_equipotential(xi, yi, q, xq, yq, c):

    xplot,yplot = compute_equipotential(xi, yi, q, xq, yq,  1)
    plt.plot(xplot, yplot, color=c, zorder=1)

    if xi == 0:
        xplot,yplot = compute_equipotential(xi, yi, q, xq, yq, -1)
        plt.plot(xplot, yplot, color=c, zorder=1)

import numpy as np
    
def setup_charges():

    # Use numpy arrays so we can use xq[(q>=0)] syntax below.
    
    q = np.zeros(2)
    xq = np.zeros(2)
    yq = np.zeros(2)

    q[0] = 1.e-6
    q[1] = -1.e-6

    xq[0] = -1.
    xq[1] = 1.

    return q, xq, yq

if __name__ == "__main__" :

    nx = 19
    lim = 4.
    if len(sys.argv) > 1: nx = int(sys.argv[1])
    if len(sys.argv) > 2: lim = float(sys.argv[2])

    fig,ax = plt.subplots()
    plt.xlabel('x')
    plt.ylabel('y')
    plt.title('equipotentials')
    plt.xlim(-lim, lim)
    plt.ylim(-lim, lim)
    ax.set_aspect('equal')

    q, xq, yq = setup_charges()
    plt.scatter(xq[(q>=0)], yq[(q>=0)], c='r', s=10, zorder=2)
    plt.scatter(xq[(q<0)], yq[(q<0)], c='b', s=10, zorder=2)

    for ix in range(nx):
        x = -0.1*ix + 0.9
        c = cm.rainbow(float(ix)/nx, 1)
        print(cm.hsv)
        print(c)
        plot_equipotential(x, 0, q, xq, yq, c)

    plt.show()