import gradio as gr
import numpy as np
import matplotlib.pyplot as plt
from scipy.integrate import solve_ivp

# Constants
# v0 = 90 * 0.44704  # initial velocity (90 mph to m/s)
# angle_launch = 32  # launch angle in degrees
# angle_spray = 20  # spray angle in degrees
Cd = 0.47  # drag coefficient for a sphere
g = 9.81  # acceleration due to gravity (m/s^2)
rho = 1.225  # air density (kg/m^3)
r = 0.07468  # radius of baseball (m)
m = 0.145  # mass of baseball (kg)
A = np.pi * r**2  # cross-sectional area (m^2)

# Constants in US customary units
# g = 32.174  # acceleration due to gravity (ft/s^2)
# rho = 0.07647  # air density (lb/ft^3)
# r = 0.245  # radius of baseball (ft)
# m = 0.32  # mass of baseball (lb)
# A = np.pi * r**2  # cross-sectional area (ft^2)

# Constants for Dr. Allen M. Nathan's calculations
omega = 1800 * 2 * np.pi / 60  # spin rate in rad/s (assuming 1800 RPM)
Cl = 0.5  # lift coefficient, this is a rough estimate and can vary
m_to_ft = 3.28084

# Magnus force due to spin
def magnus_force(vx, vy, vz, v0):
    Fm_x = Cl * rho * A * (vy * omega - vz * v0)
    Fm_y = Cl * rho * A * (vz * omega - vx * v0)
    Fm_z = Cl * rho * A * (vx * omega - vy * v0)
    return Fm_x, Fm_y, Fm_z

# Drag force
def drag(v):
    return -0.5 * Cd * A * rho * v**2

# Equations of motion
def equations_of_motion(t, y):
    x, y, z, vx, vy, vz, v0 = y
    y = y * m_to_ft
    v = np.sqrt(vx**2 + vy**2 + vz**2)
    Fm_x, Fm_y, Fm_z = magnus_force(vx, vy, vz, v0)
    dxdt = vx
    dydt = vy
    dzdt = vz
    dvxdt = drag(v) * vx / (m * v)
    dvydt = (drag(v) * vy - m * g) / (m * v)
    dvzdt = drag(v) * vz / (m * v)
    # dvxdt = (drag(v) * vx + Fm_x) / (m * v)
    # dvydt = ((drag(v) * vy + Fm_y) - m * g) / (m * v)
    # dvzdt = (drag(v) * vz + Fm_z) / (m * v)
    return [dxdt * m_to_ft, dydt * m_to_ft, dzdt, dvxdt, dvydt, dvzdt, v0]

def plot_trajectory(exit_velocity, launch_angle, spray_angle, distance_ticks, initial_x_val=0):
    v0 = exit_velocity * 0.44704
    angle_launch = launch_angle
    angle_spray = spray_angle
    vx0 = v0 * np.cos(np.radians(angle_launch)) * np.cos(np.radians(angle_spray))
    vy0 = v0 * np.sin(np.radians(angle_launch))
    vz0 = v0 * np.cos(np.radians(angle_launch)) * np.sin(np.radians(angle_spray))
    initial_conditions = [0, initial_x_val, 0, vx0, vy0, vz0, v0]
    solution = solve_ivp(equations_of_motion, [0, distance_ticks], initial_conditions, t_eval=np.linspace(0, distance_ticks, 1000))
    plt.figure(figsize=(10, 5))
    plt.plot(solution.y[0], solution.y[1])
    plt.title(f'Baseball Trajectory, distance - {solution.y[0][len(solution.y[0])-1]}')
    plt.xlabel('Horizontal Distance (ft)')
    plt.ylabel('Vertical Distance (ft)')
    plt.grid(True)
    return plt

plot = gr.Plot()
iface = gr.Interface(fn=plot_trajectory,
                     inputs=[gr.components.Slider(60, 120, step=0.5, label='Exit Velocity (mph)'),
                             gr.components.Slider(10, 50, step=1, label='Launch Angle (degrees)'),
                             gr.components.Slider(0, 45, step=1, label='Spray Direction (degrees)'),
                             gr.components.Slider(5, 20, step=0.5, label='X ticks'),
                             gr.components.Slider(0, 5, step=1, label='Elevation (ft)')],
                     outputs=[plot])
iface.launch()