using DifferentialEquations
import DifferentialEquations.DiffEqBase: AbstractODESolution, HermiteInterpolation,AbstractDiffEqInterpolation

abstract type AbstractSimulation end
export AbstractSimulation

struct Simulation{Q<:AbstractODESolution,W<:AbstractDomain,L<:AbstractArray,G<:Function,G2<:AbstractArray} <: AbstractSimulation
    sol::Q
    domain::W
    names::L
    N::G
    Ns::G2
end

function Simulation(sol::Q,domain::W) where {Q<:AbstractODESolution,W<:AbstractDomain}
    names = getfield.(domain.phase.species,:name)
    Ns = sum(hcat(sol.interp.u...)[domain.indexes[1]:domain.indexes[2],:],dims=1)
    Nderivs = sum(hcat(sol.interp.du...)[domain.indexes[1]:domain.indexes[2],:],dims=1)
    N = HermiteInterpolation(sol.interp.t,Ns,Nderivs)
    F(t::T) where {T<:Real} = N(t,nothing,Val{0},sol.prob.p,:left)
    return Simulation(sol,domain,names,F,Ns)
end

export Simulation

length(p::T) where {T<:AbstractSimulation} = 1
export length

iterate(p::T) where {T<:AbstractSimulation} = p
export iterate

Broadcast.broadcastable(p::T) where {T<:AbstractSimulation} = Ref(p)
export broadcastable

spcindex(bsol::Z,name::Q) where {Z<:Simulation,Q<:AbstractString} = findfirst(isequal(name),getfield.(bsol.domain.phase.species,:name))
export spcindex

function molefractions(bsol::Q,name::W,t::E) where {Q<:AbstractSimulation, W<:String, E<:Real}
    @assert name in bsol.names
    ind = findfirst(isequal(name),bsol.names)
    return bsol.sol(t)[ind]/bsol.N(t)
end

function molefractions(bsol::Q, t::E) where {Q<:AbstractSimulation,E<:Real}
    return bsol.sol(t)[bsol.domain.indexes[1]:bsol.domain.indexes[2]]./bsol.N(t)
end

function molefractions(bsol::Q) where {Q<:AbstractSimulation}
    @views return hcat(bsol.sol.u...)[bsol.domain.indexes[1]:bsol.domain.indexes[2],:]./bsol.Ns
end

export molefractions

getT(bsol::Simulation{Q,W,L,G}, t::K) where {W<:Union{ConstantTPDomain,ConstantTVDomain},K<:Real,Q,G,L} = bsol.domain.T
getT(bsol::Simulation{Q,W,L,G}, t::K) where {W<:Union{ConstantVDomain,ParametrizedVDomain},K<:Real,Q,G,L} = bsol.sol(t)[bsol.domain.indexes[3]]
getT(bsol::Simulation{Q,W,L,G}, t::K) where {W<:Union{ParametrizedTConstantVDomain,ParametrizedTPDomain},K<:Real,Q,G,L} = bsol.domain.T(t)
export getT
getV(bsol::Simulation{Q,W,L,G}, t::K) where {W<:ConstantTPDomain,K<:Real,Q,G,L} = bsol.N(t)*getT(bsol,t)*R /bsol.domain.P
getV(bsol::Simulation{Q,W,L,G}, t::K) where {W<:Union{ConstantVDomain,ConstantTVDomain,ParametrizedTConstantVDomain},K<:Real,Q,G,L} = bsol.domain.V
getV(bsol::Simulation{Q,W,L,G}, t::K) where {W<:ParametrizedVDomain,K<:Real,Q,G,L} = bsol.domain.V(t)
getV(bsol::Simulation{Q,W,L,G}, t::K) where {W<:ParametrizedTPDomain,K<:Real,Q,G,L} = bsol.N(t)*getT(bsol,t)*R /bsol.domain.P(t)
export getV
getP(bsol::Simulation{Q,W,L,G}, t::K) where {W<:ConstantTPDomain,K<:Real,Q,G,L} = bsol.domain.P
getP(bsol::Simulation{Q,W,L,G}, t::K) where {W<:Union{ConstantTVDomain,ParametrizedTConstantVDomain},K<:Real,Q,G,L} = 1.0e6
getP(bsol::Simulation{Q,W,L,G}, t::K) where {W<:Union{ConstantVDomain,ParametrizedVDomain},K<:Real,Q,G,L} = bsol.N(t)*R*getT(bsol,t)/getV(bsol,t)
getP(bsol::Simulation{Q,W,L,G}, t::K) where {W<:ParametrizedTPDomain,K<:Real,Q,G,L} = bsol.domain.P(t)
export getP
getC(bsol::Simulation{Q,W,L,G}, t::K) where {W<:ConstantTPDomain,K<:Real,Q,G,L} = bsol.domain.P/(R*bsol.domain.T)
getC(bsol::Simulation{Q,W,L,G}, t::K) where {W<:Union{ConstantVDomain,ConstantTVDomain,ParametrizedTConstantVDomain},K<:Real,Q,G,L} = bsol.N(t)/bsol.domain.V
getC(bsol::Simulation{Q,W,L,G}, t::K) where {W<:ParametrizedVDomain,K<:Real,Q,G,L} = bsol.N(t)/bsol.domain.V(t)
getC(bsol::Simulation{Q,W,L,G}, t::K) where {W<:ParametrizedTPDomain,K<:Real,Q,G,L} = bsol.domain.P(t)/(R*bsol.domain.T(t))
export getC
"""
calculates the rates of production/loss at a given time point
this outputs a sparse matrix of  num reactions xnum species containing the production/loss
rate of that species associated with that reaction
"""
function rops(bsol::Q,t::X) where {Q<:Simulation,X<:Real}
    ropmat = zeros(length(bsol.domain.phase.reactions),length(bsol.domain.phase.species))
    cs,kfs,krevs = calcthermo(bsol.domain,bsol.sol(t),t)[[2,9,10]]
    @simd for i in 1:length(bsol.domain.phase.reactions)
        rxn = bsol.domain.phase.reactions[i]
        R = getrate(rxn,cs,kfs,krevs)
        for ind in rxn.productinds
            ropmat[i,ind] += R
        end
        for ind in rxn.reactantinds
            ropmat[i,ind] -= R
        end
    end
    return ropmat
end

"""
calculates the rates of production/loss at a given time point for a single species
this outputs a sparse vector of length num reactions containing the production/loss
rate associated with that reaction for the given species
"""
function rops(bsol::Y,name::X,t::Z) where {Y<:Simulation, X<:AbstractString, Z<:Real}
    rop = spzeros(length(bsol.domain.phase.reactions))
    cs,kfs,krevs = calcthermo(bsol.domain,bsol.sol(t),t)[[2,9,10]]
    ind = findfirst(isequal(name),getfield.(bsol.domain.phase.species,:name))
    @assert !isa(ind,Nothing) "species $name not in species array"
    for (i,rxn) in enumerate(bsol.domain.phase.reactions)
        c = 0
        R = getrate(rxn,cs,kfs,krevs)
        c -= count(isequal(ind),rxn.reactantinds)
        c += count(isequal(ind),rxn.productinds)
        if c != 0
            rop[i] = c*R
        end
    end
    return rop
end
export rops

function getconcentrationsensitivity(bsol::Simulation{Q,W,L,G}, numerator::String, denominator::String, t::K) where {W<:Union{ConstantVDomain,ConstantTVDomain},K<:Real,Q,G,L}
    @assert numerator in bsol.names
    @assert denominator in bsol.names
    indnum = findfirst(isequal(numerator),bsol.names)
    inddeno = findfirst(isequal(denominator),bsol.names)
    Nvars = bsol.domain.indexes[end]-bsol.domain.indexes[1]+1
    Nrxns = length(bsol.domain.phase.reactions)
    arr = bsol.sol(t)
    s = arr[Nvars+Nrxns*Nvars+(inddeno-1)*Nvars+indnum]
    val = s/arr[indnum] #constant volume
    if t == 0
        return 0.0
    else
        return val
    end
end

function getconcentrationsensitivity(bsol::Simulation{Q,W,L,G}, numerator::String, denominator::String, t::K) where {W<:ConstantTPDomain,K<:Real,Q,G,L}
    @assert numerator in bsol.names
    @assert denominator in bsol.names
    indnum = findfirst(isequal(numerator),bsol.names)
    inddeno = findfirst(isequal(denominator),bsol.names)
    Nvars = bsol.domain.indexes[end]-bsol.domain.indexes[1]+1
    Nrxns = length(bsol.domain.phase.reactions)
    arr = bsol.sol(t)
    s = arr[Nvars+Nrxns*Nvars+(inddeno-1)*Nvars+indnum]
    V = getV(bsol,t)
    c = arr[indnum]/V
    val =  (s-c*sum(arr[Nvars+Nrxns*Nvars+(inddeno-1)*Nvars+1:Nvars+Nrxns*Nvars+inddeno*Nvars])*R*bsol.domain.T/bsol.domain.P)/(c*V) #known T and P
    if t == 0
        return 0.0
    else
        return val
    end
end

function getconcentrationsensitivity(bsol::Simulation{Q,W,L,G}, numerator::String, denominator::Z, t::K) where {W<:Union{ConstantVDomain,ConstantTVDomain},K<:Real,Z<:Integer,Q,G,L}
    @assert numerator in bsol.names
    indnum = findfirst(isequal(numerator),bsol.names)
    inddeno = denominator
    Nvars = bsol.domain.indexes[end]-bsol.domain.indexes[1]+1
    Nrxns = length(bsol.domain.phase.reactions)
    arr = bsol.sol(t)
    s = arr[Nvars+(inddeno-1)*Nvars+indnum]
    T = getT(bsol,t)
    P = getP(bsol,t)
    C = getC(bsol,t)
    k = bsol.domain.phase.reactions[inddeno].kinetics(T=T,P=P,C=C)
    val = s*k/arr[indnum] #constant volume
    if t == 0
        return 0.0
    else
        return val
    end
end

function getconcentrationsensitivity(bsol::Simulation{Q,W,L,G}, numerator::String, denominator::Z, t::K) where {W<:ConstantTPDomain,K<:Real,Z<:Integer,Q,G,L}
    @assert numerator in bsol.names
    indnum = findfirst(isequal(numerator),bsol.names)
    inddeno = denominator
    Nvars = bsol.domain.indexes[end]-bsol.domain.indexes[1]+1
    Nrxns = length(bsol.domain.phase.reactions)
    arr = bsol.sol(t)
    s = arr[Nvars+(inddeno-1)*Nvars+indnum]
    V = getV(bsol,t)
    T = getT(bsol,t)
    P = getP(bsol,t)
    C = getC(bsol,t)
    c = arr[indnum]/V
    k = bsol.domain.phase.reactions[inddeno].kinetics(T=T,P=P,C=C)
    val = k*(s-c*sum(arr[Nvars+(inddeno-1)*Nvars+1:Nvars+inddeno*Nvars])*R*bsol.domain.T/bsol.domain.P)/(c*V) #known T and P
    if t == 0
        return 0.0
    else
        return val
    end
end

export getconcentrationsensitivity

"""
calculate the rates of all reactions at time t
"""
function rates(bsol::Q,t::X) where {Q<:Simulation,X<:Real}
    cs,kfs,krevs = calcthermo(bsol.domain,bsol.sol(t),t)[[2,9,10]]
    return [getrate(rxn,cs,kfs,krevs) for rxn in bsol.domain.phase.reactions]
end

"""
calculate the rates of all reactions at given times ts
defaults to using bsol.sol.t if ts is not supplied
"""
function rates(bsol::Q;ts::X=Array{Float64,1}()) where {Q<:Simulation,X<:AbstractArray}
    if length(ts) == 0
        ts = bsol.sol.t
    end
    return hcat([rates(bsol,t) for t in ts]...)
end

export rates


# name = "Alpine"
# Alpine = "07493b3f-dabb-5b16-a503-4139292d7dd4"

using Alpine, Ipopt, Cbc

const ALP_IPOPT = optimizer_with_attributes(Ipopt.Optimizer, MOI.Silent() => true, "sb" => "yes")
const ALP_CBC = optimizer_with_attributes(Cbc.Optimizer, MOI.Silent() => true)

# ALP_OPT = optimizer_with_attributes(Alpine.Optimizer, "nlp_solver" => ALP_IPOPT,
# "mip_solver" => ALP_CBC,
# "loglevel" => 100)
# ALP = ALP_OPT

ALP_OPT = Alpine.Optimizer()
MOI.set(ALP_OPT, MOI.RawParameter("nlp_solver"), ALP_IPOPT)
MOI.set(ALP_OPT, MOI.RawParameter("mip_solver"), ALP_CBC)
ALP = MOI.Bridges.Constraint.SOCtoNonConvexQuad{Float64}(ALP_OPT)

push!(solvers_nlp_lowtol, (opt = ALP, mode = BilevelJuMP.ProductMode(1e-5)))
push!(solvers_nlp, (opt = ALP, mode = BilevelJuMP.ProductMode(1e-9)))
push!(solvers_nlp, (opt = ALP, mode = BilevelJuMP.ProductMode(1e-9, with_slack = true)))

push!(solvers, (opt = ALP, mode = BilevelJuMP.ProductMode(1e-9)))
push!(solvers_cached, (opt = ALP, mode = BilevelJuMP.ProductMode(1e-9)))
push!(solvers_quad, (opt = ALP, mode = BilevelJuMP.ProductMode(1e-9)))

push!(solvers_nlp_sd, (opt = ALP, mode = BilevelJuMP.StrongDualityMode(1e-9, inequality = true)))
push!(solvers_nlp_sd, (opt = ALP, mode = BilevelJuMP.StrongDualityMode(inequality = false)))
push!(solvers_nlp_sd_i, (opt = ALP, mode = BilevelJuMP.StrongDualityMode(1e-9, inequality = true)))
push!(solvers_nlp_sd_e, (opt = ALP, mode = BilevelJuMP.StrongDualityMode(inequality = false)))


# Quasistatic analysis for beam, including several timesteps

using JuliaFEM
using JuliaFEM.Preprocess
using JuliaFEM.Postprocess
using JuliaFEM.Abaqus: create_surface_elements

Logging.configure(level=Logging.DEBUG)

function create_body(mesh, name, E, nu, rho)
    body = Problem(mesh, Elasticity, name, 3)
    update!(body, "youngs modulus", E)
    update!(body, "poissons ratio", nu)
    update!(body, "density", rho)
    return body
end

""" Create boundary condition from surface set. """
function create_bc_from_surface_set(mesh, name, u1, u2, u3)
    bc = Problem(Dirichlet, string(name), 3, "displacement")
    bc.elements = create_surface_elements(mesh, name)
    update!(bc, "displacement 1", u1)
    update!(bc, "displacement 2", u2)
    update!(bc, "displacement 3", u3)
    return bc
end

""" Create boundary condition from node set. """
function create_bc_from_node_set(mesh, name, u1, u2, u3)
    bc = Problem(Dirichlet, string(name), 3, "displacement")
    bc.elements = [Element(Poi1, [nid]) for nid in mesh.node_sets[name]]
    update!(bc, "geometry", mesh.nodes)
    update!(bc, "displacement 1", u1)
    update!(bc, "displacement 2", u2)
    update!(bc, "displacement 3", u3)
    return bc
end

function create_bc(mesh, name, u1=0.0, u2=0.0, u3=0.0)
    name = Symbol(name)
    if haskey(mesh.surface_sets, name)
        return create_bc_from_surface_set(mesh, name, u1, u2, u3)
    elseif haskey(mesh.node_sets, name)
        return create_bc_from_node_set(mesh, name, u1, u2, u3)
    else
        error("Mesh does not contain node or surface set $name")
    end
