def test_masks_and_nans():
    X, Y, U, V = velocity_field()
    mask = np.zeros(U.shape, dtype=bool)
    mask[40:60, 40:60] = 1
    U = np.ma.array(U, mask=mask)
    U[:20, :20] = np.nan
    plt.streamplot(X, Y, U, V, color=U, cmap=plt.cm.Blues)

def quadrupole(a):
    """Function: quadrupole(a)
    Plot the electric field lines of two antiparallel dipoles were the horizontal and vertical
    distance between charges is 2*a."""
    plt.figure(figsize=(20, 14))
    xlim = (a+5); ylim = (a+5)
    plt.axes = plt.gca()
    plt.axes.set_xlim([-xlim,xlim])
    plt.axes.set_ylim([-ylim,ylim])
    x = np.linspace(-xlim, xlim, 100)
    y = np.linspace(-ylim, ylim, 100)
    x, y = np.meshgrid(x, y)
    E1 = mkcharge(a, a, 1, x, y)
    E2 = mkcharge(a, -a, -1, x, y)
    E3 = mkcharge(-a, a, -1, x, y)
    E4 = mkcharge(-a, -a, +1, x, y)
    Ex = E1[0] + E2[0] + E3[0] + E4[0]
    Ey = E1[1] + E2[1] + E3[1] + E4[1]
    color = np.log(np.sqrt(Ex**2 + Ey**2))
    #streamplot(x, y, Ex/sqrt(Ex**2+Ey**2), Ey/sqrt(Ex**2+Ey**2), color=color, linewidth=1, cmap=plt.cm.inferno, density=3.5, minlength=0.11, arrowstyle='->', arrowsize=1.5)
    plt.streamplot(x, y, Ex/np.sqrt(Ex**2+Ey**2), Ey/np.sqrt(Ex**2+Ey**2), color=color, linewidth=1, cmap=plt.cm.inferno, density=3.5, arrowstyle='->', arrowsize=1.5)
    #color = u + v
    #streamplot(x, y, Ex/sqrt(Ex**2+Ey**2), Ey/sqrt(Ex**2+Ey**2), color=color, linewidth=1, cmap=plt.cm.seismic, density=2, arrowstyle='->', arrowsize=1.5)
    plt.xlabel('$x$')
    plt.ylabel('$y$')
    plt.savefig('electric_quadrupole.pdf', transparent=True, bbox_inches='tight', pad_inches=0)

def plot_cylinder_lift(gamma=4.):
    # computes the velocity field on the mesh grid
    u_vortex, v_vortex = get_velocity_vortex(gamma, x_vortex, y_vortex, X, Y)

    # computes the stream-function on the mesh grid
    psi_vortex = get_stream_function_vortex(gamma, x_vortex, y_vortex, X, Y)
    # superposition of the doublet on the freestream flow
    u = u_freestream + u_doublet + u_vortex
    v = v_freestream + v_doublet + v_vortex

    # calculates the stagnation points
    x_stagn1, y_stagn1 = +math.sqrt(R**2-(gamma/(4*math.pi*u_inf))**2), -gamma/(4*math.pi*u_inf)
    x_stagn2, y_stagn2 = -math.sqrt(R**2-(gamma/(4*math.pi*u_inf))**2), -gamma/(4*math.pi*u_inf)

    # plots the streamlines
    size = 10
    pyplot.figure(figsize=(size, (y_end-y_start)/(x_end-x_start)*size))
    pyplot.xlabel('x', fontsize=16)
    pyplot.ylabel('y', fontsize=16)
    pyplot.xlim(x_start, x_end)
    pyplot.ylim(y_start, y_end)
    pyplot.streamplot(X, Y, u, v, density=2, linewidth=1, arrowsize=1.5, arrowstyle='->')
    circle = pyplot.Circle((0, 0), radius=R, color='#CD2305', alpha=0.5)
    pyplot.gca().add_patch(circle)
    pyplot.scatter(x_vortex, y_vortex, color='#CD2305', s=80, marker='o')
    pyplot.scatter([x_stagn1, x_stagn2], [y_stagn1, y_stagn2], color='g', s=80, marker='o');

