text:
function prepare_parameters_NDE_animation(𝒟train, uw_NN, vw_NN, wT_NN, f=1f-4, Nz=32)
    H = abs(𝒟train.uw.z[end] - 𝒟train.uw.z[1])
    τ = abs(𝒟train.t[:,1][end] - 𝒟train.t[:,1][1])
    u_scaling = 𝒟train.scalings["u"]
    v_scaling = 𝒟train.scalings["v"]
    T_scaling = 𝒟train.scalings["T"]
    uw_scaling = 𝒟train.scalings["uw"]
    vw_scaling = 𝒟train.scalings["vw"]
    wT_scaling = 𝒟train.scalings["wT"]
    μ_u = u_scaling.μ
    μ_v = v_scaling.μ
    σ_u = u_scaling.σ
    σ_v = v_scaling.σ
    σ_T = T_scaling.σ
    σ_uw = uw_scaling.σ
    σ_vw = vw_scaling.σ
    σ_wT = wT_scaling.σ
    uw_weights, re_uw = Flux.destructure(uw_NN)
    vw_weights, re_vw = Flux.destructure(vw_NN)
    wT_weights, re_wT = Flux.destructure(wT_NN)
    weights = Float32[uw_weights; vw_weights; wT_weights]
    D_cell = Float32.(Dᶜ(Nz, 1 / Nz))
    D_face = Float32.(Dᶠ(Nz, 1 / Nz))
    size_uw_NN = length(uw_weights)
    size_vw_NN = length(vw_weights)
    size_wT_NN = length(wT_weights)
    uw_range = 1:size_uw_NN
    vw_range = size_uw_NN + 1:size_uw_NN + size_vw_NN
    wT_range = size_uw_NN + size_vw_NN + 1:size_uw_NN + size_vw_NN + size_wT_NN
    return f, H, τ, Nz, u_scaling, v_scaling, T_scaling, uw_scaling, vw_scaling, wT_scaling, μ_u, μ_v, σ_u, σ_v, σ_T, σ_uw, σ_vw, σ_wT, weights, re_uw, re_vw, re_wT, D_cell, D_face, size_uw_NN, size_vw_NN, size_wT_NN, uw_range, vw_range, wT_range
end

function prepare_BCs(𝒟, scalings)
    uw_top = scalings.uw(𝒟.uw.coarse[end,1])
    uw_bottom = scalings.uw(𝒟.uw.coarse[1,1])
    vw_top = scalings.vw(𝒟.vw.coarse[end,1])
    vw_bottom = scalings.vw(𝒟.vw.coarse[1,1])
    wT_top = scalings.wT(𝒟.wT.coarse[end,1])
    wT_bottom = scalings.wT(𝒟.wT.coarse[1,1])
    return (uw=(top=uw_top, bottom=uw_bottom), vw=(top=vw_top, bottom=vw_bottom), wT=(top=wT_top, bottom=wT_bottom))
end

function prepare_BCs(𝒟, scalings, wT_top_function)
    uw_top = scalings.uw(𝒟.uw.coarse[end,1])
    uw_bottom = scalings.uw(𝒟.uw.coarse[1,1])
    vw_top = scalings.vw(𝒟.vw.coarse[end,1])
    vw_bottom = scalings.vw(𝒟.vw.coarse[1,1])
    wT_top = t -> scalings.wT(wT_top_function(t))
    wT_bottom = scalings.wT(𝒟.wT.coarse[1,1])
    return (uw=(top=uw_top, bottom=uw_bottom), vw=(top=vw_top, bottom=vw_bottom), wT=(top=wT_top, bottom=wT_bottom))
end

function solve_NDE_mutating(uw_NN, vw_NN, wT_NN, scalings, constants, BCs, derivatives, uvT₀, ts, timestepper, conditions)
    μ_u = scalings.u.μ
    μ_v = scalings.v.μ
    σ_u = scalings.u.σ
    σ_v = scalings.v.σ
    σ_T = scalings.T.σ
    σ_uw = scalings.uw.σ
    σ_vw = scalings.vw.σ
    σ_wT = scalings.wT.σ
    H, τ, f, Nz, g, α = constants.H, constants.τ, constants.f, constants.Nz, constants.g, constants.α
    ν₀, ν₋, Riᶜ, ΔRi, Pr = constants.ν₀, constants.ν₋, constants.Riᶜ, constants.ΔRi, constants.Pr
    D_face = derivatives.face
    D_cell = derivatives.cell

    uw = zeros(Float32, Nz+1)
    vw = similar(uw)
    wT = similar(uw)

    ∂u∂z = similar(uw)
    ∂v∂z = similar(uw)
    ∂T∂z = similar(uw)

    ν = similar(uw)
    ν_T = similar(uw)

    Ri = similar(uw)

    uw[1] = BCs.uw.bottom - scalings.uw(0f0)
    vw[1] = BCs.vw.bottom - scalings.vw(0f0)
    wT[1] = BCs.wT.bottom - scalings.wT(0f0)

    uw[end] = BCs.uw.top - scalings.uw(0f0)
    vw[end] = BCs.vw.top - scalings.vw(0f0)

    if BCs.wT.top isa Number
        wT[end] = BCs.wT.top - scalings.wT(0f0)
    else
        wT[end] = BCs.wT.top(0) - scalings.wT(0f0)
    end

    uw_interior = @view uw[2:end-1]
    vw_interior = @view vw[2:end-1]
    wT_interior = @view wT[2:end-1]

    ∂uw∂z = zeros(Float32, Nz)
    ∂vw∂z = similar(∂uw∂z)
    ∂wT∂z = similar(∂uw∂z)

    dx = zeros(Float32, 3Nz)

    function predict_flux!(uvT, u, v, T)
        uw_interior .= uw_NN(uvT)
        vw_interior .= vw_NN(uvT)
        wT_interior .= wT_NN(uvT)

        mul!(∂u∂z, D_face, u)
        mul!(∂v∂z, D_face, v)
        mul!(∂T∂z, D_face, T)

        Ri .= local_richardson.(∂u∂z, ∂v∂z, ∂T∂z, H, g, α, σ_u, σ_v, σ_T)

        ν .= ν₀ .+ ν₋ .* tanh_step.((Ri .- Riᶜ) ./ ΔRi)

        if conditions.convective_adjustment
            for i in 2:length(ν_T)-1
                ν_T[i] = ∂u∂z[i] > 0 ? ν[i] / Pr : constants.κ
            end
        else
            ν_T .= ν ./ Pr
        end

        uw_interior .-= σ_u ./ σ_uw ./ H .* @view(ν[2:end-1]) .* @view(∂u∂z[2:end-1])
        vw_interior .-= σ_v ./ σ_vw ./ H .* @view(ν[2:end-1]) .* @view(∂v∂z[2:end-1])
        wT_interior .-= σ_T ./ σ_wT ./ H .* @view(ν_T[2:end-1]) .* @view(∂T∂z[2:end-1])
    end

    function NDE!(dx, x, p, t)
        u = @view x[1:Nz]
        v = @view x[Nz + 1:2Nz]
        T = @view x[2Nz + 1:end]

        ∂u∂t = @view dx[1:Nz]
        ∂v∂t = @view dx[Nz+1:2Nz]
        ∂T∂t = @view dx[2Nz+1:end]

        predict_flux!(x, u, v, T)

        if !isa(BCs.wT.top, Number)
            wT[end] = BCs.wT.top(t * τ)
        end

        mul!(∂uw∂z, D_cell, uw)
        mul!(∂vw∂z, D_cell, vw)
        mul!(∂wT∂z, D_cell, wT)

        ∂u∂t .= -τ ./ H .* σ_uw ./ σ_u .* ∂uw∂z .+ f .* τ ./ σ_u .* (σ_v .* v .+ μ_v)
        ∂v∂t .= -τ ./ H .* σ_vw ./ σ_v .* ∂vw∂z .- f .* τ ./ σ_v .* (σ_u .* u .+ μ_u)
        ∂T∂t .= -τ ./ H .* σ_wT ./ σ_T .* ∂wT∂z
    end

    tspan = (ts[1], ts[end])
    prob = ODEProblem(NDE!, uvT₀, tspan)
    sol = Array(solve(prob, timestepper, saveat=ts))
    return sol
end

function solve_NDE_mutating_GPU(uw_NN, vw_NN, wT_NN, scalings, constants, BCs, derivatives, uvT₀, ts, tspan, timestepper)
    μ_u = scalings.u.μ
    μ_v = scalings.v.μ
    σ_u = scalings.u.σ
    σ_v = scalings.v.σ
    σ_T = scalings.T.σ
    σ_uw = scalings.uw.σ
    σ_vw = scalings.vw.σ
    σ_wT = scalings.wT.σ
    H, τ, f, Nz, g, α = constants.H, constants.τ, constants.f, constants.Nz, constants.g, constants.α
    ν₀, ν₋, Riᶜ, ΔRi, Pr = constants.ν₀, constants.ν₋, constants.Riᶜ, constants.ΔRi, constants.Pr
    D_face = derivatives.face |> gpu
    D_cell = derivatives.cell |> gpu

    uw = zeros(Float32, Nz+1)
    vw = similar(uw)
    wT = similar(uw)

    ∂u∂z = similar(uw) |> gpu
    ∂v∂z = similar(∂u∂z)
    ∂T∂z = similar(∂u∂z)

    ν = similar(∂u∂z)
    Ri = similar(∂u∂z)

    uw[1] = BCs.uw.bottom - scalings.uw(0f0)
    vw[1] = BCs.vw.bottom - scalings.vw(0f0)
    wT[1] = BCs.wT.bottom - scalings.wT(0f0)

    uw[end] = BCs.uw.top - scalings.uw(0f0)
    vw[end] = BCs.vw.top - scalings.vw(0f0)
    wT[end] = BCs.wT.top - scalings.wT(0f0)

    uw = uw |> gpu
    vw = vw |> gpu
    wT = wT |> gpu

    uw_interior = @view uw[2:end-1]
    vw_interior = @view vw[2:end-1]
    wT_interior = @view wT[2:end-1]

    ∂uw∂z = zeros(Float32, Nz) |> gpu
    ∂vw∂z = similar(∂uw∂z)
    ∂wT∂z = similar(∂uw∂z)

    # dx = zeros(Float32, 3Nz) |> gpu

    function predict_flux!(uvT, u, v, T)
        uw_interior .= uw_NN(uvT)
        vw_interior .= vw_NN(uvT)
        wT_interior .= wT_NN(uvT)

        mul!(∂u∂z, D_face, u)
        mul!(∂v∂z, D_face, v)
        mul!(∂T∂z, D_face, T)

        Ri .= local_richardson.(∂u∂z, ∂v∂z, ∂T∂z, H, g, α, σ_u, σ_v, σ_T)
        ν .= ν₀ .+ ν₋ .* tanh_step.((Ri .- Riᶜ) ./ ΔRi)

        uw_interior .-= σ_u ./ σ_uw ./ H .* @view(ν[2:end-1]) .* @view(∂u∂z[2:end-1])
        vw_interior .-= σ_v ./ σ_vw ./ H .* @view(ν[2:end-1]) .* @view(∂v∂z[2:end-1])
        wT_interior .-= σ_T ./ σ_wT ./ H .* @view(ν[2:end-1]) .* @view(∂T∂z[2:end-1]) ./ Pr
    end

    function NDE!(dx, x, p, t)
        u = @view x[1:Nz]
        v = @view x[Nz + 1:2Nz]
        T = @view x[2Nz + 1:end]

        ∂u∂t = @view dx[1:Nz]
        ∂v∂t = @view dx[Nz+1:2Nz]
        ∂T∂t = @view dx[2Nz+1:end]

        predict_flux!(x, u, v, T)

        mul!(∂uw∂z, D_cell, uw)
        mul!(∂vw∂z, D_cell, vw)
        mul!(∂wT∂z, D_cell, wT)

        ∂u∂t .= -τ ./ H .* σ_uw ./ σ_u .* ∂uw∂z .+ f .* τ ./ σ_u .* (σ_v .* v .+ μ_v)
        ∂v∂t .= -τ ./ H .* σ_vw ./ σ_v .* ∂vw∂z .- f .* τ ./ σ_v .* (σ_u .* u .+ μ_u)
        ∂T∂t .= -τ ./ H .* σ_wT ./ σ_T .* ∂wT∂z
    end

    prob = ODEProblem(NDE!, uvT₀, tspan)
    sol = Array(solve(prob, timestepper, saveat=ts))
    return sol
