#!/usr/bin/env python

import sys
from math import sqrt
import matplotlib.pyplot as plt

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*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 = sqrt(dx*dx + dy*dy)

        phi += K*q[i]/dr

    return phi

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

    xplot = [xi]
    yplot = [yi]

    while xi*xi+yi*yi < R_MAX*R_MAX \
            and abs(potential(xi, yi, q, xq, yq)) < PHI_0:

        Ex,Ey = field(xi, yi, q, xq, yq)

        E = sqrt(Ex*Ex + Ey*Ey)
        ds = DS
        if ds > D_PHI/E: ds = D_PHI/E

        xi += direction*ds*Ex/E
        yi += direction*ds*Ey/E

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

    return xplot,yplot

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

def plot_field_line(xi, yi, q, xq, yq):

    xplot,yplot = compute_field_line(xi, yi, q, xq, yq, +1)
    plt.plot(xplot, yplot, 'm-', label='$+$', zorder=1)

    xplot,yplot = compute_field_line(xi, yi, q, xq, yq, -1)
    plt.plot(xplot, yplot, 'c-', label='$-$', zorder=1)
    
if __name__ == "__main__" :

    xi = 2.0
    yi = 1.0
    if len(sys.argv) > 1: xi = float(sys.argv[1])
    if len(sys.argv) > 2: yi = float(sys.argv[2])

    q, xq, yq = setup_charges()

    # Set up the graphics.

    fig,ax = plt.subplots()
    plt.xlabel('x')
    plt.ylabel('y')
    plt.title('field line')
    plt.xlim(-4.5, 4.5)
    plt.ylim(-0.5, 8.5)
    ax.set_aspect('equal')

    # Show the charges and the starting point.
    
    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)
    plt.scatter(xi, yi, facecolors='none', edgecolors='green',
                marker='o', s=50, zorder=2)

    # Plot the field line.
    
    plot_field_line(xi, yi, q, xq, yq)

    plt.legend()
    plt.show()