function dynamics!(
	state,
	prev_state,
	params,
	states,
)
	# Unpack the parameters
	N, modes, beta, A, dt, num_steps = params

	for i in 1:num_steps
		next_state = zeros(N)

		# update the left particle (set to rest)
		state[1] = 0

		# update the right particle (set to rest)
		state[N] = 0

		# update the middle particles
		for i in 2:N-1
			a = 2 * state[i] - prev_state[i]
			b = dt * dt * (state[i+1] - 2 * state[i] + state[i-1])
			c = dt * dt * beta * ((state[i+1] - state[i])^3 + (state[i-1] - state[i])^3)
			next_state[i] = a + b + c
		end

		push!(states, next_state)

		# update the previous state
		for i in 1:N
			prev_state[i] = state[i]
		end
		# update the current state
		for i in 1:N
			state[i] = next_state[i]
		end
	end
end

function get_initial_state(
	N,
	modes,
	beta,
	A,
)
	state = zeros(N)
	amp = A * sqrt(2 / (N - 1))
	for i in 2:N-1
		state[i] = amp * sin((modes * π * (i - 1)) / (N - 1))
	end
	return state
end

function run_simulation(
	N,
	modes,
	beta,
	A,
	dt,
	num_steps,
)
	states = []
	state = get_initial_state(N, 1, beta, A)
	prev_state = zeros(N)
	for i in 1:N
		prev_state[i] = state[i]
	end
	params = (N, modes, beta, A, dt, num_steps)
	dynamics!(state, prev_state, params, states)
	return states
end

# Run the simulation
N = 32 # number of masses
beta = 1.5 # cubic string spring
A = 10 # amplitude
modes = 3 # number of modes to plot
final_time = 50 # seconds
dt = 0.5 # seconds
num_steps = Int(final_time / dt)
params = (N, modes, beta, A, dt, num_steps)
s = run_simulation(N, modes, beta, A, dt, num_steps)
println("Final state: ", s[end])

# plot the inital positions and the final positions
using Plots
# plot(get_initial_state(N, 1, beta, A), label="Initial", marker=:circle, xlabel="Mass Number", ylabel="Displacement", title="First Three Modes")
# plot!(s[end], label="Final", marker=:circle)

# animate the s array of positions
anim = @animate for i in 1:length(s)
	plot(s[i], label = "t = $(i * dt)", marker = :circle, xlabel = "Mass Number", ylabel = "Displacement", ylim = (-3, 3))
end
gif(anim, "plz.gif", fps = 30)

# function caluclate_energies_for_mode(states, mode)
#     total = []
#     kinetic = []
#     potential = []
#     prev_a_k = 0
#     for i in 1:length(states) - 1
#         total_energy = 0
#         kinetic_energy = 0
#         potential_energy = 0

#         sum = 0
#         for j in 1:length(states[i])
#             sum += states[i][j] * sin((mode * j * π) / (N - 1))
#         end
#         amp = A * sqrt(2 / (N - 1))
#         a_k = amp * sum

#         k = .5 * ((a_k - prev_a_k) / dt)^2
#         if i == 1
#             k = 0
#         end
#         p = (2 * sin((mode * π) / (2 * (N - 1)))^2 * a_k^2)


#         total_energy += k + p
#         kinetic_energy += k
#         potential_energy += p
#         push!(total, total_energy)
#         push!(kinetic, kinetic_energy)
#         push!(potential, potential_energy)

#         prev_a_k = a_k
#     end
#     return total, kinetic, potential
# end

# # calculate the energies for each mode from the displacements
# total_1, kinetic_1, potential_2 = caluclate_energies_for_mode(s, 1)
# total_2, kinetic_2, potential_2 = caluclate_energies_for_mode(s, 2)
# total_3, kinetic_3, potential_3 = caluclate_energies_for_mode(s, 3)

# # build a timestep space
# # timesteps = [i * dt for i in 1:num_steps - 1]
# # # plot the modes energies on the same y-axis
# plot(timesteps, total_1, label="Mode 1", xlabel="Time", ylabel="Energy of Modes", title="Energy Over Time (\$\\beta = 1.5\$)")
# plot!(timesteps, total_2, label="Mode 3")
# plot!(timesteps, total_3, label="Mode 5")

# savefig("modes-beta15.png")