end

function NDE_profile(uw_NN, vw_NN, wT_NN, test_file, 𝒟test, 𝒟train, trange;
                              ν₀=1f-4, ν₋=1f-1, ΔRi=1f0, Riᶜ=0.25, Pr=1f0, κ=10f0, α=2f-4, g=9.80665f0, f=1f-4,
                              OUTPUT_PATH = "",
                              modified_pacanowski_philander=false, convective_adjustment=false,
                              smooth_NN=false, smooth_Ri=false,
                              zero_weights=false,
                              loss_scalings = (u=1f0, v=1f0, T=1f0, ∂u∂z=5f-3, ∂v∂z=5f-3, ∂T∂z=5f-3),
                              timestepper=ROCK4())

    # @assert !modified_pacanowski_philander || !convective_adjustment

    @info "Preparing constants"

    Nz = length(𝒟train.u.z)

    conditions = (modified_pacanowski_philander=modified_pacanowski_philander, convective_adjustment=convective_adjustment,
                    smooth_NN=smooth_NN, smooth_Ri=smooth_Ri,
                    zero_weights=zero_weights)

    constants, scalings, derivatives, NN_constructions, weights, NN_sizes, NN_ranges, filters = prepare_parameters_NDE_training(𝒟train, uw_NN, vw_NN, wT_NN, f, Nz, g, α, ν₀, ν₋, Riᶜ, ΔRi, Pr, κ, conditions)

    H, τ, f = constants.H, constants.τ, constants.f
    D_face, D_cell = derivatives.face, derivatives.cell

    diurnal = occursin("diurnal", test_file)

    if diurnal
        wT_top_function = diurnal_fluxes([test_file], constants)[1]
        BCs = prepare_BCs(𝒟test, scalings, wT_top_function)
    else
        BCs = prepare_BCs(𝒟test, scalings)
    end

    @info "Setting up differential equations"

    if modified_pacanowski_philander
        constants_NN_only = (H=constants.H, τ=constants.τ, f=constants.f, Nz=constants.Nz, g=constants.g, α=constants.α, ν₀=0f0, ν₋=0f0, Riᶜ=constants.Riᶜ, ΔRi=constants.ΔRi, Pr=constants.Pr)
    end

    t_test = Float32.(𝒟test.t[trange] ./ constants.τ)
    uvT₀ = [scalings.u(𝒟test.uvT_unscaled[1:Nz, 1]); scalings.v(𝒟test.uvT_unscaled[Nz + 1:2Nz, 1]); scalings.T(𝒟test.uvT_unscaled[2Nz + 1:3Nz, 1])]

    u_𝒟test_scaled = scalings.u.(split_u(𝒟test.uvT_unscaled, Nz))
    v_𝒟test_scaled = scalings.v.(split_v(𝒟test.uvT_unscaled, Nz))
    T_𝒟test_scaled = scalings.T.(split_T(𝒟test.uvT_unscaled, Nz))

    ∂u∂z_𝒟test_scaled = ∂_∂z(u_𝒟test_scaled, D_face)
    ∂v∂z_𝒟test_scaled = ∂_∂z(v_𝒟test_scaled, D_face)
    ∂T∂z_𝒟test_scaled = ∂_∂z(T_𝒟test_scaled, D_face)

    @info "Solving NDEs"

    sol = solve_NDE_mutating(uw_NN, vw_NN, wT_NN, scalings, constants, BCs, derivatives, uvT₀, t_test, timestepper, conditions)

    u_sol, v_sol, T_sol = split_u(sol, Nz), split_v(sol, Nz), split_T(sol, Nz)

    ∂u∂z_sol = ∂_∂z(u_sol, D_face)
    ∂v∂z_sol = ∂_∂z(v_sol, D_face)
    ∂T∂z_sol = ∂_∂z(T_sol, D_face)

    unscaled_losses = (
        u = loss_per_tstep(u_sol, u_𝒟test_scaled),
        v = loss_per_tstep(v_sol, v_𝒟test_scaled),
        T = loss_per_tstep(T_sol, T_𝒟test_scaled),
        ∂u∂z = loss_per_tstep(∂u∂z_sol, ∂u∂z_𝒟test_scaled),
        ∂v∂z = loss_per_tstep(∂v∂z_sol, ∂v∂z_𝒟test_scaled),
        ∂T∂z = loss_per_tstep(∂T∂z_sol, ∂T∂z_𝒟test_scaled),
    )

    scaled_losses = apply_loss_scalings(unscaled_losses, loss_scalings)

    profile_losses = scaled_losses.u .+ scaled_losses.v .+ scaled_losses.T
    gradient_losses = scaled_losses.∂u∂z .+ scaled_losses.∂v∂z .+ scaled_losses.∂T∂z

    @info "Solving diffusivity-only equations"

    if modified_pacanowski_philander
        zeros_uw_NN = NN_constructions.uw(zeros(Float32, NN_sizes.uw))
        zeros_vw_NN = NN_constructions.vw(zeros(Float32, NN_sizes.vw))
        zeros_wT_NN = NN_constructions.wT(zeros(Float32, NN_sizes.wT))

        sol_mpp = solve_NDE_mutating(zeros_uw_NN, zeros_vw_NN, zeros_wT_NN, scalings, constants, BCs, derivatives, uvT₀, t_test, timestepper, conditions)

        u_sol_mpp, v_sol_mpp, T_sol_mpp = split_u(sol_mpp, Nz), split_v(sol_mpp, Nz), split_T(sol_mpp, Nz)

        ∂u∂z_sol_mpp = ∂_∂z(u_sol_mpp, D_face)
        ∂v∂z_sol_mpp = ∂_∂z(v_sol_mpp, D_face)
        ∂T∂z_sol_mpp = ∂_∂z(T_sol_mpp, D_face)

        unscaled_losses_mpp = (
            u = loss_per_tstep(u_sol_mpp, u_𝒟test_scaled),
            v = loss_per_tstep(v_sol_mpp, v_𝒟test_scaled),
            T = loss_per_tstep(T_sol_mpp, T_𝒟test_scaled),
            ∂u∂z = loss_per_tstep(∂u∂z_sol_mpp, ∂u∂z_𝒟test_scaled),
            ∂v∂z = loss_per_tstep(∂v∂z_sol_mpp, ∂v∂z_𝒟test_scaled),
            ∂T∂z = loss_per_tstep(∂T∂z_sol_mpp, ∂T∂z_𝒟test_scaled),
            )

        scaled_losses_mpp = apply_loss_scalings(unscaled_losses_mpp, loss_scalings)

        profile_losses_mpp = scaled_losses_mpp.u .+ scaled_losses_mpp.v .+ scaled_losses_mpp.T
        gradient_losses_mpp = scaled_losses_mpp.∂u∂z .+ scaled_losses_mpp.∂v∂z .+ scaled_losses_mpp.∂T∂z
    end

    if diurnal
        BCs_unscaled = (uw=(top=𝒟test.uw.coarse[end, 1], bottom=𝒟test.uw.coarse[1, 1]),
        vw=(top=𝒟test.vw.coarse[end, 1], bottom=𝒟test.uw.coarse[1, 1]),
        wT=(top=wT_top_function, bottom=𝒟test.wT.coarse[1, 1]))
    else
        BCs_unscaled = (uw=(top=𝒟test.uw.coarse[end, 1], bottom=𝒟test.uw.coarse[1, 1]),
        vw=(top=𝒟test.vw.coarse[end, 1], bottom=𝒟test.uw.coarse[1, 1]),
        wT=(top=𝒟test.wT.coarse[end, 1], bottom=𝒟test.wT.coarse[1, 1]))
    end

    ICs_unscaled = (u=𝒟test.u.coarse[:,1], v=𝒟test.v.coarse[:,1], T=𝒟test.T.coarse[:,1])

    t = 𝒟test.t[trange]

    @info "Solving k-profile parameterizations"

    sol_kpp = column_model_1D_kpp(constants, BCs_unscaled, ICs_unscaled, t, OceanTurb.KPP.Parameters())

    ∂u∂z_sol_kpp = ∂_∂z(scalings.u.(sol_kpp.U), D_face)
    ∂v∂z_sol_kpp = ∂_∂z(scalings.v.(sol_kpp.V), D_face)
    ∂T∂z_sol_kpp = ∂_∂z(scalings.T.(sol_kpp.T), D_face)

    unscaled_losses_kpp = (
        u = loss_per_tstep(scalings.u.(sol_kpp.U), u_𝒟test_scaled),
        v = loss_per_tstep(scalings.v.(sol_kpp.V), v_𝒟test_scaled),
        T = loss_per_tstep(scalings.T.(sol_kpp.T), T_𝒟test_scaled),
        ∂u∂z = loss_per_tstep(∂u∂z_sol_kpp, ∂u∂z_𝒟test_scaled),
        ∂v∂z = loss_per_tstep(∂v∂z_sol_kpp, ∂v∂z_𝒟test_scaled),
        ∂T∂z = loss_per_tstep(∂T∂z_sol_kpp, ∂T∂z_𝒟test_scaled),
        )

    scaled_losses_kpp = apply_loss_scalings(unscaled_losses_kpp, loss_scalings)

    profile_losses_kpp = scaled_losses_kpp.u .+ scaled_losses_kpp.v .+ scaled_losses_kpp.T
    gradient_losses_kpp = scaled_losses_kpp.∂u∂z .+ scaled_losses_kpp.∂v∂z .+ scaled_losses_kpp.∂T∂z

    @info "Preparing outputs"

    truth_u = 𝒟test.uvT_unscaled[1:Nz, trange]
    truth_v = 𝒟test.uvT_unscaled[Nz + 1:2Nz, trange]
    truth_T = 𝒟test.uvT_unscaled[2Nz + 1:3Nz, trange]

    truth_uw = 𝒟test.uw.coarse[:,trange]
    truth_vw = 𝒟test.vw.coarse[:,trange]
    truth_wT = 𝒟test.wT.coarse[:,trange]

    test_u = inv(scalings.u).(sol[1:Nz,:])
    test_v = inv(scalings.v).(sol[Nz + 1:2Nz, :])
    test_T = inv(scalings.T).(sol[2Nz + 1: 3Nz, :])

    test_uw = similar(truth_uw)
    test_vw = similar(truth_vw)
    test_wT = similar(truth_wT)

    for i in 1:size(test_uw, 2)
        if diurnal
            BCs_flux = (
                uw = (top=BCs.uw.top, bottom=top=BCs.uw.bottom),
                vw = (top=BCs.vw.top, bottom=top=BCs.vw.bottom),
                wT = (top=wT_scaling(wT_top_function(t[i])), bottom=BCs.wT.bottom),
                )
            test_uw[:,i], test_vw[:,i], test_wT[:,i] = predict_flux(uw_NN, vw_NN, wT_NN, @view(sol[:,i]), BCs_flux, conditions, scalings, constants, derivatives, filters)
        else
            test_uw[:,i], test_vw[:,i], test_wT[:,i] = predict_flux(uw_NN, vw_NN, wT_NN, @view(sol[:,i]), BCs, conditions, scalings, constants, derivatives, filters)
        end
    end

    test_uw .= inv(scalings.uw).(test_uw)
    test_vw .= inv(scalings.vw).(test_vw)
    test_wT .= inv(scalings.wT).(test_wT)

    truth_Ri = similar(𝒟test.uw.coarse[:,trange])
    test_Ri = similar(truth_Ri)

    for i in 1:size(truth_Ri, 2)
        truth_Ri[:,i] .= local_richardson.(D_face * 𝒟test.u.scaled[:,i], D_face * 𝒟test.v.scaled[:,i], D_face * 𝒟test.T.scaled[:,i], H, g, α, scalings.u.σ, scalings.v.σ, scalings.T.σ)
        test_Ri[:,i] .= local_richardson.(D_face * sol[1:Nz,i], D_face * sol[Nz + 1:2Nz, i], D_face * sol[2Nz + 1: 3Nz, i], H, g, α, scalings.u.σ, scalings.v.σ, scalings.T.σ)
    end

    if modified_pacanowski_philander
        test_uw_mpp = similar(truth_uw)
        test_vw_mpp = similar(truth_vw)
        test_wT_mpp = similar(truth_wT)

        for i in 1:size(test_uw_mpp, 2)

            if diurnal
                BCs_flux = (
                    uw = (top=BCs.uw.top, bottom=top=BCs.uw.bottom),
                    vw = (top=BCs.vw.top, bottom=top=BCs.vw.bottom),
                    wT = (top=wT_scaling(wT_top_function(t[i])), bottom=BCs.wT.bottom),
                    )
                test_uw_mpp[:,i], test_vw_mpp[:,i], test_wT_mpp[:,i] = predict_flux(NN_constructions.uw(zeros(Float32, NN_sizes.uw)),
                                                                            NN_constructions.vw(zeros(Float32, NN_sizes.vw)),
                                                                            NN_constructions.wT(zeros(Float32, NN_sizes.wT)),
                                                                            @view(sol_mpp[:,i]), BCs_flux, conditions, scalings, constants, derivatives, filters)
            else
                test_uw_mpp[:,i], test_vw_mpp[:,i], test_wT_mpp[:,i] = predict_flux(NN_constructions.uw(zeros(Float32, NN_sizes.uw)),
                                                                            NN_constructions.vw(zeros(Float32, NN_sizes.vw)),
                                                                            NN_constructions.wT(zeros(Float32, NN_sizes.wT)),
                                                                            @view(sol_mpp[:,i]), BCs, conditions, scalings, constants, derivatives, filters)
            end
        end

        test_uw_mpp .= inv(scalings.uw).(test_uw_mpp)
        test_vw_mpp .= inv(scalings.vw).(test_vw_mpp)
        test_wT_mpp .= inv(scalings.wT).(test_wT_mpp)
        test_u_mpp = inv(scalings.u).(sol_mpp[1:Nz,:])
        test_v_mpp = inv(scalings.v).(sol_mpp[Nz + 1:2Nz, :])
        test_T_mpp = inv(scalings.T).(sol_mpp[2Nz + 1: 3Nz, :])

        test_Ri_mpp = similar(truth_Ri)

        for i in 1:size(test_Ri_mpp,2)
            test_Ri_mpp[:,i] .=
            local_richardson.(D_face * sol_mpp[1:Nz,i],
            D_face * sol_mpp[Nz + 1:2Nz, i],
            D_face * sol_mpp[2Nz + 1: 3Nz, i], H, g, α, scalings.u.σ, scalings.v.σ, scalings.T.σ)
        end

        test_uw_NN_only = similar(truth_uw)
        test_vw_NN_only = similar(truth_vw)
        test_wT_NN_only = similar(truth_wT)

        for i in 1:size(test_uw_NN_only, 2)
            if diurnal
                BCs_flux = (
                    uw = (top=BCs.uw.top, bottom=top=BCs.uw.bottom),
                    vw = (top=BCs.vw.top, bottom=top=BCs.vw.bottom),
                    wT = (top=wT_scaling(wT_top_function(t[i])), bottom=BCs.wT.bottom),
                    )
                test_uw_NN_only[:,i], test_vw_NN_only[:,i], test_wT_NN_only[:,i] =
                    predict_flux(uw_NN, vw_NN, wT_NN, @view(sol[:,i]), BCs_flux, conditions, scalings, constants_NN_only, derivatives, filters)
            else
                test_uw_NN_only[:,i], test_vw_NN_only[:,i], test_wT_NN_only[:,i] =
                    predict_flux(uw_NN, vw_NN, wT_NN, @view(sol[:,i]), BCs, conditions, scalings, constants_NN_only, derivatives, filters)
            end
        end

        test_uw_NN_only .= inv(scalings.uw).(test_uw_NN_only)
        test_vw_NN_only .= inv(scalings.vw).(test_vw_NN_only)
        test_wT_NN_only .= inv(scalings.wT).(test_wT_NN_only)
    end

    test_u_kpp = sol_kpp.U
    test_v_kpp = sol_kpp.V
    test_T_kpp = sol_kpp.T

    test_uw_kpp = sol_kpp.UW
    test_vw_kpp = sol_kpp.VW
    test_wT_kpp = sol_kpp.WT

    test_Ri_kpp = similar(truth_Ri)

    for i in 1:size(test_Ri_kpp,2)
        test_Ri_kpp[:,i] .= local_richardson.(D_face * scalings.u.(@view(sol_kpp.U[:,i])),
        D_face * scalings.v.(@view(sol_kpp.V[:,i])),
        D_face * scalings.T.(@view(sol_kpp.T[:,i])),
        H, g, α, scalings.u.σ, scalings.v.σ, scalings.T.σ)
    end

    test_uw .= test_uw .- test_uw[1, 1]
    test_vw .= test_vw .- test_vw[1, 1]
    test_wT .= test_wT .- test_wT[1, 1]

    if modified_pacanowski_philander
        test_uw_mpp .= test_uw_mpp .- test_uw_mpp[1, 1]
        test_vw_mpp .= test_vw_mpp .- test_vw_mpp[1, 1]
        test_wT_mpp .= test_wT_mpp .- test_wT_mpp[1, 1]

        test_uw_NN_only .= test_uw_NN_only .- test_uw_NN_only[1, 1]
        test_vw_NN_only .= test_vw_NN_only .- test_vw_NN_only[1, 1]
        test_wT_NN_only .= test_wT_NN_only .- test_wT_NN_only[1, 1]
    end

    depth_profile = 𝒟test.u.z
    depth_flux = 𝒟test.uw.z

    @info "Writing outputs"

    output = Dict()

    output["depth_profile"] = 𝒟test.u.z
    output["depth_flux"] = 𝒟test.uw.z
    output["t"] = 𝒟test.t[trange]

    output["truth_u"] = truth_u
    output["truth_v"] = truth_v
    output["truth_T"] = truth_T

    output["truth_uw"] = truth_uw
    output["truth_vw"] = truth_vw
    output["truth_wT"] = truth_wT

    output["truth_Ri"] = truth_Ri

    output["test_u"] = test_u
    output["test_v"] = test_v
    output["test_T"] = test_T

    output["test_uw"] = test_uw
    output["test_vw"] = test_vw
    output["test_wT"] = test_wT

    output["test_Ri"] = test_Ri

    output["u_losses"] = scaled_losses.u
    output["v_losses"] = scaled_losses.v
    output["T_losses"] = scaled_losses.T
    output["∂u∂z_losses"] = scaled_losses.∂u∂z
    output["∂v∂z_losses"] = scaled_losses.∂v∂z
    output["∂T∂z_losses"] = scaled_losses.∂T∂z

    output["losses"] = profile_losses
    output["loss"] = mean(profile_losses)
    output["losses_gradient"] = gradient_losses
    output["loss_gradient"] = mean(gradient_losses)

    if modified_pacanowski_philander
        output["train_parameters"] = (ν₀=ν₀, ν₋=ν₋, ΔRi=ΔRi, Riᶜ=Riᶜ, Pr=Pr, loss_scalings=loss_scalings)

        output["test_u_modified_pacanowski_philander"] = test_u_mpp
        output["test_v_modified_pacanowski_philander"] = test_v_mpp
        output["test_T_modified_pacanowski_philander"] = test_T_mpp

        output["test_uw_modified_pacanowski_philander"] = test_uw_mpp
        output["test_vw_modified_pacanowski_philander"] = test_vw_mpp
        output["test_wT_modified_pacanowski_philander"] = test_wT_mpp

        output["test_uw_NN_only"] = test_uw_NN_only
        output["test_vw_NN_only"] = test_vw_NN_only
        output["test_wT_NN_only"] = test_wT_NN_only

        output["test_Ri_modified_pacanowski_philander"] = test_Ri_mpp

        output["u_losses_modified_pacanowski_philander"] = scaled_losses_mpp.u
        output["v_losses_modified_pacanowski_philander"] = scaled_losses_mpp.v
        output["T_losses_modified_pacanowski_philander"] = scaled_losses_mpp.T
        output["∂u∂z_losses_modified_pacanowski_philander"] = scaled_losses_mpp.∂u∂z
        output["∂v∂z_losses_modified_pacanowski_philander"] = scaled_losses_mpp.∂v∂z
        output["∂T∂z_losses_modified_pacanowski_philander"] = scaled_losses_mpp.∂T∂z

        output["losses_modified_pacanowski_philander"] = profile_losses_mpp
        output["loss_modified_pacanowski_philander"] = mean(profile_losses_mpp)
        output["losses_modified_pacanowski_philander_gradient"] = gradient_losses_mpp
        output["loss_modified_pacanowski_philander_gradient"] = mean(gradient_losses_mpp)
    end

    output["test_u_kpp"] = test_u_kpp
    output["test_v_kpp"] = test_v_kpp
    output["test_T_kpp"] = test_T_kpp

    output["test_uw_kpp"] = test_uw_kpp
    output["test_vw_kpp"] = test_vw_kpp
    output["test_wT_kpp"] = test_wT_kpp

    output["test_Ri_kpp"] = test_Ri_kpp

    output["u_losses_kpp"] = scaled_losses_kpp.u
    output["v_losses_kpp"] = scaled_losses_kpp.v
    output["T_losses_kpp"] = scaled_losses_kpp.T
    output["∂u∂z_losses_kpp"] = scaled_losses_kpp.∂u∂z
    output["∂v∂z_losses_kpp"] = scaled_losses_kpp.∂v∂z
    output["∂T∂z_losses_kpp"] = scaled_losses_kpp.∂T∂z

    output["losses_kpp"] = profile_losses_kpp
    output["loss_kpp"] = mean(profile_losses_kpp)
    output["losses_kpp_gradient"] = gradient_losses_kpp
    output["loss_kpp_gradient"] = mean(gradient_losses_kpp)

    if OUTPUT_PATH !== ""
        @info "Saving file"
        jldopen(OUTPUT_PATH, "w") do file
            file["NDE_profile"] = output
        end
    end

    return output
