Add initial pygame gui

3 files changed, 82 insertions, 22 deletions

diff --git a/main.py b/main.py
index d322193..1860672 100644
--- a/main.py
+++ b/main.py
@@ -42,43 +42,84 @@ def main(argv):
     simulation = Simulation(universe, body, rocket)
 
     pygame.init()
-    simulation_display = pygame.display.set_mode((400,500))  
+    clock = pygame.time.Clock()
+
+    SCREEN_WIDTH = 1024
+    SCREEN_HEIGHT = 720
+
+    simulation_display = pygame.display.set_mode((SCREEN_WIDTH,SCREEN_HEIGHT))
     while(True):
-        delta = 0.01
-        sleep(delta)
+        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_SPACE:
-                    simulation.rocket.engines_on = not simulation.rocket.engines_on
-                elif event.key == pygame.K_LEFT:
-                    sys.exit(0)
-                elif event.key == pygame.K_RIGHT:
-                    sys.exit(0)
+                if event.key == pygame.K_q:
+                    quit()
+                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)
-        draw_simulation(simulation_display, simulation)
-
-        pygame.display.update()
 
+        #TODO: draw floor, flame
+        #TODO: add support for rocket stages
         #TODO: do max load on rocket so it blows up
-        #TODO: display sim info on screen
         #TODO: allow for x movement, speed, accel etc
-        #TODO: draw floor, flame
         #TODO: allow multilanguage api for landing algorithms etc
 
 def draw_simulation(simulation_display: type[pygame.Surface], simulation: type[Simulation]) -> None:
+
+        font = pygame.font.SysFont("Comic Sans MS", 30)
+        
+        #draw background
+        def get_color_for_height(height: int) -> (int, int, int):
+            return (255, 255, 255)
+        
         simulation_display.fill(get_color_for_height(simulation.y)) #gradient for atmosphere TODO
 
-        pygame.draw.rect(simulation_display, (0, 125, 255), pygame.Rect(30, 30, 60, 60))
-        if simulation.rocket.engines_on:
-            pygame.draw.circle(simulation_display, (255, 125, 100), (60, 100), 10)
+        #draw stats text
+        simulation_display.blit(font.render("Altitude: {:.0f}m".format(simulation.y), False, (0, 0, 0)),(0,0))
+        simulation_display.blit(font.render("Speed: {:.0f}m/s".format(simulation.speed_y), False, (0, 0, 0)),(0,40))
+        simulation_display.blit(font.render("Acceleration: {:.2f}m/s2".format(simulation.acceleration_y), False, (0, 0, 0)),(0,80))
+        simulation_display.blit(font.render("Fuel: {:.0f}kg".format(simulation.rocket.fuel_mass), False, (0, 0, 0)),(0,120))
+
+        #draw rocket
+        rocket_height = 90
+        rocket_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(), 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_color = (0, 125, 255)
+
+        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)
 
-def get_color_for_height(height: int) -> (int, int, int):
-    return (255, 255, 255)
+def handle_key_press(simulation, key):
+    if key == pygame.K_SPACE:
+        simulation.rocket.engines_on = not simulation.rocket.engines_on
+    elif key == pygame.K_LEFT:
+        sys.exit(0)
+    elif key == pygame.K_RIGHT:
+        sys.exit(0)
 
 if __name__ == "__main__":
     main(sys.argv)
\ No newline at end of file
diff --git a/simulation.py b/simulation.py
index a93a263..223f931 100644
--- a/simulation.py
+++ b/simulation.py
@@ -19,6 +19,7 @@ class Simulation():
         self.rocket = rocket
         self.y = 0
         self.speed_y = 0
+        self.acceleration_y = 0
 
     #simulation logic
     def tick(self, delta: int) -> None:
@@ -58,9 +59,9 @@ class Simulation():
         total_force = upwards_force - downwards_force
         print("Total force: " + str(total_force))
 
-        acceleration = total_force / self.rocket.total_mass() #mayb we need momentum??
-        print("Acceleration: " + str(acceleration))
-        self.speed_y = self.speed_y + (acceleration * delta) #i think thir swrong
+        self.acceleration_y = total_force / self.rocket.total_mass() #mayb we need momentum??
+        print("Acceleration: " + str(self.acceleration_y))
+        self.speed_y = self.speed_y + (self.acceleration_y * delta) #i think thir swrong
 
         #update position based on velocity and delta
         self.y += self.speed_y * delta
diff --git a/stage.py b/stage.py
new file mode 100644
index 0000000..7cd2439
--- /dev/null
+++ b/stage.py
@@ -0,0 +1,18 @@
+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):
+        self.stage_mass = stage_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
+
+    #total drag coefficient is just the upper stage
+    #engines on is just the lower stage
+    #thrust is just the lower stage
+    #fuel used is just lower stage
+    #when stage separation lower stage jetissoned
\ No newline at end of file