1 files changed, 31 insertions, 24 deletions
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