def displayWindMapPlot(vdata,udata, lons, lats,):
    """ TODO add a docstring! """
    #plt.clf()
    #pc = plt.contourf(lons, lats, data, 20)
    #plt.colorbar(pc, orientation='horizontal')
    #plt.title(title)
    #plt.xlabel("longitude (degrees east)")
    #plt.ylabel("latitude (degrees north)")
    #plt.show()
    fig, ax = plt.subplots()
    # Do the plot code
    # make orthographic basemap.
    m = Basemap(projection='cyl',llcrnrlat=-40,urcrnrlat=0,\
            llcrnrlon=-20,urcrnrlon=60,resolution='l')

    X,Y=np.meshgrid(lons, lats)
    x,y=m(X,Y) #Convert to map coordinates
    #m.barbs(x,y,vdata,udata,20)
    m.quiver(x,y,vdata,udata,10)
    plt.streamplot(x,y,vdata,udata,10)
    #plt.colorbar(pc,orientation='horizontal')
    m.drawmapboundary()
    m.drawcountries()
    
    m.drawcoastlines(linewidth=1.5)
    
    fig.savefig('myimage.svg', format='svg', dpi=1200)
    plt.show()
    #m.drawparallels(parallels)
    #m.drawmeridians(meridians)
    
    
    """ Contains code for displaying data """

def octupole(a):
    """Function: octupole(a)"""
    plt.figure(figsize=(20, 14))
    xlim = (a+5); ylim = (a+5)
    plt.axes = plt.gca()
    plt.axes.set_xlim([-xlim,xlim])
    plt.axes.set_ylim([-ylim,ylim])
    x = np.linspace(-xlim, xlim, 100)
    y = np.linspace(-ylim, ylim, 100)
    x, y = np.meshgrid(x, y)
    # RR
    E1 = mkcharge(a+2*a, a, -1, x, y)
    E2 = mkcharge(a+2*a, -a, 1, x, y)
    # R
    E5 = mkcharge(a, a, 1, x, y)
    E6 = mkcharge(a, -a, -1, x, y)
    # L
    E7 = mkcharge(-a, a, -1, x, y)
    E8 = mkcharge(-a, -a, 1, x, y)
    # LL
    E3 = mkcharge(-(a+2*a), a, 1, x, y)
    E4 = mkcharge(-(a+2*a), -a, -1, x, y)
    Ex = E1[0] + E2[0] + E3[0] + E4[0] + E5[0] + E6[0] + E7[0] + E8[0]
    Ey = E1[1] + E2[1] + E3[1] + E4[1] + E5[1] + E6[1] + E7[1] + E8[1]
    #color = log(sqrt(abs(Ex) + abs(Ey)+2))
    color = np.log(np.sqrt(Ex**2 + Ey**2))
    plt.streamplot(x, y, Ex/np.sqrt(Ex**2+Ey**2), Ey/np.sqrt(Ex**2+Ey**2), color=color, linewidth=1, cmap=plt.cm.inferno, density=3.5, arrowstyle='->', arrowsize=1.5)
    plt.xlabel('$x$')
    plt.ylabel('$y$')
    plt.savefig('electric_octupole.pdf', transparent=True, bbox_inches='tight', pad_inches=0)

