import numpy as np
from scipy.integrate import solve_ivp

def equations(x, y, params):
    x, vx, y, vy, z, vz = y
    Cd, Cm, rho, A, m, g, v_wind_x, v_wind_y, omega_x, omega_y, omega_z = params
    
    v_rel_x = vx - v_wind_x
    v_rel_y = vy - v_wind_y
    v_rel_z = vz
    
    v_rel = np.sqrt(v_rel_x**2 + v_rel_y**2 + v_rel_z**2)
    drag_force = 0.5 * Cd * rho * A * v_rel**2
    
    ax_drag = -drag_force * (v_rel_x / v_rel) / m
    ay_drag = -drag_force * (v_rel_y / v_rel) / m
    az_drag = -drag_force * (v_rel_z / v_rel) / m
    
    ax_magnus = Cm * (omega_y * v_rel_z - omega_z * v_rel_y)
    ay_magnus = Cm * (omega_z * v_rel_x - omega_x * v_rel_z)
    az_magnus = Cm * (omega_x * v_rel_y - omega_y * v_rel_x)
    
    ax = ax_drag + ax_magnus
    ay = ay_drag + ay_magnus - g
    az = az_drag + az_magnus
    
    return [vx, ax, vy, ay, vz, az]

vmph = 112.5 #exit velo, mph
theta0 = 23 #launch angle, degrees

v0 = vmph * 0.44705 #convert to m/s
Cd = 0.47
Cm = 0.004 #Spin coefficient
rho = 1.225 #kg/m3, 1.225
A = 0.0042
m = 0.145
g = 9.81
v_wind = 5 #m/s
theta_wind = 30 #direction in deg
spin_rate = 2000 #RPM
omega = spin_rate * 2 * np.pi / 60 #spin rate per second
omega_x = -10
omega_y = 0
omega_z = 15.5

v_wind_x = v_wind * np.cos(np.radians(theta_wind))
v_wind_y = v_wind * np.sin(np.radians(theta_wind))
theta0_rad = np.radians(theta0)
vx0 = v0 * np.cos(theta0_rad)
vy0 = v0 * np.sin(theta0_rad)
vz0 = 0

sol = solve_ivp(equations, [0, 10], [0, vx0, 0, vy0, 0, vz0], args=([Cd, Cm, rho, A, m, g, v_wind_x, v_wind_y, omega_x, omega_y, omega_z],), dense_output=True)

t = np.linspace(0, 10, 500)
x, vx, y, vw, z, vz = sol.sol(t)

ground_hit_idx = np.where(y < 0)[0][0]
travel_distance = x[ground_hit_idx]