import matplotlib.pyplot as plt

from src.KSP_tools.calculator_dv import cal_dv


def plot_stage_performance(isp, gross_mass, dry_mass, payload, stage_name):
    stage_performance = [cal_dv(isp, dry_mass + i, gross_mass + i) for i in payload]
    plt.plot(payload, stage_performance, label=stage_name)


payload_list = [i for i in range(0, 40200, 200)]
plot_stage_performance(450.5, 20830 + 2247, 2247, payload_list, 'Centaur III')
# plot_stage_performance(438, 21000, 2800, payload_list, 'Long March 3m')
# plot_stage_performance(450, 23500 + 2900, 2900, payload_list, 'Long March 3.35m')
plot_stage_performance(453.8, 54000+4500, 5000, payload_list, 'Centaur V')
plot_stage_performance(348, 112000, 4000, payload_list, 'Falcon 9')
plot_stage_performance(355, 80600, 600, payload_list, 'Neutron S2')
plot_stage_performance(365, 100600, 600, payload_list, 'Neutron S2 MAX')
# plot_stage_performance(451, 68000, 8000, payload_list, 'Long March 10 S3')
# plot_stage_performance(442.6, 36000, 5100, payload_list, 'Long March 5 S2')
# plot_stage_performance(462, 30710, 3490, payload_list, 'DCSS 5m')
# plot_stage_performance(462, 32748, 3800, payload_list, 'ICPS')
# plot_stage_performance(460.1, 129000+14110, 14110, payload_list, 'EUS')
plt.xlabel('payload mass in kg')
plt.ylabel('delta-v in m/s')
plt.legend()
plt.show()
