export inferType, inferModule, inferExpr, inferDecls

function inferType(globalTC::GlobalTC, localTC::LocalTC, exp::Surf.TyExpr)::HMT
    ln = localTC.ln
    typetype::HMT = @switch exp begin
    @case Surf.TNew(tn)
        T(Nom(Symbol(tn, '@', ln.line, "_", symgen(globalTC))))
    @case Surf.TImplicit(t)
        T(Implicit(inferType(globalTC, localTC, t)))
    @case Surf.TQuery(label, exp)
        let ret = T(inferType(globalTC, localTC, exp))
            push!(globalTC.queries, label => ret)
            ret
        end

    @case Surf.TSym(a)
        T(Nom(a))

    @case Surf.TVar(:_)
        T(globalTC.tcstate.new_tvar())

    @case Surf.TVar(a)
        t = get(localTC.typeEnv, a, nothing)
        t === nothing && throw(MLError(
            localTC.ln,
            UnboundTypeVar(a)))
        t::HMT

    @case Surf.TApp(f, b)
        f = inferType(globalTC, localTC, f)
        b = inferType(globalTC, localTC, b)
        T(App(f, b))

    @case Surf.TTuple(args)
        T(Tup(HMT[inferType(globalTC, localTC, arg) for arg in args]))

    @case Surf.TForall(tvars, p)
        let typeEnv = localTC.typeEnv

            uniqueNames = UN[UN(k) for k in tvars]
            typeEnv = typeEnv[[un.name => T(Bound(un)) for un in uniqueNames]]
            localTC = @set localTC.typeEnv = typeEnv
            p = inferType(globalTC, localTC, p)
            T(Forall(uniqueNames, p))
        end

    @case Surf.TArrow(a, b)
        a = inferType(globalTC, localTC, a)
        b = inferType(globalTC, localTC, b)
        T(Arrow(a, b))
    end

    tcstate = globalTC.tcstate
    unify = tcstate.unify
    new_tvar = tcstate.new_tvar
    prune = tcstate.prune

    @match typetype begin
        T(x) => x
        _ => begin
            targ = new_tvar()
            unify(T(targ), typetype) || throw(MLError(ln, UnificationFail))
            prune(targ)
        end
    end
end

const _InferPostponed = CFunc{IR.Decl, Tuple{}}

function inferModule(globalTC::GlobalTC, localTC::LocalTC, decls::Vector{Surf.Decl})
    props, localTC = inferDecls(globalTC, localTC, decls, Val(true))
    IR.Decl[f() for f in props], localTC
end

function inferDecls(globalTC::GlobalTC, localTC::LocalTC, decls::Vector{Surf.Decl})
    inferDecls(globalTC, localTC, decls, Val(false))
end

function inferDecls(globalTC::GlobalTC, localTC::LocalTC, decls::Vector{Surf.Decl}, ::Val{isGlobal}) where isGlobal
    annotated :: Dict{Symbol, Vector{HM.UN}} = Dict{Symbol, Tuple}()
    loweredDeclFs = _InferPostponed[]
    @inline addDecl!(f::Function) = push!(loweredDeclFs, _InferPostponed(f))

    for decl in decls
        @switch decl begin
        @case Surf.DQuery(label, sym)
            annTy = if sym in localTC.typeEnv
                globalTC.tcstate.prune(localTC.typeEnv[sym])
            else
                throw(MLError(localTC.ln, UnboundVar(sym)))
            end
            push!(globalTC.queries, label=>annTy)
            nothing
        @case Surf.DLoc(ln)
            localTC = @set localTC.ln = ln
            nothing
        @case Surf.DOpen(exp)
            let exprTyProp = inferExpr(globalTC, localTC, exp),
                ex = exprTyProp(NoProp, localTC.localImplicits),
                tcstate = globalTC.tcstate,
                namespaceTy = tcstate.new_tvar(),
                target = App(Nom(:namespace), Tup(HMT[ex.ty, namespaceTy])),
                (fromTy, _) = instanceResolveAndType(
                    globalTC, localTC.localImplicits, target, localTC.ln)

                tcstate.unifyImplicits!(target, fromTy, HMT[]) || throw(MLError(ln, UnificationFail))

                namespaceTy = tcstate.prune(namespaceTy)
                @match namespaceTy begin
                    Tup(fields) =>
                        for field in fields
                            @match field begin
                                Tup([Nom(name), Nom(gensym), t]) =>
                                    begin
                                        localTC = @set localTC.symmap = localTC.symmap[name => gensym]
                                        localTC = @set localTC.typeEnv = localTC.typeEnv[name => t]
                                    end
                                _ => error(MLError(localTC.ln, InvalidNamespaceType(namespaceTy)))
                            end
                        end
                    _ => error(MLError(localTC.ln, InvalidNamespaceType(namespaceTy)))
                end
                IR.Perform(ex)
            end

        @case Surf.DAnn(sym, tyExpr)
            haskey(annotated, sym) && throw(MLError(localTC.ln, UnusedAnnotation(sym)))
            gensym = Symbol(symgen(globalTC), :_, sym)
            localTC = @set localTC.symmap = localTC.symmap[sym => gensym]
            annTy = inferType(globalTC, localTC, tyExpr)
            localTC = @set localTC.typeEnv = localTC.typeEnv[sym => annTy]
            @match annTy begin
                Forall(ns, _) => begin
                        annotated[sym] = ns
                    end
                Implicit(t) =>
                    if isGlobal
                        t = globalTC.tcstate.prune(t)
                        tConsHead = getTConsHead(t)
                        globalImplicits = get!(globalTC.globalImplicits, tConsHead) do
                            InstRec[]
                        end
                        push!(globalImplicits, InstRec(t, localTC.ln, gensym, false))
                        globalTC.globalImplicitDeltas[tConsHead] =
                            get(globalTC.globalImplicitDeltas, tConsHead, 0) + 1

                        annotated[sym] = HM.UN[]
                    else
                        inst = InstRec(t, localTC.ln, gensym, false)
                        localImplicits = cons(inst, localTC.localImplicits)
                        localTC = @set localTC.localImplicits = localImplicits
                        annotated[sym] = HM.UN[]
                    end
                _ => begin annotated[sym] = HM.UN[] end
            end
            nothing
        @case Surf.DBind(sym, expr)
            if sym === :_

                let exprTyProp = inferExpr(globalTC, localTC, expr),
                    localImplicits = localTC.localImplicits

                    addDecl!(() -> IR.Perform(exprTyProp(NoProp, localImplicits)))
                end
                continue
            end
            (gensym, annTy, typevars) = if haskey(annotated, sym)
                typevars = pop!(annotated, sym)
                localTC.symmap[sym], globalTC.tcstate.prune(localTC.typeEnv[sym]), typevars
            else
                gensym = Symbol(symgen(globalTC), :_, sym)
                localTC = @set localTC.symmap = localTC.symmap[sym => gensym]
                tvar = globalTC.tcstate.new_tvar()
                localTC = @set localTC.typeEnv = localTC.typeEnv[sym => tvar]
                gensym, tvar, ()
            end
            # @info :annotation sym annTy
            let localTC = if isempty(typevars)
                            localTC
                        else
                           @set localTC.typeEnv =
                                localTC.typeEnv[[un.name => T(Bound(un)) for un in typevars]]
                        end,
                exprTyProp = inferExpr(globalTC, localTC, expr),
                gensym = gensym,
                annTy = annTy,
                prune = globalTC.tcstate.prune,
                localImplicits = localTC.localImplicits

                addDecl!(
                    () ->
                    IR.Assign(
                        gensym,
                        annTy,
                        exprTyProp(InstTo(annTy), localImplicits)))
            end
        end
    end
    loweredDeclFs, localTC
end

intTypes = Dict{Int, Nom}()
floatTypes = Dict{Int, Nom}()
for bit in (8, 16, 32, 64)
    intTypes[bit] = Nom(Symbol(:i, bit))
    floatTypes[bit] = Nom(Symbol(:f, bit))
end
strT = Nom(:str)
charT = Nom(:char)
boolT = Nom(:bool)

function inferExpr(globalTC::GlobalTC, localTC::LocalTC, expr::Surf.Expr)
    tcstate = globalTC.tcstate
    unifyINST = tcstate.unifyINST
    unifyImplicits! = tcstate.unifyImplicits!
    unify = tcstate.unify
    prune = tcstate.prune
    new_tvar = tcstate.new_tvar
    localImplicits = localTC.localImplicits
    ln = localTC.ln

    @inline function applyTI(ti::TypeInfo, me::HMT)
        local implicits
        implicits = HMT[]
        @match ti begin
            NoProp => (me, implicits)
            InstTo(instTy) => begin
                unifyImplicits!(instTy, me, implicits) || begin
                    # @info :UNIFY_FAIL prune(instTy) prune(me)
                    throw(MLError(ln, UnificationFailOf(prune(instTy), prune(me))))
                end
                (instTy, implicits)
            end
        end
    end

    @inline function prop(varTy::HMT, eImpl::IR.ExprImpl)
        function propInner(ti::TypeInfo, localImplicits::InstResolCtx)
            local implicits
            varTy = prune(varTy)
            varTy, implicits = applyTI(ti, varTy)
            IR.applyImplicits(eImpl, implicits, varTy, ln, localImplicits)
        end
    end

    @switch expr begin
    @case Surf.EExt(jlex)

        function (ti::TypeInfo, localImplicits::InstResolCtx)
            local t, implicits
            t = new_tvar()
            t, implicits = applyTI(ti, t)
            IR.applyImplicits(IR.EExt(jlex), implicits, t, ln, localImplicits)
        end

    @case Surf.EQuery(label, e)
        exprTyProp = inferExpr(globalTC, localTC, e)
        queries = globalTC.queries
        function propQuery(ti::TypeInfo, localImplicits::InstResolCtx)
            e = exprTyProp(ti, localImplicits)
            push!(queries, label=>e.ty)
            e
        end
    @case Surf.ELoc(ln, expr)
        inferExpr(globalTC, (@set localTC.ln = ln), expr)

    @case Surf.EField(subject, field)

        exprTyProp = inferExpr(globalTC, localTC, subject)
        function propField(ti::TypeInfo, localImplicits::InstResolCtx)
            local eSub, retTV, retT, eRet, op_Getter
            eSub = exprTyProp(NoProp, localImplicits)
            retTV = new_tvar()
            target = App(Nom(:field), Tup(HMT[eSub.ty, Nom(field), retTV]))
            retT, eAccess = instanceResolveAndType(globalTC, localImplicits, target, ln)
            unifyImplicits!(target, retT, HMT[]) || throw(MLError(ln, UnificationFail))
            retTV, implicits = applyTI(ti, retTV)

            op_Getter = IR.Expr(ln, nothing, IR.EVar(:op_Getter))
            eRet = IR.EApp(IR.Expr(ln, nothing, IR.EApp(op_Getter, eAccess)), eSub)
            IR.applyImplicits(eRet, implicits, retTV, ln, localImplicits)
        end

    @case Surf.ETup(xs)

        let elty = CFunc{IR.Expr, Tuple{TypeInfo, InstResolCtx}},
            props = elty[elty(inferExpr(globalTC, localTC, x)) for x in xs],
            n_xs = length(xs)

            function propTup(ti::TypeInfo, localImplicits::InstResolCtx)
                local ts, tupT, elts, implicits
                ts = HMT[new_tvar() for i = 1:n_xs]
                tupT = Tup(ts)
                tupT, implicits = applyTI(ti, tupT)
                elts = IR.Expr[(prop(InstTo(t), localImplicits)) for (prop, t) in zip(props, ts)]
                IR.applyImplicits(IR.ETup(elts), implicits, tupT, ln, localImplicits)
            end
        end
    @case Surf.EApp(f, arg)
        let fProp = inferExpr(globalTC, localTC, f),
            argProp = inferExpr(globalTC, localTC, arg)

            function propApp(ti::TypeInfo, localImplicits::InstResolCtx)
                local argT, retT, eF, eArg, implicits
                argT = new_tvar()
                retT = new_tvar()
                retT, implicits = applyTI(ti, retT)
                eF = fProp(InstTo(Arrow(argT, retT)), localImplicits)
                eArg = argProp(InstTo(argT), localImplicits)
                IR.applyImplicits(IR.EApp(eF, eArg), implicits, retT, ln, localImplicits)
            end
        end

    @case Surf.EFun(n, expr)
        gensym = Symbol(symgen(globalTC), :_, n)
        localTC = @set localTC.symmap = localTC.symmap[n => gensym]
        argT = new_tvar()
        localTC = @set localTC.typeEnv = localTC.typeEnv[n => argT]
        retT = new_tvar()
        let exprTyProp = inferExpr(globalTC, localTC, expr),
            gensym = gensym

            function propFun(ti::TypeInfo, localImplicits::InstResolCtx)
                local arrowT, eBody, implicits
                arrowT, implicits = applyTI(ti, Arrow(argT, retT))
                localImplicits = @match prune(argT) begin
                    Implicit(instance) =>
                        cons(InstRec(instance, ln, gensym, false), localImplicits)
                    _ => localImplicits
                end
                eBody = exprTyProp(InstTo(retT), localImplicits)
                IR.applyImplicits(
                    IR.EFun(gensym, eBody),
                    implicits, arrowT, ln, localImplicits)
            end
        end

    @case Surf.ELet(decls, expr)
        let (loweredDeclFs, localTC) = inferDecls(globalTC, localTC, decls),
            exprTyProp = inferExpr(globalTC, localTC, expr)

            function propLet(t::TypeInfo, localImplicits::InstResolCtx)
                local eBody, loweredDecls, eI
                eBody = exprTyProp(t, localImplicits)
                loweredDecls = [f() for f in loweredDeclFs]
                eI = IR.ELet(loweredDecls, eBody)
                IR.Expr(ln, eBody.ty, eI)
            end
        end

    @case Surf.EITE(cond, arm1, arm2)
        let condProp = inferExpr(globalTC, localTC, cond),
            arm1Prop = inferExpr(globalTC, localTC, arm1),
            arm2Prop = inferExpr(globalTC, localTC, arm2),
            new_tvar = globalTC.tcstate.new_tvar

            function propITE(ti::TypeInfo, localImplicits::InstResolCtx)
                local t, eArm1, eArm2, eCond, eI, implicits
                t, implicits = applyTI(ti, new_tvar())
                ti = InstTo(t)
                eArm1 = arm1Prop(ti, localImplicits)
                eArm2 = arm2Prop(ti, localImplicits)
                eCond = condProp(InstTo(boolT), localImplicits)
                eI = IR.EITE(eCond, eArm1, eArm2)
                IR.applyImplicits(eI, implicits, t, ln, localImplicits)
            end
        end

    @case Surf.EVal(val)
        (eValI, eT) = @switch val begin
            @case ::Integer
                bit = 8 * sizeof(val)
                IR.EInt(val, bit), intTypes[bit]
            @case ::AbstractFloat
                bit = 8 * sizeof(val)
                IR.EFloat(val, bit), floatTypes[bit]
            @case ::String
                IR.EStr(val), strT
            @case ::Char
                IR.EChar(val), charT
            @case ::Bool
                IR.EBool(val), boolT
        end
        prop(eT, eValI)

    @case Surf.EVar(n)
        n in localTC.typeEnv || throw(MLError(localTC.ln, UnboundVar(n)))
        varT = prune(localTC.typeEnv[n])
        @match varT begin
            T(t) =>
                function propTVar(ti::TypeInfo, _::InstResolCtx)
                    @match ti begin
                        InstTo(t′) =>
                            begin
                                unify(t′, varT) || throw(MLError(ln, UnificationFail))
                            end
                        _ => nothing
                    end

                    IR.Expr(ln, varT, IR.ETypeVal(t))
                end
            _ =>
                let gensym = localTC.symmap[n]
                    prop(varT, IR.EVar(gensym))
                end
        end
    end
