# AtmosMassEnergyLoss
#
# Dump mass and energy loss

using ..Atmos
using ..BalanceLaws
using ..Mesh.Topologies
using ..Mesh.Grids
using ..MPIStateArrays

"""
    setup_atmos_mass_energy_loss(
        ::ClimateMachineConfigType,
        interval::String,
        out_prefix::String,
        writer = NetCDFWriter(),
        interpol = nothing,
    )

Create and return a `DiagnosticsGroup` containing the
"AtmosMassEnergyLoss" diagnostics for Atmos LES and GCM
configurations. All the diagnostics in the group will run at the
specified `interval`, be interpolated to the specified boundaries
and resolution, and written to files prefixed by `out_prefix`
using `writer`.
"""
function setup_atmos_mass_energy_loss(
    ::ClimateMachineConfigType,
    interval::String,
    out_prefix::String,
    writer = NetCDFWriter(),
    interpol = nothing,
)
    @assert isnothing(interpol)

    return DiagnosticsGroup(
        "AtmosMassEnergyLoss",
        Diagnostics.atmos_mass_energy_loss_init,
        Diagnostics.atmos_mass_energy_loss_fini,
        Diagnostics.atmos_mass_energy_loss_collect,
        interval,
        out_prefix,
        writer,
        interpol,
    )
end

# store initial mass and energy
mutable struct AtmosMassEnergyLoss
    Σρ₀::Union{Nothing, AbstractFloat}
    Σρe₀::Union{Nothing, AbstractFloat}

    AtmosMassEnergyLoss() = new(nothing, nothing)
end
const AtmosMassEnergyLoss₀ = AtmosMassEnergyLoss()

function atmos_mass_energy_loss_init(dgngrp, currtime)
    mpicomm = Settings.mpicomm
    mpirank = MPI.Comm_rank(mpicomm)
    dg = Settings.dg
    bl = dg.balance_law
    Q = Settings.Q
    FT = eltype(Q)

    energy_model(bl) isa TotalEnergyModel ||
        error("Only TotalEnergyModel supported")
    ρ_idx = varsindices(vars_state(bl, Prognostic(), FT), "ρ")
    ρe_idx = varsindices(vars_state(bl, Prognostic(), FT), :(energy.ρe))
    Σρ₀ = weightedsum(Q, ρ_idx)
    Σρe₀ = weightedsum(Q, ρe_idx)

    if mpirank == 0
        AtmosMassEnergyLoss₀.Σρ₀ = Σρ₀
        AtmosMassEnergyLoss₀.Σρe₀ = Σρe₀

        # outputs are scalars
        dims = OrderedDict()

        # set up the variables we're going to be writing
        vars = OrderedDict(
            "mass_loss" => ((), FT, Variables["mass_loss"].attrib),
            "energy_loss" => ((), FT, Variables["energy_loss"].attrib),
        )

        # create the output file
        dprefix = @sprintf("%s_%s", dgngrp.out_prefix, dgngrp.name)
        dfilename = joinpath(Settings.output_dir, dprefix)
        noov = Settings.no_overwrite
        init_data(dgngrp.writer, dfilename, noov, dims, vars)
    end

    return nothing
end

function atmos_mass_energy_loss_collect(dgngrp, currtime)
    mpicomm = Settings.mpicomm
    mpirank = MPI.Comm_rank(mpicomm)
    dg = Settings.dg
    bl = dg.balance_law
    Q = Settings.Q
    FT = eltype(Q)

    energy_model(bl) isa TotalEnergyModel ||
        error("Only TotalEnergyModel supported")
    ρ_idx = varsindices(vars_state(bl, Prognostic(), FT), "ρ")
    ρe_idx = varsindices(vars_state(bl, Prognostic(), FT), :(energy.ρe))
    Σρ = weightedsum(Q, ρ_idx)
    Σρe = weightedsum(Q, ρe_idx)

    if mpirank == 0
        δρ = (Σρ - AtmosMassEnergyLoss₀.Σρ₀) / AtmosMassEnergyLoss₀.Σρ₀
        δρe = (Σρe - AtmosMassEnergyLoss₀.Σρe₀) / AtmosMassEnergyLoss₀.Σρe₀

        varvals = OrderedDict("mass_loss" => δρ, "energy_loss" => δρe)

        # write output
        append_data(dgngrp.writer, varvals, currtime)
    end

end

function atmos_mass_energy_loss_fini(dgngrp, currtime) end


fitting3K = DataFrame(fittingType = ["PIPE",
"ELL-THD-STD",
"ELL-THD-LR",
"ELL-FLG-RD1",
"ELL-FLG-RD2",
"ELL-FLG-RD4",
"ELL-FLG-RD6",
"ELL-MIT-1",
"ELL-MIT-2",
"EL-MIT-3",
"EL45-THD-STD",
"EL45-LR",
"EL45-MIT-1",
"EL45-MIT-2",
"EL180-THD-1",
"EL180-FLG-RD1",
"EL180-FLG-RD15",
"TEE-THD-RD1",
"TEE-THD-RD15",
"TEE-FLG",
"TEE-STUB",
"TEE-THD-THRU",
"TEE-FLG-THRU",
"TEE-STUB-THRU",
"VALVE-ANGLE-45-FULL",
"VALVE-ANGLE-90-FULL",
"VALVE-GLOBE-STD-B1",
"VALVE-PLUG-BRANCH",
"VALVE-PLUG-THRU",
"VALVE-PLUG-3WAY-THRU",
"VALVE-GATE-STD-B1",
"VALVE-BALL-STD-B1",
"VALVE-DIAPH-DAM",
"VALVE-CHECK-SWING",
"VALVE-CHECK-LIFT",
"ENTRANCE-NON-RECOVER",
"EXIT-NON-RECOVER"],
Description = ["PipePipePipe",
"Elbow - 90 DegreeThreaded-Standardr/D = 1",
"Elbow - 90 DegreeThreaded Long Radiusr/D = 1.5",
"Elbow - 90 DegreeFlanged Welded Bendr/D = 1",
"Elbow - 90 DegreeFlanged Welded Bendr/D = 2",
"Elbow - 90 DegreeFlanged Welded Bendr/D = 4",
"Elbow - 90 DegreeFlanged Welded Bendr/D = 6",
"Elbow - 90 DegreeMitered 1 Weld 90 Degree",
"Elbow - 90 DegreeMitered 2 Welds 45 Degree",
"Elbow - 90 DegreeMitered 3 Welds 30 Degree",
"Elbow - 45 DegreeThreaded Standardr/D = 1",
"Elbow - 45 DegreeLong Radiusr/D = 1.5",
"Elbow - 45 DegreeMitered 1 Weld 45 Degree",
"Elbow - 45 DegreeMitered 2 Welds 22.5 Degree",
"Elbows - 180 DegreeThreaded Close Return Bendr/D = 1",
"Elbows - 180 DegreeFlangedr/D = 1",
"Elbows - 180 DegreeAllr/D = 1.5",
"TEEThreadedr/D = 1",
"TEEThreadedr/D = 1.5",
"TEEFlangedr/D = 1",
"TEEStub-in Branch",
"TEERun Through Threadedr/D = 1",
"TEEFlangedr/D = 1",
"TEEStub-in Branch",
"ValveAngle - 45 DegreeFull Line Size b = 1",
"ValveAngle - 90 DegreeFull Line Size b = 1",
"ValveGlobeStandard b = 1",
"ValvePlugBranch Flow",
"ValvePlugStraight Through",
"ValvePlugThree-Way (flow through)",
"ValveGateStandard b = 1",
"ValveBallStandard b = 1",
"ValveDiaphragmDam-Type",
"ValveSwing CheckVmin = 35r-1/2",
"ValveLift CheckVmin = 40r-1/2",
"InletNon-recoverable",
"ExitNon-recoverable"],
Fitting = ["Pipe",
"Elbow - 90 Degree",
"Elbow - 90 Degree",
"Elbow - 90 Degree",
"Elbow - 90 Degree",
"Elbow - 90 Degree",
"Elbow - 90 Degree",
"Elbow - 90 Degree",
"Elbow - 90 Degree",
"Elbow - 90 Degree",
"Elbow - 45 Degree",
"Elbow - 45 Degree",
"Elbow - 45 Degree",
"Elbow - 45 Degree",
"Elbows - 180 Degree",
"Elbows - 180 Degree",
"Elbows - 180 Degree",
"TEE",
"TEE",
"TEE",
"TEE",
"TEE",
"TEE",
"TEE",
"Valve",
"Valve",
"Valve",
"Valve",
"Valve",
"Valve",
"Valve",
"Valve",
"Valve",
"Valve",
"Valve",
"Inlet",
"Exit"],
Type = ["Pipe",
"Threaded-Standard",
"Threaded Long Radius",
"Flanged Welded Bend",
"Flanged Welded Bend",
"Flanged Welded Bend",
"Flanged Welded Bend",
"Mitered 1 Weld 90 Degree",
"Mitered 2 Welds 45 Degree",
"Mitered 3 Welds 30 Degree",
"Threaded Standard",
"Long Radius",
"Mitered 1 Weld 45 Degree",
"Mitered 2 Welds 22.5 Degree",
"Threaded Close Return Bend",
"Flanged",
"All",
"Threaded",
"Threaded",
"Flanged",
"Stub-in Branch",
"Run Through Threaded",
"Flanged",
"Stub-in Branch",
"Angle - 45 Degree",
"Angle - 90 Degree",
"Globe",
"Plug",
"Plug",
"Plug",
"Gate",
"Ball",
"Diaphragm",
"Swing Check",
"Lift Check",
"",
""],
Geometry = ["Pipe",
"r/D = 1",
"r/D = 1.5",
"r/D = 1",
"r/D = 2",
"r/D = 4",
"r/D = 6",
"",
"",
"",
"r/D = 1",
"r/D = 1.5",
"",
"",
"r/D = 1",
"r/D = 1",
"r/D = 1.5",
"r/D = 1",
"r/D = 1.5",
"r/D = 1",
"",
"r/D = 1",
"r/D = 1",
"",
"Full Line Size b = 1",
"Full Line Size b = 1",
"Standard b = 1",
"Branch Flow",
"Straight Through",
"Three-Way (flow through)",
"Standard b = 1",
"Standard b = 1",
"Dam-Type",
"Vmin = 35r-1/2",
"Vmin = 40r-1/2",
"Non-recoverable",
"Non-recoverable"],
K1 = [0,
800,
800,
800,
800,
800,
800,
1000,
800,
800,
500,
500,
500,
500,
1000,
1000,
1000,
500,
800,
800,
1000,
200,
150,
100,
950,
1000,
1500,
500,
300,
300,
300,
300,
1000,
1500,
2000,
0,
0],
Kinf = [0,
0.14,
0.071,
0.091,
0.056,
0.066,
0.075,
0.27,
0.068,
0.035,
0.071,
0.052,
0.086,
0.052,
0.23,
0.12,
0.1,
0.274,
0.14,
0.28,
0.34,
0.091,
0.017,
0,
0.25,
0.69,
1.7,
0.41,
0.084,
0.14,
0.037,
0.017,
0.69,
0.46,
2.85,
1.5,
0],
Kd = [0,
4,
4.2,
4,
3.9,
3.9,
4.2,
4,
4.1,
4.2,
4.2,
4,
4,
4,
4,
4,
4,
4,
4,
4,
4,
4,
4,
0,
4,
4,
3.6,
4,
3.9,
4,
3.9,
4,
4.9,
4,
3.8,
0,
0],
Kp = [1,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0,
0] )


#-----------------------------------------------------------------------
# MetadataTools
# https://github.com/IainNZ/MetadataTools.jl
#-----------------------------------------------------------------------
# Copyright (c) 2015: Iain Dunning
# Licensed under the MIT License
#-----------------------------------------------------------------------
# graph_stats.jl
# Calculate some rankings for package connectedness
#----------------------------------------------------------------------

using MetadataTools

# Load in all the package data from METADATA
pkgmeta = get_all_pkg()

# Build a PkgGraph with edges A->B if A directly requires B
pg = make_dep_graph(pkgmeta)
total_pkgs = size(pg)

# Count size of every connected subgraph obtained by traversing the
# dependency graph starting from each package
pkg_req_counts = Any[]
for pkg in keys(pkgmeta)
    # Don't warn about cycles, we just want total counts anyway
    g = get_pkg_dep_graph(pkg, pg)
    req_count = size(g)
    push!(pkg_req_counts, (req_count,pkg))
end
sort!(pkg_req_counts, rev=true)  # Sort descending
print_with_color(:blue, "Top packages by number of packages depended\n")
print_with_color(:blue, "on, directly or indirectly:\n")
for (count, pkg) in pkg_req_counts[1:10]
    print(rpad(pkg,30," "))
    @printf("%4d / %4d\n", count-1, total_pkgs)
end
println()

# Do the same, but only counting immediate dependencies
pkg_req_counts = Any[]
for pkg in keys(pkgmeta)
    # Don't warn about cycles, we just want total counts anyway
    g = get_pkg_dep_graph(pkg, pg, depth=2)
    req_count = size(g)
    push!(pkg_req_counts, (req_count,pkg))
