Add support for rocket stages

6 files changed, 155 insertions, 64 deletions

diff --git a/.vscode/settings.json b/.vscode/settings.json
new file mode 100644
index 0000000..86ab3eb
--- /dev/null
+++ b/.vscode/settings.json
@@ -0,0 +1,4 @@
+{
+    "python.linting.mypyEnabled": true,
+    "python.linting.enabled": true
+}
\ No newline at end of file
diff --git a/engine.py b/engine.py
index 26a5bab..bc08a43 100644
--- a/engine.py
+++ b/engine.py
@@ -1,7 +1,13 @@
 import fuel
 
 class Engine():
-    def __init__(self, name: str, thrust: int, flow_rate: float):
+    def __init__(self, name: str, isp: int, max_flow_rate: int):
         self.name = name
-        self.thrust = thrust
-        self.flow_rate = flow_rate
\ No newline at end of file
+        self.max_flow_rate = max_flow_rate
+        self.isp = isp
+
+    def thrust(self, throttle: int):
+        return self.flow_rate(throttle) * self.isp
+
+    def flow_rate(self, throttle: int):
+        return self.max_flow_rate * (throttle / 100)
\ No newline at end of file
diff --git a/main.py b/main.py
index 65bdf1d..f94780a 100644
--- a/main.py
+++ b/main.py
@@ -4,6 +4,7 @@ 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
@@ -14,14 +15,33 @@ import pygame
 from pygame.locals import *
 
 def main(argv):
-    rocket = Rocket(name="starship", 
-                    rocket_mass=240000, #thrust=245000
-                    engine=Engine(name="raptor", thrust=2.3E6, flow_rate=1000), #https://en.wikipedia.org/wiki/SpaceX_Raptor
-                    engine_number=33,
-                    fuel_type=Fuel(name="methane", energy_density=None),
-                    fuel_mass=4000000,
-                    drag_coefficient=1.18,
-                    cross_sectional_area=(math.pi * (9**2))
+    #add engine mass
+    raptor_engine = Engine(name="raptor", isp=360, max_flow_rate=9000) #https://en.wikipedia.org/wiki/SpaceX_Raptor
+    methane_fuel = Fuel(name="methane", energy_density=None) #TODO: density
+
+    #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]
                     )
     
     body = Body(name="earth",
@@ -40,6 +60,7 @@ def main(argv):
                         )
     
     simulation = Simulation(universe, body, rocket)
+    simulation.rocket.current_stage().engines_on = True
 
     pygame.init()
     pygame.display.set_caption("OSLS - Overly Simple Launch Simulator")
@@ -49,9 +70,11 @@ def main(argv):
     SCREEN_HEIGHT = 720
 
     simulation_display = pygame.display.set_mode((SCREEN_WIDTH,SCREEN_HEIGHT))
-    while(True):
-        draw_simulation(simulation_display, simulation)
-        pygame.display.update()
+    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:  
@@ -60,15 +83,18 @@ def main(argv):
             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?
-        print("delta: " + str(delta))
-        simulation.tick(delta=delta)
+        if not paused: #tick with pause messes up delta TODO: TODOODODODODODOOD TODO
+            print("delta: " + str(delta))
+            simulation.tick(delta=delta)
 