end


# Run model fitting and crossvalidation
# for real or simulated data

# Main runner for model fitting and model comparison
# NOTE: These procedures take a long time to run.
# The settings are different from the ones used in the paper (DEBUG settings).
# Comment/Uncomment the corresponding lines (around 383 and 445)
# to perform the fitting and comparison procedures as in the paper.
function runner_crossval_multipletimes_fit()

    # -------- Fit (once)
    i=""
    result_path_base, batch_fit_models_spec = model_list_main(filestr = string(i))
    runner_fmri_fit(result_path_base, batch_fit_models_spec)

    # -------- Crossvalidation (multiple runs)
    nruns = 5
    for i in 1:nruns
        result_path_base, batch_fit_models_spec = model_list_main(filestr = string(i))
        runner_fmri_3foldscrossval(result_path_base,
                                        batch_fit_models_spec)
    end
end

# Main runner for model and parameter recovery
# NOTE: You need to EDIT which data (real or simulated)
# should be used in
# runner_fmri_fit and runner_fmri_3foldscrossval!!!!
# Right now the real participants' data are used.
# See lines 367 and 428
function runner_multipletimes_fit_recovery()

    # -------- Fit for parameter recovery
    # ---- We run fit multiple times to get some error bars
    nruns = 5
    for i in 1:nruns
        result_path_base, batch_fit_models_spec = model_list_paramrecovery(filestr = string(i))
        runner_fmri_fit(result_path_base, batch_fit_models_spec)

    end

    # -------- Crossvalidation (multiple runs) for model recovery
    nruns = 5
    for i in 1:nruns
        result_path_base, batch_fit_models_spec = model_list_main(filestr = string(i))
        runner_fmri_3foldscrossval(result_path_base,
                                        batch_fit_models_spec)
    end
end

# ------------------------------------------------------------------------------
# ------------------------------------------------------------------------------
# ------- Model list -----------------------------------------------------------
# ------------------------------------------------------------------------------
# List and specifications of algorithms for Fig. 3 of paper
function model_list_main(; filestr = "_")

    result_path_base = makehomesavepath("test_main"*filestr)

    stochmax = 20.

    # # Biased Random Walk
    # --------------------
    biasedRandomWalk_prior = [
        RL_Fit.paramBias(0.0, 1.)
    ]
    biasedRandomWalk_prior_std = [.15]
    biasedRandomWalk_spec = RL_Fit.Model_spec_per_subj("Biased Random Walk (0-Model)",
                                RL_Fit.getLL_BiasedRandomWalk_V,
                                biasedRandomWalk_prior, biasedRandomWalk_prior_std,
                                string(result_path_base, "biased_random_walk/"))
    # Sarsa-λ
    # -------
    sarsaL_prior = [
        RL_Fit.paramAlpha(0.0, 1.0),
        RL_Fit.paramGamma(0.0, 1.0),
        RL_Fit.paramLambda(0.0, 1.0),
        RL_Fit.paramTemp(0.01, 10.)
        ]
    sarsaL_prior_std = [.15, .15, .15, 1.5]
    sarsaL_spec = RL_Fit.Model_spec_per_subj("Sarsa Lambda",
                    RL_Fit.getLL_SarsaLambda_V,
                    sarsaL_prior, sarsaL_prior_std,
                    string(result_path_base, "sarsa_lambda/"))
    # Prioritized Sweeping with Particle Filtering
    # --------------------------------------------
    priosweep_particle_prior = [
        RL_Fit.paramGamma(0.0, 1.0),
        RL_Fit.paramTemp(0.01, 10.),
        RL_Fit.paramSurpriseProbability(0.0, 1.),
        RL_Fit.paramStochasticity(0.00001, stochmax),
        RL_Fit.paramNrUpdateCycles(0, 9),
        RL_Fit.paramNParticles(1, 20)
        ]
    priosweep_particle_prior_std = [.15, 1.5, .15, 3., 1.35, 3.]
    priosweep_particle_spec = RL_Fit.Model_spec_per_subj("Prioritized Sweeping Particle Filter",
                                RL_Fit.getLL_Prio_Sweep_TParticleFilterJump_V,
                                priosweep_particle_prior,
                                priosweep_particle_prior_std,
                                string(result_path_base, "priosweep_particle/"))
    # Hybrid-λ-PS-PF
    # --------------
    hybrid_pspf_sarsaL_prior = [
        RL_Fit.paramGamma(0.0, 1.0),
        RL_Fit.paramSurpriseProbability(0.0, 1.),
        RL_Fit.paramStochasticity(0.00001, stochmax),
        RL_Fit.paramNrUpdateCycles(0, 9),
        RL_Fit.paramNParticles(1, 20),
        RL_Fit.paramAlpha(0.0, 1.),
        RL_Fit.paramLambda(0.0, 1.0),
        RL_Fit.paramWeight(0.0, 20.0; name="w_MB"),
        RL_Fit.paramWeight(0.0, 20.0; name="w_MF")
        ]
    hybrid_pspf_sarsaL_prior_std = [.15, .15, 3., 1.35, 3., .15, .15, 3., 3.]
    hybrid_pspf_sarsaL_spec = RL_Fit.Model_spec_per_subj("Hybrid: PS with Particle Filtering + Sarsa Lambda",
                                RL_Fit.getLL_Hybrid_Prio_Sweep_TParticleFilterJump_V,
                                hybrid_pspf_sarsaL_prior, hybrid_pspf_sarsaL_prior_std,
                                string(result_path_base, "hybrid_pspf_sarsaL/"))
    # Actor-Critic
    # ------------
    actorCriticTD_prior = [
        RL_Fit.paramAlpha(0.0, 1.; name="learning rate actor"),
        RL_Fit.paramEta(0.0, 1.; name="learning rate critic"),
        RL_Fit.paramGamma(0.0, 1.),
        RL_Fit.paramLambda(0.0, 1.; name="elig. trace actor"), # actor
        RL_Fit.paramLambda(0.0, 1.; name="elig. trace critic"), # critic
        RL_Fit.paramTemp(0.01, 10.)
    ]
    actorCriticTD_prior_std = [.15, .15, .15, .15, .15, 1.5]
    actorCriticTD_spec = RL_Fit.Model_spec_per_subj("Actor-Critic (TD-Lambda)",
                                RL_Fit.getLL_ActorCritic_TDLambda_V,
                                actorCriticTD_prior, actorCriticTD_prior_std,
                                string(result_path_base, "actor_critic_TD/"))
    # REINFORCE
    # ---------
    REINFORCE_prior = [
        RL_Fit.paramAlpha(0.0, 1.),
        RL_Fit.paramGamma(0.0, 1.),
        RL_Fit.paramTemp(0.01, 10.)
    ]
    REINFORCE_prior_std = [.15, .15, 1.5]
    REINFORCE_spec = RL_Fit.Model_spec_per_subj("REINFORCE (Williams 1992)",
                            RL_Fit.getLL_REINFORCE_V,
                            REINFORCE_prior, REINFORCE_prior_std,
                            string(result_path_base, "REINFORCE/"))
    # Hybrid Actor-Critic
    # -------------------
    hybrid_actorCritic_FWD_prior = [
        RL_Fit.paramAlpha(0.0, 1.; name="learning rate actor"),
        RL_Fit.paramAlpha(0.0, 1.; name="learning rate critic"),
        RL_Fit.paramGamma(0.0, 1.),
        RL_Fit.paramLambda(0.0, 1.; name="elig. trace actor"), # actor
        RL_Fit.paramLambda(0.0, 1.; name="elig. trace critic"), # critic
        RL_Fit.paramEta(0., 1.; name="learning rate fwd"),
        RL_Fit.paramWeight(0.0, 20.0; name="w_MB"),
        RL_Fit.paramWeight(0.0, 20.0; name="w_MF")
    ]
    hybrid_actorCritic_FWD_prior_std = [.15, .15, .15, .15, .15, .15, 3., 3.]
    hybrid_actorCritic_FWD_spec = RL_Fit.Model_spec_per_subj("Hybrid Actor-Critic with FWD",
                                RL_Fit.getLL_Hybrid_AC_FWD_V,
                                hybrid_actorCritic_FWD_prior, hybrid_actorCritic_FWD_prior_std,
                                string(result_path_base, "hybrid_actorcritic_FWD/"))

    # Forward Learner (MB only learner)
    # ---------------------------------
    forwardLearner_prior = [
        RL_Fit.paramEta(0., 1.),
        RL_Fit.paramGamma(0., 1.),
        RL_Fit.paramTemp(0.01, 10.)
    ]
    forwardLearner_prior_std = [.15, .15, 1.5]
    forwardLearner_spec = RL_Fit.Model_spec_per_subj("Forward Learner",
                                RL_Fit.getLL_Forward_Learner_V,
                                forwardLearner_prior, forwardLearner_prior_std,
                                string(result_path_base, "forward_learner/"))
    # Hybrid-0 (Glaescher et al. 2010)
    # --------------------------------
    glaescherHybrid_prior = [
        RL_Fit.paramAlpha(0., 1.),
        RL_Fit.paramGamma(0., 1.),
        RL_Fit.paramEta(0., 1.),
        RL_Fit.paramExpDecayOffset(0., 1.),
        RL_Fit.paramExpDecaySlope(-.1, +.8),
        RL_Fit.paramTemp(0.01, 10.)
    ]
    glaescherHybrid_prior_std = [.15, .15, .15, .15, .15, 1.5]
    glaescherHybrid_spec = RL_Fit.Model_spec_per_subj("Glaescher-Hybrid",
                                RL_Fit.getLL_GlaescherHybrid_V,
                                glaescherHybrid_prior, glaescherHybrid_prior_std,
                                string(result_path_base, "glaescher_hybrid/"))
    # Hybrid-λ (Daw et al. 2011)
    # --------------------------
    hybridSarsaL_prior = [
        RL_Fit.paramAlpha(0.0, 1.),
        RL_Fit.paramGamma(0., 1.),
        RL_Fit.paramLambda(0.0, 1.),
        RL_Fit.paramEta(0., 1.),
        RL_Fit.paramExpDecayOffset(0., 1.),
        RL_Fit.paramExpDecaySlope(0., 1.),
        RL_Fit.paramTemp(0.01, 10.)
    ]
    hybridSarsaL_prior_std = [.15, .15, .15, .15, .15, .15, 1.5]
    hybridSarsaL_prior_spec = RL_Fit.Model_spec_per_subj("Hybrid-SarsaLambda",
                                RL_Fit.getLL_HybridSarsaLambda_V,
                                hybridSarsaL_prior, hybridSarsaL_prior_std,
                                string(result_path_base, "hybrid_FWD_SarsaLambda/"))

    # REINFORCE with baseline
    # -----------------------
    REINFORCE_baseline_prior = [
        RL_Fit.paramAlpha(0.0, 1., name="learning rate"),
        RL_Fit.paramGamma(0.0, 1.),
        RL_Fit.paramTemp(0.01, 10.),
        RL_Fit.paramAlpha(0.0, 1.; name="learning rate baseline")
    ]
    REINFORCE_baseline_prior_std = [.15, .15, 1.5, .15]
    REINFORCE_baseline_spec = RL_Fit.Model_spec_per_subj("REINFORCE (Williams 1992) baseline",
                            RL_Fit.getLL_REINFORCE_baseline_V,
                            REINFORCE_baseline_prior, REINFORCE_baseline_prior_std,
                            string(result_path_base, "REINFORCE_baseline/"))
    # Surprise Sarsa-λ
    # ----------------
    sarsaL_modulated_prior = [
        RL_Fit.paramGamma(0.0, 1.0),
        RL_Fit.paramTemp(0.01, 10.),
        RL_Fit.paramSurpriseProbability(0.0, 1.),
        RL_Fit.paramStochasticity(0.00001, stochmax),
        RL_Fit.paramNParticles(1, 20),
        RL_Fit.paramAlpha(0.0, 1.0),
        RL_Fit.paramLambda(0.0, 1.0)
        ]
    sarsaL_modulated_prior_std = [.15, 1.5, .15, 3., 3., .15, .15]
    sarsaL_modulated_spec = RL_Fit.Model_spec_per_subj("Sarsa Lambda modulated by Particle Filtering",
                                RL_Fit.getLL_sarsaL_modulated_TParticleFilterJump_V,
                                sarsaL_modulated_prior,
                                sarsaL_modulated_prior_std,
                                string(result_path_base, "sarsaL_modulated/"))

    # Surprise Actor-Critic
    # ---------------------
    actorCriticTD_modulatedActor_prior = [
        RL_Fit.paramAlpha(0.0, 1.; name="learning rate actor"),
        RL_Fit.paramSurpriseProbability(0.0, 1.),
        RL_Fit.paramStochasticity(0.00001, stochmax),
        RL_Fit.paramEta(0.0, 1.; name="learning rate critic"),
        RL_Fit.paramGamma(0.0, 1.),
        RL_Fit.paramLambda(0.0, 1.; name="elig. trace actor"), # actor
        RL_Fit.paramLambda(0.0, 1.; name="elig. trace critic"), # critic
        RL_Fit.paramTemp(0.01, 10.),
        RL_Fit.paramNParticles(1, 20)
    ]
    actorCriticTD_modulatedActor_prior_std = [.15, .15, 3., .15, .15, .15, .15, 1.5, 3.]
    actorCriticTD_modulatedActor_spec = RL_Fit.Model_spec_per_subj("Actor-Critic (TD-Lambda) modulated actor (1 degree)",
                                RL_Fit.getLL_ActorCritic_TDLambda_modulatedActor_V,
                                actorCriticTD_modulatedActor_prior, actorCriticTD_modulatedActor_prior_std,
                                string(result_path_base, "actor_critic_TD_modulatedActor_1degree/"))
    # Surprise REINFORCE
    # ------------------
    REINFORCE_modulated_prior = [
        RL_Fit.paramAlpha(0.0, 1.),
        RL_Fit.paramSurpriseProbability(0.0, 1.),
        RL_Fit.paramStochasticity(0.00001, stochmax),
        RL_Fit.paramGamma(0.0, 1.),
        RL_Fit.paramTemp(0.01, 10.),
        RL_Fit.paramNParticles(1, 20)
    ]
    REINFORCE_modulated_prior_std = [.15, .15, 3., .15, 1.35, 3.]
    REINFORCE_modulated_spec = RL_Fit.Model_spec_per_subj("REINFORCE (Williams 1992) modulated",
                            RL_Fit.getLL_REINFORCE_modulated_V,
                            REINFORCE_modulated_prior, REINFORCE_modulated_prior_std,
                            string(result_path_base, "REINFORCE_modulated/"))

    # Hybrid Actor-Critic PF
    # ----------------------
    hybrid_actorCritic_FWD_BF_prior = [
        RL_Fit.paramAlpha(0.0, 1.; name="learning rate actor"),
        RL_Fit.paramAlpha(0.0, 1.; name="learning rate critic"),
        RL_Fit.paramGamma(0.0, 1.),
        RL_Fit.paramLambda(0.0, 1.; name="elig. trace actor"), # actor
        RL_Fit.paramLambda(0.0, 1.; name="elig. trace critic"), # critic
        RL_Fit.paramSurpriseProbability(0.0, 1.),
        RL_Fit.paramStochasticity(0.00001, stochmax),
        RL_Fit.paramNParticles(1, 20),
        RL_Fit.paramWeight(0.0, 20.0; name="w_MB"),
        RL_Fit.paramWeight(0.0, 20.0; name="w_MF")
    ]
    hybrid_actorCritic_FWD_BF_prior_std = [.15, .15, .15, .15, .15, .15, 3., 3., 3., 3.,]
    hybrid_actorCritic_FWD_BF_spec = RL_Fit.Model_spec_per_subj("Hybrid Actor-Critic with FWD - BF",
                                RL_Fit.getLL_Hybrid_AC_FWD_BF_V,
                                hybrid_actorCritic_FWD_BF_prior, hybrid_actorCritic_FWD_BF_prior_std,
                                string(result_path_base, "hybrid_actorcritic_FWD_BF/"))
    # --------------------------------------------------------------------------
    # --------------------------------------------------------------------------
    # --------------------------------------------------------------------------
    # --- Run for the paper:
    batch_fit_models_spec = [
                            actorCriticTD_spec,
                            actorCriticTD_modulatedActor_spec,
                            hybrid_actorCritic_FWD_BF_spec,
                            hybrid_actorCritic_FWD_spec,
                            REINFORCE_spec,
                            REINFORCE_modulated_spec,
                            REINFORCE_baseline_spec,
                            sarsaL_spec,
                            sarsaL_modulated_spec,
                            forwardLearner_spec,
                            priosweep_particle_spec,
                            hybridSarsaL_prior_spec,
                            glaescherHybrid_spec,
                            hybrid_pspf_sarsaL_spec,
                            biasedRandomWalk_spec
                          ]
    result_path_base, batch_fit_models_spec
