import numpy as np
import matplotlib.pyplot as plt

g = 9.81

class RocketEngine:
    def __init__(self, engine_name, engine_mass, engine_isp, engine_atmo_thrust, engine_vac_thrust):
        self.name = engine_name
        self.mass = engine_mass
        self.isp = engine_isp
                self.atmo_thrust = engine_atmo_thrust
                self.vac_thrust = engine_vac_thrust

ant = RocketEngine("Ant", 20, 310, 0.51, 2.0)
spark = RocketEngine("Spark", 100, 300, 16.2, 18.0)
lv909 = RocketEngine("Terrier", 500, 340, 14.783, 60.0)
nerv = RocketEngine("Nerv", 3000, 800, 13.875, 60.0)
nerv2x = RocketEngine("2x Nerv", 6000, 800, 27.5, 120.0)
lightbulb = RocketEngine("Lightbulb", 20000.0, 1500.0, 225.0, 400.0)
reliant = RocketEngine("Reliant", 1250, 320, 200.67, 215.0)
aerospike = RocketEngine("Aerospike", 1000, 340, 153.53, 180.0)
poodle = RocketEngine("Poodle", 1750, 350, 64.29, 250.0)
mainsail = RocketEngine("Mainsail", 6000, 310, 1379.0, 1500.0)
skipper = RocketEngine("Skipper", 3000, 320, 568.75, 650.0)
rhino = RocketEngine("Rhino", 9000, 340, 1500.0, 2000.0)
mammoth = RocketEngine("Mammoth", 15000, 315, 3746.0, 4000.0)
twinboar = RocketEngine("Twin-boar", 10000, 300, 1866.7, 2000.0)
engines = [aerospike,
           reliant,
           poodle,
           mainsail,
           skipper,
           rhino,
           twinboar,
           mammoth,
           nerv,
           ant,
           spark,
           lv909,]

payloads = [10000.0]

hs = []
colors = plt.get_cmap('jet')(np.linspace(0, 1.0, len(engines)))

plt.figure(1)

def how_many_engines(engine, more_than_twr, fuel_mass, payload_mass):
        num_engines = 0
        twr = 0
        while twr < more_than_twr:
            num_engines = num_engines + 1
            twr = (num_engines * engine.vac_thrust) / ((fuel_mass + engine.mass + payload_mass) * g) * 1000
        return 0, 1
#        return twr, num_engines

requested_twr = 0.2

for p, payload_mass in enumerate(payloads):
        print "---------------------------"
        print "payload", payload_mass, "kg"
        print "---------------------------"
        for color, engine in zip(colors, engines):
                fuel_masses = np.arange(0, 10000 * (p+1), 50)
                dv_values = np.zeros(np.size(fuel_masses))

                for i, fuel_mass in enumerate(fuel_masses):
                        twr, num_engines = how_many_engines(engine, requested_twr, fuel_mass, payload_mass)
                        engine_block_mass = engine.mass * num_engines
                        total_mass = fuel_mass + engine_block_mass + payload_mass
                        dv = g * engine.isp * (engine.vac_thrust / engine.vac_thrust) * np.log(total_mass / (total_mass - fuel_mass))
                        dv_values[i] = dv
                        #print engine.name, num_engines, fuel_mass, "kg", twr, dv, "m/s"

                subplot_arg = 110 + (p + 1)
                plt.subplot(subplot_arg)

                twr = engine.vac_thrust / ((fuel_masses[-1:] + engine.mass + payload_mass) * g) * 1000
                label = "{0}".format(engine.name, twr[0])

                h, = plt.plot(dv_values, fuel_masses + payload_mass + engine.mass, label=label, linewidth=3, color=color)
                hs.append(h)

        if p == 0:
                plt.legend(handles=hs[0:len(engines)], loc='best')

        plt.title("payload = {0} kg".format(payload_mass, requested_twr))
        plt.ylabel("total mass [kg]")
        plt.xlabel("delta-v [m/s]")
        plt.grid()

plt.show()