about summary refs log tree commit diff
path: root/rocket.py
blob: 97aaa5921fc4f1b4788c16448cfc23d1b195c9f5 (plain) (blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
from stage import Stage

class Rocket():
    def __init__(self, name: str, stages: [type[Stage]]):
        self.name = name
        self.stages = stages
        self.stages_spent = 0

    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):
        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"