import math
import load_common_resources
import copy
from matplotlib import pyplot as plt

Vc = 299792458
Ga = 9.80665

path = 'CommonResources.cfg'
common_resources_dict = load_common_resources.load_common_resources(path)


fuel_density_map = {
    'methalox': common_resources_dict['LqdMethane']['density'] * 0.4137 + common_resources_dict['LqdOxygen'][
        'density'] * 0.5863,
    'kerolox': common_resources_dict['Kerosene']['density'] * 0.3487 + common_resources_dict['LqdOxygen'][
        'density'] * 0.6513,
    'udmh/nto': common_resources_dict['UDMH']['density'] * 0.46 + common_resources_dict['NTO'][
        'density'] * 0.54,
    'solid': common_resources_dict['PBAN']['density'],
    'hydrolox': common_resources_dict['LqdHydrogen']['density'] * 0.7276 + common_resources_dict['LqdOxygen']['density']
                * 0.2724
}


def cal_fuel_burned(isp, thrust, duration):
    """
    calculate the amount of fuel burned in duration. In KG
    :param isp: in seconds
    :param thrust: in KN
    :param duration: in seconds
    :return:
    """
    fuel_burned = duration * (thrust * 1000 / isp / Ga)
    # fuel_burned *= pow(1 - pow(isp * Ga / Vc, 2), 0.5)  # Adjust for relativity speed
    return fuel_burned


class Stage:
    def __init__(self, fuel_mass, dry_mass, max_thrust, isp, fuel_type, name, throttle_map=None, map_type=None,
                 has_separator=True):
        self.fuel_mass = fuel_mass
        self.total_mass = fuel_mass + dry_mass
        self.fuel_volume = fuel_mass / fuel_density_map[fuel_type]
        self.current_mass = fuel_mass + dry_mass
        self.current_fuel_mass = fuel_mass
        self.dry_mass = dry_mass
        self.max_thrust = max_thrust
        self.fuel_type = fuel_type
        self.isp = isp
        self.is_empty = False
        self.name = name
        self.throttle_map = throttle_map
        self.current_thrust = self.max_thrust
        self.map_type = map_type
        self.has_separator = has_separator
        # demo_map = {
        #     'time_map': {0: 1, 10: 1, 20: 1, 30: 0.7, 40: 0.6, 50: 0.7, 60: 0.8, 70: 1, 90: 1, 110: 0.7},
        #     'fuel_map': {1: 1, 0.5: 1, 0.25: 0.8, 0.2: 0.7, 0.1: 0.6, 0.05: 0.55, 0: 0.55},
        # }
        # self.throttle_map = demo_map

    def cal_current_thrust(self, stage_time=0, stage_acc=0):
        throttle = 1
        if not self.throttle_map:
            return self.max_thrust
        curve = self.throttle_map.get(self.map_type)
        if not curve:
            return self.max_thrust
        if self.map_type == 'time_map':
            curve_keys = list(curve.keys())
            for i in range(len(curve)):
                try:
                    if curve_keys[i] <= stage_time < curve_keys[i + 1]:
                        throttle = curve[curve_keys[i]] + (curve[curve_keys[i+1]] - curve[curve_keys[i]]) / (curve_keys[i+1] - curve_keys[i]) * (stage_time - curve_keys[i])
                        break
                except IndexError:
                    throttle = curve[curve_keys[-1]]
        elif self.map_type == 'fuel_map':
            curve_keys = list(curve.keys())
            flag = self.current_fuel_mass / self.fuel_mass
            for i in range(len(curve)):
                try:
                    if curve_keys[i] > flag >= curve_keys[i + 1]:
                        throttle = curve[curve_keys[i]] + (curve[curve_keys[i + 1]] - curve[curve_keys[i]]) / (
                                    curve_keys[i + 1] - curve_keys[i]) * (flag - curve_keys[i])
                        break
                except IndexError:
                    throttle = curve[curve_keys[-1]]
        current_thrust = throttle * self.max_thrust
        return current_thrust

    def reset_simulation(self):
        self.is_empty = False
        self.current_mass = self.fuel_mass + self.dry_mass
        self.current_fuel_mass = self.fuel_mass

    def simulated_burn_step_by_step(self, stage_time, duration):
        current_thrust = self.cal_current_thrust(stage_time)
        now_thrust = current_thrust
        if self.is_empty:
            return True, self.current_mass
        next_fuel_burned = cal_fuel_burned(self.isp, now_thrust, duration)
        if self.current_fuel_mass - next_fuel_burned < 0:
            self.is_empty = True
            print(f'Part: {self.name} is depleted after {round(stage_time, 3)}s.')
        else:
            self.current_fuel_mass -= next_fuel_burned
        self.current_mass = self.current_fuel_mass + self.dry_mass
        return self.is_empty, self.current_mass


