#!/usr/bin/env python
# coding: utf-8

# # Prometheus - Western Engineering Rocketry Team 2022
# 
# Launched at Spaceport America Cup 2022 edition.  
# Western Engineering Rocketry Team is a rocket design team from Western University in London, Ontario (Canada).
# Permission to use flight data given by Giorgio Chassikos, 2024.
# 
# Website: [https://www.werocketry.com/](https://www.werocketry.com/)
# 
# ## Team Information:
# The Western Engineering Rocketry Team is a student-led engineering design team founded in 2016 at Western University in London, Ontario.  
# Comprised of engineering students from a variety of disciplines, the team's goal is to develop and launch rockets through the Intercollegiate Rocket Engineering Competition.
# 
# ## Flight Card Results:
# 1. **Launch Date:** June 24th, 2022, 9:17am local time
# 2. **Last Simulated Apogee:** 4190.05 m
# 3. **Official Recorded Apogee:** 3898.37 m
# 
# If we calculate the simulation error and divide it by the actual apogee, we get a relative error of 7.48%.
# 

# In[1]:


get_ipython().run_line_magic('load_ext', 'autoreload')
get_ipython().run_line_magic('autoreload', '2')


# In[2]:


from rocketpy import Function, Rocket, GenericMotor, Flight, Environment
from rocketpy.simulation.flight_data_importer import FlightDataImporter

import matplotlib.pyplot as plt
import numpy as np

plt.style.use("seaborn-v0_8-colorblind")


# ## Define the Environment

# In[3]:


env = Environment(
    latitude=32.939377,
    longitude=-106.911986,
    elevation=1401,
)
env.set_date(date=(2023, 6, 24, 9), timezone="America/Denver")

env.set_atmospheric_model(
    type="Reanalysis",
    file="../../data/weather/spaceport_america_pressure_levels_2023_hourly.nc",
    dictionary="ECMWF",
)

env.max_expected_height = 6000


# In[4]:


env.info()


# ## Define the motor

# In[5]:


motor_M1520 = GenericMotor(
    # burn specs: https://www.thrustcurve.org/simfiles/5f4294d20002e900000006b1/
    thrust_source="../../data/motors/cesaroni/Cesaroni_7579M1520-P.eng",
    burn_time=4.897,
    propellant_initial_mass=3.737,
    dry_mass=2.981,
    # casing specs: Pro98 3G Gen2 casing
    chamber_radius=0.064,
    chamber_height=0.548,
    chamber_position=0.274,
    nozzle_radius=0.027,
)


motor_M1520.info()


# ## Define the Rocket

# In[6]:


def prometheus_cd_at_ma(mach):
    """Gives the drag coefficient of the rocket at a given mach number."""
    if mach <= 0.15:
        return 0.422
    elif mach <= 0.45:
        return 0.422 + (mach - 0.15) * (0.38 - 0.422) / (0.45 - 0.15)
    elif mach <= 0.77:
        return 0.38 + (mach - 0.45) * (0.32 - 0.38) / (0.77 - 0.45)
    elif mach <= 0.82:
        return 0.32 + (mach - 0.77) * (0.3 - 0.32) / (0.82 - 0.77)
    elif mach <= 0.88:
        return 0.3 + (mach - 0.82) * (0.3 - 0.3) / (0.88 - 0.82)
    elif mach <= 0.94:
        return 0.3 + (mach - 0.88) * (0.32 - 0.3) / (0.94 - 0.88)
    elif mach <= 0.99:
        return 0.32 + (mach - 0.94) * (0.37 - 0.32) / (0.99 - 0.94)
    elif mach <= 1.04:
        return 0.37 + (mach - 0.99) * (0.44 - 0.37) / (1.04 - 0.99)
    elif mach <= 1.24:
        return 0.44 + (mach - 1.04) * (0.43 - 0.44) / (1.24 - 1.04)
    elif mach <= 1.33:
        return 0.43 + (mach - 1.24) * (0.42 - 0.43) / (1.33 - 1.24)
    elif mach <= 1.49:
        return 0.42 + (mach - 1.33) * (0.39 - 0.42) / (1.49 - 1.33)
    else:
        return 0.39


prometheus = Rocket(
    radius=0.06985,  # 5.5" diameter circle
    mass=13.93,
    inertia=(
        4.87,
        4.87,
        0.05,
    ),
    power_off_drag=prometheus_cd_at_ma,
    power_on_drag=lambda x: prometheus_cd_at_ma(x) * 1.02,  # 5% increase in drag
    center_of_mass_without_motor=0.9549,
    coordinate_system_orientation="tail_to_nose",
)

prometheus.set_rail_buttons(0.69, 0.21, 60)