end

function solve_oceananigans_modified_pacanowski_philander_nn(test_files, EXTRACTED_FILE_PATH, OUTPUT_DIR;
                                                        timestep=60, convective_adjustment=false)
    @info "Loading Training Data..."
    extracted_training_file = jldopen(EXTRACTED_FILE_PATH, "r")

    uw_NN = extracted_training_file["neural_network/uw"]
    vw_NN = extracted_training_file["neural_network/vw"]
    wT_NN = extracted_training_file["neural_network/wT"]

    train_files = extracted_training_file["training_info/train_files"]

    diurnal = occursin("diurnal", test_files[1])

    𝒟train = WindMixing.data(train_files, scale_type=ZeroMeanUnitVarianceScaling, enforce_surface_fluxes=false)

    u_scaling = 𝒟train.scalings["u"]
    v_scaling = 𝒟train.scalings["v"]
    T_scaling = 𝒟train.scalings["T"]
    uw_scaling = 𝒟train.scalings["uw"]
    vw_scaling = 𝒟train.scalings["vw"]
    wT_scaling = 𝒟train.scalings["wT"]

    scalings = (u=u_scaling, v=v_scaling, T=T_scaling, uw=uw_scaling, vw=vw_scaling, wT=wT_scaling)
    diffusivity_params = extracted_training_file["training_info/parameters"]

    close(extracted_training_file)

    if !ispath(OUTPUT_DIR)
        mkdir(OUTPUT_DIR)
    end

    for test_file in test_files
        @info "Starting $test_file"
        ds = jldopen(directories[test_file])

        f = ds["parameters/coriolis_parameter"]
        α = ds["parameters/thermal_expansion_coefficient"]
        g = ds["parameters/gravitational_acceleration"]
        Nz = 32
        Lz = ds["grid/Lz"]
        Δz = ds["grid/Δz"]

        frames = keys(ds["timeseries/t"])
        stop_time = ds["timeseries/t/$(frames[end])"]

        uw_flux = ds["parameters/boundary_condition_u_top"]
        vw_flux = 0

        if diurnal
            wT_flux = diurnal_fluxes([test_file], (; α, g))[1]
        else
            wT_flux = ds["parameters/boundary_condition_θ_top"]
        end

        T₀ = Array(ds["timeseries/T/0"][1, 1, :])

        ∂u₀∂z = ds["parameters/boundary_condition_u_bottom"]
        ∂v₀∂z = ds["parameters/boundary_condition_u_bottom"]

        constants = (; f, α, g, Nz, Lz, T₀)
        BCs = (top=(uw=uw_flux, vw=vw_flux, wT=wT_flux), bottom=(u=∂u₀∂z, v=∂v₀∂z))

        if test_file in train_files
            dir_str = "train_$test_file"
        else
            dir_str = "test_$test_file"
        end

        DIR_PATH = joinpath(OUTPUT_DIR, dir_str)

        if !ispath(DIR_PATH)
            mkdir(DIR_PATH)
        end

        BASELINE_RESULTS_PATH = joinpath(DIR_PATH, "baseline_oceananigans")
        NN_RESULTS_PATH = joinpath(DIR_PATH, "NN_oceananigans")

        oceananigans_modified_pacanowski_philander_nn(uw_NN, vw_NN, wT_NN, constants, BCs, scalings, diffusivity_params,
                                                    BASELINE_RESULTS_PATH=BASELINE_RESULTS_PATH,
                                                    NN_RESULTS_PATH=NN_RESULTS_PATH,
                                                    stop_time=stop_time, Δt=timestep,
                                                    convective_adjustment=convective_adjustment)
    end
