N = 32 # number of masses
b = 10 # cubic string spring
A = 1 # amplitude
modes = 3 # number of modes to plot
final_time = 10 # seconds
dt = .05 # seconds

function kinetic_energy(a_k, a_k_plus1, delta_t)
    return .5 * ((a_k_plus1 - a_k) / delta_t)^2
end

function potential_energy(a_k, a_k_plus1, mode, num_masses)
    w_k = 2 * sin((mode * π) / (2 * (num_masses + 1)))
    return .5 * (w_k^2 * a_k^2)
end

function calculate_energy(a_k, a_k_plus1, delta_t, mode, num_masses)
    k = kinetic_energy(a_k, a_k_plus1, delta_t)
    p = potential_energy(a_k, a_k_plus1, mode, num_masses)

    return k + p
end

function update_state!(state, prev_state, dt, beta)
    x_prev = prev_state
    x = copy(state)

    # 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
        state[i] = 2 * x[i] - x_prev[i] 
        + dt * dt * ((x[i + 1] - 2 * x[i] + x[i - 1]) 
        + dt * dt * beta * ((x[i + 1] - x[i])^3 - (x[i - 1] - x[i])^3)
        )
    end
end

function initial_state(N, modes, beta, A)
    # make the range of masses
    masses = 1:N
    # make the range of modes
    init_state = A * sin.((modes * π * masses) / (N + 1))
    init_state[1] = 0
    init_state[N] = 0
    return init_state
end

function run_simulation(N, modes, beta, A, dt, num_steps)
    data = []
    state = initial_state(N, modes, beta, A)
    prev_state = copy(state)
    println("Initial state: ", state)
    for i in 1:num_steps
        update_state!(state, prev_state, dt, beta)
        push!(data, state)
        prev_state = copy(state)
    end

    return data
end

steps = Int(final_time / dt)

p = run_simulation(N, modes, b, A, dt, steps)

using Plots
# plot the first and final position
plot(p[1], label="Initial position")
plot!(p[end], label="Final position")

# plot displacement for the first particle over time
t_span = 0:dt:final_time
plot(t_span, p, label="Displacement for first particle")