# This is a code for 8.08 - 8.S421, 1st-2nd recitation.
# You will run simulations of a Brownian particle in harmonic potential
# Compare how the outcome varies when you change the simulation time step and the algorithm implemented. 

# parse command line argument(s)
((length(ARGS)==1) && (ARGS[1] in ["1","2"])) || throw(ArgumentError("Argument 1 must be either 1 (Euler/Maruyama) or 2 (Mannella)"))

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

# load packages
using Plots
using Random
using LsqFit

# integrator for simulating a harmonic oscillator
function runHarmonicOscillator!(
    tf,             # duration of your simulation
    dt,             # time step
    seed,           # seed for random number generator
    resolution,     # resolution of your historgram
    width,          # width of your histogram
    num_cells,      # number of cells in your histogram
    num_record,     # number of data to be recorded
    data_histogram  # array for recording time
    )

    rng = MersenneTwister(seed)

    x_now, x_next = 0, 0    #state variables

    #### You might want to define more quantities or not
    if ARGS[1]=="1"
        a = 1-dt
        b = sqrt(2*dt)
    elseif ARGS[1]=="2"
        a = 1-dt+dt*dt/2
        b = sqrt(2*dt*(1-dt+dt*dt/3))
    end
    ####

    t = 0                       # current time
    Delta_t = tf / num_record   # sampling time gap
    record_index = 1            # current index of my recording

    fill!(data_histogram, 0)    # initialize your histogram

    while t < tf + Delta_t       # simulate until tf is reached, with a tiny margin to assure full recording
        ####
        ####
        #### WRITE YOUR EQUATIONS OF MOTION HERE
        x_next = a*x_now + b*randn(rng)

        x_now   = x_next
        ####
        ####
        ####

        if t > record_index*Delta_t && record_index<num_record  # save position in your histogram
            record_index += 1

            cell = Int(floor( (width/2 + x_now + resolution/2)/width*num_cells ))
            if 0 < cell < num_cells+1
                data_histogram[cell] += 1
            end
        end
        t += dt # update time
    end
end



# For your histogram, specify the resolution and the number of cells. 
resolution = 0.005
width = 20
num_cells = Int( round(width/resolution) ) + 1

x_array = -(width/2):resolution:(width/2)   # array of possible x values in your histogram
data_histogram = zeros(Float64, num_cells)  # your histogram


# theory
theory_data = exp.( -x_array.^ 2 ./ 2) ./ sqrt(2*π)  # Boltzmann distribution

plot(x_array, theory_data , linestyle=:solid, linewidth=1.5, linecolor=:black,
     title="Harmonic Oscillator", xlabel = "x", ylabel = "P(x)", legend=:outerright, label="theory")


# run simulation
tf = 2e+8               # This is the simulation duration
num_record = 2*10^8       # Number of samples to be recorded
seed = 1234             # seed for random number generator

# rough simulation
dt = 1.0
runHarmonicOscillator!(tf, dt, seed, resolution, width, num_cells,  num_record, data_histogram)
plot!(x_array, data_histogram / (num_record*resolution), linewidth=0.5, label = "dt = "*string(dt), color=:red )

# more precise simulation
dt = 0.50
runHarmonicOscillator!(tf, dt, seed, resolution, width, num_cells,  num_record, data_histogram)
plot!(x_array, data_histogram / (num_record*resolution), linewidth=0.5, label = "dt = "*string(dt), color=:blue )

# even better simulation
dt = 0.25
runHarmonicOscillator!(tf, dt, seed, resolution, width, num_cells,  num_record, data_histogram)
plot!(x_array, data_histogram / (num_record*resolution), linewidth=0.5, label = "dt = "*string(dt), color=:green )

# save histogram figure
cd(@__DIR__)
savefig("rec1-2_histo_algo"*ARGS[1]*".png")



# error estimation
gaussian(x, p) = exp.(-x.^2 ./ (2*p[1]*p[1])) ./ sqrt(2*π*p[1]*p[1])
dt_array = 0.05:0.05:0.5                    # list of dt values you will try
error = zeros(Float64, length(dt_array) )   # array for recording error

for i_dt in 1:length(dt_array)
    local dt = dt_array[i_dt]

    # run simulation with this dt
    runHarmonicOscillator!(tf, dt, seed, resolution, width, num_cells,  num_record, data_histogram)

    # fit result to a gaussian
    fit = curve_fit(gaussian, x_array, data_histogram ./ (num_record*resolution), [1.0])
    println(fit.param)

    # calculate the error in the distribution width
    error[i_dt] = sqrt((fit.param[1]-1.0)^2)
end


# plot the error
plot(dt_array, error, m=(1), xaxis=:log10, yaxis=:log10, label="error", xlabel = "dt", title="error in distribution width")

# fit the error to power law and plot the result

# define power-law function
power_law(x, p) = p[1]*x.^p[2]

# initial guess
p0 = [0.1,1.0]

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

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

# save error figure
savefig("rec1-2_err_algo"*ARGS[1]*".png")