from pygasflow.shockwave import PressureDeflectionLocus
import matplotlib.pyplot as plt

gamma = 1.4
M1 = 3
theta = 25
locus1 = PressureDeflectionLocus(M=M1, gamma=gamma)
locus2 = locus1.new_locus_from_shockwave(theta)

theta_1, pr_1 = locus1.pressure_deflection(include_mirror=False)
theta_2, pr_2 = locus2.pressure_deflection(include_mirror=True)

fig, ax = plt.subplots()
ax.plot(theta_1, pr_1, label=f"M = {locus1.M}")
ax.plot(theta_2, pr_2, label=f"M = {locus2.M}")
ax.set_xlabel(r"Deflection Angle $  heta$ [deg]")
ax.set_ylabel("Pressure Ratio to Freestream")
ax.legend()
ax.minorticks_on()
ax.grid(which='major', linestyle='-', alpha=0.7)
ax.grid(which='minor', linestyle=':', alpha=0.5)
plt.show()