def plot_streamplot():
    a = Model(1, 1)
    J = a.joint_distn
    gridshape = (20, 20)
    U = np.empty(gridshape, dtype=float)
    V = np.empty(gridshape, dtype=float)
    Y, X = np.mgrid[0.1:2:20j, 0.1:2:20j]
    for i in range(20):
        for j in range(20):
            ap = X[i, j]
            bp = Y[i, j]
            b = Model(ap, bp)
            alpha, beta = b.solve_EM(J)
            U[i, j] = alpha - ap
            V[i, j] = beta - bp

    # init the figure
    pyplot.figure(figsize=(10, 10))

    # stream plot
    speed = np.sqrt(U*U + V*V)
    lw = 5 * speed / speed.max()
    pyplot.streamplot(X, Y, U, V, density=0.8, color='k', linewidth=lw)
    #pyplot.show()
    pyplot.savefig('streamplot.png')

 def animate(n):  # del ax.collections[:]; del ax.lines[:]; ;
     ax.cla()
     if track_way == "backward":
         Time = (locstart_time - timedelta(hours=n)).strftime("%d-%b-%Y %H:%M")
     else:
         Time = (locstart_time + timedelta(hours=n)).strftime("%d-%b-%Y %H:%M")
     plt.suptitle(
         "%.f%% simulated drifters ashore\n%d days, %d m, %s" % (int(round(p)), track_days, depth, Time)
     )
     if streamline == "ON":
         plt.streamplot(
             lonpps[n], latpps[n], US[n], VS[n], color=speeds[n], arrowsize=4, cmap=plt.cm.cool, density=2.0
         )
     for j in xrange(stp_num):
         ax.plot(lon_set[j][0], lat_set[j][0], color=colors[j % 10], marker="x", markersize=4)
         if n >= len(lon_set[j]):
             ax.plot(lon_set[j][-1], lat_set[j][-1], "o", color=colors[j % 10], markersize=5)
         if n < 5:
             if n < len(lon_set[j]):
                 ax.plot(
                     lon_set[j][: n + 1], lat_set[j][: n + 1], "o-", color=colors[j % 10], markersize=4
                 )  # ,label='Depth=10m'
         if n >= 5:
             if n < len(lon_set[j]):
                 ax.plot(
                     lon_set[j][n - 4 : n + 1],
                     lat_set[j][n - 4 : n + 1],
                     "o-",
                     color=colors[j % 10],
                     markersize=4,
                 )
     draw_basemap(ax, points)  # points is using here

def capacitor(a):
    plt.figure(figsize=(20, 14), dpi=80)
    xlim = (a+2); ylim = (a+2)
    plt.axes = plt.gca()
    plt.axes.set_xlim([-xlim,xlim])
    plt.axes.set_ylim([-ylim,ylim])
    x = np.linspace(-xlim, xlim, 60)
    y = np.linspace(-ylim, ylim, 60)
    x, y = np.meshgrid(x, y)
    Lx = []
    Ly = []
    Ex = 0
    Ey = 0
    for m in range(-30,30):
        EFt = mkcharge(0.015*m, a, -1, x, y)
        Lx.append(EFt[0])
        Ly.append(EFt[1])
    for n in range(-30,30):
        EFb = mkcharge(0.015*n, -a, 1, x, y)
        Lx.append(EFb[0])
        Ly.append(EFb[1])
    for xelem in Lx:
        Ex += xelem
    for yelem in Ly:
        Ey += yelem
    color = np.log(np.sqrt(np.abs(Ex) + np.abs(Ey)))
    plt.streamplot(x, y, Ex/np.sqrt(Ex**2+Ey**2), Ey/np.sqrt(Ex**2+Ey**2), color=color, linewidth=1, cmap=plt.cm.inferno, density=2, arrowstyle='->', arrowsize=1.5)
    plt.xlabel('$x$')
    plt.ylabel('$y$')
    plt.savefig('capacitor.pdf', transparent=True, bbox_inches='tight', pad_inches=0)

def test_linewidth():
    X, Y, U, V = velocity_field()
    speed = np.hypot(U, V)
    lw = 5 * speed / speed.max()
    df = 25 / 30   # Compatibility factor for old test image
    plt.streamplot(X, Y, U, V, density=[0.5 * df, 1. * df], color='k',
                   linewidth=lw)

def test_complete():
    fig = plt.figure('Figure with a label?', figsize=(10, 6))

    plt.suptitle('Can you fit any more in a figure?')

    # make some arbitrary data
    x, y = np.arange(8), np.arange(10)
    data = u = v = np.linspace(0, 10, 80).reshape(10, 8)
    v = np.sin(v * -0.6)

    plt.subplot(3, 3, 1)
    plt.plot(list(xrange(10)))

    plt.subplot(3, 3, 2)
    plt.contourf(data, hatches=['//', 'ooo'])
    plt.colorbar()

    plt.subplot(3, 3, 3)
    plt.pcolormesh(data)

    plt.subplot(3, 3, 4)
    plt.imshow(data)

    plt.subplot(3, 3, 5)
    plt.pcolor(data)

    plt.subplot(3, 3, 6)
    plt.streamplot(x, y, u, v)

    plt.subplot(3, 3, 7)
    plt.quiver(x, y, u, v)

    plt.subplot(3, 3, 8)
    plt.scatter(x, x**2, label='$x^2$')
    plt.legend(loc='upper left')

    plt.subplot(3, 3, 9)
    plt.errorbar(x, x * -0.5, xerr=0.2, yerr=0.4)

    ###### plotting is done, now test its pickle-ability #########

    # Uncomment to debug any unpicklable objects. This is slow (~200 seconds).
