export get_F_all

function get_F_all(s::Simulation)
    F_au = zeros(Float64, N+2, 3)
    F_au_n_n= zeros(Float64, N+2, 3)
    F_au_o_n= zeros(Float64, N+2, 3)
    F_n_o_n_temp = MMatrix{2, 3, Float64, 6}(zeros(2, 3))
    F_n_o_n= zeros(Float64, N+2, 3)
    F_neutral= zeros(Float64, N+2, 3)

    F_au_n_i = zeros(Float64, N+2, 3)
    F_au_o_i= zeros(Float64, N+2, 3)
    F_n_o_i_temp = MMatrix{2, 3, Float64, 6}(zeros(2, 3))
    F_n_o_i = zeros(Float64, N+2, 3)
    F_image_i_temp = MMatrix{2, 3, Float64, 6}(zeros(2, 3))
    F_image_i = zeros(Float64, N+2, 3)
    F_ion = zeros(Float64, N+2, 3)

    F_au_n_c = zeros(Float64, N+2, 3)
    F_au_o_c = zeros(Float64, N+2, 3)
    F_coupling = zeros(Float64, N+2, 3)
    get_F_all_inner!(s, F_au, F_au_n_n , F_au_o_n,F_n_o_n_temp ,F_n_o_n,F_neutral,F_au_n_i ,F_au_o_i,F_n_o_i_temp ,F_n_o_i ,F_image_i_temp ,F_image_i ,F_ion ,F_au_n_c ,F_au_o_c ,F_coupling)
end

function get_F_all_inner!(s::Simulation, F_au, F_au_n_n , F_au_o_n,F_n_o_n_temp ,F_n_o_n,F_neutral,F_au_n_i ,F_au_o_i,F_n_o_i_temp ,F_n_o_i ,F_image_i_temp ,F_image_i ,F_ion ,F_au_n_c ,F_au_o_c ,F_coupling
)
    #Units are STILL not clear
    F_au .= get_F_au_au(s)
    F_au_n_n .= get_F_au_n_neutral(s)
    F_au_o_n .= get_F_au_o_neutral(s)
    F_n_o_n_temp .= get_F_n_o_neutral(s)
    F_n_o_n .= [F_n_o_n_temp; zeros(Float64, N, 3)]
    F_neutral .= @. F_au_n_n + F_au_o_n + F_n_o_n - F_au #add half force of gold-gold?

    F_au_n_i .= get_F_au_n_ion(s)
    F_au_o_i.= get_F_au_o_ion(s)
    F_n_o_i_temp .= get_F_n_o_ion(s)
    F_n_o_i .= [F_n_o_i_temp; zeros(Float64, N, 3)]
    F_image_i_temp .= get_F_image_ion(s)
    F_image_i .= [F_image_i_temp; zeros(Float64, N, 3)]
    F_ion .= @. F_au_n_i + F_au_o_i + F_n_o_i + F_image_i - F_au  #add half force of gold-gold?

    F_au_n_c .= get_F_au_n_coup(s)
    F_au_o_c .= get_F_au_o_coup(s)
    F_coupling .= @. F_au_n_c + F_au_o_c

    #diabatic matrix elements of NO-Au Hamiltonian
    s.dhp_neutral .= -F_neutral
    s.dhp_ion .= -F_ion
    s.dhp_coup .= -F_coupling
end
#---------------------------Au-N Coupling----------------------------------
@inline function get_F_au_n_coup_loop!(x::float_array, F::float_array, i,
     delta_xn, cutoff_test, bool_cutoff_test, s::Simulation)::float_array
    @views delta_xn =  x[i+2, :] - x[1, :]
    delta_xn[1] = @.delta_xn[1] - cell[1]*round(delta_xn[1]/cell[1])
    delta_xn[2] = @.delta_xn[2] - cell[2]*round(delta_xn[2]/cell[2])

    cutoff_test .= abs.(delta_xn) .- cutoff
    bool_cutoff_test .= cutoff_test .<= 0.0
    sum_bool = sum(bool_cutoff_test)
    if sum_bool != zero(sum_bool)
        norm_delta_xn = norm(delta_xn)
        if norm_delta_xn <= cutoff
            term1 = exp(coup_β_N*norm_delta_xn)
            term2 = (1.0 + coup_b_N*term1)^2 * norm_delta_xn
            term3 = coup_a_N * coup_b_N * coup_β_N/term2

            dV_dri = @. -term1*delta_xn*term3
            dV_drn = -dV_dri
            @views F[1, :] .= F[1, :] .- dV_drn
            @views F[i+2, :] .= F[i+2, :] .- dV_dri
        end
    end
    return F
end

@inbounds @inline function get_F_au_n_coup(s::Simulation)::float_array
    F = zeros(Float64, N+2, 3)
    delta_xn = MVector{3, Float64}(zeros(Float64, 3))
    cutoff_test = MVector{3, Float64}(zeros(Float64, 3))
    bool_cutoff_test = BitArray{1}(undef, 3)
    for i in 1:N
        F = get_F_au_n_coup_loop!(s.x, F, i, delta_xn, cutoff_test, bool_cutoff_test, s)
    end
    return F
end
#---------------------------Au-O Coupling----------------------------------
@inline function get_F_au_o_coup_loop!(x::float_array, F::float_array, i, delta_xo,
    cutoff_test, bool_cutoff_test, s::Simulation)::float_array
    @views delta_xo = x[i+2, :] .- x[2, :]
    delta_xo[1] = delta_xo[1] - cell[1]*round(delta_xo[1]/cell[1])
    delta_xo[2] = delta_xo[2] - cell[2]*round(delta_xo[2]/cell[2])

    cutoff_test .= abs.(delta_xo) .- cutoff
    bool_cutoff_test .= cutoff_test .<= 0.0
    sum_bool = sum(bool_cutoff_test)
    if sum_bool != zero(sum_bool)
        norm_delta_xo = norm(delta_xo)
        if norm_delta_xo <= cutoff
            term1 = exp(coup_β_O*norm_delta_xo)
            term2 = (1.0 + coup_b_O*term1)^2 * norm_delta_xo
            term3 = coup_a_O * coup_b_O * coup_β_O/term2

            dV_dri = -term1*delta_xo*term3
            dV_dro = -dV_dri
            @views F[2, :] .= F[2, :] .- dV_dro
            @views F[i+2, :] .= F[i+2, :] .- dV_dri
        end
    end
    return F
end

@inline @inbounds function get_F_au_o_coup(s::Simulation)::float_array
    F = zeros(Float64, N+2, 3)
    delta_xo = MVector{3, Float64}(zeros(Float64, 3))
    cutoff_test = MVector{3, Float64}(zeros(Float64, 3))
    bool_cutoff_test = BitArray{1}(undef, 3)
    for i in 1:N
        F = get_F_au_o_coup_loop!(s.x, F, i, delta_xo, cutoff_test, bool_cutoff_test, s)
    end
    return F
end

#---------------------------Ionic N O ----------------------------------

function get_F_n_o_ion(s::Simulation)::MMatrix{2, 3, Float64, 6}
    F = MMatrix{2, 3, Float64, 6}(zeros(2, 3))
    norm_r_no = norm(s.Δ_no)
    term = 2*F_i *δ_i * (exp(-2.0 * δ_i * (norm_r_no - rNO_e_i)) - exp(-δ_i*(norm_r_no - rNO_e_i)))
    F[1, :] .=  term/norm_r_no .* r_no
    F[2, :] .= - F[1, :]
    return F
end
#---------------------------Ionic Au O ----------------------------------
@inline function get_F_au_o_ion_loop(x::Array{Float64,2}, F::float_array, i::Int,
     delta_xo, cutoff_test, bool_cutoff_test, s::Simulation)::float_array
    @views delta_xo .=  x[i+2, :] .- x[2, :]
    delta_xo[1] = delta_xo[1] - cell[1]*round(delta_xo[1]/cell[1])
    delta_xo[2] = delta_xo[2] - cell[2]*round(delta_xo[2]/cell[2])

    cutoff_test .= abs.(delta_xo) .- cutoff
    bool_cutoff_test .= cutoff_test .<= 0.0
    sum_bool = sum(bool_cutoff_test)
    if sum_bool != zero(sum_bool)
        norm_delta_xo = norm(delta_xo)
        if norm_delta_xo <= cutoff
            term = -A_n * exp(-α_n*norm_delta_xo) * α_n/norm_delta_xo
            dV_dri = term * delta_xo
            dV_dro = -dV_dri
            F[2, :] .= view(F, 2, :) .- dV_dro
            F[i+2, :] .= view(F, i+2, :) .- dV_dri
        end
    end
    return F
end

@inline @inbounds function get_F_au_o_ion(s::Simulation)::float_array #check
    F = zeros(Float64, N+2, 3)
    delta_xo = MVector{3, Float64}(zeros(Float64, 3))
    cutoff_test = MVector{3, Float64}(zeros(Float64, 3))
    bool_cutoff_test = BitArray{1}(undef, 3)
    for i in 1:N
        F = get_F_au_o_ion_loop(s.x, F, i, delta_xo, cutoff_test, bool_cutoff_test, s)
    end
    return F
end

#---------------------------Ionic Au N ----------------------------------
@inline @inbounds function get_F_au_n_ion_loop(x::Array{Float64,2}, F::float_array,
     i::Int, delta_xn, cutoff_test, bool_cutoff_test, s::Simulation)::float_array #check
    @views delta_xn .=  x[i+2, :] .- x[1, :]
    delta_xn[1] = delta_xn[1] - cell[1]*round(delta_xn[1]/cell[1])
    delta_xn[2] = delta_xn[2] - cell[2]*round(delta_xn[2]/cell[2])
    cutoff_test .=  abs.(delta_xn) .- cutoff
    bool_cutoff_test .=  cutoff_test .<= 0.0
    sum_bool = sum(bool_cutoff_test)
    if sum_bool != zero(sum_bool)
        norm_delta_xn = norm(delta_xn)
        if norm_delta_xn <= cutoff
            norm_r_no = norm(s.Δ_no)
            r_diff = norm_delta_xn - rN_sur_e_i
            term1 = s.x[2, 3] - s.x[1, 3]
            term2 = exp(-β_i*r_diff)
            term3 = exp(-2.0 * β_i * r_diff)
            term4 = norm_r_no^4
            #println(term2)
            dV_xN_1 = 4.0*term2*(s.Δ_no[1])*term1^2
            dV_xN_1 = dV_xN_1/term4
            dV_xN_4 = 2.0*term3*delta_xn[1]*β_i
            dV_xN_4 = dV_xN_4/norm_delta_xn
            dV_xN_5 = 2.0*term2*delta_xn[1]*term1*term1*β_i
            dV_xN_5 = dV_xN_5/(norm_delta_xn*norm_r_no*norm_r_no)
            dV_xN = B_i*(dV_xN_1+dV_xN_4-dV_xN_5)

            dV_yN_1 = 4.0*term2*(s.Δ_no[2])*term1*term1
            dV_yN_1 = dV_yN_1/term4
            dV_yN_4 = 2.0*term3*delta_xn[2]*β_i
            dV_yN_4 = dV_yN_4/norm_delta_xn
            dV_yN_5 = 2.0*term2*delta_xn[2]*term1*term1*β_i
            dV_yN_5 = dV_yN_5/(norm_delta_xn*norm_r_no*norm_r_no)
            dV_yN = B_i*(dV_yN_1+dV_yN_4-dV_yN_5)

            dV_zN_1 = 4.0*term2*(s.Δ_no[3])*term1*term1
            dV_zN_1 = dV_zN_1/term4
            dV_zN_2 = 4.0*term2*term1
            dV_zN_2 = dV_zN_2/(norm_r_no*norm_r_no)
            dV_zN_4 = 2.0*term3*delta_xn[3]*β_i
            dV_zN_4 = dV_zN_4/norm_delta_xn
            dV_zN_5 = 2.0*term2*delta_xn[3]*term1*term1*β_i
            dV_zN_5 = dV_zN_5/(norm_delta_xn*norm_r_no*norm_r_no)
            dV_zN = B_i*(dV_zN_1+dV_zN_2+dV_zN_4-dV_zN_5)

            dV_xO_1 = 4.0*term2*(s.Δ_no[1])*term1*term1
            dV_xO_1 = dV_xO_1/term4
            dV_xO = B_i*(-dV_xO_1)

            dV_yO_1 = 4.0*term2*(s.Δ_no[2])*term1*term1
            dV_yO_1 = dV_yO_1/term4
            dV_yO = B_i*(-dV_yO_1)

            dV_zO_1 = 4.0*term2*(s.Δ_no[3])*term1*term1
            dV_zO_1 = dV_zO_1/term4
            dV_zO_2 = 4.0*term2*term1
            dV_zO_2 = dV_zO_2/(norm_r_no*norm_r_no)
            dV_zO = B_i*(-dV_zO_1-dV_zO_2)

            dV_xi_3 = 2.0*term3*delta_xn[1]*β_i
            dV_xi_3 = dV_xi_3/norm_delta_xn
            dV_xi_4 = 2.0*term2*delta_xn[1]*term1*term1*β_i
            dV_xi_4 = dV_xi_4/(norm_delta_xn*norm_r_no*norm_r_no)
            dV_xi = B_i*(-dV_xi_3+dV_xi_4)

            dV_yi_3 = 2.0*term3*delta_xn[2]*β_i
            dV_yi_3 = dV_yi_3/norm_delta_xn
            dV_yi_4 = 2.0*term2*delta_xn[2]*term1*term1*β_i
            dV_yi_4 = dV_yi_4/(norm_delta_xn*norm_r_no*norm_r_no)
            dV_yi = B_i*(-dV_yi_3+dV_yi_4)

            dV_zi_3 = 2.0*term3*delta_xn[3]*β_i
            dV_zi_3 = dV_zi_3/norm_delta_xn
            dV_zi_4 = 2.0*term2*delta_xn[3]*term1*term1*β_i
            dV_zi_4 = dV_zi_4/(norm_delta_xn*norm_r_no*norm_r_no)
            dV_zi = B_i*(-dV_zi_3+dV_zi_4)

            #------   Calculating derivatives of cutoff function

            r_diff = cutoff - rN_sur_e_i

            term2 = exp(-β_i*r_diff)
            term3 = exp(-2.0*β_i*r_diff)

            dVc_xN_1 = 4.0*term2*(s.Δ_no[1])*term1*term1
            dVc_xN_1 = dVc_xN_1/term4
            dVc_xN = B_i*(dVc_xN_1)

            dVc_yN_1 = 4.0*term2*(s.Δ_no[2])*term1*term1
            dVc_yN_1 = dVc_yN_1/term4
            dVc_yN = B_i*(dVc_yN_1)

            dVc_zN_1 = 4.0*term2*(s.x[1, 3] - s.x[2, 3])*term1*term1
            dVc_zN_1 = dVc_zN_1/term4
            dVc_zN_2 = 4.0*term2*term1
            dVc_zN_2 = dVc_zN_2/(norm_r_no*norm_r_no)
            dVc_zN = B_i*(dVc_zN_1+dVc_zN_2)

            dVc_xO_1 = 4.0*term2*(s.Δ_no[1])*term1*term1
            dVc_xO_1 = dVc_xO_1/term4
            dVc_xO = B_i*(-dVc_xO_1)

            dVc_yO_1 = 4.0*term2*(s.Δ_no[2])*term1*term1
            dVc_yO_1 = dVc_yO_1/term4
            dVc_yO = B_i*(-dVc_yO_1)

            dVc_zO_1 = 4.0*term2*(s.Δ_no[3])*term1*term1
            dVc_zO_1 = dVc_zO_1/term4
            dVc_zO_2 = 4.0*term2*term1
            dVc_zO_2 = dVc_zO_2/(norm_r_no*norm_r_no)
            dVc_zO = B_i*(-dVc_zO_1-dVc_zO_2)

            dVc_xi = 0.0
            dVc_yi = 0.0
            dVc_zi = 0.0
            F[1, 1] = F[1, 1] - dV_xN + dVc_xN
            F[1, 2] = F[1, 2] - dV_yN + dVc_yN
            F[1, 3] = F[1, 3] - dV_zN + dVc_zN

            F[2, 1] = F[2, 1] - dV_xO + dVc_xO
            F[2, 2] = F[2, 2] - dV_yO + dVc_yO
            F[2, 3] = F[2, 3] - dV_zO + dVc_zO

            F[i+2, 1] = F[i+2, 1] - dV_xi + dVc_xi
            F[i+2, 2] = F[i+2, 2] - dV_yi + dVc_yi
            F[i+2, 3] = F[i+2, 3] - dV_zi + dVc_zi
        end
    end
    return F