end
sort!(pkg_req_counts, rev=true)  # Sort descending
print_with_color(:blue, "Top packages by number of packages depended\n")
print_with_color(:blue, "on, directly:\n")
for (count, pkg) in pkg_req_counts[1:10]
    print(rpad(pkg,30," "))
    @printf("%4d / %4d\n", count-1, total_pkgs)
end
println()

# Build a graph with edges A->B if A is directly required by B
pgrev = make_dep_graph(pkgmeta, reverse=true)

# Count size of every subgraphs like above
pkg_req_counts = Any[]
for pkg in keys(pkgmeta)
    g = get_pkg_dep_graph(pkg, pgrev)
    req_count = size(g)
    push!(pkg_req_counts, (req_count,pkg))
end
sort!(pkg_req_counts, rev=true)  # Sort descending
print_with_color(:blue, "Top packages by number of packages that depend\n")
print_with_color(:blue, "on them, directly or indirectly:\n")
for (count, pkg) in pkg_req_counts[1:10]
    print(rpad(pkg,30," "))
    @printf("%4d / %4d\n", count-1, total_pkgs)
end
println()

# Same, but only direct dependencies
pkg_req_counts = Any[]
for pkg in keys(pkgmeta)
    g = get_pkg_dep_graph(pkg, pgrev, depth=2)
    req_count = size(g)
    push!(pkg_req_counts, (req_count,pkg))
end
sort!(pkg_req_counts, rev=true)  # Sort descending
print_with_color(:blue, "Top packages by number of packages that depend\n")
print_with_color(:blue, "on them, directly:\n")
for (count, pkg) in pkg_req_counts[1:10]
    print(rpad(pkg,30," "))
    @printf("%4d / %4d\n", count-1, total_pkgs)
end


function PowellOptimizer(arg0::jdouble, arg1::jdouble)
    return PowellOptimizer((jdouble, jdouble), arg0, arg1)
end

function PowellOptimizer(arg0::jdouble, arg1::jdouble, arg2::ConvergenceChecker)
    return PowellOptimizer((jdouble, jdouble, ConvergenceChecker), arg0, arg1, arg2)
end

function PowellOptimizer(arg0::jdouble, arg1::jdouble, arg2::jdouble, arg3::jdouble)
    return PowellOptimizer((jdouble, jdouble, jdouble, jdouble), arg0, arg1, arg2, arg3)
end

function PowellOptimizer(arg0::jdouble, arg1::jdouble, arg2::jdouble, arg3::jdouble, arg4::ConvergenceChecker)
    return PowellOptimizer((jdouble, jdouble, jdouble, jdouble, ConvergenceChecker), arg0, arg1, arg2, arg3, arg4)
end


module ITensorsVisualization

using GLMakie
using ITensors
using LinearAlgebra
using NetworkLayout
using SparseArrays

export @visualize

include("utils.jl")
include("visualize_tensornetwork.jl")
include("visualize_macro.jl")

end


using Oceananigans.Architectures: CPU
using Oceananigans.Fields: AbstractField, ReducedField

import Oceananigans.Fields: set!

const CubedSphereCPUField = CubedSphereField{X, Y, Z, <:CPU} where {X, Y, Z}
const CubedSphereGPUField = CubedSphereField{X, Y, Z, <:GPU} where {X, Y, Z}

const CubedSphereCPUReducedField = CubedSphereReducedField{X, Y, Z, <:CPU} where {X, Y, Z}
const CubedSphereGPUReducedField = CubedSphereReducedField{X, Y, Z, <:GPU} where {X, Y, Z}

const CubedSphereCPUFields = Union{CubedSphereCPUField, CubedSphereCPUReducedField}
const CubedSphereGPUFields = Union{CubedSphereGPUField, CubedSphereGPUReducedField}

# We need to define the function once for CPU fields then again for GPU fields to avoid the method
# ambiguity with Fields.set!.

function set!(u::CubedSphereCPUFields , v::CubedSphereCPUFields)
    for (u_face, v_face) in zip(faces(u), faces(v))
        @. u_face.data.parent = v_face.data.parent
    end
    return nothing
end

function set!(u::CubedSphereGPUField , v::CubedSphereGPUField)
    for (u_face, v_face) in zip(faces(u), faces(v))
        @. u_face.data.parent = v_face.data.parent
    end
    return nothing
end

set!(field::CubedSphereCPUFields, f::Function) = [set_face_field!(field_face, f) for field_face in faces(field)]
set!(field::CubedSphereCPUFields, f::Number) = [set_face_field!(field_face, f) for field_face in faces(field)]
set!(field::CubedSphereGPUFields, f::Function) = [set_face_field!(field_face, f) for field_face in faces(field)]
set!(field::CubedSphereGPUFields, f::Number) = [set_face_field!(field_face, f) for field_face in faces(field)]

set_face_field!(field, a) = set!(field, a)


"""
    struct Crane

Struct representing a rail-mounted gantry crane as modeled in the GCSPET. Straight forward model of crane-movement at constant `speed` (`default = 1`), with a zone of `safety` (`default = 1`) wide that must be kept free on either side of the crane.

# Functionality
- id(q): identifier and position in the order of cranes along the non-crossing axis
- speed(q): movement speed of the crane
- starting_position(q): initial position of the crane along the non-crossing axis
- zone(q): minimum and maximum position of the crane along the non-crossing axis
- safety(q): minimal distance that must be kept free on either side of the crane
"""
struct Crane
    id::Int
    l0::Int
    speed::Int
    zone::UnitRange{Int64}
    safety::Int
end

id(c::Crane) = c.id
starting_position(c::Crane) = c.l0
speed(c::Crane) = c.speed
zone(c::Crane) = c.zone
safety(c::Crane) = c.safety


@testset "equal.jl" begin
    @testset "EqualConstant()" begin
        trailer = SeaPearl.Trailer()
        x = SeaPearl.IntVar(2, 6, "x", trailer)

        constraint = SeaPearl.EqualConstant(x, 3, trailer)

        @test constraint in x.onDomainChange
        @test constraint.active.value
    end
    @testset "propagate!(::EqualConstant)" begin
        trailer = SeaPearl.Trailer()
        x = SeaPearl.IntVar(2, 6, "x", trailer)
        ax = SeaPearl.IntVarViewMul(x, 3, "3x")

        constraint = SeaPearl.EqualConstant(ax, 6, trailer)

        toPropagate = Set{SeaPearl.Constraint}()
        prunedDomains = SeaPearl.CPModification()

        @test SeaPearl.propagate!(constraint, toPropagate, prunedDomains)

        @test length(ax.domain) == 1
        @test 6 in ax.domain
        @test !(9 in ax.domain)
        @test !(10 in ax.domain)
        @test prunedDomains == SeaPearl.CPModification(
            "3x" => [9, 12, 15, 18],
            "x"  => [3, 4, 5, 6]
        )

        cons2 = SeaPearl.EqualConstant(ax, 9, trailer)

        @test !SeaPearl.propagate!(cons2, toPropagate, prunedDomains)

        @test isempty(ax.domain)

        y = SeaPearl.IntVar(2, 6, "y", trailer)
        constraint1 = SeaPearl.EqualConstant(y, 3, trailer)
        constraint2 = SeaPearl.EqualConstant(y, 4, trailer)

        toPropagate2 = Set{SeaPearl.Constraint}()
        SeaPearl.propagate!(constraint1, toPropagate2, prunedDomains)

        @test constraint2 in toPropagate2

    end

    @testset "pruneEqual!()" begin
        trailer = SeaPearl.Trailer()
        x = SeaPearl.IntVar(2, 6, "x", trailer)
        y = SeaPearl.IntVar(5, 8, "y", trailer)

        SeaPearl.pruneEqual!(y, x)

        @test length(y.domain) == 2
        @test !(8 in y.domain) && 5 in y.domain && 6 in y.domain

    end

    @testset "propagate!(::Equal)" begin
        trailer = SeaPearl.Trailer()
        x = SeaPearl.IntVar(2, 6, "x", trailer)
        y = SeaPearl.IntVar(5, 8, "y", trailer)

        constraint = SeaPearl.Equal(x, y, trailer)
        toPropagate = Set{SeaPearl.Constraint}()
        prunedDomains = SeaPearl.CPModification()

        @test SeaPearl.propagate!(constraint, toPropagate, prunedDomains)

        @test length(x.domain) == 2
        @test length(y.domain) == 2
        @test !(2 in x.domain) && 5 in x.domain && 6 in x.domain
        @test !(8 in y.domain) && 5 in y.domain && 6 in y.domain
        @test prunedDomains == SeaPearl.CPModification("x" => [2, 3, 4],"y" => [7, 8])

        # Propagation test
        z = SeaPearl.IntVar(5, 15, "z", trailer)
        constraint2 = SeaPearl.Equal(y, z, trailer)
        SeaPearl.propagate!(constraint2, toPropagate, prunedDomains)

        # Domain not reduced => not propagation
        @test !(constraint in toPropagate)
        @test !(constraint2 in toPropagate)

        # Domain reduced => propagation
        SeaPearl.remove!(z.domain, 5)
        SeaPearl.propagate!(constraint2, toPropagate, prunedDomains)
        @test constraint in toPropagate
        @test !(constraint2 in toPropagate)

        #Unfeasible test
        t = SeaPearl.IntVar(15, 30, "t", trailer)
        constraint3 = SeaPearl.Equal(z, t, trailer)
        @test !SeaPearl.propagate!(constraint3, toPropagate, prunedDomains)
    end
end


#=##############################################################################
# DESCRIPTION
Testing rotor-on-wing and wing-on-rotor modeling using a scaled-down Bertin's
wing (45-deg, swept-back, planar wing in Example 7.2, pp. 343 of Bertin's
Aerodynamics for Engineers) with two APC 10x7 propellers (one on each side).

# AUTHORSHIP
  * Author    : Eduardo J. Alvarez
  * Email     : Edo.AlvarezR@gmail.com
  * Created   : Dec 2019
  * License   : MIT
=###############################################################################

# ------------ MODULES ---------------------------------------------------------
# Load simulation engine
# import FLOWUnsteady
reload("FLOWUnsteady")
uns = FLOWUnsteady
vlm = uns.vlm

import GeometricTools
gt = GeometricTools

using PyPlot

noise = nothing
try
    # FLOWNoise https://github.com/byuflowlab/FLOWNoise (it requires PSU-WOPWOP)
    import FLOWNoise
    noise = FLOWNoise
catch e
    @warn("FLOWNoise not loaded: $e")
    println("FLOWNoise not loaded: $e")
end

# ------------ GLOBAL VARIABLES ------------------------------------------------
# Default path where to save data
extdrive_path = "/media/edoalvar/MyExtDrive/simulationdata7/"
# extdrive_path = "temps/"

# ------------ HEADERS ---------------------------------------------------------
for header_name in ["singleprop", "isolatedwing", "blownwing", "postprocessing",
                                                                        "noise"]
    include("blownwing_"*header_name*".jl")
end

# ------------ DRIVERS ---------------------------------------------------------
function run_singleprop(; xfoil=true, prompt=true)
    singleprop(; xfoil=xfoil, save_path=extdrive_path*"fvs_singleprop04/",
                                                                  prompt=prompt)
end

function run_isolatedwing(; prompt=true)
    isolatedwing(; save_path=extdrive_path*"fvs_isolatedwing10/", prompt=prompt)
end

function run_blownwing(; xfoil=true, prompt=true)
    blownwing(; xfoil=xfoil, save_path=extdrive_path*"blownwing_test05/",
                                     VehicleType=uns.VLMVehicle, nrevs=10, prompt=prompt)
end

function run_noise(; )

    read_path = "/media/edoalvar/MyExtDrive/simulationdata5/fvs_blownwing11_alloutputs_noise/aero_rawoutputs/"
    save_path = "/media/edoalvar/MyExtDrive/simulationdata5/fvs_blownwing11_alloutputs_noise/test_noise06/"
    postprocessing_noise(  save_path;
                            # ---------- AERO SIMULATION INFO ----------
                            num0=2,                         # First time step to analyze
                            nrevs=29,                       # Number of revs to analyze
                            # ---------- OBSERVERS -------------------------
                            loading=true,                   # Include loading pressure
                            sph_R=1.2*1.693/2, sph_nR=24, sph_ntht=24, # Sphere definition
                            sph_nphi=24, sph_phimax=360,
                            sph_rotation=[90, 0, 0],
                            sph_C=[1.693/4, 0, 0],
                            microphoneX=nothing,
                            ww_nummin=72*15,
                            ww_numless=72*11-2,
                            # ---------- INPUT OPTIONS ---------------------
                            read_path=read_path,
                            # ---------- OUTPUT OPTIONS --------------------
                            prompt=true, debug_paraview=true,
                            debuglvl=1,                     # WW debug level
                        )
end


using CRCBS, GraphUtils, CairoMakie, Makie, FactoryRendering

vtx_grid = initialize_dense_vtx_grid(4,4)
#  1   2   3   4
#  5   6   7   8
#  9  10  11  12
# 13  14  15  16
starts = [1,2,3,4,5,6,7,8]
goals = [13,14,15,16,9,10,11,12]
# graph = initialize_grid_graph_from_vtx_grid(vtx_grid)
graph = construct_factory_env_from_vtx_grid(vtx_grid)
mapf = init_mapf_problem(graph,starts,goals)

solver = PIBTPlanner{Float64}()
set_iteration_limit!(solver,10)
set_verbosity!(solver,4)
solution, valid = pibt!(solver, mapf)
@show valid, convert_to_vertex_lists(solution)