end

# List and specifications of algorithm for parameter recovery
function model_list_paramrecovery(; filestr = "_")

    result_path_base = makehomesavepath("test_main_recovery"*filestr)

    stochmax = 20.

    # Surprise Actor-Critic
    # ---------------------
    actorCriticTD_modulatedActor_prior = [
        RL_Fit.paramAlpha(0.0, 1.; name="learning rate actor"),
        RL_Fit.paramSurpriseProbability(0.0, 1.),
        RL_Fit.paramStochasticity(0.00001, stochmax),
        RL_Fit.paramEta(0.0, 1.; name="learning rate critic"),
        RL_Fit.paramGamma(0.0, 1.),
        RL_Fit.paramLambda(0.0, 1.; name="elig. trace actor"), # actor
        RL_Fit.paramLambda(0.0, 1.; name="elig. trace critic"), # critic
        RL_Fit.paramTemp(0.01, 10.),
        RL_Fit.paramNParticles(1, 20)
    ]
    actorCriticTD_modulatedActor_prior_std = [.15, .15, 3., .15, .15, .15, .15, 1.5, 3.]
    actorCriticTD_modulatedActor_spec = RL_Fit.Model_spec_per_subj("Actor-Critic (TD-Lambda) modulated actor (1 degree)",
                                RL_Fit.getLL_ActorCritic_TDLambda_modulatedActor_V,
                                actorCriticTD_modulatedActor_prior, actorCriticTD_modulatedActor_prior_std,
                                string(result_path_base, "actor_critic_TD_modulatedActor_1degree/"))

    # --------------------------------------------------------------------------
    # --------------------------------------------------------------------------
    # --------------------------------------------------------------------------
    # --- Run for parameter recovery:
    batch_fit_models_spec = [
                            actorCriticTD_modulatedActor_spec,
                          ]

    result_path_base, batch_fit_models_spec
end

# ------------------------------------------------------------------------------
# ------------------------------------------------------------------------------
# ------- MCMC settings and helpers --------------------------------------------
# ------------------------------------------------------------------------------

# Settings for data fitting
# Same settings as in runner_fmri_3foldscrossval (see below),
# but here: only fit, not crossvalidation
function runner_fmri_fit(result_path_base,
                                        batch_fit_models_spec
                                        # ; filestr = "_")
                                        )
    rl_module_path = "hi"
    # --------------------------------------------------------------------------
    # ------ Which data? Real or simulated? Edit!

    # Real participants' data:
    data_path = "./mcmc_rl_fit/projects/fmri/data/SARSPEIC_all_fMRI.csv"

    # Simulated data:
    # data_path = "./mcmc_rl_fit/projects/fmri/data_sim/SARSPEICZVG_all_fMRI_Sim_exporder_SurpAC_run2.csv"

    # --------------------------------------------------------------------------
    # --------------------------------------------------------------------------
    # --------------------------------------------------------------------------
    description = "Batch fit: Models / subjects"
    # ---- DEBUG settings
    max_nr_episodes = 50  # FIT to the first 10 episodes. For the main analysis, I always used 40
    max_nr_episodes = 2  # FIT to the first 10 episodes. For the main analysis, I always used 40
    nr_repetitions = 2
    top_n_fraction = 0.1
    nr_burnin = 3
    nr_samples = 3
    inter_collect_interval = 4
    # # end debug settings
    # ---------------------------
    # ---- Settings used for paper
    # max_nr_episodes = 50  # FIT to the first n episodes.
    # nr_repetitions = 50 #20 # one repetition generates nr_mcmc_samples. doing several repetitions, starting from different random initial positions helps finding isolated regions (e.g. bimodal distributions).
    # top_n_fraction = 0.001 # not important. I sometimes a the top 1percent of mcmc samples
    # nr_burnin = 1500 # important. mcmc needs to "burn in". The more params (larger space) the longer it takes to walk into the volume of interest.
    # nr_samples = 2000 # how many samples to collect from each independent repetition
    # inter_collect_interval = 10 # samples are correlated. therefore make N steps before collecting a sample. 7 to 10 seems a good value.
    # # end
    # --------------------------------------------------------------------------
    # --------------------------------------------------------------------------
    # --------------------------------------------------------------------------

    fit_spec = RL_Fit.FitSpec_multiple_models(
        batch_fit_models_spec,
        max_nr_episodes,
        nr_burnin,
        nr_samples,
        inter_collect_interval,
        top_n_fraction,
        nr_repetitions,
        # false,  #do_per_subject_fit
        true,  #do_per_subject_fit
        result_path_base
    )
    # --------------------------------------------------------------------------
    # ----------------- RUN -----------------
    RL_Fit.run_fit_multiple_models(description, data_path, fit_spec, rl_module_path)

end

# Settings for crossvalidation
# Same settings as in runner_fmri_fit (see above),
# but here: ONLY crossvalidation, not fit
function runner_fmri_3foldscrossval(result_path_base,
                                        batch_fit_models_spec
                                        )
    rl_module_path = "hi"
    # --------------------------------------------------------------------------
    # ------ Which data? Real or simulated? Edit!

    # Real participants' data:
    data_path = "./mcmc_rl_fit/projects/fmri/data/SARSPEIC_all_fMRI.csv"

    # Simulated data:
    # data_path = "./mcmc_rl_fit/projects/fmri/data_sim/SARSPEICZVG_all_fMRI_Sim_exporder_SurpAC_run2.csv"

    # --------------------------------------------------------------------------
    # --------------------------------------------------------------------------
    # --------------------------------------------------------------------------
    description = "Batch fit: Models / subjects"
    # ---- DEBUG settings:
    max_nr_episodes = 50  # FIT to the first 10 episodes. For the main analysis, I always used 40
    max_nr_episodes = 2  # FIT to the first 10 episodes. For the main analysis, I always used 40
    nr_repetitions = 2
    top_n_fraction = 0.1
    nr_burnin = 3
    nr_samples = 3
    inter_collect_interval = 4
    # # end debug settings
    # ---------------------------
    # ---- Settings used for paper
    # max_nr_episodes = 50  # FIT to the first n episodes.
    # nr_repetitions = 50 #20 # one repetition generates nr_mcmc_samples. doing several repetitions, starting from different random initial positions helps finding isolated regions (e.g. bimodal distributions).
    # top_n_fraction = 0.001 # not important. I sometimes a the top 1percent of mcmc samples
    # nr_burnin = 1500 # important. mcmc needs to "burn in". The more params (larger space) the longer it takes to walk into the volume of interest.
    # nr_samples = 2000 # how many samples to collect from each independent repetition
    # inter_collect_interval = 10 # samples are correlated. therefore make N steps before collecting a sample. 7 to 10 seems a good value.
    # # end
    # --------------------------------------------------------------------------
    # --------------------------------------------------------------------------
    # --------------------------------------------------------------------------

    fit_spec = RL_Fit.FitSpec_multiple_models(
        batch_fit_models_spec,
        max_nr_episodes,
        nr_burnin,
        nr_samples,
        inter_collect_interval,
        top_n_fraction,
        nr_repetitions,
        false,  #do_per_subject_fit
        # true,  #do_per_subject_fit
        result_path_base
    )
    # --------------------------------------------------------------------------
    # # ----------------- CROSSVALIDATE -----------------
    (dataFmriAll, all_data_stats) = RL_Fit.load_SARSPEI(data_path; max_episodes_per_subj = max_nr_episodes)
    # when running just an additional model, reuse the same (initially randomly) sampled
    # (make results comparable with previously fit algorithms)

    # I called kf_Spec = RL_Fit.get_k_folds_all_out(Set(all_data_stats.all_subject_ids), 3) once
    # and then fix the subject ID per fold
    kf_Spec = RL_Fit.get_k_folds( Set(all_data_stats.all_subject_ids),
    [
    Set([7, 2, 19, 5, 8, 12, 20])
    Set([18, 4, 13, 3, 15, 21, 6])
    Set([9, 14, 10, 16, 17, 11, 1])])
    # Otherwise shuffle(collect(1:21))... and split in 3
    CV_multi_model_spec = RL_Fit.CV_multiple_models_spec(
        batch_fit_models_spec,
        kf_Spec,
        nr_burnin,
        nr_samples,
        inter_collect_interval,
        top_n_fraction,
        nr_repetitions,
        result_path_base)
    cv_result = RL_Fit.run_cross_validation(description, dataFmriAll,
                                            CV_multi_model_spec, rl_module_path)
