import streamlit as st
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
import streamlit.components.v1 as components

# ===============================
# Streamlit Interface Setup
# ===============================
st.title("Spin Launch & Orbital Attachment Simulation (Animated)")

st.markdown(
    """
    This simulation demonstrates a two‑phase launch:
    
    1. **Spin‑Launch Phase:** A payload is accelerated from an elevated platform.
    2. **Orbital Phase:** At a specified altitude, the payload “docks” with a human‑rated upper stage by
       instantly setting its velocity to that of a circular orbit.
       
    Adjust the parameters on the sidebar and click **Run Simulation**.
    """
)

# Sidebar parameters for simulation:
st.sidebar.header("Simulation Parameters")

# Time and integration parameters:
dt = st.sidebar.number_input("Time Step (s)", min_value=0.01, max_value=1.0, value=0.1, step=0.01)
t_max = st.sidebar.number_input("Total Simulation Time (s)", min_value=100, max_value=5000, value=1000, step=100)

# Launch conditions:
initial_altitude_km = st.sidebar.number_input("Initial Altitude (km)", min_value=1, max_value=100, value=50, step=1)
initial_velocity = st.sidebar.number_input("Initial Velocity (m/s)", min_value=100, max_value=5000, value=1200, step=100)

# Docking/Orbital attachment altitude:
docking_altitude_km = st.sidebar.number_input("Docking Altitude (km)", min_value=1, max_value=500, value=100, step=1)

# Button to run simulation:
run_sim = st.sidebar.button("Run Simulation")

if run_sim:
    st.info("Running simulation... Please wait.")
    # ===============================
    # Physical Constants and Conversions
    # ===============================
    mu = 3.986e14       # Earth's gravitational parameter, m^3/s^2
    R_E = 6.371e6       # Earth's radius, m

    # Convert altitude values from km to m:
    initial_altitude = initial_altitude_km * 1000.0  # initial altitude above Earth's surface (m)
    docking_altitude = docking_altitude_km * 1000.0    # docking altitude above Earth's surface (m)
    r_dock = R_E + docking_altitude                    # docking radius from Earth's center (m)

    # ===============================
    # Initial Conditions for the Simulation
    # ===============================
    # Starting at a point along the x-axis at a radial distance (Earth's radius + initial altitude)
    initial_position = np.array([R_E + initial_altitude, 0.0])
    # Initial velocity is chosen to be radial (pointing outward) at the chosen speed.
    initial_velocity_vec = np.array([initial_velocity, 0.0])

    # ===============================
    # Define the Acceleration Function
    # ===============================
    def acceleration(pos):
        """Compute acceleration due to Earth's gravity at position pos (in m)."""
        r = np.linalg.norm(pos)
        return -mu * pos / r**3

    # ===============================
    # Run the Simulation Loop
    # ===============================
    # We'll record (time, x, y, phase)
    # phase = 1: spin‑launch phase
    # phase = 2: orbital phase (after docking)
    states = []     
    phase = 1
    t = 0.0
    pos = initial_position.copy()
    vel = initial_velocity_vec.copy()
    docking_done = False
    docking_event_time = None
    docking_event_coords = None

    while t < t_max:
        # --- Check for the Docking Event ---
        # When the payload's distance from Earth's center exceeds the docking radius
        # and it's still moving outward, trigger the docking event.
        if phase == 1 and not docking_done and np.linalg.norm(pos) >= r_dock and vel.dot(pos) > 0:
            docking_done = True
            phase = 2  # switch to orbital phase
            r_current = np.linalg.norm(pos)
            # Compute the circular orbital speed at the current radius:
            v_circ = np.sqrt(mu / r_current)
            # For a prograde orbit, choose a tangential (perpendicular) direction:
            tangential_dir = np.array([-pos[1], pos[0]]) / r_current
            vel = v_circ * tangential_dir  # instantaneous burn to circular orbit
            docking_event_time = t
            docking_event_coords = pos.copy()
            st.write(f"**Docking Event:** t = {t:.1f} s, Altitude = {(r_current - R_E)/1000:.1f} km, Circular Speed = {v_circ:.1f} m/s")
        
        # Record the current state: (time, x, y, phase)
        states.append((t, pos[0], pos[1], phase))
        
        # --- Propagate the State (Euler Integration) ---
        a = acceleration(pos)
        pos = pos + vel * dt
        vel = vel + a * dt
        t += dt

    # Convert the recorded states to a NumPy array for easier slicing.
    states = np.array(states)  # columns: time, x, y, phase

    # ===============================
    # Create the Animation Using Matplotlib
    # ===============================
    fig, ax = plt.subplots(figsize=(8, 8))
    ax.set_aspect('equal')
    ax.set_xlabel("x (km)")
    ax.set_ylabel("y (km)")
    ax.set_title("Trajectory: Spin Launch & Orbital Attachment")

    # Draw Earth as a blue circle (Earth's radius in km)
    earth = plt.Circle((0, 0), R_E / 1000, color='blue', alpha=0.3, label="Earth")
    ax.add_artist(earth)

    # Set plot limits to show the trajectory (e.g., up to Earth's radius + 300 km)
    max_extent = (R_E + 300e3) / 1000  # in km
    ax.set_xlim(-max_extent, max_extent)
    ax.set_ylim(-max_extent, max_extent)

    # Initialize the trajectory line and payload marker for animation:
    trajectory_line, = ax.plot([], [], 'r-', lw=2, label="Trajectory")
    payload_marker, = ax.plot([], [], 'ko', markersize=5, label="Payload")
    time_text = ax.text(0.02, 0.95, '', transform=ax.transAxes, fontsize=10)

    def init():
        trajectory_line.set_data([], [])
        payload_marker.set_data([], [])
        time_text.set_text('')
        return trajectory_line, payload_marker, time_text

    def update(frame):
        # Update the trajectory with all positions up to the current frame.
        t_val = states[frame, 0]
        x_vals = states[:frame+1, 1] / 1000  # convert m to km
        y_vals = states[:frame+1, 2] / 1000  # convert m to km
        trajectory_line.set_data(x_vals, y_vals)
        # Wrap coordinates in a list so that they are interpreted as sequences:
        payload_marker.set_data([states[frame, 1] / 1000], [states[frame, 2] / 1000])
        time_text.set_text(f"Time: {t_val:.1f} s")
        return trajectory_line, payload_marker, time_text

    # Create the animation. Adjust the interval (ms) for playback speed.
    anim = FuncAnimation(fig, update, frames=len(states), init_func=init, interval=20, blit=True)

    # Convert the animation to an HTML5 video.
    video_html = anim.to_html5_video()

    if video_html:
        st.markdown("### Simulation Animation")
        # Option 1: Use streamlit components to embed the HTML video
        components.html(video_html, height=500)
        # Option 2: Alternatively, use st.markdown (uncomment the following line to try it)
        # st.markdown(video_html, unsafe_allow_html=True)
    else:
        st.error("No video generated. Please ensure that ffmpeg is installed and properly configured.")

    st.markdown(
        """
        **Note:** This is a highly simplified simulation. In a real-world scenario, the spin‑launch, docking, 
        and orbital insertion phases would involve much more complex physics including aerodynamics, non‑instantaneous burns,
        and detailed guidance and control.
        """
    )