"""
    couple_models(aero::LiftingLine, stru::GEBT, flap::LiftingLineFlaps, ctrl::Trim)

Create a coupled model using a lifting line aerodynamic model, a
geometrically exact beam theory model, a lifting line control surface model,
and a controller which maintains trimmed operating conditions.  This model introduces
additional parameters corresponding to the freestream air density ``\\rho_\\infty``,
followed by the external loads ``F_{x,i}, F_{y,i}, F_{z,i}, M_{x,i}, M_{y,i},
M_{z,i}`` or displacements ``u_{x,i}, u_{y,i}, u_{z,i}, \\theta_{x,i},
\\theta_{y,i}, \\theta_{z,i}`` for each node, followed by the constant distributed
loads ``f_{x,i}, f_{y,i}, f_{z,i}, m_{x,i}, m_{y,i}, m_{z,i}`` for each beam
element (excluding aerodynamic loads), followed by the body frame linear and
angular velocities ``u, v, w, p, q, r``, and the control surface deflections
which are not defined by the controller.

**NOTE: When using this model, the local frame for each beam element should be
oriented with the x-axis along the beam's axis, the y-axis forward, and the
z-axis normal to the surface**
"""
function couple_models(aero::LiftingLine, stru::GEBT, flap::LiftingLineFlaps,
    ctrl::Trim)

    return (aero, stru, flap, ctrl)
end

# --- traits --- #

function number_of_additional_parameters(aero::LiftingLine, stru::GEBT,
    flap::LiftingLineFlaps{NF,NG,TF,TG}, ctrl::Trim{NC}) where {NA,NF,NG,NC,TF,TG}

    return 1 + 6*length(stru.icol_point) + 6*length(stru.icol_elem) + 6 + NG - NC
end

function coupling_inplaceness(::Type{<:LiftingLine}, ::Type{<:GEBT},
    ::Type{<:LiftingLineFlaps}, ::Type{<:Trim})

    return InPlace()
end

function coupling_rate_jacobian_type(::Type{LiftingLine{NA,TA}}, ::Type{<:GEBT},
    ::Type{LiftingLineFlaps{NF,NG,TF,TG}}, ::Type{<:Trim}) where {NA,NF,NG,TA,TF,TG}

    aero_model_types = TA.parameters
    flap_model_types = TF.parameters
    if all(islinear.(coupling_rate_jacobian_type.(aero_model_types, Ref(LiftingLineSection),
        flap_model_types, Ref(LiftingLineSectionControl))))
        return Linear()
    else
        return Nonlinear()
    end
end

function coupling_state_jacobian_type(::Type{LiftingLine{NA,TA}}, ::Type{<:GEBT},
    ::Type{LiftingLineFlaps{NF,NG,TF,TG}}, ::Type{<:Trim}) where {NA,NF,NG,TA,TF,TG}

    aero_model_types = TA.parameters
    flap_model_types = TF.parameters
    if all(islinear.(coupling_state_jacobian_type.(aero_model_types, Ref(LiftingLineSection),
        flap_model_types, Ref(LiftingLineSectionControl))))
        return Linear()
    else
        return Nonlinear()
    end
end

coupling_parameter_jacobian_type(::Type{<:LiftingLine}, ::Type{<:GEBT},
    ::Type{LiftingLineFlaps}, ::Type{Trim}) = Nonlinear()

function coupling_time_gradient_type(::Type{LiftingLine{NA,TA}}, ::Type{<:GEBT},
    ::Type{LiftingLineFlaps{NF,NG,TF,TG}}, ::Type{<:Trim}) where {NA,NF,NG,TA,TF,TG}

    aero_model_types = TA.parameters
    flap_model_types = TF.parameters
    if all(isempty.(coupling_time_gradient_type.(aero_model_types, Ref(LiftingLineSection),
        flap_model_types, Ref(LiftingLineSectionControl))))
        return Empty()
    elseif all(iszero.(coupling_time_gradient_type.(aero_model_types, Ref(LiftingLineSection),
        flap_model_types, Ref(LiftingLineSectionControl))))
        return Zeros()
    elseif all(isconstant.(coupling_time_gradient_type.(aero_model_types, Ref(LiftingLineSection),
        flap_model_types, Ref(LiftingLineSectionControl))))
        return Constant()
    elseif all(islinear.(coupling_time_gradient_type.(aero_model_types, Ref(LiftingLineSection),
        flap_model_types, Ref(LiftingLineSectionControl))))
        return Linear()
    else
        return Nonlinear()
    end
end

# --- methods --- #