end


struct TextBuffer{N}
    positions::GPUVector{Point{N, Float32}}
    offsets::GPUVector{Vec2f0}
    rotations::GPUVector{Vec4f0}
    colors::GPUVector{RGBAf0}
    uv_offset_width::GPUVector{Vec4f0}
    scale::GPUVector{Vec2f0}
    text::Vector{Char}
    robj::RenderObject
    range::Signal{Int}
    cursors
end

get_atlas(x::TextBuffer) = x.robj[:atlas]

function TextBuffer(pos::Point{N, <: AbstractFloat} = Point3f0(0)) where N
    positions = gpuvec(Point{N, Float32}[])
    offset = gpuvec(Vec2f0[])
    rotations = gpuvec(Vec4f0[])
    colors = gpuvec(RGBAf0[])
    uv_offset_width = gpuvec(Vec4f0[])
    scale = gpuvec(Vec2f0[])
    atlas = GLVisualize.get_texture_atlas()
    range = Signal(0)
    robj = visualize(
        (DISTANCEFIELD, positions.buffer),
        atlas = atlas,
        offset = offset.buffer,
        rotation = rotations.buffer,
        color = colors.buffer,
        uv_offset_width = uv_offset_width.buffer,
        scale = scale.buffer,
        distancefield = atlas.images,
        indices = range,
        boundingbox = Signal(AABB{Float32}())
    )
    TextBuffer{N}(
        positions,
        offset,
        rotations,
        colors,
        uv_offset_width,
        scale,
        Char[],
        robj.children[],
        range,
        nothing
    )
end

function Base.empty!(tb::TextBuffer{N}) where N
    resize!(tb.positions, 0)
    resize!(tb.offsets, 0)
    resize!(tb.rotations, 0)
    resize!(tb.colors, 0)
    resize!(tb.uv_offset_width, 0)
    resize!(tb.scale, 0)
    push!(tb.range, 0)
    Reactive.set_value!(tb.robj.boundingbox, AABB{Float32}())
    push!(tb.robj.boundingbox, AABB{Float32}())
    return
end

function Base.append!(tb::TextBuffer, startpos::StaticVector{N}, str::String, scale, color, rot, align, font = GLVisualize.defaultfont()) where N
    atlas = get_atlas(tb)
    pos = Point{N, Float32}(startpos)
    rscale = Float32(scale)
    position = GLVisualize.calc_position(str, Point2f0(0), rscale, font, atlas)
    toffset = GLVisualize.calc_offset(str, rscale, font, atlas)
    aoffsetvec = Vec2f0(alignment2num.(align))
    aoffset = align_offset(rot, Point2f0(0), position[end], atlas, rscale, font, aoffsetvec)
    aoffsetn = Point{N, Float32}(to_nd(aoffset, Val{N}, 0f0))
    uv_offset_width = Vec4f0[GLVisualize.glyph_uv_width!(atlas, c, font) for c = str]
    scale = Vec2f0[GLVisualize.glyph_scale!(atlas, c, font, rscale) for c = str]
    position = map(position) do p
        pn = qmul(rot, Point{N, Float32}(to_nd(p, Val{N}, 0f0)))
        pn .+ (pos .+ aoffsetn)
    end

    append!(tb.positions, position)
    append!(tb.offsets, Vec2f0.(toffset))
    append!(tb.rotations, fill(rot, length(position)))
    append!(tb.colors, fill(to_color(color), length(position)))
    append!(tb.uv_offset_width, uv_offset_width)
    append!(tb.scale, scale)

    bb = value(tb.robj.boundingbox)
    for (s, pos) in zip(scale, position)
        pos3d = Vec{3, Float32}(to_nd(pos, Val{3}, 0))
        bb = update(bb, pos3d)
        bb = update(bb, pos3d .+ Vec{3, Float32}(to_nd(s, Val{3}, 0)))
    end
    push!(tb.robj.boundingbox, bb)
    Reactive.set_value!(tb.robj.boundingbox, bb)
    push!(tb.range, length(tb.positions))
    return
end

function align_offset(rot, startpos, lastpos, atlas, rscale, font, align)
    xscale, yscale = GLVisualize.glyph_scale!('X', rscale)
    xmove = (lastpos-startpos)[1] + xscale
    xmove, yscale, z = qmul(rot, Vec3f0(xmove, yscale, 0))
    if isa(align, GeometryTypes.Vec)
        return -Vec2f0(xmove, yscale) .* align
    elseif align == :top
        return -Vec2f0(xmove/2f0, yscale)
    elseif align == :right
        return -Vec2f0(xmove, yscale/2f0)
    else
        error("Align $align not known")
    end
end

function alignment2num(x::Symbol)
    (x in (:hcenter, :vcenter)) && return 0.5f0
    (x in (:left, :bottom)) && return 0.0f0
    (x in (:right, :top)) && return 1.0f0
    0.0f0 # 0 default, or better to error?
end

struct RichText{C, R, F, P, A}
    text::String
    color::C
    rotation::R
    font::F
    position::P
    align::A
end


module NDTensors

using Base.Threads
using Compat
using Dictionaries
using Random
using LinearAlgebra
using StaticArrays
using HDF5
using Requires
using Strided
using TimerOutputs
using TupleTools

using Base: @propagate_inbounds, ReshapedArray

using Base.Cartesian: @nexprs

using Base.Threads: @spawn

#####################################
# Imports and exports
#
include("exports.jl")
include("imports.jl")

#####################################
# DenseTensor and DiagTensor
#
include("aliasstyle.jl")
include("similar.jl")
include("tupletools.jl")
include("dims.jl")
include("tensorstorage.jl")
include("tensor.jl")
include("contraction_logic.jl")
include("dense.jl")
include("symmetric.jl")
include("linearalgebra.jl")
include("diag.jl")
include("combiner.jl")
include("truncate.jl")
include("svd.jl")

#####################################
# BlockSparseTensor
#
include("blocksparse/blockdims.jl")
include("blocksparse/block.jl")
include("blocksparse/blockoffsets.jl")
include("blocksparse/blocksparse.jl")
include("blocksparse/blocksparsetensor.jl")
include("blocksparse/diagblocksparse.jl")
include("blocksparse/combiner.jl")
include("blocksparse/linearalgebra.jl")

#####################################
# Empty
#
include("empty.jl")

#####################################
# Deprecations
#
include("deprecated.jl")

#####################################
# A global timer used with TimerOutputs.jl
#

const timer = TimerOutput()

#####################################
# Optional block sparse multithreading
#

include("blas_get_num_threads.jl")

const _using_threaded_blocksparse = Ref(false)

function enable_threaded_blocksparse_docstring(module_name)
  return """
             $(module_name).enable_threaded_blocksparse()
             $(module_name).disable_threaded_blocksparse()

         Enable or disable block sparse multithreading.

         Returns the current state of `$(module_name).using_threaded_blocksparse()`, i.e. `true` if threaded block sparse was previously enabled, and `false` if threaded block sparse was previously disabled. This is helpful for turning block sparse threading on or off temporarily. For example:
         ```julia
         using_threaded_blocksparse = $(module_name).enable_threaded_blocksparse()
         # Run code that you want to be threaded
         if !using_threaded_blocksparse
           $(module_name).disable_threaded_blocksparse()
         end
         ```

         Note that you need to start Julia with multiple threads. For example, to start Julia with 4 threads, you can use any of the following:
         ```
         \$ julia --threads=4

         \$ julia -t 4

         \$ JULIA_NUM_THREADS=4 julia
         ```

         In addition, we have found that it is best to disable `BLAS` and `Strided` multithreading when using block sparse multithreading. You can do that with the commands `using LinearAlgebra; BLAS.set_num_threads(1)` and `$(module_name).Strided.disable_threads()`.

         See also: `$(module_name).enable_threaded_blocksparse`, `$(module_name).disable_threaded_blocksparse`, `$(module_name).using_threaded_blocksparse`.
         """
end

function _enable_threaded_blocksparse()
  current_using_threaded_blocksparse = using_threaded_blocksparse()
  if !current_using_threaded_blocksparse
    if Threads.nthreads() == 1
      println(
        "WARNING: You are trying to enable block sparse multithreading, but you have started Julia with only a single thread. You can start Julia with `N` threads with `julia -t N`, and check the number of threads Julia can use with `Threads.nthreads()`. Your system has $(Sys.CPU_THREADS) threads available to use, which you can determine by running `Sys.CPU_THREADS`.\n",
      )
    end
    if blas_get_num_threads() > 1 && Threads.nthreads() > 1
      println(
        "WARNING: You are enabling block sparse multithreading, but BLAS $(BLAS.vendor()) is currently set to use $(blas_get_num_threads()) threads. When using block sparse multithreading, we recommend setting BLAS to use only a single thread, otherwise you may see suboptimal performance. You can set it with `using LinearAlgebra; BLAS.set_num_threads(1)`.\n",
      )
    end
    if Strided.get_num_threads() > 1
      println(
        "WARNING: You are enabling block sparse multithreading, but Strided.jl is currently set to use $(Strided.get_num_threads()) threads for performing dense tensor permutations. When using block sparse multithreading, we recommend setting Strided.jl to use only a single thread, otherwise you may see suboptimal performance. You can set it with `NDTensors.Strided.disable_threads()` and see the current number of threads it is using with `NDTensors.Strided.get_num_threads()`.\n",
      )
    end
    _using_threaded_blocksparse[] = true
  end
  return current_using_threaded_blocksparse
end

function _disable_threaded_blocksparse()
  current_using_threaded_blocksparse = using_threaded_blocksparse()
  if current_using_threaded_blocksparse
    _using_threaded_blocksparse[] = false
  end
  return current_using_threaded_blocksparse
end

"""
$(enable_threaded_blocksparse_docstring(@__MODULE__))
"""
using_threaded_blocksparse() = _using_threaded_blocksparse[]

"""
$(enable_threaded_blocksparse_docstring(@__MODULE__))
"""
enable_threaded_blocksparse() = _enable_threaded_blocksparse()

"""
$(enable_threaded_blocksparse_docstring(@__MODULE__))
"""
disable_threaded_blocksparse() = _disable_threaded_blocksparse()

#####################################
# Optional TBLIS contraction backend
#

const _using_tblis = Ref(false)

using_tblis() = _using_tblis[]

function enable_tblis()
  _using_tblis[] = true
  return nothing
end

function disable_tblis()
  _using_tblis[] = false
  return nothing
end

function __init__()
  @require TBLIS = "48530278-0828-4a49-9772-0f3830dfa1e9" begin
    enable_tblis()
    include("tblis.jl")
  end
  @require Octavian = "6fd5a793-0b7e-452c-907f-f8bfe9c57db4" begin
    include("octavian.jl")
  end
end

end # module NDTensors


using egads

width = 5.0
rad   = 1.0

NULL_E = ego(0)

(context, status) = EG_open( )
if (status < 0) error("Can't open, failure code: %i", status) end

# build cylinder 1
datacyl = [0.0, 0.0,    0.0,
           0.0, 0.0,  width,
           rad]
(ebody1, status) = EG_makeSolidBody(context, CYLINDER, datacyl)
if (status != EGADS_SUCCESS) cleanup(status, context) end

# find faces
nface1_ptr = Ref{Cint}(0)
faces1_ptr = Ref{Ptr{ego}}()
egads.EG_getBodyTopos(ebody1, NULL_E, FACE, nface1_ptr, faces1_ptr)
nface1 = nface1_ptr[]

# build cylinder 2
datacyl = [0.0, 0.0, -width,
           0.0, 0.0, 0.0,
           rad]
(ebody2, status) = EG_makeSolidBody(context, CYLINDER, datacyl)
if (status != EGADS_SUCCESS) cleanup(status, context) end

# find faces
nface2_ptr = Ref{Cint}(0)
faces2_ptr = Ref{Ptr{ego}}()
egads.EG_getBodyTopos(ebody2, NULL_E, FACE, nface2_ptr, faces2_ptr)
nface2 = nface2_ptr[]

# match faces
nmatch_ptr = Ref{Cint}(0)
match_ptr  = Ref{Ptr{Cint}}()
toler = Cdouble(0.0)
EG_matchBodyFaces(ebody1, ebody2, toler, nmatch_ptr, match_ptr)
nmatch = nmatch_ptr[]
matches = fill( 0, 2*nmatch )
for ii in 1:nmatch
    matches[2*ii-1] = unsafe_load(match_ptr[],2*ii-1)
    matches[2*ii]   = unsafe_load(match_ptr[],2*ii)
end

# delete all matching faces
efaces = Vector{ ego }( nface1 + nface2 )
jj = 1
for ii in 1:nface1
    efaces[jj] = unsafe_load(faces1_ptr[],ii); jj += 1
    for kk in 1:nmatch
        if (ii == matches[2*kk-1])
            jj -= 1
        end
    end
end
for ii in 1:nface2
    efaces[jj] = unsafe_load(faces2_ptr[],ii); jj += 1
    for kk in 1:nmatch
        if (ii == matches[2*kk])
            jj -= 1
        end
    end
end
jj -= 1

# sew faces
toler = 0.0
(emodel, result) = EG_sewFaces( Cint(jj), efaces, toler, Cint(0) )
if (status != EGADS_SUCCESS) cleanup(status, context) end

# save model
status = EG_saveModel(emodel, "test.egads")
if (status != EGADS_SUCCESS) cleanup(status, context) end
status = EG_saveModel(emodel, "test.step")
if (status != EGADS_SUCCESS) cleanup(status, context) end