vtx_lists = convert_to_vertex_lists(solution)
grid_vtxs = map(v->[v[1],v[2]], mapf.env.graph.vtxs)
paths = map(p->map(v->grid_vtxs[v],p),vtx_lists)
# scene = Scene()
scene = Scene(show_axis=false,scale_plot=false,resolution=(200,200))
t0 = 0.0
tf = length(paths[1]) + 1.0
pad = 0.75
corners = FactoryRendering.bbox_corners(grid_vtxs,pad)
poly!(scene,Point2f0.(corners), color=:gray)
scene = FactoryRendering.plot_grid_world!(scene,grid_vtxs)

t = Node(t0) # This is the life signal
positions = lift(t->FactoryRendering.env_state_snapshot(paths,t),t)
x = lift(positions->map(p->p[1],positions),positions)
y = lift(positions->map(p->p[2],positions),positions)
text_positions = map(i->lift(positions->tuple(positions[i]...),positions),1:length(paths))
scatter!(scene,x,y,marker=:circle,markersize=0.8,color=:red)
for i in 1:length(paths)
    text!(scene,string(i),textsize=0.4,align=(:center,:center), position=text_positions[i])
end
for (i,g) in enumerate(goals)
    text!(scene,string(i),textsize=0.2,align=(:left,:top),position=tuple((grid_vtxs[g] .+ [-0.45,0.48])...),overdraw=true)
end
scene

record(scene, "robots_moving.webm", 0:0.1:tf; framerate = 20) do dt
    t[] = t0 + dt
end


module KitMSA

using DataFrames
using DelimitedFiles
using FastaIO
using StatsBase
using Random
using GZip

include("clean_MSA.jl")
include("read_write_MSA.jl")
include("BM_utils.jl")
include("contact_prediction.jl")

cod2amino = Dict( "ATA" => Int8(8), "ATC" => Int8(8), "ATT"=> Int8(8), "ATG"=> Int8(11),
        "ACA"=>Int8(17), "ACC"=>Int8(17), "ACG"=>Int8(17), "ACT"=> Int8(17),
        "AAC"=>Int8(12), "AAT"=>Int8(12), "AAA"=>Int8(9), "AAG"=>Int8(9),
        "AGC"=>Int8(16), "AGT"=> Int8(16), "AGA"=> Int8(15), "AGG"=> Int8(15),
        "CTA"=>Int8(10), "CTC"=>Int8(10), "CTG"=>Int8(10), "CTT"=>Int8(10),
        "CCA"=>Int8(13), "CCC"=>Int8(13), "CCG"=>Int8(13), "CCT"=>Int8(13),
        "CAC"=>Int8(7), "CAT"=>Int8(7), "CAA"=>Int8(14), "CAG"=>Int8(14),
        "CGA"=>Int8(15), "CGC"=>Int8(15), "CGG"=>Int8(15), "CGT"=>Int8(15),
        "GTA"=>Int8(18), "GTC"=>Int8(18), "GTG"=>Int8(18), "GTT"=>Int8(18),
        "GCA"=>Int8(1), "GCC"=>Int8(1), "GCG"=>Int8(1), "GCT"=>Int8(1),
        "GAC"=>Int8(3), "GAT"=>Int8(3), "GAA"=>Int8(4), "GAG"=>Int8(4),
        "GGA"=>Int8(6), "GGC"=>Int8(6), "GGG"=>Int8(6), "GGT"=>Int8(6),
        "TCA"=>Int8(16), "TCC"=>Int8(16), "TCG"=>Int8(16), "TCT"=>Int8(16),
        "TTC"=>Int8(5), "TTT"=>Int8(5), "TTA"=>Int8(10), "TTG"=>Int8(10),
        "TAC"=>Int8(20), "TAT"=>Int8(20), "TAA"=> Int8(21), "TAG"=> Int8(21),
        "TGC"=> Int8(2), "TGT"=>Int8(2) , "TGA"=> Int8(21), "TGG"=> Int8(19),
        "---" => Int8(21) )

amino2cod = Dict()
for amino in 1:21
    codons = Array{String, 1}()
    for (key, val) in cod2amino
       val == amino && push!(codons, key)
    end
    amino2cod[amino] = codons
end

export cod2amino, amino2cod

end


# Bio.Align
# =========
#
# Module for sequence alignments.
#
# This file is a part of BioJulia.
# License is MIT: https://github.com/BioJulia/Bio.jl/blob/master/LICENSE.md

module Align

export
    Operation,
    AlignmentAnchor,
    Alignment,
    AlignedSequence,
    seq2ref,
    ref2seq,
    ismatchop,
    isinsertop,
    isdeleteop,
    cigar,
    OP_MATCH,
    OP_INSERT,
    OP_DELETE,
    OP_SKIP,
    OP_SOFT_CLIP,
    OP_HARD_CLIP,
    OP_PAD,
    OP_SEQ_MATCH,
    OP_SEQ_MISMATCH,
    OP_BACK,
    OP_START,
    # alignment types
    GlobalAlignment,
    SemiGlobalAlignment,
    OverlapAlignment,
    LocalAlignment,
    EditDistance,
    HammingDistance,
    LevenshteinDistance,
    # substitution matrices
    AbstractSubstitutionMatrix,
    SubstitutionMatrix,
    DichotomousSubstitutionMatrix,
    EDNAFULL,
    PAM30,
    PAM70,
    PAM250,
    BLOSUM45,
    BLOSUM50,
    BLOSUM62,
    BLOSUM80,
    BLOSUM90,
    # alignment models
    AbstractScoreModel,
    AffineGapScoreModel,
    AbstractCostModel,
    CostModel,
    # pairwise alignment
    PairwiseAlignment,
    count_matches,
    count_mismatches,
    count_insertions,
    count_deletions,
    count_aligned,
    PairwiseAlignmentResult,
    pairalign,
    score,
    distance,
    alignment,
    hasalignment,
    # HTS
    SAMHeader,
    SAMRecord,
    SAMReader,
    SAMWriter,
    BAMRecord,
    BAMReader,
    BAMWriter,
    BAI,
    AuxDataDict,
    ismapped,
    refname,
    nextrefname,
    refindex,
    nextrefindex,
    leftposition,
    rightposition,
    nextleftposition,
    mappingquality,
    flag,
    templatelength,
    seqname,
    cigar,
    sequence,
    seqlength,
    qualities,
    header,
    isoverlapping,  # re-export from Bio.Intervals
    # SAM flags
    SAM_FLAG_PAIRED,
    SAM_FLAG_PROPER_PAIR,
    SAM_FLAG_UNMAP,
    SAM_FLAG_MUNMAP,
    SAM_FLAG_REVERSE,
    SAM_FLAG_MREVERSE,
    SAM_FLAG_READ1,
    SAM_FLAG_READ2,
    SAM_FLAG_SECONDARY,
    SAM_FLAG_QCFAIL,
    SAM_FLAG_DUP,
    SAM_FLAG_SUPPLEMENTARY

import Bio
using Bio.Seq
using Bio.Intervals
import Bio.StringFields: StringField
import IntervalTrees

include("operations.jl")
include("anchors.jl")
include("alignment.jl")
include("alignedseq.jl")

include("types.jl")
include("submat.jl")
include("models.jl")
include("pairwise/pairalign.jl")
include("hts/hts.jl")

include("deprecated.jl")

end


# RunAllSimulations
println("Spinning up model code ... ")
include("../src/new_SimulationFunctions.jl")

# Establish simulation parameters
TSTART = 0.0;
Ts = 1.0; # 0.1
TSTOP = 250.0; # 100 ; 500

# Load data dictionary
data_dictionary = DataFile(TSTART,TSTOP,Ts)

 # # # #
#println(data_dictionary["CONTROL_PARAMETER_ARRAY"])

# update control parameter array here for now
GAIN_MASTER = 0.3
num_regs = 19;
control_parameter_array=data_dictionary["CONTROL_PARAMETER_ARRAY"]
for i in 1:num_regs
  control_parameter_array[i,1]=GAIN_MASTER # Gain for transcriptional activity
end
# TGFB3_inhibition: TGFB3_Ab inhibition; TGFB3_binding
control_parameter_array[1,1]=30.0 # GAIN # _tgfb3_Ab
# NFATc_inhibition: vivit inhibition; Calcineurin_activation_NFATc
control_parameter_array[2,1]=3.0 # GAIN # NFATc_inhibition: vivit inhibition; Calcineurin_activation_NFATc
# GSK3_P_inhibition_Binding_APC_AXIN_forw: GSK3_P inhibition; Binding_APC_AXIN_forw
control_parameter_array[3,1]=0.5 # GAIN # GSK3_P_inhibition_Binding_APC_AXIN_forw
control_parameter_array[3,2]=1.0 # ORDER# GSK3_P_inhibition_Binding_APC_AXIN_forw

# Active_TCF4_inhibition_Ecad: Ecadherin_Bcatenin inhibition; TCF4_Bcatenin_complex_formation
control_parameter_array[4,1]=0.1 # GAIN # Active_TCF4_inhibition_Ecad
control_parameter_array[4,1]=0.9 # GAIN # Active_TCF4_inhibition_Ecad
# Active_TCF4_inhibition_APC: APC_AXIN_Bcatenin inhibition; TCF4_Bcatenin_complex_formation
control_parameter_array[5,1]=0.2 # GAIN # Active_TCF4_inhibition_APC
control_parameter_array[5,2]=1.0 # ORDER # Active_TCF4_inhibition_APC

# activation_Ecadherin_by_Active_NFATc: Active_NFATc induction induction_gene_Ecadherin
control_parameter_array[6,1]=0.3 # GAIN # activation_Ecadherin_by_Active_NFATc
# repression_Ecadherin_by_pSmad: pSmad repression induction_gene_Ecadherin
control_parameter_array[7,1]=5.0 # GAIN # REP_ECAD_PSMAD
control_parameter_array[7,2]=2.0 # ORDER # REP_ECAD_PSMAD
# repression_Ecadherin_by_SNAIL_SLUG: SNAIL_SLUG repression induction_gene_Ecadherin
control_parameter_array[8,1]=0.5 # GAIN_SNAIL_SLUG_ECAD #2.0
# repression_Ecadherin_by_Active_LEF1: Active_LEF1 repression induction_gene_Ecadherin
control_parameter_array[9,1]=0.6# GAIN_ACTIVE_LEF1_ECAD #2.0 0.5
control_parameter_array[9,2]=1.0# ORDER_ACTIVE_LEF1_ECAD
#  control_parameter_array[6,1]=0.3 # GAIN_BCAT_TGFB3
# activation_Ecadherin_by_ecad_virus: ecad_virus induction induction_gene_Ecadherin
control_parameter_array[10,1]=5.0 # ecad_virus ***
# activation_TGFB3_by_TCF4: Active_TCF4 induction induction_gene_TGFB3
control_parameter_array[11,1]=0.15 # GAIN_TCF4_TGFB3
control_parameter_array[11,2]=1.5 # GAIN_TCF4_TGFB3

# repression_TGFB3_by_YREG1: YREG1 repression induction_gene_TGFB3
control_parameter_array[12,1]=1.0 # YREG_TGFB3 #1.0
control_parameter_array[12,2]=1.0 # YREG_TGFB3
# repression_YREG1_by_SNAIL_SLUG: SNAIL_SLUG repression induction_gene_YREG1
control_parameter_array[13,1]=0.1 # GAIN_SNAIL_YREG1 # 0.1
control_parameter_array[13,2]=2.0 # ORDER_SNAIL_YREG1 # 2.0
# activation_LEF1_by_pSmad: pSmad induction induction_gene_LEF1
control_parameter_array[14,1]=0.8# GAIN # _PSMAD_LEF1 # 0.5 # 0.8
control_parameter_array[14,2]=2.0# ORDER# _PSMAD_LEF1       # 1.0, 2.0
# repression_LEF1_by_YREG1: YREG1 repression induction_gene_LEF1
control_parameter_array[15,1]=5.0 # GAIN YREG1_LEF1 # 0.1
control_parameter_array[15,2]=0.5 # ORDER YREG1_LEF1 # 2.0
# activation_vimentin_by_Active_LEF1: Active_LEF1 induction induction_gene_vimentin
control_parameter_array[16,1]=0.3# GAIN_LEF1_VIM
control_parameter_array[16,2]=2.0# ORDER_LEF1_VIM
# activation_SNAIL_SLUG_by_AP1_SP1_P: AP1_SP1_P induction induction_gene_SNAIL_SLUG
control_parameter_array[18,1]=1.0# SS by AP1 SP1
# activation_SNAIL_SLUG_by_ss_virus: ss_virus induction induction_gene_SNAIL_SLUG
control_parameter_array[19,1]=5.0# SS_VIRUS ***

# # # #

data_dictionary["CONTROL_PARAMETER_ARRAY"]=control_parameter_array
#println(data_dictionary["CONTROL_PARAMETER_ARRAY"])

println("Running to steady state ... ")

# # # STEADY STATE
#@show data_dictionary["RATE_CONSTANT_ARRAY"][141] # here

println(now())
# WARNING data_dicionary will be edited within EstimateSteadyState unless copy/deepcopy used!
XSS = EstimateSteadyState(deepcopy(data_dictionary))
println(now())
initial_condition_array = vec(XSS);
#println(initial_condition_array) ###
#@show data_dictionary["RATE_CONSTANT_ARRAY"][141]