end

function solve_oceananigans_modified_pacanowski_philander_nn(test_files, uw_NN_PATH, vw_NN_PATH, wT_NN_PATH, OUTPUT_DIR;
                                                        timestep=60, convective_adjustment=false)
    @info "Loading Training Data..."
    uw_file = jldopen(uw_NN_PATH, "r")
    vw_file = jldopen(vw_NN_PATH, "r")
    wT_file = jldopen(wT_NN_PATH, "r")

    uw_NN = uw_file["neural_network"]
    vw_NN = vw_file["neural_network"]
    wT_NN = wT_file["neural_network"]

    train_files = uw_file["training_info/train_files"]

    diurnal = occursin("diurnal", test_files[1])

    𝒟train = WindMixing.data(train_files, scale_type=ZeroMeanUnitVarianceScaling, enforce_surface_fluxes=false)

    u_scaling = 𝒟train.scalings["u"]
    v_scaling = 𝒟train.scalings["v"]
    T_scaling = 𝒟train.scalings["T"]
    uw_scaling = 𝒟train.scalings["uw"]
    vw_scaling = 𝒟train.scalings["vw"]
    wT_scaling = 𝒟train.scalings["wT"]

    scalings = (u=u_scaling, v=v_scaling, T=T_scaling, uw=uw_scaling, vw=vw_scaling, wT=wT_scaling)
    diffusivity_params = uw_file["training_info/parameters"]

    close(uw_file)
    close(vw_file)
    close(wT_file)

    if !ispath(OUTPUT_DIR)
        mkdir(OUTPUT_DIR)
    end

    for test_file in test_files
        @info "Starting $test_file"
        ds = jldopen(directories[test_file])

        f = ds["parameters/coriolis_parameter"]
        α = ds["parameters/thermal_expansion_coefficient"]
        g = ds["parameters/gravitational_acceleration"]
        Nz = 32
        Lz = ds["grid/Lz"]
        Δz = ds["grid/Δz"]

        frames = keys(ds["timeseries/t"])
        stop_time = ds["timeseries/t/$(frames[end])"]

        uw_flux = ds["parameters/boundary_condition_u_top"]
        vw_flux = 0

        if diurnal
            wT_flux = diurnal_fluxes([test_file], (; α, g))[1]
        else
            wT_flux = ds["parameters/boundary_condition_θ_top"]
        end

        T₀ = Array(ds["timeseries/T/0"][1, 1, :])

        ∂u₀∂z = ds["parameters/boundary_condition_u_bottom"]
        ∂v₀∂z = ds["parameters/boundary_condition_u_bottom"]

        constants = (; f, α, g, Nz, Lz, T₀)
        BCs = (top=(uw=uw_flux, vw=vw_flux, wT=wT_flux), bottom=(u=∂u₀∂z, v=∂v₀∂z))

        if test_file in train_files
            dir_str = "train_$test_file"
        else
            dir_str = "test_$test_file"
        end

        DIR_PATH = joinpath(OUTPUT_DIR, dir_str)

        if !ispath(DIR_PATH)
            mkdir(DIR_PATH)
        end

        BASELINE_RESULTS_PATH = joinpath(DIR_PATH, "baseline_oceananigans")
        NN_RESULTS_PATH = joinpath(DIR_PATH, "NN_oceananigans")

        oceananigans_modified_pacanowski_philander_nn(uw_NN, vw_NN, wT_NN, constants, BCs, scalings, diffusivity_params,
                                                    BASELINE_RESULTS_PATH=BASELINE_RESULTS_PATH,
                                                    NN_RESULTS_PATH=NN_RESULTS_PATH,
                                                    stop_time=stop_time, Δt=timestep,
                                                    convective_adjustment=convective_adjustment)
    end
end