# delete
status = EG_deleteObject(ebody1)
if (status != EGADS_SUCCESS) cleanup(status, context) end
status = EG_deleteObject(ebody2)
if (status != EGADS_SUCCESS) cleanup(status, context) end
status = EG_deleteObject(emodel)
if (status != EGADS_SUCCESS) cleanup(status, context) end

# close everything
cleanup( context )


using POLO: loss!, boost!, stepsize, smooth!, prox!, terminate, log
using LinearAlgebra
import POLO: getf, getx!, initialize!
import Base: copyto!

abstract type MultiThread <: ExecutionPolicy end

initialize!(::ProxGradient{<:MultiThread}) = nothing

mutable struct Consistent <: MultiThread
    mutex::Threads.Mutex
    k::Int
    fval::Float64
    x::Vector{Float64}
    g::Vector{Float64}

    function (::Type{Consistent})()
        return new(Threads.Mutex(),1,0.,Vector{Float64}(),Vector{Float64}())
    end
end

getk(consistent::Consistent) = consistent.k
getf(consistent::Consistent) = consistent.fval
setf!(consistent::Consistent, fval::AbstractFloat) = consistent.fval = fval
readx(consistent::Consistent) = consistent.x
readg(consistent::Consistent) = consistent.g
increment(consistent::Consistent) = consistent.k += 1

function initialize!(consistent::Consistent, x₀::AbstractVector)
    consistent.k = 1
    consistent.fval = 0.
    resize!(consistent.x,length(x₀))
    consistent.x .= x₀
    resize!(consistent.g,length(x₀))
    consistent.g .= zeros(length(x₀))
end

mutable struct Inconsistent <: MultiThread
    k::Threads.Atomic{Int}
    fval::Threads.Atomic{Float64}
    x::Vector{Threads.Atomic{Float64}}
    g::Vector{Threads.Atomic{Float64}}

    function (::Type{Inconsistent})()
        return new(Threads.Atomic{Int}(1),Threads.Atomic{Float64}(0.),Vector{Threads.Atomic{Float64}}(),Vector{Threads.Atomic{Float64}}())
    end
end

function copyto!(::IndexStyle, dest::AbstractArray{<:Threads.Atomic}, ::IndexStyle, src::AbstractArray)
    destinds, srcinds = LinearIndices(dest), LinearIndices(src)
    isempty(srcinds) || (checkbounds(Bool, destinds, first(srcinds)) && checkbounds(Bool, destinds, last(srcinds))) ||
        throw(BoundsError(dest, srcinds))
    @inbounds for i in srcinds
        dest[i][] = src[i]
    end
    return dest
end

getk(inconsistent::Inconsistent) = inconsistent.k[]
getf(inconsistent::Inconsistent) = inconsistent.fval[]
setf!(inconsistent::Inconsistent, fval::AbstractFloat) = inconsistent.fval[] = fval
readx(inconsistent::Inconsistent) = getindex.(inconsistent.x)
readg(inconsistent::Inconsistent) = getindex.(inconsistent.g)
increment(inconsistent::Inconsistent) = Threads.atomic_add!(inconsistent.k, 1)

function initialize!(inconsistent::Inconsistent, x₀::AbstractVector)
    inconsistent.k[] = 1.
    inconsistent.fval[] = 0.
    resize!(inconsistent.x,length(x₀))
    inconsistent.x .= Threads.Atomic{Float64}.(x₀)
    resize!(inconsistent.g,length(x₀))
    inconsistent.g .= Threads.Atomic{Float64}.(0.)
end

function (proxgrad::ProxGradient{<:MultiThread})(x₀::AbstractVector, loss::AbstractLoss, termination::AbstractTermination, logger::AbstractLogger)
    initialize!(proxgrad, x₀)

    Threads.@threads for i = 1:Threads.nthreads()
        kernel(proxgrad, Threads.threadid(), loss, termination, logger)
    end
end

getf(proxgrad::ProxGradient{<:MultiThread}) = getf(proxgrad.execution)
getx!(proxgrad::ProxGradient{<:MultiThread}) = nothing
getx(multithread::MultiThread) = readx(multithread)

function kernel(proxgrad::ProxGradient{<:MultiThread}, wid::Integer, loss::AbstractLoss, termination::AbstractTermination, logger::AbstractLogger)
    xlocal::Vector{Float64} = zeros(length(proxgrad.execution.x))
    glocal::Vector{Float64} = zeros(length(proxgrad.execution.g))

    read!(proxgrad, xlocal)
    flocal::Float64 = loss!(loss, xlocal, glocal)
    while true
        klocal::Int = getk(proxgrad.execution)
        if !iterate!(proxgrad, wid, klocal, flocal, glocal, termination, logger)
            return nothing
        end
        read!(proxgrad, xlocal)
        flocal = loss!(loss, xlocal, glocal)
    end
end

function iterate!(proxgrad::ProxGradient{<:MultiThread}, wid::Integer, klocal::Integer, flocal::Float64, glocal::AbstractVector, termination::AbstractTermination, logger::AbstractLogger, ::Val{false})
    k = getk(proxgrad.execution)
    x = proxgrad.execution.x
    g = proxgrad.execution.g
    if terminate(termination, k, flocal, readx(proxgrad.execution), glocal)
        setf!(proxgrad.execution, flocal)
        return false
    end
    boost!(POLO.boosting(proxgrad), wid, klocal, k, glocal, g)
    smooth!(POLO.smoothing(proxgrad), klocal, k, readx(proxgrad.execution), readg(proxgrad.execution), g)
    η = stepsize(POLO.stepsize(proxgrad), klocal, k, flocal, readx(proxgrad.execution), readg(proxgrad.execution))
    prox!(POLO.prox(proxgrad), η, readx(proxgrad.execution), readg(proxgrad.execution), x)
    increment(proxgrad.execution)
    return true
end
function iterate!(proxgrad::ProxGradient{Inconsistent}, wid::Integer, klocal::Integer, flocal::Float64, glocal::AbstractVector, termination::AbstractTermination, logger::AbstractLogger)
    res = iterate!(proxgrad, wid, klocal, flocal, glocal, termination, logger, Val{false}())
    if wid == 1
        log(logger, klocal, flocal, readx(proxgrad.execution), readg(proxgrad.execution))
    end
    return res
end

function iterate!(proxgrad::ProxGradient{<:MultiThread}, wid::Integer, klocal::Integer, flocal::Float64, glocal::AbstractVector, termination::AbstractTermination, logger::AbstractLogger, ::Val{true})
    lock(proxgrad.execution.mutex)
    res = iterate!(proxgrad, wid, klocal, flocal, glocal, termination, logger, Val{false}())
    log(logger, klocal, flocal, readx(proxgrad.execution), readg(proxgrad.execution))
    unlock(proxgrad.execution.mutex)
    return res
end
iterate!(proxgrad::ProxGradient{Consistent}, wid::Integer, klocal::Integer, flocal::Float64, glocal::AbstractVector, termination::AbstractTermination, logger::AbstractLogger) = iterate!(proxgrad, wid, klocal, flocal, glocal, termination, logger, Val{true}())

function read!(proxgrad::ProxGradient{<:MultiThread}, xlocal::AbstractVector, ::Val{false})
    xlocal .= readx(proxgrad.execution)
end
read!(proxgrad::ProxGradient{Inconsistent}, xlocal::AbstractVector) = read!(proxgrad, xlocal, Val{false}())

function read!(proxgrad::ProxGradient{<:MultiThread}, xlocal::AbstractVector, ::Val{true})
    lock(proxgrad.execution.mutex)
    read!(proxgrad, xlocal, Val{false}())
    unlock(proxgrad.execution.mutex)
end
read!(proxgrad::ProxGradient{Consistent}, xlocal::AbstractVector) = read!(proxgrad, xlocal, Val{true}())


function triangulate(polygon::Vector{Vector{Point{2, Float64}}})
    lengths = map(x-> UInt32(length(x)), polygon)
    len = UInt32(length(lengths))
    array = ccall(
        (:u32_triangulate_f64, earcut),
        Tuple{Ptr{GLTriangle}, Cint},
        (Ptr{Ptr{Float64}}, Ptr{UInt32}, UInt32),
        polygon, lengths, len
    )
    unsafe_wrap(Vector{GLTriangle}, array[1], array[2])
end

function triangulate(polygon::Vector{Vector{Point{2, Float32}}})
    lengths = map(x-> UInt32(length(x)), polygon)
    len = UInt32(length(lengths))
    array = ccall(
        (:u32_triangulate_f32, earcut),
        Tuple{Ptr{GLTriangle}, Cint},
        (Ptr{Ptr{Float32}}, Ptr{UInt32}, UInt32),
        polygon, lengths, len
    )
    unsafe_wrap(Vector{GLTriangle}, array[1], array[2])
end

function triangulate(polygon::Vector{Vector{Point{2, Int64}}})
    lengths = map(x-> UInt32(length(x)), polygon)
    len = UInt32(length(lengths))
    array = ccall(
        (:u32_triangulate_i64, earcut),
        Tuple{Ptr{GLTriangle}, Cint},
        (Ptr{Ptr{Int64}}, Ptr{UInt32}, UInt32),
        polygon, lengths, len
    )
    unsafe_wrap(Vector{GLTriangle}, array[1], array[2])
end

function triangulate(polygon::Vector{Vector{Point{2, Int32}}})
    lengths = map(x-> UInt32(length(x)), polygon)
    len = UInt32(length(lengths))
    array = ccall(
        (:u32_triangulate_i32, earcut),
        Tuple{Ptr{GLTriangle}, Cint},
        (Ptr{Ptr{Int32}}, Ptr{UInt32}, UInt32),
        polygon, lengths, len
    )
    unsafe_wrap(Vector{GLTriangle}, array[1], array[2])
end


"""
Outer construct of `FSimBase.Simulator`.
"""
function Simulator(state0, dyn, p=nothing;
        Problem=:ODE,
        solver=nothing,
        kwargs...)
    if Problem == :ODE && solver == nothing
        solver = Tsit5()
    elseif Problem == :Discrete && solver == nothing
        solver = FunctionMap()
    end
    FSimBase.Simulator(state0, dyn, p;
                       Problem=Problem,
                       solver=solver,
                       kwargs...)
end


using LinearAlgebra

function update_rpr_camera!(oldvals, camera, cam_controls, cam)
    # Create camera
    c = cam_controls
    l, u, p, fov = c.lookat[], c.upvector[], c.eyeposition[], c.fov[]
    far, near, res = c.far[], c.near[], cam.resolution[]
    new_vals = (; l, u, p, fov, far, near, res)
    new_vals == oldvals && return oldvals
    wd = norm(l - p)
    RPR.rprCameraSetSensorSize(camera, res...)
    RPR.rprCameraSetFocusDistance(camera, wd)
    lookat!(camera, p, l, u)
    RPR.rprCameraSetFarPlane(camera, far)
    RPR.rprCameraSetNearPlane(camera, near)
    h = norm(res)
    RPR.rprCameraSetFocalLength(camera, (30*h)/fov)
    # RPR_CAMERA_FSTOP
    # RPR_CAMERA_MODE
    return new_vals
end

function to_rpr_object(context, matsys, scene, plot)
    @warn("Not supported $(typeof(plot))")
    return nothing
end

function insert_plots!(context, matsys, scene, mscene::Makie.Scene, @nospecialize(plot::Combined))
    if isempty(plot.plots) # if no plots inserted, this truly is an atomic
        object = to_rpr_object(context, matsys, mscene, plot)
        if !isnothing(object)
            if object isa AbstractVector
                foreach(x-> push!(scene, x), object)
            else
                push!(scene, object)
            end
        end
    else
        for plot in plot.plots
            insert_plots!(context, matsys, scene, mscene, plot)
        end
    end
end

function to_rpr_light(context::RPR.Context, light::Makie.PointLight)
    pointlight = RPR.PointLight(context)
    map(light.position) do pos
        transform!(pointlight, Makie.translationmatrix(pos))
    end
    map(light.radiance) do r
        setradiantpower!(pointlight, red(r), green(r), blue(r))
    end
    return pointlight
end

function to_rpr_light(context::RPR.Context, light::Makie.AmbientLight)
    env_img = fill(light.color[], 1, 1)
    img = RPR.Image(context, env_img)
    env_light = RPR.EnvironmentLight(context)
    set!(env_light, img)
    return env_light
end

function to_rpr_light(context::RPR.Context, light::Makie.EnvironmentLight)
    env_light = RPR.EnvironmentLight(context)
    last_img = RPR.Image(context, light.image[])
    set!(env_light, last_img)
    setintensityscale!(env_light, light.intensity[])
    on(light.intensity) do i
        setintensityscale!(env_light, i)
    end
    on(light.image) do img
        new_img = RPR.Image(context, img)
        set!(env_light, new_img)
        RPR.release(last_img)
        last_img = new_img
    end
    return env_light
end

function to_rpr_scene(context::RPR.Context, matsys, mscene::Makie.Scene)
    scene = RPR.Scene(context)
    set!(context, scene)
    # Only set background image if it isn't set by env light, since
    # background image takes precedence
    if !any(x-> x isa Makie.EnvironmentLight, mscene.lights)
        env_img = fill(to_color(mscene.backgroundcolor[]), 1, 1)
        img = RPR.Image(context, env_img)
        RPR.rprSceneSetBackgroundImage(scene, img)
    end
    for light in mscene.lights
        rpr_light = to_rpr_light(context, light)
        push!(scene, rpr_light)
    end

    for plot in mscene.plots
        insert_plots!(context, matsys, scene, mscene, plot)
    end
    return scene
end

