# This is a module for 8.08 - 8.S308 numerics recitation.
module Rec
    #parameters of the model
    mutable struct Param
        L::Int64     # number of points in domain
        lam::Float64 # strength of nonlinearity
        nu::Float64  # diffusion constant
        D::Float64   # noise strength
        dt::Float64  # timestep
    end

    #constructor
    function setParam(L, lam, nu, D, dt)
        param = Param(L,lam,nu,D,dt)
        return param
    end

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

    #constructor
    function setState(param)
        h = zeros(Float64, param.L)
        state = State(0, h)
        return state
    end
end

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

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

    L = param.L
    dt = param.dt
    lam = param.lam
    nu = param.nu
    D = param.D

    while state.t < t_now + t_run

        ### Simulation algorithm here
        
        ###

        state.t += dt
    end
end

# function to get a movie of the roughening dynamics
function get_movie!(state, param, rng, t_gap, n_frame, in_fps)
    x_array = 1:param.L

    # create a movie of the aggregation process
    anim = @animate for frame in 1:n_frame
        run_simulation!(state, param, t_gap, rng)
        plot(x_array, state.h, xlims=(0,param.L), ylims=(minimum(state.h),maximum(state.h)))
    end

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


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

# saturating power-law
plaw_sat(x,p) = p[1] .* ((x./p[2]).^(-10.0*p[3]) .+ 1) .^ (-0.1)

# function to estimate the exponents of 
function get_exponents(lam, nu, D, dt, Ls, ts, rng, n_seeds)
    n_Ls = length(Ls)     # number of L's
    n_pts = length(ts[1]) # number of time points

    ws = zeros(Float64, (n_Ls,n_pts)) # roughness
    wsats = zeros(Float64, n_Ls)      # saturation roughnesses
    tsats = zeros(Float64, n_Ls)      # saturation times

    # cycle through different L's
    for i in 1:n_Ls 
        L = Ls[i]
        N = L
        println("running L = "*string(L))

        # start new system with this L
        param = Rec.setParam(L, lam, nu, D, dt)
        state = Rec.setState(param)
        
        # run different seeds with this L
        for j in ProgressBar(1:n_seeds)

            # re-start with an empty system
            fill!(state.h, 0.0)
            state.t = 0

            # run for a series of time gaps specified in ts[i] for this L=Ls[i]
            for k in 1:n_pts
                t_gap = ts[i][k] - state.t
                run_simulation!(state, param, t_gap, rng)

                # calculate average height and roughness
                h_avg = sum(state.h) / param.L
                w = sqrt(sum((state.h .- h_avg).^2) / param.L)
                ws[i,k] += w 
            end
        end
        # normalize roughness by # of seeds
        ws[i,:] = ws[i,:] ./ n_seeds

        # fit roughness to a saturating power law
        fit = curve_fit(plaw_sat, ts[i], ws[i,:], [10.0,ts[i][n_pts]/2.0,0.3]; lower=zeros(Float64,3).+0.1, upper=Array([1000,ts[i][n_pts],2.0]))

        # save saturation roughness and saturation time
        wsats[i] = fit.param[1]
        tsats[i] = fit.param[2]
    end

    # fit saturation roughness
    fit = curve_fit(power_law, Ls, wsats, [1.0,1.0]; lower=zeros(Float64,2).+0.1)
    α = fit.param[2]

    # fit saturation time
    fit = curve_fit(power_law, Ls, tsats, [1.0,1.0]; lower=zeros(Float64,2).+0.1)
    z = fit.param[2]

    println("alpha = "*string(α)*", z="*string(z))

    # plot scaled roughness over time to demonstrate scaling
    plot(xaxis=:log10,yaxis=:log10,ylabel="w(t) / L^"*string(round(α,digits=2)), xlabel="t/L^"*string(round(z,digits=2)))
    for i in 1:n_Ls
        plot!(ts[i] ./ (Ls[i]^z), ws[i,:] ./ (Ls[i]^α), label="L="*string(Ls[i]))
    end

    savefig("Rec10_KPZ_scaled_w.png")
end