#!/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)

import numpy as np

def setup_charges(n, s):

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

    if s != 0: np.random.seed(s)
    
    q = np.random.uniform(-1.e-6, 1.e-6, n)
    lim = 2
    xq = np.random.uniform(-lim, lim, n)
    yq = np.random.uniform(-lim, lim, n)

    if 1:
        qmean = q.sum()/n
        q -= qmean		# force zero total charge

    return q, xq, yq

if __name__ == "__main__" :

    n = 5
    s = 42
    nx = 3
    lim = 4.
    if len(sys.argv) > 1: n = int(sys.argv[1])
    if len(sys.argv) > 2: s = int(sys.argv[2])
    if len(sys.argv) > 3: nx = int(sys.argv[3])
    if len(sys.argv) > 4: lim = float(sys.argv[4])

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

    q, xq, yq = setup_charges(n, s)
    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)

    # Draw equipotentials crossing the specified lines.

    xmin = -lim
    xmax = lim
    dx = (xmax-xmin)/nx

    for ix in range(nx+1):
        x = xmin + ix*dx
        c = cm.rainbow(float(2*ix)/(2*nx+1), 1)
        plot_equipotential(x, -0.5, q, xq, yq, c)
        c = cm.rainbow(float(2*ix+1)/(2*nx+1), 1)
        plot_equipotential(x,  0.5, q, xq, yq, c)

    plt.show()