function replace_scene_rpr!(scene::Makie.Scene, screen=RPRScreen(scene); refresh=Observable(nothing))
    context = screen.context
    matsys = screen.matsys
    screen.scene = scene
    rpr_scene = to_rpr_scene(context, matsys, scene)
    screen.rpr_scene = rpr_scene

    cam_controls = Makie.cameracontrols(scene)
    cam = Makie.camera(scene)
    camera = RPR.Camera(context)
    set!(rpr_scene, camera)
    # hide Makie scene
    # scene.visible[] = false
    foreach(p -> p.visible = false, scene.plots)
    sub = campixel(scene)
    fb_size = size(scene)
    im = image!(sub, rand(RGBAf, fb_size))

    # translate!(im, 0, 0, 1000)

    clear = true
    onany(refresh, cam.projectionview) do _, _
        clear = true
        return
    end

    RPR.rprContextSetParameterByKey1u(context, RPR.RPR_CONTEXT_ITERATIONS, 1)
    cam_values = (;)
    task = @async while isopen(scene)
        cam_values = update_rpr_camera!(cam_values, camera, cam_controls, cam)
        framebuffer2 = render(screen; clear=clear, iterations=1)
        if clear
            clear = false
        end
        data = RPR.get_data(framebuffer2)
        im[1] = reverse(reshape(data, screen.fb_size); dims=2)
        sleep(1/10)
    end
    return context, task, rpr_scene
end

struct RPRBackend <: Makie.AbstractBackend end

mutable struct RPRScreen <: Makie.AbstractScreen
    context::RPR.Context
    matsys::RPR.MaterialSystem
    framebuffer1::RPR.FrameBuffer
    framebuffer2::RPR.FrameBuffer
    fb_size::Tuple{Int, Int}
    scene::Union{Nothing, Scene}
    setup_scene::Bool
    rpr_scene::Union{Nothing, RPR.Scene}
    iterations::Int
    cleared::Bool
end

function RPRScreen(scene::Scene; kw...)
    fb_size = size(scene)
    screen = RPRScreen(fb_size; kw...)
    screen.scene = scene
    return screen
end

function RPRScreen(fb_size::NTuple{2,<:Integer};
                   iterations=NUM_ITERATIONS[],
                   resource=RENDER_RESOURCE[],
                   plugin=RENDER_PLUGIN[], max_recursion=10)
    context = RPR.Context(; resource=resource, plugin=plugin)
    matsys = RPR.MaterialSystem(context, 0)
    set_standard_tonemapping!(context)
    set!(context, RPR.RPR_CONTEXT_MAX_RECURSION, UInt(max_recursion))
    framebuffer1 = RPR.FrameBuffer(context, RGBA, fb_size)
    framebuffer2 = RPR.FrameBuffer(context, RGBA, fb_size)
    set!(context, RPR.RPR_AOV_COLOR, framebuffer1)
    return RPRScreen(context, matsys, framebuffer1, framebuffer2, fb_size, nothing, false, nothing,
                     iterations, false)
end

function render(screen; clear=true, iterations=screen.iterations)
    context = screen.context
    fb_size = screen.fb_size
    framebuffer1 = screen.framebuffer1
    framebuffer2 = screen.framebuffer2
    scene = screen.scene
    if fb_size != size(scene)
        fb_size = size(scene)
        RPR.release(framebuffer1)
        RPR.release(framebuffer2)
        framebuffer1 = RPR.FrameBuffer(context, RGBA, fb_size)
        framebuffer2 = RPR.FrameBuffer(context, RGBA, fb_size)
        set!(context, RPR.RPR_AOV_COLOR, framebuffer1)
        screen.fb_size = fb_size
        screen.framebuffer1 = framebuffer1
        screen.framebuffer2 = framebuffer2
    end
    clear && clear!(framebuffer1)
    RPR.rprContextSetParameterByKey1u(context, RPR.RPR_CONTEXT_ITERATIONS, iterations)
    RPR.render(context)
    RPR.rprContextResolveFrameBuffer(context, framebuffer1, framebuffer2, false)
    return framebuffer2
end

function Makie.colorbuffer(screen::RPRScreen)
    if !screen.setup_scene
        Makie.backend_display(screen, screen.scene)
    end
    data_1d = RPR.get_data(render(screen))
    r = reverse(reshape(data_1d, screen.fb_size), dims=2)
    img = rotl90(r)
    return map(img) do color
        RGB{Colors.N0f8}(mapc(x-> clamp(x, 0, 1), color))
    end
end

function Makie.backend_display(::RPRBackend, scene::Scene)
    screen = RPRScreen(scene)
    Makie.backend_display(screen, scene)
    return screen
end

function Makie.backend_display(screen::RPRScreen, scene::Scene)
    screen.scene = scene
    rpr_scene = to_rpr_scene(screen.context, screen.matsys, scene)
    screen.rpr_scene = rpr_scene
    cam_controls = cameracontrols(scene)
    cam = camera(scene)
    rpr_camera = RPR.Camera(screen.context)
    set!(rpr_scene, rpr_camera)
    update_rpr_camera!((;), rpr_camera, cam_controls, cam)
    screen.setup_scene = true
    return screen
end

function Base.insert!(screen::RPRScreen, scene::Scene, plot::AbstractPlot)
    context = screen.context
    matsys = screen.matsys
    rpr_scene = screen.rpr_scene
    insert_plots!(context, matsys, rpr_scene, scene, plot)
    return screen
end

function Makie.backend_show(b::RPRBackend, io::IO, ::MIME"image/png", scene::Scene)
    screen = Makie.backend_display(b, scene)
    img = Makie.colorbuffer(screen)
    FileIO.save(FileIO.Stream{FileIO.format"PNG"}(Makie.raw_io(io)), img)
    return screen
end

function Base.show(io::IO, ::MIME"image/png", screen::RPRScreen)
    img = Makie.colorbuffer(screen)
    FileIO.save(FileIO.Stream{FileIO.format"PNG"}(Makie.raw_io(io)), img)
    return screen
end


module File

importall Images
importall ImmutableArrays
importall FileSystem
importall Textures
import GL
import BSP.Bsp
import ..Mesh
import ..Tree

import Base:read,sizeof

const Entities       = 1
const Planes         = 2
const Vertices       = 3
const Visibility     = 4
const Nodes          = 5
const TexInfos       = 6
const Faces          = 7
const Lightmaps      = 8
const Leaves         = 9
const LeafFaceTable  = 10
const LeafBrushTable = 11
const Edges          = 12
const FaceEdgeTable  = 13
const Models         = 14
const Brushes        = 15
const BrushSides     = 16
const Pop            = 17
const Areas          = 18
const AreaPortals    = 19

const Magic = "IBSP"

immutable Header
	signature::Uint32
	version::Uint32
end

immutable Lump
	offset::Uint32
	length::Uint32
end

immutable Plane
	normal::Vector3{Float32}
	distance::Float32
	type_::Uint32
end

typealias Vertex Vector3{Float32}

immutable Face
	plane::Uint16
	plane_side::Uint16
	first_edge::Uint32
	num_edges::Uint16
	tex_info::Uint16
	lightmap_styles::Uint32
	lightmap_offset::Uint32
end

typealias FaceEdge Int32

immutable Node
	plane::Uint32
	front::Int32
	back::Int32
	bbox_min::Vector3{Int16}
	bbox_max::Vector3{Int16}
	first_face::Uint16
	num_faces::Uint16
end

immutable Leaf
	brush_or::Uint32
	cluster::Uint16
	area::Uint16
	bbox_min::Vector3{Int16}
	bbox_max::Vector3{Int16}
	first_face::Uint16
	num_faces::Uint16
	first_brush::Uint16
	num_brushes::Uint16
end

typealias LeafFace Uint16

immutable Edge
	v1::Uint16
	v2::Uint16
end

immutable VisOffsets
	sight::Uint32
	sound::Uint32
end

immutable TexInfo
	u::Vector4{Float32}
	v::Vector4{Float32}
	flags::Uint32
	value::Uint32
	name::ASCIIString
	next::Uint32
end
function read(io, ::Type{TexInfo})
	u = read(io, Vector4{Float32})
	v = read(io, Vector4{Float32})
	flags = read(io, Uint32)
	value = read(io, Uint32)
	name = lowercase(rstrip(bytestring(read(io, Uint8, 32)), "\0"))
	next = read(io, Uint32)
	return TexInfo(u,v,flags,value,name,next)
end
sizeof(::Type{TexInfo}) = 76

read(s::IO, t::DataType) = t(ntuple(length(t.types), x->read(s, t.types[x]))...)

function build(idx::Integer, nodes::Vector{Node}, planes::Vector{Plane}, leaves::Vector{Tree.Leaf})
	node = nodes[idx]
	if node.front < 0
		front = leaves[-node.front]
	else
		front = build(node.front+1, nodes, planes, leaves)
	end
	if node.back < 0
		back = leaves[-node.back]
	else
		back = build(node.back+1, nodes, planes, leaves)
	end
	normal = planes[node.plane+1].normal
	distance = planes[node.plane+1].distance
	return Tree.Node(front, back, normal, distance)
end

function readlump(io::IO, lump::Lump, t::Type)
	seek(io, lump.offset)
	read(io, t, int(lump.length / sizeof(t)))
end

function readheightmap(name::String)
	img = imread(string("/home/jay/q2renew/textures/", name, ".height.png"))
	width, height = size(img)[2:3]
	gray = uint8(convert(Array, img))[:,:,1]
	handle = GL.GenTexture()

	# create normals from heights
	normals = Array(Uint8, 0)
	for w = 1:width
		wp = w-1>=1 ? w-1 : width
		wn = w+1<=width ? w+1 : 1
		for h = 1:height
			hp = h-1>=1 ? h-1 : height
			hn = h+1<=height ? h+1 : 1
			x = int(gray[w,hp] - gray[w,hn]) / 32.0
			y = int(gray[wp,h] - gray[wn,h]) / 32.0
			n = Vector3(x,y,1.0)
			n /= norm(n)
			n = uint8(n * (255/2) + (255/2))
			push!(normals, n[1])
			push!(normals, n[2])
			push!(normals, n[3])
			push!(normals, gray[w,h])
		end
	end

	# upload image data to GPU
	GL.BindTexture(GL.TEXTURE_2D, handle)
	GL.TexImage2D(GL.TEXTURE_2D, GL.RGBA, width, height, GL.RGBA, normals)
	GL.GenerateMipmap(GL.TEXTURE_2D)
	GL.TexParameteri(GL.TEXTURE_2D, GL.TEXTURE_MIN_FILTER, GL.LINEAR_MIPMAP_LINEAR)
	GL.TexParameteri(GL.TEXTURE_2D, GL.TEXTURE_MAG_FILTER, GL.LINEAR)
	GL.TexParameteri(GL.TEXTURE_2D, 0x84FE, 8) # anisotropic filtering
	GL.TexParameteri(GL.TEXTURE_2D, GL.TEXTURE_WRAP_S, GL.REPEAT)
	GL.TexParameteri(GL.TEXTURE_2D, GL.TEXTURE_WRAP_T, GL.REPEAT)
	GL.BindTexture(GL.TEXTURE_2D, 0)
	return handle
end