-        #TODO: draw floor, flame
-        #TODO: add support for rocket stages
+        #TODO: support rocket engine mass
+        #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
@@ -101,15 +127,20 @@ def draw_simulation(simulation_display: type[pygame.Surface], simulation: type[S
         #draw stats text
         font = pygame.font.SysFont("Comic Sans MS", 30)
 
-        simulation_display.blit(font.render("Time: {:.0f}s".format(simulation.time), False, (255, 255, 255)),(0,0))
+        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("Fuel: {:.0f}kg".format(simulation.rocket.fuel_mass), False, (255, 255, 255)),(0,160))
+        simulation_display.blit(font.render("Thrust: {:.0f}N".format(simulation.rocket.current_stage().current_thrust()), 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
-        rocket_height = 90
-        rocket_width = 60
+        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
@@ -120,21 +151,42 @@ def draw_simulation(simulation_display: type[pygame.Surface], simulation: type[S
         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(), rocket_width, None, None)
-        rocket_y = calculate_rocket_y_based_on_y_speed_accel(simulation_display.get_height(), rocket_height, simulation.speed_y, simulation.acceleration_y)
+        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)
 
-        pygame.draw.rect(simulation_display, rocket_color, pygame.Rect(rocket_x, rocket_y, rocket_width, rocket_height))
-        if simulation.rocket.engines_on and simulation.rocket.fuel_mass > 0:
-            pygame.draw.circle(simulation_display, flame_color, (rocket_x + (rocket_width / 2), rocket_y + rocket_height + flame_radius), flame_radius)
+        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_SPACE:
-        simulation.rocket.engines_on = not simulation.rocket.engines_on
+    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:
         sys.exit(0)
     elif key == pygame.K_RIGHT:
diff --git a/rocket.py b/rocket.py
index eb73c7c..97aaa59 100644
--- a/rocket.py
+++ b/rocket.py
@@ -1,33 +1,40 @@
-from engine import Engine
-from fuel import Fuel
+from stage import Stage
 
 class Rocket():
-    def __init__(self, name: str, rocket_mass: int, engine: type[Engine], engine_number: int, fuel_type: type[Fuel], fuel_mass: int, drag_coefficient: float, cross_sectional_area: float):
+    def __init__(self, name: str, stages: [type[Stage]]):
         self.name = name
-        self.rocket_mass = rocket_mass
-        self.engine = engine
-        self.engine_number = engine_number
-        self.fuel_type = fuel_type
-        self.fuel_mass = fuel_mass
-        self.drag_coefficient = drag_coefficient
-        self.cross_sectional_area = cross_sectional_area
+        self.stages = stages
+        self.stages_spent = 0
 
-        self.engines_on = True
+    def current_stage(self) -> type[Stage]:
+        return self.stages[0]
+
+    def top_stage(self) -> type[Stage]:
+        return self.stages[len(self.stages) - 1] #TODO: drag coef and cross sectional area of top stage
+
+    def perform_stage_separation(self, engines_on: bool):
+        if len(self.stages) > 1:
+            self.stages.pop(0)
+            self.stages_spent += 1
+            self.current_stage().engines_on = engines_on
 
     def total_mass(self):
-        return self.rocket_mass + self.fuel_mass
-
-    def total_thrust(self):
-        if(self.engines_on):
-            return self.engine.thrust * self.engine_number
-        else:
-            return 0
-
-    def total_fuel_used(self, delta: int):
-        if(self.engines_on):
-            return self.engine.flow_rate * self.engine_number * delta
-        else:
-            return 0
+        total_mass = 0
+        for stage in self.stages:
+            total_mass += stage.total_mass()
+        return total_mass
+    
+    def total_fuel(self):
+        fuel_mass = 0
+        for stage in self.stages:
+            fuel_mass += stage.fuel_mass
+        return fuel_mass
+
+    def s_cross_sectional_area(self):
+        return self.top_stage().cross_sectional_area
+
+    def s_drag_coefficient(self):
+        return self.top_stage().drag_coefficient
 
     def __str__(self):
         return "eue"
\ No newline at end of file
diff --git a/simulation.py b/simulation.py
index a382f63..4d6fc73 100644
--- a/simulation.py
+++ b/simulation.py
@@ -23,18 +23,18 @@ class Simulation():
 
     #simulation logic
     def tick(self, delta: int) -> None:
+        current_stage = self.rocket.current_stage()
         #calculate upwards force by fuel       
-        # TODO able to turn engine on and off 
-        fuel_used = self.rocket.total_fuel_used(delta)
-        if self.rocket.fuel_mass < fuel_used:
-            fuel_used = self.rocket.fuel_mass
-        self.rocket.fuel_mass -= fuel_used
-        print("Fuel remaining: " + str(self.rocket.fuel_mass))
+        fuel_used = current_stage.total_fuel_used(delta)
+        if current_stage.fuel_mass < fuel_used:
+            fuel_used = current_stage.fuel_mass
+        current_stage.fuel_mass -= fuel_used
+        print("Fuel remaining: " + str(current_stage.fuel_mass))
         
         #upwards_force = fuel_used * self.rocket.fuel_type.energy_density #we should calculate thrust based on this
         upwards_force = 0
         if fuel_used > 0:
-            upwards_force = self.rocket.total_thrust()
+            upwards_force = current_stage.current_thrust()
         print("Upwards force: " + str(upwards_force))
 
         print("g: " + str(self.body.g(G=self.universe.G, height=self.y)))
@@ -49,8 +49,11 @@ class Simulation():
         print("Atmosphere density: " + str(self.body.atmosphere.density_at_height(self.y, self.body.g(G=self.universe.G, height=self.y))))
 
         #https://www.grc.nasa.gov/www/k-12/airplane/drageq.html
-        drag_force = (1/2) * self.body.atmosphere.density_at_height(self.y, self.body.g(G=self.universe.G, height=self.y)) * (self.speed_y ** 2) * self.rocket.drag_coefficient * self.rocket.cross_sectional_area
-        print("Drag: " + str(drag_force)) #drag can be negative too?
+        drag_force = (1/2) * self.body.atmosphere.density_at_height(self.y, self.body.g(G=self.universe.G, height=self.y)) * (self.speed_y ** 2) * self.rocket.s_drag_coefficient() * self.rocket.s_cross_sectional_area()
+        #drag goes against speed
+        if self.speed_y < 0:
+            drag_force *= -1
+        print("Drag: " + str(drag_force))
 
         downwards_force = gravitational_force + drag_force #shouldnt delta influence, TODO: WAIT DRAG COULD BE POSITIVE OR NEGATIVE
         print("Downwards force: " + str(downwards_force))
diff --git a/stage.py b/stage.py
index 7cd2439..783e1f4 100644
--- a/stage.py
+++ b/stage.py
@@ -2,14 +2,33 @@ from engine import Engine
 from fuel import Fuel
 
 class Stage():
-    def __init__(self, stage_mass: int, engine: type[Engine], engine_number: int, fuel_type: type[Fuel], fuel_mass: int, drag_coefficient: float, cross_sectional_area: float):
+    def __init__(self, name: str, stage_mass: int, engine: type[Engine], engine_number: int, fuel_type: type[Fuel], fuel_mass: int, drag_coefficient: float, cross_sectional_area: float):
+        self.name = name
         self.stage_mass = stage_mass
         self.engine = engine
-        self.engine_number = engine_number
+        self.engine_number = engine_number 
         self.fuel_type = fuel_type
         self.fuel_mass = fuel_mass
         self.drag_coefficient = drag_coefficient
         self.cross_sectional_area = cross_sectional_area
+        
+        self.throttle = 100
+        self.engines_on = False
+
+    def total_mass(self):
+        return (self.stage_mass + self.fuel_mass)
+
+    def current_thrust(self):
+        if(self.engines_on and self.fuel_mass > 0):
+            return self.engine.thrust(self.throttle) * self.engine_number
+        else:
+            return 0
+
+    def total_fuel_used(self, delta: int):
+        if(self.engines_on):
+            return self.engine.flow_rate(self.throttle) * self.engine_number * delta
+        else:
+            return 0
 
     #total drag coefficient is just the upper stage
     #engines on is just the lower stage