using DifferentialEquations
using Plots

r_max = 160
dr = .01

sim_time = 1000.0

# STEP 1, go over each value of r and store the results
# Simulate Lorenz 63 system and investigate sensitivity to initial conditions
function tendency!(du, u, p, t)
    x,y,z = u
    σ,ρ,β = p
    
    du[1] = dx = σ*(y-x)
    du[2] = dy = x*(ρ-z) - y
    du[3] = dz = x*y - β*z
end


# make a linspace for these values of r
r_steps = Int(r_max/dr)
r_values = range(0, r_max, length=r_steps)
r_maxes = []
z_maxes = []
r_mins = []
z_mins = []
let 
    r_diverge_1 = -1
    for i in 1:r_steps
        r = r_values[i]
        p = [10.0, r, 8/3] # parameters of the Lorentz 63 system
        tspan = (0.0, sim_time)

        u0 = [1.0, 0.0, 0.0]
        prob = ODEProblem(tendency!, u0, tspan, p)
        sol = solve(prob, Tsit5(), reltol=1e-8, abstol=1e-8) # control simulation

        # iterate over to find the local maxima and minima
        z_values = sol[3, :]

        # take off the first 300 values to avoid transient
        for j in 301:length(z_values)-1
            if z_values[j] > z_values[j-1] && z_values[j] > z_values[j+1]
                push!(r_maxes, r_values[i])
                push!(z_maxes, z_values[j])
            end
            if z_values[j] < z_values[j-1] && z_values[j] < z_values[j+1]
                push!(r_mins, r_values[i])
                push!(z_mins, z_values[j])
            end

            # calcualte dz
            # if r_diverge_1 < 0 && length(z_mins) > 1 && length(z_maxes) > 1 && abs(z_mins[end] - z_maxes[end]) > 20.0
            #     r_diverge_1 = r
            #     println("Divergence starts at r = $r")
            # end
        end

        if i % 500 == 0
            println("r=$r")
        end
    end

    plt = plot(r_maxes, z_maxes, seriestype=:scatter, mc=:blue, ms=.25, ma=.25, msw=0, xlabel="r", ylabel="z", title="Bifurcation Diagram for Lorenz 63 System", legend=:topleft, grid=:true, label="Local Maxima")
    plot!(r_mins, z_mins, seriestype=:scatter, mc=:green, ms=.25, ma=.25, msw=0, xlabel="r", ylabel="z", label="Local Minima")

    # r_diverge_display = round(r_diverge_1, digits=3)
    vline!([24], label="Critical r value ≈ 24", lc=:red, lw=1.5, ls=:dash)
    vline!([122], label="Critical r value ≈ 122", lc=:orange, lw=1.5, ls=:dash)
    vline!([146], label="Critical r value ≈ 146", lc=:purple, lw=1.5, ls=:dash)


    savefig("hw3/test11.png")
end