function read(io::IO, ::Type{Bsp})

	###   Read in header and lump data   #######################################

	hdr   = read(io, Header)
	lumps = read(io, Lump, 19)

	bin_edges    = readlump(io, lumps[Edges],      Edge)
	bin_faces    = readlump(io, lumps[Faces],      Face)
	bin_leaves   = readlump(io, lumps[Leaves],     Leaf)
	bin_nodes    = readlump(io, lumps[Nodes],      Node)
	bin_planes   = readlump(io, lumps[Planes],     Plane)
	bin_texinfos = readlump(io, lumps[TexInfos],   TexInfo)
	bin_vertices = readlump(io, lumps[Vertices],   Vertex)
	bin_vis      = readlump(io, lumps[Visibility], Uint8)

	face2edge = readlump(io, lumps[FaceEdgeTable], FaceEdge)
	leaf2face = readlump(io, lumps[LeafFaceTable], LeafFace)

	seek(io, lumps[Visibility].offset)
	num_clusters = read(io, Uint32)
	vis_offsets = read(io, VisOffsets, num_clusters)

	###   Convert File.TexInfo to Mesh.Textures   ##############################

	default_normal_map = GL.GenTexture()
	GL.BindTexture(GL.TEXTURE_2D, default_normal_map)
	GL.TexImage2D(GL.TEXTURE_2D, GL.RGBA, 1, 1, GL.RGBA, Uint8[127,127,255,127])
	GL.TexParameteri(GL.TEXTURE_2D, GL.TEXTURE_MIN_FILTER, GL.NEAREST)
	GL.TexParameteri(GL.TEXTURE_2D, GL.TEXTURE_MAG_FILTER, GL.NEAREST)
	GL.TexParameteri(GL.TEXTURE_2D, GL.TEXTURE_WRAP_S, GL.REPEAT)
	GL.TexParameteri(GL.TEXTURE_2D, GL.TEXTURE_WRAP_T, GL.REPEAT)
	GL.BindTexture(GL.TEXTURE_2D, 0)

	textures = Dict{String,Mesh.Texture}()
	for texinfo = bin_texinfos
		if !has(textures, texinfo.name)
			if texinfo.flags > 0
				println(texinfo.flags, ", ", texinfo.value)
			end

			f = qopen(string("textures/", texinfo.name, ".wal"))
			img = imread(f, Textures.WAL)
			close(f)

			width = uint32(size(img)[2])
			height = uint32(size(img)[3])
			diffuse = GL.GenTexture()

			# upload image data to GPU
			GL.BindTexture(GL.TEXTURE_2D, diffuse)
			GL.TexImage2D(GL.TEXTURE_2D, GL.RGB, width, height, GL.RGB, uint8(img[:]))
			GL.GenerateMipmap(GL.TEXTURE_2D)
			GL.TexParameteri(GL.TEXTURE_2D, GL.TEXTURE_MIN_FILTER, GL.LINEAR_MIPMAP_LINEAR)
			GL.TexParameteri(GL.TEXTURE_2D, GL.TEXTURE_MAG_FILTER, GL.LINEAR)
			GL.TexParameteri(GL.TEXTURE_2D, 0x84FE, 8) # anisotropic filtering
			GL.TexParameteri(GL.TEXTURE_2D, GL.TEXTURE_WRAP_S, GL.REPEAT)
			GL.TexParameteri(GL.TEXTURE_2D, GL.TEXTURE_WRAP_T, GL.REPEAT)
			GL.BindTexture(GL.TEXTURE_2D, 0)

			normal = default_normal_map
			try
				normal = readheightmap(texinfo.name)
				print("[h] ")
			catch e
				print("[ ] ")
			end

			textures[texinfo.name] = Mesh.Texture(diffuse, normal, width, height, texinfo.flags)
			println(texinfo.name)
		end
	end

	###   Convert File.Faces to Mesh.Faces   ###################################

	faces = Array(Mesh.Face,0)
	for face = bin_faces

		indices = Array(Uint16,0)
		first = face.first_edge + 1
		hub = bin_edges[abs(face2edge[first])+1].v1
		last = first + face.num_edges - 1
		for idx = face2edge[first:last]
			v1 = bin_edges[abs(idx)+1].v1
			v2 = bin_edges[abs(idx)+1].v2
			if idx < 0
				v1, v2 = v2, v1
			end
			if v1 != hub && v2 != hub
				push!(indices, hub)
				push!(indices, v2)
				push!(indices, v1)
			end
		end

		ibo = GL.GenBuffer()
		GL.BindBuffer(GL.ELEMENT_ARRAY_BUFFER, ibo)
		GL.BufferData(GL.ELEMENT_ARRAY_BUFFER, indices, GL.STATIC_DRAW)
		GL.BindBuffer(GL.ELEMENT_ARRAY_BUFFER, 0)

		normal = bin_planes[face.plane+1].normal
		if face.plane_side != 0
			normal = -normal
		end

		texinfo = bin_texinfos[face.tex_info+1]
		u = texinfo.u
		v = texinfo.v
		tex = textures[texinfo.name]

		push!(faces, Mesh.Face(indices, ibo, normal, tex, u, v, Array(Mesh.Light,0)))
	end

	###   Convert File.Leaves to Tree.Leaves   #################################

	leaves_in_cluster = Dict{Uint16,Vector{Tree.Leaf}}()

	leaves = Array(Tree.Leaf,0)
	for leaf = bin_leaves
		leaf_faces = Array(Mesh.Face,0)
		first = leaf.first_face + 1
		last = first + leaf.num_faces - 1
		for i = leaf2face[first:last]
			face = faces[i+1]
			if contains(leaf_faces, face)
				warn("duplicate face in leaf")
			else
				push!(leaf_faces, face)
			end
		end
		c = leaf.cluster
		if !has(leaves_in_cluster, c)
			leaves_in_cluster[c] = Array(Tree.Leaf,0)
		end
		leaf = Tree.Leaf(leaf_faces)
		push!(leaves, leaf)
		push!(leaves_in_cluster[c], leaf)
	end

	###   Populate Tree.Leaves visibility info   ###############################

	faces_in_cluster = Dict{Uint16,Vector{Mesh.Face}}()
	for tup = leaves_in_cluster
		c = tup[1]
		cluster_leaves = tup[2]
		visible_faces = Array(Mesh.Face,0)
		for leaf = cluster_leaves
			for face = leaf.faces
				push!(visible_faces, face)
			end
		end
		faces_in_cluster[c] = unique(visible_faces)
	end

	faces_from_cluster = Dict{Uint16,Vector{Mesh.Face}}()
	for tup = faces_in_cluster
		c = tup[1]                   # this cluster index
		visible_faces = copy(tup[2])

		if c == uint16(-1)
			continue
		end

		v = vis_offsets[c+1].sight+1 # pvs buffer index
		k = 0                        # other cluster index
		while k < num_clusters
			if bin_vis[v] == 0
				v += 1
				k += 8 * bin_vis[v]
			else
				for bit = 0:7
					if bin_vis[v] & (1 << bit) != 0 && has(faces_in_cluster, k)
						for face = faces_in_cluster[k]
							push!(visible_faces, face)
						end
					end
					k += 1
				end
			end
			v += 1
		end
		faces_from_cluster[c] = unique(visible_faces)
	end

	for i = 1:length(leaves)
		leaf = leaves[i]
		c = bin_leaves[i].cluster
		if c != uint16(-1)
			leaf.faces = faces_from_cluster[c]
		else
			leaf.faces = faces
		end
	end

	###   Build BSP tree   #####################################################

	tree = build(1, bin_nodes, bin_planes, leaves)

	###   Parse entity data   ##################################################

	# Read entity lump
	seek(io, lumps[Entities].offset)
	entityLump = bytestring(read(io, Uint8, lumps[Entities].length-1))
	entities = Array(Dict{String, String},0)

	# Create a dictionary for each entity
	for entity = split(entityLump, ['{', '}'])
		entity = strip(entity)
		if length(entity) < 1 # TODO: remove this w/ regex
			continue
		end
		dict = Dict{String, String}()
		for field = split(entity, '\n')
			fieldStr = split(strip(field), "\" \"")
			name = lstrip(strip(fieldStr[1], '"'), "_")
			value = strip(fieldStr[2], '"')
			dict[name] = value
		end
		if length(dict) < 1
			continue
		elseif !has(dict, "classname")
			warn("entity has no classname")
			continue
		else
			push!(entities, dict)
		end
	end

	###   Build list of lights from entity data   ##############################

	lights = Array(Mesh.Light,0)
	for ent = entities
		if ent["classname"] != "light"
			continue
		end
		if has(ent, "origin")
			origin = split(ent["origin"])
			origin = Vector3{Float32}(float32(origin[1]), float32(origin[2]), float32(origin[3]))
		else
			warn("light has no origin")
			continue
		end
		if has(ent, "color")
			color = split(ent["color"])
			color = Vector3{Float32}(float32(color[1]), float32(color[2]), float32(color[3]))
			color /= max(color)
		else
			color = Vector3{Float32}(1.0, 1.0, 1.0)
		end
		if has(ent, "light")
			power = float32(ent["light"])
		else
			power = float32(300)
		end
		power *= 1.5
		push!(lights, Mesh.Light(origin, color, power))
	end

	###   Associate lights with visible faces   ################################

	for light = lights
		for face = search(tree, light.origin).faces
			lit = Mesh.islit(face, bin_vertices, light)
			if lit && !contains(face.lights, light)
				push!(face.lights, light)
			end
		end
	end

	###   Calculate leaf statistics   ##########################################

	leaf_stats = Array(Int,0)
	for leaf = leaves
		if length(leaf.faces) > 0
			push!(leaf_stats, length(leaf.faces))
		end
	end
	println("visible faces per non-empty leaf:")
	println("min:    ", min(leaf_stats))
	println("mean:   ", mean(leaf_stats))
	println("median: ", median(leaf_stats))
	println("max:    ", max(leaf_stats))

	###   Calculate light statistics   #########################################

	light_stats = Array(Int,0)
	for face = faces
		push!(light_stats, length(face.lights))
	end
	println("lights per face:")
	println("min:    ", min(light_stats))
	println("mean:   ", mean(light_stats))
	println("median: ", median(light_stats))
	println("max:    ", max(light_stats))

	return Bsp(tree, entities, bin_vertices, max(light_stats))
end

end


abstract type GraphEntity end
abstract type IPFSEntity <: GraphEntity end

struct GQLQuery
    args::Dict{AbstractString,Any}
    fields::AbstractVector{AbstractString}
end

struct Allocation <: IPFSEntity
    id::AbstractString
    ipfshash::AbstractString
    amount::Real
    created_at_epoch::Integer

    function Allocation(
        ipfshash::AbstractString, amount::AbstractString, created_at_epoch::Integer
    )
        return new("", ipfshash, togrt(amount), created_at_epoch)
    end
    function Allocation(ipfshash::AbstractString, amount::Real, created_at_epoch::Integer)
        return new("", ipfshash, amount, created_at_epoch)
    end
    function Allocation(id::AbstractString, ipfshash::AbstractString)
        return new(id, ipfshash, 0.0, 0)
    end
end

struct Indexer <: GraphEntity
    id::AbstractString
    stake::Real
    allocations::AbstractVector{Allocation}

    function Indexer(id, delegation::AbstractString, stake::AbstractString, allocation)
        return new(
            id,
            togrt(stake) + togrt(delegation),
            map(
                x -> Allocation(
                    x["subgraphDeployment"]["ipfsHash"],
                    x["allocatedTokens"],
                    x["createdAtEpoch"],
                ),
                allocation,
            ),
        )
    end
    function Indexer(id, stake::Real, allocation)
        return new(id, stake, allocation)
    end
    function Indexer(allocation)
        return new(
            "",
            0.0,
            map(x -> Allocation(x["id"], x["subgraphDeployment"]["ipfsHash"]), allocation),
        )
    end
end

struct SubgraphDeployment <: IPFSEntity
    id::AbstractString
    ipfshash::AbstractString
    signal::Real

    function SubgraphDeployment(id, ipfshash, signal::AbstractString)
        return new(id, ipfshash, togrt(signal))
    end
    SubgraphDeployment(id, ipfshash, signal::Real) = new(id, ipfshash, signal)
end

struct Repository
    indexers::AbstractVector{Indexer}
    subgraphs::AbstractVector{SubgraphDeployment}
end

verify_ipfshash(x::AbstractString) = startswith(x, "Qm") && length(x) == 46

function togrt(x)::Float64
    return parse(Float64, x) / 1e18
end

signal(s::SubgraphDeployment) = s.signal

ipfshash(x::IPFSEntity) = x.ipfshash

id(x::GraphEntity) = x.id

allocation(i::Indexer) = i.allocations

allocated_stake(a::Allocation) = a.amount


# SPD-License-Identifier: MIT

using ArgParse
using AutoGrad
using CUDA
using Dates
using HDF5
using JLD2
using Knet
using MarkovModels
using Random
using TOML

# Data loader to load batch of utterances.
include("dataloader.jl")

# Differentiable LF-MMI objective function.
include("lfmmi_loss.jl")

function train!(model, batchloader, denfsm, opts, θ, use_gpu, getlengths)
    trainmode!(model)

    acc_loss::Float64 = 0
    N = 0

    for (i, (batch_data, batch_nums, inlengths)) in enumerate(batchloader)
        seqlengths = getlengths(inlengths)
        batch_dens = union([denfsm for i in 1:size(batch_data, 3)]...)
        if use_gpu
            batch_data = CuArray(batch_data)
            batch_nums = MarkovModels.gpu(batch_nums)
        end

        L = @diff lfmmi_loss(model(batch_data), batch_nums, batch_dens, seqlengths)

        update!(value.(θ), grad.([L], θ), opts)

        acc_loss += value(L) / sum(seqlengths)
        N += 1

        # Remove the variable from the scope to free the memory.
        L = nothing
    end

    acc_loss / N
end

function test!(model, batchloader, denfsm, use_gpu, getlengths)
    testmode!(model)

    acc_loss::Float64 = 0
    N = 0

    for (i, (batch_data, batch_nums, inlengths)) in enumerate(batchloader)
        seqlengths = getlengths(inlengths)

        batch_dens = union([denfsm for i in 1:size(batch_data, 3)]...)

        if use_gpu
            batch_data = CuArray(batch_data)
            batch_nums = MarkovModels.gpu(batch_nums)
        end

        L = lfmmi_loss(model(batch_data), batch_nums, batch_dens, seqlengths)

        acc_loss += L
        N += sum(seqlengths)
    end

    acc_loss / N
end

function main(args)
    conf = TOML.parsefile(args["config"])
    epochs = conf["epochs"]
    mbsize = conf["mini_batch_size"]
    use_gpu = args["use-gpu"]
    cdir = args["checkpoint-dir"]
    ckptpath = args["from-checkpoint"]

    ckpt = if ckptpath != ""
        @debug "Using checkpoint: $ckptpath."
        load(ckptpath)
    else
        nothing
    end

    getlengths_fn = s -> getlengths(TOML.parsefile(args["modelconfig"]), s)

    train_numfsms = load(args["train_numfsms"])
    dev_numfsms = load(args["dev_numfsms"])
    denfsm = load(args["denfsm"])["fsm"]

    model = isnothing(ckpt) ? load(args["model"])["model"] : ckpt["model"]

    @debug "model = $model"

    if args["use-gpu"]
        model = model |> gpu
        denfsm = denfsm |> MarkovModels.gpu
    end

    θ = params(model)
    opts = isnothing(ckpt) ? [getoptimizer(conf) for p in θ] : ckpt["optimizers"]
    scheduler = isnothing(ckpt) ? getscheduler(conf) : ckpt["scheduler"]
    best_loss = isnothing(ckpt) ? Float32(Inf) : ckpt["best_loss"]
    startepoch = isnothing(ckpt) ? 1 : ckpt["epoch"]+1
    for epoch in startepoch:epochs
        t₀ = now()

        train_loss = 0
        dev_loss = 0
        shuffledata = epoch > conf["curriculum"] ? true : false
        if ! shuffledata
            @debug "Loading utterances sorted by length (i.e. curriculum training)."
        end

        h5open(args["train"], "r") do trainfea
            trainbl = BatchLoader(trainfea, train_numfsms, mbsize; shuffledata)
            train_loss = train!(model, trainbl, denfsm, opts, θ, use_gpu,
                                getlengths_fn)
        end

        h5open(args["dev"], "r") do devfea
            devbl = BatchLoader(devfea, dev_numfsms, mbsize)
            dev_loss = test!(model, devbl, denfsm, use_gpu, getlengths_fn)
        end

        update_scheduler!(scheduler, opts, dev_loss)

        checkpoint = Dict(
            "model" => model |> cpu,
            "optimizers" => opts,
            "epoch" => epoch,
            "best_loss" => best_loss,
            "scheduler" => scheduler
        )
        ckpt_path = joinpath(cdir, "last.jld2")
        @debug "Saving checkpoint to $ckpt_path."
        save(ckpt_path, checkpoint)

        if dev_loss < best_loss
            best_loss = dev_loss
            ckpt_path = joinpath(cdir, "best.jld2")
            @debug "New best lost is $best_loss; model saved to $ckpt_path."
            save(ckpt_path, checkpoint)
        end

        t₁ = now()
        println("epoch=$epoch/$epochs " *
                "train_loss=$(round(train_loss, digits = 4)) " *
                "dev_loss=$(round(dev_loss, digits = 4)) " *
                "epoch_duration=$((t₁ - t₀).value/1000)")
    end
