import matplotlib.pyplot as plt 
import numpy as np

# == Parameters ==
#  [Center of the circle in z-plane]
xc,yc = -0.1,0.2 
#  [Joukowski mapping parameter: \zeta = z + a*a/z]
a = 0.8
#  [Angle of Attack in z- and \zeta-planes]
alpha = 10.0*np.pi/180.0 # 10[degree] is converted to [radian]

# == Coodinate points of z=x+i*y ==
size = 5.0
x1d = np.arange(-size,size,0.01)
y1d = x1d
x,y = np.meshgrid(x1d,y1d) # Making 2D mesh of (x,y)
z = x + 1j*y # z = x + i*y (1j is an imaginary unit in Python)
zc = xc + 1j*yc # center of the circle in z-plane


# == Fixed parameters ==
U = 1.0 # Velocity of the uniform flow
R = np.sqrt((np.abs(xc)+a)*(np.abs(xc)+a)+yc*yc) # Radius of the circle in z-plane
beta = np.arcsin(yc) # complex angle of the circle center in z-plane
expa = np.exp(-1j*alpha) # = exp[-i*\alpha]
expa2 = np.exp(2j*alpha) # = exp[2i*\alpha]
Gamma = -4.0*np.pi*a*U*np.sin(alpha + beta) # Circulation determined by the Kutta condition
# Gamma = -4.0*np.pi*a*U*np.sin(alpha + beta) - 5.0 # If the circulation does not satisfy the Kutta condition...

# == Define potential in z-plane (flow around the circle with circulation) ==
W = U*expa*((z - zc) + R*R*expa2/(z - zc)) - 1j*Gamma/(2.0*np.pi)*np.log((z - zc)*expa/R) 
# == Im(W) is the streamline function ==
Psi = np.where(np.abs((z - zc))<=R, 0.0, W.imag) # Remove streamline inside the circle
# == Define surface shape in z-plane (circle with the center at zc=xc+i*yc)
body = R*np.exp(-1j*np.linspace(-np.pi,np.pi,100))+zc 

# == Joukowski mapping ==
zeta = lambda z: (z+a*a/z) 

# == Calculate lift ==
rho = 1.0
L = -rho*Gamma*U # Kutta-Joukowski Theorem !!
print('Lift=', L, ' (density=1.0 is assumed)')
# What about the drag...?

# == Plots ==
fig = plt.figure(figsize=(10,5))
# == z-plane (flow around the circle) ==
ax1 = fig.add_subplot(1,2,2)
clines = np.arange(-4.0,4.0,0.1) # Define contour lines for streamline (Im(W)=\Psi)
ax1.contour(zeta(z).real,zeta(z).imag,Psi,levels=clines,colors='blue',linestyles='-',linewidths=0.5) # Streamline
ax1.plot(zeta(body).real,zeta(body).imag,color='black') # Surface of the body
xr=3.
ax1.set_xlabel('$\\xi$')
ax1.set_ylabel('$\\eta$')
ax1.set_xlim(-xr,xr)
ax1.set_ylim(-xr,xr)
ax1.set_aspect('equal')
ax1.set_title('$\\zeta$-plane')

# == \zeta-plane (flow around the airfoil) ==
ax2 = fig.add_subplot(1,2,1)
clines = np.arange(-4.0,4.0,0.1)
ax2.contour(z.real,z.imag,Psi,levels=clines,colors='green',linestyles='-',linewidths=0.5) # Streamline
ax2.plot(body.real,body.imag,color='black') # Surface of the body
xr=3.
ax2.set_xlabel('$x$')
ax2.set_ylabel('$y$')
ax2.set_xlim(-xr,xr)
ax2.set_ylim(-xr,xr)
ax2.set_title('$z$-plane')
ax2.set_aspect('equal')
plt.show() # inline execution (e.g., jupyter notebook)
#fig.savefig('Joukowski.pdf') # if run as a main script (i.e., python airfoil.py) and get .pdf file
print('Plot completed => Joukowski.pdf')