end

@inline function get_F_au_n_ion(s::Simulation)::float_array
    F = zeros(Float64, N+2, 3)
    delta_xn = MVector{3, Float64}(zeros(Float64, 3))
    cutoff_test = MVector{3, Float64}(zeros(Float64, 3))
    bool_cutoff_test = BitArray{1}(undef, 3)
    for i in 1:N
        F = get_F_au_n_ion_loop(s.x, F, i, delta_xn, cutoff_test, bool_cutoff_test, s)
    end
    return F
end

#---------------------------Ionic Image Potential----------------------------------
function get_F_image_ion(s::Simulation)::MMatrix{2, 3, Float64, 6}  #check
    F = MMatrix{2, 3, Float64, 6}(zeros(2, 3))
    @views zcom = zcom_NO(s.x[1:2, :])
    term = (D_i*(zcom - z0_i))/(C_i^2 + (zcom-z0_i)^2)^(1.5)
    F[1, 3] = term * mass_arr[1]/(mass_arr[1] + mass_arr[2])
    F[2, 3] = term * mass_arr[2]/(mass_arr[1] + mass_arr[2])
    return -F
end
#---------------------------Ionic N-O----------------------------------
function get_F_n_o_ion(s::Simulation)::MMatrix{2, 3, Float64, 6}
    F = MMatrix{2, 3, Float64, 6}(undef)
    norm_r_no = norm(s.Δ_no)
    term = 2*F_i*δ_i*exp(-δ_i*(norm_r_no - rNO_e_i))*(1-exp(-δ_i*(norm_r_no - rNO_e_i)))
    F[1, :] .= term/norm_r_no .* s.Δ_no
    F[2, :] .= -F[1, :]
    return -F
end

#---------------------------Neutral N-O----------------------------------

@inline @inbounds function get_F_n_o_neutral(s::Simulation)::MMatrix{2, 3, Float64, 6}
    F = MMatrix{2, 3, Float64, 6}(undef)
    norm_r_no = norm(s.Δ_no)
    rdiff = norm_r_no - r_0_NO
    term = 2*F_n*δ_n*(exp(-2.0*δ_n * rdiff) - exp(-δ_n * rdiff))
    #F[1, 1] = -term * (s.x[2, 1] - s.x[1, 1]) / norm_r_no
    F[1, :] .= term/norm_r_no .* s.Δ_no
    F[2, :] .= -F[1, :]
    return F
end
#---------------------------Neutral AU_O----------------------------------
@inline function get_F_au_o_neutral_loop(x::float_array, i::Int, F::float_array,
     delta_xo, cutoff_test, bool_cutoff_test, s::Simulation)::float_array
    @views delta_xo .= x[i+2, :] .- x[2, :]
    delta_xo[1] = delta_xo[1] - cell[1]*round(delta_xo[1]/cell[1])
    delta_xo[2] = delta_xo[2] - cell[2]*round(delta_xo[2]/cell[2])

    cutoff_test .= abs.(delta_xo) .- cutoff
    bool_cutoff_test .= cutoff_test .<= 0.0
    sum_bool = sum(bool_cutoff_test)
    if sum_bool != zero(sum_bool)
        norm_delta_xo = norm(delta_xo)
        if norm_delta_xo <= cutoff
            term = -A_n * exp(-α_n * norm_delta_xo) * α_n/norm_delta_xo
            dV_dxi = term * delta_xo
            dV_dxo = -dV_dxi

            F[2,:] .= view(F, 2, :) .- dV_dxo
            F[i+2,:] .= view(F, i+2, :) .- dV_dxi
        end
    end
    return F
end

@inbounds @inline function get_F_au_o_neutral(s::Simulation)::float_array
    F = zeros(Float64, N+2, 3)
    delta_xo = MVector{3, Float64}(zeros(Float64, 3))
    cutoff_test = MVector{3, Float64}(zeros(Float64, 3))
    bool_cutoff_test = BitArray{1}(undef, 3)
    for i in 1:N
        F = get_F_au_o_neutral_loop(s.x, i, F, delta_xo, cutoff_test, bool_cutoff_test, s)
    end
    return F
end

#---------------------------Neutral AU_N----------------------------------

@inline function get_F_au_n_neutral_loop(x::Array{Float64,2}, i::Int, F::float_array,
     delta_xn, cutoff_test, bool_cutoff_test, s::Simulation)::float_array #maybe check
    @views delta_xn .=  x[i+2, :] .- x[1, :]
    delta_xn[1] = delta_xn[1] - cell[1]*round(delta_xn[1]/cell[1])
    delta_xn[2] = delta_xn[2] - cell[2]*round(delta_xn[2]/cell[2])

    cutoff_test .= abs.(delta_xn) .- cutoff
    bool_cutoff_test .=  cutoff_test .<= 0.0
    sum_bool = sum(bool_cutoff_test)
    if sum_bool != zero(sum_bool)
        norm_delta_xn = norm(delta_xn)
        if norm_delta_xn <= cutoff
            term = -B_n * exp(-β_n * norm_delta_xn) * β_n/norm_delta_xn
            dV_dxi =  term * delta_xn
            dV_dxN = -dV_dxi

            F[1,:] .= view(F, 1, :) .- dV_dxN
            F[i+2,:] .= view(F, i+2, :) .- dV_dxi
        end
    end
    return F
end

@inbounds @inline function get_F_au_n_neutral(s::Simulation)::float_array #check
    F = zeros(Float64, N+2, 3)
    delta_xn = MVector{3, Float64}(zeros(Float64, 3))
    cutoff_test = MVector{3, Float64}(zeros(Float64, 3))
    bool_cutoff_test = BitArray{1}(undef, 3)
    for i in 1:N
        F = get_F_au_n_neutral_loop(s.x, i, F, delta_xn, cutoff_test, bool_cutoff_test, s)
    end
    return F
end

#---------------------------GOLD AU AU LATTICE----------------------------------
@inline function F_au_au_loop_if(x::float_array, F::float_array, i::Int, j::Int,
     xm, xi, temp,r, qf_1, s::Simulation)::float_array
    if s.nn_arr[i, j] != 0
        m = s.nn_arr[i, j] + 2
        @views xm .= x[m, :]
        @views xi .= x[i+2, :]
        temp .= xm .- xi
        temp .= temp./cell
        temp .= temp .- floor.(temp .+ 0.5)
        @views r .=  temp .* cell .- r0[:, j]
        if j == 1 || j== 7
            mul!(qf_1, d1_new, r)
            F[i+2, :] .= view(F, i+2, :) .- qf_1
        elseif j == 2 || j== 8
            mul!(qf_1, d2_new, r)
            F[i+2, :] .= view(F, i+2, :) .- qf_1
        elseif j == 3 || j== 9
            mul!(qf_1, d3_new, r)
            F[i+2, :] .= view(F, i+2, :) .- qf_1
        elseif j == 4 || j== 10
            mul!(qf_1, d4_new, r)
            F[i+2, :] .= view(F, i+2, :) .- qf_1
        elseif j == 5|| j== 11
            mul!(qf_1, d5_new, r)
            F[i+2, :] .= view(F, i+2, :) .- qf_1
        elseif j == 6 || j== 12
            mul!(qf_1, d6_new, r)
            F[i+2, :] .= view(F, i+2, :) .- qf_1
        end
    end
    return F
end
@inline @inbounds function F_au_au_loop(x::float_array, F::float_array,
     s::Simulation)::float_array #check
    xm = MVector{3, Float64}(zeros(Float64, 3))
    xi = MVector{3, Float64}(zeros(Float64, 3))
    temp = MVector{3, Float64}(zeros(Float64, 3))
    r = MVector{3, Float64}(zeros(Float64, 3))
    qf_1 = MVector{3, Float64}(zeros(Float64, 3))
    for i in 1:N
        for j in 1:12
            F .= F_au_au_loop_if(x, F, i, j, xm, xi, temp,r, qf_1, s)
        end
    end
    return F
end

@inline function get_F_au_au(s::Simulation)::float_array
    F = zeros(Float64, N+2, 3)
    F = F_au_au_loop(s.x, F, s)
    return F
end


import Base: +, setindex!, filter!

# ------------ Parameter store ------------ #

struct ParameterStore
    params::Dict{Address,Any}
    ParameterStore() = new(Dict{Address, Any}())
    ParameterStore(d::Dict{Address, Any}) = new(d)
end
haskey(ps::ParameterStore, addr) = haskey(ps.params, addr)
setindex!(ps::ParameterStore, val, addr) = ps.params[addr] = val

Zygote.@adjoint ParameterStore(params) = ParameterStore(params), store_grad -> (nothing,)

function +(a::ParameterStore, b::ParameterStore)
    params = Dict{Address, Any}()
    for (k, v) in Iterators.flatten((a.params, b.params))
        if !haskey(params, k)
            params[k] = v
        else
            params[k] += v
        end
    end
    ParameterStore(params)
end

# ------------ Backpropagation contexts ------------ #

abstract type BackpropagationContext <: ExecutionContext end

# Learnable parameters
mutable struct ParameterBackpropagateContext{T <: Trace} <: BackpropagationContext
    tr::T
    weight::Float64
    initial_params::Parameters
    params::ParameterStore
    param_grads::Gradients
end
ParameterBackpropagate(tr::T, init, params) where T <: Trace = ParameterBackpropagateContext(tr, 0.0, init, params, Gradients())
ParameterBackpropagate(tr::T, init, params, param_grads::Gradients) where {T <: Trace, K <: UnconstrainedSelection} = ParameterBackpropagateContext(tr, 0.0, init, params, param_grads)

# Choice sites
mutable struct ChoiceBackpropagateContext{T <: Trace, K <: UnconstrainedSelection} <: BackpropagationContext
    tr::T
    weight::Float64
    initial_params::Parameters
    params::ParameterStore
    choice_grads::Gradients
    select::K
end
ChoiceBackpropagate(tr::T, init, params, choice_grads) where {T <: Trace, K <: UnconstrainedSelection} = ChoiceBackpropagateContext(tr, 0.0, init, params, choice_grads, UnconstrainedAllSelection())
ChoiceBackpropagate(tr::T, init, params, choice_grads, sel::K) where {T <: Trace, K <: UnconstrainedSelection} = ChoiceBackpropagateContext(tr, 0.0, init, params, choice_grads, sel)

# ------------ Learnable ------------ #

read_parameter(ctx::K, addr::Address) where K <: BackpropagationContext = read_parameter(ctx, ctx.params, addr)
read_parameter(ctx::K, params::ParameterStore, addr::Address) where K <: BackpropagationContext = get_top(ctx.initial_params, addr)

Zygote.@adjoint function read_parameter(ctx, params, addr)
    ret = read_parameter(ctx, params, addr)
    fn = param_grad -> begin
        state_grad = nothing
        params_grad = ParameterStore(Dict{Address, Any}(addr => param_grad))
        (state_grad, params_grad, nothing)
    end
    return ret, fn
end

# ------------ Call sites ------------ #

# Grads for learnable parameters.
simulate_call_pullback(params, param_grads, cl::T, args) where T <: CallSite = cl.ret

Zygote.@adjoint function simulate_call_pullback(params, param_grads, cl::HierarchicalCallSite, args)
    ret = simulate_call_pullback(params, param_grads, cl, args)
    fn = ret_grad -> begin
        arg_grads = accumulate_parameter_gradients!(params, param_grads, cl, ret_grad)
        (nothing, nothing, nothing, arg_grads)
    end
    return ret, fn
end

# Utility.
merge(tp1::Tuple{}, tp2::Tuple{}) = tp1
merge(tp1::Tuple{Nothing}, tp2::Tuple{Nothing}) where T = tp1
merge(tp1::NTuple{N, Float64}, tp2::NTuple{N, Float64}) where N = [tp1[i] + tp2[i] for i in 1 : N]
merge(tp1::Array{Float64}, tp2::NTuple{N, Float64}) where N = [tp1[i] + tp2[i] for i in 1 : N]

Zygote.@adjoint function simulate_call_pullback(params, param_grads, cl::VectorizedCallSite{typeof(markov)}, args)
    ret = simulate_call_pullback(params, param_grads, cl, args)
    fn = ret_grad -> begin
        arg_grads = accumulate_parameter_gradients!(params, param_grads, get_sub(cl, cl.len), ret_grad)
        for i in (cl.len - 1) : -1 : 1
            arg_grads = accumulate_parameter_gradients!(params, param_grads, get_sub(cl, i), arg_grads)
        end
        (nothing, nothing, nothing, arg_grads)
    end
    return ret, fn
end

