import numpy as np
from matplotlib import pyplot as plt


def cal_fuel_consumption_ratio(thrust, ISP):
    return float(thrust) / (9.80665 * float(ISP))


def get_thrust(fuel_percentage):
    global thrust_curve
    fuel_level = [i for i in thrust_curve]
    for i in range(len(fuel_level) - 1):
        if fuel_level[i] >= fuel_percentage > fuel_level[i + 1]:
            prev_fuel = fuel_level[i]
            nex_fuel = fuel_level[i + 1]
            prev_thrust = thrust_curve[prev_fuel]
            next_thrust = thrust_curve[nex_fuel]
            if prev_thrust == next_thrust:
                return prev_thrust
            k = (prev_thrust - next_thrust) / (prev_fuel - nex_fuel)
            result = prev_thrust - k * abs(prev_fuel - fuel_percentage)
            return result
    return 0


def output_thrust_curve():
    global thrust_curve
    print("thrustCurve")
    print("{")
    for item in thrust_curve:
        print(f'    key = {item} {thrust_curve[item]}')
    print("}")

"""
thrust_curve = {
    1.0000: 0.95,
    0.9997: 0.95,
    0.9500: 0.98,
    0.8000: 1.0,
    0.7500: 0.98,
    0.7000: 0.95,
    0.6500: 0.9,
    0.6000: 0.85,
    0.5500: 0.8,
    0.5000: 0.75,
    0.4700: 0.764,
    0.4400: 0.778,
    0.4000: 0.7966666666666666,
    0.3500: 0.82,
    0.3200: 0.826,
    0.3000: 0.83,
    0.2700: 0.826,
    0.2500: 0.84,
    0.2300: 0.824,
    0.2000: 0.80,
    0.1700: 0.77,
    0.1500: 0.75,
    0.1300: 0.67,
    0.1000: 0.55,
    0.0900: 0.5,
    0.0800: 0.45,
    0.0700: 0.40,
    0.0600: 0.35714285714285715,
    0.0500: 0.3142857142857143,
    0.0400: 0.27142857142857146,
    0.0350: 0.25,
    0.0300: 0.23,
    0.0200: 0.19,
    0.0150: 0.17,
    0.0100: 0.15,
    0.0050: 0.12916666666666665,
    0.0028: 0.12,
    0.0018: 0.08,
    0.0005: 0.05,
    0.0000: 0.03}
"""
f = open('data/thrust_curve.txt', 'r')
char = f.readlines()
f.close()

thrust_curve = {}
for line in char:
    temp = line.strip().replace('key = ', '').split(' ')
    thrust_curve[float(temp[0])] = float(temp[1])
    # print(temp)


print(thrust_curve)


dry_mass = 90
fuel_mass = 500
thrust = 14750
ISP = 268

plt.subplot(2, 2, 1)
plt.title("Thrust Level to Fuel Comsumption")
x = [1 - i for i in thrust_curve]
y = [thrust_curve[i] for i in thrust_curve]

plt.plot(x, y)

plt.ylabel('Thrust Level')

plt.subplot(2, 2, 2)
plt.title("Acceleraion to Fuel Comsumption")
twr = [thrust * thrust_curve[i] / (dry_mass + i * fuel_mass) for i in thrust_curve]

plt.plot(x, twr)

plt.xlabel('Fuel Consumed')
plt.ylabel('Acceleraion')

remaining_fuel_mass = fuel_mass

time_list = []
thrust_list = []
acc_list = []

now_step_time = 0
now_step_thrust = 1
now_step_fuel_percentage = 1
while remaining_fuel_mass >= 0:
    time_list.append(now_step_time)
    thrust_list.append(get_thrust(now_step_fuel_percentage))
    now_step_acc = thrust_list[-1] * thrust / (dry_mass + fuel_mass * now_step_fuel_percentage)
    acc_list.append(now_step_acc)
    # print(now_step_time, now_step_thrust, now_step_fuel_percentage, now_step_acc, remaining_fuel_mass)
    now_step_time += 0.01
    remaining_fuel_mass -= cal_fuel_consumption_ratio(thrust * now_step_thrust, ISP) * 0.01
    now_step_fuel_percentage = remaining_fuel_mass / fuel_mass
    now_step_thrust = get_thrust(now_step_fuel_percentage)
    if now_step_acc == 0:
        print('error')
        print('remaining_fuel_mass', remaining_fuel_mass)
        print('now_step_fuel_percentage', now_step_fuel_percentage)
        break
plt.subplot(2, 2, 4)
plt.title("Acceleraion to Time")

plt.plot(time_list, acc_list)

plt.xlabel('Time in second')
plt.ylabel('Acceleraion')

plt.subplot(2, 2, 3)
plt.title("Thrust to Time")

plt.plot(time_list, thrust_list)

plt.xlabel('Time in second')
plt.ylabel('Thrust level')

plt.show()
print(time_list[-1])
output_thrust_curve()
