import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation, PillowWriter
import gradio as gr
import math

# Earth constants
G = 6.67430e-11
M = 5.972e24
R = 6.371e6

ESCAPE_VELOCITY = math.sqrt(2 * G * M / R) / 1000  # km/s

def simulate_escape(v0_kms):
    v0 = v0_kms * 1000
    dt = 1
    steps = 200

    r = R
    v = v0
    r_vals = []

    for _ in range(steps):
        a = -G * M / (r ** 2)
        v += a * dt
        r += v * dt
        r_vals.append(r / 1000)
        if r < R:
            break

    fig, ax = plt.subplots()
    ax.set_xlim(0, len(r_vals))
    ax.set_ylim(R/1000, max(r_vals) + 500)
    ax.set_xlabel("Time (s)")
    ax.set_ylabel("Distance from Earth center (km)")
    ax.set_title("Escape Velocity vs Free Fall (Earth)")

    point, = ax.plot([], [], "ro")

    def update(frame):
        point.set_data([frame], [r_vals[frame]])  # ✅ FIX HERE
        return point,

    ani = FuncAnimation(fig, update, frames=len(r_vals), interval=50)

    output_path = "escape_velocity.gif"
    ani.save(output_path, writer=PillowWriter(fps=20))
    plt.close()

    status = (
        "🚀 Object ESCAPES Earth gravity!"
        if v0_kms >= ESCAPE_VELOCITY
        else "⬇️ Object falls back to Earth (Free Fall)"
    )

    explanation = f"""
**Escape Velocity (Earth):** {ESCAPE_VELOCITY:.2f} km/s  
**Initial Speed:** {v0_kms:.2f} km/s  

{status}

**Physics Explanation**
- Gravity decreases with distance
- Escape occurs when kinetic energy ≥ gravitational potential energy
- Otherwise, the object returns in free fall
"""

    return output_path, explanation

with gr.Blocks() as demo:
    gr.Markdown("## 🌍 Escape Velocity & Free Fall Animation")
    speed = gr.Slider(1, 15, value=8, step=0.5, label="Initial Speed (km/s)")
    btn = gr.Button("Run Simulation")
    img = gr.Image(type="filepath")
    text = gr.Markdown()

    btn.click(simulate_escape, speed, outputs=[img, text])

demo.launch()