# -*- coding: utf-8 -*-
"""
Created on Fri Jan  8 16:52:25 2021

@author: Daniel Powell
"""

import numpy as np
from matplotlib import pyplot as plt


def forced_vortex(r, c, R):
    """
    Computes and returns the tangential velocity at radial distance r for a
    solid body rotation.

    Parameters
    ----------
    r : float
        radial position (m).
    c : float
        constant of proportionality (/s).

    Returns
    -------
    u_t: float
        tangential velocity (m/s).

    """
    u_t = c * r / R ** 2

    return u_t


def irrotational_vortex(r, c):
    """
    Computes and returns the tangential velocity at radial distance r for an
    irrotational vortex.

    Parameters
    ----------
    r : float
        radial position (m). Must be greater than 0.
    c : float
        constant (m^2/s).

    Returns
    -------
    u_t: float
        tangential velocity (m/s).

    """
    u_t = c / r

    return u_t


def rankine_vortex(r, c, R):
    """
    Computes and returns the tangential velocity at radial distance r for a
    Rankine vortex, with maximum velocity at r=R

    Parameters
    ----------
    r : float
        radial position (m). Must be greater than 0.
    c : float
        constant (m^2/s).
    R : float
        radial position of maximum velocity (m). Must be greater than 0.

    Returns
    -------
    u_t: float
        tangential velocity (m/s).

    """
    if np.abs(r) < R:
        u_t = forced_vortex(r, c, R)
    else:
        u_t = irrotational_vortex(r, c)

    return u_t


r_max = 2.0
r = np.linspace(-r_max, r_max, num=10000)
c = 1.0
R = 0.3

u_rankine_max = rankine_vortex(R, c, R)
max_ext = 1.5

u_forced = np.zeros_like(r)
u_irro = np.zeros_like(r)
u_rankine = np.zeros_like(r)


for i in range(len(r)):
    u_forced[i] = forced_vortex(r[i], c, R)
    u_irro[i] = irrotational_vortex(r[i], c)
    u_rankine[i] = rankine_vortex(r[i], c, R)

    if np.abs(u_forced[i]) > u_rankine_max * max_ext:
        u_forced[i] = np.nan
        
    if np.abs(u_irro[i]) > u_rankine_max * max_ext:
        u_irro[i] = np.nan


fig, ax = plt.subplots()
plt.plot(r / R, u_forced / u_rankine_max, '--', label="Rigid Body Vortex")
plt.plot(r / R, u_irro / u_rankine_max, ':', label="Irrotational Vortex")
plt.plot(r / R, u_rankine / u_rankine_max, 'k-', label="Rankine Vortex")
plt.xlabel(r"$\frac{r}{R}$", fontsize=16)
plt.ylabel(r"$\frac{u_t}{u_{max}}$", fontsize=16)
plt.legend(loc='lower right')
#ax.set_aspect('equal')
plt.tight_layout()
# plt.savefig("rankine_vortex.eps")