writedlm("../output/xss.dat", XSS)

data_dictionary["INITIAL_CONDITION_ARRAY"] = initial_condition_array;
xss_data_dict = deepcopy(data_dictionary)

# # # # TGFB SIMULATION
#
# println("Running TGFB simulation ... ")
#
# # not required for first simulation
# #data_dictionary= deepcopy(xss_data_dict) # reinitialize dictionary with steady state dictionary
#
# (TSIM,X,ICA) =  SolveModelAddTGFB(TSTART,TSTOP,Ts,data_dictionary)
#
# writedlm("/Users/dbassen/work/varner_gould_emt_model/emt_grn_bassen_model/Results/TGFbSimulation/TSIM.dat", TSIM)
# writedlm("/Users/dbassen/work/varner_gould_emt_model/emt_grn_bassen_model/Results/TGFbSimulation/X.dat", X)

# # # VEGF SIMULATION
## WARNING, SIMULATION DATAFILE MAY BE LEAKING
# println("Running VEGF simulation ... ")
#
# data_dictionary = deepcopy(xss_data_dict) # reinitialize dictionary with steady state dictionary
#
# (TSIM,X,ICA) =  SolveModelAddVEGF(TSTART,TSTOP,Ts,data_dictionary)
#
# writedlm("/Users/dbassen/work/varner_gould_emt_model/emt_grn_bassen_model/Results/VEGFSimulation/TSIM.dat", TSIM)
# writedlm("/Users/dbassen/work/varner_gould_emt_model/emt_grn_bassen_model/Results/VEGFSimulation/X.dat", X)
#
#@show xss_data_dict["RATE_CONSTANT_ARRAY"][141]

simulations_name = ["edit_file_none.jl","edit_file_tgfb.jl","edit_file_mapkin.jl","edit_file_ss.jl","edit_file_ssecad.jl","edit_file_ss_tgfb3ab.jl","edit_file_ss_dnlef1.jl","edit_file_tgfb3.jl","edit_file_dnsmad.jl","edit_file_tgfb3_dnlef1.jl","edit_file_tgfb_tgfb3ab.jl"]
#simulations_name = ["edit_file_tgfb.jl","edit_file_vegf.jl","edit_file_vegf_vi.jl","edit_file_tgfb_vi.jl","edit_file_tgfb_vegf.jl","edit_file_vegf_vi.jl","edit_file_tgfb_vegf_vi.jl","edit_file_yreg_ox.jl","edit_file_yreg_ko.jl","edit_file_yreg_ox_tgfb.jl","edit_file_yreg_ko_tgfb.jl"] # Robustness sims
include("spin_up_edits.jl")
#include("robustness_spin_up_edits.jl")
sim_data = [] # Keep data in RAM
i_simdat = 1
for sim in simulations
  name = simulations_name[i_simdat]
  msg = string("Running ",name," simulation ... ")
  println(msg)
  data_dictionary_active = deepcopy(xss_data_dict)
  (TSIM,X) = SolveModelGeneric_ram(TSTART,TSTOP,Ts,data_dictionary_active,sim)

  outdir="../output/Fig2Simulations/"
  tpath= string(outdir,"TSIM",name[10:end-3],".dat")
  xpath= string(outdir,"X",name[10:end-3],".dat")
  push!(sim_data,X)
  writedlm(tpath, TSIM)
  writedlm(xpath, X)
  i_simdat = i_simdat + 1
end


# Convert between different parameterizations.

module Transform

using Compat

# TODO: don't import Model; transformations should operate on
# generic ParamSets
using ..SensitiveFloats
import ..Log

import ..Model: NUM_SOURCE_TYPES, ParamSet, NUM_BANDS, NUM_COLOR_COMPONENTS

export DataTransform, ParamBounds, ParamBox, SimplexBox,
       get_mp_transform, enforce_bounds!,
       VariationalParams, FreeVariationalParams

# A vector of variational parameters.  The outer index is
# of celestial objects, and the inner index is over individual
# parameters for that object (referenced using ParamIndex).
const VariationalParams{T <: Number} = Vector{Vector{T}}
const FreeVariationalParams{T <: Number} = Vector{Vector{T}}

#####################################################################################
# this is essentially a compatibility layer since Model has gotten rid of this code
# it's messy, but this doesn't really matter too much, since Transforms.jl is going
# the way of the dinosaur

struct CanonicalParams <: ParamSet
    pos::Vector{Int}
    gal_frac_dev::Int
    gal_axis_ratio::Int
    gal_angle::Int
    gal_radius_px::Int
    flux_loc::Vector{Int}
    flux_scale::Vector{Int}
    color_mean::Matrix{Int}
    color_var::Matrix{Int}
    a::Matrix{Int}
    k::Matrix{Int}
    CanonicalParams() =
        new([1, 2], 3, 4, 5, 6,
            collect(7:(7+NUM_SOURCE_TYPES-1)),  # flux_loc
            collect((7+NUM_SOURCE_TYPES):(7+2NUM_SOURCE_TYPES-1)), # flux_scale
            reshape((7+2NUM_SOURCE_TYPES):(7+2NUM_SOURCE_TYPES+(NUM_BANDS-1)*NUM_SOURCE_TYPES-1), (NUM_BANDS-1, NUM_SOURCE_TYPES)),  # color_mean
            reshape((7+2NUM_SOURCE_TYPES+(NUM_BANDS-1)*NUM_SOURCE_TYPES):(7+2NUM_SOURCE_TYPES+2*(NUM_BANDS-1)*NUM_SOURCE_TYPES-1), (NUM_BANDS-1, NUM_SOURCE_TYPES)),  # color_var
            reshape((7+2NUM_SOURCE_TYPES+2*(NUM_BANDS-1)*NUM_SOURCE_TYPES):(7+3NUM_SOURCE_TYPES+2*(NUM_BANDS-1)*NUM_SOURCE_TYPES-1), (NUM_SOURCE_TYPES, 1)),  # a
            reshape((7+3NUM_SOURCE_TYPES+2*(NUM_BANDS-1)*NUM_SOURCE_TYPES):(7+3NUM_SOURCE_TYPES+2*(NUM_BANDS-1)*NUM_SOURCE_TYPES+NUM_COLOR_COMPONENTS*NUM_SOURCE_TYPES-1), (NUM_COLOR_COMPONENTS, NUM_SOURCE_TYPES))) # k
end
const ids = CanonicalParams()
Base.length(::Type{CanonicalParams}) = 6 + 3*NUM_SOURCE_TYPES + 2*(NUM_BANDS-1)*NUM_SOURCE_TYPES + NUM_COLOR_COMPONENTS*NUM_SOURCE_TYPES

struct UnconstrainedParams <: ParamSet
    pos::Vector{Int}
    gal_frac_dev::Int
    gal_axis_ratio::Int
    gal_angle::Int
    gal_radius_px::Int
    flux_loc::Vector{Int}
    flux_scale::Vector{Int}
    color_mean::Matrix{Int}
    color_var::Matrix{Int}
    a::Matrix{Int}
    k::Matrix{Int}
    UnconstrainedParams() =
        new([1, 2], 3, 4, 5, 6,
            collect(7:(7+NUM_SOURCE_TYPES-1)),  # flux_loc
            collect((7+NUM_SOURCE_TYPES):(7+2NUM_SOURCE_TYPES-1)), # flux_scale
            reshape((7+2NUM_SOURCE_TYPES):(7+2NUM_SOURCE_TYPES+(NUM_BANDS-1)*NUM_SOURCE_TYPES-1), (NUM_BANDS-1, NUM_SOURCE_TYPES)),  # color_mean
            reshape((7+2NUM_SOURCE_TYPES+(NUM_BANDS-1)*NUM_SOURCE_TYPES):(7+2NUM_SOURCE_TYPES+2*(NUM_BANDS-1)*NUM_SOURCE_TYPES-1), (NUM_BANDS-1, NUM_SOURCE_TYPES)),  # color_var
            reshape((7+2NUM_SOURCE_TYPES+2*(NUM_BANDS-1)*NUM_SOURCE_TYPES):
                    (7+2NUM_SOURCE_TYPES+2*(NUM_BANDS-1)*NUM_SOURCE_TYPES+(NUM_SOURCE_TYPES-1)-1), (NUM_SOURCE_TYPES - 1, 1)),  # a
            reshape((7+2NUM_SOURCE_TYPES+2*(NUM_BANDS-1)*NUM_SOURCE_TYPES+(NUM_SOURCE_TYPES-1)):
                    (7+2NUM_SOURCE_TYPES+2*(NUM_BANDS-1)*NUM_SOURCE_TYPES+(NUM_SOURCE_TYPES-1)+(NUM_COLOR_COMPONENTS-1)*NUM_SOURCE_TYPES-1), (NUM_COLOR_COMPONENTS-1, NUM_SOURCE_TYPES))) # k
end
const ids_free = UnconstrainedParams()
Base.length(::Type{UnconstrainedParams}) =  6 + 2*NUM_SOURCE_TYPES + 2*(NUM_BANDS-1)*NUM_SOURCE_TYPES + (NUM_COLOR_COMPONENTS-1)*NUM_SOURCE_TYPES + NUM_SOURCE_TYPES-1

################################
# Elementary functions.

"""
Unconstrain x in the unit interval to lie in R.
"""
function inv_logit(x::Number)
    @assert(x >= 0)
    @assert(x <= 1)
    -log(1.0 / x - 1)
end

"""
Convert x in R to lie in the unit interval.
"""
function logit(x::Number)
    1.0 / (1.0 + exp(-x))
end

"""
Convert an (n - 1)-vector of real numbers to an n-vector on the simplex, where
the last entry implicitly has the untransformed value 1.
"""
function constrain_to_simplex(x::Vector{T}) where {T<:Number}
    m = maximum(x)
    if m == Inf
        # If more than 1 entry in x is Inf, it may be because the
        # the last entry in z is 0. Here we set all those entries to the
        # same value, though that may not be strictly correct.
        z   = T[ x_entry .== Inf ? one(T) : zero(T) for x_entry in x]
        z ./= sum(z)
        push!(z, 0)
        return z
    else
        z = exp.(x .- m)
        z_sum = sum(z) + exp(-m)
        z ./= z_sum
        push!(z, inv(z_sum)*exp(-m))
        return z
    end
end

"""
Convert an n-vector on the simplex to an (n - 1)-vector in R^{n -1}.  Entries
are expressed relative to the last element.
"""
function unconstrain_simplex(z::Vector{T}) where {T<:Number}
    n = length(z)
    T[ log(z[i]) - log(z[n]) for i = 1:(n - 1)]
end

################################
# The transforms for Celeste.

struct ParamBox
    lb::Float64  # lower bound
    ub::Float64  # upper bound
    scale::Float64

    function ParamBox(lb, ub, scale)
        @assert lb > -Inf # Not supported
        @assert scale > 0.0
        @assert lb < ub
        new(lb, ub, scale)
    end
end

struct SimplexBox
    lb::Float64  # lower bound
    scale::Float64
    n::Int

    function SimplexBox(lb, scale, n)
        @assert n >= 2
        @assert 0.0 <= lb < 1 / n
        new(lb, scale, n)
    end
end

# The vector of transform parameters for a Symbol.
const ParamBounds = Dict{Symbol, Union{Vector{ParamBox}, Vector{SimplexBox}}}

###############################################
# Functions for a "free transform".

function unbox_parameter(param::Number, pb::ParamBox)
    positive_constraint = (pb.ub == Inf)

    # exp and the logit functions handle infinities correctly, so
    # parameters can equal the bounds.
    @assert(pb.lb .<= param .<= pb.ub,
                    string("unbox_parameter: param outside bounds: ",
                                 "$param ($(pb.lb), $(pb.ub))"))

    if positive_constraint
        return log(param - pb.lb) * pb.scale
    else
        param_bounded = (param - pb.lb) / (pb.ub - pb.lb)
        return inv_logit(param_bounded) * pb.scale
    end
end

function box_parameter(free_param::Number, pb::ParamBox)
    positive_constraint = (pb.ub == Inf)
    if positive_constraint
        return exp(free_param / pb.scale) + pb.lb
    else
        return logit(free_param / pb.scale) * (pb.ub - pb.lb) + pb.lb
    end
end

"""
Convert an unconstrained (n-1)-vector to a simplicial n-vector, z, such that
  - sum(z) = 1
  - z >= sb.lb
See notes for a derivation and reasoning.
"""
function simplexify_parameter(free_param::Vector{T},
                              sb::SimplexBox) where {T<:Number}
    @assert length(free_param) == (sb.n - 1)

    # Broadcasting doesn't work with DualNumbers and Floats. :(
    # z_sim is on an unconstrained simplex.
    z_sim = constrain_to_simplex(T[ p / sb.scale for p in free_param ])
    param = T[ (1 - sb.n * sb.lb) * p + sb.lb for p in z_sim ]

    param
end

