# This is a module for 8.08 - 8.S421 numerics recitation.

module Rec
    #parameters of the model
    mutable struct Param
        dt::Float64             # time step
        
        N::Int64                # number of monomers
        k::Float64              # nearest neighbor interaction strength
        mu::Float64             # mobility
        kT::Float64             # thermal energy
    end

    #constructor
    function setParam(dt, N, k, mu, kt)

        param = Param(dt, N, k, mu, kt)
        return param
    end

    #state variables of the model
    mutable struct State
        t::Float64                  # current time
        x::Array{Float64, 1}        # x-coordinates
        y::Array{Float64, 1}        # y-coordinates

        x_next::Array{Float64, 1}   # next x-coordinates
        y_next::Array{Float64, 1}   # next y-coordinates

        fx::Array{Float64, 1}       # force along x-direction
        fy::Array{Float64, 1}       # force along y-direction
    end

    #constructor
    function setState(param)
        x = zeros(Float64, param.N)
        y = zeros(Float64, param.N)

        x_next = zeros(Float64, param.N)
        y_next = zeros(Float64, param.N)

        fx = zeros(Float64, param.N)
        fy = zeros(Float64, param.N)

        state = State(0, x, y, x_next, y_next, fx, fy)
        return state
    end
end

# calculate interaction forces between all monomers
function get_force!(param, state)
    fill!(state.fx, 0)
    fill!(state.fy, 0)

    for n in 1:param.N-1
        state.fx[n] += - param.k*(state.x[n] - state.x[n+1])
        state.fy[n] += - param.k*(state.y[n] - state.y[n+1])

        state.fx[n+1] += param.k*(state.x[n] - state.x[n+1])
        state.fy[n+1] += param.k*(state.y[n] - state.y[n+1])
    end
end 

# equilibriation while fixing n=1 and n=N monomers
function run_equilibration!(state, param, t_run, rng)
    t_now = state.t     # time at the beginning of your run
    dt = param.dt       # time step
    mudt = param.mu*dt

    while state.t < t_now + t_run
        get_force!(param, state)

        for n in 2:param.N-1
            state.x_next[n] = state.x[n] + mudt*state.fx[n] + sqrt(2mudt*param.kT)*randn(rng)
            state.y_next[n] = state.y[n] + mudt*state.fy[n] + sqrt(2mudt*param.kT)*randn(rng) 
        end

        state.x = copy(state.x_next)
        state.y = copy(state.y_next)
        state.t += dt
    end
end


# run pulling experiment
function run_Jarzynski!(state, param, t_run, pull_dist, w_record, expw_record, rng)

    t_now = state.t     # time at the beginning of your simulation
    dt = param.dt       # time step
    mudt = param.mu*dt

    v = pull_dist/t_run      # pulling speed
    work = 0.0               # work done

    Delta_t = t_run / length(w_record)  # interval between recording w and expw
    index = 1                           # how many recordings we have done

    while state.t < t_now + t_run
        get_force!(param, state)

        for n in 2:param.N-1
            state.x_next[n] = state.x[n] + mudt*state.fx[n] + sqrt(2mudt*param.kT)*randn(rng)
            state.y_next[n] = state.y[n] + mudt*state.fy[n] + sqrt(2mudt*param.kT)*randn(rng) 
        end

        ### Jarznski

        # pull end
        state.x_next[param.N] = state.x[param.N] + v*dt

        work -= state.fx[param.N] * v*dt
        ###

        state.x = copy(state.x_next)
        state.y = copy(state.y_next)

        state.t += dt
        if state.t - t_now > index*Delta_t
            ### record
            w_record[index] += work
            expw_record[index] += exp(-work/param.kT)
            
            ### record
            index += 1
        end
    end
end