# This is a code for 8.08 - 8.S421 recitation.
# You will try measurement of free energy using Jarzynski equality

# You will pull a polymer and measure the work you did during the process
# This will be compared to the free energy difference of the polymer before and after undergoing the process.

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

# for plotting
using Plots
using Random

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

#random number generator
rng = MersenneTwister(1234)

# parameters
dt = 1e-2       # time step
N = 11          # number of monomers
k = 1           # nearest neighbor interaction strength
mu = 1          # mobility
kT = 1          # thermal energy
pull_dist = 2   # final distance after extension

num_ensemble = 1000                         # number of independent ensemble samples
num_record = 10                             # number of work or exp(-w) samples to be gathered
work_record = zeros(Float64, num_record)    # array for work
expw_record = zeros(Float64, num_record)    # array for exp(-w)

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

#set state
state = Rec.setState(param)

t_eq = 400         # equilibriation time
t_Jarzynski = 200  # time for pulling experiment

for i_ens in 1:num_ensemble     # ensemble averaging
    if (i_ens*10)%num_ensemble==0
        println(string(i_ens)*"/"*string(num_ensemble)*" is done")
    end

    ### run simulation
    global state = Rec.setState(param) # reset state to all x and y at 0
    run_equilibration!(state, param, t_eq, rng)
    run_Jarzynski!(state, param, t_Jarzynski, pull_dist, work_record, expw_record, rng)
    ###
end

work_record ./= num_ensemble
expw_record ./= num_ensemble

#histogram(work_record)
dist_samples = (pull_dist/num_record):(pull_dist/num_record):pull_dist
plot( dist_samples, work_record, m=(4), linewidth=0.5, color=:black, label="work", xlabel="end-to-end distance")
plot!( dist_samples, -log.(expw_record), m=(4), linewidth=0.5, color=:red, label="Jarzynski measurement")
plot!( dist_samples,  k ./ (2*(N-1)) .* (dist_samples) .^ 2, linewidth=1.5, color=:blue ,label="free energy")

savefig("rec4_Jarzynski.png")