#    recursive_pickle(fig)

    result_fh = BytesIO()
    pickle.dump(fig, result_fh, pickle.HIGHEST_PROTOCOL)

    plt.close('all')

    # make doubly sure that there are no figures left
    assert_equal(plt._pylab_helpers.Gcf.figs, {})

    # wind back the fh and load in the figure
    result_fh.seek(0)
    fig = pickle.load(result_fh)

    # make sure there is now a figure manager
    assert_not_equal(plt._pylab_helpers.Gcf.figs, {})

    assert_equal(fig.get_label(), 'Figure with a label?')

def test_startpoints():
    X, Y, U, V = velocity_field()
    start_x = np.linspace(X.min(), X.max(), 10)
    start_y = np.linspace(Y.min(), Y.max(), 10)
    start_points = np.column_stack([start_x, start_y])
    plt.streamplot(X, Y, U, V, start_points=start_points)
    plt.plot(start_x, start_y, 'ok')

def test_masks_and_nans():
    X, Y, U, V = velocity_field()
    mask = np.zeros(U.shape, dtype=bool)
    mask[40:60, 40:60] = 1
    U[:20, :20] = np.nan
    U = np.ma.array(U, mask=mask)
    with np.errstate(invalid='ignore'):
        plt.streamplot(X, Y, U, V, color=U, cmap=plt.cm.Blues)

 def plot_both(self, c_points=200, g_points=200):
     """ Plots the grad quiver plot over the domain of the function. """
     # Latex
     plt.rc('text', usetex=True)
     plt.rc('font', family='serif')
     ## Domain Correction
     # Lower x0 Limit
     if np.isfinite(self.domain[0][0]):
         x0_lim_lower = self.domain[0][0]
     else:
         x0_lim_lower = -10.0
     # Upper x0 Limit
     if np.isfinite(self.domain[0][1]):
         x0_lim_upper = self.domain[0][1]
     else:
         x0_lim_upper = +10.0
     # Lower x1 Limit
     if np.isfinite(self.domain[1][0]):
         x1_lim_lower = self.domain[1][0]
     else:
         x1_lim_lower = -10.0
     # Upper x1 Limit
     if np.isfinite(self.domain[1][1]):
         x1_lim_upper = self.domain[1][1]
     else:
         x1_lim_upper = +10.0
     ## Lines
     x0c = np.linspace(x0_lim_lower, x0_lim_upper, c_points)
     x1c = np.linspace(x1_lim_lower, x1_lim_upper, c_points)
     x0g = np.linspace(x0_lim_lower, x0_lim_upper, g_points)
     x1g = np.linspace(x1_lim_lower, x1_lim_upper, g_points)
     ## Meshes
     X0c, X1c = np.meshgrid(x0c, x1c)
     X0g, X1g = np.meshgrid(x0g, x1g)
     ## Combined
     Xc = np.array([X0c, X1c])
     Xg = np.array([X0g, X1g])
     ## Calculate Costs
     cost = self.cost(Xc)
     ## Renormalise
     cost_norm = np.log(cost - np.min(cost) + 1)
     ## Calculate Grad
     grad = self.grad(Xg)
     ## Renormalise
     grad_norm = grad / np.linalg.norm(grad, axis=0)
     ## Plot
     plt.figure()
     plt.contourf(X0c, X1c, cost_norm, 50)
     plt.scatter(self.min[..., 0], self.min[..., 1], c='w', marker='x')
     plt.streamplot(X0g, X1g, -grad_norm[0], -grad_norm[1], density=4.0, color='k')
     plt.scatter(self.min[0], self.min[1], c='w', marker='x')
     plt.grid()
     plt.title(self.latex_name + "\n" + self.latex_cost)
     plt.subplots_adjust(top=0.8)
     plt.xlabel('$x_0$')
     plt.ylabel('$x_1$')
     plt.xlim([x0_lim_lower, x0_lim_upper])
     plt.ylim([x1_lim_lower, x1_lim_upper])

