using Plots

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


# get the intial positions
function calculate_x_i(mass_num, node_num, num_masses, amplitutude)
    return A *
        sqrt(2 / (num_masses + 1)) * 
        sin((mass_num * node_num * π) / (num_masses + 1))
end

# dynamics calculations
function dynamics!(
    mass_displacement, # 2d array
    params)
    (beta, delta_t, num_masses, num_steps) = params

    for step in 1:num_steps-1
        if step == 1
            continue
        end
        for mass_num in 2:num_masses-1

            if step == 1
                prev_mass_pos = 0
            else
                prev_mass_pos = mass_displacement[mass_num, step - 1]
            end

            right_mass_pos = mass_displacement[mass_num + 1, step]
            left_mass_pos = mass_displacement[mass_num - 1, step]

            mass_pos = mass_displacement[mass_num, step]

            xs[mass_num, step + 1] = caluclate_next_displacement(
                mass_pos, prev_mass_pos, 
                left_mass_pos, right_mass_pos, beta, delta_t)
        end

        # update the end points
        mass_displacement[1, step + 1] = caluclate_next_displacement(
            0, 0, 
            0, mass_displacement[2, step], beta, delta_t)
        mass_displacement[num_masses, step + 1] = caluclate_next_displacement(
            0, 0, 
            mass_displacement[num_masses - 1, step], 0, beta, delta_t)
    end

end

function caluclate_next_displacement(
    mass_pos, prev_mass_pos,
    left_mass_pos, right_mass_pos, 
    beta, delta_t
)
    # println(mass_pos, " " , prev_mass_pos, " ", left_mass_pos, " ", right_mass_pos, '\n')
    return 2 * mass_pos - prev_mass_pos
        + delta_t^(2) * (right_mass_pos + left_mass_pos - 2 * mass_pos)
        + beta * delta_t^(2) * ((right_mass_pos - mass_pos)^3 + (left_mass_pos - mass_pos)^3)
end


# energy calcuations, after the dynamics
function calculate_a_k(k, num_masses, delta_t, xs_n, beta)
    sum = 0
    for i in 1:num_masses
        sum += xs_n[i] * sin((k * i * π) / (num_masses + 1))
    end
    return sqrt(2 / (num_masses + 1)) * sum
end

function calculate_energy(a_k, a_k_plus1, delta_t, mode, num_masses)
    kinetic = .5 * ((a_k_plus1 - a_k) / delta_t)^2

    # sum over the three modes
    w_k = 2 * sin((mode * π) / (2 * (num_masses + 1)))
    potential = .5 * (w_k^2 * a_k^2)

    return kinetic + potential
end


# do the simulation
num_steps = Int(final_time / dt)
params = (b, dt, N, num_steps)
# build the 2d array of mass displacements to time
xs = zeros(N, num_steps)
# fill in the initial displacements
for mass_num in 2:N-1
    xs[mass_num, 1] = calculate_x_i(mass_num, 1, N, A)
end
dynamics!(xs, params)

# println(xs)


# plot these displacements over time as an animation
a = @animate for i in 1:num_steps
    plot(xs[:, i], legend=false, 
    marker = :circle, xlabel="Mass Number", ylabel="Displacement",
    yaxis = (-30, 30), title="Displacement Over Time"
    )
end

gif(a, "fpu.gif", fps=15)

# plot the first two timespets positions
# plot(xs[:, 1], label="t=0", marker=:circle, xlabel="Mass Number", ylabel="Displacement", title="First Two Time Steps")
# plot!(xs[:, 2], label="t=1", marker=:circle)
# plot!(xs[:, 3], label="t=2", marker=:circle)
# plot!(xs[:, 4], label="t=3", marker=:circle)
# plot!(xs[:, 5], label="t=4", marker=:circle)
# plot!(xs[:, 6], label="t=5", marker=:circle)
# plot!(xs[:, 7], label="t=6", marker=:circle)


# # calculate the energies for each mode from the displacements 
# energies = zeros(modes, num_steps)
# for mode in 1:modes
#     energies[mode, :] = calculate_energy_for_mode(mode, xs, N, dt, b)
# end

# make a range time steps
# time = range(0, final_time, length=num_steps)

# println("e:", length(energies[1, :]))
# println("t:", length(time))

# plot the energies for each mode
# scatter(time, energies[1, :], label="Mode 1", marker=:circle, xlabel="Time", ylabel="Energy", title="Energy for First Three Modes")
# scatter!(time, energies[2, :], label="Mode 2", marker=:circle)
# scatter!(time, energies[3, :], label="Mode 3", marker=:circle)