function get_coupling_inputs!(y, aero::LiftingLine{NA,TA}, stru::GEBT,
    flap::LiftingLineFlaps{NF,NG,TF,TG}, ctrl::Trim{NC}, dx, x, p, t) where {NA,NF,NG,NC,TA,TF,TG}

    # extract number of state variables, inputs, and parameters
    nxa = number_of_states(aero) # number of aerodynamic states
    nxs = number_of_states(stru) # number of structural states
    nxf = number_of_states(flap) # number of control surface states
    nxc = number_of_states(ctrl) # number of control surface states
    nya = number_of_inputs(aero) # number of aerodynamic inputs
    nys = number_of_inputs(stru) # number of structural inputs
    nyf = number_of_inputs(flap) # number of control surface inputs
    nyc = number_of_inputs(ctrl) # number of control surface inputs
    npa = number_of_parameters(aero) # number of aerodynamic parameters
    nps = number_of_parameters(stru) # number of structural parameters
    npf = number_of_parameters(flap) # number of control surface parameters
    npc = number_of_parameters(ctrl) # number of control surface parameters
    npadd = number_of_additional_parameters(aero, stru, flap, ctrl) # number of additional parameters

    # get indices for state variables, inputs, and parameters
    ixa = 1:nxa # indices of aerodynamic states
    iya = 1:nya # indices of aerodynamic inputs
    ipa = 1:npa # indices of aerodynamic parameters
    ixs = nxa + 1 : nxa + nxs # indices of structural states
    iys = nya + 1 : nya + nys # indices of structural inputs
    ips = npa + 1 : npa + nps # indices of structural parameters
    ixf = nxa + nxs + 1 : nxa + nxs + nxf # indices of control surface states
    iyf = nya + nys + 1 : nya + nys + nyf # indices of control surface inputs
    ipf = npa + nps + 1 : npa + nps + npf # indices of control surface parameters
    ixc = nxa + nxs + nxf + 1 : nxa + nxs + nxf + nxc # indices of control surface states
    iyc = nya + nys + nyf + 1 : nya + nys + nyf + nyc # indices of control surface inputs
    ipc = npa + nps + npf + 1 : npa + nps + npf + npc # indices of control surface parameters
    ipadd = npa + nps + npf + npc + 1 : npa + nps + npf + npc + npadd # indices of additional parameters
    ixas = state_indices(aero.models) # indices of aerodynamic states for each section
    iyas = input_indices(aero.models) # indices of aerodynamic inputs for each section
    ipas = parameter_indices(aero.models) # indices of aerodynamic parameters for each section
    ixfs = state_indices(flap.models) # indices of control surface states for each section
    iyfs = input_indices(flap.models) # indices of control surface inputs for each section
    ipfs = parameter_indices(flap.models) # indices of control surface parameters for each section

    # separate state variables, inputs, and parameters
    dxa = view(dx, ixa) # aerodynamic rates
    xa = view(x, ixa) # aerodynamic states
    ya = view(y, iya) # aerodynamic inputs
    pa = view(p, ipa) # aerodynamic parameters
    dxs = view(dx, ixs) # structural rates
    xs = view(x, ixs) # structural states
    ys = view(y, iys) # structural inputs
    ps = view(p, ips) # structural parameters
    dxf = view(dx, ixf) # control surface rates
    xf = view(x, ixf) # control surface states
    yf = view(y, iyf) # control surface inputs
    pf = view(p, ipf) # control surface parameters
    dxc = view(dx, ixc) # controller rates
    xc = view(x, ixc) # controller states
    yc = view(y, iyc) # controller inputs
    pc = view(p, ipc) # controller parameters
    padd = view(p, ipadd) # additional parameters for coupled model
    dxas = view.(Ref(dxa), ixas) # aerodynamic rates for each section
    xas = view.(Ref(xa), ixas) # aerodynamic states for each section
    yas = view.(Ref(ya), iyas) # aerodynamic inputs for each section
    pas = view.(Ref(pa), ipas) # aerodynamic parameters for each section
    dxfs = view.(Ref(dxf), ixfs) # control surface rates for each section
    xfs = view.(Ref(xf), ixfs) # control surface states for each section
    yfs = view.(Ref(yf), iyfs) # control surface inputs for each section
    pfs = view.(Ref(pf), ipfs) # control surface parameters for each section

    # number of points and elements
    npoint = length(stru.icol_point)
    nelem = length(stru.icol_elem)

    # controller state variables
    δc = xc

    # global parameters
    ρ = padd[1]
    ur = padd[1 + 6*npoint + 6*nelem + 1]
    vr = padd[1 + 6*npoint + 6*nelem + 2]
    wr = padd[1 + 6*npoint + 6*nelem + 3]
    pr = padd[1 + 6*npoint + 6*nelem + 4]
    qr = padd[1 + 6*npoint + 6*nelem + 5]
    rr = padd[1 + 6*npoint + 6*nelem + 6]
    δp = view(padd, 1 + 6*npoint + 6*nelem + 6 + 1 : 1 + 6*npoint + 6*nelem + 6 + NG - NC)

    # commanded control deflections
    T = promote_type(eltype(dx), eltype(x), eltype(p))
    dδ = zeros(SVector{NG,T})
    iδu = 1
    for i = 1:NG
        if i in ctrl.state_indices
            dδ = setindex(dδ, δc[iδu], i)
            iδu += 1
        else
            dδ = setindex(dδ, 0, i)
        end
    end

    # commanded control deflections
    T = promote_type(eltype(x), eltype(p))
    δ = zeros(SVector{NG,T})
    iδu = iδp = 1
    for i = 1:NG
        if i in ctrl.state_indices
            δ = setindex(δ, δc[iδu], i)
            iδu += 1
        else
            δ = setindex(δ, δp[iδp], i)
            iδp += 1
        end
    end

    # body linear and angular velocity
    V = SVector(ur, vr, wr)
    Ω = SVector(pr, qr, rr)

    # construct assembly from structural parameters
    assembly = gxbeam_assembly(ps, npoint, nelem, stru.start, stru.stop)

    # initialize total forces and moments
    Ftot = @SVector zeros(3)
    Mtot = @SVector zeros(3)

    # save prescribed point loads/displacements
    for ip = 1:npoint
        yoff = 6*(ip-1)
        poff = 1 + 6*(ip-1)
        y[iys[yoff+1]] = padd[poff+1]
        y[iys[yoff+2]] = padd[poff+2]
        y[iys[yoff+3]] = padd[poff+3]
        y[iys[yoff+4]] = padd[poff+4]
        y[iys[yoff+5]] = padd[poff+5]
        y[iys[yoff+6]] = padd[poff+6]
    end

    # add prescribed loads to total forces and moments
    for ip = 1:npoint
        # point location
        point = assembly.points[ip]
        # prescribed condition identities
        prescribed_forces = SVector{6}(view(stru.displacement, :, ip)) .== false
        # point loads
        poff = 1 + 6*(ip-1)
        F1 = ifelse(prescribed_forces[1], padd[poff+1], zero(padd[poff+1]))
        F2 = ifelse(prescribed_forces[2], padd[poff+2], zero(padd[poff+2]))
        F3 = ifelse(prescribed_forces[3], padd[poff+3], zero(padd[poff+3]))
        M1 = ifelse(prescribed_forces[4], padd[poff+4], zero(padd[poff+4]))
        M2 = ifelse(prescribed_forces[5], padd[poff+5], zero(padd[poff+5]))
        M3 = ifelse(prescribed_forces[6], padd[poff+6], zero(padd[poff+6]))
        Fp = SVector(F1, F2, F3)
        Mp = SVector(M1, M2, M3)
        Ftot += Fp
        Mtot += cross(point, Fp) + Mp
    end

    # calculate aerodynamic inputs, distributed loads, and element properties
    for i = 1:NA
        # models for this section
        section_aero = aero.models[i]
        section_stru = LiftingLineSection()
        section_flap = flap.models[i]
        section_ctrl = LiftingLineSectionControl()
        section_models = (section_aero, section_stru, section_flap, section_ctrl)

        # model dimensions for this section
        Nxai = number_of_states(section_aero)
        Nyai = number_of_inputs(section_aero)
        Npai = number_of_parameters(section_aero)
        Nxsi = number_of_states(section_stru)
        Nysi = number_of_inputs(section_stru)
        Npsi = number_of_parameters(section_stru)
        Nxfi = number_of_states(section_flap)
        Nyfi = number_of_inputs(section_flap)
        Npfi = number_of_parameters(section_flap)
        Nxci = number_of_states(section_ctrl)
        Nyci = number_of_inputs(section_ctrl)
        Npci = number_of_parameters(section_ctrl)

        # local section properties
        icol = stru.icol_elem[i]
        element = assembly.elements[i]
        u_elem = SVector(xs[icol], xs[icol+1], xs[icol+2]) # linear displacement
        θ_elem = SVector(xs[icol+3], xs[icol+4], xs[icol+5]) # angular displacement
        F_elem = SVector(xs[icol+6], xs[icol+7], xs[icol+8]) .* stru.force_scaling # internal force
        M_elem = SVector(xs[icol+9], xs[icol+10], xs[icol+11]) .* stru.force_scaling # internal moment
        P_elem = SVector(xs[icol+12], xs[icol+13], xs[icol+14]) .* stru.mass_scaling # linear momentum
        H_elem = SVector(xs[icol+15], xs[icol+16], xs[icol+17]) .* stru.mass_scaling # angular momentum

        dF_elem = SVector(dxs[icol+6], dxs[icol+7], dxs[icol+8]) .* stru.force_scaling # internal force
        dM_elem = SVector(dxs[icol+9], dxs[icol+10], dxs[icol+11]) .* stru.force_scaling # internal moment
        dP_elem = SVector(dxs[icol+12], dxs[icol+13], dxs[icol+14]) .* stru.mass_scaling # linear momentum
        dH_elem = SVector(dxs[icol+15], dxs[icol+16], dxs[icol+17]) .* stru.mass_scaling # angular momentum

        # convert rotation parameter to Wiener-Milenkovic parameters
        scaling = GXBeam.rotation_parameter_scaling(θ_elem)
        θ_elem *= scaling # angular displacement (Wiener-Milenkovic parameters)

        # length and location of the element
        ΔL = element.L # element length
        pe = element.x + u_elem # location of element (in body frame)

        # local to body transformation
        CtCab = GXBeam.get_C(θ_elem)'*element.Cab

        # local structural to aerodynamic transformation
        R = @SMatrix [0 -1 0; 1 0 0; 0 0 1]

        # local freestream linear and angular velocities
        vi = -R*(CtCab'*V + GXBeam.element_linear_velocity(element, P_elem, H_elem)) # local linear freestream velocity
        ωi = R*GXBeam.element_angular_velocity(element, P_elem, H_elem) # local angular freestream velocity

        # local freestream linear and angular accelerations
        dvi = -R*GXBeam.element_linear_velocity(element, dP_elem, dH_elem)
        dωi = R*GXBeam.element_angular_velocity(element, dP_elem, dH_elem)

        # NOTE: We assume that the local beam frame y-axis is oriented towards
        # the leading edge and the z-axis is oriented up
        # NOTE: We assume the local aerodynamic frame is oriented with the x-axis
        # in the chordwise direction and the z-axis in the (airfoil) normal
        # direction.

        # section state variable rates
        dxai = SVector{Nxai}(dxas[i]) # aerodynamic state variables
        dxsi = vcat(dvi, dωi) # structural state variables
        dxfi = SVector{Nxfi}(dxfs[i]) # control surface state variables
        dxci = SMatrix{1,NG}(getindex.(flap.gains, i))* dδ # controller state variables
        dxi = vcat(dxai, dxsi, dxfi, dxci)

        # section state variables
        xai = SVector{Nxai}(xas[i]) # aerodynamic state variables
        xsi = vcat(vi, ωi) # structural state variables
        xfi = SVector{Nxfi}(xfs[i]) # control surface state variables
        xci = SMatrix{1,NG}(getindex.(flap.gains, i)) * δ # controller state variables
        xi = vcat(xai, xsi, xfi, xci)

        # section parameters
        pai = SVector{Npai}(pas[i]) # aerodynamic parameters
        psi = SVector(ρ) # structural parameters
        pfi = SVector{Npfi}(pfs[i]) # control surface parameters
        pci = SVector{0,Float64}() # controller parameters
        pi = vcat(pai, psi, pfi, pci)

        # section inputs
        yi = get_coupling_inputs(section_models, dxi, xi, pi, t)

        # separate inputs
        yai = view(yi, 1:Nyai)
        ysi = view(yi, Nyai + 1 : Nyai + Nysi)
        yfi = view(yi, Nyai + Nysi + 1 : Nyai + Nysi + Nyfi)
        yci = view(yi, Nyai + Nysi + Nyfi + 1 : Nyai + Nysi + Nyfi + Nyci)

        # save local inputs
        y[iya[iyas[i]]] .= yai # aerodynamic inputs
        y[iyf[iyfs[i]]] .= yfi # control surface inputs

        # section aerodynamic loads (in body frame)
        fi = CtCab*R'*SVector(yi[Nyai+1], yi[Nyai+2], yi[Nyai+3])
        mi = CtCab*R'*SVector(yi[Nyai+4], yi[Nyai+5], yi[Nyai+6])

        # add constant distributed loads (in body frame)
        poff = 1 + 6*npoint + 6*(i-1)
        fi += SVector(padd[poff+1], padd[poff+2], padd[poff+3])
        mi += SVector(padd[poff+4], padd[poff+5], padd[poff+6])

        # save distributed loads for this element (in body frame)
        yoff = 6*npoint + 6*(i-1)
        y[iys[yoff+1 : yoff+3]] = fi
        y[iys[yoff+4 : yoff+6]] = mi

        # add distributed loads to total forces and moments
        Ftot += ΔL*fi
        Mtot += cross(pe, ΔL*fi) + ΔL*mi

    end

    # save body frame linear/angular velocities
    yoff = 6*npoint + 6*nelem
    y[iys[yoff+1]] = ur
    y[iys[yoff+2]] = vr
    y[iys[yoff+3]] = wr
    y[iys[yoff+4]] = pr
    y[iys[yoff+5]] = qr
    y[iys[yoff+6]] = rr

    # save trim model inputs
    loads = SVector(Ftot[1], Ftot[2], Ftot[3], Mtot[1], Mtot[2], Mtot[3])

    for i = 1:NC
        idx = ctrl.force_indices[i]
        y[iyc[i]] = loads[idx]
    end

    return y
end

# --- performance overloads --- #

# TODO

# --- convenience methods --- #

function set_additional_parameters!(padd, model::LiftingLine, stru::GEBT,
    flap::LiftingLineFlaps{NF,NG,TF,TG}, ctrl::Trim{NC}; rho, point_conditions,
    element_loads, u, v, w, p, q, r, delta) where {NF,NG,NC,TF,TG}

    np = length(stru.icol_point)
    ne = length(stru.icol_elem)

    padd[1] = rho

    for ip = 1:np
        padd[1+6*(ip-1)+1] = point_conditions[6*(ip-1)+1]
        padd[1+6*(ip-1)+2] = point_conditions[6*(ip-1)+2]
        padd[1+6*(ip-1)+3] = point_conditions[6*(ip-1)+3]
        padd[1+6*(ip-1)+4] = point_conditions[6*(ip-1)+4]
        padd[1+6*(ip-1)+5] = point_conditions[6*(ip-1)+5]
        padd[1+6*(ip-1)+6] = point_conditions[6*(ip-1)+6]
    end

    for ie = 1:ne
        padd[1+6*np+6*(ie-1)+1] = element_loads[6*(ie-1)+1]
        padd[1+6*np+6*(ie-1)+2] = element_loads[6*(ie-1)+2]
        padd[1+6*np+6*(ie-1)+3] = element_loads[6*(ie-1)+3]
        padd[1+6*np+6*(ie-1)+4] = element_loads[6*(ie-1)+4]
        padd[1+6*np+6*(ie-1)+5] = element_loads[6*(ie-1)+5]
        padd[1+6*np+6*(ie-1)+6] = element_loads[6*(ie-1)+6]
    end

    padd[1 + 6*np + 6*ne + 1] = u
    padd[1 + 6*np + 6*ne + 2] = v
    padd[1 + 6*np + 6*ne + 3] = w
    padd[1 + 6*np + 6*ne + 4] = p
    padd[1 + 6*np + 6*ne + 5] = q
    padd[1 + 6*np + 6*ne + 6] = r
    padd[1 + 6*np + 6*ne + 6 + 1 : 1 + 6*np + 6*ne + 6 + NG - NC] = delta

    return padd
end

function separate_additional_parameters(model::LiftingLine, stru::GEBT,
    flap::LiftingLineFlaps{NF,NG,TF,TG}, ctrl::Trim{NC}, padd) where {NF,NG,NC,TF,TG}

    np = length(stru.icol_point)
    ne = length(stru.icol_elem)

    rho = padd[1]

    point_conditions = reshape(view(padd, 1 + 1 : 1 + 6*np), 6, np)

    element_loads = reshape(view(padd, 1 + 6*np + 1 : 1 + 6*np + 6*ne), 6, ne)

    u = padd[1 + 6*np + 6*ne + 1]
    v = padd[1 + 6*np + 6*ne + 2]
    w = padd[1 + 6*np + 6*ne + 3]
    p = padd[1 + 6*np + 6*ne + 4]
    q = padd[1 + 6*np + 6*ne + 5]
    r = padd[1 + 6*np + 6*ne + 6]
    delta = view(padd, 1 + 6*np + 6*ne + 6 + 1 : 1 + 6*np + 6*ne + 6 + NG - NC)

    return (rho = rho, point_conditions = point_conditions, element_loads = element_loads,
        u = u, v = v, w = w, p = p, q = q, r = r, delta = delta)
end


function runfig5()
    p = Params(Ln = 2000, W = 2000, Ls = 20, Ws = 20, scale = 40, λ = 5, α = 0,
    Δ = 1, d = 0, τ = 1)

    p = reconstruct(p, μN = 0.000001)
    sp = spectrumsweepb(p, 0:.25:5, true)
    savespectrum("fig5", p, sp[1], sp[2])

    p = reconstruct(p, μN = 0.6)
    sp = spectrumsweepb(p, 0:.25:5, true)
    savespectrum("fig5", p, sp[1], sp[2])

    p = reconstruct(p, μN = 1.192)
    sp = spectrumsweepb(p, 0:.25:5, true)
    savespectrum("fig5", p, sp[1], sp[2])

    p = reconstruct(p, EZ = SA[0, 2, 0], μN = 0.6)
    sp = spectrumsweepα(p, 0:.25:4, true)
    savespectrum("fig5", p, sp[1], sp[2])

    p = reconstruct(p, EZ = SA[0, 4, 0], μN = 1.192)
    sp = spectrumsweepα(p, 0:.25:4, true)
    savespectrum("fig5", p, sp[1], sp[2])
end


struct SingleSpecies <: DiscreteSpeciesPool
    species::Vector{Symbol}

end

"""
    DiscreteUnipartiteSpeciesPool
"""
struct DiscreteUnipartiteSpeciesPool <: DiscreteSpeciesPool
    species::Vector{Symbol}
    metaweb::Matrix
    function DiscreteUnipartiteSpeciesPool(
        s::Vector{Symbol},
        adjmat::T,
    ) where {U<:Number,T<:AbstractArray{U,2}}
        new(s, adjmat)
    end
end

function DiscreteUnipartiteSpeciesPool(
    en::T,
) where {T<:EcologicalNetworks.UnipartiteNetwork}
    DiscreteUnipartiteSpeciesPool(en.S, Matrix(en.edges))
end

function DiscreteUnipartiteSpeciesPool(; numspecies = 30, connectance = 0.1)
    nm = EcologicalNetworks.nichemodel(numspecies, connectance)
    return DiscreteUnipartiteSpeciesPool(nm)
end
function DiscreteUnipartiteSpeciesPool(s::Vector{ST}) where {ST<:String}
    DiscreteUnipartiteSpeciesPool(Symbol.(s), zeros(Int32, length(s), length(s)))
end
function DiscreteUnipartiteSpeciesPool(s::Vector{ST}) where {ST<:Symbol}
    DiscreteUnipartiteSpeciesPool(s, zeros(Int32, length(s), length(s)))
end

function DiscreteUnipartiteSpeciesPool(
    s::Vector{String},
    adjmat::T,
) where {U<:Number,T<:AbstractArray{U,2}}
    DiscreteUnipartiteSpeciesPool(Symbol.(s), adjmat)
end

struct DiscreteKpartiteSpeciesPool <: DiscreteSpeciesPool
    species::Vector{Symbol}
    metaweb::Matrix
    partitions::Int
end

struct ContinuousUnipartiteSpeciesPool <: ContinuousSpeciesPool
    traitdists::Any
end

struct ContinuousKpartiteSpeciesPool <: ContinuousSpeciesPool
    traitdists::Any
    partitions::Int
end

numspecies(sp::T) where {T<:DiscreteSpeciesPool} = length(species(sp))
species(sp::T) where {T<:DiscreteSpeciesPool} = sp.species
metaweb(sp::T) where {T<:DiscreteSpeciesPool} = sp.metaweb

# some interfaces to EN generators
"""
    layers(speciespool::ST; element=Biomass, dims=(100,100))

    Converts a species pool to a set of layers of type `element`.
    Should be able to apply generic function. Mostly used to generate initial conditions.
"""
function layers(
    speciespool::ST;
    element = Biomass,
    dims = (100, 100),
) where {ST<:DiscreteSpeciesPool}
    specieslist = species(speciespool)
    layerlist = []
    for sp in specieslist
        thislayer = zeros(element, dims...)
        push!(layerlist, thislayer)
    end
    return layerlist
end
layers(
    speciespool::SPT,
    element::ET;
    dims = (100, 100),
) where {SPT<:DiscreteSpeciesPool,ET<:Measurement} =
    layers(speciespool, element = element, dims = dims)

function layernames(sp::SPT) where {SPT<:DiscreteSpeciesPool}
    return Symbol.(sp.species)
end

"""
    speciespool(net::UnipartiteNetwork)
"""
function speciespool(net::UnipartiteNetwork)
    DiscreteUnipartiteSpeciesPool([Symbol(i) for i in net.S], Matrix(net.edges))
end

"""
    interesting idea but getting names is better for filtering
    because returning a sub-metaweb losess information about interactions

"""
function Base.filter(f::Function, sp::T) where {T<:DiscreteSpeciesPool}

    newspecieslist = []
    indecies = []

    for (i, spec) in enumerate(species(sp))
        if f(String(spec)) == true
            push!(newspecieslist, spec)
            push!(indecies, i)
        end
    end

    newmetaweb = sp.metaweb[indecies, indecies]

    return T(newspecieslist, newmetaweb)
end


#-------------------------------------------------------------------
#* EMSO Model Library (EML) Copyright (C) 2004 - 2007 ALSOC.
#*
#* This LIBRARY is free software; you can distribute it and/or modify
#* it under the therms of the ALSOC FREE LICENSE as available at
#* http://www.enq.ufrgs.br/alsoc.
#*
#* EMSO Copyright (C) 2004 - 2007 ALSOC, original code
#* from http://www.rps.eng.br Copyright (C) 2002-2004.
#* All rights reserved.
#*
#* EMSO is distributed under the therms of the ALSOC LICENSE as
#* available at http://www.enq.ufrgs.br/alsoc.
#*--------------------------------------------------------------------
#* Author: Gerson Balbueno Bicca
#* $Id: DoublePipeIncr.mso								$
#*------------------------------------------------------------------
type Details_Main
	Details_Main()=begin
		N=outers.N
		new(
			DanaInteger(Dict{Symbol,Any}(
				:Brief=>"Number of zones",
				:Default=>2
			)),
			area(Dict{Symbol,Any}(
				:Brief=>"Total Exchange Surface Area"
			)),
			fill(power(Dict{Symbol,Any}(
				:Brief=>"Incremental Duty",
				:Default=>7000,
				:Lower=>1e-6,
				:Upper=>1e10
			)),(N)),
			power(Dict{Symbol,Any}(
				:Brief=>"Total Duty",
				:Default=>7000,
				:Lower=>1e-6,
				:Upper=>1e10
			)),
			fill(heat_trans_coeff(Dict{Symbol,Any}(
				:Brief=>"Incremental Overall Heat Transfer Coefficient Clean",
				:Default=>1,
				:Lower=>1e-6,
				:Upper=>1e10
			)),(N)),
			fill(heat_trans_coeff(Dict{Symbol,Any}(
				:Brief=>"Incremental Overall Heat Transfer Coefficient Dirty",
				:Default=>1,
				:Lower=>1e-6,
				:Upper=>1e10
			)),(N)),
			[:N,],
			[:A,:Q,:Qtotal,:Uc,:Ud,]
		)
	end
	N::DanaInteger
	A::area
	Q::Array{power}
	Qtotal::power
	Uc::Array{heat_trans_coeff}
	Ud::Array{heat_trans_coeff}
	parameters::Array{Symbol,1}
	variables::Array{Symbol,1}
	attributes::Dict{Symbol,Any}
end
export Details_Main
function atributes(in::Details_Main,_::Dict{Symbol,Any})
	fields::Dict{Symbol,Any}=Dict{Symbol,Any}()
	fields[:Pallete]=false
	fields[:Brief]="to be documented"
	fields[:Info]="to be documented"
	drive!(fields,_)
	return fields
end
Details_Main(_::Dict{Symbol,Any})=begin
	newModel=Details_Main()
	newModel.attributes=atributes(newModel,_)
	newModel
end


using Statistics
using CUDA
using Oceananigans.Architectures
using Oceananigans.AbstractOperations: AbstractOperation

using OffsetArrays: OffsetArray
using Oceananigans.ImmersedBoundaries: ImmersedBoundaryGrid

import Base: getindex, size, show, minimum, maximum
import Statistics: mean

import Oceananigans.Fields: AbstractField, Field, minimum, maximum, mean, location, short_show, set!
import Oceananigans.Grids: new_data
import Oceananigans.BoundaryConditions: FieldBoundaryConditions

struct CubedSphereFaces{E, F}
    faces :: F
end

function CubedSphereFaces(faces::F) where F
    E = typeof(faces[1])
    return CubedSphereFaces{E, F}(faces)
end

@inline Base.getindex(f::CubedSphereFaces, i::Int) = @inbounds f.faces[i]

#####
##### Dispatch the world / insane type unions
#####

const CubedSphereData = CubedSphereFaces{<:OffsetArray}

# Some dispatch foo to make a type union for CubedSphereFaceField...
#
# Conformal cubed sphere grid wrapped in ImmersedBoundaryGrid:
const ImmersedConformalCubedSphereFaceGrid = ImmersedBoundaryGrid{FT, TX, TY, TZ, <:ConformalCubedSphereFaceGrid} where {FT, TX, TY, TZ}

# CubedSphereFaceField:
const NonImmersedCubedSphereFaceField = AbstractField{LX, LY, LZ, <:ConformalCubedSphereFaceGrid} where {LX, LY, LZ}
const ImmersedCubedSphereFaceField    = AbstractField{LX, LY, LZ, <:ImmersedConformalCubedSphereFaceGrid} where {LX, LY, LZ}

const CubedSphereFaceField = Union{NonImmersedCubedSphereFaceField{LX, LY, LZ},
                                      ImmersedCubedSphereFaceField{LX, LY, LZ}} where {LX, LY, LZ}

# CubedSphereField

# Flavors of CubedSphereField
const CubedSphereField{LX, LY, LZ} =
    Union{Field{LX, LY, LZ, <:Nothing, <:ConformalCubedSphereGrid},
          Field{LX, LY, LZ, <:AbstractOperation, <:ConformalCubedSphereGrid}}

const CubedSphereAbstractField{LX, LY, LZ} = AbstractField{LX, LY, LZ, <:ConformalCubedSphereGrid}

const AbstractCubedSphereField{LX, LY, LZ} =
    Union{CubedSphereAbstractField{LX, LY, LZ},
                  CubedSphereField{LX, LY, LZ}}

#####
##### new data
#####

function new_data(FT, grid::ConformalCubedSphereGrid, (LX, LY, LZ))
    faces = Tuple(new_data(FT, face_grid, (LX, LY, LZ)) for face_grid in grid.faces)
    return CubedSphereFaces{typeof(faces[1]), typeof(faces)}(faces)
end

#####
##### FieldBoundaryConditions
#####

function FieldBoundaryConditions(grid::ConformalCubedSphereGrid, (LX, LY, LZ); user_defined_bcs...)

    faces = Tuple(
        inject_cubed_sphere_exchange_boundary_conditions(
            FieldBoundaryConditions(face_grid, (LX, LY, LZ); user_defined_bcs...),
            face_index,
            grid.face_connectivity
        )
        for (face_index, face_grid) in enumerate(grid.faces)
    )

    return CubedSphereFaces{typeof(faces[1]), typeof(faces)}(faces)
end

#####
##### Utils
#####

function Base.show(io::IO, field::Union{CubedSphereField, AbstractCubedSphereField})
    LX, LY, LZ = location(field)
    arch = architecture(field)
    A = typeof(arch)
    return print(io, "$(typeof(field).name.wrapper) at ($LX, $LY, $LZ)\n",
          "├── architecture: $A\n",
          "└── grid: $(short_show(field.grid))")
end

@inline function interior(field::AbstractCubedSphereField)
    faces = Tuple(interior(face_field_face) for face_field in faces(field))
    return CubedSphereFaces(faces)
end

Base.size(field::CubedSphereField) = size(field.data)
Base.size(data::CubedSphereData) = (size(data.faces[1])..., length(data.faces))

@inline get_face(field::CubedSphereField, face_index) =
    Field(location(field), get_face(field.grid, face_index);
          data = get_face(field.data, face_index),
          boundary_conditions = get_face(field.boundary_conditions, face_index))

faces(field::AbstractCubedSphereField) = Tuple(get_face(field, face_index) for face_index in 1:length(field.data.faces))

#####
##### set!
#####

function cubed_sphere_set!(u, v)
    for face = 1:length(u.grid.faces)
        set!(get_face(u, face), get_face(v, face))
    end
    return nothing
end

# Resolve ambiguities
set!(u::CubedSphereField, v) = cubed_sphere_set!(u, v)
set!(u::CubedSphereField, v::Function) = cubed_sphere_set!(u, v)
set!(u::CubedSphereField, v::Union{Array, CuArray, OffsetArray}) = cubed_sphere_set!(u, v)
set!(u::CubedSphereField, v::CubedSphereField) = cubed_sphere_set!(u, v)

#####
##### Random utils
#####

minimum(field::AbstractCubedSphereField; dims=:) = minimum(minimum(face_field; dims) for face_field in faces(field))
maximum(field::AbstractCubedSphereField; dims=:) = maximum(maximum(face_field; dims) for face_field in faces(field))
mean(field::AbstractCubedSphereField; dims=:) = mean(mean(face_field; dims) for face_field in faces(field))

minimum(f::Function, field::AbstractCubedSphereField; dims=:) = minimum(minimum(f, face_field; dims) for face_field in faces(field))
maximum(f::Function, field::AbstractCubedSphereField; dims=:) = maximum(maximum(f, face_field; dims) for face_field in faces(field))
mean(f::Function, field::AbstractCubedSphereField; dims=:) = mean(mean(f, face_field; dims) for face_field in faces(field))

λnodes(field::CubedSphereFaceField{LX, LY, LZ}) where {LX, LY, LZ} = λnodes(LX(), LY(), LZ(), field.grid)
φnodes(field::CubedSphereFaceField{LX, LY, LZ}) where {LX, LY, LZ} = φnodes(LX(), LY(), LZ(), field.grid)


function get_constant_normalized_provider(obj::GravityFields, arg0::jint, arg1::jint)
    return jcall(obj, "getConstantNormalizedProvider", NormalizedSphericalHarmonicsProvider, (jint, jint), arg0, arg1)
end

function get_constant_unnormalized_provider(obj::GravityFields, arg0::jint, arg1::jint)
    return jcall(obj, "getConstantUnnormalizedProvider", UnnormalizedSphericalHarmonicsProvider, (jint, jint), arg0, arg1)
end

function get_normalized_provider(obj::GravityFields, arg0::jint, arg1::jint)
    return jcall(obj, "getNormalizedProvider", NormalizedSphericalHarmonicsProvider, (jint, jint), arg0, arg1)
end

function get_ocean_tides_waves(obj::GravityFields, arg0::jint, arg1::jint)
    return jcall(obj, "getOceanTidesWaves", List, (jint, jint), arg0, arg1)
end

function get_unnormalized_provider(obj::GravityFields, arg0::jint, arg1::jint)
    return jcall(obj, "getUnnormalizedProvider", UnnormalizedSphericalHarmonicsProvider, (jint, jint), arg0, arg1)
end


using Random
Random.seed!(209)

@testset "0D grids" begin
  @test Space((1,)).graph == Agents.Graph(1)
end

@testset "1D grids" begin
  a = Space((5,1))
  ae = collect(Agents.LightGraphs.edges(a.graph))

  b = Space((5,1), periodic=true)
  be = collect(Agents.LightGraphs.edges(b.graph))

  @test ae == [Agents.LightGraphs.Edge(1,2),Agents.LightGraphs.Edge(2,3), Agents.LightGraphs.Edge(3,4), Agents.LightGraphs.Edge(4,5)]
  @test be == [Agents.LightGraphs.Edge(1,2), Agents.LightGraphs.Edge(1,5), Agents.LightGraphs.Edge(2,3), Agents.LightGraphs.Edge(3,4), Agents.LightGraphs.Edge(4,5)]
end

@testset "2D grids" begin
  @test Space((2,1)).graph == Agents.Graph(2,1)

  a = Space((2,3))
  b = Space((2,3), periodic=true) # 2D grid

  @test a.dimensions == (2,3)
  @test b.dimensions == (2,3)
  @test length(a.agent_positions) == 6
  @test length(b.agent_positions) == 6

  @test Agents.nv(a) == 6
  @test Agents.ne(a) == 7
  @test Agents.nv(b) == 6
  @test Agents.ne(b) == 9

  ae = collect(Agents.LightGraphs.edges(a.graph))
  be = collect(Agents.LightGraphs.edges(b.graph))

  @test ae == [Agents.LightGraphs.Edge(1,2), Agents.LightGraphs.Edge(1,3), Agents.LightGraphs.Edge(2,4), Agents.LightGraphs.Edge(3,4), Agents.LightGraphs.Edge(3,5), Agents.LightGraphs.Edge(4,6), Agents.LightGraphs.Edge(5,6)]
  @test be == [Agents.LightGraphs.Edge(1,2), Agents.LightGraphs.Edge(1,3), Agents.LightGraphs.Edge(1,5), Agents.LightGraphs.Edge(2,4), Agents.LightGraphs.Edge(2,6), Agents.LightGraphs.Edge(3,4), Agents.LightGraphs.Edge(3,5), Agents.LightGraphs.Edge(4,6), Agents.LightGraphs.Edge(5,6)]
end

@testset "2D Moore" begin
  a = Space((3, 3), moore=true)
  @test Agents.nv(a) == 9
  @test Agents.ne(a) == 20

  a = Space((3, 4), moore=true)
  @test Agents.nv(a) == 12
  @test Agents.ne(a) == 29

  b = Space((3, 2), moore=true, periodic=true)
  @test Agents.nv(b) == 6
  @test Agents.ne(b) == 15

  b = Space((3, 3), moore=true, periodic=true)
  @test Agents.nv(b) == 9
  @test Agents.ne(b) == 36
end

@testset "3D grid" begin
  g1 = Space((2,3,2))
  g2 = Space((2,3,3))
  g3 = Space((2,3,2), periodic=true)
  g4 = Space((2,3,3), periodic=true)
  @test Agents.ne(g1) == 20
  @test Agents.ne(g2) == 33
  @test Agents.ne(g3) == 24
  @test Agents.ne(g4) == 39
end

@testset "grid coord/vertex conversions" begin
  @test coord2vertex((2,2,1), (3,4,1)) == 5
  @test coord2vertex((2,2), (3,4)) == 5
  @test coord2vertex((2,2,1), (3,4,1)) == 5
  @test coord2vertex((2,2), (3,4)) == 5
  @test vertex2coord(5, (3,4,1)) == (2,2,1)
  @test vertex2coord(5, (3,4)) == (2,2)

  @test coord2vertex((2,2,1), (2,3,3)) == 4
  @test coord2vertex((2,2,2), (2,3,3)) == 10
  @test coord2vertex((2,2,3), (2,3,3)) == 16
  @test coord2vertex(((1,3,1)), (2,3,3)) == 5
  @test coord2vertex((1,3,2), (2,3,3)) == 11
  @test coord2vertex((1,3,3), (2,3,3)) == 17

  @test vertex2coord(5, (2,3,3)) == (1,3,1)
  @test vertex2coord(7, (2,3,3)) == (1,1,2)
  @test vertex2coord(13, (2,3,3)) == (1,1,3)
  @test vertex2coord(13, (2,3,3)) == (1,1,3)
  @test vertex2coord(15, (2,3,3)) == (1,2,3)
  @test vertex2coord(18, (2,3,3)) == (2,3,3)
  @test vertex2coord(18, (2,3,3)) == (2,3,3)
end

@testset "nodes" begin
  space = Space((3,3))
  model = ABM(Agent1, space)
  for node in nodes(model)
    if rand() > 0.7
      add_agent!(node, model)
    end
  end
  @test nodes(model) == 1:9
  @test nodes(model, by=:random) == [8, 9, 1, 7, 6, 5, 3, 2, 4]
  @test nodes(model, by=:population) == [4, 5, 8, 9, 1, 2, 3, 6, 7]
  @test length(get_node_contents(4, model)) > length(get_node_contents(7, model))
  @test_throws ErrorException nodes(model, by=:notreal)
end


using HDF5, GeoData, Test, Statistics, Dates
using GeoData: Time, window, name
testpath = joinpath(dirname(pathof(GeoData)), "../test/")
include(joinpath(testpath, "test_utils.jl"))

# TODO example files without a login requirement
path1 = joinpath(testpath, "data/SMAP_L4_SM_gph_20160101T223000_Vv4011_001.h5")
path2 = joinpath(testpath, "data/SMAP_L4_SM_gph_20160102T223000_Vv4011_001.h5")

@testset "stack" begin
    stack = SMAPstack(path1)

    @testset "conversion to GeoArray" begin
        smaparray = stack["soil_temp_layer1"][Lat(), Lon()]
        @test smaparray isa GeoArray{Float32,2}
        @test dims(smaparray) isa Tuple{<:Lon{<:Array{Float32,1}}, <:Lat{<:Array{Float32,1}}}
        @test span(smaparray) isa Tuple{Irregular{Tuple{Float32,Float32}},Irregular{Tuple{Float32,Float32}}}
        @test span(smaparray) == (Irregular((-180.0f0, 180.0f0)), Irregular((-85.04456f0, 85.04456f0)))
        @test bounds(smaparray) == ((-180.0f0, 180.0f0), (-85.04456f0, 85.04456f0))
        @test index(smaparray) isa Tuple{Vector{Float32},Vector{Float32}}
        @test refdims(smaparray) isa Tuple{<:Ti}
        @test missingval(smaparray) == -9999.0
        @test smaparray[1] == -9999.0
        @test name(smaparray) == :soil_temp_layer1
        dt = DateTime(2016, 1, 1, 22, 30)
        step_ = Hour(3)
        @test refdims(stack) ==
            (Ti(dt:step_:dt; mode=Sampled(Ordered(), Regular(step_), Intervals(Start()))),)
        # Currently empty
        @test metadata(smaparray) isa SMAPstackMetadata
    end

    @testset "conversion to GeoStack" begin
        # Stack Constructors
        # This uses too much ram! There is a lingering memory leak in HDF5.
        # geostack = GeoStack(stack)
        # @test Symbol.(Tuple(keys(stack))) == keys(geostack)
        geostack = GeoStack(stack; keys=(:baseflow_flux, :snow_mass, :soil_temp_layer1))
        @test keys(geostack) == (:baseflow_flux, :snow_mass, :soil_temp_layer1)
    end

    if VERSION > v"1.1-"
        @testset "copy" begin
            geoarray = zero(stack[:soil_temp_layer1])
            @test geoarray isa GeoArray
            @test geoarray != stack[:soil_temp_layer1]
            copy!(geoarray, stack, :soil_temp_layer1)
            @test geoarray == stack[:soil_temp_layer1]
        end
    end

    @testset "window" begin
        windowedstack = SMAPstack(path1; window=(Lat(1:5), Lon(1:5), Ti(1)))
        @test window(windowedstack) == (Lat(1:5), Lon(1:5), Ti(1))
        windowedarray = windowedstack[:soil_temp_layer1]
        @test size(windowedarray) == (5, 5)
        @test windowedarray[1:3, 2:2] == reshape([-9999.0, -9999.0, -9999.0], 3, 1)
        @test windowedarray[1:3, 2] == [-9999.0, -9999.0, -9999.0]
        @test windowedarray[1, 2] == -9999.0
        windowedstack = SMAPstack(path1; window=(Lat(1:5), Lon(1:5), Ti(1:1)))
        windowedarray = windowedstack[:soil_temp_layer1]
        @test windowedarray[1:3, 2:2, 1] == reshape([-9999.0, -9999.0, -9999.0], 3, 1)
        @test windowedarray[1:3, 2, 1] == [-9999.0, -9999.0, -9999.0]
        @test windowedarray[1, 2, 1] == -9999.0
        windowedstack = SMAPstack(path1; window=(Ti(1),))
        windowedarray = windowedstack[:soil_temp_layer1]
        @test windowedarray[1:3, 2:2] == reshape([-9999.0, -9999.0, -9999.0], 3, 1)
        @test windowedarray[1:3, 2] == [-9999.0, -9999.0, -9999.0]
        @test windowedarray[1, 2] == -9999.0
    end

    @testset "show" begin
        sh1 = sprint(show, stack[:soil_temp_layer1])
        # Test but don't lock this down too much
        @test occursin("GeoArray", sh1)
        @test occursin("Latitude", sh1)
        @test occursin("Longitude", sh1)
        @test occursin("Time", sh1)
        sh2 = sprint(show, stack[:soil_temp_layer1][Lat(Between(0, 100)), Lon(Between(1, 100))])
        # Test but don't lock this down too much
        @test occursin("GeoArray", sh2)
        @test occursin("Latitude", sh2)
        @test occursin("Longitude", sh2)
        @test occursin("Time", sh2)
    end

end

@testset "series" begin
    series = SMAPseries([path1, path2]);
    val.(dims(series))
    @test series[1] isa SMAPstack
    @test first(bounds(series, Ti)) == DateTime(2016, 1, 1, 22, 30)
    @test last(bounds(series, Ti)) == DateTime(2016, 1, 3, 1, 30)
    modified_series = modify(Array, series)
    stackkeys = keys(modified_series[1])
    @test typeof(modified_series) <: GeoSeries{<:GeoStack{<:NamedTuple{stackkeys,<:Tuple{<:GeoArray{Float32,2,<:Tuple,<:Tuple,<:Array{Float32,2}},Vararg}}}}
end


using Test
using Printf
using Random
using Statistics
using LinearAlgebra
using Logging

using CUDA
using JLD2
using FFTW
using OffsetArrays
using SeawaterPolynomials

using Oceananigans
using Oceananigans.Architectures
using Oceananigans.Grids
using Oceananigans.Operators
using Oceananigans.Advection
using Oceananigans.BoundaryConditions
using Oceananigans.Fields
using Oceananigans.Coriolis
using Oceananigans.Buoyancy
using Oceananigans.Forcings
using Oceananigans.Solvers
using Oceananigans.Models
using Oceananigans.Simulations
using Oceananigans.Diagnostics
using Oceananigans.OutputWriters
using Oceananigans.TurbulenceClosures
using Oceananigans.AbstractOperations
using Oceananigans.Logger
using Oceananigans.Utils
using Oceananigans.Architectures: device # to resolve conflict with CUDA.device

using Dates: DateTime, Nanosecond
using TimesDates: TimeDate
using Statistics: mean
using LinearAlgebra: norm
using NCDatasets: Dataset
using KernelAbstractions: @kernel, @index, Event

import Oceananigans.Fields: interior
import Oceananigans.Utils: launch!, datatuple

Logging.global_logger(OceananigansLogger())

#####
##### Testing parameters
#####

float_types = (Float32, Float64)

         archs = (CPU(),)
@hascuda archs = (GPU(),)

closures = (
    :IsotropicDiffusivity,
    :AnisotropicDiffusivity,
    :AnisotropicBiharmonicDiffusivity,
    :TwoDimensionalLeith,
    :SmagorinskyLilly,
    :BlasiusSmagorinsky,
    :RozemaAnisotropicMinimumDissipation,
    :VerstappenAnisotropicMinimumDissipation
)

#####
##### Run tests!
#####

CUDA.allowscalar(true)

include("utils_for_runtests.jl")

group = get(ENV, "TEST_GROUP", :all) |> Symbol

@testset "Oceananigans" begin
    if group == :unit || group == :all
        @testset "Unit tests" begin
            include("test_grids.jl")
            include("test_operators.jl")
            include("test_boundary_conditions.jl")
            include("test_fields.jl")
            include("test_averaged_field.jl")
            include("test_kernel_computed_field.jl")
            include("test_halo_regions.jl")
            include("test_solvers.jl")
            include("test_pressure_solvers.jl")
            include("test_coriolis.jl")
            include("test_buoyancy.jl")
            include("test_surface_waves.jl")
            include("test_utils.jl")
        end
    end

    if group == :time_stepping_1 || group == :all
        @testset "Model and time stepping tests (part 1)" begin
            include("test_incompressible_models.jl")
            include("test_time_stepping.jl")
        end
    end

    if group == :time_stepping_2 || group == :all
        @testset "Model and time stepping tests (part 2)" begin
            include("test_boundary_conditions_integration.jl")
            include("test_forcings.jl")
            include("test_turbulence_closures.jl")
            include("test_dynamics.jl")
        end
    end

    if group == :simulation || group == :all
        @testset "Simulation tests" begin
            include("test_simulations.jl")
            include("test_diagnostics.jl")
            include("test_output_writers.jl")
            include("test_abstract_operations_computed_field.jl")
        end
    end

    if group == :regression || group == :all
        include("test_regression.jl")
    end

    if group == :scripts || group == :all
        @testset "Scripts" begin
            include("test_validation.jl")
        end
    end

    if group == :convergence
        include("test_convergence.jl")
    end

    if group == :shallow_water || group == :all
        include("test_shallow_water_models.jl")
    end
end


using Pkg

if !haskey(Pkg.installed(), "GaussDCA")
    Pkg.add(PackageSpec(url="https://github.com/carlobaldassi/GaussDCA.jl"))
end

msa = map(Residue, rand(1:21,100,20))
dca = gaussdca(msa, min_separation=2)

@test  isnan(dca[1,1])
@test  isnan(dca[1,2])
@test !isnan(dca[1,3])


module FTensorKit
using GeometryBasics, LinearAlgebra

include("utils/Lattices.jl")
include("utils/PerspectiveMap.jl")
include("utils/SwapGates.jl")
include("eincode.jl")
include("FTN.jl")

end


include("regression.jl")
include("classification.jl")
include("multiclass.jl")
include("event.jl")
include("generic_likelihood.jl")
function pdf(l::Likelihood{T},y::Real,f::Real) where T
    @error "pdf not implemented"
end

@inline logpdf(l::Likelihood{T},y::Real,f::Real) where T = log(pdf(l,y,f))

isaugmented(::Likelihood) = false

Base.length(::Likelihood) = 1

Base.iterate(l::Likelihood) = (l,nothing)
Base.iterate(l::Likelihood, ::Any) = nothing


module ViscosityDrop

using Oceananigans
using Optim
using JLD2
using DelimitedFiles

using Oceananigans.Grids
include("simulation/grid.jl")
export make_grid

include("simulation/boundary_conditions.jl")
export no_slip_bc

using Oceananigans.Advection
include("simulation/model.jl")
export make_model

include("simulation/buoyancy.jl")
export set_buoyancy!

using Oceananigans.AbstractOperations
using Oceananigans.Fields
using Oceananigans.OutputWriters
include("simulation/simulation.jl")
export make_simulation, run!

include("simulation/wrapper.jl")
export do_all

include("post_process/circle.jl")
include("post_process/fit_functions.jl")
include("post_process/get_results.jl")
export get_results

end # module


module Example

using Twinkle, Twinkle.FlexUI
using Colors
using StaticArrays
using GeometryBasics
import Luxor

println("Start [$(splitext(basename(@__FILE__))[1])]")

app = App()
prop!(app, :title, "Twinkle Example - Luxor")

#---------------------------------------------------------------
# Define Variables
#---------------------------------------------------------------

angle = addVariable!(app, Variable(name="angle", type="flota64",value=8))
count = addVariable!(app, Variable(name="count", type="flota64",value=16))

image = addVariable!(app, Variable(
    name="image",
    type="image",
    value="",
))

#---------------------------------------------------------------
# Define Controls
#---------------------------------------------------------------

ui = VContainer(
    H1Label("Value Controls"),
    Slider(
        text="Angle (π over value)",
        trailing_text="[\$()]",
        min=2,
        max=16,
        value=8,
        variable="angle"
    ),
    Slider(
        text="Count",
        trailing_text="[\$()]",
        min=1,
        max=30,
        value=16,
        variable="count"
    ),
    H1Label("Static Image Example"),
    Image(
        source="\$(image)",
        # height="50%",
    ),
    H1Label("Image Viewer allowing Pan (left-drag) and Zoom (wheel)"),
    PanZoom(
        style="width: 100%; height=400px;",
        content = Container(
            direction = "row warp",
            children = [
                Image(
                    source="\$(image)",
                ),
            ]
        ),
    ),

)
FlexUI.controls!(app, ui)

#---------------------------------------------------------------
# the render function
#---------------------------------------------------------------
function render()
    Luxor.Drawing(600, 400, :svg)
    Luxor.origin()
    Luxor.background("white")

    for θ in range(0, step=π/angle[], length=count[])
        Luxor.gsave()
        Luxor.scale(0.2)
        Luxor.rotate(θ)
        Luxor.translate(350, 0)
        Luxor.julialogo(action=:fill, bodycolor=Luxor.randomhue())
        Luxor.grestore()
    end

    Luxor.gsave()
    Luxor.scale(0.3)
    Luxor.juliacircles()
    Luxor.grestore()

    Luxor.translate(150, -150)
    Luxor.scale(0.3)
    Luxor.julialogo()
    Luxor.finish()

    image["svg"] = Luxor.preview()
end
renderFunction!(app, render)

#---------------------------------------------------------------
# Run the application
#---------------------------------------------------------------
run(app)

println("End [$(splitext(basename(@__FILE__))[1])]")

end # module


# ExactReach, ai2, maxSens

using NeuralVerification, LazySets, Test, LinearAlgebra, GLPKMathProgInterface
import NeuralVerification: ReLU, Id

macro no_error(ex)
    quote
        try $(esc(ex))
            true
        catch
            false
        end
    end
end

small_nnet_file = "$(@__DIR__)/../examples/networks/small_nnet.nnet"
small_nnet_id_file = "$(@__DIR__)/../examples/networks/small_nnet_id.nnet"

mnist1_file = "$(@__DIR__)/../examples/networks/mnist1.nnet"
mnist2_file = "$(@__DIR__)/../examples/networks/mnist2.nnet"
mnist3_file = "$(@__DIR__)/../examples/networks/mnist3.nnet"
mnist4_file = "$(@__DIR__)/../examples/networks/mnist4.nnet"

small_nnet = mnist1 = read_nnet(small_nnet_file, last_layer_activation = ReLU())
small_nnet_id = mnist1 = read_nnet(small_nnet_id_file, last_layer_activation = Id())

mnist1 = read_nnet(mnist1_file, last_layer_activation = Id())
mnist2 = read_nnet(mnist2_file, last_layer_activation = Id())
mnist3 = read_nnet(mnist3_file, last_layer_activation = Id())
mnist4 = read_nnet(mnist4_file, last_layer_activation = Id())

println("###### Problem type - input:HPolytope, output:HPolytope ######")
println("###### Network: small_nnet, problem: holds              ######")

in_hyper  = Hyperrectangle(low = [-0.9], high = [0.9])
in_hpoly  = convert(HPolytope, in_hyper)
#out_superset    = Hyperrectangle(low = [30.0], high = [80.0])

problem_holds    = Problem(small_nnet, in_hyper, HPolytope([HalfSpace([1.], 100.)]))

solver = FastLin(10, 10.0, 1.0)
println("$(typeof(solver)) - small_nnet")
timed_result =@timed solve(solver, problem_holds)
println(" - Time: " * string(timed_result[2]) * " s")
println(" - Output: ")
println(timed_result[1])
println("")

solver = FastLip(10, 10.0, 1.0)
println("$(typeof(solver)) - small_nnet")
timed_result =@timed solve(solver, problem_holds)
println(" - Time: " * string(timed_result[2]) * " s")
println(" - Output: ")
println(timed_result[1])
println("")

optimizer = GLPKSolverMIP()
solver = MIPVerify(optimizer)
println("$(typeof(solver)) - small_nnet")
timed_result =@timed solve(solver, problem_holds)
println(" - Time: " * string(timed_result[2]) * " s")
println(" - Output: ")
println(timed_result[1])
println("")

optimizer = GLPKSolverMIP()
solver = ILP(optimizer, 1)
println("$(typeof(solver)) - small_nnet")
timed_result =@timed solve(solver, problem_holds)
println(" - Time: " * string(timed_result[2]) * " s")
println(" - Output: ")
println(timed_result[1])
println("")

println("###### Problem type - input:HPolytope, output:HPolytope ######")

# entry 23 in MNIST datset
input_center = [0.0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,121,254,136,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,13,230,253,248,99,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,4,118,253,253,225,42,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,61,253,253,253,74,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,32,206,253,253,186,9,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,211,253,253,239,69,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,254,253,253,133,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,142,255,253,186,8,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,149,229,254,207,21,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,54,229,253,254,105,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,152,254,254,213,26,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,112,251,253,253,26,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,29,212,253,250,149,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,36,214,253,253,137,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,75,253,253,253,59,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,93,253,253,189,17,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,224,253,253,84,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,43,235,253,126,1,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,99,248,253,119,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,225,235,49,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0]
output_center = [-1311.1257826380004,4633.767704436501,-654.0718535670002,-1325.349417307,1175.2361184373997,-1897.8607293569007,-470.3405972940001,830.8337987382,-377.7467076115001,572.3674015264198]

in_epsilon = 1.0 #0-255
out_epsilon = 10.0 #logit domain

input_low = input_center .- in_epsilon
input_high = input_center .+ in_epsilon

output_low = output_center .- out_epsilon
output_high = output_center .+ out_epsilon

inputSet = Hyperrectangle(low=input_low, high=input_high)

A = Matrix(undef, 2, 1)
A = [1.0, -1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0 ,0.0]'
b = [0.0]
outputSet = HPolytope(A, b)

problem_mnist1 = Problem(mnist1, inputSet, outputSet)
## MNIST1 ##
print("\n\n\n")
println("###### Network: mnist1                                  ######")
print("\n\n\n")

solver = FastLin(10, 10.0, 1.0)
println("$(typeof(solver)) - mnist1")
timed_result =@timed solve(solver, problem_mnist1)
print(" - Time: " * string(timed_result[2]) * " s")
print(" - Output: ")
print(timed_result[1])
println("")

solver = FastLip(10, 10.0, 1.0)
println("$(typeof(solver)) - mnist1")
timed_result =@timed solve(solver, problem_mnist1)
print(" - Time: " * string(timed_result[2]) * " s")
print(" - Output: ")
print(timed_result[1])
println("")

optimizer = GLPKSolverMIP()
solver = MIPVerify(optimizer)
println("$(typeof(solver)) - mnist1")
timed_result =@timed solve(solver, problem_mnist1)
println(" - Time: " * string(timed_result[2]) * " s")
println(" - Output: ")
println(timed_result[1])
println("")

optimizer = GLPKSolverMIP()
solver = ILP(optimizer, 1)
println("$(typeof(solver)) - mnist1")
timed_result =@timed solve(solver, problem_mnist1)
println(" - Time: " * string(timed_result[2]) * " s")
println(" - Output: ")
println(timed_result[1])
println("")

problem_mnist2 = Problem(mnist2, inputSet, outputSet)
## MNIST1 ##
print("\n\n\n")
println("###### Network: mnist2                                  ######")
print("\n\n\n")

solver = FastLin(10, 10.0, 1.0)
println("$(typeof(solver)) - mnist2")
timed_result =@timed solve(solver, problem_mnist2)
print(" - Time: " * string(timed_result[2]) * " s")
print(" - Output: ")
print(timed_result[1])
println("")

solver = FastLip(10, 10.0, 1.0)
println("$(typeof(solver)) - mnist2")
timed_result =@timed solve(solver, problem_mnist2)
print(" - Time: " * string(timed_result[2]) * " s")
print(" - Output: ")
print(timed_result[1])
println("")

optimizer = GLPKSolverMIP()
solver = MIPVerify(optimizer)
println("$(typeof(solver)) - mnist2")
timed_result =@timed solve(solver, problem_mnist2)
println(" - Time: " * string(timed_result[2]) * " s")
println(" - Output: ")
println(timed_result[1])
println("")

optimizer = GLPKSolverMIP()
solver = ILP(optimizer, 1)
println("$(typeof(solver)) - mnist2")
timed_result =@timed solve(solver, problem_mnist2)
println(" - Time: " * string(timed_result[2]) * " s")
println(" - Output: ")
println(timed_result[1])
println("")

problem_mnist3 = Problem(mnist3, inputSet, outputSet)
## MNIST1 ##
print("\n\n\n")
println("###### Network: mnist3                                  ######")
print("\n\n\n")

solver = FastLin(10, 10.0, 1.0)
println("$(typeof(solver)) - mnist3")
timed_result =@timed solve(solver, problem_mnist3)
print(" - Time: " * string(timed_result[2]) * " s")
print(" - Output: ")
print(timed_result[1])
println("")

solver = FastLip(10, 10.0, 1.0)
println("$(typeof(solver)) - mnist3")
timed_result =@timed solve(solver, problem_mnist3)
print(" - Time: " * string(timed_result[2]) * " s")
print(" - Output: ")
print(timed_result[1])
println("")

optimizer = GLPKSolverMIP()
solver = MIPVerify(optimizer)
println("$(typeof(solver)) - mnist3")
timed_result =@timed solve(solver, problem_mnist3)
println(" - Time: " * string(timed_result[2]) * " s")
println(" - Output: ")
println(timed_result[1])
println("")

optimizer = GLPKSolverMIP()
solver = ILP(optimizer, 1)
println("$(typeof(solver)) - mnist3")
timed_result =@timed solve(solver, problem_mnist3)
println(" - Time: " * string(timed_result[2]) * " s")
println(" - Output: ")
println(timed_result[1])
println("")

problem_mnist4 = Problem(mnist4, inputSet, outputSet)
## MNIST1 ##
print("\n\n\n")
println("###### Network: mnist4                                  ######")
print("\n\n\n")

solver = FastLin(10, 10.0, 1.0)
println("$(typeof(solver)) - mnist4")
timed_result =@timed solve(solver, problem_mnist4)
print(" - Time: " * string(timed_result[2]) * " s")
print(" - Output: ")
print(timed_result[1])
println("")

solver = FastLip(10, 10.0, 1.0)
println("$(typeof(solver)) - mnist4")
timed_result =@timed solve(solver, problem_mnist4)
print(" - Time: " * string(timed_result[2]) * " s")
print(" - Output: ")
print(timed_result[1])
println("")

#=
optimizer = GLPKSolverMIP()
solver = MIPVerify(optimizer)
println("$(typeof(solver)) - mnist4")
timed_result =@timed solve(solver, problem_mnist4)
println(" - Time: " * string(timed_result[2]) * " s")
println(" - Output: ")
println(timed_result[1])
println("")
=#
optimizer = GLPKSolverMIP()
solver = ILP(optimizer, 1)
println("$(typeof(solver)) - mnist4")
timed_result =@timed solve(solver, problem_mnist4)
println(" - Time: " * string(timed_result[2]) * " s")
println(" - Output: ")
println(timed_result[1])
println("")

# ACAS
print("\n\n\n")
println("###### Problem type - input:HPolytope, output:HPolytope ######")
print("\n\n\n")
println("###### Network: acas                                    ######")
print("\n\n\n")

#acas_file = "$(@__DIR__)/../examples/networks/ACASXU_run2a_1_1_tiny_4.nnet"
acas_file = "$(@__DIR__)/../examples/networks/ACASXU_run2a_4_5_batch_2000.nnet"
#ACASXU_run2a_4_5_batch_2000.nnet

acas_nnet = read_nnet(acas_file, last_layer_activation = Id())

# ACAS PROPERTY 10 - modified
# Original input range:
# LOWER BOUND: array([[ 0.21466922,  0.11140846, -0.4999999 ,  0.3920202 ,  0.15      ]])
# UPPER BOUND: array([[ 0.58819589,  0.4999999 , -0.49840835,  0.66474747,  0.65      ]])
A0 = Matrix{Float64}(I, 5, 5)
A1 = -Matrix{Float64}(I, 5, 5)
A = vcat(A0, A1)

#b_lower = [0.21466922, 0.11140846, -0.4999999, 0.52838384, 0.4]
#b_upper = [0.58819589, 0.4999999 , -0.4999999, 0.52838384, 0.4]

#b_lower = [ 0.21466922,  0.11140846, -0.4999999 ,  0.3920202 ,  0.15      ]
#b_upper = [ 0.58819589,  0.4999999 , -0.49840835,  0.66474747,  0.65      ]

#b_lower = [ 0.21466922,  0.11140846, -0.4999999 ,  0.3920202 ,  0.20      ]
#b_upper = [ 0.58819589,  0.4999999 , -0.49840835,  0.66474747,  0.65      ]

b_lower = [ 0.21466922,  0.11140846, -0.4999999 ,  0.3920202 ,  0.4      ]
b_upper = [ 0.58819589,  0.4999999 , -0.49840835,  0.66474747,  0.4      ]

in_hyper  = Hyperrectangle(low = b_lower, high = b_upper)
inputSet = convert(HPolytope, in_hyper)

#inputSet = HPolytope(A, b)

#A = Matrix(undef, 2, 1)
#A = [1.0, 0.0, 0.0, 0.0, -1.0]'
#b = [0.0]
#outputSet = HPolytope(A, b)

outputSet = HPolytope([HalfSpace([1.0, 0.0, 0.0, 0.0, -1.0], 0.0)])

problem_polytope_polytope_acas = Problem(acas_nnet, in_hyper, outputSet)

print("\n\n\n")
println("###### Network: acas                                  ######")
print("\n\n\n")

solver = FastLin(10, 10.0, 1.0)
println("$(typeof(solver)) - acas")
timed_result =@timed solve(solver, problem_polytope_polytope_acas)
print(" - Time: " * string(timed_result[2]) * " s")
print(" - Output: ")
print(timed_result[1])
println("")

solver = FastLip(10, 10.0, 1.0)
println("$(typeof(solver)) - acas")
timed_result =@timed solve(solver, problem_polytope_polytope_acas)
print(" - Time: " * string(timed_result[2]) * " s")
print(" - Output: ")
print(timed_result[1])
println("")

#=
optimizer = GLPKSolverMIP()
solver = MIPVerify(optimizer)
println("$(typeof(solver)) - acas")
timed_result =@timed solve(solver, problem_polytope_polytope_acas)
println(" - Time: " * string(timed_result[2]) * " s")
println(" - Output: ")
println(timed_result[1])
println("")
=#
optimizer = GLPKSolverMIP()
solver = ILP(optimizer, 1)
println("$(typeof(solver)) - acas")
timed_result =@timed solve(solver, problem_polytope_polytope_acas)
println(" - Time: " * string(timed_result[2]) * " s")
println(" - Output: ")
println(timed_result[1])
println("")


using StaticArrays
# -----------------------------------------------------
#
# Datasets from Time Series Analysis by Box and Jenkins
#
# -----------------------------------------------------

const seriesA = SA[17.0, 16.6, 16.3, 16.1, 17.1, 16.9, 16.8, 17.4, 17.1, 17.0, 16.7, 17.4, 17.2, 17.4,
                17.4, 17.0, 17.3, 17.2, 17.4, 16.8, 17.1, 17.4, 17.4, 17.5, 17.4, 17.6, 17.4, 17.3,
                17.0, 17.8, 17.5, 18.1, 17.5, 17.4, 17.4, 17.1, 17.6, 17.7, 17.4, 17.8, 17.6, 17.5,
                16.5, 17.8, 17.3, 17.3, 17.1, 17.4, 16.9, 17.3, 17.6, 16.9, 16.7, 16.8, 16.8, 17.2,
                16.8, 17.6, 17.2, 16.6, 17.1, 16.9, 16.6, 18.0, 17.2, 17.3, 17.0, 16.9, 17.3, 16.8,
                17.3, 17.4, 17.7, 16.8, 16.9, 17.0, 16.9, 17.0, 16.6, 16.7, 16.8, 16.7, 16.4, 16.5,
                16.4, 16.6, 16.5, 16.7, 16.4, 16.4, 16.2, 16.4, 16.3, 16.4, 17.0, 16.9, 17.1, 17.1,
                16.7, 16.9, 16.5, 17.2, 16.4, 17.0, 17.0, 16.7, 16.2, 16.6, 16.9, 16.5, 16.6, 16.6,
                17.0, 17.1, 17.1, 16.7, 16.8, 16.3, 16.6, 16.8, 16.9, 17.1, 16.8, 17.0, 17.2, 17.3,
                17.2, 17.3, 17.2, 17.2, 17.5, 16.9, 16.9, 16.9, 17.0, 16.5, 16.7, 16.8, 16.7, 16.7,
                16.6, 16.5, 17.0, 16.7, 16.7, 16.9, 17.4, 17.1, 17.0, 16.8, 17.2, 17.2, 17.4, 17.2,
                16.9, 16.8, 17.0, 17.4, 17.2, 17.2, 17.1, 17.1, 17.1, 17.4, 17.2, 16.9, 16.9, 17.0,
                16.7, 16.9, 17.3, 17.8, 17.8, 17.6, 17.5, 17.0, 16.9, 17.1, 17.2, 17.4, 17.5, 17.9,
                17.0, 17.0, 17.0, 17.2, 17.3, 17.4, 17.4, 17.0, 18.0, 18.2, 17.6, 17.8, 17.7, 17.2,
                17.4]

const seriesB = SA[460, 457, 452, 459, 462, 459, 463, 479, 493, 490, 492, 498, 499, 497, 496,
                490, 489, 478, 487, 491, 487, 482, 479, 478, 479, 477, 479, 475, 479, 476,
                476, 478, 479, 477, 476, 475, 475, 473, 474, 474, 474, 465, 466, 467, 471,
                471, 467, 473, 481, 488, 490, 489, 489, 485, 491, 492, 494, 499, 498, 500,
                497, 494, 495, 500, 504, 513, 511, 514, 510, 509, 515, 519, 523, 519, 523,
                531, 547, 551, 547, 541, 545, 549, 545, 549, 547, 543, 540, 539, 532, 517,
                527, 540, 542, 538, 541, 541, 547, 553, 559, 557, 557, 560, 571, 571, 569,
                575, 580, 584, 585, 590, 599, 603, 599, 596, 585, 587, 585, 581, 583, 592,
                592, 596, 596, 595, 598, 598, 595, 595, 592, 588, 582, 576, 578, 589, 585,
                580, 579, 584, 581, 581, 577, 577, 578, 580, 586, 583, 581, 576, 571, 575,
                575, 573, 577, 582, 584, 579, 572, 577, 571, 560, 549, 556, 557, 563, 564,
                567, 561, 559, 553, 553, 553, 547, 550, 544, 541, 532, 525, 542, 555, 558,
                551, 551, 552, 553, 557, 557, 548, 547, 545, 545, 539, 539, 535, 537, 535,
                536, 537, 543, 548, 546, 547, 548, 549, 553, 553, 552, 551, 550, 553, 554,
                551, 551, 545, 547, 547, 537, 539, 538, 533, 525, 513, 510, 521, 521, 521,
                523, 516, 511, 518, 517, 520, 519, 519, 519, 518, 513, 499, 485, 454, 462,
                473, 482, 486, 475, 459, 451, 453, 446, 455, 452, 457, 449, 450, 435, 415,
                398, 399, 361, 383, 393, 385, 360, 364, 365, 370, 374, 359, 335, 323, 306,
                333, 330, 336, 328, 316, 320, 332, 320, 333, 344, 339, 350, 351, 350, 345,
                350, 359, 375, 379, 376, 382, 370, 365, 367, 372, 373, 363, 371, 369, 376,
                387, 387, 376, 385, 385, 380, 373, 382, 377, 376, 379, 386, 387, 386, 389,
                394, 393, 409, 411, 409, 408, 393, 391, 388, 396, 387, 383, 388, 382, 384,
                382, 383, 383, 388, 395, 392, 386, 383, 377, 364, 369, 355, 350, 353, 340,
                350, 349, 358, 360, 360, 366, 359, 356, 355, 367, 357, 361, 355, 348, 343,
                330, 340, 339, 331, 345, 352, 346, 352, 357]

const seriesB2 = SA[445, 448, 450, 447, 451, 453, 454, 454, 459, 440, 446, 443, 443, 440, 439, 435,
                435, 436, 435, 435, 435, 433, 429, 428, 425, 427, 425, 422, 409, 407, 423, 422,
                417, 421, 424, 414, 419, 429, 426, 425, 424, 425, 425, 424, 425, 421, 414, 410,
                411, 406, 406, 413, 411, 410, 405, 409, 410, 405, 401, 401, 401, 414, 419, 425,
                423, 411, 414, 420, 412, 415, 412, 412, 411, 412, 409, 407, 408, 415, 413, 413,
                410, 405, 410, 412, 413, 411, 411, 409, 406, 407, 410, 408, 408, 409, 410, 409,
                405, 406, 405, 407, 409, 407, 409, 425, 425, 428, 436, 442, 442, 433, 435, 433,
                435, 429, 439, 437, 439, 438, 435, 433, 437, 437, 444, 441, 440, 441, 439, 439,
                438, 437, 441, 442, 441, 437, 427, 423, 424, 428, 428, 431, 425, 423, 420, 426,
                418, 416, 419, 418, 416, 419, 425, 421, 422, 422, 417, 420, 417, 418, 419, 419,
                417, 419, 422, 423, 422, 421, 421, 419, 418, 421, 420, 413, 413, 408, 409, 415,
                415, 420, 420, 424, 426, 423, 423, 425, 431, 436, 436, 440, 436, 443, 445, 439,
                443, 445, 450, 461, 471, 467, 462, 456, 464, 463, 465, 464, 456, 460, 458, 453,
                453, 449, 447, 453, 450, 459, 457, 453, 455, 453, 450, 456, 461, 463, 463, 461,
                465, 473, 473, 475, 499, 485, 491, 496, 504, 504, 509, 511, 524, 525, 541, 531,
                529, 530, 531, 527, 525, 519, 514, 509, 505, 513, 525, 519, 519, 522, 522]

const seriesC = SA[26.6, 27.0, 27.1, 27.1, 27.1, 27.1, 26.9, 26.8, 26.7, 26.4, 26.0, 25.8, 25.6, 25.2,
                25.0, 24.6, 24.2, 24.0, 23.7, 23.4, 23.1, 22.9, 22.8, 22.7, 22.6, 22.4, 22.2, 22.0,
                21.8, 21.4, 20.9, 20.3, 19.7, 19.4, 19.3, 19.2, 19.1, 19.0, 18.9, 18.9, 19.2, 19.3,
                19.3, 19.4, 19.5, 19.6, 19.6, 19.6, 19.6, 19.6, 19.7, 19.9, 20.0, 20.1, 20.2, 20.3,
                20.6, 21.6, 21.9, 21.7, 21.3, 21.2, 21.4, 21.7, 22.2, 23.0, 23.8, 24.6, 25.1, 25.6,
                25.8, 26.1, 26.3, 26.3, 26.2, 26.0, 25.8, 25.6, 25.4, 25.2, 24.9, 24.7, 24.5, 24.4,
                24.4, 24.4, 24.4, 24.4, 24.3, 24.4, 24.4, 24.4, 24.4, 24.4, 24.5, 24.5, 24.4, 24.3,
                24.2, 24.2, 24.0, 23.9, 23.7, 23.6, 23.5, 23.5, 23.5, 23.5, 23.5, 23.7, 23.8, 23.8,
                23.9, 23.9, 23.8, 23.7, 23.6, 23.4, 23.2, 23.0, 22.8, 22.6, 22.4, 22.0, 21.6, 21.3,
                21.2, 21.2, 21.1, 21.0, 20.9, 21.0, 21.0, 21.1, 21.2, 21.1, 20.9, 20.8, 20.8, 20.8,
                20.8, 20.9, 20.8, 20.8, 20.7, 20.7, 20.8, 20.9, 21.2, 21.4, 21.7, 21.8, 21.9, 22.2,
                22.5, 22.8, 23.1, 23.4, 23.8, 24.1, 24.6, 24.9, 24.9, 25.1, 25.0, 25.0, 25.0, 25.0,
                24.9, 24.8, 24.7, 24.6, 24.5, 24.5, 24.5, 24.5, 24.5, 24.5, 24.5, 24.4, 24.4, 24.2,
                24.2, 24.1, 24.1, 24.0, 24.0, 24.0, 23.9, 23.8, 23.8, 23.7, 23.7, 23.6, 23.7, 23.6,
                23.6, 23.6, 23.5, 23.5, 23.4, 23.3, 23.3, 23.3, 23.4, 23.4, 23.3, 23.2, 23.3, 23.3,
                23.2, 23.1, 22.9, 22.8, 22.6, 22.4, 22.2, 21.8, 21.3, 20.8, 20.2, 19.7, 19.3, 19.1,
                19.0, 18.8]

const seriesD = SA[8.0, 8.0, 7.4, 8.0, 8.0, 8.0, 8.0, 8.8, 8.4, 8.4, 8.0, 8.2, 8.2, 8.2,
                8.4, 8.4, 8.4, 8.6, 8.8, 8.6, 8.6, 8.6, 8.6, 8.6, 8.8, 8.9, 9.1, 9.5,
                8.5, 8.4, 8.3, 8.2, 8.1, 8.3, 8.4, 8.7, 8.8, 8.8, 9.2, 9.6, 9.0, 8.8,
                8.6, 8.6, 8.8, 8.8, 8.6, 8.6, 8.4, 8.3, 8.4, 8.3, 8.3, 8.1, 8.2, 8.3,
                8.5, 8.1, 8.1, 7.9, 8.3, 8.1, 8.1, 8.1, 8.4, 8.7, 9.0, 9.3, 9.3, 9.5,
                9.3, 9.5, 9.5, 9.5, 9.5, 9.5, 9.5, 9.9, 9.5, 9.7, 9.1, 9.1, 8.9, 9.3,
                9.1, 9.1, 9.3, 9.5, 9.3, 9.3, 9.3, 9.9, 9.7, 9.1, 9.3, 9.5, 9.4, 9.0,
                9.0, 8.8, 9.0, 8.8, 8.6, 8.6, 8.0, 8.0, 8.0, 8.0, 8.6, 8.0, 8.0, 8.0,
                7.6, 8.6, 9.6, 9.6, 10.0, 9.4, 9.3, 9.2, 9.5, 9.5, 9.5, 9.9, 9.9, 9.5,
                9.3, 9.5, 9.5, 9.1, 9.3, 9.5, 9.3, 9.1, 9.3, 9.1, 9.5, 9.4, 9.5, 9.6,
                10.2, 9.8, 9.6, 9.6, 9.4, 9.4, 9.4, 9.4, 9.6, 9.6, 9.4, 9.4, 9.0, 9.4,
                9.4, 9.6, 9.4, 9.2, 8.8, 8.8, 9.2, 9.2, 9.6, 9.6, 9.8, 9.8, 10.0, 10.0,
                9.4, 9.8, 8.8, 8.8, 8.8, 8.8, 9.6, 9.6, 9.6, 9.2, 9.2, 9.0, 9.0, 9.0,
                9.4, 9.0, 9.0, 9.4, 9.4, 9.6, 9.4, 9.6, 9.6, 9.6, 10.0, 10.0, 9.6, 9.2,
                9.2, 9.2, 9.0, 9.0, 9.6, 9.8, 10.2, 10.0, 10.0, 10.0, 9.4, 9.2, 9.6, 9.7,
                9.7, 9.8, 9.8, 9.8, 10.0, 10.0, 8.6, 9.0, 9.4, 9.4, 9.4, 9.4, 9.4, 9.6,
                10.0, 10.0, 9.8, 9.8, 9.7, 9.6, 9.4, 9.2, 9.0, 9.4, 9.6, 9.6, 9.6, 9.6,
                9.6, 9.6, 9.0, 9.4, 9.4, 9.4, 9.6, 9.4, 9.6, 9.6, 9.8, 9.8, 9.8, 9.6,
                9.2, 9.6, 9.2, 9.2, 9.6, 9.6, 9.6, 9.6, 9.6, 9.6, 10.0, 10.0, 10.4, 10.4,
                9.8, 9.0, 9.6, 9.8, 9.6, 8.6, 8.0, 8.0, 8.0, 8.0, 8.4, 8.8, 8.4, 8.4,
                9.0, 9.0, 9.4, 10.0, 10.0, 10.0, 10.2, 10.0, 10.0, 9.6, 9.0, 9.0, 8.6,
                9.0, 9.6, 9.6, 9.0, 9.0, 8.9, 8.8, 8.7, 8.6, 8.3, 7.9, 8.5, 8.7, 8.9, 9.1,
                9.1, 9.1]

const seriesE = SA[101, 82, 66, 35, 31, 7, 20, 92, 154, 125,
                85, 68, 38, 23, 10, 24, 83, 132, 131, 118,
                90, 67, 60, 47, 41, 21, 16, 6, 4, 7,
                14, 34, 45, 43, 48, 42, 28, 10, 8, 2,
                0, 1, 5, 12, 14, 35, 46, 41, 30, 24,
                16, 7, 4, 2, 8, 17, 36, 50, 62, 67,
                71, 48, 28, 8, 13, 57, 122, 138, 103, 86,
                63, 37, 24, 11, 15, 40, 62, 98, 124, 96,
                66, 64, 54, 39, 21, 7, 4, 23, 55, 94,
                96, 77, 59, 44, 47, 30, 16, 7, 37, 74]

const seriesF = SA[47, 64, 23, 71, 38, 64, 55, 41, 59, 48, 71, 35, 57, 40, 58,
                44, 80, 55, 37, 74, 51, 57, 50, 60, 45, 57, 50, 45, 25, 59,
                50, 71, 56, 74, 50, 58, 45, 54, 36, 54, 48, 55, 45, 57, 50,
                62, 44, 64, 43, 52, 38, 59, 55, 41, 53, 49, 34, 35, 54, 45,
                68, 38, 50, 60, 39, 59, 40, 57, 54, 23]

const seriesG = SA[112, 118, 132, 129, 121, 135, 148, 148, 136, 119, 104, 118,
                115, 126, 141, 135, 125, 149, 170, 170, 158, 133, 114, 140,
                145, 150, 178, 163, 172, 178, 199, 199, 184, 162, 146, 166,
                171, 180, 193, 181, 183, 218, 230, 242, 209, 191, 172, 194,
                196, 196, 236, 235, 229, 243, 264, 272, 237, 211, 180, 201,
                204, 188, 235, 227, 234, 264, 302, 293, 259, 229, 203, 229,
                242, 233, 267, 269, 270, 315, 364, 347, 312, 274, 237, 278,
                284, 277, 317, 313, 318, 374, 413, 405, 355, 306, 271, 306,
                315, 301, 356, 348, 355, 422, 465, 467, 404, 347, 305, 336,
                340, 318, 362, 348, 363, 435, 491, 505, 404, 359, 310, 337,
                360, 342, 406, 396, 420, 472, 548, 559, 463, 407, 362, 405,
                417, 391, 419, 461, 472, 535, 622, 606, 508, 461, 390, 432]

# ------------------------------------------------------
#
# Datasets from Introduction Time Series and Forecasting
# by Brockwell and Davis
#
# ------------------------------------------------------

const dowj = SA[110.94, 110.69, 110.43, 110.56, 110.75, 110.84, 110.46, 110.56, 110.46, 110.05,
            109.6, 109.31, 109.31, 109.25, 109.02, 108.54, 108.77, 109.02, 109.44, 109.38,
            109.53, 109.89, 110.56, 110.56, 110.72, 111.23, 111.48, 111.58, 111.9, 112.19,
            112.06, 111.96, 111.68, 111.36, 111.42, 112, 112.22, 112.7, 113.15, 114.36,
            114.65, 115.06, 115.86, 116.4, 116.44, 116.88, 118.07, 118.51, 119.28, 119.79,
            119.7, 119.28, 119.66, 120.14, 120.97, 121.13, 121.55, 121.96, 122.26, 123.79,
            124.11, 124.14, 123.37, 123.02, 122.86, 123.02, 123.11, 123.05, 123.05, 122.83,
            123.18, 122.67, 122.73, 122.86, 122.67, 122.09, 122, 121.23]

const wine = SA[464, 675, 703, 887, 1139, 1077, 1318, 1260, 1120, 963, 996, 960, 530, 883,
            894, 1045, 1199, 1287, 1565, 1577, 1076, 918, 1008, 1063, 544, 635, 804, 980,
            1018, 1064, 1404, 1286, 1104, 999, 996, 1015, 615, 722, 832, 977, 1270, 1437,
            1520, 1708, 1151, 934, 1159, 1209, 699, 830, 996, 1124, 1458, 1270, 1753, 2258,
            1208, 1241, 1265, 1828, 809, 997, 1164, 1205, 1538, 1513, 1378, 2083, 1357, 1536,
            1526, 1376, 779, 1005, 1193, 1522, 1539, 1546, 2116, 2326, 1596, 1356, 1553, 1613,
            814, 1150, 1225, 1691, 1759, 1754, 2100, 2062, 2012, 1897, 1964, 2186, 966, 1549,
            1538, 1612, 2078, 2137, 2907, 2249, 1883, 1739, 1828, 1868, 1138, 1430, 1809, 1763,
            2200, 2067, 2503, 2141, 2103, 1972, 2181, 2344, 970, 1199, 1718, 1683, 2025, 2051,
            2439, 2353, 2230, 1852, 2147, 2286, 1007, 1665, 1642, 1525, 1838, 1892, 2920, 2572,
            2617, 2047]

const lake = SA[10.38, 11.86, 10.97, 10.8, 9.79, 10.39, 10.42, 10.82, 11.4, 11.32, 11.44, 11.68, 11.17,
            10.53, 10.01, 9.91, 9.14, 9.16, 9.55, 9.67, 8.44, 8.24, 9.1, 9.09, 9.35, 8.82, 9.32, 9.01,
            9, 9.8, 9.83, 9.72, 9.89, 10.01, 9.37, 8.69, 8.19, 8.67, 9.55, 8.92, 8.09, 9.37, 10.13,
            10.14, 9.51, 9.24, 8.66, 8.86, 8.05, 7.79, 6.75, 6.75, 7.82, 8.64, 10.58, 9.48, 7.38,
            6.9, 6.94, 6.24, 6.84, 6.85, 6.9, 7.79, 8.18, 7.51, 7.23, 8.42, 9.61, 9.05, 9.26, 9.22,
            9.38, 9.1, 7.95, 8.12, 9.75, 10.85, 10.41, 9.96, 9.61, 8.76, 8.18, 7.21, 7.13, 9.1, 8.25,
            7.91, 6.89, 5.96, 6.8, 7.68, 8.38, 8.52, 9.74, 9.31, 9.89, 9.96, ]


struct AEDTrialParam
    max_ia::Int
    M::Int
    warmup::Int
    nevents::Vector{Vector{Int}}
    delta::Float64
    bounds::Matrix{Float64}
    h::Vector{Tuple{Float64,Bool}}
    utility::Function

    function AEDTrialParam(max_ia::Int, M::Int, warmup::Int, d::Vector{Vector{Int}},
                        delta::Float64, bounds::Matrix{Float64};
                        h=[(0.00025, true), (0.00025, false)], utility::Function=sponsor_utility)
        if M <= warmup
            @error "Total number of iterations must be larger than warmup iterations"
        end

        new(max_ia, M, warmup, d, delta, bounds, h, utility)
    end
end

struct AEDTrial
    B::Int
    s1_model::Model
    s2_models::Dict{String,Model} ## bi-model, Sc only, S only

    s1_prior::NTuple{3,AbstractPrior}
    s2_priors::Dict{String,NTuple{3,AbstractPrior}}

    s1_param::AEDTrialParam
    s2_params::Dict{String,AEDTrialParam}
end

function extract_subset(x::BiPHWeibullData, sel::BitArray{1})
    BiPHWeibullData(x.id[sel], x.entry[sel], x.time[sel], x.status[sel], x.b[sel], x.z[sel])
end

function extract_subset(x::WeibullPHData, sel::BitArray{1})
    WeibullPHData(x.id[sel], x.entry[sel], x.time[sel], x.status[sel], x.z[sel,:])
end

function find_bounds(A::Vector{Int}, fsim::NTuple{3,ForwardSim}, cef::Float64, cef_t2::Float64; h=0.001, ngrid=50, utility=gsd_utility)
    if length(A) == 2
        upper, obj = SA.grid_search([ngrid, ngrid], [0.90, 0.90], [1.0, 1.0],
                                    x -> objective(A, reshape(collect(x), 1, 2), cef, 1.0, fsim[1], fsim[2], fsim[3]; h=h, utility=utility)[1])
        collect(upper), objective(A, reshape(collect(upper), 1, 2), cef, cef_t2, fsim[1], fsim[2], fsim[3]; h=h, utility=utility)
    else
        tmp = quantile(fsim[2].post_pd[:,1,A[1]], 1-cef)
        upper = ones(Float64, 1, 2)
        upper[1, A[1]] = tmp
        upper[1,:], objective(A, upper, cef, 1.0, fsim[1], fsim[2], fsim[3]; h=h, utility=utility)
    end
end

function get_selection(tpf, tpd, ia_data, fsim, fsim0, fsim1, fsim_dbl, fsim0_dbl, fsim1_dbl, cef, cef_t2, h, utility; allow_Sc_only=true)
    B = size(fsim.post_pd, 1)
    upperF, objF = find_bounds([1,2], (fsim, fsim0, fsim1), cef, cef_t2; h=h, utility=utility)
    upperSc, objSc = find_bounds([1], (fsim_dbl, fsim0_dbl, fsim1_dbl), cef, cef_t2; h=h, utility=utility)
    upperS, objS = find_bounds([2], (fsim_dbl, fsim0_dbl, fsim1_dbl), cef, cef_t2; h=h, utility=utility)

    n = length(ia_data.time)
    objFut = objective(Vector{Int}(undef, 0), zeros(Float64, 1, 2), zeros(Float64, B, 1, 2), n.*ones(Int, B, 1), fsim.param; h=h)

    A = [Vector{Int}(undef,0), [1,2], [2], [1]]
    obj = [objFut, objF[1], objS[1]]

    if allow_Sc_only
        obj = vcat(obj, objSc[1])
    end

    upper = [[0.0, 0.0], upperF, upperS, upperSc]
    k = findmax(obj)[2]

    @debug "selection", A[k], obj, upper, k, cef, cef_t2, h, allow_Sc_only

    A[k], obj[k], upper[k], h, allow_Sc_only
end

function select_population(B::Int, pd1, x_bi, ia_data, ia_time, m::Dict{String,Model}, p::Dict{String,NTuple{3,Tp}},
                           tpar::Dict{String,AEDTrialParam}, old_fa_bounds; utility=gsd_utility) where Tp <: AbstractPrior
    tpf = tpar["full"]
    tpd = tpar["double"]

    x0_bi = fit(tpf.M, ia_data; prior=p["full"][2], warmup=tpf.warmup)
    x1_bi = fit(tpf.M, ia_data; prior=p["full"][3], warmup=tpf.warmup)

    fsim = forward_simulate(1, ia_time, tpf.delta, tpf.nevents, ia_data, m["full"], x_bi, x_bi.prior; B=B)
    fsim0 = forward_simulate(1, ia_time, tpf.delta, tpf.nevents, ia_data, m["full"], x0_bi, x_bi.prior; B=B)
    fsim1 = forward_simulate(1, ia_time, tpd.delta, tpf.nevents, ia_data, m["full"], x1_bi, x_bi.prior; B=B)

    prior_dbl = p["double"][1]
    fsim_dbl = forward_simulate(1, ia_time, tpd.delta, tpd.nevents, ia_data, m["double"], x_bi, prior_dbl; B=B)
    fsim0_dbl = forward_simulate(1, ia_time, tpd.delta, tpd.nevents, ia_data, m["double"], x0_bi, prior_dbl; B=B)
    fsim1_dbl = forward_simulate(1, ia_time, tpd.delta, tpd.nevents, ia_data, m["double"], x1_bi, prior_dbl; B=B)

    cef = alpha(fsim0.post_pd, old_fa_bounds)
    cef_t2 = 1-alpha(fsim1.post_pd, old_fa_bounds)

    [get_selection(tpf, tpd, ia_data, fsim, fsim0, fsim1, fsim_dbl, fsim0_dbl, fsim1_dbl, cef, cef_t2, hh[1], utility; allow_Sc_only=hh[2]) for hh in tpf.h]
end

check_bounds(pd, bnds) = pd .> bnds

function fit_model(M, warmup, data, prior, delta)
    x = fit(M, data; warmup=warmup, prior=prior)
    beta = get_outcome(x)
    mean(beta .<= delta, dims=1)[1,:], x
end

function split(m::Model{T,S}) where {T <: DoubleWeibullParam, S <: DoubleWeibullFixedParam}
    Model(m.param.x, m.fparam.x, m.generate_data), Model(m.param.y, m.fparam.y, m.generate_data)
end

function stage2(id, A, obj, fa_bnds, h, res, ia_data, s1_data, s2_params, s2_models, s2_priors)
    @info id, A

    ## Stage 2
    if A == [1,2]
        s2p = s2_params["full"]
        s2_model = s2_models["full"]
        s2_prior = s2_priors["full"][1]
        nev2 = s2p.nevents[1]
        subpop = "full"
    else
        s2p = s2_params["double"]
        nev2 = sum(s2p.nevents[1])
        if A == [1]
            s2_model = split(s2_models["double"])[1]
            s2_prior = s2_priors["double"][1].x
            s1_data = extract_subpop(false, s1_data)
            subpop = "bm-neg"
        else
            s2_model = split(s2_models["double"])[2]
            s2_prior = s2_priors["double"][1].y
            s1_data = extract_subpop(true, s1_data)
            subpop = "bm-pos"
        end
    end

    s2_data = s2_model.generate_data(s2_model.param, s2_model.fparam)

    ## all subjects from IA but with post-IA follow-up
    s1_data = extract_subset(s1_data, s1_data.id .<= maximum(ia_data.id))

    ## first stage 2 subject enrolled after last IA subject
    s2_data.entry .+= maximum(s1_data.entry)

    new_data = append_data(s1_data, s2_data; add_ids=true)
    fa_time, fa_data = select_ia_data(nev2, new_data; n_max=400)

    pd = [NaN, NaN]
    tmp, x2 = fit_model(s2p.M, s2p.warmup, fa_data, s2_prior, s2p.delta)
    pd[A] .= tmp

    reject = check_bounds(pd, fa_bnds)

    push!(res, [id, 2, h[1], h[2], length(fa_data.time), fa_time, subpop, reject[1], reject[2],
                Int(any(reject)), pd[1], pd[2], sum(fa_data.status)])
end

function run_trial(trial::AEDTrial, id::Int)
    s1p = trial.s1_param

    ## Stage 1
    s1_model = trial.s1_model
    s1_data = s1_model.generate_data(s1_model.param, s1_model.fparam)

    ## initial 2-stage GSD bounds
    bounds = trial.s1_param.bounds

    ## stage 1 model fit
    ia_time, ia_data = select_ia_data(s1p.nevents[1], s1_data)
    pd, x1 = fit_model(s1p.M, s1p.warmup, ia_data, trial.s1_prior[1], s1p.delta)
    reject = check_bounds(pd, bounds[1,:])

    ## empty DataFrame to hold results
    res = DataFrame(trial_id=Vector{Int}(undef,0), ia=Vector{Int}(undef,0), h=Vector{Float64}(undef,0), allowSc=Vector{Bool}(undef,0),
                    n=Vector{Int}(undef,0),
                    time=Vector{Float64}(undef,0), subpop=Vector{String}(undef,0), rejectSc=Vector{Bool}(undef,0), rejectS=Vector{Bool}(undef,0),
                    decision=Vector{Int}(undef,0), pdSc=Vector{Float64}(undef,0), pdS=Vector{Float64}(undef,0), nevents=Vector{Int}(undef,0))

    ## add stage 1 results to DataFrame for every "h" value
    for h in trial.s1_param.h
        push!(res, [id, 1, h[1], h[2], length(ia_data.time), ia_time, "full", reject[1], reject[2], Int(any(reject)), pd[1], pd[2], sum(ia_data.status)])
    end

    @info id, reject

    if !any(reject)
        ## select population
        sel_vec = select_population(trial.B, pd, x1, ia_data, ia_time, trial.s2_models, trial.s2_priors,
                                    trial.s2_params, reshape(bounds[2,:], 1, 2); utility=s1p.utility)

        for k in 1:length(sel_vec)
            sv = sel_vec[k]
            A, obj, fa_bnds, h, allow_Sc_only = sv[1], sv[2], sv[3], sv[4], sv[5]

            if length(A) == 0
                ## res.subpop[1] = "fut"
                push!(res, [id, 2, h, allow_Sc_only, length(ia_data.time), ia_time, "fut", false, false, 0, NaN, NaN, sum(ia_data.status)])
            else
                stage2(id, A, obj, fa_bnds, (h, allow_Sc_only), res, ia_data, s1_data, trial.s2_params, trial.s2_models, trial.s2_priors)
            end
        end
    end

    res
end

function summary(x::DataFrame)
    res = by(x, [:h, :trial_id]) do zz
        tmp = zz[end,:]
        [tmp.n tmp.time tmp.nevents tmp.subpop (tmp.rejectSc | tmp.rejectS) (tmp.rejectSc & tmp.rejectS) tmp.rejectSc tmp.rejectS tmp.pdSc tmp.pdS]
    end
    names!(res, [:h, :trial_id, :n, :time, :nevents, :subpop, :any, :both, :Sc, :S, :pdSc, :pdS])
    res
end


#
println("Ae) Test of the CI part for an internally generated neon multiplet without Breit interaction.")
#
settings = AsfSettings(AsfSettings(),selectLevelsCI = true, selectedLevelsCI = [1,2, 4,5, 7,8], breitCI=false,
                                     selectSymmetriesCI = true, selectedSymmetriesCI = [ LevelSymmetry(1//2,Basics.plus),
                                                          LevelSymmetry(1//2,Basics.minus),  LevelSymmetry(5//2,Basics.plus)], jjLS = LSjjSettings(false) )

wa = Atomic.Computation(Atomic.Computation(), name="xx", grid=JAC.Radial.Grid(true), nuclearModel=Nuclear.Model(26.),
                        configs=[Configuration("[Ne] 3s^2 3p^5"), Configuration("[Ne] 3s 3p^6"), Configuration("[Ne] 3s^2 3p^4 3d")],  #,
                                 # Configuration("[Ne] 3s^2 3p^2 3d^3"), Configuration("[Ne] 3s 3p^2 3d^4"), Configuration("[Ne] 3p^2 3d^5")],
                        asfSettings=settings  )

@time wb = perform(wa)


########## Vanilla GBP messages: Factor to variable ##########
function messageFactorVariable(gbp::ContinuousModel)
    @inbounds Threads.@threads for i in gbp.graph.iterateFactor
        for j in gbp.graph.rowptr[i]
            Mrow = gbp.graph.meanIndirect[j]; Vrow = gbp.graph.varianceIndirect[j]; q = gbp.graph.toVariable.nzval[j]
            for k in gbp.graph.rowptr[i]
                if j != k
                    Mrow -= gbp.graph.coefficient[k] * gbp.inference.meanVariableFactor[k]
                    Vrow += gbp.graph.coefficient[k]^2 * gbp.inference.varianceVariableFactor[k]
                end
            end
            gbp.inference.meanFactorVariable[q] = Mrow / gbp.graph.coefficient[j]
            gbp.inference.varianceFactorVariable[q] = Vrow / (gbp.graph.coefficient[j]^2)
        end
    end
end

########## Vanilla GBP means: Factor to variable ##########
function meanFactorVariable(gbp::ContinuousModel)
    @inbounds Threads.@threads for i in gbp.graph.iterateFactor
        for j in gbp.graph.rowptr[i]
            Mrow = gbp.graph.meanIndirect[j]
            for k in gbp.graph.rowptr[i]
                if j != k
                    Mrow -= gbp.graph.coefficient[k] * gbp.inference.meanVariableFactor[k]
                end
            end
            gbp.inference.meanFactorVariable[gbp.graph.toVariable.nzval[j]] = Mrow / gbp.graph.coefficient[j]
        end
    end
end

########## Vanilla GBP variances: Factor to variable ##########
function varianceFactorVariable(gbp::ContinuousModel)
    @inbounds Threads.@threads for i in gbp.graph.iterateFactor
        for j in gbp.graph.rowptr[i]
            Vrow = gbp.graph.varianceIndirect[j]
            for k in gbp.graph.rowptr[i]
                if j != k
                    Vrow += gbp.graph.coefficient[k]^2 * gbp.inference.varianceVariableFactor[k]
                end
            end
            gbp.inference.varianceFactorVariable[gbp.graph.toVariable.nzval[j]] = Vrow / (gbp.graph.coefficient[j]^2)
        end
    end
end

########## Vanilla GBP damp messages: Factor to variable ##########
function messageDampFactorVariable(gbp::ContinuousModel)
    @inbounds Threads.@threads for i in gbp.graph.iterateFactor
        for j in gbp.graph.rowptr[i]
            Mrow = gbp.graph.meanIndirect[j]; Vrow = gbp.graph.varianceIndirect[j]; q = gbp.graph.toVariable.nzval[j]
            for k in gbp.graph.rowptr[i]
                if j != k
                    Mrow -= gbp.graph.coefficient[k] * gbp.inference.meanVariableFactor[k]
                    Vrow += gbp.graph.coefficient[k]^2 * gbp.inference.varianceVariableFactor[k]
                end
            end
            gbp.inference.meanFactorVariable[q] = gbp.graph.alphaNew[j] * Mrow / gbp.graph.coefficient[j] + gbp.graph.alphaOld[j] * gbp.inference.meanFactorVariable[q]
            gbp.inference.varianceFactorVariable[q] = Vrow / (gbp.graph.coefficient[j]^2)
        end
    end
end

########## Vanilla GBP damp means: Factor to variable ##########
function meanDampFactorVariable(gbp::ContinuousModel)
    @inbounds Threads.@threads for i in gbp.graph.iterateFactor
        for j in gbp.graph.rowptr[i]
            Mrow = gbp.graph.meanIndirect[j]; q = gbp.graph.toVariable.nzval[j]
            for k in gbp.graph.rowptr[i]
                if j != k
                    Mrow -= (gbp.graph.coefficient[k] * gbp.inference.meanVariableFactor[k])
                end
            end
            gbp.inference.meanFactorVariable[q] = gbp.graph.alphaNew[j] * Mrow / gbp.graph.coefficient[j] + gbp.graph.alphaOld[j] * gbp.inference.meanFactorVariable[q]
        end
    end
end

############# Vanilla GBP messages: Variable to factor ##########
function messageVariableFactor(gbp::ContinuousModel)
    @inbounds Threads.@threads for i in gbp.graph.iterateVariable
        for j in gbp.graph.colptr[i]
            Mcol = gbp.graph.meanDirect[i]; Wcol = gbp.graph.weightDirect[i]; q = gbp.graph.toFactor.nzval[j]
            for k in gbp.graph.colptr[i]
                if j != k
                    Mcol += gbp.inference.meanFactorVariable[k] / gbp.inference.varianceFactorVariable[k]
                    Wcol += 1 / gbp.inference.varianceFactorVariable[k]
                end
            end
            gbp.inference.varianceVariableFactor[q] = 1 / Wcol
            gbp.inference.meanVariableFactor[q] = Mcol / Wcol
        end
    end
end

############# Vanilla GBP means: Variable to factor ##########
function meanVariableFactor(gbp::ContinuousModel)
    @inbounds Threads.@threads for i in gbp.graph.iterateVariable
        for j in gbp.graph.colptr[i]
            Mcol = gbp.graph.meanDirect[i]; q = gbp.graph.toFactor.nzval[j]
            for k in gbp.graph.colptr[i]
                if j != k
                    Mcol += gbp.inference.meanFactorVariable[k] / gbp.inference.varianceFactorVariable[k]
                end
            end
            gbp.inference.meanVariableFactor[q] = Mcol * gbp.inference.varianceVariableFactor[q]
        end
    end
end

############# Vanilla GBP variances: Variable to factor ##########
function varianceVariableFactor(gbp::ContinuousModel)
    @inbounds Threads.@threads for i in gbp.graph.iterateVariable
        for j in gbp.graph.colptr[i]
            Wcol = gbp.graph.weightDirect[i]
            for k in gbp.graph.colptr[i]
                if j != k
                    Wcol += 1 / gbp.inference.varianceFactorVariable[k]
                end
            end
            gbp.inference.varianceVariableFactor[gbp.graph.toFactor.nzval[j]] = 1 / Wcol
        end
    end
end


struct Coordinate{M,R1,S} <: TaskMap{M,R1,S}
    ind::Int
end
struct CoordinateT{M,R1,S} <: TaskMapT{M,R1,S}
    base_map::Coordinate{M,R1,S}
end
CoordinateT{M,R1,S}(i::Int) where {M,R1,S} = CoordinateT{M,R1,S}(Coordinate{M,R1,S}(i))

task_map_emb(::EmbRep, ::EmbRep, xm, task_map::Coordinate{ℝ{m}, R1}) where m =
    SA[xm[task_map.ind]]


using GLVisualize, GeometryTypes, GLAbstraction, Colors, FileIO

if !isdefined(:runtests)
	window = glscreen()
end
const static_example = true

# loadasset is defined as `loadasset(name) = load(assetphat(name))`
# load comes from FileIO, which automatically detects the file format and loads it
# in general, load will try to load the mesh as is, which may return different
# meshtypes. This is why it's wrapped into GLNormalMesh, to guarantee some type
# stability.
# Formats supported are currently: obj, stl, ply, off and 2DM
msh = GLNormalMesh(loadasset("cat.obj"))
v = vertices(msh)
f = faces(msh)
colors = RGBA{Float32}[RGBA{Float32}(rand(), rand(), rand(), 0.5) for i=1:length(v)]

colored_mesh = GLNormalVertexcolorMesh(
    vertices=v, faces=f,
    color=colors
)

_view(visualize(colored_mesh, is_fully_opaque=false), window, camera=:perspective)
_view(visualize(msh, :lines, thickness=1f0), window, camera=:perspective)

if !isdefined(:runtests)
	@async renderloop(window)
end


using DFTK
using Printf
using HDF5
using Dates
import FFTW
import JSON
using LinearAlgebra
import DFTK: HybridMixing
import DFTK: ScfDiagtol, ScfDefaultCallback

function write_scfres(h5gr, scfres; store_wfn=false)
    for symbol in (:ρin, :ρ)
        hasproperty(scfres, symbol) || continue
        object = getproperty(scfres, symbol)
        h5gr["$(symbol)_real", "shuffle", (), "compress", 3] = object.real
        h5gr["$(symbol)_fourier", "shuffle", (), "compress", 3] = object.fourier
    end

    for symbol in (:εF, )
        hasproperty(scfres, symbol) || continue
        object = getproperty(scfres, symbol)
        object === nothing && continue
        h5gr[string(symbol)] = object
    end

    for symbol in (:eigenvalues, :occupation)
        hasproperty(scfres, symbol) || continue
        h5gr[string(symbol)] = hcat(getproperty(scfres, symbol)...)
    end

    for key in keys(scfres.energies)
        h5gr["energy_$(key)"] = scfres.energies[key]
    end

    if store_wfn
        for symbol in (:ψ, )
            hasproperty(scfres, symbol) || continue
            object = getproperty(scfres, symbol)
            object === nothing && continue

            for ik in 1:length(object)
                h5gr["ψ_$(ik)_fourier", "shuffle", (), "compress", 3] = object[ik]
            end
        end
    end
end

function load_scfres(basis, fileprefix)
    has_results(fileprefix) || error("Need result in $fileprefix")

    h5open(fileprefix * ".hdf5", "r") do h5f
        res = h5f["result"]
        n_kpoints = length(basis.kpoints)
        n_bands = size(res["ψ_1_fourier"], 2)

        # Checks for consistency
        @assert size(res["ρ_real"]) == basis.fft_size
        @assert size(res["occupation"]) == (n_bands, n_kpoints)
        @assert size(res["eigenvalues"]) == (n_bands, n_kpoints)

        # Read state
        εF = read(res, "εF")
        ρ = from_real(basis, res["ρ_real"][:, :, :])
        ψ = Matrix{ComplexF64}[]
        occupation = Vector{Float64}[]
        eigenvalues = Vector{Float64}[]
        for ik = 1:n_kpoints
            n_fft = length(G_vectors(basis.kpoints[ik]))
            @assert size(res["ψ_$(ik)_fourier"]) == (n_fft, n_bands)
            push!(ψ, res["ψ_$(ik)_fourier"][:, :])

            push!(occupation, res["occupation"][:, ik])
            push!(eigenvalues, res["eigenvalues"][:, ik])
        end

        # Construct Hamiltonian and compute energies
        energies, ham = energy_hamiltonian(basis, ψ, occupation;
                                           ρ=ρ, eigenvalues=eigenvalues, εF=εF)

        (ham=ham, basis=basis, energies=energies, converged=true, ρ=ρ,
         eigenvalues=eigenvalues, occupation=occupation, εF=εF, n_iter=nothing, ψ=ψ)
    end
end

function has_results(fileprefix)
    !isfile(fileprefix * ".hdf5") && return false
    try
        return h5open(fileprefix * ".hdf5", "r") do h5f
            exists(h5f, "result")
        end
    catch
    end
    false
end

function run_dftk(lattice, atoms; Ecut=20, tol=1e-11, mixing=DFTK.SimpleMixing(),
                  maxiter=50, kspacing=0.3, smearing=Smearing.None(),
                  temperature=0, fileprefix="dftk",
                  determine_bands=nelec -> nelec // 2 + sqrt(nelec // 2),
                  stepsize=10, randomise_guess=false, kgrid=nothing,
                  supersampling=1.5, ensure_odd_fft_size=false,
                  find_occupation=find_occupation, random_ampl=1e-2,
                  diagtol_max=1e-3,
                 )
    model = model_DFT(lattice, atoms, [:gga_x_pbe, :gga_c_pbe],
                      smearing=smearing, temperature=temperature)
    isnothing(kgrid) && (kgrid = kgrid_size_from_minimal_spacing(lattice, kspacing))

    fft_size = determine_grid_size(model, Ecut, supersampling=supersampling)
    ensure_odd_fft_size && (fft_size = Vec3(fft_size .+ iseven.(fft_size)))
    basis = PlaneWaveBasis(model, Ecut, kgrid=kgrid, kshift=[0, 0, 0], fft_size=fft_size)
    n_bands = ceil(Int, determine_bands(model.n_electrons))

    if has_results(fileprefix)
        @info "Loading existing results instead of SCF"
        return load_scfres(basis, fileprefix)
    end
    n_blas = Int(ccall((BLAS.@blasfunc(openblas_get_num_threads), BLAS.libblas), Cint, ()))

    println()
    println("hostname       = $(read(`hostname`, String))")
    println("started on     = $(Dates.now())")
    println("julia threads  = $(Threads.nthreads())")
    println("BLAS threads   = $(n_blas)")
    println()
    println("temperature    = $(basis.model.temperature)")
    println("smearing       = $(basis.model.smearing)")
    if isdiag(lattice)
        println("diag(lattice)  = $(round.(diag(lattice), sigdigits=4))")
    else
        println("lattice        = $(round.(lattice, sigdigits=4))")
    end
    println("Ecut           = $Ecut")
    println("fft_size       = $(basis.fft_size)")
    println("kgrid          = $kgrid")
    println("irreducible k  = $(length(basis.kpoints))")
    println("n_bands        = $(n_bands)")
    println("n_electr       = $(basis.model.n_electrons)")
    println("mixing         = $mixing")
    flush(stdout)

    # Collection of data to dump later
    save_dict = Dict(
        "temperature"  => basis.model.temperature,
        "smearing"     => string(basis.model.smearing),
        "diag_lattice" => diag(lattice),
        "lattice"      => lattice,
        "Ecut"         => Ecut,
        "fft_size"     => basis.fft_size,
        "kgrid"        => kgrid,
        "n_kirred"     => length(basis.kpoints),
        "n_bands"      => n_bands,
        "n_electr"     => basis.model.n_electrons,
        "mixing"       => string(mixing),
        "α"            => mixing.α,
        "energies"     => Float64[],
    )

    # Empty the existing HDF5 file
    isdir(dirname(fileprefix)) || mkdir(dirname(fileprefix))
    close(h5open(fileprefix * ".hdf5", "w"))

    default_callback = ScfDefaultCallback()
    function callback(info)
        default_callback(info)
        flush(stdout)
        key = info.stage == :finalize ? "result" : "neval_$(info.n_iter)"
        store_wfn = info.stage == :finalize
        h5open(fileprefix * ".hdf5", "r+") do h5f
            write_scfres(g_create(h5f, key), info, store_wfn=store_wfn)
        end
        push!(save_dict["energies"], info.energies.total)
    end

    ρ0 = guess_density(basis)
    if randomise_guess
        # Add randomness and adjust electron count
        ρ0.real .= abs.(ρ0.real .+ random_ampl * rand(size(ρ0.real)))
        ρ0.real .*= basis.model.n_electrons ./ (sum(ρ0.real) ./ length(ρ0.real) .* basis.model.unit_cell_volume)
    end
    scfres = self_consistent_field(basis, ρ=ρ0, tol=tol, mixing=mixing,
                                   callback=callback,
                                   determine_diagtol=ScfDiagtol(ratio_ρdiff=0.04,
                                                                diagtol_max=diagtol_max,
                                                                diagtol_min=1e-12),
                                   solver=scf_nlsolve_solver(stepsize),
                                   maxiter=maxiter, n_bands=n_bands,
                                   find_occupation=find_occupation,
                                   enforce_symmetry=true)

    println(scfres.energies)
    flush(stdout)

    open(fp -> JSON.print(fp, save_dict), fileprefix * ".json", "w")
    scfres
end