# Julia translation of http://nbviewer.jupyter.org/github/barbagroup/AeroPython/blob/master/lessons/03_Lesson03_doublet.ipynb
# Lesson 3 doublet

using PyPlot
using Distributions

close("all")
meshgrid(x,y) = (repmat(x',length(y),1),repmat(y,1,length(x)))

N = 50                                # number of points in each direction
x_start, x_end = -2.0, 2.0            # boundaries in the x-direction
y_start, y_end = -1.0, 1.0            # boundaries in the y-direction
x = linspace(x_start, x_end, N)    # creates a 1D-array with the x-coordinates
y = linspace(y_start, y_end, N)    # creates a 1D-array with the y-coordinates

X,Y=meshgrid(x,y)

kappa = 1.0

x_doublet, y_doublet = 0.0, 0.0   # location of the source

function get_velocity_doublet(strength, xd, yd, X, Y):
    """Returns the velocity field generated by a doublet.

    Arguments
    ---------
    strength -- strength of the doublet.
    xd, yd -- coordinates of the doublet.
    X, Y -- mesh grid.
    """
    u = - strength/(2*pi).*((X-xd).^2-(Y-yd).^2)./((X-xd).^2+(Y-yd).^2).^2
    v = - strength/(2*pi)*2.0.*(X-xd).*(Y-yd)./((X-xd).^2+(Y-yd).^2).^2

    return u, v
  end

function get_stream_function_doublet(strength, xd, yd, X, Y):
    """Returns the stream-function generated by a doublet.

    Arguments
    ---------
    strength -- strength of the doublet.
    xd, yd -- coordinates of the doublet.
    X, Y -- mesh grid.
    """
    psi = - strength/(2*pi).*(Y-yd)./((X-xd).^2+(Y-yd).^2)

    return psi
  end

  # computes the velocity field on the mesh grid
u_doublet, v_doublet = get_velocity_doublet(kappa, x_doublet, y_doublet, X, Y)

# computes the stream-function on the mesh grid
psi_doublet = get_stream_function_doublet(kappa, x_doublet, y_doublet, X, Y)


  # plots the streamlines
  size = 10
  PyPlot.figure(figsize=(size, (y_end-y_start)/(x_end-x_start)*size))
  PyPlot.grid(true)
  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_doublet, v_doublet, density=2, linewidth=1, arrowsize=1, arrowstyle="->")
  PyPlot.scatter(x_doublet, y_doublet, color="#CD2305", s=80, marker="o")

# Freestream and Doublet

  u_inf = 1.0

  # computes the freestream velocity field
  u_freestream = u_inf * ones((N, N))
  v_freestream = zeros((N, N))

  # computes the stream-function
  psi_freestream = u_inf * Y

  u = u_freestream + u_doublet
  v = v_freestream + v_doublet
  psi = psi_freestream + psi_doublet


  # plots the streamlines
  size = 10
  PyPlot.figure(figsize=(size, (y_end-y_start)/(x_end-x_start)*size))
  PyPlot.grid(true)
  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, arrowstyle="->")
  PyPlot.contour(X, Y, psi, levels=[0.], colors="#CD2305", linewidths=2, linestyles="solid");
  PyPlot.scatter(x_doublet, y_doublet, color="#CD2305", s=80, marker="o")

  # calculates the stagnation points
x_stagn1, y_stagn1 = +sqrt(kappa/(2*pi*u_inf)), 0
x_stagn2, y_stagn2 = -sqrt(kappa/(2*pi*u_inf)), 0

# adds the stagnation points to the figure
PyPlot.scatter([x_stagn1, x_stagn2], [y_stagn1, y_stagn2], color="g", s=80, marker="o");


# computes the pressure coefficient field
cp = 1.0 - (u.^2+v.^2)./u_inf^2

# plots the pressure coefficient field
size = 10
PyPlot.figure(figsize=(1.1*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)
contf = PyPlot.contourf(X, Y, cp, levels=linspace(-2.0, 1.0, 100), extend="both")
cbar = PyPlot.colorbar(contf)
cbar[:set_label](L"$C_p$", fontsize=16)
cbar[:set_ticks]([-2.0, -1.0, 0.0, 1.0])
PyPlot.scatter(x_doublet, y_doublet, color="#CD2305", s=80, marker="o")
PyPlot.contour(X, Y, psi, levels=[0.], colors="#CD2305", linewidths=2, linestyles="solid");
PyPlot.scatter([x_stagn1, x_stagn2], [y_stagn1, y_stagn2], color="g", s=80, marker="o");
