# Import necessary packages
import simpy
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D  # for 3D plotting
from PIL import Image
import gradio as gr
import io

# Function for jet flight simulation using SimPy with 3D plot
def run_jet_simulation(max_altitude, climb_rate, acceleration, max_speed_kmph):
    max_speed_mps = max_speed_kmph * 1000 / 3600

    time_log = []
    altitude_log = []
    speed_log = []

    def jet_process(env):
        altitude = 0
        speed = 0
        while True:
            if altitude < max_altitude:
                altitude += climb_rate
            if speed < max_speed_mps:
                speed += acceleration
            time_log.append(env.now)
            altitude_log.append(min(altitude, max_altitude))
            speed_log.append(min(speed, max_speed_mps))
            yield env.timeout(1)
            if altitude >= max_altitude and speed >= max_speed_mps:
                break

    env = simpy.Environment()
    env.process(jet_process(env))
    env.run(until=300)

    # 3D Plotting
    fig = plt.figure(figsize=(8, 6))
    ax = fig.add_subplot(111, projection='3d')
    ax.plot(time_log, altitude_log, speed_log, color='blue', label='Jet Climb Path')
    ax.set_xlabel('Time (s)')
    ax.set_ylabel('Altitude (m)')
    ax.set_zlabel('Speed (m/s)')
    ax.set_title('🛩️ 3D Jet Climb Simulation')
    ax.legend()

    # Save to buffer
    buf = io.BytesIO()
    plt.savefig(buf, format='png')
    buf.seek(0)
    return Image.open(buf)

# Gradio Interface
demo = gr.Interface(
    fn=run_jet_simulation,
    inputs=[
        gr.Number(label="Max Altitude (meters)", value=10000),
        gr.Number(label="Climb Rate (m/s)", value=50),
        gr.Number(label="Acceleration (m/s²)", value=20),
        gr.Number(label="Max Speed (km/h)", value=2500),
    ],
    outputs=gr.Image(type="pil"),
    title="🛩️ 3D Fighter Jet Simulation using SimPy",
    description="Simulates a fighter jet’s altitude and speed climb in 3D space using SimPy and Matplotlib."
)

demo.launch()