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

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

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

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

def color(phi):
    phiscale = 0.1*PHI_0
    if phi >= 0:
        phis = max(1.e-3, min(1., phi/phiscale))
        cc = 1 + 0.5*math.log10(phis)/3				# 0.5 <= cc <= 1
    else:
        phis = max(1.e-3, min(1., -phi/phiscale))
        cc = -0.5*math.log10(phis)/3				# 0 <= cc < 0.5

    #print('phis =', phis, 'cc =', cc)
    return cm.rainbow(cc)

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

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:
            x = event.xdata
            y = event.ydata
            pot = potential(x, y, self.q, self.xq, self.yq)
            print('start:', x, y, pot)
            c = color(pot)
            self.func(self.ax, self.q, self.xq, self.yq, c, x, y)
            self.fig.canvas.draw_idle()

if __name__ == "__main__" :

    n = 5
    s = 42
    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: lim = float(sys.argv[3])

    # Set up graphics.
    
    fig,ax = plt.subplots(figsize=(6,6))
    ax.set_title('equipotentials')
    ax.set_xlabel('x')
    ax.set_ylabel('y')
    plim = 1.5*lim
    ax.set_xlim(-plim, plim)
    ax.set_ylim(-plim, plim)
    ax.set_aspect('equal')

    # Show the charges, red positive.
    
    q, xq, yq = setup_charges(n, s)
    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 equipotentials selected by mouse click.

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