Zygote.@adjoint function simulate_call_pullback(params, param_grads, cl::VectorizedCallSite{typeof(plate)}, args)
    ret = simulate_call_pullback(params, param_grads, cl, args)
    fn = ret_grad -> begin
        arg_grads = accumulate_parameter_gradients!(params, param_grads, get_sub(cl, cl.len), ret_grad[1])
        for i in 2 : cl.len
            new = accumulate_parameter_gradients!(params, param_grads, get_sub(cl, i), ret_grad[i])
            arg_grads = merge(arg_grads, new)
        end
        (nothing, nothing, nothing, arg_grads)
    end
    return ret, fn
end

# Grads for choices with differentiable logpdfs.
simulate_choice_pullback(params, choice_grads, choice_selection, cl::T, args) where T <: CallSite = cl.ret

Zygote.@adjoint function simulate_choice_pullback(params, choice_grads, choice_selection, cl, args)
    ret = simulate_choice_pullback(params, choice_grads, choice_selection, cl, args)
    fn = ret_grad -> begin
        arg_grads, choice_vals, choice_grads = choice_gradients(params, choice_grads, choice_selection, cl, ret_grad)
        (nothing, nothing, nothing, (choice_vals, choice_grads), arg_grads)
    end
    return ret, fn
end

# ------------ Accumulate gradients ------------ #

function accumulate_parameter_gradients!(initial_params, param_grads, cl::HierarchicalCallSite, ret_grad, scaler::Float64 = 1.0)
    fn = (args, params) -> begin
        ctx = ParameterBackpropagate(cl.trace, initial_params, params, param_grads)
        ret = ctx(cl.fn, args...)
        (ctx.weight, ret)
    end
    blank = ParameterStore()
    _, back = Zygote.pullback(fn, cl.args, blank)
    arg_grads, ps_grad = back((1.0, ret_grad))
    if !(ps_grad isa Nothing)
        for (addr, grad) in ps_grad.params
            push!(param_grads, addr, scaler .* grad)
        end
    end
    return arg_grads
end

function accumulate_parameter_gradients!(initial_params, param_grads, cl::HierarchicalCallSite, ret_grad::Tuple, scaler::Float64 = 1.0)
    fn = (args, params) -> begin
        ctx = ParameterBackpropagate(cl.trace, initial_params, params, param_grads)
        ret = ctx(cl.fn, args...)
        (ctx.weight, ret)
    end
    blank = ParameterStore()
    _, back = Zygote.pullback(fn, cl.args, blank)
    arg_grads, ps_grad = back((1.0, ret_grad...))
    if !(ps_grad isa Nothing)
        for (addr, grad) in ps_grad.params
            push!(param_grads, addr, scaler .* grad)
        end
    end
    return arg_grads
end

function accumulate_parameter_gradients!(initial_params, param_grads, cl::VectorizedCallSite{typeof(markov)}, ret_grad, scaler::Float64 = 1.0) where T <: CallSite
    addr = gensym()
    v_params = parameters(addr => initial_params)
    tr = Trace()
    add_call!(tr, addr, cl)
    fn = (args, params) -> begin
        ctx = ParameterBackpropagate(tr, v_params, params, param_grads)
        ret = ctx(markov, addr, cl.fn, cl.len, args...)
        (ctx.weight, ret)
    end
    blank = ParameterStore()
    _, back = Zygote.pullback(fn, cl.args, blank)
    arg_grads, ps_grad = back((1.0, ret_grad))
    if !(ps_grad isa Nothing)
        for (addr, grad) in ps_grad.params
            push!(param_grads, addr, scaler .* grad)
        end
    end
    return arg_grads
end

function accumulate_parameter_gradients!(initial_params, param_grads, cl::VectorizedCallSite{typeof(plate)}, ret_grad, scaler::Float64 = 1.0) where T <: CallSite
    addr = gensym()
    v_params = parameters(addr => initial_params)
    tr = Trace()
    add_call!(tr, addr, cl)
    fn = (args, params) -> begin
        ctx = ParameterBackpropagate(tr, v_params, params, param_grads)
        ret = ctx(plate, addr, cl.fn, args)
        (ctx.weight, ret)
    end
    blank = ParameterStore()
    _, back = Zygote.pullback(fn, cl.args, blank)
    arg_grads, ps_grad = back((1.0, ret_grad))
    if !(ps_grad isa Nothing)
        for (addr, grad) in ps_grad.params
            push!(param_grads, addr, scaler .* grad)
        end
    end
    return arg_grads
end

# ------------ Compute choice gradients ------------ #

function filter!(choice_grads, cl::HierarchicalCallSite, grad_tr::NamedTuple, sel::K) where K <: UnconstrainedSelection
    values = ConstrainedHierarchicalSelection()
    for (k, v) in cl.trace.choices
        has_top(sel, k) && begin
            push!(values, k, v.val)
            push!(choice_grads, k, grad_tr.choices[k].val)
        end
    end
    return values
end

function choice_gradients(initial_params, choice_grads, choice_selection, cl, ret_grad)
    call = cl.fn
    fn = (args, trace) -> begin
        ctx = ChoiceBackpropagate(trace, initial_params, ParameterStore(), choice_grads, choice_selection)
        ret = ctx(call, args...)
        (ctx.weight, ret)
    end
    _, back = Zygote.pullback(fn, cl.args, cl.trace)
    arg_grads, grad_ref = back((1.0, ret_grad))
    choice_vals = filter!(choice_grads, cl, grad_ref, choice_selection)
    return arg_grads, choice_vals, choice_grads
end

# ------------ Convenience getters ------------ #

function get_choice_gradients(cl::T, ret_grad) where T <: CallSite
    choice_grads = Gradients()
    choice_selection = UnconstrainedAllSelection()
    _, vals, _ = choice_gradients(Parameters(), choice_grads, choice_selection, cl, ret_grad)
    return vals, choice_grads
end

function get_choice_gradients(params, cl::T, ret_grad) where T <: CallSite
    choice_grads = Gradients()
    choice_selection = UnconstrainedAllSelection()
    _, vals, _ = choice_gradients(params, choice_grads, choice_selection, cl, ret_grad)
    return vals, choice_grads
end

function get_choice_gradients(sel::K, cl::T, ret_grad) where {T <: CallSite, K <: UnconstrainedSelection}
    choice_grads = Gradients()
    _, vals, _ = choice_gradients(Parameters(), choice_grads, sel, cl, ret_grad)
    return vals, choice_grads
end

function get_choice_gradients(sel::K, params, cl::T, ret_grad) where {T <: CallSite, K <: UnconstrainedSelection}
    choice_grads = Gradients()
    _, vals, _ = choice_gradients(params, choice_grads, sel, cl, ret_grad)
    return vals, choice_grads
end

function get_parameter_gradients(params, cl::HierarchicalCallSite, ret_grad, scaler::Float64 = 1.0)
    param_grads = Gradients()
    accumulate_parameter_gradients!(params, param_grads, cl, ret_grad, scaler)
    return param_grads
end

function get_parameter_gradients(params, cl::VectorizedCallSite, ret_grad, scaler::Float64 = 1.0)
    param_grads = Gradients()
    accumulate_parameter_gradients!(params, param_grads, cl, ret_grad, scaler)
    for k in keys(param_grads.tree)
        return param_grads.tree[k]
    end
end

# ------------ train ------------ #

function train(params, fn::Function, args...; opt = ADAM(0.05, (0.9, 0.8)), iters = 1000)
    for i in 1 : iters
        _, cl = simulate(params, fn, args...)
        grads = get_parameter_gradients(params, cl, 1.0)
        params = update_parameters(opt, params, grads)
    end
    return params
end

function train(sel::K, params, fn::Function, args...; opt = ADAM(0.05, (0.9, 0.8)), iters = 1000) where K <: ConstrainedSelection
    for i in 1 : iters
        _, cl, _ = generate(sel, params, fn, args...)
        grads = get_parameter_gradients(params, cl, 1.0)
        params = update_parameters(opt, params, grads)
    end
    return params
end

# ------------ includes ------------ #

include("hierarchical/backpropagate.jl")
include("plate/backpropagate.jl")
include("markov/backpropagate.jl")
include("factor/backpropagate.jl")

# ------------ Documentation ------------ #

@doc(
"""
```julia
mutable struct ParameterBackpropagateContext{T <: Trace} <: BackpropagationContext
    tr::T
    weight::Float64
    initial_params::Parameters
    params::ParameterStore
    param_grads::Gradients
end
```
`ParameterBackpropagateContext` is used to compute the gradients of parameters with respect to following objective:

Outer constructors:
```julia
ParameterBackpropagate(tr::T, params) where T <: Trace = ParameterBackpropagateContext(tr, 0.0, params, Gradients())
ParameterBackpropagate(tr::T, params, param_grads::Gradients) where {T <: Trace, K <: UnconstrainedSelection} = ParameterBackpropagateContext(tr, 0.0, params, param_grads)
```
""", ParameterBackpropagateContext)

@doc(
"""
```julia
mutable struct ChoiceBackpropagateContext{T <: Trace} <: BackpropagationContext
    tr::T
    weight::Float64
    initial_params::Parameters
    params::ParameterStore
    param_grads::Gradients
end
```
`ChoiceBackpropagateContext` is used to compute the gradients of choices with respect to following objective:

Outer constructors:
```julia
ChoiceBackpropagate(tr::T, init_params, params, choice_grads) where {T <: Trace, K <: UnconstrainedSelection} = ChoiceBackpropagateContext(tr, 0.0, params, choice_grads, UnconstrainedAllSelection())
ChoiceBackpropagate(tr::T, init_params, params, choice_grads, sel::K) where {T <: Trace, K <: UnconstrainedSelection} = ChoiceBackpropagateContext(tr, 0.0, params, choice_grads, sel)
```
""", ChoiceBackpropagateContext)

@doc(
"""
```julia
gradients = get_choice_gradients(params, cl::T, ret_grad) where T <: CallSite
gradients = get_choice_gradients(cl::T, ret_grad) where T <: CallSite
```

Returns a `Gradients` object which tracks the gradients with respect to the objective of random choices with differentiable `logpdf` in the program.
""", get_choice_gradients)

@doc(
"""
```julia
gradients = get_parameter_gradients(params, cl::T, ret_grad, scaler::Float64 = 1.0) where T <: CallSite
```

Returns a `Gradients` object which tracks the gradients of the objective with respect to parameters in the program.
""", get_parameter_gradients)


# Modified scheme of Bassi and Rebay (BR2)

"""
  Penalty for BR2
"""
struct BR2Penalty{Tsol, Tres} <: AbstractDiffusionPenalty

  # primal method
  delta_w_n::Matrix{Tsol}
  qL::Array{Tres, 3}
  qR::Array{Tres, 3}
  t1L::Array{Tres, 3}
  t1R::Array{Tres, 3}
  t2L::Array{Tres, 3}
  t2R::Array{Tres, 3}

  # diff method
  t1_dotL::Array{Tres, 5}
  t1_dotR::Array{Tres, 5}

  # the first L/R indicates if this is element kappa or nu
  # the second L/R indices if the derivative is wrt qL or qR
  t2LL::Array{Tres, 5}
  t2LR::Array{Tres, 5}
  t2RL::Array{Tres, 5}
  t2RR::Array{Tres, 5}

  t3LL::Array{Tres, 5}
  t3LR::Array{Tres, 5}
  t3RL::Array{Tres, 5}
  t3RR::Array{Tres, 5}

  # re-use t1_dotL,R
  t4RL::Array{Tres, 5}
  t4RR::Array{Tres, 5}

  function BR2Penalty{Tsol, Tres}(mesh, sbp, opts) where {Tsol, Tres}

    @unpack mesh numDofPerNode numNodesPerFace numNodesPerElement dim

    delta_w_n = zeros(Tsol, numDofPerNode, numNodesPerFace)
    qL = zeros(Tres, numDofPerNode, numNodesPerElement, dim)
    qR = zeros(Tres, numDofPerNode, numNodesPerElement, dim)
    t1L = zeros(Tres, numDofPerNode, numNodesPerElement, dim)
    t1R = zeros(Tres, numDofPerNode, numNodesPerElement, dim)
    t2L = zeros(Tres, numDofPerNode, numNodesPerFace, dim)
    t2R = zeros(Tres, numDofPerNode, numNodesPerFace, dim)

    t1_dotL = zeros(Tres, numDofPerNode, numDofPerNode, dim, numNodesPerFace,
                   numNodesPerElement)
    t1_dotR = zeros(Tres, numDofPerNode, numDofPerNode, dim, numNodesPerFace,
                     numNodesPerElement)

    t2LL = zeros(Tres, numDofPerNode, numDofPerNode, dim, numNodesPerElement,
                      numNodesPerElement)
    t2LR = zeros(Tres, numDofPerNode, numDofPerNode, dim, numNodesPerElement,
                      numNodesPerElement)

    t2RL = zeros(Tres, numDofPerNode, numDofPerNode, dim, numNodesPerElement,
                      numNodesPerElement)
    t2RR = zeros(Tres, numDofPerNode, numDofPerNode, dim, numNodesPerElement,
                      numNodesPerElement)

    t3LL = zeros(Tres, numDofPerNode, numDofPerNode, dim, numNodesPerElement,
                       numNodesPerElement)
    t3LR = zeros(Tres, numDofPerNode, numDofPerNode, dim, numNodesPerElement,
                       numNodesPerElement)

    t3RL = zeros(Tres, numDofPerNode, numDofPerNode, dim, numNodesPerElement,
                       numNodesPerElement)
    t3RR = zeros(Tres, numDofPerNode, numDofPerNode, dim, numNodesPerElement,
                       numNodesPerElement)

    t4RL = zeros(Tres, numDofPerNode, numDofPerNode, dim, numNodesPerFace,
                     numNodesPerElement)
    t4RR = zeros(Tres, numDofPerNode, numDofPerNode, dim, numNodesPerFace,
                       numNodesPerElement)

    return new(delta_w_n, qL, qR, t1L, t1R, t2L, t2R,
               t1_dotL, t1_dotR, t2LL, t2LR, t2RL, t2RR, t3LL, t3LR, t3RL, t3RR,
               t4RL, t4RR)
  end

end

