# 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)

# 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");


gamma = 4.0                      # strength of the vortex
x_vortex, y_vortex = 0.0, 0.0    # location of the vortex

function get_velocity_sink(strength, xs, ys, X, Y):
    """Returns the velocity field generated by a sink.

    Arguments
    ---------
    strength -- strength of the sink.
    xs, ys -- coordinates of the sink.
    X, Y -- mesh grid.
    """
    u = strength/(2*pi).*(X-xs)./((X-xs).^2+(Y-ys).^2)
    v = strength/(2*pi).*(Y-ys)./((X-xs).^2+(Y-ys).^2)

    return u, v
  end

function get_stream_function_sink(strength, xs, ys, X, Y):
      """Returns the stream-function generated by a sink.

      Arguments
      ---------
      strength -- strength of the sink.
      xs, ys -- coordinates of the sink.
      X, Y -- mesh grid.
      """
      psi = strength/(2*pi)*atan2((Y-ys), (X-xs))

      return psi
    end

    function get_velocity_vortex(strength, xv, yv, X, Y):
          """Returns the velocity field generated by a vortex.

          Arguments
          ---------
          strength -- strength of the vortex.
          xv, yv -- coordinates of the vortex.
          X, Y -- mesh grid.
          """
          u = + strength/(2*pi).*(Y-yv)./((X-xv).^2+(Y-yv).^2)
          v = - strength/(2*pi).*(X-xv)./((X-xv).^2+(Y-yv).^2)

          return u, v
        end

    function get_stream_function_vortex(strength, xv, yv, X, Y):
        """Returns the stream-function generated by a vortex.

        Arguments
        ---------
        strength -- strength of the vortex.
        xv, yv -- coordinates of the vortex.
        X, Y -- mesh grid.
        """
        psi = strength/(4*pi).*log((X-xv).^2+(Y-yv).^2)

        return psi
      end

    # 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 and the vortex on the freestream flow
u = u_freestream + u_doublet + u_vortex
v = v_freestream + v_doublet + v_vortex
psi = psi_freestream + psi_doublet + psi_vortex

# calculates cylinder radius
R = sqrt(kappa/(2*pi*u_inf))

# calculates the stagnation points
x_stagn1, y_stagn1 = +sqrt(R^2-(gamma/(4*pi*u_inf))^2), -gamma/(4*pi*u_inf)
x_stagn2, y_stagn2 = -sqrt(R^2-(gamma/(4*pi*u_inf))^2), -gamma/(4*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=13000, marker="o")
PyPlot.scatter([x_stagn1, x_stagn2], [y_stagn1, y_stagn2], color="g", s=80, marker="o");


# calculates the surface tangential velocity on the cylinder
theta = linspace(0, 2*pi, 100)
u_theta = -2*u_inf*sin(theta) - gamma/(2*pi*R)

# computes the surface pressure coefficient
cp = 1.0 - (u_theta./u_inf).^2

# if there was no vortex
u_theta_no_vortex = -2*u_inf*sin(theta)
cp_no_vortex = 1.0 - (u_theta_no_vortex./u_inf).^2

q=0.5*1.225*u_inf^2
drag1 = - cp.*q .* 2*pi .* sin(2*pi) .* R
DRAG = sum(drag1)

lift1  = cp.*q .* 2*pi .* cos(2*pi) .* R
LIFT = sum(lift1)

println(DRAG)
println(LIFT)

# plots the surface pressure coefficient
size = 6
PyPlot.figure(figsize=(size, size))
PyPlot.grid(true)
PyPlot.xlabel(L"$theta$", fontsize=18)
PyPlot.ylabel(L"$C_p$", fontsize=18)
PyPlot.xlim(minimum(theta), maximum(theta))
PyPlot.plot(theta, cp, color="#CD2305", linewidth=2, linestyle="-", label = "With Vortex")
PyPlot.plot(theta, cp_no_vortex, color="g", linewidth=2, linestyle="-",label = "Without Vortex")
PyPlot.legend(loc = "best");