class Vessel:
    def __init__(self, stage_list, name, payload=0, duration=0.01):
        self.stage_list = stage_list
        self.duration = duration
        self.payload = payload
        self.name = name

    def reset_simulation(self):
        for stages in self.stage_list:
            for stage in stages:
                stage.reset_simulation()

    def get_current_mass(self):
        current_mass = 0
        for stages in self.stage_list:
            for stage in stages:
                if not stage.is_empty:
                    current_mass += stage.current_mass
        return current_mass + self.payload

    def get_current_thrust(self):
        current_thrust = 0
        stage_flag = False
        for stages in self.stage_list:
            if stage_flag:
                break
            for stage in stages:
                if stage.is_empty:
                    continue
                else:
                    current_thrust += stage.max_thrust
                    stage_flag = True
        return current_thrust

    @staticmethod
    def is_depleted(stages):
        for stage in stages:
            if not stage.is_empty:
                return False
        return True

    def stage_burn(self, stages, stage_time):
        result = []
        for stage in stages:
            stage_is_empty, stage_mass = stage.simulated_burn_step_by_step(stage_time, self.duration)
            result.append((stage_is_empty, stage_mass))
        return result

    @staticmethod
    def stage_info(stages):
        stage_thrust, stage_fuel, gross_mass, dry_mass = 0, 0, 0, 0
        for stage in stages:
            stage_thrust += stage.max_thrust
            stage_fuel += stage.fuel_mass
            gross_mass += stage.total_mass
            dry_mass += stage.dry_mass
        print(f'  Total Vacuum Thrust: {stage_thrust}KN, Burned Mass: {stage_fuel/1000}ton, '
              f'Gross Mass: {gross_mass/1000}ton, '
              f'Dry Mass: {dry_mass/1000}ton')

    def run_sim_new(self):
        self.vessel_info()
        self.reset_simulation()
        print('Start Simulation')
        stage_burn_time, timer, dv, stage_dv = 0, 0, 0, 0
        now_stage_id = 0
        stage_acc_result = {}

        # time based simulation
        while True:
            try:
                now_stage = self.stage_list[now_stage_id]
                if f'stage{now_stage_id}' not in stage_acc_result:
                    stage_acc_result[f'stage{now_stage_id}'] = []
            except IndexError:
                break
            stage_burn_result = self.stage_burn(now_stage, stage_burn_time)
            timer += self.duration
            stage_burn_time += self.duration
            step_dv = self.get_current_thrust() * 1000 / self.get_current_mass() * self.duration
            step_acc = self.get_current_thrust() * 1000 / self.get_current_mass() / Ga
            stage_acc_result[f'stage{now_stage_id}'].append(step_acc)

            dv += step_dv
            stage_dv += step_dv

            if self.is_depleted(now_stage):
                print(f'stage{now_stage_id}:\n  burn time: {round(stage_burn_time, 3)}s, dv: {round(stage_dv, 3)}m/s ')
                start_twr = round(stage_acc_result[f'stage{now_stage_id}'][0], 3)
                end_twr = round(stage_acc_result[f'stage{now_stage_id}'][-2], 3)
                max_twr = round(max(stage_acc_result[f'stage{now_stage_id}']), 3)
                print(f'  start twr: {start_twr}  end twr: {end_twr}  max_twr: {max_twr}')
                self.stage_info(now_stage)
                stage_burn_time, stage_dv = 0, 0
                now_stage_id += 1
        timer = round(timer, 3)
        dv = round(dv, 3)
        print(f'Total Engine Burn Time: {timer}s, Total dv: {dv}m/s')

    def vessel_info(self):
        lift_off_mass = self.get_current_mass() / 1000
        lift_off_thrust = self.get_current_thrust()
        print(f'Info of {self.name}: ')
        print(f'  Vessel Lift-off Mass: {lift_off_mass} ton, Vessel Lift-off Vacuum Thrust: {lift_off_thrust} KN')
        print(f'  Payload Mass: {self.payload} KG')
        print('-'*80)


