Started preparation to implement x movement

1 files changed, 14 insertions, 4 deletions

diff --git a/stage.py b/stage.py
index 0df1d52..b3140c8 100644
--- a/stage.py
+++ b/stage.py
@@ -1,8 +1,10 @@
+import math
+
 from engine import Engine
 from fuel import Fuel
 
 class Stage():
-    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):
+    def __init__(self, name: str, stage_mass: int, engine: type[Engine], engine_number: int, max_engine_gimbaling_angle: 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
@@ -12,17 +14,22 @@ class Stage():
         self.drag_coefficient = drag_coefficient
         self.cross_sectional_area = cross_sectional_area
         
+        self.max_engine_gimbaling_angle = max_engine_gimbaling_angle
+        self.gimbal = 0
         self.throttle = 100
         self.engines_on = False
 
     def total_mass(self):
         return (self.stage_mass + self.fuel_mass)
 
-    def current_thrust(self, g: float):
+    def current_thrust(self, g: float) -> (float, float):
         if(self.engines_on and self.fuel_mass > 0):
-            return self.engine.thrust(self.throttle, g) * self.engine_number
+            total_thrust = self.engine.thrust(self.throttle, g) * self.engine_number
+            thrust_x = math.fabs(math.sin(math.radians(self.gimbal)) * total_thrust)
+            thrust_y = math.fabs(math.cos(math.radians(self.gimbal)) * total_thrust)
+            return (thrust_x, thrust_y)
         else:
-            return 0
+            return (0, 0)
 
     def total_fuel_used(self, delta: int):
         if(self.engines_on):
@@ -30,6 +37,9 @@ class Stage():
         else:
             return 0
 
+    def convert_y_component_to_total_with_gimbal(self, y_value):
+        return math.fabs(y_value / math.cos(math.radians(self.gimbal)))
+
     #total drag coefficient is just the upper stage
     #engines on is just the lower stage
     #thrust is just the lower stage