"""
  Applies the penalty for the SBP-SAT generalization of the modified scheme
  of Bassi and Rebay (BR2).
"""
function applyPenalty(penalty::BR2Penalty{Tsol, Tres}, sbp, params::AbstractParamType,
                      sbpface,
                      diffusion::AbstractDiffusion, iface::Interface,
                      delta_w::AbstractMatrix{Tsol}, theta::AbstractMatrix{Tres},
                      qL_el::AbstractMatrix, qR_el::AbstractMatrix,
                      wL::AbstractMatrix, wR::AbstractMatrix,
                      coordsL::AbstractMatrix, coordsR::AbstractMatrix,
                      nrm_face::AbstractMatrix,
                      alphas::AbstractVector,
                      dxidxL::Abstract3DArray, dxidxR::Abstract3DArray,
                      jacL::AbstractVector, jacR::AbstractVector,
                      res1L::AbstractMatrix, res1R::AbstractMatrix,
                      res2L::AbstractMatrix, res2R::AbstractMatrix
                     ) where {Tsol, Tres}

  fill!(res1L, 0); fill!(res1R, 0)
  fill!(res2L, 0); fill!(res2R, 0)
  numDofPerNode, numNodesPerFace = size(delta_w)
  numNodesPerElement = length(jacL)
  dim = size(nrm_face, 1)

  @unpack penalty delta_w_n qL qR t1L t1R t2L t2R
  fill!(qL, 0); fill!(qR, 0)

  #--------------------------
  # apply T1
  # multiply by normal vector, then R^T B
  #alpha_g = 1/(dim + 1)  # = 1/number of faces of a simplex
  @simd for d1=1:dim
    @simd for j=1:numNodesPerFace
      @simd for k=1:numDofPerNode
        delta_w_n[k, j] = 0.25*delta_w[k, j]*nrm_face[d1, j]
      end
    end

    qL_d = sview(qL, :, :, d1); qR_d = sview(qR, :, :, d1)
    interiorFaceIntegrate!(sbpface, iface, delta_w_n, qL_d, qR_d)

    # interiorFaceIntegrates -= the second output
    @simd for j=1:numNodesPerElement
      @simd for k=1:numDofPerNode
        qR_d[k, j] = -qR_d[k, j]
      end
    end
  end

  # apply Lambda matrix
  applyDiffusionTensor(diffusion, sbp, params, qL_el, wL, coordsL, dxidxL, jacL, iface.elementL, sbpface, iface.faceL,
                       qL, t1L)
  applyDiffusionTensor(diffusion, sbp, params, qR_el, wR, coordsR, dxidxR, jacR,
                       iface.elementR, sbpface, iface.faceR,
                       qR, t1R)

  # apply inverse mass matrix, then apply B*Nx*R*t2L_x + B*Ny*R*t2L_y
  @simd for d1=1:dim
    @simd for j=1:numNodesPerElement
      facL = (1/alphas[1])*jacL[j]/sbp.w[j]
      facR = (1/alphas[2])*jacR[j]/sbp.w[j]
      #facL = (3)*jacL[j]/sbp.w[j]
      #facR = (3)*jacR[j]/sbp.w[j]

      @simd for k=1:numDofPerNode
        t1L[k, j, d1] *= facL
        t1R[k, j, d1] *= facR
      end
    end

    #TODO: make t2L 2D rather than 3D
    t1L_d = ro_sview(t1L, :, :, d1); t1R_d = ro_sview(t1R, :, :, d1)
    t2L_d = sview(t2L, :, :, d1);    t2R_d = sview(t2R, :, :, d1)
    interiorFaceInterpolate!(sbpface, iface, t1L_d, t1R_d, t2L_d, t2R_d)

    @simd for j=1:numNodesPerFace
      @simd for k=1:numDofPerNode
        val = sbpface.wface[j]*nrm_face[d1, j]*(t2L_d[k, j] + t2R_d[k, j])
        res1L[k, j] += val
        res1R[k, j] -= val  # delta_w is reversed for elementR
      end
    end

  end  # end d1

  #----------------------------
  # apply T2 and T3
  @simd for j=1:numNodesPerFace
    @simd for k=1:numDofPerNode
      val1 =  0.5*sbpface.wface[j]*theta[k, j]
      val2 = -0.5*sbpface.wface[j]*delta_w[k, j]
      res1L[k, j] += val1
      res1R[k, j] += val1
      res2L[k, j] += val2
      res2R[k, j] -= val2  # delta_w is reversed for elementR
    end
  end

  return nothing
end

"""
  Applies the Dirichlet penalty matrix T_gammaD.

  **Inputs**

   * penalty: [`AbstractDiffusionPenalty`](@ref)
   * sbp
   * sbpface
   * diffusion: [`AbstractDiffusion`](@ref)
   * bndry: `Boundary` object
   * delta_w: values to multiply the penalty against, `numDofPerNode` x
              `numNodesPerFace`
   * wL: entropy variables for the element, `numDofPerNode` x `numNodesPerElement`
   * nrm_face: (scaled) normal vector at each face node, `dim` x
               `numNodesPerFace`
   * alpha: the alpha_gk parameter (Number)
   * jacL: mapping jacobian determinant, `numNodesPerElement`

  **Inputs/Outputs**

   * res1L: `numDofPerNode` x `numNodesPerFace` array to overwrite with
            the result

"""
#=
function applyDirichletPenalty(penalty::BR2Penalty{Tsol, Tres}, sbp, sbpface,
                      diffusion::AbstractDiffusion, bndry::Boundary,
                      delta_w::AbstractMatrix{Tsol},
                      wL::AbstractMatrix,
                      nrm_face::AbstractMatrix,
                      alpha::Number,
                      jacL::AbstractVector,
                      res1L::AbstractMatrix,
                     ) where {Tsol, Tres}
=#
function applyDirichletPenalty(penalty::BR2Penalty{Tsol, Tres}, sbp,
                      params::AbstractParamType, sbpface,
                      diffusion::AbstractDiffusion, bndry::Boundary,
                      delta_w::AbstractMatrix,
                      uL::AbstractMatrix,
                      wL::AbstractMatrix,
                      coordsL::AbstractMatrix,
                      nrm_face::AbstractMatrix,
                      alpha::Number,
                      dxidxL::Abstract3DArray,
                      jacL::AbstractVector,
                      res1L::AbstractMatrix)  where {Tsol, Tres}

  fill!(res1L, 0)
  numDofPerNode, numNodesPerFace = size(delta_w)
  numNodesPerElement = length(jacL)
  dim = size(nrm_face, 1)

  @unpack penalty delta_w_n qL qR t1L t1R t2L t2R
  fill!(qL, 0); fill!(qR, 0)

  # multiply by normal vector, then R^T B
  #alpha_g = 1/(dim + 1)  # = 1/number of faces of a simplex
  @simd for d1=1:dim
    @simd for j=1:numNodesPerFace
      @simd for k=1:numDofPerNode
        delta_w_n[k, j] = delta_w[k, j]*nrm_face[d1, j]
      end
    end

    qL_d = sview(qL, :, :, d1); qR_d = sview(qR, :, :, d1)
    boundaryFaceIntegrate!(sbpface, bndry.face, delta_w_n, qL_d)
  end

  # apply Lambda matrix
  applyDiffusionTensor(diffusion, sbp, params, uL, wL, coordsL, dxidxL, jacL,
                       bndry.element, sbpface, bndry.face,
                       qL, t1L)

  # apply inverse mass matrix, then apply B*Nx*R*t2L_x + B*Ny*R*t2L_y
  @simd for d1=1:dim
    @simd for j=1:numNodesPerElement
      facL = (1/alpha)*jacL[j]/sbp.w[j]
      #facL = (3)*jacL[j]/sbp.w[j]
      #facR = (3)*jacR[j]/sbp.w[j]

      @simd for k=1:numDofPerNode
        t1L[k, j, d1] *= facL
      end
    end

    #TODO: make t2L 2D rather than 3D
    t1L_d = ro_sview(t1L, :, :, d1)
    t2L_d = sview(t2L, :, :, d1)
    boundaryFaceInterpolate!(sbpface, bndry.face, t1L_d, t2L_d)

    @simd for j=1:numNodesPerFace
      @simd for k=1:numDofPerNode
        val = sbpface.wface[j]*nrm_face[d1, j]*t2L_d[k, j]
        res1L[k, j] += val
      end
    end

  end  # end d1

  return nothing
end

"""
  Differentiated version of BR2 penalty
"""
function applyPenalty_diff(penalty::BR2Penalty{Tsol, Tres}, sbp,
                      params::AbstractParamType, sbpface,
                      diffusion::AbstractDiffusion, iface::Interface,
                      delta_w::AbstractMatrix{Tsol},
                      delta_w_dotL::Abstract4DArray, delta_w_dotR::Abstract4DArray,
                      theta::AbstractMatrix{Tres},
                      theta_dotL::Abstract4DArray, theta_dotR::Abstract4DArray,
                      qL_el::AbstractMatrix, qR_el::AbstractMatrix,
                      wL::AbstractMatrix, wR::AbstractMatrix,
                      coordsL::AbstractMatrix, coordsR::AbstractMatrix,
                      nrm_face::AbstractMatrix,
                      alphas::AbstractVector,
                      dxidxL::Abstract3DArray, dxidxR::Abstract3DArray,
                      jacL::AbstractVector, jacR::AbstractVector,
                      res1L_dotL::Abstract4DArray, res1L_dotR::Abstract4DArray,
                      res1R_dotL::Abstract4DArray, res1R_dotR::Abstract4DArray,
                      res2L::AbstractMatrix, res2R::AbstractMatrix,
                      res2L_dotL::Abstract4DArray, res2L_dotR::Abstract4DArray,
                      res2R_dotL::Abstract4DArray, res2R_dotR::Abstract4DArray,
                     ) where {Tsol, Tres}

  numDofPerNode, numNodesPerFace = size(delta_w)
  numNodesPerElement = length(jacL)
  dim = size(nrm_face, 1)
  add = SummationByParts.Add()

  # nodes in the stencil of R
  nodesL = sview(sbpface.perm, :, iface.faceL)
  nodesR = sview(sbpface.perm, :, iface.faceR)

  # the first L/R indicates if this is element kappa or nu
  # the second L/R indices if the derivative is wrt qL or qR
  @unpack penalty delta_w_n qL qR t1_dotL t1_dotR t2LL t2LR t2RL t2RR
  @unpack penalty t3LL t3LR t3RL t3RR t4RL t4RR
  t4LL = t1_dotL; t4LR = t1_dotR

  fill!(qL, 0); fill!(qR, 0)

  #--------------------------
  # T1

  # compute primal quantity needed later
  # multiply by normal vector, then R^T B
  #alpha_g = 1/(dim + 1)  # = 1/number of faces of a simplex
  @simd for d1=1:dim
    @simd for j=1:numNodesPerFace
      @simd for k=1:numDofPerNode
        delta_w_n[k, j] = 0.25*delta_w[k, j]*nrm_face[d1, j]
      end
    end

    qL_d = sview(qL, :, :, d1); qR_d = sview(qR, :, :, d1)
    interiorFaceIntegrate!(sbpface, iface, delta_w_n, qL_d, qR_d)

    # interiorFaceIntegrates -= the second output
    @simd for j=1:numNodesPerElement
      @simd for k=1:numDofPerNode
        qR_d[k, j] = -qR_d[k, j]
      end
    end
  end

  # compute derivative quantity
  # apply B and [Nx, Ny]
  @simd for q=1:numNodesPerElement
    @simd for p=1:numNodesPerFace
      @simd for d=1:dim
        fac = 0.25*sbpface.wface[p]*nrm_face[d, p]
        @simd for j=1:numDofPerNode
          @simd for i=1:numDofPerNode
            t1_dotL[i, j, d, p, q] = fac*delta_w_dotL[i, j, p, q]
            t1_dotR[i, j, d, p, q] = fac*delta_w_dotR[i, j, p, q]
          end
        end
      end
    end
  end

  # interpolate back to volume nodes
  fill!(t2LL, 0); fill!(t2RL, 0); fill!(t2LR, 0); fill!(t2RR, 0)
  interiorFaceIntegrate_jac!(sbpface, iface, t1_dotL, t1_dotL, t2LL, t2RL, add,
                             add, false)
  interiorFaceIntegrate_jac!(sbpface, iface, t1_dotR, t1_dotR, t2LR, t2RR, add,
                             add, false)

  # apply diffusion tensor
  applyDiffusionTensor_diff(diffusion, sbp, params, qL_el, wL, coordsL, dxidxL,
                            jacL, iface.elementL, sbpface,
                            iface.faceL, qL, t2LL, t2LR, t3LL, t3LR)
  # reverse RR and RL here because Lambda is a function of qR, not qL
  applyDiffusionTensor_diff(diffusion, sbp, params, qR_el, wR, coordsR, dxidxR,
                            jacR, iface.elementR, sbpface,
                            iface.faceR, qR, t2RR, t2RL, t3RR, t3RL)

  # apply inverse mass matrix (only nodes needed for interpolation later)
  @simd for q=1:numNodesPerElement
    # elementL
    @simd for p in nodesL
      facL = (1/alphas[1])*jacL[p]/sbp.w[p]
      @simd for d=1:dim
        @simd for j=1:numDofPerNode
          @simd for i=1:numDofPerNode
            t3LL[i, j, d, p, q] *= facL
            t3LR[i, j, d, p, q] *= facL
          end
        end
      end
    end
    # elementR
    @simd for p in nodesR
      facR = (1/alphas[2])*jacR[p]/sbp.w[p]
      @simd for d=1:dim
        @simd for j=1:numDofPerNode
          @simd for i=1:numDofPerNode
            t3RL[i, j, d, p, q] *= facR
            t3RR[i, j, d, p, q] *= facR
          end
        end
      end
    end
  end  # end q

  # interpolate back to the face
  fill!(t4LL, 0); fill!(t4RL, 0); fill!(t4LR, 0); fill!(t4RR, 0)
  interiorFaceInterpolate_jac!(sbpface, iface, t3LL, t3RL, t4LL, t4RL)
  interiorFaceInterpolate_jac!(sbpface, iface, t3LR, t3RR, t4LR, t4RR)

  # apply B * [Nx, Ny] and sum