end

function create_interface(mesh, slave_surface::String, master_surface::String)
    interface = Problem(Mortar, "interface between $slave_surface and $master_surface", 3, "displacement")
    interface.properties.dual_basis = true
    slave_elements = create_surface_elements(mesh, Symbol(slave_surface))
    master_elements = create_surface_elements(mesh, Symbol(master_surface))
    nslaves = length(slave_elements)
    nmasters = length(master_elements)
    info("$nslaves slaves, $nmasters masters")
    update!(slave_elements, "master elements", master_elements)
    interface.elements = [slave_elements; master_elements]
    return interface
end

function create_interface(mesh, slave::Problem, master::Problem)
    slave_surface = slave.name * "_TO_" * master.name
    master_surface = master.name * "_TO_" * slave.name
    return create_interface(mesh, slave_surface, master_surface)
end

""" Convert Mesh object from quadratic to linear. """
function to_linear!(mesh)
    mapping = Dict(:Tet10 => :Tet4, :Tri6 => :Tri3)
    nnodes = Dict(:Tet4 => 4, :Tri3 => 3)
    for elid in keys(mesh.elements)
        eltype = mesh.element_types[elid]
        if haskey(mapping, eltype)
            mesh.element_types[elid] = mapping[eltype]
            nnodes_new = nnodes[mapping[eltype]]
            mesh.elements[elid] = mesh.elements[elid][1:nnodes_new]
        end
    end
end

# start of simulation
#mesh = abaqus_read_mesh("beam_3d_1st_order_tetra_30mm.inp")
mesh = abaqus_read_mesh("beam_3d_2nd_order_tetra_30mm.inp")
#to_linear!(mesh)
info("element sets = ", collect(keys(mesh.element_sets)))
info("surface sets = ", collect(keys(mesh.surface_sets)))

# parts
beam = create_body(mesh, "beam", 210.0e3, 0.3, 7.85e-9)

# boundary conditions
bc1 = create_bc(mesh, "fixed")

# load
load = Problem(Elasticity, "pressure load", 3)
load.elements = create_surface_elements(mesh, :load)
update!(load, "surface pressure", 0.0 => 0.0)
update!(load, "surface pressure", 1.0 => 20.0)

# give some information about surface load
nload = length(load.elements)
info("$nload elements in load surface")
area = 0.0
time = 0.0
for element in load.elements
    for ip in get_integration_points(element)
        detJ = element(ip, time, Val{:detJ})
        area += ip.weight*detJ
    end
end
info("load surface area: $area")

# solution
isfile("results_quasistatic.h5") && rm("results_quasistatic.h5")
isfile("results_quasistatic.xmf") && rm("results_quasistatic.xmf")
solver = Solver(Linear, beam, bc1, load)
solver.xdmf = Xdmf("results")

for time in [0.0, 1.0]
    solver.time = time
    solver()
end

# calculate stresses in integration points and use least-squares fitting to
# extrapolate results to nodes

