# This is a module for 8.08 - 8.S308 numerics recitation.
module Rec
    #parameters of the model
    mutable struct Param
        L::Int64
        r::Float64
    end

    #constructor
    function setParam(L, r)
        param = Param(L,r)
        return param
    end

    #state variables of the model
    mutable struct State
        t::Float64                  # current time
        h::Array{Int64, 1}          # height array
        img::Array{Int8,2}          # image array
    end

    #constructor
    function setState(param)
        h = zeros(Int64, param.L)
        img = zeros(Int8, param.L, Int(1e4))
        state = State(0, h, img)
        return state
    end
end

function run_simulation!(state, param, t_run, rng)

    t_now = state.t # time at the beginning of your simulation

    dt = 1/(param.r*param.L) # timestep

    while state.t < t_now + t_run
        # I selected a site for you
        i = rand(rng, 1:param.L)

        ### CONSTRUCT YOUR BALLISTIC DEPOSITION ALGORITHM HERE
        
        ###

        state.t += dt
    end
end

function get_movie!(state, param, rng, t_gap, n_frame, in_fps)

    # create a movie of the aggregation process
    anim = @animate for frame in 1:n_frame
        run_simulation!(state, param, t_gap, rng)
        heatmap(transpose(state.img[:,1:maximum(state.h)+1].*(-1)))
    end

    name = "Rec9_ballistic_deposition.gif"
    gif(anim, name, fps=in_fps)
end

# for fitting roughness data to a power law
using LsqFit
power_law(x, p) = p[1]*x.^p[2] # define power-law function
p0 = [0.1,1.0] # initial guess for power-law fits

function plot_roughness!(state, param, rng, t_gap, n_pts, n_seeds)

    # range for roughness plot
    ts = t_gap:t_gap:n_pts*t_gap
    ws = zeros(Float64,n_pts)

    # iterate over multiple seeds
    for j in 1:n_seeds

        # re-start with an empty system
        state.h = zeros(Int64, param.L)
        state.t = 0
        state.img = zeros(Int8, param.L, Int(1e4))

        # repeatedly simulate with time intervals t_gap
        for i in 1:n_pts
            run_simulation!(state, param, t_gap, rng)

            # average height
            h_avg = sum(state.h) / param.L

            # average roughness
            w = sqrt(sum((state.h .- h_avg).^2) / param.L)
            ws[i] += w
        end
    end

    # normalize roughness by # of seeds
    ws = ws ./ n_seeds

    # fit roughness to a power law
    fit = curve_fit(power_law, ts, ws, p0)

    # plot roughness and its power-law fit
    plot(ts,ws,m=(2),xaxis=:log10,yaxis=:log10,label="w(t)",xlabel="t")
    plot!(ts, power_law(ts, fit.param), label="~ t^"*string(round(fit.param[2],digits=2)))
    
    savefig("Rec9_ballistic_deposition_w.png")
end