if __name__ == '__main__':
    # demo_map = {
    #     'time_map': {0: 1, 10: 1, 20: 1, 29: 1, 30: 0.9, 31: 0.80, 32: 0.79, 33: 0.79, 173: 0.79, 174: 0.8, 175: 0.8,
    #                  176: 0.9, 177: 1, 180: 1, 190: 1},
    #     # 'fuel_map': {1: 1, 0.5: 1, 0.25: 0.8, 0.2: 0.7, 0.1: 0.6, 0.05: 0.55, 0: 0.55},
    # }
    demo_map = {
        'time_map': {0: 1, 10: 1, 20: 1, 59: 1, 60: 0.85, 61: 0.73, 62: 0.73, 63: 0.73, 173: 0.73, 174: 0.73, 175: 0.85,
                     176: 0.9, 177: 1, 180: 1, 190: 1},
    }

    # # CBC: 173s     S1: 227s    S2: 212s    S3: 800s
    # cbc_fuel_mass = 510000
    # cbc_dry_mass = 50000
    # cbc_thrust = 1397 * 7
    # cbc = Stage(cbc_fuel_mass * 2, cbc_dry_mass*2 + 150000, cbc_thrust * 2, 338.2, 'kerolox', 'Side Booster')
    # S1 = Stage(cbc_fuel_mass + 70000, 47000, cbc_thrust, 338.2, 'kerolox', 'Center Core', throttle_map=demo_map, map_type='time_map')
    # S2 = Stage(180000, 15000, 2920, 352.3, 'kerolox', 'YF-100M Kerolox Stage')
    # S3 = Stage(60000, 8000, 276.324, 451, 'hydrolox', 'YF-75E Hydrolox Stage')
    # cz_10 = Vessel([[cbc, S1], [S2], [S3]], name='CZ-10', payload=27000)
    # # cz_10 = Vessel([[cbc, S1], [S2]], name='CZ-10', payload=70000)
    # cz_10.run_sim_new()

    booster = Stage(1040000, 100000, 1397*14, 338.2, 'kerolox', 'Side Booster')
    S1 = Stage(680000, 60000, 1397*7, 338.2, 'kerolox', 'Stage 1')
    S2 = Stage(185000, 20000, 2900, 352.3, 'kerolox', 'Stage 2')
    S3 = Stage(58500, 11500, 276.324, 451, 'hydrolox', 'Stage 3')
    cz_10 = Vessel([[booster, S1], [S2], [S3]], name='CZ-10', payload=27000)
    cz_10.run_sim_new()

    # S1 = Stage(339000, 30000, 1397*4, 338.2, 'kerolox', 'Stage 1')
    # S2 = Stage(48000, 6000, 360, 341.2, 'kerolox', 'Stage 2')
    # cz_12 = Vessel([[S1], [S2]], name='CZ-12', payload=10000)
    # cz_12.run_sim_new()

    booster = Stage(72000*4, 7500*4, 1340*4, 335, 'kerolox', 'CZ-7A Booster')
    Stage1 = Stage(144000, 15000, 1340*2, 335, 'kerolox', 'CZ-7A Stage1')
    Stage2 = Stage(61500, 8000, 180*4, 342, 'kerolox', 'CZ-7A Stage2')
    Stage3 = Stage(18250, 3050, 83*2, 438, 'hydrolox', 'CZ-7A Stage3')

    cz_7a = Vessel([[booster, Stage1], [Stage2], [Stage3]], name='CZ-7A', payload=8000)
    cz_7a.run_sim_new()

    print('done')