function NDE_profile_oceananigans(FILE_DIR, train_files, test_files;
                                  ν₀=1f-1, ν₋=1f-4, ΔRi=1f-1, Riᶜ=0.25f0, Pr=1,
                                  loss_scalings=(u=1f0, v=1f0, T=1f0, ∂u∂z=5f-3, ∂v∂z=5f-3, ∂T∂z=5f-3),
                                  OUTPUT_PATH="")
    @assert length(test_files) == 1
    𝒟train = WindMixing.data(train_files, scale_type=ZeroMeanUnitVarianceScaling, enforce_surface_fluxes=false)
    𝒟test = WindMixing.data(test_files, scale_type=ZeroMeanUnitVarianceScaling, enforce_surface_fluxes=false)

    @info "Reading files"

    BASELINE_SOL_PATH = joinpath(FILE_DIR, "baseline_oceananigans.jld2")
    NDE_SOL_PATH = joinpath(FILE_DIR, "NN_oceananigans.jld2")

    baseline_sol = jldopen(BASELINE_SOL_PATH)
    NDE_sol = jldopen(NDE_SOL_PATH)

    frames = keys(baseline_sol["timeseries/t"])

    @assert length(frames) == length(𝒟test.t)

    @info "Loading constants"

    Nz = baseline_sol["grid/Nz"]
    α = baseline_sol["buoyancy/model/equation_of_state/α"]
    g = baseline_sol["buoyancy/model/gravitational_acceleration"]
    f = 1f-4
    t = 𝒟test.t
    zC = baseline_sol["grid/zC"][2:end-1]
    zF = baseline_sol["grid/zF"][2:end-1]
    H = abs(zF[1])
    constants = (; Nz, α, g, f, H)
    train_parameters = (; ν₀, ν₋, ΔRi, Riᶜ, Pr, loss_scalings)

    derivatives_dimensionless = (cell=Float32.(Dᶜ(Nz, 1 / Nz)), face=Float32.(Dᶠ(Nz, 1 / Nz)))

    u_scaling = 𝒟train.scalings["u"]
    v_scaling = 𝒟train.scalings["v"]
    T_scaling = 𝒟train.scalings["T"]
    uw_scaling = 𝒟train.scalings["uw"]
    vw_scaling = 𝒟train.scalings["vw"]
    wT_scaling = 𝒟train.scalings["wT"]

    scalings = (u=u_scaling, v=v_scaling, T=T_scaling, uw=uw_scaling, vw=vw_scaling, wT=wT_scaling)

    @info "Loading solutions"

    truth_u = 𝒟test.u.coarse
    truth_v = 𝒟test.v.coarse
    truth_T = 𝒟test.T.coarse

    truth_uw = 𝒟test.uw.coarse
    truth_vw = 𝒟test.vw.coarse
    truth_wT = 𝒟test.wT.coarse

    test_u_mpp = similar(truth_u)
    test_v_mpp = similar(truth_u)
    test_T_mpp = similar(truth_u)

    test_uw_mpp = similar(truth_uw)
    test_vw_mpp = similar(truth_uw)
    test_wT_mpp = similar(truth_uw)

    test_u = similar(truth_u)
    test_v = similar(truth_u)
    test_T = similar(truth_u)

    test_uw = similar(truth_uw)
    test_vw = similar(truth_uw)
    test_wT = similar(truth_uw)

    for i in 1:size(truth_u,2)
        test_u_mpp[:,i] .= baseline_sol["timeseries/u/$(frames[i])"][:]
        test_v_mpp[:,i] .= baseline_sol["timeseries/v/$(frames[i])"][:]
        test_T_mpp[:,i] .= baseline_sol["timeseries/T/$(frames[i])"][:]
        test_uw_mpp[:,i] .= baseline_sol["timeseries/uw/$(frames[i])"][:]
        test_vw_mpp[:,i] .= baseline_sol["timeseries/vw/$(frames[i])"][:]
        test_wT_mpp[:,i] .= baseline_sol["timeseries/wT/$(frames[i])"][:]

        test_u[:,i] .= NDE_sol["timeseries/u/$(frames[i])"][:]
        test_v[:,i] .= NDE_sol["timeseries/v/$(frames[i])"][:]
        test_T[:,i] .= NDE_sol["timeseries/T/$(frames[i])"][:]
        test_uw[:,i] .= NDE_sol["timeseries/uw/$(frames[i])"][:]
        test_vw[:,i] .= NDE_sol["timeseries/vw/$(frames[i])"][:]
        test_wT[:,i] .= NDE_sol["timeseries/wT/$(frames[i])"][:]
    end

    close(baseline_sol)
    close(NDE_sol)

    test_uw_NN_only = similar(truth_uw)
    test_vw_NN_only = similar(truth_uw)
    test_wT_NN_only = similar(truth_uw)

    for i in 1:size(test_uw_NN_only,2)
        uw_total = @view test_uw[:, i]
        vw_total = @view test_vw[:, i]
        wT_total = @view test_wT[:, i]

        uw_mpp = @view test_uw_mpp[:, i]
        vw_mpp = @view test_vw_mpp[:, i]
        wT_mpp = @view test_wT_mpp[:, i]

        test_uw_NN_only[:, i] .= uw_total .+ uw_mpp
        test_vw_NN_only[:, i] .= vw_total .+ vw_mpp
        test_wT_NN_only[:, i] .= wT_total .+ wT_mpp
    end

    @inline function local_richardson(∂u∂z, ∂v∂z, ∂T∂z, g, α)
        Bz = g * α * ∂T∂z
        S² = ∂u∂z ^2 + ∂v∂z ^2
        return Bz / S²
    end

    @inline function local_richardson(∂u∂z, ∂v∂z, ∂T∂z, H, g, α, σ_u, σ_v, σ_T)
        Bz = H * g * α * σ_T * ∂T∂z
        S² = (σ_u * ∂u∂z) ^2 + (σ_v * ∂v∂z) ^2
        return Bz / S²
    end

    D_face = Float32.(Dᶠ(Nz, zC[2] - zC[1]))
    D_face_dimensionless = derivatives_dimensionless.face

    truth_Ri = local_richardson.(∂_∂z(truth_u, D_face), ∂_∂z(truth_v, D_face), ∂_∂z(truth_T, D_face), g, α)
    test_Ri = local_richardson.(∂_∂z(test_u, D_face), ∂_∂z(test_v, D_face), ∂_∂z(test_T, D_face), g, α)
    test_Ri_modified_pacanowski_philander = local_richardson.(∂_∂z(test_u_mpp, D_face), ∂_∂z(test_v_mpp, D_face), ∂_∂z(test_T_mpp, D_face), g, α)

    diurnal = occursin("diurnal", test_files[1])

    if diurnal
        wT_top_function = diurnal_fluxes(test_files, constants)[1]
        BCs_unscaled = (uw=(top=𝒟test.uw.coarse[end, 1], bottom=𝒟test.uw.coarse[1, 1]),
        vw=(top=𝒟test.vw.coarse[end, 1], bottom=𝒟test.uw.coarse[1, 1]),
        wT=(top=wT_top_function, bottom=𝒟test.wT.coarse[1, 1]))
    else
        BCs_unscaled = (uw=(top=𝒟test.uw.coarse[end, 1], bottom=𝒟test.uw.coarse[1, 1]),
        vw=(top=𝒟test.vw.coarse[end, 1], bottom=𝒟test.uw.coarse[1, 1]),
        wT=(top=𝒟test.wT.coarse[end, 1], bottom=𝒟test.wT.coarse[1, 1]))
        @show 𝒟test.wT.coarse[end, 1]
    end

    ICs_unscaled = (u=𝒟test.u.coarse[:,1], v=𝒟test.v.coarse[:,1], T=𝒟test.T.coarse[:,1])

    trange = 1:1153
    t = 𝒟test.t[trange]

    @info "Solving k-profile parameterizations"

    sol_kpp = column_model_1D_kpp(constants, BCs_unscaled, ICs_unscaled, t, OceanTurb.KPP.Parameters())

    test_u_kpp = sol_kpp.U
    test_v_kpp = sol_kpp.V
    test_T_kpp = sol_kpp.T

    test_uw_kpp = sol_kpp.UW
    test_vw_kpp = sol_kpp.VW
    test_wT_kpp = sol_kpp.WT

    test_Ri_kpp = similar(truth_Ri)

    for i in 1:size(test_Ri_kpp,2)
        test_Ri_kpp[:,i] .= local_richardson.(D_face * scalings.u.(@view(sol_kpp.U[:,i])),
        D_face * scalings.v.(@view(sol_kpp.V[:,i])),
        D_face * scalings.T.(@view(sol_kpp.T[:,i])),
                            H, g, α, scalings.u.σ, scalings.v.σ, scalings.T.σ)
    end

    @info "Calculating Losses"

    truth_u_scaled = scalings.u.(split_u(𝒟test.uvT_unscaled, Nz))
    truth_v_scaled = scalings.v.(split_v(𝒟test.uvT_unscaled, Nz))
    truth_T_scaled = scalings.T.(split_T(𝒟test.uvT_unscaled, Nz))

    baseline_u_scaled = scalings.u.(test_u_mpp)
    baseline_v_scaled = scalings.v.(test_v_mpp)
    baseline_T_scaled = scalings.T.(test_T_mpp)

    test_u_scaled = scalings.u.(test_u)
    test_v_scaled = scalings.v.(test_v)
    test_T_scaled = scalings.T.(test_T)

    truth_∂u∂z_scaled = ∂_∂z(truth_u_scaled, D_face_dimensionless)
    truth_∂v∂z_scaled = ∂_∂z(truth_v_scaled, D_face_dimensionless)
    truth_∂T∂z_scaled = ∂_∂z(truth_T_scaled, D_face_dimensionless)

    baseline_∂u∂z_scaled = ∂_∂z(baseline_u_scaled, D_face_dimensionless)
    baseline_∂v∂z_scaled = ∂_∂z(baseline_v_scaled, D_face_dimensionless)
    baseline_∂T∂z_scaled = ∂_∂z(baseline_T_scaled, D_face_dimensionless)

    test_∂u∂z_scaled = ∂_∂z(test_u_scaled, D_face_dimensionless)
    test_∂v∂z_scaled = ∂_∂z(test_v_scaled, D_face_dimensionless)
    test_∂T∂z_scaled = ∂_∂z(test_T_scaled, D_face_dimensionless)

    u_loss_unscaled = loss_per_tstep(truth_u_scaled, test_u_scaled)
    u_loss_unscaled = loss_per_tstep(truth_u_scaled, test_u_scaled)
    u_loss_unscaled = loss_per_tstep(truth_u_scaled, test_u_scaled)

    u_loss = loss_per_tstep(truth_u_scaled, test_u_scaled)
    u_loss = loss_per_tstep(truth_u_scaled, test_u_scaled)
    u_loss = loss_per_tstep(truth_u_scaled, test_u_scaled)

    unscaled_losses = (
        u = loss_per_tstep(truth_u_scaled, test_u_scaled),
        v = loss_per_tstep(truth_v_scaled, test_v_scaled),
        T = loss_per_tstep(truth_T_scaled, test_T_scaled),
        ∂u∂z = loss_per_tstep(truth_∂u∂z_scaled, test_∂u∂z_scaled),
        ∂v∂z = loss_per_tstep(truth_∂v∂z_scaled, test_∂v∂z_scaled),
        ∂T∂z = loss_per_tstep(truth_∂T∂z_scaled, test_∂T∂z_scaled),
        )

    scaled_losses = apply_loss_scalings(unscaled_losses, loss_scalings)

    profile_losses = scaled_losses.u .+ scaled_losses.v .+ scaled_losses.T
    gradient_losses = scaled_losses.∂u∂z .+ scaled_losses.∂v∂z .+ scaled_losses.∂T∂z

    profile_loss = mean(profile_losses)
    loss_gradient = mean(gradient_losses)

    unscaled_losses_mpp = (
        u = loss_per_tstep(truth_u_scaled, baseline_u_scaled),
        v = loss_per_tstep(truth_v_scaled, baseline_v_scaled),
        T = loss_per_tstep(truth_T_scaled, baseline_T_scaled),
        ∂u∂z = loss_per_tstep(truth_∂u∂z_scaled, baseline_∂u∂z_scaled),
        ∂v∂z = loss_per_tstep(truth_∂v∂z_scaled, baseline_∂v∂z_scaled),
        ∂T∂z = loss_per_tstep(truth_∂T∂z_scaled, baseline_∂T∂z_scaled),
        )

    scaled_losses_mpp = apply_loss_scalings(unscaled_losses_mpp, loss_scalings)

    profile_losses_mpp = scaled_losses_mpp.u .+ scaled_losses_mpp.v .+ scaled_losses_mpp.T
    gradient_losses_mpp = scaled_losses_mpp.∂u∂z .+ scaled_losses_mpp.∂v∂z .+ scaled_losses_mpp.∂T∂z

    profile_loss_mpp = mean(profile_losses_mpp)
    loss_gradient_mpp = mean(gradient_losses_mpp)

    ∂u∂z_sol_kpp = ∂_∂z(scalings.u.(sol_kpp.U), D_face)
    ∂v∂z_sol_kpp = ∂_∂z(scalings.v.(sol_kpp.V), D_face)
    ∂T∂z_sol_kpp = ∂_∂z(scalings.T.(sol_kpp.T), D_face)

    unscaled_losses_kpp = (
        u = loss_per_tstep(scalings.u.(sol_kpp.U), truth_u_scaled),
        v = loss_per_tstep(scalings.v.(sol_kpp.V), truth_v_scaled),
        T = loss_per_tstep(scalings.T.(sol_kpp.T), truth_T_scaled),
        ∂u∂z = loss_per_tstep(∂u∂z_sol_kpp, truth_∂u∂z_scaled),
        ∂v∂z = loss_per_tstep(∂v∂z_sol_kpp, truth_∂v∂z_scaled),
        ∂T∂z = loss_per_tstep(∂T∂z_sol_kpp, truth_∂T∂z_scaled),
        )

    scaled_losses_kpp = apply_loss_scalings(unscaled_losses_kpp, loss_scalings)

    profile_losses_kpp = scaled_losses_kpp.u .+ scaled_losses_kpp.v .+ scaled_losses_kpp.T
    gradient_losses_kpp = scaled_losses_kpp.∂u∂z .+ scaled_losses_kpp.∂v∂z .+ scaled_losses_kpp.∂T∂z

    profile_loss_kpp = mean(profile_losses_kpp)
    loss_gradient_kpp = mean(gradient_losses_kpp)

    @info "Writing outputs"

    output = Dict(
           "depth_profile" => zC,
              "depth_flux" => zF,
                       "t" => t,
        "train_parameters" => train_parameters,

        "truth_u" => truth_u,
        "truth_v" => truth_v,
        "truth_T" => truth_T,

        "test_u" => test_u,
        "test_v" => test_v,
        "test_T" => test_T,

        "test_u_modified_pacanowski_philander" => test_u_mpp,
        "test_v_modified_pacanowski_philander" => test_v_mpp,
        "test_T_modified_pacanowski_philander" => test_T_mpp,

        "test_u_kpp" => test_u_kpp,
        "test_v_kpp" => test_v_kpp,
        "test_T_kpp" => test_T_kpp,

        "truth_uw" => truth_uw,
        "truth_vw" => truth_vw,
        "truth_wT" => truth_wT,

        "test_uw" => test_uw,
        "test_vw" => test_vw,
        "test_wT" => test_wT,

        "test_uw_modified_pacanowski_philander" => test_uw_mpp,
        "test_vw_modified_pacanowski_philander" => test_vw_mpp,
        "test_wT_modified_pacanowski_philander" => test_wT_mpp,

        "test_uw_kpp" => test_uw_kpp,
        "test_vw_kpp" => test_vw_kpp,
        "test_wT_kpp" => test_wT_kpp,

        "test_uw_NN_only" => test_uw_NN_only,
        "test_vw_NN_only" => test_vw_NN_only,
        "test_wT_NN_only" => test_wT_NN_only,

                                     "truth_Ri" => truth_Ri,
                                      "test_Ri" => test_Ri,
        "test_Ri_modified_pacanowski_philander" => test_Ri_modified_pacanowski_philander,
                                  "test_Ri_kpp" => test_Ri_kpp,

           "u_losses" => scaled_losses.u,
           "v_losses" => scaled_losses.v,
           "T_losses" => scaled_losses.T,
        "∂u∂z_losses" => scaled_losses.∂u∂z,
        "∂v∂z_losses" => scaled_losses.∂v∂z,
        "∂T∂z_losses" => scaled_losses.∂T∂z,

           "u_losses_modified_pacanowski_philander" => scaled_losses_mpp.u,
           "v_losses_modified_pacanowski_philander" => scaled_losses_mpp.v,
           "T_losses_modified_pacanowski_philander" => scaled_losses_mpp.T,
        "∂u∂z_losses_modified_pacanowski_philander" => scaled_losses_mpp.∂u∂z,
        "∂v∂z_losses_modified_pacanowski_philander" => scaled_losses_mpp.∂v∂z,
        "∂T∂z_losses_modified_pacanowski_philander" => scaled_losses_mpp.∂T∂z,

                                               "losses" => scaled_losses.u .+ scaled_losses.v .+ scaled_losses.T,
                                                 "loss" => profile_loss,
                                      "losses_gradient" => scaled_losses.∂u∂z .+ scaled_losses.∂v∂z .+ scaled_losses.∂T∂z,
                                        "loss_gradient" => loss_gradient,
                 "losses_modified_pacanowski_philander" => scaled_losses_mpp.u .+ scaled_losses_mpp.v .+ scaled_losses_mpp.T,
                   "loss_modified_pacanowski_philander" => profile_loss_mpp,
        "losses_modified_pacanowski_philander_gradient" => scaled_losses_mpp.∂u∂z .+ scaled_losses_mpp.∂v∂z .+ scaled_losses_mpp.∂T∂z,
          "loss_modified_pacanowski_philander_gradient" => loss_gradient_mpp,

           "u_losses_kpp" => scaled_losses_kpp.u,
           "v_losses_kpp" => scaled_losses_kpp.v,
           "T_losses_kpp" => scaled_losses_kpp.T,
        "∂u∂z_losses_kpp" => scaled_losses_kpp.∂u∂z,
        "∂v∂z_losses_kpp" => scaled_losses_kpp.∂v∂z,
        "∂T∂z_losses_kpp" => scaled_losses_kpp.∂T∂z,

                 "losses_kpp" => scaled_losses_kpp.u .+ scaled_losses_kpp.v .+ scaled_losses_kpp.T,
                   "loss_kpp" => profile_loss_kpp,
        "losses_kpp_gradient" => scaled_losses_kpp.∂u∂z .+ scaled_losses_kpp.∂v∂z .+ scaled_losses_kpp.∂T∂z,
          "loss_kpp_gradient" => loss_gradient_kpp,
    )

    if OUTPUT_PATH !== ""
        @info "Writing file"
        jldopen(OUTPUT_PATH, "w") do file
            file["NDE_profile"] = output
        end
    end

    return output