#  fill!(res1L_dotL, 0); fill!(res1L_dotR, 0)  #TODO: tile this
#  fill!(res1R_dotL, 0); fill!(res1R_dotR, 0)
  @simd for q=1:numNodesPerElement
    @simd for p=1:numNodesPerFace
      @simd for j=1:numDofPerNode
        @simd for i=1:numDofPerNode
          res1L_dotL[i, j, p, q] = 0
          res1L_dotR[i, j, p, q] = 0
          res1R_dotL[i, j, p, q] = 0
          res1R_dotR[i, j, p, q] = 0
        end
      end

      @simd for d=1:dim
        fac = sbpface.wface[p]*nrm_face[d, p]
        @simd for j=1:numDofPerNode
          @simd for i=1:numDofPerNode
            valL = fac*(t4LL[i, j, d, p, q] + t4RL[i, j, d, p, q])
            valR = fac*(t4LR[i, j, d, p, q] + t4RR[i, j, d, p, q])
            #println("t2L_d_dot = ", real(t4LL[i, j, d, p, q]))
            #println("t2R_d_dot = ", real(t4RL[i, j, d, p, q]))
            #println("val = ", real(valL))
            res1L_dotL[i, j, p, q] +=  valL
            res1L_dotR[i, j, p, q] +=  valR
            res1R_dotL[i, j, p, q] += -valL  # delta_w is reversed for elementR
            res1R_dotR[i, j, p, q] += -valR
          end
        end
      end
    end
  end

  # T2 and T3
  @simd for q=1:numNodesPerElement
    @simd for p=1:numNodesPerFace
      @simd for j=1:numDofPerNode
        @simd for i=1:numDofPerNode
          res2L_dotL[i, j, p, q] = 0
          res2L_dotR[i, j, p, q] = 0
          res2R_dotL[i, j, p, q] = 0
          res2R_dotR[i, j, p, q] = 0
        end
      end
      fac1 =  0.5*sbpface.wface[p]
      fac2 = -0.5*sbpface.wface[p]
      @simd for j=1:numDofPerNode
        @simd for i=1:numDofPerNode
          res1L_dotL[i, j, p, q] += fac1*theta_dotL[i, j, p, q]
          res1L_dotR[i, j, p, q] += fac1*theta_dotR[i, j, p, q]
          res1R_dotL[i, j, p, q] += fac1*theta_dotL[i, j, p, q]
          res1R_dotR[i, j, p, q] += fac1*theta_dotR[i, j, p, q]

          res2L_dotL[i, j, p, q] += fac2*delta_w_dotL[i, j, p, q]
          res2L_dotR[i, j, p, q] += fac2*delta_w_dotR[i, j, p, q]
          res2R_dotL[i, j, p, q] -= fac2*delta_w_dotL[i, j, p, q]
          res2R_dotR[i, j, p, q] -= fac2*delta_w_dotR[i, j, p, q]
        end
      end
    end
  end

  # also need res2 as output
  fill!(res2L, 0); fill!(res2R, 0)
  @simd for j=1:numNodesPerFace
    @simd for k=1:numDofPerNode
      val2 = -0.5*sbpface.wface[j]*delta_w[k, j]
      res2L[k, j] =  val2
      res2R[k, j] = -val2  # delta_w is reversed for elementR
    end
  end

  return false
end

"""
  Differentiated version of BR2 penalty
"""
function applyDirichletPenalty_diff(penalty::BR2Penalty{Tsol, Tres}, sbp,
                      params::AbstractParamType,
                      sbpface, diffusion::AbstractDiffusion, bndry::Boundary,
                      delta_w::AbstractMatrix{Tsol},
                      delta_w_dot::Abstract4DArray,
                      qL_el::AbstractMatrix,
                      wL::AbstractMatrix,
                      coordsL::AbstractMatrix,
                      nrm_face::AbstractMatrix,
                      alpha::Number,
                      dxidxL::Abstract3DArray,
                      jacL::AbstractVector,
                      res_dot::Abstract4DArray,
                     ) where {Tsol, Tres}

  numDofPerNode, numNodesPerFace = size(delta_w)
  numNodesPerElement = length(jacL)
  dim = size(nrm_face, 1)
  add = SummationByParts.Add()

  # nodes in the stencil of R
  nodesL = sview(sbpface.perm, :, bndry.face)

  # the first L/R indicates if this is element kappa or nu
  # the second L/R indices if the derivative is wrt qL or qR
  @unpack penalty delta_w_n qL t1_dotL t2LL t3LL
  t4LL = t1_dotL

  fill!(qL, 0);

  # compute primal quantity needed later
  # multiply by normal vector, then R^T B
  #alpha_g = 1/(dim + 1)  # = 1/number of faces of a simplex
  @simd for d1=1:dim
    @simd for j=1:numNodesPerFace
      @simd for k=1:numDofPerNode
        delta_w_n[k, j] = delta_w[k, j]*nrm_face[d1, j]
      end
    end

    qL_d = sview(qL, :, :, d1)
    boundaryFaceIntegrate!(sbpface, bndry.face, delta_w_n, qL_d)
  end

  # compute derivative quantity
  # apply B and [Nx, Ny]
  @simd for q=1:numNodesPerElement
    @simd for p=1:numNodesPerFace
      @simd for d=1:dim
        fac = sbpface.wface[p]*nrm_face[d, p]
        @simd for j=1:numDofPerNode
          @simd for i=1:numDofPerNode
            t1_dotL[i, j, d, p, q] = fac*delta_w_dot[i, j, p, q]
          end
        end
      end
    end
  end

  # interpolate back to volume nodes
  fill!(t2LL, 0)
  boundaryFaceIntegrate_jac!(sbpface, bndry.face, t1_dotL, t2LL, include_quadrature=false)

  # apply diffusion tensor
  applyDiffusionTensor_diff(diffusion, sbp, params, qL_el, wL, coordsL, dxidxL,
                            jacL, bndry.element, sbpface,
                            bndry.face, qL, t2LL, t3LL)

  # apply inverse mass matrix (only nodes needed for interpolation later)
  @simd for q=1:numNodesPerElement
    # elementL
    @simd for p in nodesL
      facL = (1/alpha)*jacL[p]/sbp.w[p]
      @simd for d=1:dim
        @simd for j=1:numDofPerNode
          @simd for i=1:numDofPerNode
            t3LL[i, j, d, p, q] *= facL
          end
        end
      end
    end
  end  # end q

  # interpolate back to the face
  fill!(t4LL, 0)
  boundaryFaceInterpolate_jac!(sbpface, bndry.face, t3LL, t4LL)

  # apply B * [Nx, Ny] and sum
  @simd for q=1:numNodesPerElement
    @simd for p=1:numNodesPerFace
      @simd for j=1:numDofPerNode
        @simd for i=1:numDofPerNode
          res_dot[i, j, p, q] = 0
        end
      end

      @simd for d=1:dim
        fac = sbpface.wface[p]*nrm_face[d, p]
        @simd for j=1:numDofPerNode
          @simd for i=1:numDofPerNode
            res_dot[i, j, p, q] +=  fac*t4LL[i, j, d, p, q]
          end
        end
      end
    end
  end

  return nothing
end


######################################################################
# waypoint.jl
######################################################################

# TODO: figure out export strategy
export WaypointPolicy

# to simplify some notation in this file
#typealias VVF Vector{Vector{Float64}}
const VVF = Vector{Vector{Float64}}
export VVF

type WaypointPolicy <: Policy
	current_index::Int
	waypoints::VVF

	function WaypointPolicy(waypoints::VVF)
		new(2, waypoints)
	end
end

# TODO: actually require the third index
function action(m::SearchDomain, x::Vehicle, o::Float64, f::AbstractFilter, p::WaypointPolicy)
	dx = p.waypoints[p.current_index][1] - x.x
	dy = p.waypoints[p.current_index][2] - x.y
	p.current_index += 1

	# Just wrap around? Or maybe we should stay still?
	if p.current_index > length(p.waypoints)
		p.current_index = 1
	end

	return dx, dy, 0.0
end


#=  Type definitions for  site  and  site_collection  and related functions for the infinite sites model.
A site stores statistics for a single site, while a site_collection stores statistics for all of the sites of a trial.
The update_sites1() function is called on every generation to update all sites in the site collection,

GLOBAL VARIABLES:  The global variables declared in this file are used only for correctness checking via the @assert macro.
  There are equivalent fields in site_collection, and results depend only on these fields, not on global variables.
  Thus, all global variables and global variable computations can be removed without affecting functionality.
=#
export site_type, site,
    site_collection, sc_push!, update_sites!, update_neutral, update_selected,
    print_summary, average_time_to_extinction, average_time_to_fixation,
    fixed_fraction, average_fitness_fixed, average_fitness_extinct, average_fitness_all,
    sites_per_gen, heterozygosity_per_gen

#=
Stores the properties of an "site" (possibly deleterious or advantageous).
=#

mutable struct site_type
  identifier::Int64   # integer for this site, not sure if this is needed
  start_gen::Int64    # the generation (time step) when the site was generated
  final_gen::Int64    # the generation when the site went extinct.  Should be zero while site is evolving
  fitness_coefficient::Float64    # The selection coefficient of the site (allele)
  sum_gens::Int64     # sum of the number of generations where counts and heterozygosity are accumulated
  sum_counts::Int64   # the current accumulated sum of the counts of the mutant allele per generation for infinite sites
  sum_sq_counts::Int64   # the current accumulated sum of squares of the counts per generation for infinite sites
  sum_heteroz::Float64  # the current accumulated sum of heterozygosities for infinite sites
  sum_sq_heteroz::Float64  # the current accumulated sum of squares ofheterozygosities for infinite sites
  previous_allele_freq::Int64  # Allele frequence of the previous generation
end

@doc """ function site()
  Initializes the properties of the site.
  Initial values for sum_counts, sum_sq_counts, sum_heteroz, sum_sq_heteroz are added.
  This means that the number of generations is one larger than the number of increments in update_innovations!(),
    and 1 is added to sum_gens in sites_per_gen() and heterozygosity_per_gen().
"""
function site( id::Int64, N::Int64, start_gen::Int64, fitness_coef::Float64=1.0 )
  # See the comments on global variables at the beginning of this file.
  global sum_counts += 1
  global sum_sq_counts += 1
  initial_heteroz = 1.0 - watterson_homozygosity([N-1,1])
  global sum_heteroz += initial_heteroz
  global sum_sq_heteroz += initial_heteroz^2
  global sum_gens += 1
  return site_type( id, start_gen, 0, fitness_coef, 1, 1, 1, initial_heteroz, initial_heteroz^2, 1  )
end

#=
Stores a collection of sites.
Sites are partitioned into 3 subsets: active, fixed, and extinct.
=#
mutable struct site_collection
  N::Int64                # popsize
  list::Dict{Int64,site_type} #    Stores a collection of sites
  active::BitSet   # indices of actively evolving sites
  fixed::BitSet   # indices of sites that have fixed
  extinct::BitSet   # indices of sites that have gone extinct
  fix_minimum::Float64  # minimum fraction of popsize for fixation for infinite alleles
  sum_gens::Int64     # sum of counts of generations
  sum_counts::Int64   # sum of counts of mutant alleles for infinite sites
  sum_sq_counts::Float64   # sum of counts of mutant alleles for infinite sites
  sum_heteroz::Float64  # sum of heterozygosities for infinite sites
  sum_sq_heteroz::Float64  # sum of heterozygosities squared for infinite sites
  sum_generations::Int64  # Total number of generations corresponding to sum_counts and sum_heteroz
  count_fixed_del::Int64  # Number of fixed mutations that are deleterious (fit_coef < 1.0)
  count_fixed_adv::Int64  # Number of fixed mutations that are advantageous (fit_coef > 1.0)
  in_use::Bool      # If false, not used
end

# Constructor for a new empty site collection with a value for fix_minimum
function site_collection( N::Int64, fix_min::Float64, in_use::Bool=true )
  global sum_counts = 0
  global sum_sq_counts = 0
  global sum_heteroz = 0.0
  global sum_sq_heteroz = 0.0
  global sum_gens = 0
  site_collection( N, Dict{Int64,site_type}(), BitSet(), BitSet(), BitSet(), fix_min, 0, 0, 0.0, 0.0, 0.0, 0,0,0, in_use )
end

@doc """ sc_push!()
 Adds an site to the collection.
"""
function sc_push!( innov_collection::site_collection, innov::site_type )
  if !innov_collection.in_use
    return
  end
  innov_collection.list[innov.identifier] = innov
  if innov.final_gen == 0
    Base.push!( innov_collection.active, innov.identifier )
  elseif innov.history[end] > 0
    Base.push!( innov_collection.fixed, innov.identifier )
  else
    Base.push!( innov_collection.extinct, innov.identifier )
  end
  innov_collection
end

@doc """ fix_test()
The most basic test for fixation, namely that the allele frequency is the population size.
This is used in the infitite sites model.
"""
function fix_test( N::Int64, new_allele_freq::Int64 )
  return N == new_allele_freq
end

@doc """ function update_sites!( )
 Update all active sites, and make some extinct and fixed
 This version is used only by infsites.jl.
"""
function update_sites!( sc::site_collection, g::Int64, N::Int64 )
  if !sc.in_use
    return
  end
  for index in sc.active  # updates sites to the next generation
    # update_selected does the equivalent of proportional selection
    new_allele_freq = update_selected( index, N, sc.list[index].previous_allele_freq, sc.list[index].fitness_coefficient )
    innov = sc.list[index]
    if new_allele_freq > 0 && new_allele_freq < N
      global sum_counts += new_allele_freq
      global sum_sq_counts += new_allele_freq^2
      innov.sum_gens += 1
      innov.sum_counts += new_allele_freq
      innov.sum_sq_counts += new_allele_freq^2
      heteroz =  1.0 - watterson_homozygosity([N-new_allele_freq, new_allele_freq])
      global sum_heteroz += heteroz
      global sum_sq_heteroz += heteroz^2
      global sum_gens +=1
      innov.sum_heteroz += heteroz
      innov.sum_sq_heteroz += heteroz^2
    end
    innov.previous_allele_freq = new_allele_freq  # save for the next call to update_selected.
    if new_allele_freq == 0 ||  new_allele_freq >= N
      sc.sum_gens += innov.sum_gens
      sc.sum_counts += innov.sum_counts
      sc.sum_sq_counts += innov.sum_sq_counts
      sc.sum_heteroz += innov.sum_heteroz
      sc.sum_sq_heteroz += innov.sum_sq_heteroz
      sc.sum_generations += innov.sum_gens
      Base.pop!( sc.active,index)
      innov.final_gen = g
      if new_allele_freq == 0  # extinction
        Base.push!( sc.extinct,index)
      elseif new_allele_freq >= N # fixation
        Base.push!( sc.fixed,index)
      end
    end
  end
end

@doc """ function update_neutral()
  Updates site according to the Wright-Fisher model of drift.
  Since mutation has already happened by the time a new site is established, there is no mutation.
"""
function update_neutral( site::Int64, N::Int64, old_allele_freq::Int64 )
  p = Float64(old_allele_freq)/N
  new_allele_freq = rand(Binomial(N,p))
  return new_allele_freq
end

@doc """ function update_selected()
  Updates site according to the Wright-Fisher model of drift with selection.
  Since mutation has already happened by the time a new site is established, there is no mutation.
  This function is called by update_sites!() in site_collection.jl.

"""
function update_selected( site::Int64, N::Int64, old_allele_freq::Int64, select_coef::Float64 )
  p = min(1.0,(select_coef*Float64(old_allele_freq)/N))
  new_allele_freq = rand(Binomial(N,p))
  return new_allele_freq
