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

module Rec
    #parameters of the model
    mutable struct Param
        dt::Float64             # time step
        
        k::Float64              # interaction strength
        lam::Float64            # strength of external potential
        kT1::Float64            # thermal energy of reservoir 1
        kT2::Float64            # thermal energy of reservoir 2
    end

    #constructor
    function setParam(dt, k, lam, kT1, kT2)

        param = Param(dt, k, lam, kT1, kT2)
        return param
    end

    #state variables of the model
    mutable struct State
        t::Float64                  # current time
        x::Array{Float64, 1}        # x-coordinates
        x_next::Array{Float64, 1}   # next x-coordinates
        x_prev::Array{Float64, 1}   # previous x-coordinates
        fx::Array{Float64, 1}       # force
    end

    #constructor
    function setState(param)
        x = zeros(Float64, 2)
        x_next = zeros(Float64, 2)
        x_prev = zeros(Float64, 2)
        fx = zeros(Float64, 2)
        state = State(0, x, x_next, x_prev, fx)
        return state
    end
end

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

    state.fx[1] += -param.k*(state.x[1] - state.x[2]) - param.lam*state.x[1]
    state.fx[2] += -param.k*(state.x[2] - state.x[1]) - param.lam*state.x[2]
end 

# equilibriation while fixing n=1 and n=N monomers
function run_simulation!(state, param, t_run, rng)
    t_now = state.t     # time at the beginning of your run
    dt = param.dt       # time step
    sqrt2kT1dt = sqrt(2*param.kT1*dt)
    sqrt2kT2dt = sqrt(2*param.kT2*dt)

    while state.t < t_now + t_run
        state.x_prev = copy(state.x)
        state.x = copy(state.x_next)
        get_force!(param, state)

        state.x_next[1] = state.x[1] + dt*state.fx[1] + sqrt2kT1dt*randn(rng)
        state.x_next[2] = state.x[2] + dt*state.fx[2] + sqrt2kT2dt*randn(rng)

        state.t += dt
    end
end

# make movie of dynamics
function get_movie!(state, param, rng, t_gap, n_frame, in_fps)

    # estimate maximum distance based on parameters
    xmax = 3*sqrt(param.k^2 * (param.kT1+param.kT2)/(2param.lam*(param.k+param.lam)*(2param.k+param.lam)) + max(param.kT1,param.kT2)/(param.k+param.lam))

    anim = @animate for frame in 1:n_frame
        run_simulation!(state, param, t_gap, rng)

        plot( [state.x[1]], [0], xlims=(-xmax, xmax), ylims=(-xmax/5, xmax/5), ms=8.0, aspectratio=1, marker=:circle, label="T1="*string(param.kT1), color=:orange, legend=:outertopright)
        plot!([state.x[2]], [0], xlims=(-xmax, xmax), ylims=(-xmax/5, xmax/5), ms=8.0, aspectratio=1, marker=:circle, label="T2="*string(param.kT2), color=:blue)

    end

    name = "Rec5_movie.gif"
    gif(anim, name, fps=in_fps, loop=1)
end

# get density and current in (x1,x2) space
function get_densities!(state, param, rng, n_sample, dt_sample, xmax, P, J)

    t_now = state.t 
    dt = param.dt
    t_run = n_sample*dt_sample

    Nbinx = size(P,1)
    dx = 2xmax/Nbinx    # resolution of histogram
    NbinxJ = size(J,1)
    dxJ = 2*xmax/NbinxJ # resolution of J histogram

    while state.t < t_now + t_run
        run_simulation!(state, param, dt_sample, rng)

        i1 = Int(floor((state.x[1]+xmax)/dx)) + 1
        i2 = Int(floor((state.x[2]+xmax)/dx)) + 1

        if i1>0 && i1<=Nbinx && i2>0 && i2<=Nbinx
            P[i1,i2] += 1.0
        end

        i1 = Int(floor((state.x[1]+xmax)/dxJ)) + 1
        i2 = Int(floor((state.x[2]+xmax)/dxJ)) + 1

        if i1>0 && i1<=NbinxJ && i2>0 && i2<=NbinxJ
            J[i1,i2,1] += 0.5*(state.x_next[1] - state.x_prev[1])/dt
            J[i1,i2,2] += 0.5*(state.x_next[2] - state.x_prev[2])/dt
        end
    end 
    P /= n_sample
    J /= n_sample

    # plot density
    x_array = -xmax+dx/2:dx:xmax-dx/2
    heatmap(x_array, x_array, P; aspect_ratio=1, colorbar=true)

    ## plot current

    # create mask to remove arrows of magnitude 0
    mag = sqrt.(J[:,:,1].^2 + J[:,:,2].^2)
    mask = mag.>0

    # separate 1,2 currents
    J1 = J[:,:,1]
    J2 = J[:,:,2]

    # plot arrows on log scale
    scale_factor = 0.05*log10.(mag[mask]./minimum(mag[mask]))
    
    # x grid
    x_array_j = -xmax+dxJ/2:dxJ:xmax-dxJ/2
    X1 = repeat(x_array_j', NbinxJ, 1)
    X2 = repeat(x_array_j, 1, NbinxJ)

    # plot current arrows
    quiver!(X1[mask], X2[mask], quiver=(J1[mask]./scale_factor,J2[mask]./scale_factor); color=:white)

    savefig("rec5_density.png")


end