import matplotlib.pyplot as plt

from src.KSP_tools.calculator_dv import cal_dv

centaur_isp = 450.5
centaur_gross_mass = 20830 + 2247
centaur_dry_mass = 2247

lm_isp = 438
lm_gross_mass = 21000
lm_dry_mass = 2800

lmh_isp = 450
lmh_gross_mass = 23500 + 3000
lmh_dry_mass = 3000

payload_list = [i for i in range(0, 10000, 100)]
centaur_performance = [cal_dv(centaur_isp, centaur_dry_mass + i, centaur_gross_mass + i) for i in payload_list]
lm_performance = [cal_dv(lm_isp, lm_dry_mass + i, lm_gross_mass + i) for i in payload_list]
lmh_performance = [cal_dv(lmh_isp, lmh_dry_mass + i, lmh_gross_mass + i) for i in payload_list]

plt.plot(centaur_performance, label='Centaur')
plt.plot(lm_performance, label='Long March')
plt.plot(lmh_performance, label='Long March 3.35')
plt.legend()
plt.show()
