# Import libraries
import math
import numpy as np
import matplotlib.pyplot as plt
import gradio as gr

# Constants
G = 6.67430e-11      # Gravitational constant (m^3 kg^-1 s^-2)
parsec = 3.0857e16   # One parsec in meters
kpc_in_meters = 1e3 * parsec  # One kiloparsec in meters
M_sun = 1.98847e30   # Mass of the sun in kg

# Function to calculate gravitational force between two masses
def gravitational_force(m1, m2, r):
    return G * m1 * m2 / r**2

# Function to convert galactic coordinates to Cartesian coordinates
def galactic_to_cartesian(l, b, R):
    # Convert angles from degrees to radians
    l_rad = np.radians(l)
    b_rad = np.radians(b)
    x = R * np.cos(b_rad) * np.cos(l_rad)
    y = R * np.cos(b_rad) * np.sin(l_rad)
    z = R * np.sin(b_rad)
    return x, y, z

# Function to run the simulation and generate plots
def run_simulation(mass, l1, b1, l2, b2):
    # Fixed parameters
    travel_time = 5  # Travel time in seconds
    R = 8.2 * kpc_in_meters  # Distance from Galactic center to the object in meters

    # Mass of the black hole (Sagittarius A*)
    black_hole_mass = 4e6 * M_sun  # Mass of Sagittarius A*

    # Convert galactic coordinates to Cartesian coordinates
    x1, y1, z1 = galactic_to_cartesian(l1, b1, R)
    x2, y2, z2 = galactic_to_cartesian(l2, b2, R)

    # Calculate distance between starting and ending points
    dx_total = x2 - x1
    dy_total = y2 - y1
    dz_total = z2 - z1
    distance = np.sqrt(dx_total**2 + dy_total**2 + dz_total**2)

    # Calculate required velocity components
    velocity_x = dx_total / travel_time
    velocity_y = dy_total / travel_time
    velocity_z = dz_total / travel_time
    velocity_magnitude = np.sqrt(velocity_x**2 + velocity_y**2 + velocity_z**2)

    # Time array
    time_steps = np.linspace(0, travel_time, num=100)

    # Initialize arrays for position, velocity, and force
    x_positions = x1 + velocity_x * time_steps
    y_positions = y1 + velocity_y * time_steps
    z_positions = z1 + velocity_z * time_steps
    positions = np.sqrt(x_positions**2 + y_positions**2 + z_positions**2)

    # Calculate gravitational force at each time step
    forces = gravitational_force(mass, black_hole_mass, positions)

    # Plotting the variables over time
    fig, axs = plt.subplots(3, 1, figsize=(12, 8))

    # Position Plot
    axs[0].plot(time_steps, positions / parsec, label='Distance from Galactic Center')
    axs[0].set_ylabel('Distance (parsecs)')
    axs[0].set_title('Object Position Over Time')
    axs[0].legend()

    # Velocity Plot (constant in this case)
    axs[1].plot(time_steps, [velocity_magnitude]*len(time_steps), label='Velocity', color='orange')
    axs[1].set_ylabel('Velocity (m/s)')
    axs[1].set_title('Object Velocity Over Time')
    axs[1].legend()

    # Gravitational Force Plot
    axs[2].plot(time_steps, forces, label='Gravitational Force', color='green')
    axs[2].set_xlabel('Time (s)')
    axs[2].set_ylabel('Force (N)')
    axs[2].set_title('Gravitational Force Over Time')
    axs[2].legend()

    plt.tight_layout()

    # Return the required velocity and the plot figure
    required_velocity_str = f"Required Velocity: {velocity_magnitude:.2e} m/s"
    return required_velocity_str, fig

# Physics equations in LaTeX format
physics_equations = r"""
### Physics Equations Used:

1. **Gravitational Force:**

   $$ F = \frac{G \cdot m_1 \cdot m_2}{r^2} $$

   - \( F \): Gravitational force (N)
   - \( G \): Gravitational constant (\(6.67430 \times 10^{-11}\) m\(^3\) kg\(^{-1}\) s\(^{-2}\))
   - \( m_1, m_2 \): Masses of the objects (kg)
   - \( r \): Distance between the centers of the two masses (m)

2. **Required Velocity:**

   $$ v = \frac{d}{t} $$

   - \( v \): Velocity (m/s)
   - \( d \): Distance traveled (m)
   - \( t \): Time taken (s)

3. **Distance Between Two Points in 3D Space:**

   $$ d = \sqrt{ (x_2 - x_1)^2 + (y_2 - y_1)^2 + (z_2 - z_1)^2 } $$

   - \( (x_1, y_1, z_1) \): Coordinates of the starting point
   - \( (x_2, y_2, z_2) \): Coordinates of the ending point

4. **Conversion from Galactic Coordinates to Cartesian Coordinates:**

   $$ \begin{aligned}
   x &= R \cos b \cos l \\
   y &= R \cos b \sin l \\
   z &= R \sin b
   \end{aligned} $$

   - \( l \): Galactic longitude (radians)
   - \( b \): Galactic latitude (radians)
   - \( R \): Distance from the Galactic center (m)
"""

# Define the Gradio interface
interface = gr.Interface(
    fn=run_simulation,
    inputs=[
        gr.Number(value=1e5, label="Mass of the object (kg)"),
        gr.Number(value=0, label="Starting Galactic Longitude (degrees)"),
        gr.Number(value=0, label="Starting Galactic Latitude (degrees)"),
        gr.Number(value=10, label="Ending Galactic Longitude (degrees)"),
        gr.Number(value=5, label="Ending Galactic Latitude (degrees)")
    ],
    outputs=[
        gr.Textbox(label="Required Velocity"),
        gr.Plot(label="Simulation Plots")
    ],
    title="AI of DornHai Bracelets (Inspired by Taklee Genesis film)",
    description=(
        "Simulate variables during interstellar travel with advanced technology.\n\n"
        + physics_equations
    )
)

# Launch the interface
if __name__ == "__main__":
    interface.launch()