#!/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(ax, q, xq, yq, xi, yi):

    xplot,yplot = compute_field_line(xi, yi, q, xq, yq, +1)
    #print(xplot[0], yplot[0], 'to', xplot[-1], yplot[-1])
    ax.plot(xplot, yplot, 'm-', zorder=1)

    xplot,yplot = compute_field_line(xi, yi, q, xq, yq, -1)
    #print(xplot[0], yplot[0], 'to', xplot[-1], yplot[-1])
    ax.plot(xplot, yplot, 'c-', zorder=1)

#--------------------------------------------------------------------------------

class click_to_call:
    def __init__(self, fig, ax, q, xq, yq, func):
        self.fig = fig
        self.ax = ax
        self.q = q
        self.xq = xq
        self.yq = yq
        self.func = func
        self.cid = self.fig.canvas.mpl_connect('button_press_event', self.click_plot)

    def click_plot(self, event):
        if event.inaxes is not None:
            #print(event)
            #print('start:', event.xdata, event.ydata)
            self.func(self.ax, self.q, self.xq, self.yq, event.xdata, event.ydata)
            self.fig.canvas.draw_idle()

if __name__ == "__main__" :

    # Set up graphics.
    
    fig,ax = plt.subplots(figsize=(6,6))
    ax.set_title('field lines')
    ax.set_xlabel('x')
    ax.set_ylabel('y')
    ax.set_xlim(-5, 5)
    ax.set_ylim(-5, 5)
    ax.set_aspect('equal')

    # Show the charges, red positive.

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

    # Plot field lines through points selected by mouse click.

    cc = click_to_call(fig, ax, q, xq, yq, plot_field_line)
    plt.show()