"""
Invert the transformation simplexify_parameter() by converting an n-vector
on a simplex to R^{n - 1}.
"""
function unsimplexify_parameter(param::Vector{T},
                                sb::SimplexBox) where {T<:Number}
    @assert length(param) == sb.n
    @assert all(param .>= sb.lb)
    @assert(abs(sum(param) - 1) < 1e-14, abs(sum(param) - 1))

    # z_sim is on an unconstrained simplex.
    # Broadcasting doesn't work with DualNumbers and Floats. :(
    z_sim = T[ (p - sb.lb) / (1 - sb.n * sb.lb) for p in param ]
    free_param = T[ p * sb.scale for p in unconstrain_simplex(z_sim) ]

    free_param
end

##################
# Derivatives

"""
Return derivatives of an unscaled transform from free parameters to a simplex.

Args:
  - z_sim: A vector in a simplex (NB: the function returns derivatives of the
           function f(unconstrained) = simplex, but it the answer is expressed
           in terms of the output of the function.)

Returns:
  - jacobian: n by (n -1) matrix of derivatives of the simplex output (in rows)
              wrt the free parameters (in columns)
  - hessian_vec: An n-length vector of the hessian of each simplex output
                 parameter with respect to the (n-1) free input parameters.
"""
function simplex_derivatives(z_sim::Vector{T}) where {T<:Number}
    n = length(z_sim)
    hessian_vec = Vector{Array{T}}(n)

    for i = 1:n
        hessian_vec[i] = Matrix{T}(n - 1, n - 1)
        for j=1:(n - 1), k=1:(n - 1)
            if j != k
                if (j == i)
                    hessian_vec[i][j, k] = -z_sim[i] * z_sim[k] * (1 - 2 * z_sim[i])
                elseif (k == i)
                    hessian_vec[i][j, k] = -z_sim[i] * z_sim[j] * (1 - 2 * z_sim[i])
                else
                    hessian_vec[i][j, k] = 2 * z_sim[i] * z_sim[j] * z_sim[k]
                end
            else # j == k
                if i == j # All equal
                    hessian_vec[i][j, k] = z_sim[i] * (1 - z_sim[j]) * (1 - 2 * z_sim[k])
                else # j == k, but both are different from i
                    hessian_vec[i][j, k] = - z_sim[i] * z_sim[j] * (1 - 2 * z_sim[j])
                end
            end
        end
    end

    jacobian =
        T[ z_sim[i] * (i == j) - z_sim[i] * z_sim[j] for i=1:n, j=1:(n - 1) ]

    jacobian, hessian_vec
end

"""
Return the derivative and hessian of a simplex transform given the constrained
parameters.

Args:
  - param: The constrained parameter (NB: the derivatives are expressed
               as a function of the constrained parameterd despite being
                     the derivative of the function unconstrained -> constrained)
    - sb: A box simplex constraint
"""
function box_simplex_derivatives(param::Vector{T},
                                 sb::SimplexBox) where {T<:Number}

    @assert length(param) == sb.n

    # z_sim is on an unconstrained simplex.
    # Broadcasting doesn't work with DualNumbers and Floats. :(
    z_sim = T[ (p - sb.lb) / (1 - sb.n * sb.lb) for p in param ]

    jacobian, hessian_vec = simplex_derivatives(z_sim)
    for i in 1:sb.n
        hessian_vec[i] *= (sb.scale ^ 2) * (1 - sb.n * sb.lb)
    end
    jacobian *= sb.scale * (1 - sb.n * sb.lb)

    jacobian, hessian_vec
end

"""
Return the derivative and hessian of a box transform given the constrained
parameters.

Args:
  - param: The constrained parameter (NB: the derivatives are expressed
               as a function of the constrained parameterd despite being
                     the derivative of the function unconstrained -> constrained)
    - pb: A box constraint
"""
function box_derivatives(param::Number, pb::ParamBox)
    if pb.ub == Inf
        centered_param = param - pb.lb
        return pb.scale * centered_param, pb.scale ^ 2 * centered_param
    else
        param_range = pb.ub - pb.lb
        centered_param = (param - pb.lb) / param_range
        derivative = param_range * centered_param * (1 - centered_param) / pb.scale
        return derivative, derivative * (1 - 2 * centered_param) / pb.scale
    end
end

"""
A datatype containing derivatives of the transform from free to constrained
parameters.

Members:
  dparam_dfree: The Jacobian of the transformation
                contrainted_param = f(free param)
  d2param_dfree2: A vector of hessians.  Each element is the Hessian of one
                  component of the aforementioned f()
"""
struct TransformDerivatives{T<:Number}
    dparam_dfree::Matrix{T}
    d2param_dfree2::Vector{Matrix{T}}
    Sa::Int

    # TODO: use sparse matrices?
    function TransformDerivatives{T}(Sa::Int) where {T<:Number}
        dparam_dfree =
            zeros(T, Sa * length(CanonicalParams), Sa * length(UnconstrainedParams))
        d2param_dfree2 = Vector{Matrix{T}}(Sa * length(CanonicalParams))
        for i in 1:(Sa * length(CanonicalParams))
            d2param_dfree2[i] =
                zeros(T, Sa * length(UnconstrainedParams), Sa * length(UnconstrainedParams))
        end
        new(dparam_dfree, d2param_dfree2, Sa)
    end
end

"""
Populate a TransformDerivatives object in place.

Args:
  - bounds: A vector containing one ParamBounds for each active source in the
            same order as active_sources.
  - transform_derivatives: TransformDerivatives to be populated.

Returns:
  Update transform_derivatives in place.
"""
function get_transform_derivatives!(
        vp::VariationalParams,
        active_sources::Vector{Int},
        bounds::Vector{ParamBounds},
        transform_derivatives::TransformDerivatives)

    @assert transform_derivatives.Sa == length(active_sources)

    for param in fieldnames(ids), sa = 1:length(active_sources)
        s = active_sources[sa]
        constraint_vec = bounds[sa][param]

        if isa(constraint_vec[1], ParamBox) # It is a box constraint
            @assert(length(constraint_vec) ==
                    length(getfield(ids_free, param)) ==
                    length(getfield(ids, param)))

            # Get each components' derivatives one by one.
            for ind = 1:length(constraint_vec)
                @assert isa(constraint_vec[ind], ParamBox)
                vp_ind = getfield(ids, param)[ind]
                vp_free_ind = getfield(ids_free, param)[ind]

                jac, hess = box_derivatives(vp[s][vp_ind], constraint_vec[ind])

                vp_sf_ind = length(CanonicalParams) * (sa - 1) + vp_ind
                vp_free_sf_ind = length(UnconstrainedParams) * (sa - 1) + vp_free_ind

                transform_derivatives.dparam_dfree[vp_sf_ind, vp_free_sf_ind] = jac
                transform_derivatives.d2param_dfree2[vp_sf_ind][vp_free_sf_ind, vp_free_sf_ind] = hess
            end
        else # It is a simplex constraint

            # If a param is not a box constraint, it must have all simplex constraints.
            @assert all([ isa(constraint, SimplexBox)  for constraint in constraint_vec])

            param_size = size(getfield(ids, param))
            @assert length(constraint_vec) == param_size[2]
            for col=1:(param_size[2])
                vp_free_ind = getfield(ids_free, param)[:, col]
                vp_ind = getfield(ids, param)[:, col]
                vp_sf_ind = length(CanonicalParams) * (sa - 1) + vp_ind
                vp_free_sf_ind = length(UnconstrainedParams) * (sa - 1) + vp_free_ind

                jac, hess = Transform.box_simplex_derivatives(
                    vp[s][vp_ind], constraint_vec[col])

                transform_derivatives.dparam_dfree[vp_sf_ind, vp_free_sf_ind] = jac
                for row in 1:(param_size[1])
                    transform_derivatives.d2param_dfree2[
                         vp_sf_ind[row]][vp_free_sf_ind, vp_free_sf_ind] = hess[row]
                end
            end
        end
    end
end

function get_transform_derivatives(vp::VariationalParams,
                                   active_sources::Vector{Int},
                                   bounds::Vector{ParamBounds})
  transform_derivatives =
    TransformDerivatives{Float64}(length(active_sources))
  get_transform_derivatives!(vp, active_sources, bounds, transform_derivatives)
  transform_derivatives
end

######################
# Functions to take actual parameter vectors.

"""
Convert between variational parameter vectors and unconstrained parameters.

Args:
  - vp: The vector of contrainted variationl parameters
  - vp_free: The vector of uncontrainted parameters for optimization
  - bounds: ParamBounds describing the transformation
  - to_unconstrained: If true, converts vp to vp_free.  If false, converts
                      vp_free to vp.

Returns:
  If to_unconstrained is true, updates vp_free in place.
  If to_unconstrained is false, updates vp in place.
"""
function perform_transform!(vp::Vector{T},
                            vp_free::Vector{T},
                            bounds::ParamBounds,
                            to_unconstrained::Bool) where {T<:Number}
    for (param, constraint_vec) in bounds
        is_box = isa(bounds[param], Array{ParamBox})
        if is_box
            # Apply a box constraint to each parameter.
            @assert(length(getfield(ids, param)) == length(getfield(ids_free, param)) ==
                length(bounds[param]))
            for ind in 1:length(getfield(ids, param))
                constraint = constraint_vec[ind]
                free_ind = getfield(ids_free, param)[ind]
                vp_ind = getfield(ids, param)[ind]
                to_unconstrained ?
                    vp_free[free_ind] = unbox_parameter(vp[vp_ind], constraint):
                    vp[vp_ind] = box_parameter(vp_free[free_ind], constraint)
            end
        else
            # Apply a simplex constraint to each parameter.
            @assert isa(bounds[param], Array{SimplexBox})
            param_size = size(getfield(ids, param))

            # Each column is a simplex and should have its own
            # simplicial constraints.
            @assert length(bounds[param]) == param_size[2]
            for col in 1:(param_size[2])
                free_ind = getfield(ids_free, param)[:, col]
                vp_ind = getfield(ids, param)[:, col]
                constraint = constraint_vec[col]
                to_unconstrained ?
                    vp_free[free_ind] = unsimplexify_parameter(vp[vp_ind], constraint):
                    vp[vp_ind] = simplexify_parameter(vp_free[free_ind], constraint)
            end
        end
    end
end

function free_to_vp!(vp_free::Vector{T},
                     vp::Vector{T},
                     bounds::ParamBounds) where {T<:Number}
    perform_transform!(vp, vp_free, bounds, false)
end

function vp_to_free!(vp::Vector{T},
                     vp_free::Vector{T},
                     bounds::ParamBounds) where {T<:Number}
    perform_transform!(vp, vp_free, bounds, true)
end

#####################
# Conversion to and from variational parameter vectors and arrays.

"""
Transform VariationalParams to an array.

Args:
  - vp = variational parameters
  - omitted_ids = ids in ParamIndex

There is probably no use for this function, since you'll only be passing
trasformations to the optimizer, but I'll include it for completeness.
"""
function free_vp_to_array(vp::FreeVariationalParams{T},
                          omitted_ids::Vector{Int}) where {T<:Number}

    left_ids = setdiff(1:length(UnconstrainedParams), omitted_ids)
    new_P = length(left_ids)
    S = length(vp)
    x_new = zeros(T, new_P, S)

    for p1 in 1:length(left_ids), s=1:S
        p0 = left_ids[p1]
        x_new[p1, s] = vp[s][p0]
    end

    x_new
end

"""
Transform a parameter vector to variational parameters in place.

Args:
 - xs: A (param x sources) matrix created from free variational parameters.
 - vp_free: Free variational parameters.  Only the ids in kept_ids
            will be updated.
 - kept_ids: Ids to update (from ids_free)

Returns:
 - Update vp_free in place.
"""
function array_to_free_vp!(xs::Array{T},
                           vp_free::FreeVariationalParams{T},
                           kept_ids::Vector{Int}) where {T}
    for s in 1:length(vp_free), p1 in 1:length(kept_ids)
        vp_free[s][kept_ids[p1]] = xs[p1, s]
    end
end

#########################
# Define the exported variables.

"""
Functions to move between a single source's variational parameters and a
transformation of the data for optimization.

to_vp: A function that takes transformed parameters and returns
             variational parameters
from_vp: A function that takes variational parameters and returned
                 transformed parameters
to_vp!: A function that takes (transformed paramters, variational parameters)
                and updates the variational parameters in place
from_vp!: A function that takes (variational paramters, transformed parameters)
                    and updates the transformed parameters in place
...
transform_sensitive_float: A function that takes (sensitive float, model
    parameters) where the sensitive float contains partial derivatives with
    respect to the variational parameters and returns a sensitive float with total
    derivatives with respect to the transformed parameters.
bounds: The bounds for each parameter and each object in ElboArgs.
active_sources: The sources that are being optimized.    Only these sources'
    parameters are transformed into the parameter vector.
"""
struct DataTransform
    bounds::Vector{ParamBounds}
    active_sources::Vector{Int}
    S::Int
    # TODO: Maybe this should be initialized with ElboArgs with optional
    # custom bounds. Or maybe it should be part of ElboArgs with one transform
    # per celestial object rather than a single object containing an array of
    # transforms.
    function DataTransform(bounds::Vector{ParamBounds},
                           active_sources=collect(1:length(bounds)),
                           S=length(bounds))
        @assert length(bounds) == length(active_sources)
        @assert maximum(active_sources) <= S
        new(bounds, active_sources, S)
    end
end

function from_vp!(dt::DataTransform,
                  vp::VariationalParams{T},
                  vp_free::VariationalParams{T}) where {T}
    active_sources = dt.active_sources
    bounds = dt.bounds
    for i in eachindex(active_sources)
        vp_to_free!(vp[active_sources[i]], vp_free[i], bounds[i])
    end
    return nothing
