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_1w, pr_1w, theta_1s, pr_1s = locus1.pressure_deflection_split(
    include_mirror=False, mode="region")
theta_2sup, pr_2sup, theta_2sub, pr_2sub = locus2.pressure_deflection_split(
    include_mirror=True, mode="sonic")

fig, ax = plt.subplots()
ax.plot(theta_1w, pr_1w, label="M1 weak")
ax.plot(theta_1s, pr_1s, label="M1 strong")
ax.plot(theta_2sup, pr_2sup, label="M2 (M3 supersonic)")
ax.plot(theta_2sub, pr_2sub, label="M2 (M3 subsonic)")
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()