end


text:
# shock capturing using Local-Discontinuous Galerkin

function calcShockCapturing(mesh::AbstractMesh, sbp::AbstractOperator,
                             eqn::EulerData, opts,
                             capture::LDGShockCapturing{Tsol, Tres},
                             shockmesh::ShockedElements) where {Tsol, Tres}

  # shockmesh should be updated before this function is called
  #TODO: need to get entropy variables conversion function from somewhere

  entropy_vars = capture.entropy_vars
  flux = capture.flux
  diffusion = capture.diffusion

  computeEntropyVariables(mesh, eqn, capture, shockmesh, entropy_vars)

  computeThetaVolumeContribution(mesh, sbp, eqn, opts, capture, shockmesh)
  #TODO: need to get flux_func from somewhere
  computeThetaFaceContribution(mesh, sbp, eqn, opts, capture, shockmesh, flux)
  # this needs to apply Minv too
  # TODO: need to get diffusion func that computes matrix-vector products
  #       with the Cji matrices
  computeQFromTheta(mesh, sbp, eqn, opts, capture, shockmesh, diffusion)

  # update the residual
  computeQVolumeTerm(mesh, sbp, eqn, opts, capture, shockmesh)
  # TODO: need to get flux_func_q from somewhere
  computeQFaceTerm(mesh, sbp, eqn, opts, capture, shockmesh, flux)

  return nothing
end

"""
  Gets the entropy variables at the volume nodes of the elements that have
  a shock in them and their neighbors.

  **Inputs**

   * eqn

  **Inputs/Outputs**

   * capture: an [`LDGShockCapturing`](@ref) object, `capture.w_el` is
              overwritten (and possibly re-allocated)
"""
function computeEntropyVariables(mesh, eqn, capture::LDGShockCapturing{Tsol, Tres},
                                 shockmesh::ShockedElements,
                                 entropy_vars::AbstractVariables
                                ) where {Tsol, Tres}

  @simd for i=1:shockmesh.numEl
    i_full = shockmesh.elnums_all[i]
    @simd for j=1:mesh.numNodesPerElement
      q_i = sview(eqn.q, :, j, i_full)
      w_i = sview(capture.w_el, :, j, i)
      convertToEntropy(entropy_vars, eqn.params, q_i, w_i)
    end
  end

  return nothing
end

#------------------------------------------------------------------------------
# Volume contribution to Theta

function computeThetaVolumeContribution(mesh, sbp, eqn, opts,
                capture::LDGShockCapturing{Tsol, Tres},
                shockmesh::ShockedElements)  where {Tsol, Tres}

  # q_j stores theta_j in this function
  fill!(capture.q_j, 0)

  # temporary storage
  wxi_i = zeros(Tres, mesh.numDofPerNode, mesh.numNodesPerElement, mesh.dim)
  op = SummationByParts.Subtract()

  for i=1:shockmesh.numShock
    i_full = shockmesh.elnums_all[i]
    w_i = ro_sview(capture.w_el, :, :, i)

    # compute Qx^T w
    dxidx_i = ro_sview(mesh.dxidx, :, :, :, i_full)
    theta_i = sview(capture.q_j, :, :, :, i)
    applyQxTransposed(sbp, w_i, dxidx_i,  wxi_i, theta_i, op)
  end

  return nothing
end

#------------------------------------------------------------------------------
# Face contribution to Theta

function computeThetaFaceContribution(mesh::AbstractMesh{Tmsh}, sbp, eqn, opts,
          capture::LDGShockCapturing{Tsol, Tres}, shockmesh::ShockedElements,
          flux::AbstractLDGFlux) where {Tsol, Tres, Tmsh}

  w_faceL = zeros(Tsol, mesh.numDofPerNode, mesh.numNodesPerFace)
  w_faceR = zeros(Tsol, mesh.numDofPerNode, mesh.numNodesPerFace)
  flux_j = zeros(Tres, mesh.numDofPerNode)
  fluxD = zeros(Tres, mesh.numDofPerNode, mesh.numNodesPerFace, mesh.dim)
  nrm = zeros(Tmsh, mesh.dim)
  aux_vars = Tres[]
  for i=1:shockmesh.numInterfaces
    iface_i = shockmesh.ifaces[i]
    iface_red = iface_i.iface  # iface in reduced numbering scheme
    iface_idx = shockmesh.ifaces[i].idx_orig
    wL = sview(capture.w_el, :, :, iface_red.elementL)
    wR = sview(capture.w_el, :, :, iface_red.elementR)

    # interpolate v to faces
    interiorFaceInterpolate!(mesh.sbpface, iface_red, wL, wR, w_faceL,
                             w_faceR)

    # apply flux function
    for j=1:mesh.numNodesPerFace
      wL_j = ro_sview(w_faceL, :, j)
      wR_j = ro_sview(w_faceR, :, j)

      # get the normalized unit vector
      for d=1:mesh.dim
        nrm[d] = mesh.nrm_face[d, j, iface_idx]
      end
      normalize_vec(eqn.params, nrm)

      applyFlux(eqn.params, flux, wL_j, wR_j, aux_vars, nrm, flux_j)

      # apply normal vector components here because SBP only applied R^T B
      for d=1:mesh.dim
        for k=1:mesh.numDofPerNode
          fluxD[k, j, d] = mesh.nrm_face[d, j, iface_idx]*flux_j[k]
        end
      end
    end  # end j

    # apply R^T B
    for d=1:mesh.dim
      flux_d = ro_sview(fluxD, :, :, d)
      resL = sview(capture.q_j, :, :, d, iface_red.elementL)
      resR = sview(capture.q_j, :, :, d, iface_red.elementR)
      interiorFaceIntegrate!(mesh.sbpface, iface_red, flux_d, resL, resR)
    end

  end  # end i

  return nothing
end

#------------------------------------------------------------------------------
# Computing q_j

function computeQFromTheta(mesh, sbp, eqn, opts,
                           capture::LDGShockCapturing{Tsol, Tres},
                           shockmesh::ShockedElements, diffusion_func) where {Tsol, Tres}

  # epsilon = 0 for neighbor elements, so q = 0 there, only do shocked elements
  theta_i = zeros(Tres, mesh.numDofPerNode, mesh.numNodesPerElement, mesh.dim)
  for i=1:shockmesh.numShock
    i_full = shockmesh.elnums_all[i]

    # this copy is unfortunate, but the alternative is allocating 2 q_j arrays.
    # also apply the inverse mass matrix
    @simd for d=1:mesh.dim
      @simd for j=1:mesh.numNodesPerElement
        val = mesh.jac[j, i_full]/sbp.w[j]
        @simd for k=1:mesh.numDofPerNode
          theta_i[k, j, d] = val*capture.q_j[k, j, d, i]
        end
      end
    end

    w_i = ro_sview(capture.w_el, :, :, i)
    q_i = sview(capture.q_j, :, :, :, i)
    applyDiffusionTensor(diffusion_func, w_i, i, theta_i, q_i)
  end

  # zero out the rest of q_i because epsilon is zero there
  q_neighbor = sview(capture.q_j, :, :, :, (shockmesh.numShock+1):shockmesh.numEl)
  fill!(q_neighbor, 0)
  return nothing
end

#------------------------------------------------------------------------------
# volume contribution to residual

function computeQVolumeTerm(mesh, sbp, eqn, opts, capture::LDGShockCapturing{Tsol, Tres},
                            shockmesh::ShockedElements) where {Tsol, Tres}

  op = SummationByParts.Subtract()
  qxi_i = zeros(Tres, mesh.numDofPerNode, mesh.numNodesPerElement, mesh.dim)
  for i=1:shockmesh.numShock
    i_full = shockmesh.elnums_all[i]
    # compute Qx^T q_x
    q_i = ro_sview(capture.q_j, :, :, :, i)
    dxidx_i = ro_sview(mesh.dxidx, :, :, :, i_full)
    res_i = sview(eqn.res, :, :, i_full)
    applyQxTransposed(sbp, q_i, dxidx_i,  qxi_i, res_i, op)
  end

  return nothing
end

#------------------------------------------------------------------------------
# face contribution to res

function computeQFaceTerm(mesh::AbstractMesh{Tmsh}, sbp, eqn, opts,
                          capture::LDGShockCapturing{Tsol, Tres},
                          shockmesh::ShockedElements, flux::AbstractLDGFlux
                         ) where {Tsol, Tres, Tmsh}

  w_faceL = zeros(Tsol, mesh.numDofPerNode, mesh.numNodesPerFace)
  w_faceR = zeros(Tsol, mesh.numDofPerNode, mesh.numNodesPerFace)
  q_faceL = zeros(Tres, mesh.numDofPerNode, mesh.numNodesPerFace, mesh.dim)
  q_faceR = zeros(Tres, mesh.numDofPerNode, mesh.numNodesPerFace, mesh.dim)
  qL_j = zeros(Tres, mesh.numDofPerNode, mesh.dim)
  qR_j = zeros(Tres, mesh.numDofPerNode, mesh.dim)
  fluxD = zeros(Tres, mesh.numDofPerNode, mesh.dim)
  flux_scaled = zeros(Tres, mesh.numDofPerNode, mesh.numNodesPerFace)
  aux_vars = Tres[]
  nrm = zeros(Tmsh, mesh.dim)
  for i=1:shockmesh.numInterfaces
    iface_red = shockmesh.ifaces[i].iface
    idx_orig = shockmesh.ifaces[i].idx_orig

    # interpolate v to face
    wL = ro_sview(capture.w_el, :, :, iface_red.elementL)
    wR = ro_sview(capture.w_el, :, :, iface_red.elementR)
    interiorFaceInterpolate!(mesh.sbpface, iface_red, wL, wR, w_faceL,
                             w_faceR)

    # interpolate q_j to face
    for d=1:mesh.dim
      #TODO: make sure q_j has enough space, and is zeroed out, for all the
      #      neighbor elements
      qL = ro_sview(capture.q_j, :, :, d, iface_red.elementL)
      qR = ro_sview(capture.q_j, :, :, d, iface_red.elementR)
      q_faceL_d = sview(q_faceL, :, :, d)
      q_faceR_d = sview(q_faceR, :, :, d)
      interiorFaceInterpolate!(mesh.sbpface, iface_red, qL, qR, q_faceL_d,
                               q_faceR_d)
    end

    # apply flux function
    fill!(flux_scaled, 0)
    for j=1:mesh.numNodesPerFace
      wL_j = ro_sview(w_faceL, :, j)
      wR_j = ro_sview(w_faceR, :, j)
      # get q_j in all directions at this node
      # TODO: make type-stable strided view
      @simd for d=1:mesh.dim
        @simd for k=1:mesh.numDofPerNode
          qL_j[k, d] = q_faceL[k, j, d]
          qR_j[k, d] = q_faceR[k, j, d]
        end
      end

      # get the normalized unit vector
      for d=1:mesh.dim
        nrm[d] = mesh.nrm_face[d, j, idx_orig]
      end
      normalize_vec(eqn.params, nrm)

      applyFlux(eqn.params, flux, qL_j, qR_j, wL_j, wR_j, aux_vars, nrm, fluxD)

      # multiply by N because SBP doesn't, also do sum over dimensions
      @simd for d=1:mesh.dim
        N_d = mesh.nrm_face[d, j, idx_orig]
        @simd for k=1:mesh.numDofPerNode
          flux_scaled[k, j] += N_d*fluxD[k, d]
        end
      end

    end  # end j

    # apply R^T B
    resL = sview(eqn.res, :, :, shockmesh.elnums_all[iface_red.elementL])
    resR = sview(eqn.res, :, :, shockmesh.elnums_all[iface_red.elementR])
    interiorFaceIntegrate!(mesh.sbpface, iface_red, flux_scaled, resL, resR)
  end  # end i

  return nothing
