# This is a code for 8.08 - 8.S421 recitation.
# You will simulate polymer in equlibrium fluid

# for installing required packages
# using Pkg
# Pkg.add("Plots")
# Pkg.add("Random")

# for plotting
using Plots
using Random
using ProgressBars
using LsqFit

cd(@__DIR__)
include("my_Rec8_module.jl")

#random number generator
rng = MersenneTwister(1234)

# parameters
dt = 1e-2   # time step
N = 100     # number of monomers
k = 1       # nearest neighbor interaction strength
mu = 1      # mobility
kT = 1      # thermal energy

#set parameters
param = Rec.setParam(dt, N, k, mu, kT)

#set state
state = Rec.setState(param)

for n in 1:param.N  # initial configuration exteded along the x-axis.
    state.y[n] = 0
    state.x[n] = n-param.N/2
end

# for generating animation of your simulation
get_movie!(state, param, rng, 100, 150, 20)

# for measuring energy of your polymer
average_energy = 0
num_records = 100
delta_t = 30

for i in 1:num_records  
    # get average energy by measuring energy with the time interval of delta_t
    run_simulation!(state, param, delta_t, rng)
    global average_energy += getEnergy(state, param)
end
average_energy /= num_records

println("my energy is about "*string(round(average_energy; digits=2))*" kT")

Task - Make sure your polymer is well-behaving and check whether your energy is NkT or not.

# If you are done early - Check the end-to-end distance distribution and compare it with your expectation
num_records = 2000
delta_t = 100

distmax = 60
resolution = 1
Nbin = Int(round(2*distmax / resolution))
dx_histo = zeros(Float64, Nbin)
dy_histo = zeros(Float64, Nbin)
println("measuring length distribution")
for i in ProgressBar(1:num_records)
    run_simulation!(state, param, delta_t, rng)

    dx = state.x[N] - state.x[1]
    dy = state.y[N] - state.y[1]
    if dx<distmax && dx>=-distmax
        i_dx = Int(floor((dx + distmax)/resolution)) + 1
        dx_histo[i_dx] += 1
    end
    if dy<distmax && dy>=-distmax
        i_dy = Int(floor((dy + distmax)/resolution)) + 1
        dy_histo[i_dy] += 1
    end
end

x_array = range(-distmax+resolution/2, step=resolution, length=Nbin)
dx_histo /= (num_records*resolution)
dy_histo /= (num_records*resolution)
plot(x_array, exp.(-x_array.^2 ./ (2*(N-1))) ./ sqrt(2*π*(N-1)), color="black", label="Gaussian")
plot!(x_array, dx_histo, label="P(dx)")
plot!(x_array, dy_histo, label="P(dy)")

cd(@__DIR__)
savefig("rec3_Pofdx.png")




# If you still have time - Check N dependence of the end-to-end distance.

Ns = 10:10:120
t_burnin = 10000
num_records = 200
delta_t = 200

avg_length2s = zeros(Float64, length(Ns))

println("measuring average length vs. N")
for i_N in ProgressBar(1:length(Ns))
    global N = Ns[i_N]
    global param = Rec.setParam(dt, N, k, mu, kT)
    global state = Rec.setState(param)

    # burn-in
    run_simulation!(state, param, t_burnin, rng)

    for i in 1:num_records
        run_simulation!(state, param, delta_t, rng)
        avg_length2s[i_N] += (state.x[N] - state.x[1])^2 + (state.y[N] - state.y[1])^2
    end
    avg_length2s[i_N] /= num_records
end

plot(Ns, avg_length2s, m=(2), xaxis=:log10, yaxis=:log10, xlabel="N", ylabel="R^2")

# fit the error to power law and plot the result
power_law(x, p) = p[1]*x.^p[2] # define power-law function

# initial guess
p0 = [1.0,1.0]


# fit error to a power law
fit = curve_fit(power_law, Ns, avg_length2s, p0)

# plot the result
plot!(Ns, power_law(Ns, fit.param), label=string(round(fit.param[1],digits=2))*" N^"*string(round(fit.param[2],digits=2)))

savefig("rec3_R2vsN.png")