end

function N_inf_sites( sc::site_collection )
  if !sc.in_use
    return
  end
  sum_N = 0
  for index in sc.active
    sum_N += sc.list[index].history[end]
  end
  sum_N
end

@doc """ check_globals_against_site_collection()
  Check that the values of global variables approximately agree with the corresponding variables in sc.
"""
function check_globals_against_site_collection( sc::site_collection )
  global sum_counts
  global sum_sq_counts
  global sum_heteroz
  global sum_sq_heteroz
  global sum_gens
  @assert isapprox(sc.sum_gens,sum_gens)
  @assert isapprox(sc.sum_counts,sum_counts)
  @assert isapprox(sc.sum_sq_counts,sum_sq_counts)
  @assert isapprox(sc.sum_heteroz,sum_heteroz)
  @assert isapprox(sc.sum_sq_heteroz,sum_sq_heteroz)
end

function print_summary( sc::site_collection; print_lists::Bool=false )
  println("print summary of site collection: ")
  global sum_counts
  global sum_sq_counts
  global sum_heteroz
  global sum_sq_heteroz
  global sum_gens

  #global sum_counts
  #println("psum sum_counts: ",sum_counts)
  if !sc.in_use
    return
  end
  if print_lists
    println("active list: ",sc.active)
    println("fixed list: ",sc.fixed)
    println("extinct list: ",sc.extinct)
  end
  println("number active: ",length(sc.active))
  println("number fixed: ",length(sc.fixed))
  println("number extinct: ",length(sc.extinct))
  println("fixed fraction: ",fixed_fraction(sc))
  println("sum_counts: ",sc.sum_counts,"  glb: ",sum_counts)
  println("sum_sq_counts: ",sc.sum_sq_counts,"  glb: ",sum_sq_counts)
  println("sum_heteroz: ",sc.sum_heteroz,"  glb: ",sum_heteroz)
  println("sum_sq_heteroz: ",sc.sum_sq_heteroz,"  glb: ",sum_sq_heteroz)
  println("sum_gens: ",sc.sum_generations,"  glb: ",sum_gens)
  (avg,stderr) = sites_per_gen(sc)
  println("avg sites per generation: ",avg)
  println("stderr sites per generation: ",stderr)
  (avg,stderr) = heterozygosity_per_gen(sc)
  println("avg per generation heterozygosity: ",avg)
  println("stderr per generation heterozygosity: ",stderr)
  println("avg time to fixation: ",average_time_to_fixation(sc))
  println("avg time to extinction: ",average_time_to_extinction(sc))
  println("avg fitness fixed: ",average_fitness_fixed(sc))
  println("avg fitness extinct: ",average_fitness_extinct(sc))
  println("avg fitness all: ",average_fitness_all(sc))
end

function average_time_to_extinction( innov_collection::site_collection )
  if !innov_collection.in_use
    return 0.0
  end
  if length(innov_collection.extinct) == 0
    println("no extinct sites")
    return -1.0
  end
  #print("average_time_to_extinction:  ")
  sum = 0
  for i in innov_collection.extinct
    sum += innov_collection.list[i].final_gen - innov_collection.list[i].start_gen
    #println("id: ",innov_collection.list[i].identifier,"  time: ",(innov_collection.list[i].final_gen-innov_collection.list[i].start_gen),"  sum: ",sum)
  end
  return Float64(sum)/length(innov_collection.extinct)
end

function average_time_to_fixation( innov_collection::site_collection )
  if !innov_collection.in_use
    return 0.0
  end
  if length(innov_collection.fixed) == 0
    println("no fixed sites")
    return -1.0
  end
  sum = 0
  for i in innov_collection.fixed
    sum += innov_collection.list[i].final_gen - innov_collection.list[i].start_gen
  end
  return Float64(sum)/length(innov_collection.fixed)
end

# Fraction of fixed sites out of fixed plus extinct sites
function fixed_fraction( innov_collection::site_collection )
  if !innov_collection.in_use
    return 0.0
  end
  if  length(innov_collection.extinct) == 0 && length(innov_collection.fixed) == 0
    println("no extinct and no fixed sites")
    return -1.0
  end
  return length(innov_collection.fixed)/Float64(length(innov_collection.fixed)+length(innov_collection.extinct))
end

function average_fitness_fixed( sc::site_collection )
  if !sc.in_use
    return 0.0
  end
  if  length(sc.fixed) == 0
    return -1.0
  end
  sum = 0.0
  for i in sc.fixed
    sum += sc.list[i].fitness_coefficient
  end
  return sum/length(sc.fixed)
end

function average_fitness_extinct( sc::site_collection )
  if !sc.in_use
    return 0.0
  end
  if  length(sc.extinct) == 0
    return -1.0
  end
  sum = 0.0
  for i in sc.extinct
    sum += sc.list[i].fitness_coefficient
  end
  return sum/length(sc.extinct)
end

function average_fitness_all( sc::site_collection )
  if !sc.in_use
    return 0.0
  end
  if  length(sc.extinct) == 0
    return -1.0
  end
  sum = 0.0
  for i in sc.extinct
    sum += sc.list[i].fitness_coefficient
  end
  for i in sc.fixed
    sum += sc.list[i].fitness_coefficient
  end
  return sum/(length(sc.extinct) + length(sc.fixed))
end

@doc """ function count_adv_del_fixed( )
  Counts the number of fixed sites that are advantageous, and the number that are deleterious at the end of the run.
  Stores these in site collection sc, and also returns them as a pair.
"""
function count_adv_del_fixed( sc::site_collection )
  #println("count_adv_del_fixed")
  if !sc.in_use
    return 0.0
  end
  if  length(sc.fixed) == 0
    return 0, 0
  end
  count_adv = 0
  count_del = 0
  for i in sc.fixed
    fit_coef = sc.list[i].fitness_coefficient
    #println("i: ",i,"  fit_coef: ",fit_coef)
    if fit_coef > 1.0
      count_adv += 1
    end
    if fit_coef < 1.0
      count_del += 1
    end
  end
  sc.count_fixed_adv = count_adv
  sc.count_fixed_del = count_del
  return count_adv, count_del
end

function sites_per_gen( sc::site_collection )
  global sum_gens
  global sum_counts
  global sum_sq_counts
  global counts_var = sqrt(sum_sq_counts/(sum_gens-1) - sum_counts^2/sum_gens/(sum_gens-1))/sum_gens
  sumgens = sc.sum_generations
  average = sc.sum_counts/sumgens
  coef_var = sqrt(sc.sum_sq_counts/(sumgens-1) - sc.sum_counts^2/sumgens/(sumgens-1))/sumgens
  @assert isapprox( coef_var, counts_var )
  if coef_var >= 0
    stderr = coef_var
  else
    stderr = -1.0
  end
  return (average, stderr)
end

function heterozygosity_per_gen( sc::site_collection )
  global sum_gens
  global sum_heteroz
  global sum_sq_heteroz
  heteroz_var = sqrt(sum_sq_heteroz/(sum_gens-1) - sum_heteroz^2/sum_gens/(sum_gens-1))/sum_gens
  sumgens = sc.sum_generations
  average = sc.sum_heteroz/sumgens
  coef_var = sqrt(sc.sum_sq_heteroz/(sumgens-1) - sc.sum_heteroz^2/sumgens/(sumgens-1))/sumgens
  @assert isapprox( coef_var, heteroz_var )
  if coef_var >= 0
    stderr = coef_var
  else
    stderr = -1.0
  end
  return (average, stderr)
end


"""
    write_smac2_par([...])

Writes a P-Smac2 parameter file. Not all relevant parameters implemented yet!
"""
function write_smac2_par(x, y, z,
                         euler_angle_0, euler_angle_1, euler_angle_2,
                         xy_size, z_depth, xy_pix::Integer,
                         input_file, output_file, path,
                         effect_module::Integer=0, effect_flag::Integer=0,
                         ν_obs::Real=1.44e9, cosmology::Integer=0,
                         CR_pmin::Real=10.0, CR_pmax::Real=1.e7)

    open(path * "smac2.par", "w") do f
        write(f,
        """%% PSmac Parameter File %%

        Cosmology $cosmology         % set to 0 for non-cosmological sims.  Cosmo.h = 1

        %% Image Properties, comoving %%

        Center_X $x
        Center_Y $y
        Center_Z $z

        Use_Barycenter 0    % 1 - set to, 2 - set relative to barycenter

        XYSize $xy_size        % set these two ==0 to project whole box
        ZDepth $z_depth
        XYPix $xy_pix

        %% Projection, see description at EOF %%
        Euler_Angle_0 $euler_angle_0    % [deg] along x : (0,90,90)
        Euler_Angle_1 $euler_angle_1    % [deg] along y : (0,90,0)
        Euler_Angle_2 $euler_angle_2     % [deg] along z : (0,0,0)

        %% HEALPIX %%

        NSide 32
        Rmin 1000
        Rmax 1000000

        %% I/O Options %%

        Input_File $input_file
        N_IOTasks 1

        NoClobber 0
        Output_File $output_file

        %% Effect Selection & Options %%

        Effect_Module $effect_module
        Effect_Flag $effect_flag

        E_min 5e4             % Energy Range [eV]: Xray, Gamma, Synchro
        E_max 1e10

        Freq $ν_obs          % Frequency [Hz]: Sz,Synchro

        a_cr 2.375          % CR spectral index
        X_cr 0.01           % CR Normalisation rel. thermal
        IntrRM 0            % Toggle Intrinsic RM
        PitchAngInt 1       % Toggle Pitch angle integration

        t_turb 1e7          % Reacceleration Timescale Cassano 05
        eta_t 0.26          % Fraction turb. Energy in magnetosonic waves
        turb_scale 100      % scale to extrapolate turbulent velocity to

        CR_Emin $CR_pmin       % Minimum energy for BP_REAL_CRs, as in OpenGadget3 parameterfile
        CR_Emax $CR_pmax       % Maximum energy for BP_REAL_CRs, as in OpenGadget3 parameterfile

        %Gadget
        UnitLength_in_cm 			3.085678e21        %  1.0 kpc
        UnitMass_in_g 				1.989e43           %  1.0e10 solar masses
        UnitVelocity_in_cm_per_s 	1e5                %  1 km/sec
        """
        )
    end
end

"""
    read_smac2_info(filename::String)

Returns the image info in a `mappingParameters` struct.
"""
function read_smac2_info(filename::String)

    f = FITS(filename)

    center = [read_key(f[1], 23)[2],
              read_key(f[1], 24)[2],
              read_key(f[1], 25)[2] ]

    xy_size = read_key(f[1], 27)[2]
    z_size  = read_key(f[1], 28)[2]
    Npixels = read_key(f[1], 29)[2]

    close(f)

    return mappingParameters(center=center,
									x_size=xy_size,
									y_size=xy_size,
									z_size=z_size,
									Npixels=Npixels)

end

"""
    read_smac2_image(filename::String)

Returns the image of a Smac2 FITS file.
"""
function read_smac2_image(filename::String, num_image::Integer=1)

    f = FITS(filename)

    image = read(f[1])[:,:,num_image]

    close(f)

    return image

end


# ================================
# Hybrid flowpipe
# ================================

"""
    HybridFlowpipe{N, D, FT<:AbstractFlowpipe, VOA<:VectorOfArray{N, D, Vector{FT}}} <: AbstractFlowpipe

Type that wraps a vector of flowpipes of the same type, such that they are
contiguous in time.

### Fields

- `Fk`  -- vector of flowpipes
- `ext` -- (optional, default: empty) dictionary for extensions

### Notes

The evaluation functions (in time) for this type do not assume that the flowpipes are contiguous in time.
That is, the final time of the `i`-th flowpipe does not match the start time of the `i+1`-th flowpipe.
"""
struct HybridFlowpipe{N, RT<:AbstractReachSet{N}, FT<:AbstractFlowpipe} <: AbstractFlowpipe
    Fk::VectorOfArray{RT, 2, Vector{FT}}
    ext::Dict{Symbol, Any}
end

function HybridFlowpipe(Fk::Vector{FT}) where {N, RT<:AbstractReachSet{N}, FT<:Flowpipe{N, RT}}
    voa = VectorOfArray{RT, 2, Vector{FT}}(Fk)
    ext = Dict{Symbol, Any}()
    return HybridFlowpipe{N, RT, FT}(voa, ext)
end

function HybridFlowpipe(Fk::Vector{FT}, ext::Dict{Symbol, Any}) where {N, RT<:AbstractReachSet{N}, FT<:Flowpipe{N, RT}}
    voa = VectorOfArray{RT, 2, Vector{FT}}(Fk)
    return HybridFlowpipe{N, RT, FT}(voa, ext)
end

function HybridFlowpipe(Fk::Vector{SFT}) where {N, RT, FT<:Flowpipe{N, RT}, NT<:Number, SFT<:ShiftedFlowpipe{FT, NT}}
    voa = VectorOfArray{RT, 2, Vector{SFT}}(Fk)
    ext = Dict{Symbol, Any}()
    return HybridFlowpipe{N, RT, SFT}(voa, ext)
end

# interface functions
array(fp::HybridFlowpipe) = fp.Fk
flowpipe(fp::HybridFlowpipe) = fp
numtype(::HybridFlowpipe{N}) where {N} = N
setrep(::Type{HybridFlowpipe{N, RT, FT}}) where {N, RT, FT} = setrep(RT)
setrep(::HybridFlowpipe{N, RT, FT}) where {N, RT, FT} = setrep(RT)
rsetrep(::Type{HybridFlowpipe{N, RT, FT}}) where {N, RT, FT} = RT
numrsets(fp::HybridFlowpipe) = mapreduce(length, +, fp)

# indexing: fp[j, i] returning the j-th reach-set of the i-th flowpipe
Base.getindex(fp::HybridFlowpipe, I::Int...) = getindex(fp.Fk, I...)

function tspan(fp::HybridFlowpipe)
    ti = minimum(tstart, fp)
    tf = maximum(tend, fp)
    return TimeInterval(ti, tf)
end

function Base.similar(fp::HybridFlowpipe{N, RT, FT}) where {N, RT, FT}
    return HybridFlowpipe(Vector{FT}())
end

function overapproximate(fp::HybridFlowpipe, args...)
    return HybridFlowpipe([overapproximate(F, args...) for F in fp], fp.ext)
end

function project(fp::HybridFlowpipe, args...)
    return [project(F, args...) for F in fp]
end

