import sys import math from random import randint from engine import Engine from fuel import Fuel from stage import Stage from rocket import Rocket from atmosphere import Atmosphere from body import Body from universe import Universe from simulation import Simulation import pygame from pygame.locals import * def main(argv): raptor_engine = Engine(name="raptor", isp=360, max_flow_rate=931) #https://en.wikipedia.org/wiki/SpaceX_Raptor methane_fuel = Fuel(name="methane") #TODO: more #https://en.wikipedia.org/wiki/SpaceX_Starship first_stage = Stage(name="superheavy booster", stage_mass=180000, engine=raptor_engine, engine_number=33, fuel_type=methane_fuel, fuel_mass=3600000, drag_coefficient=1.18, cross_sectional_area=(math.pi * (9**2)) ) second_stage = Stage(name="starship", stage_mass=80000, engine=raptor_engine, engine_number=6, fuel_type=methane_fuel, fuel_mass=1200000, drag_coefficient=1.18, cross_sectional_area=(math.pi * (9**2)) ) rocket = Rocket(name="starship launch system", stages=[first_stage, second_stage], payload_mass=100 ) body = Body(name="earth", density=5.51, radius=6371000, atmosphere=Atmosphere( avg_sea_level_pressure=101325, molar_mass_air=0.02896, standard_temp=288.15 ) ) universe = Universe(name="conventional", G=6.67E-11 ) simulation = Simulation(universe, body, rocket) simulation.rocket.current_stage().engines_on = True pygame.init() pygame.display.set_caption("OSLS - Overly Simple Launch Simulator") clock = pygame.time.Clock() SCREEN_WIDTH = 1024 SCREEN_HEIGHT = 720 simulation_display = pygame.display.set_mode((SCREEN_WIDTH,SCREEN_HEIGHT)) paused = False while True: if not paused: draw_simulation(simulation_display, simulation) pygame.display.update() for event in pygame.event.get(): if event.type == pygame.QUIT: pygame.quit() quit() elif event.type == pygame.KEYDOWN: if event.key == pygame.K_q: quit() elif event.key == pygame.K_SPACE: paused = not paused else: handle_key_press(simulation, event.key) delta = clock.tick(60) / 1000 #60fps #are we using delta in the simulation tick everywhere needed? if not paused: #tick with pause messes up delta TODO: TODOODODODODODOOD TODO print("delta: " + str(delta)) simulation.tick(delta=delta) #TODO: draw floor, flame (continuity) #TODO: do max load on rocket so it blows up #TODO: allow for x movement, speed, accel etc #TODO: allow multilanguage api for landing algorithms etc #TODO: probs need cloud sprite that spawns and despawns as well as floor sprite def linear_gradient(start_color, end_color, length, value_at): return [ int(start_color[j] + (float(value_at)/(length-1))*(end_color[j]-start_color[j])) for j in range(3) ] def draw_simulation(simulation_display: type[pygame.Surface], simulation: type[Simulation]) -> None: #draw background def get_color_for_height(height: float) -> (int, int, int): if height < 70000: return linear_gradient((31,118,194), (0, 0, 0), 70000, int(height)) else: return (0, 0, 0) #gradient for atmosphere simulation_display.fill(get_color_for_height(simulation.y)) #draw clouds and stars #draw clouds (we need continuity TODO) #if simulation.y < 20000 and randint(0, 100) < 5: # pygame.draw.circle(simulation_display, (255, 255, 255), (randint(0, simulation_display.get_width()), randint(0, simulation_display.get_height())), 30) #draw stars if simulation.y > 30000: for _ in range(100): simulation_display.set_at((randint(0, simulation_display.get_width()), randint(0, simulation_display.get_height())), (255, 255, 255)) #draw stats text font = pygame.font.SysFont("Comic Sans MS", 30) simulation_display.blit(font.render("Simulation time: {:.0f}s".format(simulation.time), False, (255, 255, 255)),(0,0)) simulation_display.blit(font.render("Altitude: {:.0f}m".format(simulation.y), False, (255, 255, 255)),(0,40)) simulation_display.blit(font.render("Speed: {:.0f}m/s".format(simulation.speed_y), False, (255, 255, 255)),(0,80)) simulation_display.blit(font.render("Acceleration: {:.2f}m/s2".format(simulation.acceleration_y), False, (255, 255, 255)),(0,120)) simulation_display.blit(font.render("Thrust: {:.0f}N".format(simulation.rocket.current_stage().current_thrust(simulation.body.g(simulation.universe.G, simulation.y))), False, (255, 255, 255)),(0,160)) simulation_display.blit(font.render("Fuel in stage: {:.0f}kg".format(simulation.rocket.current_stage().fuel_mass), False, (255, 255, 255)),(0,200)) simulation_display.blit(font.render("Stage mass: {:.0f}kg".format(simulation.rocket.current_stage().total_mass()), False, (255, 255, 255)),(0,240)) simulation_display.blit(font.render("Rocket mass: {:.0f}kg".format(simulation.rocket.total_mass()), False, (255, 255, 255)),(0,280)) simulation_display.blit(font.render("Stage number: {:.0f}".format(simulation.rocket.stages_spent), False, (255, 255, 255)),(0,320)) simulation_display.blit(font.render("Throttle: {:.0f}%".format(simulation.rocket.current_stage().throttle), False, (255, 255, 255)),(0,360)) #draw rocket first_stage_height = 90 #TODO first_stage_width = 60 def calculate_rocket_y_based_on_y_speed_accel(display_height: int, rocket_height: int, speed_y: float, accel_y: float) -> int: top = display_height / 5 - (rocket_height / 2) #in the case we are accelerating positively bottom = display_height - (top * 2) return bottom def calculate_rocket_x_based_on_x_speed_accel(display_width: int, rocket_width: int, speed_x: float, accel_x: float) -> int: return display_width / 2 - (rocket_width / 2) rocket_x = calculate_rocket_x_based_on_x_speed_accel(simulation_display.get_width(), first_stage_width, None, None) rocket_y = calculate_rocket_y_based_on_y_speed_accel(simulation_display.get_height(), first_stage_height, simulation.speed_y, simulation.acceleration_y) rocket_color = (244, 67, 54) flame_radius = 10 flame_color = (255, 125, 100) i = simulation.rocket.stages_spent stage_height = first_stage_height / (i + 1) stage_y = rocket_y + first_stage_height - stage_height for _ in simulation.rocket.stages: stage_width = first_stage_width / (i + 1) stage_x = rocket_x + i * (stage_width / 2) pygame.draw.rect(simulation_display, rocket_color, pygame.Rect(stage_x, stage_y, stage_width, stage_height)) stage_y -= stage_height / 2 stage_height /= 2 i += 1 #draw flame if simulation.rocket.current_stage().engines_on and simulation.rocket.current_stage().fuel_mass > 0: pygame.draw.circle(simulation_display, flame_color, (rocket_x + (first_stage_width / 2), rocket_y + first_stage_height + flame_radius), flame_radius) def handle_key_press(simulation, key): if key == pygame.K_x: simulation.rocket.current_stage().engines_on = not simulation.rocket.current_stage().engines_on elif key == pygame.K_z: simulation.rocket.perform_stage_separation(True) elif key == pygame.K_DOWN: current_stage = simulation.rocket.current_stage() if current_stage.throttle > 0: current_stage.throttle -= 1 elif key == pygame.K_UP: current_stage = simulation.rocket.current_stage() if current_stage.throttle < 100: current_stage.throttle += 1 elif key == pygame.K_LEFT: return None #sys.exit(0) elif key == pygame.K_RIGHT: return None #sys.exit(0) if __name__ == "__main__": main(sys.argv)