def plot_stresses_cartesian(chi, F, R_min):#, R_max, R_steps):
    Y, X = np.mgrid[-5:5:1000j, -5:5:1000j]
    R = np.sqrt(X**2. + Y**2.)

    # pre_X = np.linspace(-R_max, R_max, R_steps) 
    # X, Y = np.meshgrid(pre_X, pre_X)
    # R = np.sqrt(X**2. + Y**2.)

    cart_betas1 = stress_beta1_cartesian(X, Y, chi, F, R_min)
    beta_X1s = to_unit_vector_x(cart_betas1)
    beta_Y1s = to_unit_vector_y(cart_betas1)
    beta_X1s[R<1] = 1
    beta_Y1s[R<1] = 0

    cart_betas2 = stress_beta2_cartesian(X, Y, chi, F, R_min)
    beta_X2s = to_unit_vector_x(cart_betas2)
    beta_Y2s = to_unit_vector_y(cart_betas2)
    beta_X2s[R<1] = 0
    beta_Y2s[R<1] = 1

    fig = plt.figure(figsize=(5, 5))
    ax=fig.add_subplot(111)

    # streamplot
    # ax=fig.add_subplot(111)
    ax.set_title('Stress Trajectories')
    plt.streamplot(X, Y, beta_X1s, beta_Y1s, arrowstyle='-', density=2.7, color='b', minlength=0.9)
    plt.streamplot(X, Y, beta_X2s, beta_Y2s, arrowstyle='-', density=1.1, color='olive', minlength=0.9)

    ax.add_patch(Circle((0, 0), radius=1, zorder=10, facecolor='k', linewidth=2))

    plt.axis("image")
    # ax.add_patch(Circle((6.55, 0), radius=.2, zorder=10, facecolor='r', linewidth=1))
    # ax.add_patch(Circle((-6.55, 0), radius=.2, zorder=10, facecolor='r', linewidth=1))
    # ax.add_patch(Circle((-np.sqrt(2)/2., np.sqrt(2)/2.), radius=.2, zorder=10, facecolor='r', linewidth=1))
    # ax.add_patch(Circle((-np.sqrt(2)/2., -np.sqrt(2)/2.), radius=.2, zorder=10, facecolor='r', linewidth=1))

    # quiver
    # plt.quiver(X, Y, beta_X1s, beta_Y1s, pivot='middle', headwidth=2, headlength=4, color='b')
    # plt.quiver(X, Y, beta_X2s, beta_Y2s, pivot='middle', headwidth=2, headlength=4, color='g')
    for label in ax.xaxis.get_ticklabels() + ax.yaxis.get_ticklabels():
        # label is a Text instance
        label.set_color('k')
        # label.set_rotation(45)
        label.set_fontsize(12)

    for line in ax.xaxis.get_ticklines() + ax.yaxis.get_ticklines():
        # line is a Line2D instance
        line.set_color('k')
        line.set_markersize(5)
        line.set_markeredgewidth(2)

    plt.xlabel(r'$\chi = $'+str(round(chi, 1)) + ', ' + r'$F = $'+ str(round(F, 1)))
    plt.ylim(-5, 5)
    plt.xlim(-5, 5)
    # plt.tight_layout()
    # plt.savefig("stream.png", dpi=300)
    plt.show()

def stream(i,a,b):
    '''
    dy=c+i-y
    dc=a+b*y-c
    '''
    y, c = np.mgrid[0:30:100j, 0:30:100j]
    dy=c+i-y
    dc=a+b*y-c
    plt.streamplot(c,y,dc,dy)
    plt.show()

 def animate(n): #del ax.collections[:]; del ax.lines[:]; ax.cla();ax.clf()
     ax.cla()
     ax.plot(drifter_points['lon'],drifter_points['lat'],'bo-',markersize=6,label='Drifter')
     if streamline == 'ON':
         plt.streamplot(lonpps[n],latpps[n],US[n],VS[n], color=speeds[n],arrowsize=4,cmap=plt.cm.cool,density=2.0)              
     if n==0:#facecolor=colors[i]'''
         ax.annotate(an2,xy=(dr_points['lon'][-1],dr_points['lat'][-1]),xytext=(dr_points['lon'][-1]+0.01*track_days,
                     dr_points['lat'][-1]+0.01*track_days),fontsize=6,arrowprops=dict(arrowstyle="fancy"))
     ax.plot(model_points['lon'][:n+1],model_points['lat'][:n+1],'ro-',markersize=6,label=MODEL)
     draw_basemap(ax, points)  # points is using here