# LazySets interface
function LazySets.ρ(d::AbstractVector, fp::HybridFlowpipe)
    return maximum(ρ(d, F) for F in array(fp))
end

#function LazySets.σ(d::AbstractVector, fp::HybridFlowpipe)
#    error("not implemented")
#end

Base.:⊆(F::HybridFlowpipe, X::LazySet) = all(fp ⊆ X for fp in F)
Base.:⊆(F::HybridFlowpipe, Y::AbstractLazyReachSet) = all(fp ⊆ set(Y) for fp in F)

# evaluation for scalars
function (fp::HybridFlowpipe{N, RT})(t::Number) where {N, RT<:AbstractReachSet{N}}
    if t ∉ tspan(fp)
        throw(ArgumentError("time $t does not belong to the time span, " *
                            "$(tspan(fp)), of the given flowpipe"))
    end
    vec = Vector{RT}()
    for (k, fk) in enumerate(array(fp)) # loop over flowpipes
        if t ∈ tspan(fk)
            for Ri in fk  # loop over reach-sets for this flowpipe
                if t ∈ tspan(Ri)
                    push!(vec, Ri)
                end
            end
        end
    end
    return vec
end

# evaluation for time intervals
function (fp::HybridFlowpipe{N, RT})(dt::TimeInterval) where {N, RT<:AbstractReachSet{N}}
    if !(dt ⊆ tspan(fp)) # ⊊ doesn't work see, IntervalArithmetic#409
        throw(ArgumentError("time interval $dt does not belong to the time span, " *
                            "$(tspan(fp)), of the given flowpipe"))
    end
    vec = Vector{RT}()
    for (k, fk) in enumerate(array(fp)) # loop over flowpipes
        dtint = dt ∩ tspan(fk)
        if !isempty(dtint)
            for R in fk(dtint)  # loop over reach-sets for this flowpipe
                push!(vec, R)
            end
        end
    end
    return vec
end


import Base: promote_rule

# SparseMatrixCSC
promote_rule(::Type{SparseMatrixCSC{Tv,Ti}}, ::Type{Matrix{T}}) where {Tv,Ti,T} =
    Matrix{promote_type(T, Tv)}

# IMatrix
promote_rule(
    ::Type{IMatrix{N,T}},
    ::Type{PermMatrix{Tv,Ti,Vv,Vi}},
) where {N,T,Tv,Ti,Vv,Vi} = (TT = promote_type(T, Tv); PermMatrix{TT,Ti,Vector{TT},Vi})
promote_rule(::Type{IMatrix{N,T}}, ::Type{SparseMatrixCSC{Tv,Ti}}) where {N,T,Tv,Ti} =
    SparseMatrixCSC{promote_type(T, Tv),Ti}
promote_rule(::Type{IMatrix{M,TA}}, ::Type{Matrix{TB}}) where {M,TA,TB} = Array{TB,2}

# PermMatrix
promote_rule(
    ::Type{PermMatrix{TvA,TiA}},
    ::Type{SparseMatrixCSC{TvB,TiB}},
) where {TvA,TiA,TvB,TiB} = SparseMatrixCSC{promote_type(TvA, TvB),promote_type(TiA, TiB)}
promote_rule(::Type{PermMatrix{Tv,Ti}}, ::Type{Matrix{T}}) where {Tv,Ti,T} =
    Array{promote_type(Tv, T),2}


matmul_sizes(s::Integer) = (s,s,s)
matmul_sizes(m_k_n::Tuple{Vararg{<:Integer, 3}}) = m_k_n

"""
    runbench(backend, TA, TB, TC; kwargs...)
"""
function runbench(backend::Symbol, ::Type{TA}, ::Type{TB}, ::Type{TC};
                  kwargs...) where {TA, TB, TC}
    return runbench(Val(backend), TA, TB, TC; kwargs...)
end

"""
    runbench(::Val{:CUDA}, TA, TB, TC; kwargs...)
"""
function runbench(::Val{:CUDA}, ::Type{TA}, ::Type{TB}, ::Type{TC};
                  estimator::F = Base.minimum,
                  libs = all_libs(),
                  sizes = [2^n for n in 7:14]) where {TA, TB, TC, F}
    sizes_vec = collect(sizes)
    _sizes_column = Vector{Int}(undef, 0)
    _libraries_column = Vector{Symbol}(undef, 0)
    _tflops_column = Vector{Float64}(undef, 0)
    _times_column = Vector{Float64}(undef, 0)
    p = ProgressMeter.Progress(length(sizes_vec))
    warmed_up = falses(length(sizes_vec))
    for (j, sz) ∈ enumerate(sizes_vec)
        showvalues = []
        push!(showvalues, (:Size, sz))
        GC.gc()
        CUDA.reclaim()
        M = sz
        K = sz
        N = sz
        for lib in libs
            benchfunc! = getbenchfunc(lib)

            if !(warmed_up[j])
                A_h = rand(TA, (M, K)) / sqrt(TA(K));
                B_h = rand(TB, (K, N)) / sqrt(TB(K));
                C_h = rand(TC, (M, N));
                A = CUDA.CuArray(A_h);
                B = CUDA.CuArray(B_h);
                C = CUDA.CuArray(C_h);
                benchfunc!(C, A, B) # warmup
                warmed_up[j] = true
            end

            A_h = rand(TA, (M, K)) / sqrt(TA(K));
            B_h = rand(TB, (K, N)) / sqrt(TB(K));
            C_h = rand(TC, (M, N));
            A = CUDA.CuArray(A_h);
            B = CUDA.CuArray(B_h);
            C = CUDA.CuArray(C_h);
            trial = benchfunc!(C, A, B)
            point_estimate = estimator(trial)
            t_ns = Base.time(point_estimate) # time in nanoseconds
            t_ps = 1_000 * t_ns # time in picoseconds
            flops_factor = 2
            tflops = flops_factor * M * K * N / t_ps # teraflops
            push!(showvalues, (lib, "$(round(tflops; digits = 2)) TFLOPS"))
            push!(_sizes_column, sz)
            push!(_libraries_column, lib)
            push!(_tflops_column, tflops)
            push!(_times_column, t_ns)
        end
        ProgressMeter.next!(p; showvalues = showvalues)
        GC.gc()
        CUDA.reclaim()
    end
    df = DataFrames.DataFrame(
        Size = _sizes_column,
        Library = _libraries_column,
        TFLOPS = _tflops_column,
        Time = _times_column,
    )
    benchmark_result = BenchmarkResult{TA, TB, TC}(df)
    return benchmark_result
end


__precompile__(false)
module Dolark

    import YAML
    import Dolang
    import Dolo
    import OrderedCollections: OrderedDict
    using StaticArrays

    # extracts annotated yaml structure from model file
    # the model file is supposed to contain two documents
    # - a valid dolo model (a.k.a. agent)
    # - a description of aggregate model
    function get_hmodel_source(txt::AbstractString)
        cons = YAML.Constructor()
        YAML.add_multi_constructor!((c,s,m)->m, cons, "tag:yaml.org")
        YAML.add_multi_constructor!((c,s,m)->m, cons, "!")
        data = YAML.load_all(txt, cons)
        agent, model = data
        return (;agent, model)
    end

    struct HModel{ID, AgentType}
        filename
        source
        symbols
        calibration
        exogenous
        equations
        factories
        agent::AgentType # dolo model
    end

    import Base: show
    Base.show(io::IO, model::HModel) = print(io, "HModel")

    function get_symbols(yaml_node)
        symbols_src = yaml_node[:symbols]
        exogenous = [Symbol(e.value) for e in symbols_src[:exogenous]]
        aggregate = [Symbol(e.value) for e in symbols_src[:aggregate]]
        parameters = [Symbol(e.value) for e in symbols_src[:parameters]]
        return (;exogenous, aggregate, parameters)
    end

    function get_h_equations(yaml_node)
        projection = yaml_node[:projection].value
        equilibrium = yaml_node[:equilibrium].value
        return (;projection, equilibrium)
    end

    ### the next two functions are almost duplicates

    function gen_projection_ff(hmodel)
        projs_block = Dolang.parse_assignment_block(hmodel.equations.projection)

        eqs = OrderedDict{Symbol, Any}()
        for ch in projs_block.children
            lhs, rhs = ch.children
            name = Dolang.convert(Expr, lhs)
            eqs[name] = Dolang.convert(Expr, rhs)
        end

        arguments = OrderedDict{Symbol, Vector{Symbol}}()
        arguments[:y] = hmodel.symbols.aggregate
        arguments[:z] = hmodel.symbols.exogenous
        arguments[:p] = hmodel.symbols.parameters

        preamble = OrderedDict()

        return Dolang.FunctionFactory(eqs, arguments, preamble, :projection)
    end

    function gen_projection_ff(agent, symbols, equations)
        projs_block = Dolang.parse_assignment_block(equations.projection)

        eqs = OrderedDict{Symbol, Any}()
        for ch in projs_block.children
            lhs, rhs = ch.children
            name = Dolang.convert(Expr, lhs)
            eqs[name] = Dolang.convert(Expr, rhs)
        end

        arguments = OrderedDict{Symbol, Vector{Symbol}}()
        arguments[:y] = [Dolang.stringify(e,0) for e in symbols.aggregate]
        arguments[:z] = [Dolang.stringify(e,0) for e in symbols.exogenous]
        arguments[:p] = [Dolang.stringify(e) for e in symbols.parameters]

        preamble = OrderedDict()

        return Dolang.FunctionFactory(eqs, arguments, preamble, :projection)
    end

    import Dolang: Tree

    function gen_equilibrium_ff(agent, symbols, equations)

        eq_block = Dolang.parse_equation_block(equations.equilibrium)

        eqs = OrderedDict{Symbol, Any}()
        for (i,ch) in enumerate(eq_block.children)
            eq = Tree("sub", [ch.children[2], ch.children[1]])
            name = Symbol("out_", i, "_")
            eqs[name] = Dolang.convert(Expr, eq)
        end

        arguments = OrderedDict{Symbol, Vector{Symbol}}()
        arguments[:s] = [Dolang.stringify(e,0) for e in agent.symbols[:states]]
        arguments[:x] = [Dolang.stringify(e,0) for e in agent.symbols[:controls]]
        arguments[:y] = [Dolang.stringify(e,0) for e in symbols.aggregate]
        arguments[:z] = [Dolang.stringify(e,0) for e in symbols.exogenous]
        arguments[:Y] = [Dolang.stringify(e,1) for e in symbols.aggregate]
        arguments[:Z] = [Dolang.stringify(e,1) for e in symbols.exogenous]
        arguments[:p] = [Dolang.stringify(e) for e in symbols.parameters]
        preamble = OrderedDict()

        return Dolang.FunctionFactory(eqs, arguments, preamble, :equilibrium)

    end

    module ModelSpace

    end

    # import Dolark.ModelSpace

    function HModel(filename::AbstractString)

        txt = open(f->read(f, String), filename)

        source = get_hmodel_source(txt)
        symbols = get_symbols(source.model)
        calibration = Dolo.get_calibration(source.model)
        exogenous = Dolo.get_exogenous(source.model, symbols.exogenous, calibration.flat)
        equations = get_h_equations(source.model)

        agent = Dolo.Model(source.agent)
        AgentType = typeof(agent)
        ID = gensym()

        projection = gen_projection_ff(agent, symbols, equations)
        equilibrium = gen_equilibrium_ff(agent, symbols, equations)

        factories = (;projection, equilibrium)

        hmodel = HModel{ID, AgentType}(filename, source, symbols, calibration, exogenous, equations, factories, agent)

        tm = typeof(hmodel)

        funp = Dolang.gen_generated_gufun(factories.projection; funname=:projection, dispatch=tm)
        fune = Dolang.gen_generated_gufun(factories.equilibrium; funname=:equilibrium, dispatch=tm)

        Core.eval(Dolark, funp)
        Core.eval(Dolark, fune)

        return hmodel

    end

    struct DModel
        hmodel
        F
        G
    end

    function discretize(hmodel::HModel, sol_agent)

        F = Dolo.Euler(hmodel.agent)
        G = Dolo.distG(hmodel.agent, sol_agent)

        DModel(hmodel, F, G)

    end

    struct Unknown
        μ
        p
        x::Dolo.MSM
        y
    end

    using Dolang: _get_oorders

    include("temp.jl")

end # module


struct LatentDiffEqModel{M,E,D}

    model_type::M
    encoder::E
    decoder::D

    @doc raw"""
        LatentDiffEqModel(model_type, encoder_layers, decoder_layers)
    Construct a LatentDiffEqModel.
    # Arguments
    `model_type`: GOKU() or LatentODE()\
    `encoder_layers`: contains `feature_extractor`, `pattern_extractor` and `latent_in` layers.\
    `decoder_layers`: contains `latent_out`, `diffeq` and `reconstructor` layers.\
    To see an example of how to build these layers, please refer to the `default_layers` function.
    """
    function LatentDiffEqModel(model_type, encoder_layers, decoder_layers)

        encoder = Encoder(model_type, encoder_layers)
        decoder = Decoder(model_type, decoder_layers)
        M, E, D = typeof(model_type), typeof(encoder), typeof(decoder)
        new{M, E, D}(model_type, encoder, decoder)
    end
end

function (model::LatentDiffEqModel)(x, t, variational=false)

    # Get encoded latent initial states and parameters
    μ, logσ² = model.encoder(x)

    # Sample from distributions
    l̃ = variational ? sample(μ, logσ², model) : μ

    # Get predicted output
    X̂ = model.decoder(l̃, t)

    return X̂, μ, logσ²
end

Flux.@functor LatentDiffEqModel

struct Encoder{M,FE,PE,LI}

    model_type::M

    feature_extractor::FE
    pattern_extractor::PE
    latent_in::LI

    @doc raw"""
        Encoder(model_type, encoder_layers)
    Construct an encoder of a LatentDiffEqModel.
    # Arguments
    `model_type`: GOKU() or LatentODE()\
    `encoder_layers`: contains `feature_extractor`, `pattern_extractor` and `latent_in` layers.  Please refer to `default_layers` for an example of how to construct this.
    """
    function Encoder(model_type, encoder_layers)
        M = typeof(model_type)
        FE, PE, LI = typeof.(encoder_layers)
        new{M,FE,PE,LI}(model_type, encoder_layers...)
    end
end

function (encoder::Encoder)(x)

    # Pass every time frame independently through the feature extractor
    fe_out = apply_feature_extractor(encoder, x)

    # Process sequentially with the pattern extractor
    pe_out = apply_pattern_extractor(encoder, fe_out)

    # Pass trough a last layer before sampling
    μ, logσ² = apply_latent_in(encoder, pe_out)

    return μ, logσ²
end