end

function to_vp!(dt::DataTransform,
                vp_free::FreeVariationalParams{T},
                vp::VariationalParams{T}) where {T<:Number}
    active_sources = dt.active_sources
    bounds = dt.bounds
    for i in eachindex(active_sources)
        free_to_vp!(vp_free[i], vp[active_sources[i]], bounds[i])
    end
    return nothing
end

function from_vp(dt::DataTransform,
                 vp::VariationalParams{T}) where {T}
    vp_free = [zeros(T, length(ids_free)) for _ in 1:length(dt.active_sources)]
    from_vp!(dt, vp, vp_free)
    vp_free
end

function to_vp(dt::DataTransform, vp_free::FreeVariationalParams{T}) where {T}
    n_active_sources = length(dt.active_sources)
    @assert(n_active_sources == dt.S,
            string("to_vp is not supported when active_sources is a ",
                   "strict subset of all sources."))
    vp = [zeros(T, length(CanonicalParams)) for _ in 1:n_active_sources]
    to_vp!(dt, vp_free, vp)
    vp
end

function vp_to_array(dt::DataTransform,
                     vp::VariationalParams{T},
                     omitted_ids::Vector{Int}) where {T}
    vp_trans = from_vp(dt, vp)
    free_vp_to_array(vp_trans, omitted_ids)
end

function array_to_vp!(dt::DataTransform,
                      xs::Array{T},
                      vp::VariationalParams{T},
                      kept_ids::Vector{Int}) where {T}
    # This needs to update vp in place so that variables in omitted_ids
    # stay at their original values.
    vp_trans = from_vp(dt, vp)
    array_to_free_vp!(xs, vp_trans, kept_ids)
    to_vp!(dt, vp_trans, vp)
end