def velFieldPlot(folder, t_0, t_end, iter):
    # Go to current directory
    curDir = os.getcwd()
    os.chdir(curDir)

    # Prepare image
    fig = plt.figure()
    fig.add_subplot(111)

    # Read data
    filenameX = folder + "/" + str(iter) + "_X.txt"
    filenameY = folder + "/" + str(iter) + "_Y.txt"
    filenameU = folder + "/" + str(iter) + "_U.txt"
    filenameV = folder + "/" + str(iter) + "_V.txt"
    nParticles, X = readFile(filenameX)
    nParticles, Y = readFile(filenameY)
    nParticles, U = readFile(filenameU)
    nParticles, V = readFile(filenameV)

    # Interpolate onto grid
    U_grid = interpolate(X,Y,U)
    V_grid = interpolate(X,Y,V)

    # Compute the extent of the coordinates
    xmin = np.min(X)
    ymin = np.min(Y)
    xmax = np.max(X)
    ymax = np.max(Y)

    # Make an XY grid corresponding to the interpolated locations of U and V
    Y_grid, X_grid = np.mgrid[xmin:xmax:400j, ymin:ymax:400j]

    # Make a speed grid
    speed = np.sqrt(U_grid**2 + V_grid**2)
    print(speed.max())
    if(speed.max() != 0):
        lw = 3*speed/speed.max()
    else:
        lw = 1
        print("Maximum speed is 0.")


    # Make a streamfunction plot
    plt.streamplot(X_grid, Y_grid, U_grid.T, V_grid.T,          # data
                   density=[0.5, 1],
                   color='DarkRed',
                   linewidth=lw)

    plt.title('Streamline Plot, at t = 0.00001 sec')
    plt.xlabel("x")
    plt.ylabel("y")
    plt.show(fig)


    return

def plot_fisher(c:(0,4,0.1)=1):
    gx, gy = np.meshgrid(np.linspace(-0.5, 1.5, 11), \
                         np.linspace(-1, 1, 11))
    dx, dy = fisher([gx, gy], 0, c)

    plt.figure(figsize=[9,6])
    plt.streamplot(gx, gy, dx, dy, color=col[0])
    plt.xlabel('U')
    plt.ylabel('V')
    plt.xlim(-0.5,1.5)
    plt.ylim(-1,1)
    plt.show()

def plot_streamplot(r,z,Br,Bz,title='',xlab = '',ylab=''):
    
    """
    This function creates a streamplot with four parameters (r,z,Br,Bz).
    """
    
    plt.streamplot(r, z, Br, Bz)
    plt.title(title, fontsize=16)
    plt.xlabel(xlab, fontsize=14)
    plt.ylabel(ylab, fontsize=14)
    plt.xticks(fontsize=8)
    plt.yticks(fontsize=8)
    plt.legend()

def PhasePortrait(p):
    Y, X = np.mgrid[0.1:0.3:5000j, 0:0.72:5000j]
    K1, K2, K3, KM1, KM3 = p
    U, V = [K1*(1 - X - 2*Y) - KM1*X - K3*X*(1 - X - 2*Y) + KM3*Y - K2*((1 - X - 2*Y)**2)*X,
            K3*X*(1 - X - 2*Y) - KM3*Y]
    #U, V = [K1*(1 - X - Y) - KM1*X - K3*X + KM3*Y - K2*((1 - X - Y)**2)*X,
    #        K3*X - KM3*Y]
    plt.streamplot(X, Y, U, V, density = [2, 2])
    plt.xlim([0, 0.72])
    plt.ylim([0.1, 0.3])
    plt.xlabel('x')
    plt.ylabel('y')
    plt.show()