from pygasflow import *
from pygasflow.atd.newton import pressure_coefficient
import numpy as np
import matplotlib.pyplot as plt
M_inf = 8
gamma = 1.4
a = 0.769
y = np.linspace(0, 2)
# To get theta: rewrite the paraboloid in terms of y=f(x). 
# Find its tangent by differentiating it by x. Substitute x with the 
# paraboloid equation. Apply the arctan.
theta = np.arctan(1 / (2 * a * np.sqrt(y**2)))
Cp_newtonian = pressure_coefficient(theta)
Cp_mod_newtonian = pressure_coefficient(theta, Mfs=M_inf, gamma=gamma)
# from the definition of Cp
ps_p_inf_newt = Cp_newtonian * M_inf**2 * gamma / 2 + 1
ps_p_inf_mod_newt = Cp_mod_newtonian * M_inf**2 * gamma / 2 + 1
# compute ratios across the shockwave
shock = normal_shockwave_solver(
    "mu", M_inf, gamma=gamma, to_dict=True)
# compute the flow quantities downstream of the normal shockwave
flow2 = isentropic_solver("m", shock["md"], gamma=gamma, to_dict=True)
p2_pt2 = 1 / flow2["pr"]
# NOTE: alternatively, rayleigh_pitot_formula can be used
pt2_p_inf = p2_pt2 * shock["pr"]
# final results
ps_pt2_newtonian = ps_p_inf_newt / pt2_p_inf
ps_pt2_mod_newtonian = ps_p_inf_mod_newt / pt2_p_inf
# plot
plt.figure()
plt.plot(y, ps_pt2_newtonian, label="Newtonian")
plt.plot(y, ps_pt2_mod_newtonian, label="Modified Newtonian")
plt.xlabel("Y")
plt.ylabel(r"$p_{s} \, / \, p_{t2}$")
plt.legend()
plt.grid(which="major", visible=True, linestyle="--")
plt.xlim(0, 2)
plt.ylim(0, 1.2)
plt.suptitle("Surface-pressure distribution")
plt.title("over the paraboloid: $x = 0.769 y^{2} - 1$")
plt.show()