end

#------------------------------------------------------------------------------
# Flux functions

"""
  Flux function for w, used to compute theta for LDG

  Note that this is different from a regular 2 point flux function becaues
  it operates on the entropy variables.
"""
function applyFlux(params::ParamType, obj::LDG_ESFlux, wL::AbstractVector,
                   wR::AbstractVector,
                   aux_vars, nrm::AbstractVector{Tmsh}, flux::AbstractVector
                  ) where {Tmsh}

  dim = length(nrm)

  # compute the factor in front of the jump term
  fac = zero(Tmsh)
  for d=1:dim
    fac += nrm[d]
  end
  fac *= obj.beta

  for i=1:size(flux, 1)
    flux[i] = 0.5*(wL[i] + wR[i]) + fac*(wL[i] - wR[i])
  end

  return nothing
end

"""
  This method computes the numerical flux function for q, which is used to
  update the residual.  This function computes the flux in all three
  directions at the same time.
"""
function applyFlux(params::ParamType, flux::LDG_ESFlux, qL::AbstractMatrix, qR::AbstractMatrix,
                   wL::AbstractVector, wR::AbstractVector, aux_vars,
                   nrm::AbstractVector, fluxD::AbstractMatrix{Tres}) where {Tres}

  numDofPerNode, dim = size(qL)

  @simd for d1=1:dim
    # do average and jumpt in v terms
    @simd for k=1:numDofPerNode
      fluxD[k, d1] = 0.5*(qL[k, d1] + qR[k, d1]) - flux.alpha*nrm[d1]*(wL[k] - wR[k])
    end

    # do the jump in Q term (which has unusual structure)
    @simd for d2=1:dim
      @simd for k=1:numDofPerNode
        fluxD[k, d1] -= flux.beta*nrm[d2]*(qL[k, d2] - qR[k, d2])
      end
    end
  end

  return nothing
end

ext: jl


text:
struct ZeroSource
end

function calcSource(source::ZeroSource,
                    state::HydroSnapshot,
                    t::Float64,
                    dt::Float64)::Array{Conserved,1}
    return [Conserved(0,0,0) for x in state.cells]
end

max_stars_repo_name: bolverk/Sirion

max_issues_repo_name: bolverk/Sirion

max_forks_repo_name: bolverk/Sirion


text:
module Wandb

using PyCall, Conda, Random, Requires

const wandb = PyNULL()

function __init__()
    copy!(wandb, pyimport("wandb"))
    @info "Using Wandb version $(wandb.__version__)"

    msg = Wandb.wandb.sdk.internal.update.check_available(wandb.__version__)
    if msg !== nothing
        @info "Trying to update Wandb to the latest stable release"
        run(`$(PyCall.pyprogramname) -m pip install wandb --upgrade`)
    end

    @require FluxTraining = "7bf95e4d-ca32-48da-9824-f0dc5310474f" begin
        include("fluxtraining.jl")
    end

    @require MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195" begin
        include("mpi.jl")
    end
end

using Base.CoreLogging: CoreLogging, AbstractLogger, LogLevel, Info, handle_message, shouldlog, min_enabled_level,
                        catch_exceptions

# Base functions like `log`, `Image`, etc.
include("main.jl")
# Wandb Artifacts API
include("artifacts.jl")
# AbstractLogger interface
include("corelogging.jl")
# HyperParameter Tuning: Sweep/Agent API
include("sweep.jl")

export WandbLogger, WandbArtifact, WandbHyperParameterSweep, update_config!, get_config, save

end


text:
################################################################################
#
#       GrpAb/Map.jl : Functions for maps between
#                      finitely generated abelian groups
#
# This file is part of Hecke.
#
# Copyright (c) 2015, 2016, 2017: Claus Fieker, Tommy Hofmann
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
# * Redistributions of source code must retain the above copyright notice, this
#   list of conditions and the following disclaimer.
#
# * Redistributions in binary form must reproduce the above copyright notice,
#   this list of conditions and the following disclaimer in the documentation
#   and/or other materials provided with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#
#
#  Copyright (C) 2015, 2016, 2017 Tommy Hofmann, Claus Fieker
#
################################################################################

export haspreimage, hasimage, hom, kernel, image, isinjective, issurjective,
       isbijective

################################################################################
#
#  Check for existence of (pre)image
#
################################################################################

