using DifferentialEquations
using Plots

r_max = 160
r_steps = 320

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_values = range(0, r_max, length=r_steps)
sols = []
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

    push!(sols, sol)

    println("r=$r")
end

# STEP 2, map data to arrays where plane crosses the x-axis
r_maxes = []
z_maxes = []
r_mins = []
z_mins = []
for i in 1:r_steps
    println("i: ", length(sols[i].t))
    z_values = sols[i][3, :]
    # iterate over to find the local maxima and minima
    # 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
    end
end

# println("r_maxes: ", r_maxes)
# println("z_maxes: ", z_maxes)

# STEP 3, plot the bifurcation diagram
plot(r_maxes, z_maxes, seriestype = :scatter, mc=:blue, ms=.25, ma=0.25, label="Maxima")
plot!(r_mins, z_mins, seriestype = :scatter, mc=:green, ms=.25, ma=0.25, label="Minima")

savefig("hw3/test3.png")