using Plots # for plotting trajectory

# simulation parameters
Δt = 0.01 # time step
y_min = 0.0
Δθ = 0.01 # in degrees
θ_start = 0.0 # in degrees
θ_end = 90.0 # 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 isothermic!(
    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
            ℯ^(-y_i/y_o)                                # 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 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
en

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

# arrays for holidng range and angle data
Θ_steps = Int64((θ_end - θ_start) / Δθ)
ranges_adi = zeros(Θ_steps + 1)
angles_adi = zeros(Θ_steps + 1)
ranges_iso = zeros(Θ_steps + 1)
angles_iso = zeros(Θ_steps + 1)
ranges_nod = zeros(Θ_steps + 1)
angles_nod = zeros(Θ_steps + 1)

for i in 1:Θ_steps+1
    # arrays to store steps
    θ = θ_start + (i-1) * Δθ

    local x_adi = [0.0]
    local y_adi = [0.0]
    local v_x_adi = [v_0 * cosd(θ)]
    local v_y_adi = [v_0 * sind(θ)]
    local t_adi = [0.0]

    local x_iso = [0.0]
    local y_iso = [0.0]
    local v_x_iso = [v_0 * cosd(θ)]
    local v_y_iso = [v_0 * sind(θ)]
    local t_iso = [0.0]

    local x_nod = [0.0]
    local y_nod = [0.0]
    local v_x_nod = [v_0 * cosd(θ)]
    local v_y_nod = [v_0 * sind(θ)]
    local t_nod = [0.0]

    # simulate the trajectory
    adiabatic!(x_adi, y_adi, v_y_adi, v_x_adi, t_adi)
    isothermic!(x_iso, y_iso, v_y_iso, v_x_iso, t_iso)
    nodensity!(x_nod, y_nod, v_y_nod, v_x_nod, t_nod)

    # interpolate the last point that's underground
    interpolate!(x_adi, y_adi)
    interpolate!(x_iso, y_iso)
    interpolate!(x_nod, y_nod)

    # store the range and angle
    ranges_adi[i] = x_adi[end]
    angles_adi[i] = θ
    ranges_iso[i] = x_iso[end]
    angles_iso[i] = θ
    ranges_nod[i] = x_nod[end]
    angles_nod[i] = θ
end

max_i_adi = argmax(ranges_adi)
max_i_iso = argmax(ranges_iso)
max_i_nod = argmax(ranges_nod)

# print the max range and angle
println("Max range (adiabatic): ", ranges_adi[max_i_adi], " at angle ", angles_adi[max_i_adi], " (T=", T_0, "K)")
println("Max range (isothermic): ", ranges_iso[max_i_iso], " at angle ", angles_iso[max_i_iso], " (T=", T_0, "K)")
println("Max range (no density model): ", ranges_nod[max_i_nod], " at angle ", angles_nod[max_i_nod], " (T=", T_0, "K)")

# plot the range over angles
plot_title = "Range over angle (T=$T_0 K)"
plot(angles_adi, ranges_adi, xlabel="angle (degrees)", ylabel="range (m)", title=plot_title, label="adiabatic", lw=2, color=:blue)
plot!(angles_iso, ranges_iso, label="isothermic", lw=2, color=:red)
plot!(angles_nod, ranges_nod, label="no density model", lw=2, color=:green)