doc"""
***
    haspreimage(M::GrpAbFinGenMap, a::GrpAbFinGenElem) -> Bool, GrpAbFinGenElem

> Returns whether $a$ is in the image of $M$. If so, the second return value is
> an element $b$ with $M(b) = a$.
"""
function haspreimage(M::GrpAbFinGenMap, a::GrpAbFinGenElem)
  if isdefined(M, :imap)
    return true, preimage(M, a)
  end

  m = vcat(M.map, rels(codomain(M)))
  fl, p = cansolve(m', a.coeff')
  if fl
    return true, domain(M)(sub(p', 1:1, 1:ngens(domain(M))))
  else
    return false, domain(M)[1]
  end
end

# Note that a map can be a partial function. The following function
# checks if an element is in the domain of definition.
#
# Here is an example:
# S, mS = sub(...), so ms: S -> G
# h = inv(mS)
# Now h is a partial function on G with domain of definition the image of mS.
# Then hasimage(h, x) would check if x is in the image of mS.
function hasimage(M::GrpAbFinGenMap, a::GrpAbFinGenElem)
  if isdefined(M, :map)
    return true, image(M, a)
  end

  m = vcat(M.imap, rels(domain(M)))
  fl, p = cansolve(m, a.coeff)
  if fl
    return true, codomain(M)(sub(p, 1:1, 1:ngens(codomain(M))))
  else
    return false, codomain(M)[1]
  end
end

################################################################################
#
#  Homomorphisms specified on arbitrary sets of generators
#
################################################################################

#TODO: Should check consistency
doc"""
    hom(A::Array{GrpAbFinGenElem, 1}, B::Array{GrpAbFinGenElem, 1}) -> Map
> Creates the homomorphism $A[i] \mapsto B[i]$
"""
function hom(A::Array{GrpAbFinGenElem, 1}, B::Array{GrpAbFinGenElem, 1}; check::Bool = false)
  GA = parent(A[1])
  GB = parent(B[1])
  @assert length(B)==length(A)
  @assert length(A)>0
  if (check)
    m = vcat([x.coeff for x in A])
    m = vcat(m, rels(parent(A[1])))
    T, i = nullspace(m')
    T = T'
    T = sub(T, 1:rows(T), 1:length(A))
    n = vcat([x.coeff for x in B])
    n = T*n
    if !cansolve(rels(parent(B[1]))', n')[1]
      error("Data does not define a homomorphism")
    end
  end

  M = vcat([hcat(A[i].coeff, B[i].coeff) for i = 1:length(A)])
  RA = rels(GA)
  M = vcat(M, hcat(RA, zero_matrix(FlintZZ, rows(RA), cols(B[1].coeff))))
  H = hnf(M)
  H = sub(H, 1:ngens(GA), ngens(GA)+1:ngens(GA)+ngens(GB))
  h = GrpAbFinGenMap(GA, GB, H)
  return h
end

doc"""
    hom(G::GrpAbFinGen, B::Array{GrpAbFinGenElem, 1}) -> Map

> Creates the homomorphism which maps `G[i]` to `B[i]`.
"""
function hom(G::GrpAbFinGen, B::Array{GrpAbFinGenElem, 1})
  GB = parent(B[1])
  M = vcat([B[i].coeff for i = 1:length(B)])
  h = GrpAbFinGenMap(G, GB, M)
  return h
end

# TODO: Extend the check to non-endomorphisms
function hom(A::GrpAbFinGen, B::GrpAbFinGen, M::fmpz_mat, check::Bool = true)
  if check #totally wrong if A != B
    if A == B
      G = A
      images = [ G([M[i, j] for j in 1:ngens(G)]) for i in 1:ngens(G) ]
      a = 0 * G[1]
      for i in 1:nrels(G)
        for j in 1:ngens(G)
          a = a + rels(G)[i, j] * images[j]
        end
        if !iszero(a)
          error("Matrix does not define a morphism of abelian groups")
        end
      end
    end
  end

  return GrpAbFinGenMap(A, B, M)
end

function hom(A::GrpAbFinGen, B::GrpAbFinGen, M::fmpz_mat, Minv, check::Bool = true)
  if check
    if A == B
      G = A
      images = [ G([M[i, j] for j in 1:ngens(G)]) for i in 1:ngens(G) ]
      a = 0 * G[1]
      for i in 1:nrels(G)
        for j in 1:ngens(G)
          a = a + G.rels[i, j] * images[j]
        end
        if !iszero(a)
          error("Matrix does not define a morphism of abelian groups")
        end
      end
    end
  end

  return GrpAbFinGenMap(A, B, M, Minv)
end

################################################################################
#
#  Kernel, image, cokernel
#
################################################################################

#TODO: store and reuse on map. Maybe need to change map

doc"""
    kernel(h::GrpAbFinGenMap) -> GrpAbFinGen, Map

Let $G$ be the domain of $h$. This functions returns an abelian group $A$ and an
injective morphism $f \colon A \to G$, such that the image of $f$ is the kernel
of $h$.
"""

function kernel(h::GrpAbFinGenMap)
  G = domain(h)
  H = codomain(h)
  m=zero_matrix(FlintZZ, rows(h.map)+rows(rels(H)), cols(h.map))
  for i=1:rows(h.map)
    for j=1:cols(h.map)
      m[i,j]=h.map[i,j]
    end
  end
  if !issnf(H)
    for i=1:nrels(H)
      for j=1:ngens(H)
        m[rows(h.map)+i,j]=H.rels[i,j]
      end
    end
  else
    for i=1:length(H.snf)
      m[rows(h.map)+i,i]=H.snf[i]
    end
  end
  hn, t = hnf_with_transform(m)
  for i = 1:rows(hn)
    if iszero_row(hn, i)
      return sub(G, sub(t, i:rows(t), 1:ngens(G)), false)
    end
  end
  error("Something went terribly wrong in kernel computation")
end

doc"""
    image(h::GrpAbFinGenMap) -> GrpAbFinGen, Map

Let $G$ be the codomain of $h$. This functions returns an abelian group $A$ and
an injective morphism $f \colon A \to G$, such that the image of $f$ is the
image of $h$.
"""
function image(h::GrpAbFinGenMap)
  G = domain(h)
  H = codomain(h)
  hn = hnf(vcat(h.map, rels(H)))
  im = GrpAbFinGenElem[]
  for i = 1:rows(hn)
    if !iszero_row(hn, i)
      push!(im, H(sub(hn, i:i, 1:ngens(H))))
    else
      break
    end
  end
  return sub(H, im, false)  # too much, this is sub in hnf....
end

doc"""
    cokernel(h::GrpAbFinGenMap) -> GrpAbFinGen, Map

Let $G$ be the codomain of $h$. This functions returns an abelian group $A$ and
a morphism $f \colon G \to A$, such that $A$ is the quotient of $G$ with
respect to the image of $h$.
"""
function cokernel(h::GrpAbFinGenMap)
  S,mS=image(h)
  return quo(codomain(h), GrpAbFinGenElem[mS(g) for g in gens(S)], false)
end

################################################################################
#
#  Surjectivity
#
################################################################################

doc"""
    issurjective(h::GrpAbFinGenMap) -> Bool

Returns whether $h$ is surjective.
"""
function issurjective(A::GrpAbFinGenMap)
  if isfinite(codomain(A))
    H, mH = image(A)
    return (order(codomain(A)) == order(H))::Bool
  else
    return (order(cokernel(A)) == 1)::Bool
  end
end

################################################################################
#
#  Injectivity
#
################################################################################

doc"""
    isinjective(h::GrpAbFinGenMap) -> Bool

Returns whether $h$ is injective.
"""
function isinjective(A::GrpAbFinGenMap)
  K = kernel(A)[1]
  return isfinite(K) && isone(order(K))
end

################################################################################
#
#  Bijectivity
#
################################################################################

doc"""
    isbijective(h::GrpAbFinGenMap) -> Bool

Returns whether $h$ is bijective.
"""
function isbijective(A::GrpAbFinGenMap)
  return isinjective(A) && issurjective(A)
end

###############################################################################
#
#  Compose and Squash for abelian group maps
#
###############################################################################

function _compose(f::GrpAbFinGenMap, g::GrpAbFinGenMap)
  @assert domain(g)==codomain(f)

  M=f.map*g.map
  return GrpAbFinGenMap(domain(f), codomain(g), M)

end


text:
"""
    methodlearner(method, data, model[; kwargs...]) -> Learner

Create a `Learner` to train a model for learning method `method` using
`data`. See also [`Learner`](#).

## Keyword arguments

- `isbackbone = true`: Whether `model` is a backbone or a complete model. If `true`,
    a model will be constructed using [`methodmodel`](#)`(method, model)`.
- `validdata = nothing`: Validation data container. If none is given, `data` will be
    randomly split into training and validation data.
- `optimizer = ADAM()`: Optimizer passed to `Learner`.
- `lossfn = `[`methodlossfn`](#)`(method)`: Loss function passed to `Learner`.

Any other keyword arguments will be passed to [`methoddataloaders`](#).
"""
function methodlearner(
        method::LearningMethod,
        data,
        backbone,
        callbacks...;
        isbackbone=true,
        pctgval=0.2,
        batchsize=16,
        validdata=nothing,
        validbsfactor=2,
        optimizer=ADAM(),
        lossfn=methodlossfn(method),
        dlkwargs=(;),
    )
    model = isbackbone ? methodmodel(method, backbone) : backbone
    dls = if isnothing(validdata)
        methoddataloaders(data, method, batchsize;
            validbsfactor=validbsfactor, pctgval=pctgval, dlkwargs...)
    else
        methoddataloaders(data, validdata, method, batchsize;
            validbsfactor=validbsfactor, dlkwargs...)
    end
    return Learner(model, dls, optimizer, lossfn, callbacks...)
end


text:
export disassemble
export load

"""
    load(filename) -> Vector{Int}
Loads machine code file and return as an integer array
"""
function load(filename::AbstractString)
    parse.(Int, filter(!isempty, readlines(filename)))
end

"""
    disassemble(filename::AbstractString)
Print out the disassembled contents of `filename`. We assume the file
contains Little Man Computer numerical codes (3 digits each).
"""
function disassemble(filename::AbstractString)
    lines = readlines(filename)
    disassemble(parse.(Int, lines))
end

"""
    disassemble(instruction::Integer)
Disassemble single instruction and produce the mnemonic and operand
"""
function disassemble(instruction::Integer)
    if instruction == 0
        "HLT"
    elseif instruction < 100
        "DAT $code"
    elseif instruction > 900
        mnemonics[instruction]
    else
        opcode   = div(instruction, 100) * 100
        mnemonic = mnemonics[opcode]
        operand  = rem(instruction, 100)

        string(mnemonic, " ", operand)
    end
end

"""
    disassemble(program::Vector{Int})
Disassemble array of instructions for Little Man Computer.
"""
function disassemble(program::Vector{<:Integer})
    lines = map(disassemble, program)

    for (i, line) in enumerate(lines)
       println(lpad(i-1, 3), " ", line)
    end
end


text:
using PersistenceDiagrams
using Aqua

Aqua.test_all(
    PersistenceDiagrams;
    ambiguities=(; exclude=[in]), # from MLJModelInterface
)


text:
module PencilArrays

using MPI
using OffsetArrays
using Reexport
using StaticPermutations
using TimerOutputs

import Base: @propagate_inbounds
import LinearAlgebra

include("Pencils/Pencils.jl")
@reexport using .Pencils

import .Pencils:
    get_comm,
    permutation,
    range_local,
    range_remote,
    size_local,
    size_global,
    topology

export PencilArray, GlobalPencilArray, PencilArrayCollection, ManyPencilArray
export pencil
export gather
export global_view
export ndims_extra, ndims_space, extra_dims, sizeof_global

# Type definitions
include("arrays.jl")       # PencilArray
include("multiarrays.jl")  # ManyPencilArray
include("global_view.jl")  # GlobalPencilArray
include("cartesian_indices.jl")  # PermutedLinearIndices, PermutedCartesianIndices

include("array_interface.jl")
include("broadcast.jl")
include("gather.jl")

include("Transpositions/Transpositions.jl")
export Transpositions

include("PencilIO/PencilIO.jl")

end


text:
import LinearAlgebra.mul!

function add!(zd1::VTUFile, zd2::VTUFile)
	add!(zd1.data,zd2.data)
	return nothing
end

function minus!(zd1::VTUFile, zd2::VTUFile)
	minus!(zd1.data,zd2.data)
	return nothing
end

function add!(zd1::VTUFile, a::Number)
	add!(zd1.data, a)
	return nothing
end

function mul!(zd1::VTUFile, c::Number)
	mul!(zd1.data, c)
	return nothing
end

function mul!(zd1::VTUFile, zd2::VTUFile, zd3::VTUFile)
	mul!(zd1.data, zd2.data, zd3.data)
	return nothing
end

function mul!(zd1::VTUFile, zd2::VTUFile, fac::Float64)
	mul!(zd1.data, zd2.data, fac)
	return nothing
end

function pow!(zd1::VTUFile, a::Number)
	pow!(zd1.data, a)
	return nothing
end

function div!(zd1::VTUFile, a::Number)
	div!(zd1.data, a)
	return nothing
end

function div!(zd1::VTUFile, zd2::VTUFile, zd3::VTUFile)
	div!(zd1.data, zd2.data, zd3.data)
	return nothing
end

function max!(zd1::VTUFile, zd2::VTUFile)
	max!(zd1.data, zd2.data)
	return nothing
end

function min!(zd1::VTUFile, zd2::VTUFile)
	min!(zd1.data, zd2.data)
	return nothing
end

import Base: +,-,*,/,^

function +(tpf1::VTUFile, tpf2::VTUFile)
	ret = similar(tpf1)
	add!(ret,tpf1)
	add!(ret,tpf2)
	return ret
end

function -(tpf1::VTUFile, tpf2::VTUFile)
	ret = similar(tpf1)
	add!(ret,tpf1)
	minus!(ret,tpf2)
	return ret
end

function *(tpf1::VTUFile, tpf2::VTUFile)
	ret = similar(tpf1)
	mul!(ret,tpf1,tpf2)
	return ret
end

function +(tpf::VTUFile, a::Number)
	ret = similar(tpf)
	add!(ret,tpf)
	add!(ret,a)
	return ret
end
+(a::Number,tpf::VTUFile) = tpf+a

function -(tpf::VTUFile, a::T) where T<:Number
	ret = similar(tpf)
	add!(ret,tpf)
	add!(ret,-a)
	return ret
end

function -(a::T,tpf::VTUFile) where T<:Number
	ret = similar(tpf)
	add!(ret,a)
	minus!(ret,tpf)
	return ret
end

function *(tpf::VTUFile, a::T) where T<:Number
	ret = similar(tpf)
	add!(ret,tpf)
	mul!(ret,a)
	return ret
end
*(a::T,tpf::VTUFile) where T<:Number = tpf*a

function /(tpf::VTUFile, a::T) where T<:Number
	ret = similar(tpf)
	add!(ret,tpf)
	div!(ret,a)
	return ret
end

function /(tpf1::VTUFile, tpf2::VTUFile)
	ret = similar(tpf1)
	div!(ret,tpf1,tpf2)
	return ret
end

function ^(tpf::VTUFile, a::T) where T<:Number
	ret = similar(tpf)
	add!(ret,tpf)
	pow!(ret,a)
	return ret
end

max_stars_repo_name: nagelt/VTUFileHandler

max_issues_repo_name: nagelt/VTUFileHandler

max_forks_repo_name: nagelt/VTUFileHandler


text:
@recipe function f(pomdp::UAVchasePOMDP, h::AbstractPOMDPHistory{MDPState})
    mdp = pomdp.mdp
    ratio --> :equal
    # xlabel := "x (m)"
    # ylabel := "y (m)"
    # zlabel := "z (m)"
    xlim --> (-1, 12)
    ylim --> (-1, 10)
    zlim --> (0, 5)
    @series begin
        seriestype := :path
        label := "UAV path"
        linewidth := 2
        x = [s.uavPose[1] for s in state_hist(h)[1:end-1]]
        y = [s.uavPose[2] for s in state_hist(h)[1:end-1]]
        z = [s.uavPose[3] for s in state_hist(h)[1:end-1]]
        # color --> :orange
        color := :orange
        x, y, z
    end
    @series begin
        seriestype := :scatter
        label := "current UAV"
        pos = state_hist(h)[end-1].uavPose
        color := :orange
        [pos[1]], [pos[2]], [pos[3]]
    end
    @series begin
        label := "belief"
        pomdp, belief_hist(h)[end]
    end
    @series begin
        seriestype := :path
        label := "target path"
        linewidth := 2
        x = [s.targetPose[1] for s in state_hist(h)[1:end-1]]
        y = [s.targetPose[2] for s in state_hist(h)[1:end-1]]
        z = [mdp.target_height for s in state_hist(h)[1:end-1]]
        # color --> :green
        color := :green
        x, y, z
    end
    @series begin
        seriestype := :path
        label := "true target path"
        linewidth := 2
        x = [s.targetPose_true[1] for s in state_hist(h)[1:end-1]]
        y = [s.targetPose_true[2] for s in state_hist(h)[1:end-1]]
        z = [mdp.target_height for s in state_hist(h)[1:end-1]]
        # color --> :blue
        color := :blue
        x, y, z
    end
    @series begin
        if length(h) == 1
            color := :red
        elseif observation_hist(h)[end-1].if_observed
            color := :green
        else
            color := :red
        end
        seriestype := :scatter
        label := "estimated target"
        pos = state_hist(h)[end-1].targetPose
        [pos[1]], [pos[2]], [mdp.target_height]
    end
    @series begin
        color := :blue
        seriestype := :scatter
        label := "true target"
        pos = state_hist(h)[end-1].targetPose_true
        [pos[1]], [pos[2]], [mdp.target_height]
    end
end

@recipe function f(pomdp::UAVchasePOMDP, pc::ParticleCollection{MDPState})
    mdp = pomdp.mdp
    seriestype := :scatter
    x = [p.targetPose[1] for p in particles(pc)]
    y = [p.targetPose[2] for p in particles(pc)]
    z = [mdp.target_height for i in 1:n_particles(pc)]
    markersize --> [10.0*sqrt(weight(pc,i)) for i in 1:n_particles(pc)]
    color --> :black
    # markersize --> 0.1
    x, y, z
end
# not for visual but for particle filter
# @recipe function f(s::MDPState)
#     [s.targetPose[1]], [s.targetPose[2]]
# end


text:
######### CmdStan program example  ###########

using CmdStan, Test, Statistics

ProjDir = dirname(@__FILE__)
cd(ProjDir) do

  bernoullimodel = "
  data {
    int<lower=1> N;
    int<lower=0,upper=1> y[N];
  }
  parameters {
    real<lower=0,upper=1> theta;
  }
  model {
    theta ~ beta(1,1);
    y ~ bernoulli(theta);
  }
  "

  bernoullidata = Dict("N" => 10, "y" => [0, 1, 0, 1, 0, 0, 0, 0, 0, 1])
  inittheta = [Dict("theta" => 0.6), Dict("theta" => 0.4), Dict("theta" => 0.2), Dict("theta" => 0.1)]

  global stanmodel
  stanmodel = Stanmodel(name="bernoulli2", model=bernoullimodel,
    output_format=:array, num_warmup=1000, random=CmdStan.Random(seed=-1));

  global rc, sim
  rc, sim, cnames = stan(stanmodel, bernoullidata, ProjDir,
    init=inittheta, CmdStanDir=CMDSTAN_HOME, summary=false)

  if rc == 0
    println()
    println("Test 0.2 <= mean(theta[1]) <= 0.5)")
    @test 0.2 <= round.(mean(sim[:,8,:]), digits=1) <= 0.5
  end

end # cd