prometheus.add_motor(motor=motor_M1520, position=0)
nose_cone = prometheus.add_nose(length=0.742, kind="Von Karman", position=2.229)
fin_set = prometheus.add_trapezoidal_fins(
    n=3,
    span=0.13,
    root_chord=0.268,
    tip_chord=0.136,
    position=0.273,
    sweep_length=0.066,
)
drogue = prometheus.add_parachute(
    "Drogue",
    cd_s=1.6 * np.pi * 0.3048**2,  # Cd = 1.6, D_chute = 24 in
    trigger="apogee",
)
main = prometheus.add_parachute(
    "Main",
    cd_s=2.2 * np.pi * 0.9144**2,  # Cd = 2.2, D_chute = 72 in
    trigger=457.2,  # 1500 ft
)


# In[7]:


prometheus.info()
prometheus.draw()
prometheus.plots.drag_curves()


# ## Simulate the Flight

# In[8]:


test_flight = Flight(
    rocket=prometheus,
    environment=env,
    inclination=80,
    heading=75,
    rail_length=5.18,
)


# In[9]:


# test_flight.prints.initial_conditions()
# test_flight.prints.surface_wind_conditions()
# test_flight.prints.launch_rail_conditions()
test_flight.prints.out_of_rail_conditions()
test_flight.prints.burn_out_conditions()
test_flight.prints.apogee_conditions()
# test_flight.prints.events_registered()
# test_flight.prints.impact_conditions()
test_flight.prints.stability_margin()
test_flight.prints.maximum_values()
# test_flight.prints.numerical_integration_settings()


# ## Read the Telemetry Data

# In[10]:


columns_map = {
    "time": "time",
    "altitude": "z",
    "height": "altitude",
    "acceleration": "acceleration",
    "pressure": "pressure",
    "accel_x": "ax",
    "accel_y": "ay",
    "accel_z": "az",
    "latitude": "latitude",
    "longitude": "longitude",
}


cots_altimeter_flight = FlightDataImporter(
    name="Telemetry Mega",
    paths="../../data/prometheus/2022-06-24-serial-5115-flight-0001-TeleMetrum.csv",
    columns_map=columns_map,
    units=None,
    interpolation="linear",
    extrapolation="zero",
    delimiter=",",
    encoding="utf-8",
)


# ## Compare Simulation and Telemetry Data

# In[11]:


real_apogee = cots_altimeter_flight.altitude.max
rocketpy_apogee = test_flight.apogee - test_flight.env.elevation
a_error = abs(real_apogee - rocketpy_apogee)
r_error = a_error / real_apogee * 100

print(f"RocketPy apogee: {rocketpy_apogee:.2f} m")
print(f"Real apogee: {real_apogee:.2f} m")
print(f"Absolute error: {a_error:.2f} m")
print(f"Relative error: {r_error:.2f}%")


# In[13]:


# Altitude comparison
Function.compare_plots(
    [
        (test_flight.altitude, "RocketPy"),
        (cots_altimeter_flight.altitude, "Telemetry"),
    ],
    title="Altitude Comparison",
    xlabel="Time (s)",
    ylabel="Altitude (m)",
    lower=0,
    # upper=30,
)

# Pressure comparison
Function.compare_plots(
    [
        (test_flight.pressure, "RocketPy"),
        (cots_altimeter_flight.pressure, "Telemetry"),
    ],
    title="Pressure Comparison",
    xlabel="Time (s)",
    ylabel="Pressure (Pa)",
    lower=0,
    # upper=30,
)

# Latitude comparison
# Function.compare_plots(
#     [
#         (test_flight.latitude, "RocketPy"),
#         (cots_altimeter_flight.latitude, "Telemetry"),
#     ],
#     title="Latitude Comparison",
#     xlabel="Time (s)",
#     ylabel="Latitude (deg)",
#     lower=0,
#     # upper=30,
# )

# Longitude comparison
# Function.compare_plots(
#     [
#         (test_flight.longitude, "RocketPy"),
#         (cots_altimeter_flight.longitude, "Telemetry"),
#     ],
#     title="Longitude Comparison",
#     xlabel="Time (s)",
#     ylabel="Longitude (deg)",
#     lower=0,
#     # upper=30,
# )

# Convert "deceleration" to "acceleration" for comparison
absolute_cots_acceleration = lambda x: abs(cots_altimeter_flight.acceleration(x))

# Acceleration comparison
Function.compare_plots(
    [
        (test_flight.acceleration, "RocketPy"),
        (Function(absolute_cots_acceleration), "Telemetry"),
    ],
    title="Acceleration Comparison",
    xlabel="Time (s)",
    ylabel="Acceleration (m/s^2)",
    lower=0,
    upper=30,
)