""" Return stress tensor. """
function get_stress_tensor(element, ip, time)
    haskey(element, "displacement") || return nothing
    gradu = element("displacement", ip, time, Val{:Grad})
    eps = 0.5*(gradu' + gradu)
    E = element("youngs modulus", ip, time)
    nu = element("poissons ratio", ip, time)
    mu = E/(2.0*(1.0+nu))
    la = E*nu/((1.0+nu)*(1.0-2.0*nu))
    S = la*trace(eps)*I + 2.0*mu*eps
    return S
end

""" Return stress vector in "ABAQUS" order 11, 22, 33, 12, 23, 13. """
function get_stress(element, ip, time)
    S = get_stress_tensor(element, ip, time)
    return [S[1,1], S[2,2], S[3,3], S[1,2], S[2,3], S[1,3]]
end

""" Return principal stresses. """
function get_stress_principal(element, ip, time)
    S = get_stress_tensor(element, ip, time)
    return sort(eigvals(S))
end

""" Make least squares fit for some field to nodes. """
function lsq_fit(elements, field, time)
    A = SparseMatrixCOO()
    b = SparseMatrixCOO()
    volume = 0.0
    for element in elements
        gdofs = get_connectivity(element)
        # increase integration order by 1 from default
        for ip in get_integration_points(element, 1)
	        detJ = element(ip, time, Val{:detJ})
	        w = ip.weight*detJ
	        N = element(ip, time)
	        f = field(element, ip, time)
	        add!(A, gdofs, gdofs, w*kron(N', N))
	        for i=1:length(f)
                add!(b, gdofs, w*f[i]*N, i)
            end
            volume += w
	    end
    end
    info("Mass matrix for least-squares fit is assembled. Total volume to fit: $volume")
    A = sparse(A)
    b = sparse(b)
    A = 1/2*(A + A')

    SparseArrays.droptol!(A, 1.0e-6)
    SparseArrays.dropzeros!(A)
    nz = get_nonzero_rows(A)
    F = ldltfact(A[nz,nz])

    x = zeros(size(b)...)
    x[nz, :] = F \ b[nz, :]

    nodal_values = Dict(i => vec(x[i,:]) for i in nz)
    return nodal_values
end

xdmf = get(solver.xdmf)

using JuliaFEM: new_dataitem, new_child, set_attribute, add_child, save!

""" Store values to xml/h5. """
function add_field_to_xdmf!(xdmf, field_name, xmlpath, xdmfpath, data::Dict; field_type="Vector")
    node_ids = sort(collect(keys(data)))
    info("Storing field $field_name to xdmfpath $xdmfpath, xmlpath $xmlpath")
    for (i, nid) in enumerate(node_ids)
        println("$nid -> $(data[nid])")
        if i > 2
            println("...")
            break
        end
    end
    datavec = hcat([data[nid] for nid in node_ids]...)
    dataitem = new_dataitem(xdmf, xdmfpath, datavec)
    frame = read(xdmf, xmlpath)
    attribute = new_child(frame, "Attribute")
    set_attribute(attribute, "Name", field_name)
    set_attribute(attribute, "Center", "Node")
    set_attribute(attribute, "AttributeType", field_type)
    add_child(attribute, dataitem)
end

# calculate and update principle stresses to xdmf file, also calculate minimum
# and maximum principle components over time

minp = Dict()
maxp = Dict()

for (i, time) in enumerate([0.0, 1.0])
    S = lsq_fit(beam.elements, get_stress, time)
    #Sp = lsq_fit(beam.elements, get_stress_principal)

    # Calculate principal stresses in nodes
    Sp = Dict()
    for (nid, s) in S
        # order is: 11, 22, 33, 12, 23, 13
        stress_tensor = [
            s[1] s[4] s[6]
            s[4] s[2] s[5]
            s[6] s[5] s[3]]
        principle = sort(eigvals(stress_tensor))
        Sp[nid] = principle

        # minimum and maximum of principle stress over time
        if haskey(minp, nid)
            minp[nid] = min(minp[nid], minimum(principle))
        else
            minp[nid] = minimum(principle)
        end
        if haskey(maxp, nid)
            maxp[nid] = max(maxp[nid], maximum(principle))
        else
            maxp[nid] = maximum(principle)
        end
    end

    xml_path = "/Domain/Grid/Grid[$i]"
    xdmf_path = "/Results/Time $time/Nodal Fields/Stress"
    add_field_to_xdmf!(xdmf, "Stress", xml_path, xdmf_path, S; field_type="Tensor6")
    xdmf_path = "/Results/Time $time/Nodal Fields/Principal Stress"
    add_field_to_xdmf!(xdmf, "Principal Stress", xml_path, xdmf_path, Sp; field_type="Vector")
end

save!(xdmf)
close(xdmf.hdf)

# calculate mean stress and stress amplitude

Sa = Dict()
Sm = Dict()
node_ids = keys(minp)
for nid in node_ids
    Sa[nid] = (maxp[nid] - minp[nid]) / 2
    Sm[nid] = (maxp[nid] + minp[nid]) / 2
end

# create a new xdmf result output and save mean stress and stress amplitude there

using JuliaFEM: new_element, get_all_elements, get_element_type,
                filter_by_element_type, get_element_id

isfile("results_cycle.h5") && rm("results_cycle.h5")
isfile("results_cycle.xmf") && rm("results_cycle.xmf")
xdmf = Xdmf("results_cycle")
domain = new_child(xdmf, "Domain")
frame = new_child(domain, "Grid")

# 1. save geometry
X_ = solver("geometry", solver.time)
node_ids = sort(collect(keys(X_)))
X = hcat([X_[nid] for nid in node_ids]...)
data_node_ids = new_dataitem(xdmf, "/Node IDs", node_ids)
geometry = new_element("Geometry", Dict("Type" => "XYZ"))
add_child(geometry, new_dataitem(xdmf, "/Geometry", X))
add_child(frame, geometry)

# 2. save topology
nid_mapping = Dict(j=>i for (i, j) in enumerate(node_ids))
all_elements = get_all_elements(solver)
nelements = length(all_elements)
debug("Saving topology: $nelements elements total.")
element_types = unique(map(get_element_type, all_elements))

xdmf_element_mapping = Dict(
    "Poi1" => "Polyvertex",
    "Seg2" => "Polyline",
    "Tri3" => "Triangle",
    "Quad4" => "Quadrilateral",
    "Tet4" => "Tetrahedron",
    "Pyramid5" => "Pyramid",
    "Wedge6" => "Wedge",
    "Hex8" => "Hexahedron",
    "Seg3" => "Edge_3",
    "Tri6" => "Tri_6",
    "Quad8" => "Quad_8",
    "Tet10" => "Tet_10",
    "Pyramid13" => "Pyramid_13",
    "Wedge15" => "Wedge_15",
    "Hex20" => "Hex_20")

for element_type in element_types
    elements = filter_by_element_type(element_type, all_elements)
    nelements = length(elements)
    info("Xdmf save: $nelements elements of type $element_type")
    sort!(elements, by=get_element_id)
    element_ids = map(get_element_id, elements)
    element_conn = map(element -> [nid_mapping[j]-1 for j in get_connectivity(element)], elements)
    element_conn = hcat(element_conn...)
    element_code = split(string(element_type), ".")[end]
    dataitem = new_dataitem(xdmf, "/Topology/$element_code/Element IDs", element_ids)
    dataitem = new_dataitem(xdmf, "/Topology/$element_code/Connectivity", element_conn)
    topology = new_element("Topology")
    set_attribute(topology, "TopologyType", xdmf_element_mapping[element_code])
    set_attribute(topology, "NumberOfElements", length(elements))
    add_child(topology, dataitem)
    add_child(frame, topology)
end

add_field_to_xdmf!(xdmf, "Stress (Amplitude)", "/Domain/Grid", "/Results/Stress (Amplitude)", Sa; field_type="Scalar")
add_field_to_xdmf!(xdmf, "Stress (Mean)", "/Domain/Grid", "/Results/Stress (Mean", Sm; field_type="Scalar")

save!(xdmf)
close(xdmf.hdf)


# # Root Mean Squared Fluctuation (RMSF)
#
# md # [![](https://mybinder.org/badge_logo.svg)](@__BINDER_ROOT_URL__cookbook/notebooks/03_RMSF.ipynb)
# md # [![](https://img.shields.io/badge/show-nbviewer-579ACA.svg)](@__NBVIEWER_ROOT_URL__cookbook/notebooks/03_RMSF.ipynb)
#
#
# ## Problem description
#
# The [Root Mean Squared Fluctuation (RMSF)](https://en.wikipedia.org/wiki/Mean_squared_displacement)
# is a common way to measure residue flexibility in a structural ensemble.
# It is a measure of how far is the residue moving from its average position
# in the group of structures. Usually, we represent a residue position with
# the spatial coordinates of its alpha carbon.
#
# The protein structures should be previously superimposed to calculate the
# RMSF, for example, by using the `superimpose` function of the
# [`PDB` module of `MIToS`](@ref Module-PDB). In this example, we are going
# to measure the RMSF of each residue from an NMR ensemble using the
# `rmsf` function.
#
# The structure superimposition could be the most complicated step of the
# process, depending on the input data. In particular, it structures come
# from different PDB structures or homologous proteins can require the use
# of external programs,
# as [MAMMOTH-mult](https://ub.cbm.uam.es/software/online/mamothmult.php) or
# [MUSTANG](https://lcb.infotech.monash.edu/mustang/) among others,
# tailored for this task.
#
# In this case, we are going to use an NMR ensemble. Therefore, we are not
# going to need to superimpose the structures as NMR models have the
# same protein sequence and are, usually, well-aligned.
#
#
# ## MIToS solution

import MIToS
using MIToS.PDB
using Plots

# Lets read the NMR ensemble:

pdb_file   = abspath(pathof(MIToS), "..", "..", "test", "data", "1AS5.pdb")
pdb_res = read(pdb_file, PDBFile, occupancyfilter=true)
# md nothing # hide

# We set `occupancyfilter` to `true` to ensure that we have one single set of
# coordinates for each atom. That filter isn't essential for NMR structures,
# but It can avoid multiple alpha carbons in crystallographic structures with
# disordered atoms.
# We can get an idea of the alpha carbon positions by plotting these residues:

scatter(pdb_res, legend=false)

# As we saw in the previous plot, the structure doesn't need to be
# superimposed. Now, we are going to separate each model into different
# vectors, storing each vector into a `Dict`:

models = Dict{String,Vector{PDBResidue}}()
for res in pdb_res
	push!(get!(models, res.id.model, []), res)
end

# Then, we simply need to collect all the PDB models in the values
# of the `Dict`, to get the vector of `PDBResidues` vectors required
# to calculate the RMSF.

pdb_models = collect(values(models))
# md nothing # hide

# And, finally, call the `rmsf` function on the list of structures. It is
# important that all the vectors has the same number of `PDBResidue`s. This
# function assumes that the nth element of each vector corresponds to the same
# residue:

RMSF = rmsf(pdb_models)

# This return the vector of RMSF values for each residue, calculated using
# the coordinates of the alpha carbons.
# You can plot this vector to get an idea of the which are the most flexible
# position in your structure:

plot(RMSF, legend=false, xlab="Residue", ylab="RMSF [Å]")


#File which defined all the algorithms. Each algorithm takes in an ABCtype

"""
    runabc(ABCsetup::ABCtype, targetdata; progress = false, verbose = false, parallel = true)

Run ABC with ABCsetup defining the algorithm and inputs to algorithm, targetdata is the data we wish to fit the model to and will be used as an input for the simulation function defined in ABCsetup. If progress is set to `true` a progress meter will be shown. Inference will be run in parallel via multithreading if `parallel = true`. The environmental variable JULIA_NUM_THREADS needs to be set prior to launching a julia session.
"""
function runabc(ABCsetup::ABCRejection, targetdata; progress = false, verbose = false, parallel = false)

  #initalize array of particles
  particlesall = Array{ParticleRejection}(undef, ABCsetup.maxiterations)

  if progress
    p = Progress(ABCsetup.nparticles, 1, "Running ABC rejection algorithm...", 30)
  end

  if parallel
    Printf.@printf("Preparing to run in parallel on %i processors\n", nthreads())

    particles = Array{ParticleRejection}(undef, ABCsetup.maxiterations)
    distvec = zeros(Float64, ABCsetup.maxiterations) #store distances in an array
    i = Atomic{Int64}(0)
    cntr = Atomic{Int64}(0)
    @threads for its = 1:ABCsetup.maxiterations

      if i[] > ABCsetup.nparticles
        break
      end

      #get new proposal parameters
      newparams = getproposal(ABCsetup.prior, ABCsetup.nparams)
      #simulate with new parameters
      dist, out = ABCsetup.simfunc(newparams, ABCsetup.constants, targetdata)
      #keep track of all particles incase we don't reach nparticles with dist < ϵ
      particlesall[its] = ParticleRejection(newparams, dist, out)

      #if simulated data is less than target tolerance accept particle
      if dist < ABCsetup.ϵ
        particles[its] = ParticleRejection(newparams, dist, out)
        distvec[its] = dist
        atomic_add!(i, 1)
      end
      atomic_add!(cntr,1)

    end
    # Remove particles that are still #undef and corresponding distances
    idx = [isassigned(particles,ii) for ii in eachindex(particles)]
    particles = particles[idx]
    distvec = distvec[idx]
    i = length(particles)    # Number of accepted particles
    its = cntr[]    # Total number of simulations

  else
    Printf.@printf("Preparing to run in serial on %i processor\n", 1)

    particles = Array{ParticleRejection}(undef, ABCsetup.nparticles)
    distvec = zeros(Float64, ABCsetup.nparticles) #store distances in an array
    i = 1 #set particle indicator to 1
    its = 0 #keep track of number of iterations
    while (i < (ABCsetup.nparticles + 1)) & (its < ABCsetup.maxiterations)

      its += 1
      #get new proposal parameters
      newparams = getproposal(ABCsetup.prior, ABCsetup.nparams)
      #simulate with new parameters
      dist, out = ABCsetup.simfunc(newparams, ABCsetup.constants, targetdata)
      #keep track of all particles incase we don't reach nparticles with dist < ϵ
      particlesall[its] = ParticleRejection(newparams, dist, out)

      #if simulated data is less than target tolerance accept particle
      if dist < ABCsetup.ϵ
        particles[i] = ParticleRejection(newparams, dist, out)
        distvec[i] = dist
        i +=1
        if progress
          next!(p)
        end
      end
    end
    i -= 1    # Correct to total number of particels
  end

  if i < ABCsetup.nparticles
    @warn "Only accepted $(i) particles with ϵ < $(ABCsetup.ϵ). \n\tYou may want to increase ϵ or increase maxiterations. \n\t Resorting to taking the $(ABCsetup.nparticles) particles with smallest distance"
    d = map(p -> p.distance, particlesall)
    particles = particlesall[sortperm(d)[1:ABCsetup.nparticles]]
    distvec = map(p -> p.distance, particles)
  elseif i > ABCsetup.nparticles
    particles = particles[1:ABCsetup.nparticles]
    distvec = distvec[1:ABCsetup.nparticles]
  end

  out = ABCrejectionresults(particles, its, ABCsetup, distvec)
  return out
end

function runabc(ABCsetup::ABCRejectionModel, targetdata; progress = false, verbose = false)

  ABCsetup.nmodels > 1 || error("Only 1 model specified, use ABCRejection method to estimate parameters for a single model")

  #initalize array of particles
  particles = Array{ParticleRejectionModel}(undef,
                ABCsetup.Models[1].nparticles)

  i = 1 #set particle indicator to 1
  its = 0 #keep track of number of iterations
  distvec = zeros(Float64, ABCsetup.Models[1].nparticles) #store distances in an array

  if progress
    p = Progress(ABCsetup.Models[1].nparticles, 1, "Running ABC rejection algorithm...", 30)
  end

  while (i < (ABCsetup.Models[1].nparticles + 1)) & (its < ABCsetup.Models[1].maxiterations)

    its += 1
    #sample uniformly from models
    model = rand(1:ABCsetup.nmodels)
    #get new proposal parameters
    newparams = getproposal(ABCsetup.Models[model].prior, ABCsetup.Models[model].nparams)
    #simulate with new parameters
    dist, out = ABCsetup.Models[model].simfunc(newparams, ABCsetup.Models[model].constants, targetdata)

    #if simulated data is less than target tolerance accept particle
    if dist < ABCsetup.Models[1].ϵ
      particles[i] = ParticleRejectionModel(newparams, model, dist, out)
      distvec[i] = dist
      i +=1
      if progress
        next!(p)
      end
    end
  end

  i > ABCsetup.Models[1].nparticles || error("Only accepted $(i) particles with ϵ < $(ABCsetup.Models[1].ϵ). \n\tDecrease ϵ or increase maxiterations ")

  out = ABCrejectionmodelresults(particles, its, ABCsetup, distvec)
  return out
end

function runabc(ABCsetup::ABCSMC, targetdata; verbose = false, progress = false, parallel = false)

  #run first population with parameters sampled from prior
  if verbose
    println("##################################################")
    println("Use ABC rejection to get first population")
  end
  ABCrejresults = runabc(ABCRejection(ABCsetup.simfunc, ABCsetup.nparams,
                  ABCsetup.ϵ1, ABCsetup.prior; nparticles = ABCsetup.nparticles,
                  maxiterations = ABCsetup.maxiterations,
                  constants = ABCsetup.constants), targetdata,
                  progress = progress, parallel = parallel);

  oldparticles, weights = setupSMCparticles(ABCrejresults, ABCsetup)
  ABCsetup.kernel.kernel_parameters = (maximum(ABCrejresults.parameters, dims = 1) - minimum(ABCrejresults.parameters, dims = 1) ./2)[:]
  ϵ = quantile(ABCrejresults.dist, ABCsetup.α) # set new ϵ to αth quantile
  ϵvec = [ϵ] #store epsilon values
  numsims = [ABCrejresults.numsims] #keep track of number of simualtions
  particles = Array{ParticleSMC}(undef, ABCsetup.nparticles) #define particles array

  if verbose
    println("Running ABC SMC... \n")
  end

  if parallel
    Printf.@printf("Preparing to run in parallel on %i processors\n", nthreads())
  else
    Printf.@printf("Preparing to run in serial on %i processor\n", 1)
  end

  popnum = 1
  finalpop = false

  if sum(ABCrejresults.dist .< ABCsetup.ϵT) == ABCsetup.nparticles
      @warn "Target ϵ reached with ABCRejection algorithm, no need to use ABC SMC algorithm, returning ABCRejection output..."
      return ABCrejresults
  end

  while (ϵ > ABCsetup.ϵT) & (sum(numsims) < ABCsetup.maxiterations)

    if progress
      p = Progress(ABCsetup.nparticles, 1, "ABC SMC population $(popnum), new ϵ: $(round(ϵ, digits = 2))...", 30)
    end

    if parallel

      # Arrays initialised with length maxiterations to enable use of unique index ii
      particles = Array{ParticleSMC}(undef, ABCsetup.maxiterations)
      distvec = zeros(Float64, ABCsetup.maxiterations)
      i = Atomic{Int64}(0)
      its = Atomic{Int64}(0)

      @threads for ii = 1:ABCsetup.maxiterations

        if i[] > ABCsetup.nparticles
          break
        end

        j = wsample(1:ABCsetup.nparticles, weights)
        particle = oldparticles[j]
        newparticle = perturbparticle(particle, ABCsetup.kernel)
        priorp = priorprob(newparticle.params, ABCsetup.prior)
        if priorp == 0.0 #return to beginning of loop if prior probability is 0
          continue
        end

        #simulate with new parameters
        dist, out = ABCsetup.simfunc(newparticle.params, ABCsetup.constants, targetdata)

        #if simulated data is less than target tolerance accept particle
        if dist < ϵ
          particles[ii] = newparticle
          distvec[ii] = dist
          particles[ii].other = out
          particles[ii].distance = dist
          atomic_add!(i, 1)
        end

        atomic_add!(its,1)
      end

      # Remove particles that are still #undef and corresponding distances
      idx = [isassigned(particles,ii) for ii in eachindex(particles)]
      particles = particles[idx][1:ABCsetup.nparticles]
      distvec = distvec[idx][1:ABCsetup.nparticles]
      its = its[]

    else
      particles = Array{ParticleSMC}(undef, ABCsetup.nparticles)
      distvec = zeros(Float64, ABCsetup.nparticles)
      i = 1 #set particle indicator to 1
      its = 1

      while i < ABCsetup.nparticles + 1

        j = wsample(1:ABCsetup.nparticles, weights)
        particle = oldparticles[j]
        newparticle = perturbparticle(particle, ABCsetup.kernel)
        priorp = priorprob(newparticle.params, ABCsetup.prior)
        if priorp == 0.0 #return to beginning of loop if prior probability is 0
          continue
        end

        #simulate with new parameters
        dist, out = ABCsetup.simfunc(newparticle.params, ABCsetup.constants, targetdata)

        #if simulated data is less than target tolerance accept particle
        if dist < ϵ
          particles[i] = newparticle
          particles[i].other = out
          particles[i].distance = dist
          distvec[i] = dist
          i += 1
          if progress
            next!(p)
          end
        end
        its += 1
      end
    end

    particles, weights = smcweights(particles, oldparticles, ABCsetup.prior, ABCsetup.kernel)
    ABCsetup.kernel.kernel_parameters = ABCsetup.kernel.calculate_kernel_parameters(particles)
    oldparticles = particles

    if finalpop == true
      break
    end

    ϵ = quantile(distvec, ABCsetup.α)

    if ϵ < ABCsetup.ϵT
      ϵ = ABCsetup.ϵT
      push!(ϵvec, ϵ)
      push!(numsims, its)
      popnum = popnum + 1
      finalpop = true
      continue
    end

    push!(ϵvec, ϵ)
    push!(numsims, its)

    if ((( abs(ϵvec[end - 1] - ϵ )) / ϵvec[end - 1]) < ABCsetup.convergence) == true
      if verbose
        println("\nNew ϵ is within $(round(ABCsetup.convergence * 100, digits=2))% of previous population, stop ABC SMC\n")
      end
      break
    end

    popnum = popnum + 1

    if verbose
      println("##################################################")
      show(ABCSMCresults(particles, numsims, ABCsetup, ϵvec))
      println("##################################################\n")
    end
  end
  if sum(numsims) >= ABCsetup.maxiterations
    if verbose
      println("\nPassed maxiterations=$(ABCsetup.maxiterations), stop ABC SMC\n")
    end
  end

  out = ABCSMCresults(particles, numsims, ABCsetup, ϵvec)
  return out
end

"""
    runabc(ABCsetup::ABCtype, targetdata; progress = false, verbose = false)

When the SMC algorithms are used, a print out at the end of each population will be made if verbose = true.
"""
function runabc(ABCsetup::ABCSMCModel, targetdata; verbose = false, progress = false)

  ABCsetup.nmodels > 1 || error("Only 1 model specified, use ABCSMC method to estimate parameters for a single model")

  #run first population with parameters sampled from prior
  if verbose
    println("##################################################")
    println("Use ABC rejection to get first population")
  end
  ABCrejresults = runabc(ABCRejectionModel(
            map(x -> x.simfunc, ABCsetup.Models),
            map(x -> x.nparams, ABCsetup.Models),
            ABCsetup.Models[1].ϵ1,
            map(x -> x.prior, ABCsetup.Models),
            constants = map(x -> x.constants, ABCsetup.Models),
            nparticles = ABCsetup.Models[1].nparticles,
            maxiterations = ABCsetup.Models[1].maxiterations),
            targetdata);

  oldparticles, weights = setupSMCparticles(ABCrejresults, ABCsetup)
  ϵ = quantile(ABCrejresults.dist, ABCsetup.α) # set new ϵ to αth quantile
  ϵvec = [ϵ] #store epsilon values
  numsims = [ABCrejresults.numsims] #keep track of number of simulations
  particles = Array{ParticleSMCModel}(undef, ABCsetup.nparticles) #define particles array
  weights, modelprob = getparticleweights(oldparticles, ABCsetup)
  ABCsetup = modelselection_kernel(ABCsetup, oldparticles)

  modelprob = ABCrejresults.modelfreq

  if verbose
    println("Run ABC SMC \n")
  end

  popnum = 1
  finalpop = false

  if verbose
    println(ABCSMCmodelresults(oldparticles, numsims, ABCsetup, ϵvec))
  end

  if sum(ABCrejresults.dist .< ABCsetup.ϵT) == ABCsetup.nparticles
      @warn "Target ϵ reached with ABCRejection algorithm, no need to use ABC SMC algorithm, returning ABCRejection output..."
      return ABCrejresults
  end

  while (ϵ >= ABCsetup.ϵT) & (sum(numsims) <= ABCsetup.maxiterations)

    i = 1 #set particle indicator to 1
    particles = Array{ParticleSMCModel}(undef, ABCsetup.nparticles)
    distvec = zeros(Float64, ABCsetup.nparticles)
    its = 1

    if progress
      p = Progress(ABCsetup.nparticles, 1, "ABC SMC population $(popnum), new ϵ: $(round(ϵ, digits = 2))...", 30)
    end
    while i < ABCsetup.nparticles + 1

      #draw model from previous model probabilities
      mstar = wsample(1:ABCsetup.nmodels, modelprob)
      #perturb model
      mdoublestar = perturbmodel(ABCsetup, mstar, modelprob)
      # sample particle with correct model
      j = wsample(1:ABCsetup.nparticles, weights[mdoublestar, :])
      particletemp = oldparticles[j]
      #perturb particle
      newparticle = perturbparticle(particletemp, ABCsetup.Models[mdoublestar].kernel)
      #calculate priorprob
      priorp = priorprob(newparticle.params, ABCsetup.Models[mdoublestar].prior)

      if priorp == 0.0 #return to beginning of loop if prior probability is 0
        continue
      end

      #simulate with new parameters
      dist, out = ABCsetup.Models[mdoublestar].simfunc(newparticle.params, ABCsetup.Models[mdoublestar].constants, targetdata)

      #if simulated data is less than target tolerance accept particle
      if dist < ϵ
        particles[i] = newparticle
        particles[i].other = out
        particles[i].distance = dist
        distvec[i] = dist
        i += 1
        if progress
          next!(p)
        end
      end

      its += 1
    end

    particles, weights = smcweightsmodel(particles, oldparticles, ABCsetup, modelprob)
    weights, modelprob = getparticleweights(particles, ABCsetup)
    ABCsetup = modelselection_kernel(ABCsetup, particles)
    oldparticles = particles

    if finalpop == true
      break
    end

    ϵ = quantile(distvec, ABCsetup.α)

    if ϵ < ABCsetup.ϵT
      ϵ = ABCsetup.ϵT
      push!(ϵvec, ϵ)
      push!(numsims, its)
      popnum = popnum + 1
      finalpop = true
      continue
    end

    push!(ϵvec, ϵ)
    push!(numsims, its)

    if verbose
      println("##################################################")
      println(ABCSMCmodelresults(particles, numsims, ABCsetup, ϵvec))
      println("##################################################\n")
    end

    if ((( abs(ϵvec[end - 1] - ϵ )) / ϵvec[end - 1]) < ABCsetup.convergence) == true
      println("New ϵ is within $(round(ABCsetup.convergence * 100, digits = 2))% of previous population, stop ABC SMC")
      break
    end

    popnum = popnum + 1
  end

  out = ABCSMCmodelresults(particles, numsims, ABCsetup, ϵvec)
  return out
end


function stonemask(
    x::Vector{Float64},
    fs::Int,
    temporal_positions::Vector{Float64},
    f0::Vector{Float64})

    refined_f0 = similar(f0)

    @ccall libworld.StoneMask(
        x::Ptr{Float64},
        length(x)::Cint,
        fs::Cint,
        temporal_positions::Ptr{Float64},
        f0::Ptr{Float64},
        length(f0)::Cint,
        refined_f0::Ptr{Float64})::Cvoid

    refined_f0
end


"""
    smooth_vars_and_write_to_netcdf!(output_file,input_file,vars_to_smooth,window_h,window_t)
Take a netcdf file, a list of variables and two integers and performs a moving mean smoothing of these variables. Write the smoothed fields into a netcdf file called output_file
"""
function smooth_vars_and_write_to_netcdf!(output_file,input_file,vars_to_smooth,window_h,window_t)

    ds_orig = Dataset(input_file)
    ds_new = Dataset(output_file,"c", attrib = OrderedDict(
        "history"                   => "Filtered data from file: $input_file, using $window_h points in the space dimensions and $window_t points in the time dimension"
    ))
    # Dimensions
    for (dim,dim_size) in ds_orig.dim
        @info "Writing $dim coordinate in netcdf file"
        flush(stdout)
        ds_new.dim[dim] = dim_size # unlimited dimension
        v = ds_orig[dim]
        create_var = defVar(ds_new,dim, eltype(v.var), dimnames(v),attrib = v.var.attrib)
        create_var[:] = v[:]
    end
    # Declare variables
    for current_var in vars_to_smooth
        @info "Processing $current_var"
        flush(stdout)
        v = ds_orig[current_var]
        create_var = defVar(ds_new,current_var, eltype(v.var), dimnames(v),attrib = v.var.attrib)
        create_var[:] = filter_array(v[:],window_h,window_t)
        @info "Finished processing $current_var"
        flush(stdout)
    end
    close(ds_new)
    close(ds_orig)
end

const dimensions_of_APE_variables = Dict(
    "rho0"    => ("z","t"),
    "N2"      => ("z","t"),
    "B"       => ("x","y","z","t"),
    "QCONVEC" => ("x","y","z","t"),
    "QRAD"    => ("x","y","z","t"),
    "APE"     => ("t",),
    "KE"      => ("t",),
    "APERAD"  => ("t",),
    "APEDIA"  => ("t",),
    "APEWN2"  => ("t",),
    "APEUb2"  => ("t",),
    "APEVb2"  => ("t",),
    "APErate" => ("t",),
    "APESURF" => ("t",)
)

"""
Create NetCDF for storing ape budget variables
"""
function create_APE_netcdf(filename,var_size)

    Dataset(filename,"c") do ds

        # Dimensions

        ds.dim["x"] = var_size[1]
        ds.dim["y"] = var_size[2]
        ds.dim["z"] = var_size[3]
        ds.dim["time"] = Inf # unlimited dimension

        # Declare variables

        ncx = defVar(ds,"x", Float32, ("x",), deflatelevel=1,  attrib = OrderedDict(
            "units"                     => "m",
        ))

        ncy = defVar(ds,"y", Float32, ("y",), deflatelevel=1, attrib = OrderedDict(
            "units"                     => "m",
        ))

        ncz = defVar(ds,"z", Float32, ("z",), deflatelevel=1,  attrib = OrderedDict(
            "units"                     => "m",
            "long_name"                 => "height",
        ))

        nctime = defVar(ds,"time", Float32, ("time",), deflatelevel=1,  attrib = OrderedDict(
            "units"                     => "d",
            "long_name"                 => "time",
        ))

        ncAPE = defVar(ds,"APE", Float32, ("time",), deflatelevel=1,  attrib = OrderedDict(
            "long_name"                 => "Available Potential Energy",
            "units"                     => "J/m^2     ",
        ))

        ncKE = defVar(ds,"KE", Float32, ("time",), deflatelevel=1,  attrib = OrderedDict(
            "long_name"                 => "Kinetic energy",
            "units"                     => "J/m^2     ",
        ))
        ncAPERAD = defVar(ds,"APERAD", Float32, ("time",), deflatelevel=1,  attrib = OrderedDict(
            "long_name"                 => "Available Potential Energy production by radiative processes",
            "units"                     => "W/m^2     ",
        ))
        ncAPEDIA = defVar(ds,"APEDIA", Float32, ("time",), deflatelevel=1,  attrib = OrderedDict(
            "long_name"                 => "Available Potential Energy production by radiative processes other than radiative and surface fluxes",
            "units"                     => "W/m^2     ",
        ))
        ncAPEWN2 = defVar(ds,"APEWN2", Float32, ("time",), deflatelevel=1,  attrib = OrderedDict(
            "long_name"                 => "Available Potential Energy conversion into kinetic energy",
            "units"                     => "W/m^2     ",
        ))
        ncAPEUB2 = defVar(ds,"APEUb2", Float32, ("time",), deflatelevel=1,  attrib = OrderedDict(
            "long_name"                 => "Available Potential Energy advection by the U component of the wind",
            "units"                     => "W/m^2     ",
        ))
        ncAPEVB2 = defVar(ds,"APEVb2", Float32, ("time",), deflatelevel=1,  attrib = OrderedDict(
            "long_name"                 => "Available Potential Energy advection by the V component of the wind",
            "units"                     => "W/m^2     ",
        ))
        ncAPESURF = defVar(ds,"APESURF", Float32, ("time",), deflatelevel=1,  attrib = OrderedDict(
            "long_name"                 => "Available Potential Energy production by surface fluxes",
            "units"                     => "W/m^2     ",
        ))
        ncAPErate = defVar(ds,"APErate", Float32, ("time",), deflatelevel=1,  attrib = OrderedDict(
            "long_name"                 => "Available Potential Energy rate of change",
            "units"                     => "W/m^2     ",
        ))

        ncrho0 = defVar(ds,"rho0", Float32, ("z","time"), deflatelevel=1, shuffle=true, attrib = OrderedDict(
            "units"                     => "kg/m^3",
            "long_name"                 => "Horizontally averaged density",
        ))

        ncN2 = defVar(ds,"N2", Float32, ("z","time"), deflatelevel=1, shuffle=true, attrib = OrderedDict(
            "units"                     => "1/s^2",
            "long_name"                 => "Brunt-Vaisala frequency",
        ))

        ncQCONVEC = defVar(ds,"QCONVEC", Float32, ("x", "y", "z", "time"), deflatelevel=1, shuffle=true, attrib = OrderedDict(
            "long_name"                 => "Convective heating departure from horizontal mean",
            "units"                     => "m/s^3       ",
        ))

        ncQRAD = defVar(ds,"QRAD", Float32, ("x", "y", "z", "time"), deflatelevel=1, shuffle=true, attrib = OrderedDict(
            "long_name"                 => "Radiative heating rate departure from horizontal mean",
            "units"                     => "m/s^3     ",
        ))

        ncB = defVar(ds,"B", Float32, ("x", "y", "z", "time"), deflatelevel=1, shuffle=true, attrib = OrderedDict(
            "long_name"                 => "Buoyancy rate departure from horizontal mean",
            "units"                     => "m/s^2     ",
        ))

        for itime = 1:var_size[end]
            nctime[itime] = itime * 3600.0
        end

    end

end

"""
Create NetCDF for storing ape budget variables
"""
function set_netcdf_var!(filename,var,data)
    @info "Started writing: $var into file: $filename"
    n_dimensions = ndims(data)
    last_dimension_size = last(size(data))
    Dataset(filename,"a") do ds
        #for ind in 1:last_dimension_size
        #    #@show ndims(data),last_dimension_size, ind,  selectdim(data,n_dimensions,ind)
        #    @info "Writing step: " ind
        #    selectdim(ds[var],n_dimensions,ind)[:] =  selectdim(data,n_dimensions,ind)
        #end
        if n_dimensions == 1
            ds[var][:] = data[:]
        elseif n_dimensions == 2
            ds[var][:,:] = data[:,:]
        elseif n_dimensions == 3
            ds[var][:,:,:] = data[:,:,:]
        elseif n_dimensions == 4
            ds[var][:,:,:,:] = data[:,:,:,:]
        end
        @info "Finished writing: $var into file: $filename"
    end
end


# This file was generated, do not modify it. # hide
RidgeRegressor = @load RidgeRegressor pkg="MultivariateStats"
KNNRegressor = @load KNNRegressor


using HDF5, JLD
using PyCall
@pyimport matplotlib2tikz

include("CenterOfMass.jl")
include("setup_parameters.jl")

save_plots = false
do_belief = false

plt[:close]("all")
pygui(!save_plots)

folder = "multi_robot"

type TrialData
  belief::Histogram
  estimate
  error
  applied_f
  f_hat
  csqmi
  robots_measurement
  robots
end

@load "$(folder)/data.25_iter.0_4" data_array csqmi_ratios thetas actuator_limit

forces = map(x->x.applied_f, data_array)
num_measurement = map(x->length(x.robots_measurement), data_array)
num_robot = map(x->length(x.robots), data_array)
actuator_limits = actuator_limit * num_measurement
limits = forces - actuator_limits .>-1e-6
partial_limits = forces - actuator_limit .>-1e-6

#######################
# plot normalized error
#######################
println("normalized error")

figure()
errors = map(x->x.error, data_array)
mean_errors = mean(errors, 2)

colors = ["b", "g", "r", "c"]

indices = collect(1:size(mean_errors,3))
for ii = 1:size(mean_errors, 1)
#for ii = 1:3
  plot(indices, mean_errors[ii, 1, :][:], color = colors[ii], linewidth = 2.0)
end
legend(map(string, csqmi_ratios))

cutoff = 0.15
for ii = 1:size(errors, 1)
  n = 2
  stds = std(errors[ii,:,:], 2)[:]/sqrt(size(errors,2))
  es = mean_errors[ii,1,:][:]
  fill([indices;reverse(indices)], [es+stds;reverse(es-stds)], color = colors[ii],
    alpha=0.2, linewidth=0.0)
end
#plot([0;25],[cutoff;cutoff],linestyle="--", color="k")
#plot([9;9],[0.0;cutoff],linestyle="--", color="k")
#plot([13;13],[0.0;cutoff],linestyle="--", color="k")
#plot([17;17],[0.0;cutoff],linestyle="--", color="k")

if false
for jj = 1:size(errors, 1)
  for ii = 1:size(errors, 2)
    plot(1:size(errors,3), errors[jj, ii,:][:], color = colors[jj], alpha = 0.2,
    linewidth = 1.0)

    #robot_added_inds = find(diff(num_robot[jj, ii, :][:]))+1
    #scatter(robot_added_inds, errors[jj, ii, robot_added_inds], color =
    #colors[jj], alpha = 0.2, marker = "1")

    #for kk = 2:size(errors, 3)
      #if
        #scatter(kk, errors[jj, ii, kk], color = colors[jj], alpha = 0.2)
      #end
    #end
  end
end
end

#ylim(0.0, 0.8)
xlim(0.0, size(errors, 3))

xlabel("Iteration")
ylabel("Normalized error")

matplotlib2tikz.save("$(folder)/convergence.tex", figureheight="\\figureheight",
  figurewidth="\\figurewidth")

#######################
# plot number of robots
#######################
println("number robots")

figure()

mean_robots = mean(num_robot, 2)
for ii = 1:size(mean_robots, 1)
  plot(indices, mean_robots[ii, 1, :][:], color = colors[ii], linewidth = 2.0)
end

legend(map(string, csqmi_ratios))

for ii = 1:size(num_robot, 2)
  for jj = 1:size(num_robot, 1)
    plot(1:size(errors,3), num_robot[jj, ii,:][:], color = colors[jj], alpha = 0.2,
    linewidth = 1.0)
  end
end

xlabel("Iteration")
ylabel("Num. robots")

matplotlib2tikz.save("$(folder)/num_robots.tex", figureheight="\\figureheight",
  figurewidth="\\figurewidth")

#############
# plot belief
#############
println("belief")

if save_plots && do_belief
  for (ii, csqmi_ratio) = enumerate(csqmi_ratios)
    for jj = 1:size(errors,2)
      trial_folder = "$(folder)/$(csqmi_ratio)_$(jj)"
      mkpath(trial_folder)

      for kk = 1:size(errors,3)
        plt[:close]("all")

        println("belief: ratio=$(csqmi_ratio), trial=$(jj), iter=$(kk)")
        data = data_array[ii,jj,kk]
        robots = data.robots
        belief = data.belief
        robots_measurement = data.robots_measurement
        theta = thetas[:,jj]
        com_p = theta[1:2]
        mass = theta[3]

        figure()
        plot_attachment_points(attachment_ps)
        plot_occupied_points(attachment_ps[setdiff(robots, robots_measurement)])
        cloud = to_cloud(belief, interior_q)
        scaled_ps = (15*cloud[4,:]'/maximum(cloud[4,:]))
        scatter3D(cloud[1,:]', cloud[2,:]', cloud[3,:]', "z", scaled_ps.^2, "purple", alpha = 0.5)

        scatter3D([com_p[1]], [com_p[2]], [mass], "z", 200, "red", marker="*", alpha=1)

        plot_measurement_points(attachment_ps[robots_measurement])

        if kk > 1
          old_robots = data_array[ii,jj,kk-1].robots
          if length(robots) > length(old_robots)
            plot_new_point(attachment_ps[setdiff(robots, old_robots)[1]])
          end
        end
        fill3d(hcat(circle_ps...), alpha=0.2)

        xlabel("X")
        ylabel("Y")
        zlabel("M")
        #fig[:axes][1][:get_yaxis]()[:set_visible](false)
        #fig[:axes][1][:get_xaxis]()[:set_visible](false)
        #axis("off")
        xlim(-1.0,1.0)
        ylim(-1.0,1.0)
        zlim(0.0,1.405)
        savefig("$(trial_folder)/belief_$(kk).png", pad_inches=0.01, bbox_inches="tight")
      end
    end
  end
end

pygui(true)


#####
##### Sampling: high-level interface and building blocks
#####

"Significant digits to display for reporting."
const REPORT_SIGDIGITS = 3

####
#### parts unaffected by warmup
####

"""
$(TYPEDEF)
A log density bundled with an RNG and options for sampling. Contains the parts of the
problem which are not changed during warmup.
# Fields
$(FIELDS)
"""
struct SamplingLogDensity{D,R,L<:AbstractProbabilityModel{D},O,S}
    "Random number generator."
    rng::R
    "Log density."
    ℓ::L
    """
    Algorithm used for sampling, also contains the relevant parameters that are not affected
    by adaptation. See eg [`NUTS`](@ref).
    """
    algorithm::O
    "Reporting warmup information and chain progress."
    reporter::S
end

####
#### warmup building blocks
####

###
### warmup state
###

"""
$(TYPEDEF)
Representation of a warmup state. Not part of the API.
# Fields
$(FIELDS)
"""
struct WarmupState{D,T,L,Tκ <: KineticEnergy, Tϵ <: Union{Real,Nothing}}
    z::PhasePoint{D,T,L}
    κ::Tκ
    ϵ::Tϵ
end

function Base.show(io::IO, warmup_state::WarmupState)
    @unpack κ, ϵ = warmup_state
    ϵ_display = ϵ ≡ nothing ? "unspecified" : "≈ $(round(ϵ; sigdigits = REPORT_SIGDIGITS))"
    print(io, "adapted sampling parameters: stepsize (ϵ) $(ϵ_display), $(κ)")
end

###
### warmup interface and stages
###

"""
$(SIGNATURES)
Return the *results* and the *next warmup state* after warming up/adapting according to
`warmup_stage`, starting from `warmup_state`.
Use `nothing` for a no-op.
"""
function warmup!(tree, chain, sampling_logdensity::SamplingLogDensity, warmup_stage::Nothing, warmup_state)
    nothing, warmup_state
end

random_position!(rng::AbstractRNG, q::PtrVector) = rand!(rng, q, -2.0, 2.0)

"""
$(SIGNATURES)
Helper function to create random starting positions in the `[-2,2]ⁿ` box.
"""
function random_position(rng::AbstractRNG, sptr::StackPointer, ::Static{P}) where {P}
    sptr, q = PtrVector{P,Float64}(sptr)
    sptr, random_position!(rng, q)
end

"Docstring for initial warmup arguments."
const DOC_INITIAL_WARMUP_ARGS =
"""
- `q`: initial position. *Default*: random (uniform [-2,2] for each coordinate).
- `κ`: kinetic energy specification. *Default*: Gaussian with identity matrix.
- `ϵ`: a scalar for initial stepsize, or `nothing` for heuristic finders.
"""

"""
$(SIGNATURES)
Create an initial warmup state from a random position.
# Keyword arguments
$(DOC_INITIAL_WARMUP_ARGS)
"""
function initialize_warmup_state(rng, sptr::StackPointer, ℓ::AbstractProbabilityModel{D}; q = nothing, κ = nothing, ϵ = nothing) where {D}
    if κ ≡ nothing
        sptr, κ′ = GaussianKineticEnergy(sptr, dimension(ℓ))
    else
        κ′ = κ
    end
    # @show κ′.W.diag, κ′.M⁻¹.diag
    # @show pointer(κ′.W.diag), pointer(κ′.M⁻¹.diag)
    tree = Tree{D,Float64}(sptr)
    # println("Allocated tree:")
    # @show κ′.W.diag, κ′.M⁻¹.diag
    # @show pointer(κ′.W.diag), pointer(κ′.M⁻¹.diag)
    treeptr, flag = undefined_z(tree)
    # @show treeptr
    LT = aligned_offset(tree)
    p = PtrVector{D,Float64}(treeptr)
    # println("About to generate a random position.")
    # @show κ′.W.diag, κ′.M⁻¹.diag
    # @show pointer(κ′.W.diag), pointer(κ′.M⁻¹.diag)
    q′ = q ≡ nothing ? random_position!(rng, PtrVector{D,Float64}(treeptr + LT)) : q
    ∇ℓq = PtrVector{D,Float64}(treeptr + 2LT)
    # println("About to evaluate the logdensity:")
    # @show κ′.W.diag, κ′.M⁻¹.diag
    # @show pointer(κ′.W.diag), pointer(κ′.M⁻¹.diag)
    Q = evaluate_ℓ!(tree.sptr, ∇ℓq, ℓ, q′)
    # println("Just evaluated the logdensity:")
    # @show κ′.W.diag, κ′.M⁻¹.diag
    # @show pointer(κ′.W.diag), pointer(κ′.M⁻¹.diag)
    tree, WarmupState(PhasePoint(Q, p, flag), κ′, ϵ)
end

"""
$(TYPEDEF)
Find a local optimum (using quasi-Newton methods).
It is recommended that this stage is applied so that the initial stepsize selection happens
in a region which is at least plausible.
"""
Base.@kwdef struct FindLocalOptimum{T}
    """
    Add `-0.5 * magnitude_penalty * sum(abs2, q)` to the log posterior **when finding the local
    optimum**. This can help avoid getting into high-density edge areas of the posterior
    which are otherwise not typical (eg multilevel models).
    """
    magnitude_penalty::T = 1e-4
    """
    Maximum number of iterations in the optimization algorithm. Recall that we don't need to
    find the mode, or even a local mode, just be in a reasonable region.
    """
    iterations::Int = 50
    # FIXME allow custom algorithm, tolerance, etc
end
@noinline ThrowOptimizationError(str) = throw(str)
function warmup!(
    tree::Tree, chain, tree_statistics, ϵs, sampling_logdensity::SamplingLogDensity{D},
    local_optimization::FindLocalOptimum, warmup_state
) where {D}
    @unpack rng, ℓ, reporter = sampling_logdensity
    @unpack magnitude_penalty, iterations = local_optimization
    @unpack z, κ, ϵ = warmup_state
    @unpack Q, p, flag = z
    @unpack q, ℓq, ∇ℓq = Q
    report(reporter, "finding initial optimum")
    for _ in 1:100
        ℓq = QuasiNewtonMethods.proptimize!(tree.sptr, ℓ, q, ∇ℓq, ℓq, magnitude_penalty, iterations)#+100)
        # @show q
        # @show ℓq
        # @show ∇ℓq
        isfinite(ℓq) && return WarmupState(PhasePoint(EvaluatedLogDensity(q, ℓq, ∇ℓq), p, flag), κ, ϵ)
        random_position!(rng, q)
        ℓq = evaluate_ℓ!(tree.sptr, ∇ℓq, ℓ, q).ℓq
        magnitude_penalty += magnitude_penalty
    end
    ThrowOptimizationError("Optimization failed to converge, returning $ℓq; Thread ID: $(Threads.threadid()).")
    # fg! = function(F, G, q)
        # ℓq, ∇ℓq = logdensity_and_gradient(ℓ, q)
        # if G ≠ nothing
            # @. G = -∇ℓq - q * magnitude_penalty
        # end
        # -ℓq - (0.5 * magnitude_penalty * sum(abs2, q))
    # end
    # objective = NLSolversBase.OnceDifferentiable(NLSolversBase.only_fg!(fg!), q)
    # opt = Optim.optimize(objective, q, Optim.LBFGS(),
                         # Optim.Options(; iterations = iterations))
    # sp, q = PtrVector{D,Float64}(sp)
    # ∇ℓq = PtrVector{D,Float64}(pointer(sp,Float64))
    # sptr is set ahead by proptimize! to store optim and gradient.
end
Base.length(::FindLocalOptimum) = 0
function warmup!(
    tree::Tree, chain, tree_statistics, ϵs,
    sampling_logdensity, stepsize_search::InitialStepsizeSearch, warmup_state
)
    @unpack rng, ℓ, reporter = sampling_logdensity
    @unpack z, κ, ϵ = warmup_state
    # @argcheck ϵ ≡ nothing "stepsize ϵ manually specified, won't perform initial search"
    rand_p!(rng, z.p, κ)
    ϵ = find_initial_stepsize(stepsize_search, local_acceptance_ratio(tree.sptr, Hamiltonian(κ, ℓ), z))
    report(reporter, "found initial stepsize",
           ϵ = round(ϵ; sigdigits = REPORT_SIGDIGITS))
    WarmupState(z, κ, ϵ)
end
Base.length(::InitialStepsizeSearch) = 0
"""
$(TYPEDEF)
Tune the step size `ϵ` during sampling, and the metric of the kinetic energy at the end of
the block. The method for the latter is determined by the type parameter `M`, which can be
1. `Diagonal` for diagonal metric (the default),
2. `Symmetric` for a dense metric,
3. `Nothing` for an unchanged metric.
# Results
A `NamedTuple` with the following fields:
- `chain`, a vector of position vectors
- `tree_statistics`, a vector of tree statistics for each sample
- `ϵs`, a vector of step sizes for each sample
# Fields
$(FIELDS)
"""
struct TuningNUTS{M,D}
    "Number of samples."
    N::Int
    "Dual averaging parameters."
    stepsize_adaptation::D
    """
    Regularization factor for normalizing variance. An estimated covariance matrix `Σ` is
    rescaled by `λ` towards ``σ²I``, where ``σ²`` is the median of the diagonal. The
    constructor has a reasonable default.
    """
    λ::Float64
    function TuningNUTS{M}(N::Integer, stepsize_adaptation::D,
                           λ = 5.0/N) where {M <: Union{Nothing,Diagonal,Symmetric},D}
        # @argcheck N ≥ 20        # variance estimator is kind of meaningless for few samples
        # @argcheck λ ≥ 0
        new{M,D}(N, stepsize_adaptation, λ)
    end
end

function Base.show(io::IO, tuning::TuningNUTS{M}) where {M}
    @unpack N, stepsize_adaptation, λ = tuning
    print(io, "Stepsize and metric tuner, $(N) samples, $(M) metric, regularization $(λ)")
end

Base.length(t::TuningNUTS) = t.N

# """
# $(SIGNATURES)
# Form a matrix from positions (`q`), with each column containing a position.
# """
# position_matrix(chain) = reduce(hcat, chain)

# """
# $(SIGNATURES)
# Estimate the inverse metric from the chain.
# In most cases, this should be regularized, see [`regularize_M⁻¹`](@ref).
# """
# sample_M⁻¹(::Type{Diagonal}, chain) = Diagonal(vec(var(position_matrix(chain); dims = 2)))

# sample_M⁻¹(::Type{Symmetric}, chain) = Symmetric(cov(position_matrix(chain); dims = 2))

# function store_sample!(_∇ℓq::PtrVector{P,T,L}, _q::PtrVector{P,T,L}, Q::EvaluatedLogDensity) where {P,T,L}
    # @unpack q, ∇ℓq = Q
    # @inbounds for l in 1:L
        # _q[l] = q[l]
        # _∇ℓq[l] = ∇ℓq[l]
    # end
    # EvaluatedLogDensity(
        # _q, Q.ℓq, _∇ℓq
    # )
# end

function warmup!(
    tree::Tree{D,T,L}, chain₊::AbstractMatrix{T},
    tree_statistics::AbstractVector{TreeStatisticsNUTS}, ϵs::AbstractVector{T},
    sampling_logdensity::SamplingLogDensity{D},
    tuning::TuningNUTS{M},
    warmup_state
) where {D,T,L,M}
    @unpack rng, ℓ, algorithm, reporter = sampling_logdensity
    @unpack z, κ, ϵ = warmup_state
    @unpack N, stepsize_adaptation, λ = tuning
    # L = VectorizationBase.align(D, T)
    chain_ptr = pointer(chain₊)
    chain = DynamicPtrMatrix{T}(chain_ptr, (D, N), L)
    # tree_statistics = Vector{TreeStatisticsNUTS}(undef, N)
    H = Hamiltonian(κ, ℓ)
    ϵ_state = initial_adaptation_state(stepsize_adaptation, ϵ)
    # ϵs = Vector{Float64}(undef, N)
    mcmc_reporter = make_mcmc_reporter(reporter, N; tuning = "stepsize and Diagonal{T,PtrVector{}} metric")
    # sp, ∇ℓq = PtrVector{D,T}(sp)
    for n in 1:N
        ϵ = current_ϵ(ϵ_state)
        # @assert ϵ > 1e-10 "Current ϵ: $ϵ; final: $(final_ϵ(ϵ_state))"
        if ϵ < 1e-10
            @show Threads.threadid(), z
            Q = evaluate_ℓ!(tree.sptr, z.Q.∇ℓq, sampling_logdensity.ℓ, z.Q.q)
            @show Threads.threadid(), Q
            throw(AssertionError("Current ϵ: $ϵ; final: $(final_ϵ(ϵ_state))"))
        end
        ϵs[n] = ϵ
        z, stats = sample_tree(rng, tree, algorithm, H, z, ϵ)
        copyto!( PtrVector{D,T}( chain_ptr ), z.Q.q ) # relocate to base of stack
        # @show z.Q.q
        chain_ptr += L*sizeof(T)
        tree_statistics[n] = stats
        ϵ_state = adapt_stepsize(stepsize_adaptation, ϵ_state, stats.acceptance_rate)
        reporting(mcmc_reporter) && report(mcmc_reporter, n; ϵ = round(ϵ; sigdigits = REPORT_SIGDIGITS))
    end
    # κ = GaussianKineticEnergy(regularize_M⁻¹(sample_M⁻¹(M, chain), λ))
    # sp, κ = GaussianKineticEnergy(sp, chain, λ, Val{D}())
    if M ≢ Nothing
        GaussianKineticEnergy!(κ, chain, λ)
        reporting(mcmc_reporter) && report(mcmc_reporter, "adaptation finished", adapted_kinetic_energy = κ)
    end
    # (chain = chain, tree_statistics = tree_statistics, ϵs = ϵs)
    WarmupState(z, κ, final_ϵ(ϵ_state))
end

function mcmc!(tree::Tree{D,T,L}, chain::AbstractMatrix, tree_statistics::AbstractVector{TreeStatisticsNUTS}, sampling_logdensity::SamplingLogDensity{D}, N, warmup_state) where {D,T,L}
    @unpack rng, ℓ, algorithm, reporter = sampling_logdensity
    @unpack z, κ, ϵ = warmup_state

    # tree_statistics = Vector{TreeStatisticsNUTS}(undef, N)
    H = Hamiltonian(κ, ℓ)
    mcmc_reporter = make_mcmc_reporter(reporter, N)
    chain_ptr = pointer(chain)
    for n in 1:N
        z, tree_statistics[n] = sample_tree(rng, tree, algorithm, H, z, ϵ)
        copyto!( PtrVector{D,T}( chain_ptr ), z.Q.q )
        chain_ptr += L*sizeof(T)

        report(mcmc_reporter, n)
    end
    # tree_statistics
end

"""
$(SIGNATURES)
Helper function for constructing the “middle” doubling warmup stages in
[`default_warmup_stages`](@ref).
"""
function _doubling_warmup_stages(M, stepsize_adaptation, middle_steps,
                                 doubling_stages::Val{D}) where {D}
    ntuple(d -> TuningNUTS{M}(middle_steps << (d - 1), stepsize_adaptation), Val(D))
end

"""
$(SIGNATURES)
A sequence of warmup stages:
1. find the local optimum using `local_optimization`,
2. select an initial stepsize using `stepsize_search` (default: based on a heuristic),
3. tuning stepsize with `init_steps` steps
4. tuning stepsize and covariance: first with `middle_steps` steps, then repeat with twice
   the steps `doubling_stages` times
5. tuning stepsize with `terminating_steps` steps.
`M` (`Diagonal`, the default or `Symmetric`) determines the type of the metric adapted from
the sample.
This is the suggested tuner of most applications.
Use `nothing` for `local_optimization` or `stepsize_adaptation` to skip the corresponding
step.
"""
function default_warmup_stages(;
                               local_optimization = FindLocalOptimum(),
                               stepsize_search = InitialStepsizeSearch(),
                               M::Type{<:Union{Diagonal,Symmetric}} = Diagonal,
                               stepsize_adaptation = DualAveraging(),
                               init_steps = 75, middle_steps = 25, doubling_stages::Val{DS} = Val(5),
                               terminating_steps = 50) where {DS}
    (local_optimization, stepsize_search,
     TuningNUTS{Nothing}(init_steps, stepsize_adaptation),
     _doubling_warmup_stages(M, stepsize_adaptation, middle_steps, Val(DS))...,
     TuningNUTS{Nothing}(terminating_steps, stepsize_adaptation))
end

"""
$(SIGNATURES)
A sequence of warmup stages for fixed stepsize:
1. find the local optimum using `local_optimization`,
2. tuning covariance: first with `middle_steps` steps, then repeat with twice
   the steps `doubling_stages` times
Very similar to [`default_warmup_stages`](@ref), but omits the warmup stages with just
stepsize tuning.
"""
function fixed_stepsize_warmup_stages(;
                                      local_optimization = FindLocalOptimum(),
                                      M::Type{<:Union{Diagonal,Symmetric}} = Diagonal,
                                      middle_steps = 25, doubling_stages::Val{DS} = 5) where {DS}
    (local_optimization,
     _doubling_warmup_stages(M, FixedStepsize(), middle_steps, Val(DS))...)
end

"""
$(SIGNATURES)
Helper function for implementing warmup.
!!! note
    Changes may imply documentation updates in [`mcmc_keep_warmup`](@ref).
"""
@generated function _warmup!(tree::Tree, chain, tree_statistics, ϵs, sampling_logdensity, stages::T, warmup_state_0) where {T}
    N = length(T.parameters)
    q = quote end
    for n in 1:N
        warmup_call_q = quote
            $(Symbol(:warmup_state_,n)) = warmup!(tree, chain, tree_statistics, ϵs, sampling_logdensity, stages[$n], $(Symbol(:warmup_state_,n-1)))
        end
        push!(q.args, warmup_call_q)
    end
    push!(q.args, Symbol(:warmup_state_,N))
    q
end


@testset "Automata" begin

    using KnuthBendix

    abt = KnuthBendix.Alphabet(['a', 'e', 'b', 'p'])
    ta = KnuthBendix.Automaton(abt)
    σ = KnuthBendix.initialstate(ta)
    @test sprint(show, ta) isa String

    @test ta isa KnuthBendix.AbstractAutomaton
    @test ta isa KnuthBendix.Automaton
    @test ta isa KnuthBendix.Automaton{4, Word{UInt16}}
    @test σ isa KnuthBendix.AbstractState
    @test σ isa KnuthBendix.State
    @test σ isa KnuthBendix.State{4, Word{UInt16}}

    @test isone(KnuthBendix.name(σ))
    @test !KnuthBendix.isterminal(σ)
    @test isone(KnuthBendix.rightrule(σ))

    @test length(KnuthBendix.states(ta)) == 1
    @test length(KnuthBendix.inedges(σ)) == 0
    @test length(KnuthBendix.outedges(σ)) == 4
    @test length(KnuthBendix.stateslengths(ta)) == 1
    @test length(σ) == 0

    KnuthBendix.declarerightrule!(σ, one(KnuthBendix.name(σ)))
    @test KnuthBendix.isterminal(σ)
    @test isone(KnuthBendix.rightrule(σ))

    push!(ta, Word([1]))
    KnuthBendix.addedge!(ta, 1, 1, 2)
    @test length(KnuthBendix.inedges(KnuthBendix.states(ta)[2])) == 1
    @test KnuthBendix.walk(ta, Word([1])) == (1, KnuthBendix.states(ta)[2])
    @test_throws BoundsError KnuthBendix.addedge!(ta, 5, 1, 2)
    @test_throws BoundsError KnuthBendix.addedge!(ta, 1, 3, 2)
    @test_throws BoundsError KnuthBendix.addedge!(ta, 1, 1, 3)
    @test_throws BoundsError KnuthBendix.removeedge!(ta, 5, 1, 2)
    @test_throws BoundsError KnuthBendix.removeedge!(ta, 1, 3, 2)
    @test_throws BoundsError KnuthBendix.removeedge!(ta, 1, 1, 3)
    KnuthBendix.removeedge!(ta, 1, 1, 2)
    @test KnuthBendix.walk(ta, Word([1])) == (0, KnuthBendix.states(ta)[1])

    KnuthBendix.addedge!(ta, 1, 1, 2)
    deleteat!(ta, 2)
    @test length(KnuthBendix.states(ta)) == 1
    @test length(KnuthBendix.stateslengths(ta)) == 1
    @test KnuthBendix.isfailstate(KnuthBendix.outedges(σ)[1])

    A = KnuthBendix.Alphabet(['a', 'e', 'b', 'p'])
    KnuthBendix.set_inversion!(A, 'a', 'e')
    KnuthBendix.set_inversion!(A, 'b', 'p')

    a = Word([1,2])
    b = Word([2,1])
    c = Word([3,4])
    d = Word([4,3])
    ε = one(a)
    ba = Word([3,1])
    ab = Word([1,3])
    be = Word([3,2])
    eb = Word([2,3])
    pa = Word([4,1])
    ap = Word([1,4])
    pe = Word([4,2])
    ep = Word([2,4])

    lenlexord = KnuthBendix.LenLex(A)
    rsc = KnuthBendix.RewritingSystem([a=>ε, b=>ε, c=>ε, d=>ε, ba=>ab, be=>eb, pa=>ap, pe=>ep], lenlexord)
    ia = KnuthBendix.makeindexautomaton(rsc)

    testword = Word([1,1,1,1,1,2,2,2,3,4,2,2,3,3,3,4,4,4,4,3,4,3,4,1,2,1,1,1,1,1,1,1,2,1,3,4])
    @test KnuthBendix.rewrite_from_left(testword, rsc) == KnuthBendix.rewrite_from_left(testword, ia)

    w = Word([1,3,4,1,4,4,1,1,4,2,3,2,4,2,2,3,1,2,1])
    @test KnuthBendix.rewrite_from_left(w, rsc) == KnuthBendix.rewrite_from_left(w, ia)

    @test !isempty(ia)
    @test isempty(empty!(ia))

    rsd = KnuthBendix.RewritingSystem([a=>ε, b=>ε], lenlexord)
    iad = KnuthBendix.makeindexautomaton(rsd)
    @test KnuthBendix.rewrite_from_left(testword, rsd) == KnuthBendix.rewrite_from_left(testword, iad)
end


const use_return  = haskey(ENV, "USE_RETURN" ) ? parse(Bool, ENV["USE_RETURN"] ) : false
const USE_GPU     = haskey(ENV, "USE_GPU"    ) ? parse(Bool, ENV["USE_GPU"]    ) : false
const do_viz      = haskey(ENV, "DO_VIZ"     ) ? parse(Bool, ENV["DO_VIZ"]     ) : false
const do_save     = haskey(ENV, "DO_SAVE"    ) ? parse(Bool, ENV["DO_SAVE"]    ) : false
const do_save_viz = haskey(ENV, "DO_SAVE_VIZ") ? parse(Bool, ENV["DO_SAVE_VIZ"]) : false
const nx          = haskey(ENV, "NX"         ) ? parse(Int , ENV["NX"]         ) : 64 - 1
const ny          = haskey(ENV, "NY"         ) ? parse(Int , ENV["NY"]         ) : 64 - 1
const nz          = haskey(ENV, "NZ"         ) ? parse(Int , ENV["NZ"]         ) : 64 - 1
###
using ParallelStencil
using ParallelStencil.FiniteDifferences3D
@static if USE_GPU
    @init_parallel_stencil(CUDA, Float64, 3)
else
    @init_parallel_stencil(Threads, Float64, 3)
end
using Plots, Printf, Statistics, LinearAlgebra, MAT
# CPU functions
@views av_xza(A) = (A[1:end-1,:,1:end-1] .+ A[1:end-1,:,2:end] .+ A[2:end,:,1:end-1] .+ A[2:end,:,2:end]).*0.25
@views av_zi(A)  = (A[2:end-1,2:end-1,2:end-2] .+ A[2:end-1,2:end-1,3:end-1]).*0.5
@views av_za(A)  = (A[:,:,1:end-1] .+ A[:,:,2:end]).*0.5
@views inn(A)    =  A[2:end-1,2:end-1,2:end-1]

@parallel function smooth!(A2::Data.Array, A::Data.Array, fact::Data.Number)
    @inn(A2)   = @inn(A) + 1.0/6.1/fact*(@d2_xi(A) + @d2_yi(A) + @d2_zi(A))
    return
end

@parallel function compute_maxloc!(Musτ2::Data.Array, Musτ::Data.Array)
    @inn(Musτ2) = @maxloc(Musτ)
    return
end

macro Mu_eff() esc(:(1.0/(1.0/@all(Musτ) + 1.0/(G*dt)))) end
@parallel function compute_iter_params!(dτ_Rho::Data.Array, Gdτ::Data.Array, Musτ::Data.Array, Vpdτ::Data.Number, G::Data.Number, dt::Data.Number, Re::Data.Number, r::Data.Number, max_lxyz::Data.Number)
    @all(dτ_Rho) = Vpdτ*max_lxyz/Re/@Mu_eff()
    @all(Gdτ)    = Vpdτ^2/@all(dτ_Rho)/(r+2)
    return
end

@parallel_indices (ix,iy,iz) function assign_τ!(τxx::Data.Array, τyy::Data.Array, τzz::Data.Array, τxy::Data.Array,  τxz::Data.Array,  τyz::Data.Array, τxx_o::Data.Array, τyy_o::Data.Array, τzz_o::Data.Array, τxy_o::Data.Array, τxz_o::Data.Array, τyz_o::Data.Array)
    if (ix<=size(τxx,1) && iy<=size(τxx,2) && iz<=size(τxx,3))  τxx_o[ix,iy,iz] = τxx[ix,iy,iz]  end
    if (ix<=size(τyy,1) && iy<=size(τyy,2) && iz<=size(τyy,3))  τyy_o[ix,iy,iz] = τyy[ix,iy,iz]  end
    if (ix<=size(τzz,1) && iy<=size(τzz,2) && iz<=size(τzz,3))  τzz_o[ix,iy,iz] = τzz[ix,iy,iz]  end
    if (ix<=size(τxy,1) && iy<=size(τxy,2) && iz<=size(τxy,3))  τxy_o[ix,iy,iz] = τxy[ix,iy,iz]  end
    if (ix<=size(τxz,1) && iy<=size(τxz,2) && iz<=size(τxz,3))  τxz_o[ix,iy,iz] = τxz[ix,iy,iz]  end
    if (ix<=size(τyz,1) && iy<=size(τyz,2) && iz<=size(τyz,3))  τyz_o[ix,iy,iz] = τyz[ix,iy,iz]  end
    return
end

macro inn_yz_Gdτ(ix,iy,iz) esc(:( Gdτ[$ix  ,$iy+1,$iz+1] )) end
macro inn_xz_Gdτ(ix,iy,iz) esc(:( Gdτ[$ix+1,$iy  ,$iz+1] )) end
macro inn_xy_Gdτ(ix,iy,iz) esc(:( Gdτ[$ix+1,$iy+1,$iz  ] )) end
macro inn_yz_Mus(ix,iy,iz) esc(:( Mus[$ix  ,$iy+1,$iz+1] )) end
macro inn_xz_Mus(ix,iy,iz) esc(:( Mus[$ix+1,$iy  ,$iz+1] )) end
macro inn_xy_Mus(ix,iy,iz) esc(:( Mus[$ix+1,$iy+1,$iz  ] )) end
macro av_xyi_Gdτ(ix,iy,iz) esc(:( (Gdτ[$ix  ,$iy  ,$iz+1] + Gdτ[$ix+1,$iy  ,$iz+1] + Gdτ[$ix  ,$iy+1,$iz+1] + Gdτ[$ix+1,$iy+1,$iz+1])*0.25 )) end
macro av_xzi_Gdτ(ix,iy,iz) esc(:( (Gdτ[$ix  ,$iy+1,$iz  ] + Gdτ[$ix+1,$iy+1,$iz  ] + Gdτ[$ix  ,$iy+1,$iz+1] + Gdτ[$ix+1,$iy+1,$iz+1])*0.25 )) end
macro av_yzi_Gdτ(ix,iy,iz) esc(:( (Gdτ[$ix+1,$iy  ,$iz  ] + Gdτ[$ix+1,$iy+1,$iz  ] + Gdτ[$ix+1,$iy  ,$iz+1] + Gdτ[$ix+1,$iy+1,$iz+1])*0.25 )) end
macro av_xyi_Mus(ix,iy,iz) esc(:( (Mus[$ix  ,$iy  ,$iz+1] + Mus[$ix+1,$iy  ,$iz+1] + Mus[$ix  ,$iy+1,$iz+1] + Mus[$ix+1,$iy+1,$iz+1])*0.25 )) end
macro av_xzi_Mus(ix,iy,iz) esc(:( (Mus[$ix  ,$iy+1,$iz  ] + Mus[$ix+1,$iy+1,$iz  ] + Mus[$ix  ,$iy+1,$iz+1] + Mus[$ix+1,$iy+1,$iz+1])*0.25 )) end
macro av_yzi_Mus(ix,iy,iz) esc(:( (Mus[$ix+1,$iy  ,$iz  ] + Mus[$ix+1,$iy+1,$iz  ] + Mus[$ix+1,$iy  ,$iz+1] + Mus[$ix+1,$iy+1,$iz+1])*0.25 )) end
macro inn_yz_Gr(ix,iy,iz) esc(:( @inn_yz_Gdτ($ix,$iy,$iz)/(G*dt) )) end
macro inn_xz_Gr(ix,iy,iz) esc(:( @inn_xz_Gdτ($ix,$iy,$iz)/(G*dt) )) end
macro inn_xy_Gr(ix,iy,iz) esc(:( @inn_xy_Gdτ($ix,$iy,$iz)/(G*dt) )) end
macro av_xyi_Gr(ix,iy,iz) esc(:( @av_xyi_Gdτ($ix,$iy,$iz)/(G*dt) )) end
macro av_xzi_Gr(ix,iy,iz) esc(:( @av_xzi_Gdτ($ix,$iy,$iz)/(G*dt) )) end
macro av_yzi_Gr(ix,iy,iz) esc(:( @av_yzi_Gdτ($ix,$iy,$iz)/(G*dt) )) end
@parallel_indices (ix,iy,iz) function compute_Pt_τ!(Pt::Data.Array, τxx::Data.Array, τyy::Data.Array, τzz::Data.Array, τxy::Data.Array, τxz::Data.Array, τyz::Data.Array, τxx_o::Data.Array, τyy_o::Data.Array, τzz_o::Data.Array, τxy_o::Data.Array, τxz_o::Data.Array, τyz_o::Data.Array, Vx::Data.Array, Vy::Data.Array, Vz::Data.Array, Mus::Data.Array, Gdτ::Data.Array, r::Data.Number, G::Data.Number, dt::Data.Number, _dx::Data.Number, _dy::Data.Number, _dz::Data.Number)
    if (ix<=size(Pt,1)  && iy<=size(Pt,2)  && iz<=size(Pt,3))    Pt[ix,iy,iz] = Pt[ix,iy,iz] - r*Gdτ[ix,iy,iz]*(_dx*(Vx[ix+1,iy,iz] - Vx[ix,iy,iz]) + _dy*(Vy[ix,iy+1,iz] - Vy[ix,iy,iz]) + _dz*(Vz[ix,iy,iz+1] - Vz[ix,iy,iz]))  end
    if (ix<=size(τxx,1) && iy<=size(τxx,2) && iz<=size(τxx,3))  τxx[ix,iy,iz] = ( τxx[ix,iy,iz] + τxx_o[ix,iy,iz]*@inn_yz_Gr(ix,iy,iz) + 2.0*@inn_yz_Gdτ(ix,iy,iz)*(_dx*(Vx[ix+1,iy+1,iz+1] - Vx[ix,iy+1,iz+1])))/(1.0 + @inn_yz_Gdτ(ix,iy,iz)/@inn_yz_Mus(ix,iy,iz) + @inn_yz_Gr(ix,iy,iz))  end
    if (ix<=size(τyy,1) && iy<=size(τyy,2) && iz<=size(τyy,3))  τyy[ix,iy,iz] = ( τyy[ix,iy,iz] + τyy_o[ix,iy,iz]*@inn_xz_Gr(ix,iy,iz) + 2.0*@inn_xz_Gdτ(ix,iy,iz)*(_dy*(Vy[ix+1,iy+1,iz+1] - Vy[ix+1,iy,iz+1])))/(1.0 + @inn_xz_Gdτ(ix,iy,iz)/@inn_xz_Mus(ix,iy,iz) + @inn_xz_Gr(ix,iy,iz))  end
    if (ix<=size(τzz,1) && iy<=size(τzz,2) && iz<=size(τzz,3))  τzz[ix,iy,iz] = ( τzz[ix,iy,iz] + τzz_o[ix,iy,iz]*@inn_xy_Gr(ix,iy,iz) + 2.0*@inn_xy_Gdτ(ix,iy,iz)*(_dz*(Vz[ix+1,iy+1,iz+1] - Vz[ix+1,iy+1,iz])))/(1.0 + @inn_xy_Gdτ(ix,iy,iz)/@inn_xy_Mus(ix,iy,iz) + @inn_xy_Gr(ix,iy,iz))  end
    if (ix<=size(τxy,1) && iy<=size(τxy,2) && iz<=size(τxy,3))  τxy[ix,iy,iz] = ( τxy[ix,iy,iz] + τxy_o[ix,iy,iz]*@av_xyi_Gr(ix,iy,iz) + 2.0*@av_xyi_Gdτ(ix,iy,iz)*(0.5*(_dy*(Vx[ix+1,iy+1,iz+1] - Vx[ix+1,iy,iz+1]) + _dx*(Vy[ix+1,iy+1,iz+1] - Vy[ix,iy+1,iz+1])))) / (1.0 + @av_xyi_Gdτ(ix,iy,iz)/@av_xyi_Mus(ix,iy,iz) + @av_xyi_Gr(ix,iy,iz))  end
    if (ix<=size(τxz,1) && iy<=size(τxz,2) && iz<=size(τxz,3))  τxz[ix,iy,iz] = ( τxz[ix,iy,iz] + τxz_o[ix,iy,iz]*@av_xzi_Gr(ix,iy,iz) + 2.0*@av_xzi_Gdτ(ix,iy,iz)*(0.5*(_dz*(Vx[ix+1,iy+1,iz+1] - Vx[ix+1,iy+1,iz]) + _dx*(Vz[ix+1,iy+1,iz+1] - Vz[ix,iy+1,iz+1])))) / (1.0 + @av_xzi_Gdτ(ix,iy,iz)/@av_xzi_Mus(ix,iy,iz) + @av_xzi_Gr(ix,iy,iz))  end
    if (ix<=size(τyz,1) && iy<=size(τyz,2) && iz<=size(τyz,3))  τyz[ix,iy,iz] = ( τyz[ix,iy,iz] + τyz_o[ix,iy,iz]*@av_yzi_Gr(ix,iy,iz) + 2.0*@av_yzi_Gdτ(ix,iy,iz)*(0.5*(_dz*(Vy[ix+1,iy+1,iz+1] - Vy[ix+1,iy+1,iz]) + _dy*(Vz[ix+1,iy+1,iz+1] - Vz[ix+1,iy,iz+1])))) / (1.0 + @av_yzi_Gdτ(ix,iy,iz)/@av_yzi_Mus(ix,iy,iz) + @av_yzi_Gr(ix,iy,iz))  end
    return
end

macro av_xi_dτ_Rho(ix,iy,iz)  esc(:( (dτ_Rho[$ix  ,$iy+1,$iz+1] + dτ_Rho[$ix+1,$iy+1,$iz+1])*0.5 )) end
macro av_yi_dτ_Rho(ix,iy,iz)  esc(:( (dτ_Rho[$ix+1,$iy  ,$iz+1] + dτ_Rho[$ix+1,$iy+1,$iz+1])*0.5 )) end
macro av_zi_dτ_Rho(ix,iy,iz)  esc(:( (dτ_Rho[$ix+1,$iy+1,$iz  ] + dτ_Rho[$ix+1,$iy+1,$iz+1])*0.5 )) end
@parallel_indices (ix,iy,iz) function compute_V!(Vx::Data.Array, Vy::Data.Array, Vz::Data.Array, Pt::Data.Array, τxx::Data.Array, τyy::Data.Array, τzz::Data.Array, τxy::Data.Array, τxz::Data.Array, τyz::Data.Array, dτ_Rho::Data.Array, _dx::Data.Number, _dy::Data.Number, _dz::Data.Number, nx_1, nx_2, ny_1, ny_2, nz_1, nz_2)
    if (ix<=nx_1 && iy<=ny_2 && iz<=nz_2)  Vx[ix+1,iy+1,iz+1] = Vx[ix+1,iy+1,iz+1] + (_dx*(τxx[ix+1,iy,iz] - τxx[ix,iy,iz]) + _dy*(τxy[ix,iy+1,iz] - τxy[ix,iy,iz]) + _dz*(τxz[ix,iy,iz+1] - τxz[ix,iy,iz]) - _dx*(Pt[ix+1,iy+1,iz+1] - Pt[ix,iy+1,iz+1]))*@av_xi_dτ_Rho(ix,iy,iz)  end
    if (ix<=nx_2 && iy<=ny_1 && iz<=nz_2)  Vy[ix+1,iy+1,iz+1] = Vy[ix+1,iy+1,iz+1] + (_dy*(τyy[ix,iy+1,iz] - τyy[ix,iy,iz]) + _dx*(τxy[ix+1,iy,iz] - τxy[ix,iy,iz]) + _dz*(τyz[ix,iy,iz+1] - τyz[ix,iy,iz]) - _dy*(Pt[ix+1,iy+1,iz+1] - Pt[ix+1,iy,iz+1]))*@av_yi_dτ_Rho(ix,iy,iz)  end
    if (ix<=nx_2 && iy<=ny_2 && iz<=nz_1)  Vz[ix+1,iy+1,iz+1] = Vz[ix+1,iy+1,iz+1] + (_dz*(τzz[ix,iy,iz+1] - τzz[ix,iy,iz]) + _dx*(τxz[ix+1,iy,iz] - τxz[ix,iy,iz]) + _dy*(τyz[ix,iy+1,iz] - τyz[ix,iy,iz]) - _dz*(Pt[ix+1,iy+1,iz+1] - Pt[ix+1,iy+1,iz]))*@av_zi_dτ_Rho(ix,iy,iz)  end
    return
end

@parallel_indices (ix,iy,iz) function compute_Res!(∇V::Data.Array, Rx::Data.Array, Ry::Data.Array, Rz::Data.Array, Vx::Data.Array, Vy::Data.Array, Vz::Data.Array, Pt::Data.Array, τxx::Data.Array, τyy::Data.Array, τzz::Data.Array, τxy::Data.Array, τxz::Data.Array, τyz::Data.Array, _dx::Data.Number, _dy::Data.Number, _dz::Data.Number)
    if (ix<=size(∇V,1) && iy<=size(∇V,2) && iz<=size(∇V,3))  ∇V[ix,iy,iz] = _dx*(Vx[ix+1,iy,iz] - Vx[ix,iy,iz]) + _dy*(Vy[ix,iy+1,iz] - Vy[ix,iy,iz]) + _dz*(Vz[ix,iy,iz+1] - Vz[ix,iy,iz]) end
    if (ix<=size(Rx,1) && iy<=size(Rx,2) && iz<=size(Rx,3))  Rx[ix,iy,iz] = _dx*(τxx[ix+1,iy,iz] - τxx[ix,iy,iz]) + _dy*(τxy[ix,iy+1,iz] - τxy[ix,iy,iz]) + _dz*(τxz[ix,iy,iz+1] - τxz[ix,iy,iz]) - _dx*(Pt[ix+1,iy+1,iz+1] - Pt[ix,iy+1,iz+1]) end
    if (ix<=size(Ry,1) && iy<=size(Ry,2) && iz<=size(Ry,3))  Ry[ix,iy,iz] = _dy*(τyy[ix,iy+1,iz] - τyy[ix,iy,iz]) + _dx*(τxy[ix+1,iy,iz] - τxy[ix,iy,iz]) + _dz*(τyz[ix,iy,iz+1] - τyz[ix,iy,iz]) - _dy*(Pt[ix+1,iy+1,iz+1] - Pt[ix+1,iy,iz+1]) end
    if (ix<=size(Rz,1) && iy<=size(Rz,2) && iz<=size(Rz,3))  Rz[ix,iy,iz] = _dz*(τzz[ix,iy,iz+1] - τzz[ix,iy,iz]) + _dx*(τxz[ix+1,iy,iz] - τxz[ix,iy,iz]) + _dy*(τyz[ix,iy+1,iz] - τyz[ix,iy,iz]) - _dz*(Pt[ix+1,iy+1,iz+1] - Pt[ix+1,iy+1,iz]) end
    return
end

@parallel_indices (iy,iz) function bc_x!(A::Data.Array)
    A[  1, iy,  iz] = A[    2,   iy,   iz]
    A[end, iy,  iz] = A[end-1,   iy,   iz]
    return
end

@parallel_indices (ix,iz) function bc_y!(A::Data.Array)
    A[ ix,  1,  iz] = A[   ix,    2,   iz]
    A[ ix,end,  iz] = A[   ix,end-1,   iz]
    return
end

@parallel_indices (ix,iy) function bc_z!(A::Data.Array)
    A[ ix,  iy,  1] = A[   ix,   iy,    2]
    A[ ix,  iy,end] = A[   ix,   iy,end-1]
    return
end

@views function Stokes3D_()
    # Physics
    lx, ly, lz = 10.0, 10.0, 10.0  # domain extends
    ξ          = 1.0               # Maxwell relaxation time
    μs0        = 1.0               # matrix viscosity
    μsi        = 1e-3              # inclusion viscosity
    G          = 1.0               # elastic shear modulus
    εbg        = 1.0               # background strain-rate
    dt         = μs0/(G*ξ)
    # Numerics
    nt         = 1#5                 # number of time steps
    iterMax    = 200#2e5               # maximum number of pseudo-transient iterations
    nout       = 2000               # error checking frequency
    Re         = 6.0π              # Reynolds number
    r          = 1.0               # Bulk to shear elastic modulus ratio
    CFL        = 0.8/sqrt(3)       # CFL number
    ε          = 1e-8              # nonlinear absolute tolerence
    # nx, ny, nz = 127, 127, 127     # numerical grid resolution; should be a mulitple of 32-1 for optimal GPU perf
    # Derived numerics
    dx, dy, dz = lx/nx, ly/ny, lz/nz      # cell sizes
    max_lxyz   = max(lx,ly,lz)
    Vpdτ       = min(dx,dy,dz)*CFL
    _dx, _dy, _dz = 1.0/dx, 1.0/dy, 1.0/dz
    xc, yc, zc = LinRange(dx/2, lx-dx/2, nx), LinRange(dy/2, ly-dy/2, ny), LinRange(dz/2, lz-dz/2, nz)
    # Array allocations
    Pt         = @zeros(nx  ,ny  ,nz  )
    dτ_Rho     = @zeros(nx  ,ny  ,nz  )
    Gdτ        = @zeros(nx  ,ny  ,nz  )
    ∇V         = @zeros(nx  ,ny  ,nz  )
    Vy         = @zeros(nx  ,ny+1,nz  )
    τxx        = @zeros(nx  ,ny-2,nz-2)
    τyy        = @zeros(nx-2,ny  ,nz-2)
    τzz        = @zeros(nx-2,ny-2,nz  )
    τxy        = @zeros(nx-1,ny-1,nz-2)
    τxz        = @zeros(nx-1,ny-2,nz-1)
    τyz        = @zeros(nx-2,ny-1,nz-1)
    τxx_o      = @zeros(nx  ,ny-2,nz-2)
    τyy_o      = @zeros(nx-2,ny  ,nz-2)
    τzz_o      = @zeros(nx-2,ny-2,nz  )
    τxy_o      = @zeros(nx-1,ny-1,nz-2)
    τxz_o      = @zeros(nx-1,ny-2,nz-1)
    τyz_o      = @zeros(nx-2,ny-1,nz-1)
    Rx         = @zeros(nx-1,ny-2,nz-2)
    Ry         = @zeros(nx-2,ny-1,nz-2)
    Rz         = @zeros(nx-2,ny-2,nz-1)
    Mus2       = @zeros(nx  ,ny  ,nz  )
    Musτ       = @zeros(nx  ,ny  ,nz  )
    # Initial conditions
    Rad2       =  zeros(nx  ,ny  ,nz  )
    Vx         =  zeros(nx+1,ny  ,nz  )
    Vz         =  zeros(nx  ,ny  ,nz+1)
    Rad2      .= [(((ix-1)*dx +0.5*dx -0.5*lx)^2 + ((iy-1)*dy +0.5*dy -0.5*ly)^2 + ((iz-1)*dz +0.5*dz -0.5*lz)^2) for ix=1:size(Rad2,1), iy=1:size(Rad2,2), iz=1:size(Rad2,3)]
    Vx         = Data.Array( -εbg.*[((ix-1)*dx -0.5*lx) for ix=1:size(Vx,1), iy=1:size(Vx,2), iz=1:size(Vx,3)] )
    Vz         = Data.Array(  εbg.*[((iz-1)*dz -0.5*lz) for ix=1:size(Vz,1), iy=1:size(Vz,2), iz=1:size(Vz,3)] )
    Mus        = μs0*ones(nx,ny,nz)
    Mus[Rad2.<1.0] .= μsi
    Mus        = Data.Array(Mus)
    Mus2      .= Mus
    for ism=1:10#15
        @parallel smooth!(Mus2, Mus, 1.0)
        Mus, Mus2 = Mus2, Mus
    end
    Musτ      .= Mus
    @parallel compute_maxloc!(Musτ, Mus)
    @parallel (1:size(Musτ,2), 1:size(Musτ,3)) bc_x!(Musτ)
    @parallel (1:size(Musτ,1), 1:size(Musτ,3)) bc_y!(Musτ)
    @parallel (1:size(Musτ,1), 1:size(Musτ,2)) bc_z!(Musτ)
    nx_1, nx_2, ny_1, ny_2, nz_1, nz_2 = nx-1, nx-2, ny-1, ny-2, nz-1, nz-2
    # Preparation of visualisation
    if do_viz || do_save_viz
        ENV["GKSwstype"]="nul"; if do_viz !ispath("../../figures") && mkdir("../../figures") end
        y_sl2, y_sl = Int(ceil((ny-2)/2)), Int(ceil(ny/2))
        xc, zc  = dx/2:dx:(lx-dx/2), dz/2:dz:(lz-dz/2)
        xv, zv  = 0:dx:lx, 0:dz:lz
    end
    # Time loop
    @parallel compute_iter_params!(dτ_Rho, Gdτ, Musτ, Vpdτ, G, dt, Re, r, max_lxyz)
    t=0.0; ittot=0; evo_t=[]; evo_τzz=[]; t_tic = 0.0
    for it = 1:nt
        err=2*ε; iter=0; err_evo1=[]; err_evo2=[]
        @parallel assign_τ!(τxx, τyy, τzz, τxy, τxz, τyz, τxx_o, τyy_o, τzz_o, τxy_o, τxz_o, τyz_o)
        # Pseudo-transient iteration
        while err > ε && iter <= iterMax
            if (it==1 && iter==11) t_tic = Base.time() end
            @parallel compute_Pt_τ!(Pt, τxx, τyy, τzz, τxy, τxz, τyz, τxx_o, τyy_o, τzz_o, τxy_o, τxz_o, τyz_o, Vx, Vy, Vz, Mus, Gdτ, r, G, dt, _dx, _dy, _dz)
            @parallel compute_V!(Vx, Vy, Vz, Pt, τxx, τyy, τzz, τxy, τxz, τyz, dτ_Rho, _dx, _dy, _dz, nx_1, nx_2, ny_1, ny_2, nz_1, nz_2)
            @parallel (1:size(Vy,2), 1:size(Vy,3)) bc_x!(Vy)
            @parallel (1:size(Vz,2), 1:size(Vz,3)) bc_x!(Vz)
            @parallel (1:size(Vx,1), 1:size(Vx,3)) bc_y!(Vx)
            @parallel (1:size(Vz,1), 1:size(Vz,3)) bc_y!(Vz)
            @parallel (1:size(Vx,1), 1:size(Vx,2)) bc_z!(Vx)
            @parallel (1:size(Vy,1), 1:size(Vy,2)) bc_z!(Vy)
            iter += 1
            if iter % nout == 0
                @parallel compute_Res!(∇V, Rx, Ry, Rz, Vx, Vy, Vz, Pt, τxx, τyy, τzz, τxy, τxz, τyz, _dx, _dy, _dz)
                Vmin, Vmax = minimum(Vx), maximum(Vx)
                Pmin, Pmax = minimum(Pt), maximum(Pt)
                norm_Rx    = norm(Rx)/(Pmax-Pmin)*lx/sqrt(length(Rx))
                norm_Ry    = norm(Ry)/(Pmax-Pmin)*lx/sqrt(length(Ry))
                norm_Rz    = norm(Rz)/(Pmax-Pmin)*lx/sqrt(length(Rz))
                norm_∇V    = norm(∇V)/(Vmax-Vmin)*lx/sqrt(length(∇V))
                # norm_Rx = norm(Rx)/length(Rx); norm_Ry = norm(Ry)/length(Ry); norm_Rz = norm(Rz)/length(Rz); norm_∇V = norm(∇V)/length(∇V)
                err = maximum([norm_Rx, norm_Ry, norm_Rz, norm_∇V])
                if isnan(err) error("NaN") end
                push!(err_evo1,maximum([norm_Rx, norm_Ry, norm_Rz, norm_∇V])); push!(err_evo2,iter)
                @printf("Step = %d, iter = %d, err = %1.3e [norm_Rx=%1.3e, norm_Ry=%1.3e, norm_Rz=%1.3e, norm_∇V=%1.3e] \n", it, iter, err, norm_Rx, norm_Ry, norm_Rz, norm_∇V)
            end
        end
        ittot += iter; t += dt
        # push!(evo_t, t); push!(evo_τzz, maximum(τzz))
    end
    # Performance
    wtime    = Base.time() - t_tic
    A_eff    = (10*2 + 6*1 + 3)/1e9*nx*ny*nz*sizeof(Data.Number) # Effective main memory access per iteration [GB] (Lower bound of required memory access: Te has to be read and written: 2 whole-array memaccess; Ci has to be read: : 1 whole-array memaccess)
    wtime_it = wtime/(ittot-10)                           # Execution time per iteration [s]
    T_eff    = A_eff/wtime_it                             # Effective memory throughput [GB/s]
    @printf("Total iters = %d (%d steps), time = %1.3e sec (@ T_eff = %1.2f GB/s) \n", ittot, nt, wtime, round(T_eff, sigdigits=3))
    # Visualisation
    if do_viz
        p1 = heatmap(xc, zc, Array(Pt)[:,y_sl,:]', aspect_ratio=1, xlims=extrema(xc), zlims=extrema(zc), c=:viridis, title="Pressure")
        p2 = heatmap(xc, zv, Array(Vz)[:,y_sl,:]', aspect_ratio=1, xlims=extrema(xc), zlims=extrema(zv), c=:viridis, title="Vz")
        p4 = heatmap(xc[2:end-1], xv[2:end-1], log10.(abs.(Array(Rz)[:,y_sl2,:]')), aspect_ratio=1,  xlims=extrema(xc[2:end-1]), zlims=extrema(zv[2:end-1]), c=:viridis, title="log10(Rz)")
        #p5 = plot(err_evo2,err_evo1, legend=false, xlabel="# iterations", ylabel="log10(error)", linewidth=2, markershape=:circle, markersize=3, labels="max(error)", yaxis=:log10)
        p3 = plot(evo_t, evo_τzz, legend=false, xlabel="time", ylabel="max(τzz)", linewidth=0, markershape=:circle, framestyle=:box, markersize=3)
            #plot!(evo_t, 2.0.*εbg.*μs0.*(1.0.-exp.(.-evo_t.*G./μs0)), linewidth=2.0) # analytical solution
        plot(p1, p2, p4, p3)
        savefig("../../figures/Stokes_3D_ve_perf_$(nx).png")
    end
    if do_save
        !ispath("../../output") && mkdir("../../output")
        open("../../output/out_Stokes3D_ve_perf.txt","a") do io
            println(io, "1 $(nx) $(ny) $(nz) $(ittot) $(wtime) $(A_eff) $(wtime_it) $(T_eff)")
        end
    end
    if do_save_viz
        !ispath("../../out_visu") && mkdir("../../out_visu")
        matwrite("../../out_visu/Stokes_3D_ve_perf.mat", Dict("Pt_3D"=> Array(Pt), "Mus_3D"=> Array(Mus), "Txz_3D"=> Array(τxz), "Vz_3D"=> Array(Vz), "dx_3D"=> dx, "dy_3D"=> dy, "dz_3D"=> dz); compress = true)
    end
    return xc, yc, zc, Pt
end

if use_return
    xc, yc, zc, P = Stokes3D_();
else
    Stokes3D = begin Stokes3D_(); return; end
end


########################
# multi-threading task
########################

@inline function subspace_mul_generic_task!(S::AbstractVector, indices, U::AbstractMatrix, subspace, range, offset)
    D = StrideArrays.static_length(indices)
    y = StrideArray{T}(undef, (D, ))
    if D == 4
        for s in range
            k = subspace[s]
            subspace_mul_kernel_generic_4x4!(S, y, indices, U, k, offset)
        end
    else
        for s in range
            k = subspace[s]
            subspace_mul_kernel_generic!(S, y, indices, U, k, offset)
        end
    end
    return
end

@inline function subspace_mul_generic_task!(S::AbstractMatrix, indices, U::AbstractMatrix, subspace, range, offset)
    D = StrideArrays.static_length(indices)
    stride = size(S, 1)
    y = StrideArray{T}(undef, (D, ))
    if D == 4
        for (b, s) in range
            k = subspace[s]
            subspace_mul_kernel_generic_4x4!(S, y, indices, U, k, b, offset)
        end
    else
        for (b, s) in range
            k = subspace[s]
            subspace_mul_kernel_generic!(S, y, indices, U, k, b, offset)
        end
    end
    return
end

@inline function subspace_mul_complex_task!(Sr::AbstractMatrix{T}, indices, U, subspace, range::CartesianPartition, offset) where {T <: Base.HWReal}
    D = StrideArrays.static_length(indices)
    y = StrideArray{T}(undef, (D, )), StrideArray{T}(undef, (D, ))

    for s in range
        k = subspace[s]
        subspace_mul_kernel_complex!(Sr, y, indices, U, k, offset)
    end
    return
end

@inline function subspace_mul_complex_task!(Sr::AbstractArray{T, 3}, indices, U, subspace, range::CartesianPartition, offset) where {T <: Base.HWReal}
    D = ArrayInterface.static_length(indices)
    C = StrideArray{T}(undef, (D, StaticInt{8}())), StrideArray{T}(undef, (D, StaticInt{8}()))
    Bmax = size(Sr, 2)

    for (_b, s) in range
        k = subspace[s]
        subspace_mul_kernel_complex!(Sr, C, indices, U, k, _b, Bmax, offset)
    end
    return
end


"""
    GMC_NS_Ensemble

Abstract supertype for model ensembles intended to be subjected to Galilean Monte Carlo Nested Sampling ('GMC_NS'). Subtyped ensemble structs must implement the following fields as their interface with 'GMC_NS':

    mutable struct Minimal_Ensemble <: GMC_NS_Ensemble
        path::String #ensemble records serialized here

        models::Vector{<:GMC_NS_Model_Record} #records of model paths & lhs
        model_initλ::Function #GMC_NS_Model constructor

        contour<:AbstractFloat #current log_Li[end], init minimum([record.log_Li for record in models])
        log_Li::Vector{<:AbstractFloat} #likelihood of lowest-ranked model at iterate i, init [-Inf]
        log_Xi::Vector{<:AbstractFloat} #amt of prior mass included in ensemble contour at Li, init [0]
        log_wi::Vector{<:AbstractFloat} #width of prior mass covered in this iterate, init [-Inf]
        log_Liwi::Vector{<:AbstractFloat} #evidentiary weight of the iterate, init [-Inf]
        log_Zi::Vector{<:AbstractFloat} #ensemble evidence, init [typemin(FloatType)]
        Hi::Vector{<:AbstractFloat} #ensemble information, init [0]

        obs::Any #observations formatted for GMC_NS_Model likelihood Function
        priors::Vector{<:Distribution} #prior distributions on GMC_NS_Model parameters, formatted for model_initλ
        constants::Vector #fixed parameters to be passed to GMC_NS_Model constructor

        sample_posterior::Bool #switch for posterior_samples collection
        posterior_samples::Vector{<:GMC_NS_Model_Record}

        #GMC tuning settings
        GMC_tune_ι<:Integer #maximum number of initial τ tuning steps
        GMC_tune_ρ<:AbstractFloat #use ensemble size*ρ GMC samples per initial tuning step
        GMC_tune_μ<:Integer #tune τ on the basis of this number of the most recent steps
        GMC_tune_α<:AbstractFloat #tune τ such that the acceptance rate is between this value and 1.
        GMC_tune_β<:AbstractFloat #multiply τ by 1+β (sampling acceptance rate=1.) or 1-β (rate <α)

        #GMC Sampler settings
        GMC_timestep_η<:AbstractFloat #>0. to apply an perturbation term randomly drawn from Normal(0, η*τ) to τ at each step (can assist equilibration)
        GMC_reflect_η<:AbstractFloat #>0. to apply perturbation to reflection vectors, keep value small (can assist equilibration)
        GMC_exhaust_σ<:AbstractFloat #apply this scale parameter to the ensemble size to obtain the maximum number of GMC iterates to perform before terminating nested sampling
        GMC_chain_κ<:Integer #maximum chain length

        t_counter<:Integer #counter for next instantiated trajectory number
    end

"""
abstract type GMC_NS_Ensemble end

"""
    show(e<:GMC_NS_Ensemble)
    show(io,e<:GMC_NS_Ensemble)

Display the ensemble path, likelihood histogram, contour, maximum log likelihood, and log evidence (marginal likelihood).
"""
function Base.show(io::IO, e::GMC_NS_Ensemble; progress=false)
	livec=[model.log_Li for model in e.models]
	maxLH=maximum(livec)
	printstyled(io, "$(typeof(e)) @ $(e.path)\n", bold=true)
	msg = @sprintf "Contour: %3.6e MaxLH:%3.3e log Evidence:%3.6e" e.contour maxLH e.log_Zi[end]
    println(io, msg)
    if length(unique(livec)) > 1
        try
            hist=UnicodePlots.histogram(livec, title="Ensemble Likelihood Distribution")
            show(io, hist)
            println()
            progress && return(nrows(hist.graphics)+6)
        catch
            printstyled("HISTOGRAM UNAVAILABLE", color=:green)
            println()
            progress && return 3
        end
    else
        printstyled("HISTOGRAM UNAVAILABLE", color=:green)
        println()
        progress && return 3
    end
end


using VLConstraintBasedModelGenerationUtilities
using DataFrames

# setup path to sequence file -
path_to_gene_sequence_file = "/Users/jeffreyvarner/Desktop/julia_work/VLConstraintBasedModelGenerationUtilities.jl/test/data/G-MZ373340.fasta"

# load -
gene_table = build_gene_table(path_to_gene_sequence_file) |> check

# get the sequence -
gene_seq = gene_table[!,:gene_sequence][1]
gene_seq_name = gene_table[!,:id][1]

# transcription table -
transcription_table = build_transcription_reaction_table(gene_seq_name, gene_seq; polymeraseSymbol=:RdRp) |> check

# setup path to protein sequence file -
path_to_protein_sequence_file = "/Users/jeffreyvarner/Desktop/julia_work/VLConstraintBasedModelGenerationUtilities.jl/test/data/P-MZ373340.fasta"

# load -
protein_table = build_protein_table(path_to_protein_sequence_file) |> check

# build the translation table -
translation_dictionary = Dict{String,DataFrame}()
translation_array = Array{DataFrame,1}()
(number_of_proteins, _) = size(protein_table)
for protein_index = 1:number_of_proteins

    local_seq = protein_table[protein_index,:protein_sequence]
    local_seq_name = protein_table[protein_index,:id]
    translation_table = build_translation_reaction_table(local_seq_name, local_seq) |> check

    # grab -
    # translation_dictionary[local_seq_name] = translation_table
    push!(translation_array, translation_table)
end

# add together -
master_translation_table = reduce(vcat, translation_array)

# lastly - let's build the transport reactions -
transport_reaction_table = build_transport_reaction_table() |> check

# append reactions into master reaction table -
master_reaction_table = reduce(vcat, [transcription_table, master_translation_table, transport_reaction_table])


applyblock(f, x::Union{Block,Loop}) = typeof(x)(f(x.body))
applyblock(f, x::If) = If(f(x.t), f(x.f))
applyblock(f, x) = x

walk(x, inner, outer) = outer(applyblock(xs -> map(inner, xs), x))

postwalk(f, x) = walk(x, x -> postwalk(f, x), f)
prewalk(f, x)  = walk(f(x), x -> prewalk(f, x), identity)

# remove nops

nops(x) = prewalk(x -> applyblock(is -> filter(i -> i ≠ nop, is), x), x)

# dead code

function deadcode(x)
  prewalk(x) do x
    applyblock(x) do is
      i = findfirst(b -> b isa Branch && !b.cond, is)
      i == 0 ? is : is[1:i]
    end
  end
end

mapbranches(f, b, level = -1) = applyblock(is -> mapbranches.(f, is, level+1), b)
mapbranches(f, b::Branch, level = -1) = f(b, level)

function branches_to(b)
  result = false
  mapbranches(b) do b, l
    b.level==l && (result = true)
    b
  end
  return result
end

# After removing a block that has no branches to it, decrement all of the
# branches that point to any level above it.
function decbranches(i)
  mapbranches(i) do b, l
    b.level > l ? Branch(b.cond, b.level-1) : b
  end
end

# After adding a block increment all branches pointing to a level above it.
function incbranches(i)
  mapbranches(i) do b, l
    b.level >= l ? Branch(b.cond, b.level+1) : b
  end
end

# recover ifs

function makeifs(code)
  prewalk(code) do x
    x isa Block || return x
    i = findfirst(x -> x isa Branch, x.body)
    i != 0 && x.body[i].cond || return x
    cond = x.body[1:i-1]
    cond[end] == Op(i32, :eqz) ? pop!(cond) : push!(cond, Op(i32, :eqz))
    fixbranches = incbranches(Block(x.body[i+1:end])).body
    return Block([cond..., If(fixbranches, [])])
  end
end

# remove unused blocks
# TODO: collapse blocks with the same head

function rmblocks(code)
  prewalk(code) do x
    applyblock(x) do is
      is′ = []
      for i in is
        i isa Block && !branches_to(i) ?
          append!(is′, decbranches(i).body) :
          push!(is′, i)
      end
      is′
    end
  end
end

optimise(b) = b |> deadcode |> makeifs |> rmblocks


using AutoML
using Base.Test

# write your own tests here
@test 1 == 1

include("evaluation/mlogloss.jl")