end

function parse_commandline()
    s = ArgParseSettings()
    @add_arg_table s begin
        "--checkpoint-dir", "-c"
            default = "checkpoints"
            help = "directory where to save training checkpoint"
        "--from-checkpoint", "-f"
            default = ""
            help = "start training from the provided checkpoint"
        "--use-gpu", "-g"
            action = :store_true
            help = "train the model on gpu"
        "modelfile"
            required = true
            help = "model definition file"
        "config"
            required = true
            help = "training configuration file (TOML)"
        "modelconfig"
            required = true
            help = "model configuration file (TOML)"
        "train"
            required = true
            help = "training set"
        "dev"
            required = true
            help = "development set"
        "train_numfsms"
            required = true
            help = "numerator fsms (i.e. alignment graphs) for the training set"
        "dev_numfsms"
            required = true
            help = "numerator fsms (i.e. alignment graphs) for the dev set"
        "denfsm"
            required = true
            help = "denominator graph"
        "model"
            required = true
            help = "input model saved in JLD2 format"
    end
    s.description = """
    Train a neural network using the LF-MMI objective function.
    of utterances to use is read from the standard input.
    """
    parse_args(s)
end

args = parse_commandline()

# This file should define:
#   * `getlengths`
#   * `getoptimizer`
#   * `getscheduler`
include(args["modelfile"])
main(args)


@info "Setting up for job..."
#JOB FILEPATHS
sib_wms_path = "/bench/PhD/NGS_binaries/BMI/sib_nuc_position_sequences.fa_wms.tr"
rys_wms_path = "/bench/PhD/NGS_binaries/BMI/rys_nuc_position_sequences.fa_wms.tr"

sib_df_binary = "/bench/PhD/NGS_binaries/BMI/sib_diff_positions"
rys_df_binary = "/bench/PhD/NGS_binaries/BMI/rys_diff_positions"
combined_df_binary = "/bench/PhD/NGS_binaries/BMI/combined_diff_positions"

sib_code_binary = "/bench/PhD/NGS_binaries/BMI/sib_diff_codes"
rys_code_binary = "/bench/PhD/NGS_binaries/BMI/rys_diff_codes"
combined_code_binary = "/bench/PhD/NGS_binaries/BMI/combined_diff_codes"

sib_diff_bg = "/bench/PhD/NGS_binaries/BMI/sib_diff_bg"
rys_diff_bg = "/bench/PhD/NGS_binaries/BMI/rys_diff_bg"
combined_diff_bg = "/bench/PhD/NGS_binaries/BMI/combined_diff_bg"

sib_e_pth = "/bench/PhD/NGS_binaries/BMI/sib_e"
rys_e_pth = "/bench/PhD/NGS_binaries/BMI/rys_e"
combined_e_pth = "/bench/PhD/NGS_binaries/BMI/combined_e"

#JOB CONSTANTS
const ensemble_size = 500
const no_sources = 8
const source_min_bases = 3
const source_max_bases = 10
@assert source_min_bases < source_max_bases
const source_length_range= source_min_bases:source_max_bases
const mixing_prior = .07
@assert mixing_prior >= 0 && mixing_prior <= 1

@info "Loading master libraries..."
using Distributed, Distributions, Serialization

@info "Adding workers..."
no_local_procs=2
worker_pool=addprocs(no_local_procs, topology=:master_worker)

@info "Loading libraries everywhere..."
@everywhere using BioMotifInference, Random
Random.seed!(myid()*10000)

@info "Assembling uninformative source priors..."
sib_ui_sp = BioMotifInference.assemble_source_priors(no_sources, Vector{Matrix{Float64}}())
rys_ui_sp = BioMotifInference.assemble_source_priors(no_sources, Vector{Matrix{Float64}}())
combined_ui_sp = BioMotifInference.assemble_source_priors(no_sources, Vector{Matrix{Float64}}())

@info "Loading BGHMM likelihood matrix binaries..."
sib_matrix=deserialize(sib_diff_bg)
rys_matrix=deserialize(rys_diff_bg)
combined_matrix=deserialize(combined_diff_bg)

@info "Loading coded observation sets..."
sib_obs = deserialize(sib_code_binary)
rys_obs = deserialize(rys_code_binary)
combined_obs = deserialize(combined_code_binary)

@info "Assembling sib IPM ensemble on uninformative priors..."
isfile(string(sib_e_pth,'/',"ens")) ? (sib_e = deserialize(string(sib_e_pth,'/',"ens"))) :
    (sib_e = BioMotifInference.IPM_Ensemble(worker_pool, sib_e_pth, ensemble_size, sib_ui_sp, (falses(0,0), mixing_prior), sib_matrix, sib_obs, source_length_range, posterior_switch=true); serialize(string(sib_e_pth,'/',"ens"),sib_e))
sib_e=[]; Base.GC.gc();

@info "Assembling rys IPM ensemble on uninformative priors..."
isfile(string(rys_e_pth,'/',"ens")) ? (rys_e = deserialize(string(rys_e_pth,'/',"ens"))) :
    (rys_e = BioMotifInference.IPM_Ensemble(worker_pool, rys_e_pth, ensemble_size, rys_ui_sp, (falses(0,0), mixing_prior), rys_matrix, rys_obs, source_length_range, posterior_switch=true); serialize(string(rys_e_pth,'/',"ens"),rys_e))
rys_e=[]; Base.GC.gc();

@info "Assembling combined IPM ensemble on uninformative priors..."
isfile(string(combined_e_pth,'/',"ens")) ? (combined_e = deserialize(string(combined_e_pth,'/',"ens"))) :
    (combined_e = BioMotifInference.IPM_Ensemble(worker_pool, combined_e_pth, ensemble_size, combined_ui_sp, (falses(0,0), mixing_prior), combined_matrix, combined_obs, source_length_range, posterior_switch=true); serialize(string(combined_e_pth,'/',"ens"),combined_e))
combined_e=[]; Base.GC.gc();

rmprocs(worker_pool)

@info "Job done!"


#####
##### `AbstractDifferential`
#####

"""
The subtypes of `AbstractDifferential` define a custom \"algebra\" for chain
rule evaluation that attempts to factor various features like complex derivative
support, broadcast fusion, zero-elision, etc. into nicely separated parts.

All subtypes of `AbstractDifferential` implement the following operations:

`+(a, b)`: linearly combine differential `a` and differential `b`

`*(a, b)`: multiply the differential `a` by the differential `b`

`Base.conj(x)`: complex conjugate of the differential `x`

`extern(x)`: convert `x` into an appropriate non-`AbstractDifferential` type for
use outside of `ChainContext`.

Valid arguments to these operations are `T` where `T<:AbstractDifferential`, or
where `T` has proper `+` and `*` implementations.

Additionally, all subtypes of `AbstractDifferential` support `Base.iterate` and
`Base.Broadcast.broadcastable(x)`.
"""
abstract type AbstractDifferential end

Base.:+(x::AbstractDifferential) = x

"""
    extern(x)

Return `x` converted to an appropriate non-`AbstractDifferential` type, for use
with external packages that might not handle `AbstractDifferential` types.

Note that this function may return an alias (not necessarily a copy) to data
wrapped by `x`, such that mutating `extern(x)` might mutate `x` itself.
"""
@inline extern(x) = x

@inline Base.conj(x::AbstractDifferential) = x

"""
    refine_differential(𝒟::Type, der)

Converts, if required, a differential object `der`
(e.g. a `Number`, `AbstractDifferential`, `Matrix`, etc.),
to another  differential that is more suited for the domain given by the type 𝒟.
Often this will behave as the identity function on `der`.
"""
refine_differential(::Any, der) = der  # most of the time leave it alone.


abstract type InitialVelocityDistribution{N} end

"""
    (::InitialVelocityDistribution{N})(
        positions::AbstractVector{SVectorF{N}},
        mass_densities::AbstractVector{Float},
    ) -> Vector{SVectorF{N}}

Compute the fluid particle velocities corresponding to this initial velocity
distribution.
"""
function (::InitialVelocityDistribution) end

include("velocity/static.jl")


module ADMIXTURE

using DelimitedFiles, Glob, Pkg.Artifacts

export ADMIXTURE_EXE
export admixture, read_PQ

if Sys.isapple()
    ADMIXTURE_EXE = joinpath(artifact"ADMIXTURE", "dist/admixture_macosx-1.3.0/admixture")
elseif Sys.islinux()
    ADMIXTURE_EXE = joinpath(artifact"ADMIXTURE", "dist/admixture_linux-1.3.0/admixture")
else
    error("ADMIXTURE only supports Linux and MacOS")
end

"""
    admixture(inputFile, K; kwargs...)

Calculate the admixture estimates assuming `K` populations. Input file can be
`.bed` (together with `.bim`, binary PLINK format), `geno` (STRUCTURE format),
or `ped` (12-coding PLINK format). The keyword arguments match those of
ADMIXTURE software. For details, see the ADMIXTUR manual
<http://dalexander.github.io/admixture/admixture-manual.pdf>

# Keyword arguments
- `B`: compute standard errors by bootstrap; default is `nothing`.
- `j`: number of threads to be used; default is `nothing`.
- `s`: random seed; integer or `:time`; default is `nothing`.
- `C`: major convergence criterion; float or integer; default is `nothing`.
- `c`: minor convergence criterion; float or integer; default is `nothing`.
- `qn`: secant pairs to be used in quasi-Newton acceleration; default is 3.
- `l`: tuning parameter in penalized estimation; default is `nothing`.
- `e`: tuning parameter in penalized estimation; default is `nothing`.
- `supervised`: use supervised training (requires `.pop` file); default is false
"""
function admixture(
    inputFile :: AbstractString,
    K         :: Integer;
    B         :: Union{Nothing,Integer} = nothing,
    j         :: Union{Nothing,Integer} = nothing,
    cv        :: Union{Nothing,Integer} = nothing,
    s         :: Union{Nothing,Integer,Symbol} = nothing,
    C         :: Union{Nothing,Real} = nothing,
    c         :: Union{Nothing,Real} = nothing,
    qn        :: Union{Nothing,Integer} = nothing,
    l         :: Union{Nothing,Number} = nothing,
    e         :: Union{Nothing,Number} = nothing,
    supervised::Bool = false
    )
    inputPath = inputFile |> abspath |> normpath
    args  = sizehint!(String[], 12)
    isnothing(B)  || push!(args, isinteger(B) ? "-B$B" : "-B")
    isnothing(j)  || push!(args, "-j$j")
    isnothing(cv) || push!(args, "--cv=$cv")
    isnothing(s)  || push!(args, isinteger(s) ? "-s $s" : "-s time")
    isnothing(C)  || push!(args, "-C $C")
    isnothing(c)  || push!(args, "-c $c")
    isnothing(qn) || push!(args, qn == 0 ? "-a none" : "-a qn$qn")
    isnothing(l)  || isnothing(e) || push!(args, "-l $l", "-e $e")
    !supervised   || push!(args, "--supervised")
    cmd  = `$ADMIXTURE_EXE $inputPath $K $args`
    @info "ADMIXTURE command:\n$cmd\n"
    @info "Output directory: $(pwd())\n"
    run(cmd)
    outputPath = joinpath(pwd(), splitpath(inputPath)[end][1:end-4] * ".$K")
    P, Q = read_PQ(outputPath, args=args)
end

"""
    read_PQ(outputFile)

Read the P (allele frequencies of the inferred ancestral populations) and
Q (the ancestry fractions) estimates output by the ADMIXTURE program.

# Keyword arguments
- `T`:: output data type; default is `Float64`.
- `args`:: parsed ADMIXTURE arguments as a `Vector{String}`; default is empty.
"""
function read_PQ(
    outputFile :: AbstractString;
    T          :: Type = Float64,
    args       :: Vector{String} = String[]
    )
    pfile = normpath(abspath(outputFile)) * ".P"
    qfile = normpath(abspath(outputFile)) * ".Q"
    if any(arg -> occursin("-l ", arg), args)
        # penalized estimates have different file name endings
        # assume unique file ending with P.lambda.* and Q.lambda.*
        pfile = glob("$(basename(pfile)).lambda.*", pwd())[1]
        qfile = glob("$(basename(qfile)).lambda.*", pwd())[1]
    end
    P = readdlm(pfile, ' ', T)
    Q = readdlm(qfile, ' ', T)
    P, Q
end

end # module