Flux.@functor Encoder

struct Decoder{M,LI,D,R}

    model_type::M

    latent_out::LI
    diffeq::D
    reconstructor::R

    @doc raw"""
        Decoder(model_type, decoder_layers)
    Construct a decoder of a LatentDiffEqModel.
    # Arguments
    `model_type`: GOKU() or LatentODE()\
    `decoder_layers`: contains `latent_out`, `diffeq` and `reconstructor` layers.  Please refer to `default_layers` for an example of how to construct this.
    """
    function Decoder(model_type, decoder_layers)
        M = typeof(model_type)
        LI, D, R = typeof.(decoder_layers)
        new{M,LI,D,R}(model_type, decoder_layers...)
    end
end

function (decoder::Decoder)(l̃, t)

    # Pass sampled latent states through a latent_out layer
    l̂ = apply_latent_out(decoder, l̃)

    # Integrate differential equations
    ẑ = diffeq_layer(decoder, l̂, t)

    # Apply reconstructor independently to each time frame
    x̂ = apply_reconstructor(decoder, ẑ)

    return x̂, ẑ, l̂
end

Flux.@functor Decoder


include("../intcode.jl")

tape = set_memory!(Tape(parse_data(readline(open("input.txt")))), 100000)

@async run_program!(tape)
output = []
while true
    ch = take!(tape.outch)
    if ch == CodeFinished()
        break
    end
    push!(output, ch)
end

instr = reshape(output[1:end], 3, div(length(output[1:end]), 3))

println(sum(instr[3, :].==2))

screen = transpose(reshape(instr[3, :], 41, 24))
for i=1:24
    for j=1:41
        val = screen[i,j]
        if val == 1 ch = '|' elseif val == 2 ch = '#' elseif val == 3 ch = '=' elseif val == 4 ch= 'O' else ch = ' ' end
        print(ch)
    end
    println()
end


using Statistics

if Juno.isactive()
    colors = Juno.syntaxcolors()
    colors = Dict(k => parse(Colorant, "#"*Base.hex(v, 6, false)) for (k, v) in colors)

    juno_palette = unique([colors[k] for k in keys(colors) if k ∉ ["background", "variable"]])

    colvec = sort(HSV.(juno_palette), lt=(a,b) -> a.h < b.h)
    filter!(c -> c.s > 0.5*mean(c -> c.s, colvec), colvec)
    grad = Vector{eltype(colvec)}()
    for i = 1:length(colvec)-1
        append!(grad, range(colvec[i], stop = colvec[i+1]))
    end

    global _themes
    _themes[:juno] = PlotTheme(
        bg = colors["background"],
        bginside = colors["background"],
        fg = colors["variable"],
        fgtext = colors["variable"],
        fgguide = colors["variable"],
        fglegend = colors["variable"],
        palette = expand_palette(colors["background"], juno_palette; lchoices = [57], cchoices = [100]),
        colorgradient = grad
    )
end


module StochasticIntegrators

    using HDF5
    using LinearAlgebra: dot, mul!
    using OffsetArrays
    using Reexport
    using SimpleSolvers

    @reexport using GeometricIntegrators.Common
    @reexport using GeometricIntegrators.Config
    @reexport using GeometricIntegrators.Utils

    import GeometricIntegrators.Equations: AbstractEquationSDE, AbstractEquationPSDE,
                                           SDE, PSDE, SPSDE, get_function_tuple

    import GeometricIntegrators.Integrators
    import GeometricIntegrators.Integrators: Integrator, Parameters

    import GeometricIntegrators.Integrators: integrate!, integrate_common!, get_internal_variables

    import GeometricIntegrators.Integrators: IntegratorCache, CacheDict, CacheType

    import GeometricIntegrators.Integrators: AbstractTableauERK, AbstractTableauIRK, CoefficientsRK

    import GeometricIntegrators.Integrators: create_internal_stage_vector, create_internal_stage_matrix,
                                             create_internal_stage_vector_with_zero, create_nonlinear_solver

    import GeometricIntegrators.Solutions

    import GeometricIntegrators.Solutions: AtomicSolution, Solution, ParallelSolution

    import GeometricIntegrators.Solutions: TimeSeries, DataSeries, SDataSeries, PDataSeries

    import GeometricIntegrators.Solutions: hdf5, timesteps, nsave, counter, offset, lastentry,
                                           DEFAULT_NSAVE, DEFAULT_NWRITE

    import GeometricIntegrators.Solutions: get_initial_conditions, get_initial_conditions!, set_initial_conditions!,
                                           get_solution, get_solution!, set_solution!,
                                           get_data!, set_data!, update!, compute_timeseries!,
                                           createHDF5, create_hdf5, create_hdf5!, save_attributes,
                                           init_timeteps_in_hdf5, init_solution_in_hdf5, copy_solution_to_hdf5

    import GeometricIntegrators.Tableaus: CoefficientsGLRK, CoefficientsLobattoIIIA, CoefficientsLobattoIIIB, CoefficientsLobattoIIID

    const DEFAULT_SCONV = :strong

    export SemiMartingale, WienerProcess, generate_wienerprocess!, conv

    include("solutions/wienerprocess.jl")

    export AtomicSolution, AtomicSolutionSDE, AtomicSolutionPSDE
    export update!, cut_periodic_solution!,
           get_increments, get_increments!, set_increments!

    include("solutions/atomic_solution_sde.jl")
    include("solutions/atomic_solution_psde.jl")

    export Solution, StochasticSolution
    export SolutionSDE, SSolutionSDE, PSolutionSDE, SolutionPSDE, SSolutionPSDE, PSolutionPSDE
    export hdf5, timesteps, nsave, counter, offset, lastentry
    export get_initial_conditions, get_initial_conditions!, set_initial_conditions!,
           get_solution, get_solution!, set_solution!,
           create_hdf5, create_hdf5!

    include("solutions/solution.jl")
    include("solutions/solution_sde.jl")
    include("solutions/solution_psde.jl")
    include("solutions/solution_hdf5.jl")

    include("solutions/atomic_solution_constructors.jl")

    export Parameters

    export Integrator, StochasticIntegrator
    export SDEIntegrator, PSDEIntegrator,
           StochasticIntegratorRK, StochasticIntegratorPRK

    export nstages, noisedims, integrate!

    export IntegratorSERK, TableauSERK
    export IntegratorSIRK, TableauSIRK
    export IntegratorSIPRK, TableauSIPRK
    export IntegratorSISPRK, TableauSISPRK
    export IntegratorWERK, TableauWERK
    export IntegratorWIRK, TableauWIRK

    include("integrators/abstract_integrators.jl")
    include("integrators/abstract_runge_kutta.jl")
    include("integrators/integrator_cache.jl")

    include("integrators/integrators_serk.jl")
    include("integrators/integrators_sirk.jl")
    include("integrators/integrators_siprk.jl")
    include("integrators/integrators_sisprk.jl")
    include("integrators/integrators_werk.jl")
    include("integrators/integrators_wirk.jl")
    include("integrators/common.jl")
    include("integrators/integrators.jl")

    include("tableaus/tableaus_sirk.jl")

    export TableauStochasticGLRK, TableauStochasticDIRK

    include("tableaus/tableaus_siprk.jl")

    export TableauStochasticStoermerVerlet, TableauStochasticSymplecticEuler

    include("tableaus/tableaus_sisprk.jl")

    export TableauStochasticLobattoIIIABD2, TableauModifiedStochasticStoermerVerlet

    include("tableaus/tableaus_serk.jl")

    export TableauPlaten, TableauBurrageR2, TableauBurrageCL
    export TableauBurrageE1, TableauBurrageG5, TableauStochasticHeun
    export TableauStochasticEuler

    include("tableaus/tableaus_werk.jl")

    export TableauRoesslerRS1, TableauRoesslerRS2

    include("tableaus/tableaus_wirk.jl")

    export TableauSRKw1, TableauSRKw2

end


struct MockMLJ <: Convention end

@testset "convention" begin
    set_convention(ScientificTypesBase.NoConvention())
    c = ""
    @test_logs (:warn, "No convention specified. Did you forget to use the `set_convention` function?") (c = ST.convention())
    @test c isa ST.NoConvention

    set_convention(MockMLJ())
    c = ST.convention()
    @test c isa MockMLJ
end

@testset "trait" begin
    TRAIT_FUNCTION_GIVEN_NAME[:table] = Tables.istable

    isjunk(::Any)     = false
    isjunk(s::String) = s == "junk" ? true : false
    TRAIT_FUNCTION_GIVEN_NAME[:junk] = isjunk

    X = [1,2,3]
    @test ST.trait(X) == :other
    X = (x = [1,2,3],
         y = [5,5,7])
    @test ST.trait(X) == :table
    TRAIT_FUNCTION_GIVEN_NAME[:named_tuple] = X -> X isa NamedTuple
    @test_throws Exception ST.trait(X)
end

@testset "nonmissing" begin
    U = Union{Missing,Int}
    @test nonmissing(U) == Int
end

@testset "table" begin
    T0 = Table(Continuous)
    @test T0 == Table{K} where K<:AbstractVector{<:Continuous}
    T1 = Table(Continuous, Count)
    @test T1 == Table{K} where K<:Union{AbstractVector{<:Continuous}, AbstractVector{<:Count}}
    T2 = Table(Continuous, Union{Missing,Continuous})
    @test T2 == Table{K} where K<:Union{AbstractVector{<:Union{Missing,Continuous}}}
end


module Coping

using Codex
using DataFrames

export julliet2scores

ans_mapping = Dict(
    "Nooit" => 1,
    "Bijna nooit" => 2,
    "Soms" => 3,
    "Vaak" => 4,
    "Altijd" => 5
)

ans2num(ans)::Int = return ans_mapping[ans]

score(responses...) = sum(ans2num.(responses))

"""
    julliet2scores(df::DataFrame)

In 2021-07, returns 289×6 DataFrame.

For Julliet, 2018 and 2019 are the same.
2020 is a bit different, but it shouldn't matter for this code, see `diff`
For more information, see the e-mail with the subject `Julliet`.

This file is not tested in-depth, because the procedure is the same as delta.
"""
function julliet2scores(df::DataFrame)
    cols = [
        :id,
        :completed_at,
        [:v1, :v5, :v8, :v11, :v14] => ByRow(score) => :problem_focused,
        [:v3, :v7, :v10, :v13, :v4] => ByRow(score) => :emotion_focused,
        [:v2, :v6, :v9, :v15, :v12] => ByRow(score) => :seeking_support
    ]
    select!(df, cols...)
    disallowmissing!(df)
end

end # module


export GridWorldBase

using MacroTools:@forward
using Random

"""
    GridWorldBase{O} <: AbstractArray{Bool, 3}

A basic representation of grid world.
The first dimension uses multi-hot encoding to encode objects in a tile.
The second and third dimension means the height and width of the grid.
"""
struct GridWorldBase{O} <: AbstractArray{Bool, 3}
    world::BitArray{3}
    objects::O
end

get_object(w::GridWorldBase) = w.objects

function GridWorldBase(objects::Tuple{Vararg{AbstractObject}}, x::Int, y::Int)
    world = BitArray{3}(undef, length(objects), x, y)
    fill!(world, false)
    GridWorldBase(world, objects)
end

@forward GridWorldBase.world Base.size, Base.getindex, Base.setindex!

@generated function Base.to_index(::GridWorldBase{O}, x::X) where {X<:AbstractObject, O}
    i = findfirst(X .=== O.parameters)
    isnothing(i) && error("unknow object $x")
    :($i)
end

Base.setindex!(w::GridWorldBase, v::Bool, o::AbstractObject, x::Int, y::Int) = setindex!(w.world, v, Base.to_index(w, o), x, y)
Base.setindex!(w::GridWorldBase, v::Bool, o::AbstractObject, i::CartesianIndex{2}) = setindex!(w, v, o, i[1], i[2])

Base.getindex(w::GridWorldBase, o::AbstractObject, x::Int, y::Int) = getindex(w.world, Base.to_index(w, o), x, y)
Base.getindex(w::GridWorldBase, o::AbstractObject, i::CartesianIndex{2}) = getindex(w, o, i[1], i[2])
Base.getindex(w::GridWorldBase, o::AbstractObject, x::Colon, y::Colon) = getindex(w.world, Base.to_index(w, o), x, y)

#####
# utils
#####

switch!(world::GridWorldBase, x, src::CartesianIndex{2}, dest::CartesianIndex{2}) = world[x, src], world[x, dest] = world[x, dest], world[x, src]

function switch!(world::GridWorldBase, src::CartesianIndex{2}, dest::CartesianIndex{2})
    for x in axes(world, 1)
        switch!(world, x, src, dest)
    end
end

function Random.rand(f::Function, w::GridWorldBase; max_try=typemax(Int), rng=Random.GLOBAL_RNG)
    inds = CartesianIndices((size(w, 2), size(w, 3)))
    for _ in 1:max_try
        pos = rand(rng, inds)
        f(view(w, :, pos)) && return pos
    end
    @warn "a rare case happened when sampling from GridWorldBase"
    return nothing
end

#####
# get_agent_view
#####

get_agent_view_inds((i, j), (m, n), ::Left) = CartesianIndices((i-(n-1)÷2:i+(n-(n-1)÷2)-1, j-m+1:j))
get_agent_view_inds((i, j), (m, n), ::Right) = CartesianIndices((i-(n-1)÷2:i+(n-(n-1)÷2)-1, j:j+m-1))
get_agent_view_inds((i, j), (m, n), ::Up) = CartesianIndices((i-m+1:i, j-(n-1)÷2:j+(n-(n-1)÷2)-1))
get_agent_view_inds((i, j), (m, n), ::Down) = CartesianIndices((i:i+m-1, j-(n-1)÷2:j+(n-(n-1)÷2)-1))

ind_map((i,j), (m, n), ::Left) = (m-j+1, i)
ind_map((i,j), (m, n), ::Right) = (j, n-i+1)
ind_map((i,j), (m, n), ::Up) = (m-i+1, n-j+1)
ind_map((i,j), (m, n), ::Down) = (i,j)

function get_agent_view!(v::AbstractArray{Bool,3}, a::AbstractArray{Bool,3}, p::CartesianIndex, dir::LRUD)
    view_size = (size(v, 2), size(v, 3))
    grid_size = (size(a,2),size(a,3))
    inds = get_agent_view_inds(p.I, view_size, dir)
    valid_inds = CartesianIndices(grid_size)
    for ind in CartesianIndices(inds)
        if inds[ind] ∈ valid_inds
            v[:, ind_map(ind.I, view_size, dir)...] .= a[:, inds[ind]]
        end
    end
    v
end