using Plots # for plotting trajectory

# simulation parameters
Δt = 0.001 # time step
x_max = 60.0 # in feet
v_0 = 95.0 # mph
v_0 = v_0 * 5280 / 3600 # convert mph to feet/s
y_0 = 6.0 # in feet
ω = 0.0 # in rad/s
m = 149 # in grams
S_0_over_m = 4.1 * 10^(-4) # unitless
B_ref_over_m = 4.0 * 10^(-5) # in m-1, at 300K
g = 32.2 # in ft/s^2

function dynamics!(
    x::Vector{Float64}, y::Vector{Float64}, 
    v_y::Vector{Float64}, v_x::Vector{Float64}, 
    t::Vector{Float64}) 
    while x[end] <= x_max
        # decompose previous positions and velocities 
        x_i = x[end]
        y_i = y[end]
        v_x_i = v_x[end]
        v_y_i = v_y[end]

        # calculate new positions
        x_new = x_i + v_x_i * Δt
        y_new = y_i + v_y_i * Δt

        # calculate drag force
        v_i = sqrt(v_x_i^2 + v_y_i^2)
        F_x = - B_ref_over_m * v_x_i * v_i
        F_y = - g + S_0_over_m * v_x_i * ω

        # calculate new velocities
        v_x_new = v_x_i + F_x * Δt
        v_y_new = v_y_i + F_y * Δt

        # store new positions and velocities
        push!(x, x_new)
        push!(y, y_new)
        push!(v_x, v_x_new)
        push!(v_y, v_y_new)
        push!(t, t[end] + Δt)
    end
end

# interpolate the last point that's past the x_max
function interpolate!(x::Vector{Float64}, y::Vector{Float64})
    if x[end] <= x_max
        return
    end

    x_i = x[end-1]
    y_i = y[end-1]

    x_f = x[end]
    y_f = y[end]

    m = (y_f - y_i) / (x_f - x_i)
    b = y_i - m * x_i
    
    x_new = x_max
    y_new = m * x_new + b

    x[end] = x_new
    y[end] = y_new
end

# run the simulation with ω = 0
x = [0.0]
y = [y_0]
v_x = [v_0] # pitch only in x direction
v_y = [0.0]
t = [0.0]
dynamics!(x, y, v_y, v_x, t)
interpolate!(x, y)
plot(x, y, label="ω = 0", xlabel="x (feet)", ylabel="y (feet)", lw=2, title="Trajectory of a Fastball with Backspin")
println("y-value at x_max (60ft) for ω = 0:\t", y[end], " feet")

# run the simulation with ω = 1000 rpm
ω = 1000 * 2 * π / 60 # convert rpm to rad/s
x = [0.0]
y = [y_0]
v_x = [v_0] # pitch only in x direction
v_y = [0.0]
t = [0.0]
dynamics!(x, y, v_y, v_x, t)
interpolate!(x, y)
println("y-value at x_max (60ft) for ω = 1000 rpm:\t", y[end], " feet")
plot!(x, y, label="ω = 1000 rpm", lw=2)