using Plots # for plotting trajectory

# simulation parameters
Δt = 0.01 # time step
y_min = 0.0
θ_to_use = [45, 35] # in degrees
v_0 = 700.0 # in m/s

# constants
B_ref_over_m = 4.0 * 10^(-5) # in m-1, at 300K
T_ref = 300.0 # in kelvin
T_0 = 300.0 # in kelvin
g = 9.8 # in m/s^2

# isothermic parameters
y_o = 10 * 10^4 # k_BT/mg in meter

# adiabatic parameters
α = 2.5 # for air
a = 6.5 * 10^(-3) # in kelvin/meter

function adiabatic!(
    x::Vector{Float64}, y::Vector{Float64}, 
    v_y::Vector{Float64}, v_x::Vector{Float64}, 
    t::Vector{Float64}) 
    while y[end] >= y_min
        # 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_drag = - B_ref_over_m * (T_0 / T_ref)^α *     # temperature variation
            (1 - ((a * y_i) / T_0))^α                   # density/altitude variation
        F_drag_x = F_drag * v_x_i * v_i
        F_drag_y = F_drag * v_y_i * v_i

        # calculate new velocities
        v_x_new = v_x_i + F_drag_x * Δt
        v_y_new = v_y_i + F_drag_y * Δt - g * Δ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

function nodensity!(
    x::Vector{Float64}, y::Vector{Float64}, 
    v_y::Vector{Float64}, v_x::Vector{Float64}, 
    t::Vector{Float64}) 
    while y[end] >= y_min
        # 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_drag = - B_ref_over_m         # coefficient of drag alone
        F_drag_x = F_drag * v_x_i * v_i
        F_drag_y = F_drag * v_y_i * v_i

        # calculate new velocities
        v_x_new = v_x_i + F_drag_x * Δt
        v_y_new = v_y_i + F_drag_y * Δt - g * Δ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 underground
function interpolate!(x::Vector{Float64}, y::Vector{Float64})
    if y[end] == 0
        return # no nothing if y is perfectly on 0
    end

    # calculate x_l, the interpolated x value at y=0
    r = -y[end-1] / y[end]
    x_l = (x[end-1] + r * x[end]) / (1 + r)

    # set final values in the array to interpolated point on ground (y=0    )
    x[end] = x_l
    y[end] = 0.0
end

# setup empty plot to add to
p = plot(xlabel="x (m)", ylabel="y (m)", title="Cannon Shell Trajectory", xlim=(0, 30000), xticks=0:5000:30000, legend=:topright, lw=2)
for θ in θ_to_use
    # arrays to store the trajectory
    x = [0.0]
    y = [0.0]
    v_x = [v_0 * cosd(θ)]
    v_y = [v_0 * sind(θ)]
    t = [0.0]

    # run the simulation
    adiabatic!(x, y, v_y, v_x, t)
    interpolate!(x, y)
    plot_label = "adiabatic, θ = $θ"
    plot!(x, y, label=plot_label, lw=2)

    # reset arrays
    x = [0.0]
    y = [0.0]
    v_x = [v_0 * cosd(θ)]
    v_y = [v_0 * sind(θ)]
    t = [0.0]

    nodensity!(x, y, v_y, v_x, t)
    interpolate!(x, y)
    plot_label = "nodensity, θ = $θ"
    plot!(x, y, label=plot_label, lw=2, linestyle=:dash)
end

# display the plot
display(p)