# This is a code for 8.08 - 8.S421, 2nd recitation.
# You will run simulations of a Brownian particle with spatially-varying temperature, 
# and compare the density to the theoretical expectation.

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

# load packages
using Plots
using Random


####
#### DEFINE YOUR TEMPERATURE FUNCTION AND ITS DERIVATIVE HERE
T(x,L) = 2.0 + sin(2*π*x/L)
Tp(x,L) = (2*π/L)*cos(2*π*x/L)
####

# integrator for simulating a harmonic oscillator
function runBrownianTofx!(
    tf,             # duration of your simulation
    h,              # time step
    alpha,          # discretization
    seed,           # seed for random number generator
    L,              # length of system 
    num_cells,      # number of cells in your histogram
    num_record,     # number of data to be recorded
    data_histogram  # array for recording time
    )

    # create random number generator
    rng = MersenneTwister(seed)

    x_now, x_next = 0, 0    #state variables

    ####
    #### You might want to define more quantities here
    sqrt2h = sqrt(2*h)
    halpha = h*alpha
    ####

    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 = x_now + sqrt2h * sqrt(T(x_now,L)) * randn(rng) + halpha * Tp(x_now,L)

        ####

        # implement periodic boundary conditions
        if x_next>=L
        	x_next = x_next - L
        elseif x_next<0
        	x_next = x_next + L 
        end

        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((x_now/L)*num_cells )) + 1
            data_histogram[cell] += 1
        end
        t += h # update time
    end
end

# simulation parameters
L = 10.0               # system size
h = 0.1                # timestep
tf = 1e+7              # simulation duration
num_record = Int(tf)   # number of snapshots to take
num_cells = Int(L*10)  # number of bins in the histogram
seed = 1234            # seed for random number generator

resolution = L / Float64(num_cells)
x_array = range(resolution/2, step=resolution, length=num_cells) # array of possible x values in your histogram
data_histogram = zeros(Float64, num_cells)                       # your histogram

## density plotting
plot(xlabel="x", ylabel="P(x)", title="h="*string(h)) # initialize plot

# run simulation with different discretizations
alphas = [0,0.5] # Ito, Stratonovich discretizations
colors = [:red,:blue]
for i_alpha in 1:length(alphas)
	alpha = alphas[i_alpha]
	runBrownianTofx!(tf, h, alpha, seed, L, num_cells, num_record, data_histogram)
	plot!(x_array, data_histogram ./ (num_record*resolution), linewidth=0.5, label="alpha = "*string(alpha)*", simulation", color=colors[i_alpha])

	####
	#### ADD THEORETICAL PREDICTION HERE
	theory_data = T.(x_array,L).^(alpha-1)
	####

	theory_data = theory_data ./ (sum(theory_data) * resolution) # normalize

	#plot theory
	plot!(x_array, theory_data, linewidth=0.5, label="alpha = "*string(alpha)*", theory", color=colors[i_alpha], ls=:dash)
end

# save histogram figure
cd(@__DIR__)
savefig("rec2_density_Tofx.png")