# Given a sensitive float with derivatives with respect to all the
# constrained parameters, calculate derivatives with respect to
# the unconstrained parameters.
#
# Note that all the other functions in DeterministicVI calculated derivatives with
# respect to the constrained parameterization.
function transform_sensitive_float!(dt::DataTransform,
                                    sf_free::SensitiveFloat,
                                    sf::SensitiveFloat,
                                    vp::VariationalParams{T},
                                    active_sources::Vector{Int}) where {T}

    if isnan(sf.v[])
       error("sf has NaN value:", sf_free.v[])
    end

    if any(isnan.(sf.d))
       error("sf has NaN derivatives:", sf_free.d[])
    end

    if any(isnan.(sf.h))
       error("sf has NaN hessian:", sf_free.h)
    end

    n_active_sources = length(active_sources)
    @assert size(sf.d) == (length(CanonicalParams), n_active_sources)

    transform_derivatives = get_transform_derivatives(vp, active_sources, dt.bounds)

    d, h = sf.d, sf.h
    free_d, free_h = sf_free.d, sf_free.h
    dparam_dfree = transform_derivatives.dparam_dfree
    d2param_dfree2 = transform_derivatives.d2param_dfree2

    sf_free.v[] = sf.v[]

    copy!(free_d, reshape(dparam_dfree' * d, length(UnconstrainedParams), n_active_sources))

    At_mul_B!(free_h, dparam_dfree, h * dparam_dfree)

    for i in eachindex(d)
        X, n = d2param_dfree2[i], d[i]
        for j in eachindex(X)
            free_h[j] += X[j] * n
        end
    end

    symmetrize!(free_h)

    if isnan(sf_free.v[])
        error("sf_free has NaN value:", sf_free.v[])
    end

    if any(isnan.(sf_free.d))
        error("sf_free has NaN derivatives:", sf_free.d[])
    end

    if any(isnan.(sf_free.h))
        error("sf_free has NaN hessian:", sf_free.h)
    end

    return sf_free
end

# Ensure exact symmetry, which is necessary for
# some numerical linear algebra routines.
function symmetrize!(A, c = 0.5)
    for i in 1:size(A, 1), j in 1:i
        A[i, j] = A[j, i] = c * (A[i, j] + A[j, i])
    end
    return A
end

function get_mp_transform(vp::VariationalParams{T},
                          active_sources::Vector{Int};
                          loc_scale=1.0,
                          loc_width=1e-4) where {T<:Number}
    bounds = Vector{ParamBounds}(length(active_sources))

    # Note that, for numerical reasons, the bounds must be on the scale
    # of reasonably meaningful changes.
    for si in 1:length(active_sources)
        s = active_sources[si]
        bounds[si] = ParamBounds()
        bounds[si][:pos] = Vector{ParamBox}(2)
        pos = vp[s][ids.pos]
        for axis in 1:2
            bounds[si][:pos][axis] =
                ParamBox(pos[axis] - loc_width, pos[axis] + loc_width, loc_scale)
        end
        bounds[si][:flux_loc] = Vector{ParamBox}(NUM_SOURCE_TYPES)
        bounds[si][:flux_scale] = Vector{ParamBox}(NUM_SOURCE_TYPES)
        for i in 1:NUM_SOURCE_TYPES
            bounds[si][:flux_loc][i] = ParamBox(-1.0, 10., 1.0)
            bounds[si][:flux_scale][i] = ParamBox(1e-4, 0.1, 1.0)
        end
        bounds[si][:color_mean] = Vector{ParamBox}(4 * NUM_SOURCE_TYPES)
        bounds[si][:color_var] = Vector{ParamBox}(4 * NUM_SOURCE_TYPES)
        for ind in 1:length(ids.color_mean)
            bounds[si][:color_mean][ind] = ParamBox(-10., 10., 1.0)
            bounds[si][:color_var][ind] = ParamBox(1e-4, 1., 1.0)
        end
        bounds[si][:gal_frac_dev] = ParamBox[ ParamBox(1e-2, 1 - 1e-2, 1.0) ]
        bounds[si][:gal_axis_ratio] = ParamBox[ ParamBox(1e-2, 1 - 1e-2, 1.0) ]
        bounds[si][:gal_angle] = ParamBox[ ParamBox(-10.0, 10.0, 1.0) ]
        bounds[si][:gal_radius_px] = ParamBox[ ParamBox(0.1, 70., 1.0) ]

        bounds[si][:a] = Vector{SimplexBox}(1)
        bounds[si][:a][1] = SimplexBox(0.005, 1.0, 2)

        bounds[si][:k] = Vector{SimplexBox}(NUM_SOURCE_TYPES)
        for i in 1:NUM_SOURCE_TYPES
            bounds[si][:k][i] = SimplexBox(0.01 / NUM_COLOR_COMPONENTS, 1.0, NUM_COLOR_COMPONENTS)
        end
    end

    DataTransform(bounds, active_sources, length(vp))
end

"""
Put the variational parameters within the bounds of the transform.
Returns:
  Updates vp in place.
"""
function enforce_bounds!(vp::VariationalParams{T},
                         active_sources::Vector{Int},
                         transform::DataTransform) where {T<:Number}
    for sa in eachindex(transform.active_sources), (param, constraint_vec) in transform.bounds[sa]
        s = active_sources[sa]
        is_box = isa(constraint_vec, Array{ParamBox})
        if is_box
            # Box parameters.
            for ind in 1:length(getfield(ids, param))
                constraint = constraint_vec[ind]
                if !(constraint.lb <=
                         vp[s][getfield(ids, param)[ind]] <=
                         constraint.ub)
                    # Don't set the value to exactly the lower bound to avoid Inf
                    diff = constraint.ub - constraint.lb
                    epsilon = diff == Inf ? 1e-12: diff * 1e-12
                    vp[s][getfield(ids, param)[ind]] =
                        min(vp[s][getfield(ids, param)[ind]], constraint.ub - epsilon)
                    vp[s][getfield(ids, param)[ind]] =
                        max(vp[s][getfield(ids, param)[ind]], constraint.lb + epsilon)
                end
            end
        else
            # Simplex parameters.
            param_size = size(getfield(ids, param))
            for col in 1:param_size[2]
                constraint = constraint_vec[col]
                param_sum = zero(T)
                for row in 1:param_size[1]
                    if !(constraint.lb <= vp[s][getfield(ids, param)[row, col]] <= 1.0)
                        # Don't set the value to exactly the lower bound to avoid Inf
                        epsilon = (1.0 - constraint.lb) * 1e-12
                        vp[s][getfield(ids, param)[row, col]] =
                            min(vp[s][getfield(ids, param)[row, col]], 1.0 - epsilon)
                        vp[s][getfield(ids, param)[row, col]] =
                            max(vp[s][getfield(ids, param)[row, col]],
                                    constraint.lb + epsilon)
                    end
                    param_sum += vp[s][getfield(ids, param)[row, col]]
                end
                if param_sum != 1.0
                    # Normalize in a way that maintains the lower bounds
                    rescale =
                      (1 - constraint.n * constraint.lb) /
                      (param_sum - constraint.n * constraint.lb)
                    vp[s][getfield(ids, param)[:, col]] =
                      constraint.lb +
                      rescale * (vp[s][getfield(ids, param)[:, col]] -
                                 constraint.lb)
                end
            end
        end
    end
end

end


### Contains all old constructors for compatibility issues

deprecation_warning = "Deprecated constructor, use VGP(X,y,kernel,likelihood,inference) or SVGP(X,y,kernel,likelihood,inference,m) in the future"

function BatchGPRegression(X::AbstractArray{T},y::AbstractArray;Autotuning::Bool=true,nEpochs::Integer = 10, kernel=0,noise::T=1e-3,verbose::Integer=0) where {T<:Real}
    @warn deprecation_warning
    model = VGP(X,y,kernel,GaussianLikelihood(noise),AnalyticVI(ϵ=ϵ),verbose=verbose,optimizer=Autotuning,atfrequency=AutotuningFrequency,IndependentPriors=IndependentGPs)
    for i in 1:model.nLatent
        model.μ[i] = μ_init[1]*ones(length(model.μ[i]))
    end
    return model
end

function SparseGPRegression(X::AbstractArray{T},y::AbstractArray;Stochastic::Bool=false,AdaptiveLearningRate::Bool=true, Autotuning::Bool=false,OptimizeIndPoints::Bool=false, nEpochs::Integer = 10000,batchsize::Integer=-1,κ_s::T=1.0,τ_s::Integer=100, kernel=0,noise::T=1e-3,m::Integer=0,AutotuningFrequency::Integer=2, ϵ::Real=1e-5,μ_init::Array{Float64,1}=[0.0],SmoothingWindow::Integer=5, verbose::Integer=0) where {T<:Real}
    @warn deprecation_warning
    model = SVGP(X,y,kernel,GaussianLikelihood(noise),Stochastic ? AnalyticSVI(batchsize,ϵ=ϵ,optimizer=AdaptiveLearningRate ? ALRSVI() : InverseDecay(τ=τ_s,κ=κ_s)) : AnalyticVI(ϵ=ϵ),m,verbose=verbose,optimizer=Autotuning,atfrequency=AutotuningFrequency,IndependentPriors=IndependentGPs)
    for i in 1:model.nLatent
        model.μ[i] = μ_init[1]*ones(length(model.μ[i]))
    end
    return model
end

function MultiClass(X::AbstractArray{T},y::AbstractArray;Autotuning::Bool=false,nEpochs::Integer = 200, KStochastic::Bool = false, nClassesUsed::Int=0,
                                kernel=0,AutotuningFrequency::Integer=2,IndependentGPs::Bool=false,
                                ϵ::Real=T(1e-5),μ_init::Vector{T}=zeros(T,1),verbose::Integer=0) where {T<:Real}
    @warn deprecation_warning
    model = VGP(X,y,kernel,LogisticSoftMaxLikelihood(),AnalyticVI(ϵ=ϵ),verbose=verbose,optimizer=Autotuning,atfrequency=AutotuningFrequency,IndependentPriors=IndependentGPs)
    for i in 1:model.nLatent
        model.μ[i] = μ_init[1]*ones(length(model.μ[i]))
    end
    return model
end

function SparseMultiClass(X::AbstractArray{T},y::AbstractArray;Stochastic::Bool=false,KStochastic::Bool=false,nClassesUsed::Int=0,AdaptiveLearningRate::Bool=true, Autotuning::Bool=false,OptimizeIndPoints::Bool=false, IndependentGPs::Bool=true, nEpochs::Integer = 10000,KSize::Int64=-1,batchsize::Integer=-1,κ_s::T=T(0.51),τ_s::Integer=1, kernel=0,m::Integer=0, AutotuningFrequency::Integer=2, ϵ::Real=1e-5,μ_init::Vector{T}=zeros(T,1),SmoothingWindow::Integer=5, verbose::Integer=0) where {T<:Real}
    @warn deprecation_warning
    model = SVGP(X,y,kernel,LogisticSoftMaxLikelihood(),Stochastic ? AnalyticSVI(batchsize,ϵ=ϵ,optimizer=AdaptiveLearningRate ? ALRSVI() : InverseDecay(τ=τ_s,κ=κ_s)) : AnalyticVI(ϵ=ϵ),m,verbose=verbose,optimizer=Autotuning,atfrequency=AutotuningFrequency,IndependentPriors=IndependentGPs)
    for i in 1:model.nLatent
        model.μ[i] = μ_init[1]*ones(length(model.μ[i]))
    end
    return model
end

function BatchBSVM(X::AbstractArray{T},y::AbstractArray;Autotuning::Bool=false,nEpochs::Integer = 200, kernel=0,AutotuningFrequency::Integer=1, ϵ::T=1e-5,μ_init::Vector{T}=ones(T,1),verbose::Integer=0) where {T<:Real}
    @warn deprecation_warning
    model = VGP(X,y,kernel,BayesianSVM(),AnalyticVI(ϵ=ϵ),verbose=verbose,optimizer=Autotuning,atfrequency=AutotuningFrequency,IndependentPriors=IndependentGPs)
    for i in 1:model.nLatent
        model.μ[i] = μ_init[1]*ones(length(model.μ[i]))
    end
    return model
end

function SparseBSVM(X::AbstractArray{T},y::AbstractArray;Stochastic::Bool=false,AdaptiveLearningRate::Bool=true,Autotuning::Bool=false,OptimizeIndPoints::Bool=false, nEpochs::Integer = 10000,batchsize::Integer=-1,κ_s::Real=1.0,τ_s::Integer=100,kernel=0,m::Integer=0,AutotuningFrequency::Integer=1,ϵ::Real=1e-5,μ_init::Array{Float64,1}=[0.0],SmoothingWindow::Integer=5,verbose::Integer=0) where {T<:Real}
    @warn deprecation_warning
    model = SVGP(X,y,kernel,BayesianSVM(),Stochastic ? AnalyticSVI(batchsize,ϵ=ϵ,optimizer=AdaptiveLearningRate ? ALRSVI() : InverseDecay(τ=τ_s,κ=κ_s)) : AnalyticVI(ϵ=ϵ),m,verbose=verbose,optimizer=Autotuning,atfrequency=AutotuningFrequency,IndependentPriors=IndependentGPs)
    for i in 1:model.nLatent
        model.μ[i] = μ_init[1]*ones(length(model.μ[i]))
    end
    return model
end

function BatchXGPC(X::AbstractArray{T},y::AbstractArray;Autotuning::Bool=false,nEpochs::Integer = 200, kernel=0,AutotuningFrequency::Integer=1, ϵ::T=1e-5,μ_init::Vector{T}=ones(T,1),verbose::Integer=0) where {T<:Real}
    @warn deprecation_warning
    model = VGP(X,y,kernel,LogisticLikelihood(),AnalyticVI(ϵ=ϵ),verbose=verbose,optimizer=Autotuning,atfrequency=AutotuningFrequency,IndependentPriors=IndependentGPs)
    for i in 1:model.nLatent
        model.μ[i] = μ_init[1]*ones(length(model.μ[i]))
    end
    return model
end

function SparseXGPC(X::AbstractArray{T},y::AbstractArray;Stochastic::Bool=false,AdaptiveLearningRate::Bool=true,Autotuning::Bool=false,OptimizeIndPoints::Bool=false, nEpochs::Integer = 10000,batchsize::Integer=-1,κ_s::Real=1.0,τ_s::Integer=100,kernel=0,m::Integer=0,AutotuningFrequency::Integer=1,ϵ::Real=1e-5,μ_init::Array{Float64,1}=[0.0],SmoothingWindow::Integer=5,verbose::Integer=0) where {T<:Real}
    @warn deprecation_warning
    model = SVGP(X,y,kernel,LogisticLikelihood(),Stochastic ? AnalyticSVI(batchsize,ϵ=ϵ,optimizer=AdaptiveLearningRate ? ALRSVI() : InverseDecay(τ=τ_s,κ=κ_s)) : AnalyticVI(ϵ=ϵ),m,verbose=verbose,optimizer=Autotuning,atfrequency=AutotuningFrequency,IndependentPriors=IndependentGPs)
    for i in 1:model.nLatent
        model.μ[i] = μ_init[1]*ones(length(model.μ[i]))
    end
    return model
end

function BatchStudentT(X::AbstractArray{T},y::AbstractArray;Autotuning::Bool=false, nEpochs::Integer = 200, kernel=0,AutotuningFrequency::Integer=1, ϵ::Real=1e-5,μ_init::Vector{T}=ones(T,1),verbose::Integer=0,ν::T=5.0) where {T<:Real}
    @warn deprecation_warning
    model = VGP(X,y,kernel,StudentTLikelihood(ν),AnalyticVI(ϵ=ϵ),verbose=verbose,optimizer=Autotuning,atfrequency=AutotuningFrequency,IndependentPriors=IndependentGPs)
    for i in 1:model.nLatent
        model.μ[i] = μ_init[1]*ones(length(model.μ[i]))
    end
    return model
end

function SparseStudentT(X::AbstractArray{T},y::AbstractArray;Stochastic::Bool=false,AdaptiveLearningRate::Bool=true, Autotuning::Bool=false,OptimizeIndPoints::Bool=false, nEpochs::Integer = 10000,batchsize::Integer=-1,κ_s::Float64=1.0,τ_s::Integer=100, kernel=0,m::Integer=0,AutotuningFrequency::Integer=1, ϵ::Real=1e-5,μ_init::Array{Float64,1}=[0.0],SmoothingWindow::Integer=5, verbose::Integer=0,ν::Real=5.0) where {T<:Real}
    @warn deprecation_warning
    model = SVGP(X,y,kernel,StudentTLikelihood(ν),Stochastic ? AnalyticSVI(batchsize,ϵ=ϵ,optimizer=AdaptiveLearningRate ? ALRSVI() : InverseDecay(τ=τ_s,κ=κ_s)) : AnalyticVI(ϵ=ϵ),m,verbose=verbose,optimizer=Autotuning,atfrequency=AutotuningFrequency,IndependentPriors=IndependentGPs)
    for i in 1:model.nLatent
        model.μ[i] = μ_init[1]*ones(length(model.μ[i]))
    end
    return model
end


using CUDA.CUBLAS
using CUDA.CUDNN

import CUDA: libcudnn, @checked, @argout
import CUDA.CUDNN:
    @check, @runtime_ccall, @workspace, handle,
    cudnnStatus_t, cudnnConvolutionDescriptor_t,
    ConvDesc, TensorDesc, FilterDesc,
    CuArray, CuVector, CUDNNFloat, cdsize,
    cudnnConvolutionForward,
    cudnnGetConvolutionForwardWorkspaceSize, cudnnConvolutionFwdAlgo_t,
    cudnnConvolutionBackwardData,
    cudnnGetConvolutionBackwardDataWorkspaceSize, cudnnConvolutionBwdDataAlgo_t,
    cudnnConvolutionBackwardFilter,
    cudnnGetConvolutionBackwardFilterWorkspaceSize, cudnnConvolutionBwdFilterAlgo_t
import NNlib: depthwiseconv!, ∇depthwiseconv_filter!, ∇depthwiseconv_data!

const CuOrVector = Union{CuVector, Vector}
const CuOrMatrix = Union{CuMatrix, Matrix}
const CuOrArray = Union{CuArray, Array}

# Implement conversion to dense, diagonal matrix.

diagonal(x::CuArray{T, 1}) where T<:Real = convert(CuArray, Diagonal(x))

# Implement GPU support for depthwise separable convolutions.

@checked function cudnnGetConvolutionGroupCount(convDesc, count)
    @runtime_ccall(
        (:cudnnGetConvolutionGroupCount, libcudnn()),
        cudnnStatus_t,
        (cudnnConvolutionDescriptor_t, Ptr{Cint}),
        convDesc,
        count
    )
end

@checked function cudnnSetConvolutionGroupCount(convDesc, count)
    @runtime_ccall(
        (:cudnnSetConvolutionGroupCount, libcudnn()),
        cudnnStatus_t,
        (cudnnConvolutionDescriptor_t, Cint),
        convDesc,
        count
    )
end

function ConvDesc(T, cdims::DepthwiseConvDims)
    cd = Ref{cudnnConvolutionDescriptor_t}()
    CUDNN.cudnnCreateConvolutionDescriptor(cd)
    N = NNlib.spatial_dims(cdims)
    CUDNN.cudnnSetConvolutionNdDescriptor(
        cd[],
        N,
        # Asymmetric padding is not supported.
        cdsize(NNlib.padding(cdims)[1:2:end], N),
        cdsize(NNlib.stride(cdims), N),
        cdsize(NNlib.dilation(cdims), N),
        NNlib.flipkernel(cdims),
        CUDNN.cudnnDataType(T)
    )
    # Set number of groups equal to number of channels to get a depthwise
    # separable convolution.
    cudnnSetConvolutionGroupCount(cd[], NNlib.channels_in(cdims))
    this = ConvDesc(cd[])
    CUDNN.finalizer(CUDNN.unsafe_free!, this)
    return this
end

function cudnnConvolutionForward(
    y::CuArray{T, N},
    x::CuArray{T, N},
    w::CuArray{T, N},
    cdims::DepthwiseConvDims;
    algo=0,
    alpha=1,
    beta=0
) where {T, N}
    @workspace size = @argout(
        cudnnGetConvolutionForwardWorkspaceSize(
            handle(),
            TensorDesc(x),
            FilterDesc(w),
            ConvDesc(T, cdims),
            TensorDesc(y),
            cudnnConvolutionFwdAlgo_t(algo),
            out(Ref{Csize_t}())
        )
    )[] workspace -> begin
        cudnnConvolutionForward(
            handle(),
            Ref(T(alpha)),
            TensorDesc(x),
            x,
            FilterDesc(w),
            w,
            ConvDesc(T,cdims),
            cudnnConvolutionFwdAlgo_t(algo),
            workspace,
            sizeof(workspace),
            Ref(T(beta)),
            TensorDesc(y),
            y
        )
    end
    return y
end

function cudnnConvolutionBackwardData(
    dx::CuArray{T, N},
    w::CuArray{T, N},
    dy::CuArray{T, N},
    cdims::DepthwiseConvDims;
    algo=0,
    alpha=1,
    beta=0
) where {T, N}
    @workspace size = @argout(
        cudnnGetConvolutionBackwardDataWorkspaceSize(
            handle(),
            FilterDesc(w),
            TensorDesc(dy),
            ConvDesc(T, cdims),
            TensorDesc(dx),
            cudnnConvolutionBwdDataAlgo_t(algo),
            out(Ref{Csize_t}())
        )
    )[] workspace -> begin
        cudnnConvolutionBackwardData(
            handle(),
            Ref(T(alpha)),
            FilterDesc(w),
            w,
            TensorDesc(dy),
            dy,
            ConvDesc(T, cdims),
            cudnnConvolutionBwdDataAlgo_t(algo),
            workspace,
            sizeof(workspace),
            Ref(T(beta)),
            TensorDesc(dx),
            dx
        )
    end
    return dx
end

function cudnnConvolutionBackwardFilter(
    dw::CuArray{T, N},
    x::CuArray{T, N},
    dy::CuArray{T, N},
    cdims::DepthwiseConvDims;
    algo=0,
    alpha=1,
    beta=0
) where {T, N}
    @workspace size = @argout(
        cudnnGetConvolutionBackwardFilterWorkspaceSize(
            handle(),
            TensorDesc(x),
            TensorDesc(dy),
            ConvDesc(T, cdims),
            FilterDesc(dw),
            cudnnConvolutionBwdFilterAlgo_t(algo),
            out(Ref{Csize_t}())
        )
    )[] workspace -> begin
        cudnnConvolutionBackwardFilter(
            handle(),
            Ref(T(alpha)),
            TensorDesc(x),
            x,
            TensorDesc(dy),
            dy,
            ConvDesc(T, cdims),
            cudnnConvolutionBwdFilterAlgo_t(algo),
            workspace,
            sizeof(workspace),
            Ref(T(beta)),
            FilterDesc(dw),
            dw
        )
    end
    return dw
end

function depthwiseconv!(
    y::CuArray{T},
    x::CuArray{T},
    w::CuArray{T},
    cdims::DepthwiseConvDims;
    alpha=1,
    algo=0
) where T<:CUDNNFloat
    cudnnConvolutionForward(
        y,
        x,
        w,
        cdims;
        alpha=alpha,
        algo=algo
    )
    return y
end

function ∇depthwiseconv_filter!(
    dw::CuArray{T},
    x::CuArray{T},
    dy::CuArray{T},
    cdims::DepthwiseConvDims;
    alpha=1,
    algo=0
) where T<:CUDNNFloat
    cudnnConvolutionBackwardFilter(
        dw,
        x,
        dy,
        cdims;
        alpha=alpha,
        algo=algo
    )
    return dw
end

function ∇depthwiseconv_data!(
    dx::CuArray{T},
    dy::CuArray{T},
    w::CuArray{T},
    cdims::DepthwiseConvDims;
    alpha=1,
    algo=0
) where T<:CUDNNFloat
    cudnnConvolutionBackwardData(
        dx,
        w,
        dy,
        cdims;
        alpha=alpha,
        algo=algo
    )
end


struct PushRecoveryController{M<:MomentumBasedController, S, E<:SE3PDController, L}
    lowlevel::M
    robotmass::Float64
    gravitymag::Float64

    end_effector_tasks::Vector{SpatialAccelerationTask}
    linmomtask::LinearMomentumRateTask
    pelvistask::AngularAccelerationTask
    jointtasks::Dict{JointID, JointAccelerationTask{Revolute{Float64}}}

    statemachine::S
    end_effector_controllers::Vector{E}
    linear_momentum_controller::L
    pelvisgains::PDGains{Diagonal{Float64, SVector{3, Float64}}, Diagonal{Float64, SVector{3, Float64}}}
    jointgains::Dict{JointID, PDGains{Float64, Float64}}

    comref::Point3D{SVector{3, Float64}}
    jointrefs::Dict{JointID, Float64}

    active_contact_points::Dict{BodyID, Vector{SPoint3D{Float64}}}
end

function PushRecoveryController(
        lowlevel::MomentumBasedController,
        pelvis::RigidBody,
        nominalstate::MechanismState,
        statemachine,
        end_effector_controllers::Vector{<:SE3PDController},
        linear_momentum_controller,
        joint_regularization::Float64 = 0.05,
        linear_momentum_weight::Float64 = 1.0,
        pelvisgains::PDGains = default_pelvis_gains(),
        jointgains = Dict(JointID(j) => critically_damped_gains(100.0) for j in tree_joints(lowlevel.state.mechanism)),
        comref::Point3D = center_of_mass(nominalstate) - FreeVector3D(root_frame(lowlevel.state.mechanism), 0., 0., 0.05),
        jointrefs = Dict(JointID(j) => configuration(nominalstate, j)[1] for j in tree_joints(lowlevel.state.mechanism) if joint_type(j) isa Revolute))
    mechanism = lowlevel.state.mechanism
    world = root_body(mechanism)
    worldframe = root_frame(mechanism)
    m = mass(mechanism)
    g = norm(mechanism.gravitational_acceleration)

    regularize!.(Ref(lowlevel), tree_joints(mechanism), joint_regularization)

    end_effector_control_joints = Set{Joint}()
    end_effector_tasks = SpatialAccelerationTask[]
    for controller in end_effector_controllers
        base = findbody(mechanism, controller.base)
        body = findbody(mechanism, controller.body)
        task = SpatialAccelerationTask(mechanism, path(mechanism, base, body))
        push!(end_effector_tasks, task)
        for joint in task.path
            push!(end_effector_control_joints, joint)
        end
        addtask!(lowlevel, task, Parameter(lowlevel.qpmodel, val=controller.weight))
    end

    linmomtask = LinearMomentumRateTask(mechanism, centroidal_frame(lowlevel))
    addtask!(lowlevel, linmomtask, linear_momentum_weight)

    pelvistask = AngularAccelerationTask(mechanism, path(mechanism, world, pelvis))
    addtask!(lowlevel, pelvistask, 10.0)

    revolutejoints = filter(j -> joint_type(j) isa Revolute, tree_joints(mechanism))
    positioncontroljoints = setdiff(revolutejoints, end_effector_control_joints)
    jointtasks = Dict(JointID(j) => JointAccelerationTask(j) for j in positioncontroljoints)
    addtask!.(Ref(lowlevel), collect(values(jointtasks)), 10.0)

    active_contact_points = Dict{BodyID, Vector{SPoint3D{Float64}}}()

    PushRecoveryController(
        lowlevel, m, g,
        end_effector_tasks, linmomtask, pelvistask, jointtasks,
        statemachine, end_effector_controllers, linear_momentum_controller, pelvisgains, jointgains,
        comref, jointrefs,
        active_contact_points)
end

function (controller::PushRecoveryController)(τ::AbstractVector, t::Number, state::MechanismState)
    # Call state machine
    controller.statemachine(t, state)

    # Update active contact points
    update_active_contacts!(controller.active_contact_points, controller.lowlevel.contacts)

    # End effector control
    for i in eachindex(controller.end_effector_tasks)
        setdesired!(controller.end_effector_tasks[i], controller.end_effector_controllers[i](t, state))
    end

    # Linear momentum control
    c = center_of_mass(state)
    h = momentum(state)
    ċ = FreeVector3D(h.frame, linear(h) / controller.robotmass)
    l̇_des = controller.linear_momentum_controller(t, c, ċ, controller.active_contact_points, state)
    world_to_centroidal = Transform3D(c.frame, centroidal_frame(controller.lowlevel), -c.v)
    l̇_des = transform(l̇_des, world_to_centroidal)
    setdesired!(controller.linmomtask, l̇_des)

    # Pelvis orientation control
    pelvis = target(controller.pelvistask.path)
    Hpelvis = transform_to_root(state, pelvis)
    Tpelvis = transform(twist_wrt_world(state, pelvis), inv(Hpelvis))
    ωd_pelvis_des = FreeVector3D(Tpelvis.frame, pd(controller.pelvisgains, rotation(Hpelvis), Tpelvis.angular))
    setdesired!(controller.pelvistask, ωd_pelvis_des)

    # Joint position control
    for jointid in keys(controller.jointtasks)
        task = controller.jointtasks[jointid]
        gains = controller.jointgains[jointid]
        ref = controller.jointrefs[jointid]
        v̇_joint_des = pd(gains, configuration(state, jointid)[1], ref, velocity(state, jointid)[1], 0.0)
        setdesired!(task, v̇_joint_des)
    end

    controller.lowlevel(τ, t, state)
    τ
end

function update_active_contacts!(active_contact_points::AbstractDict{BodyID, <:AbstractVector{<:Point3D}}, contacts)
    for active_points in values(active_contact_points)
        empty!(active_points)
    end
    for (body, points) in contacts
        bodyid = BodyID(body)
        active_points_body = get!(valtype(active_contact_points), active_contact_points, bodyid)
        for contactpoint in points
            if QPControl.isenabled(contactpoint)
                push!(active_points_body, contactpoint.position)
            end
        end
    end
    nothing
end

function default_pelvis_gains()
    gains_xy = critically_damped_gains(100.0)
    gains_z = critically_damped_gains(50.0)
    PDGains(Diagonal(SVector(gains_xy.k, gains_xy.k, gains_z.k)), Diagonal(SVector(gains_xy.d, gains_xy.d, gains_z.d)))
end


# random functions that dispatch either to CURAND or GPUArrays' generic RNG

using Random

export rand_logn!, rand_poisson!

# native RNG

"""
    CUDA.RNG()

A random number generator using `rand()` in a device kernel.

See also: [CUDA.Philox2x32](@ref)
"""
mutable struct RNG <: AbstractRNG
    seed::UInt32
    counter::UInt32

    function RNG(seed::Integer)
        new(seed%UInt32, 0)
    end
end

RNG() = RNG(Random.rand(UInt32))

function Random.seed!(rng::RNG, seed::Integer)
    rng.seed = seed % UInt32
    rng.counter = 0
end

Random.seed!(rng::RNG) = seed!(rng, Random.rand(UInt32))

function Random.rand!(rng::RNG, A::AnyCuArray)
    function kernel(A::AbstractArray{T}, seed::UInt32, counter::UInt32) where {T}
        device_rng = Random.default_rng()

        # initialize the state
        @inbounds Random.seed!(device_rng, seed, counter)

        # grid-stride loop
        threadId = threadIdx().x
        offset = (blockIdx().x - 1) * blockDim().x
        while offset < length(A)
            i = threadId + offset
            if i <= length(A)
                @inbounds A[i] = Random.rand(device_rng, T)
            end

            offset += (blockDim().x - 1) * gridDim().x
        end

        return
    end

    kernel = @cuda launch=false name="rand!" kernel(A, rng.seed, rng.counter)
    config = launch_configuration(kernel.fun; max_threads=64)
    threads = max(32, min(config.threads, length(A)))
    blocks = min(config.blocks, cld(length(A), threads))
    kernel(A, rng.seed, rng.counter; threads=threads, blocks=blocks)

    new_counter = Int64(rng.counter) + length(A)
    overflow, remainder = fldmod(new_counter, typemax(UInt32))
    rng.seed += overflow     # XXX: is this OK?
    rng.counter = remainder

    A
end

# TODO: `randn!`; cannot reuse from Base or RandomNumbers, as those do scalar indexing

# generic functionality

function Random.rand!(rng::Union{RNG,CURAND.RNG,GPUArrays.RNG}, A::AbstractArray{T}) where {T}
    B = CuArray{T}(undef, size(A))
    Random.rand!(rng, B)
    copyto!(A, B)
end

function Random.rand(rng::Union{RNG,CURAND.RNG,GPUArrays.RNG}, T::Type)
    assertscalar("scalar rand")
    A = CuArray{T}(undef, 1)
    Random.rand!(rng, A)
    A[]
end

# RNG-less interface

# the interface is split in two levels:
# - functions that extend the Random standard library, and take an RNG as first argument,
#   will only ever dispatch to CURAND and as a result are limited in the types they support.
# - functions that take an array will dispatch to either CURAND or GPUArrays
# - non-exported functions are provided for constructing GPU arrays from only an eltype

const curand_rng = CURAND.default_rng
gpuarrays_rng() = GPUArrays.default_rng(CuArray)

function seed!(seed=Base.rand(UInt64))
    Random.seed!(curand_rng(), seed)
    Random.seed!(gpuarrays_rng(), seed)
end

# CURAND in-place
Random.rand!(A::CURAND.UniformArray) = Random.rand!(curand_rng(), A)
Random.randn!(A::CURAND.NormalArray; kwargs...) = Random.randn!(curand_rng(), A; kwargs...)
rand_logn!(A::CURAND.LognormalArray; kwargs...) = CURAND.rand_logn!(curand_rng(), A; kwargs...)
rand_poisson!(A::CURAND.PoissonArray; kwargs...) = CURAND.rand_poisson!(curand_rng(), A; kwargs...)

# CURAND out-of-place
rand(T::CURAND.UniformType, dims::Dims) = Random.rand(curand_rng(), T, dims)
randn(T::CURAND.NormalType, dims::Dims; kwargs...) = Random.randn(curand_rng(), T, dims; kwargs...)
rand_logn(T::CURAND.LognormalType, dims::Dims; kwargs...) = CURAND.rand_logn(curand_rng(), T, dims; kwargs...)
rand_poisson(T::CURAND.PoissonType, dims::Dims; kwargs...) = CURAND.rand_poisson(curand_rng(), T, dims; kwargs...)

# support all dimension specifications
rand(T::CURAND.UniformType, dim1::Integer, dims::Integer...) =
    Random.rand(curand_rng(), T, Dims((dim1, dims...)))
randn(T::CURAND.NormalType, dim1::Integer, dims::Integer...; kwargs...) =
    Random.randn(curand_rng(), T, Dims((dim1, dims...)); kwargs...)
rand_logn(T::CURAND.LognormalType, dim1::Integer, dims::Integer...; kwargs...) =
    CURAND.rand_logn(curand_rng(), T, Dims((dim1, dims...)); kwargs...)
rand_poisson(T::CURAND.PoissonType, dim1::Integer, dims::Integer...; kwargs...) =
    CURAND.rand_poisson(curand_rng(), T, Dims((dim1, dims...)); kwargs...)

# GPUArrays in-place
Random.rand!(A::AnyCuArray) = Random.rand!(gpuarrays_rng(), A)
Random.randn!(A::AnyCuArray; kwargs...) =
    error("CUDA.jl does not support generating normally-distributed random numbers of type $(eltype(A))")
# FIXME: GPUArrays.jl has a randn! nowadays, but it doesn't work with e.g. Cuint
rand_logn!(A::AnyCuArray; kwargs...) =
    error("CUDA.jl does not support generating lognormally-distributed random numbers of type $(eltype(A))")
rand_poisson!(A::AnyCuArray; kwargs...) =
    error("CUDA.jl does not support generating Poisson-distributed random numbers of type $(eltype(A))")

# GPUArrays out-of-place
rand(T::Type, dims::Dims) = Random.rand!(CuArray{T}(undef, dims...))
randn(T::Type, dims::Dims; kwargs...) = Random.randn!(CuArray{T}(undef, dims...); kwargs...)
rand_logn(T::Type, dims::Dims; kwargs...) = rand_logn!(CuArray{T}(undef, dims...); kwargs...)
rand_poisson(T::Type, dims::Dims; kwargs...) = rand_poisson!(CuArray{T}(undef, dims...); kwargs...)

# support all dimension specifications
rand(T::Type, dim1::Integer, dims::Integer...) =
    Random.rand!(CuArray{T}(undef, dim1, dims...))
randn(T::Type, dim1::Integer, dims::Integer...; kwargs...) =
    Random.randn!(CuArray{T}(undef, dim1, dims...); kwargs...)
rand_logn(T::Type, dim1::Integer, dims::Integer...; kwargs...) =
    rand_logn!(CuArray{T}(undef, dim1, dims...); kwargs...)
rand_poisson(T::Type, dim1::Integer, dims::Integer...; kwargs...) =
    rand_poisson!(CuArray{T}(undef, dim1, dims...); kwargs...)

# untyped out-of-place
rand(dim1::Integer, dims::Integer...) =
    Random.rand(curand_rng(), Dims((dim1, dims...)))
randn(dim1::Integer, dims::Integer...; kwargs...) =
    Random.randn(curand_rng(), Dims((dim1, dims...)); kwargs...)
rand_logn(dim1::Integer, dims::Integer...; kwargs...) =
    CURAND.rand_logn(curand_rng(), Dims((dim1, dims...)); kwargs...)
rand_poisson(dim1::Integer, dims::Integer...; kwargs...) =
    CURAND.rand_poisson(curand_rng(), Dims((dim1, dims...)); kwargs...)