text:
abstract type AbstractPeriodicOrbitProblem end

# Periodic orbit computations by finite differences
abstract type AbstractPOFDProblem <: AbstractPeriodicOrbitProblem end
# Periodic orbit computations by shooting
abstract type AbstractShootingProblem <: AbstractPeriodicOrbitProblem end

# get the number of time slices
@inline getM(pb::AbstractPeriodicOrbitProblem) = pb.M

# update a problem with arguments
function updateForBS(prob::AbstractPeriodicOrbitProblem, F, dF, hopfpt, ζr, M, orbitguess_a, period) end

# update a section with a problem, we do nothing by default
updateSection!(prob::AbstractPeriodicOrbitProblem, x, par) = @warn "Not yet implemented!"

Base.size(pb::AbstractPOFDProblem) = (pb.M, pb.N)
onGpu(pb::AbstractPOFDProblem) = pb.ongpu
hasHessian(pb::AbstractPOFDProblem) = pb.d2F == nothing
isInplace(pb::AbstractPOFDProblem) = pb.isinplace

function applyF(pb::AbstractPOFDProblem, dest, x, p)
	if isInplace(pb)
		pb.F(dest, x, p)
	else
		dest .= pb.F(x, p)
	end
	dest
end

function applyJ(pb::AbstractPOFDProblem, dest, x, p, dx)
	if isInplace(pb)
		pb.J(dest, x, p, dx)
	else
		dest .= apply(pb.J(x, p), dx)
	end
	dest
end
####################################################################################################
"""
$(TYPEDEF)

Define a type to interface the Jacobian of the Shooting Problem with the Floquet computation methods. If we use the same code as for `newton` (see below) in `continuation`, it is difficult to tell the eigensolver not to use the jacobian but instead the monodromy matrix.

$(TYPEDFIELDS)
"""
struct FloquetWrapper{Tpb, Tjacpb, Torbitguess, Tp}
	pb::Tpb
	jacpb::Tjacpb
	x::Torbitguess
	par::Tp
end
FloquetWrapper(pb, x, par) = FloquetWrapper(pb, dx -> pb(x, par, dx), x, par)

# jacobian evaluation
(shjac::FloquetWrapper)(dx) = apply(shjac.jacpb, dx)

# specific linear solver to dispatch
struct FloquetWrapperLS{T} <: AbstractLinearSolver
	solver::T
end

(ls::FloquetWrapperLS)(J, rhs; kwargs...) = ls.solver(J, rhs; kwargs...)
(ls::FloquetWrapperLS)(J::FloquetWrapper, rhs; kwargs...) = ls.solver(J.jacpb, rhs; kwargs...)
(ls::FloquetWrapperLS)(J::FloquetWrapper, rhs1, rhs2) = ls.solver(J.jacpb, rhs1, rhs2)

# this is for the use of MatrixBLS
LinearAlgebra.hcat(shjac::FloquetWrapper, dR) = hcat(shjac.jacpb, dR)

####################################################################################################
# newton wrapper
"""
$(SIGNATURES)

This is the Newton-Krylov Solver for computing a periodic orbit using (Standard / Poincaré) Shooting method.
Note that the linear solver has to be apropriately set up in `options`.

# Arguments

Similar to [`newton`](@ref) except that `prob` is either a [`ShootingProblem`](@ref) or a [`PoincareShootingProblem`](@ref). These two problems have specific options to be tuned, we refer to their link for more information and to the tutorials.

- `prob` a problem of type `<: AbstractShootingProblem` encoding the shooting functional G.
- `orbitguess` a guess for the periodic orbit. See [`ShootingProblem`](@ref) and See [`PoincareShootingProblem`](@ref) for information regarding the shape of `orbitguess`.
- `par` parameters to be passed to the functional
- `options` same as for the regular [`newton`](@ref) method.

# Optional argument
- `linearPO` Specify the choice of the linear algorithm, which must belong to `[:autodiffMF, :MatrixFree, :autodiffDense, :FiniteDifferences]`. This is used to select a way of inverting the jacobian dG
    - For `:MatrixFree`, we use an iterative solver (e.g. GMRES) to solve the linear system. The jacobian was specified by the user in `prob`.
    - For `:autodiffMF`, we use iterative solver (e.g. GMRES) to solve the linear system. We use Automatic Differentiation to compute the (matrix-free) derivative of `x -> prob(x, p)`.
    - For `:autodiffDense`. Same as for `:autodiffMF` but the jacobian is formed as a dense Matrix. You can use a direct solver or an iterative one using `options`.
    - For `:FiniteDifferencesDense`, same as for `:autodiffDense` but we use Finite Differences to compute the jacobian of `x -> prob(x, p)` using the `δ = 1e-8` which can be passed as an argument.

# Output:
- solution
- history of residuals
- flag of convergence
- number of iterations
"""
function newton(prob::AbstractShootingProblem, orbitguess, par, options::NewtonPar; linearPO = :MatrixFree, δ = 1e-8, kwargs...)
	@assert linearPO in [:autodiffMF, :MatrixFree, :autodiffDense, :FiniteDifferencesDense]
	if linearPO == :autodiffDense
		# jac = (x, p) -> ForwardDiff.jacobian(z -> prob(z, p), x)
		_J = prob(Val(:JacobianMatrix), orbitguess, par)
		jac = (x, p) -> (prob(Val(:JacobianMatrixInplace), _J, x, p); _J);
	elseif linearPO == :autodiffMF
		jac = (x, p) -> (dx -> ForwardDiff.derivative(z -> prob((@. x + z * dx), p), 0))
	elseif linearPO == :FiniteDifferencesDense
		jac = (x, p) -> finiteDifferences(z -> prob(z, p), x; δ = 1e-8)
	else
		jac = (x, p) -> (dx -> prob(x, p, dx))
	end
	return newton(prob, jac, orbitguess, par, options; kwargs...)
end

"""
$(SIGNATURES)

This is the deflated Newton-Krylov Solver for computing a periodic orbit using a (Standard / Poincaré) Shooting method.

# Arguments

Similar to [`newton`](@ref) except that `prob` is either a [`ShootingProblem`](@ref) or a [`PoincareShootingProblem`](@ref).

# Optional argument
- `linearPO` Specify the choice of the linear algorithm, which must belong to `[:autodiffMF, :MatrixFree, :autodiffDense, :FiniteDifferences]`. This is used to select a way of inverting the jacobian dG
    - For `:MatrixFree`, we use an iterative solver (e.g. GMRES) to solve the linear system. The jacobian was specified by the user in `prob`.
    - For `:autodiffMF`, we use iterative solver (e.g. GMRES) to solve the linear system. We use Automatic Differentiation to compute the (matrix-free) derivative of `x -> prob(x, p)`.
    - For `:autodiffDense`. Same as for `:autodiffMF` but the jacobian is formed as a dense Matrix. You can use a direct solver or an iterative one using `options`.
    - For `:FiniteDifferencesDense`, same as for `:autodiffDense` but we use Finite Differences to compute the jacobian of `x -> prob(x, p)` using the `δ = 1e-8` which can be passed as an argument.

# Output:
- solution
- history of residuals
- flag of convergence
- number of iterations
"""
function newton(prob::AbstractShootingProblem, orbitguess::vectype, par, options::NewtonPar{T, S, E}, defOp::DeflationOperator{T, Tf, vectype}; linearPO = :MatrixFree, kwargs...) where {T, Tf, vectype, S, E}
	@assert linearPO in [:autodiffMF, :MatrixFree, :autodiffDense, :FiniteDifferencesDense]
	if linearPO == :autodiffDense
		# jac = (x, p) -> ForwardDiff.jacobian(z -> prob(z, p), x)
		_J = prob(Val(:JacobianMatrix), orbitguess, par)
		jac = (x, p) -> (prob(Val(:JacobianMatrixInplace), _J, x, p); _J);
	elseif linearPO == :autodiffMF
		jac = (x, p) -> (dx -> ForwardDiff.derivative(z -> prob((@. x + z * dx), p), 0))
	elseif linearPO == :FiniteDifferencesDense
		jac = (x, p) -> finiteDifferences(z -> prob(z, p), x; δ = 1e-8)
	else
		jac = (x, p) -> (dx -> prob(x, p, dx))
	end
	return newton(prob, jac, orbitguess, par, options, defOp; kwargs...)
end

####################################################################################################
# Continuation for shooting problems

"""
$(SIGNATURES)

This is the continuation method for computing a periodic orbit using a (Standard / Poincaré) Shooting method.

# Arguments

Similar to [`continuation`](@ref) except that `prob` is either a [`ShootingProblem`](@ref) or a [`PoincareShootingProblem`](@ref). By default, it prints the period of the periodic orbit.

# Optional argument
- `linearPO` Specify the choice of the linear algorithm, which must belong to `[:autodiffMF, :MatrixFree, :autodiffDense, :FiniteDifferences]`. This is used to select a way of inverting the jacobian dG
    - For `:MatrixFree`, we use an iterative solver (e.g. GMRES) to solve the linear system. The jacobian was specified by the user in `prob`.
    - For `:autodiffMF`, we use iterative solver (e.g. GMRES) to solve the linear system. We use Automatic Differentiation to compute the (matrix-free) derivative of `x -> prob(x, p)`.
    - For `:autodiffDense`. Same as for `:autodiffMF` but the jacobian is formed as a dense Matrix. You can use a direct solver or an iterative one using `options`.
    - For `:FiniteDifferencesDense`, same as for `:autodiffDense` but we use Finite Differences to compute the jacobian of `x -> prob(x, p)` using the `δ = 1e-8` which can be passed as an argument.
- `updateSectionEveryStep = 0` updates the section every `updateSectionEveryStep` step during continuation
"""

function continuation(prob::AbstractShootingProblem, orbitguess, par, lens::Lens, contParams::ContinuationPar, linearAlgo::AbstractBorderedLinearSolver; linearPO = :MatrixFree, updateSectionEveryStep = 0, kwargs...)
	@assert linearPO in [:autodiffMF, :MatrixFree, :autodiffDense, :FiniteDifferencesDense]

	options = contParams.newtonOptions

	if computeEigenElements(contParams)
		contParams = @set contParams.newtonOptions.eigsolver = FloquetQaD(contParams.newtonOptions.eigsolver)
	end

	_finsol = get(kwargs, :finaliseSolution, nothing)
	_finsol2 = isnothing(_finsol) ? (z, tau, step, contResult; kwargs...) ->
		begin
			modCounter(step, updateSectionEveryStep) == 1 && updateSection!(prob, z.u, setParam(contResult, z.p))
			true
		end :
		(z, tau, step, contResult; prob = prob, kwargs...) ->
			begin
				modCounter(step, updateSectionEveryStep) == 1 && updateSection!(prob, z.u, setParam(contResult, z.p))
				_finsol(z, tau, step, contResult; prob = prob, kwargs...)
			end

	if linearPO == :autodiffDense
		# jac = (x, p) -> FloquetWrapper(prob, ForwardDiff.jacobian(z -> prob(z, p), x), x, p)
		_J = prob(Val(:JacobianMatrix), orbitguess, par)
		jac = (x, p) -> (prob(Val(:JacobianMatrixInplace), _J, x, p); FloquetWrapper(prob, _J, x, p));
	elseif linearPO == :FiniteDifferencesDense
		jac = (x, p) -> FloquetWrapper(prob, finiteDifferences(z -> prob(z, p), x), x, p)
	elseif linearPO == :autodiffMF
		jac = (x, p) -> FloquetWrapper(prob, (dx -> ForwardDiff.derivative(z -> prob((@. x + z * dx), p), 0)), x, p)
	elseif linearPO == :FiniteDifferences
		jac = (x, p) -> FloquetWrapper(prob, dx -> (prob(x .+ δ .* dx, p) .- prob(x, p)) ./ δ, x, p)
	else
		jac = (x, p) -> FloquetWrapper(prob, x, p)
	end

	# we have to change the Bordered linearsolver to cope with our type FloquetWrapper
	linearAlgo = @set linearAlgo.solver = FloquetWrapperLS(linearAlgo.solver)

	branch, u, τ = continuation(
		prob, jac,
		orbitguess, par, lens,
		(@set contParams.newtonOptions.linsolver = FloquetWrapperLS(options.linsolver)), linearAlgo;
		printSolution = (x, p) -> (period = getPeriod(prob, x, set(par, lens, p)),),
		finaliseSolution = _finsol2,
		kwargs...,)
	return setproperties(branch; type = :PeriodicOrbit, functional = prob), u, τ
end

"""
$(SIGNATURES)

This is the continuation routine for computing a periodic orbit using a (Standard / Poincaré) Shooting method.

# Arguments

Similar to [`continuation`](@ref) except that `prob` is either a [`ShootingProblem`](@ref) or a [`PoincareShootingProblem`](@ref).

- `printPeriod` boolean to print the period of the solution. This is useful for `prob::PoincareShootingProblem` as this information is not easily available.
"""
function continuation(prob::AbstractShootingProblem, orbitguess, par, lens::Lens, _contParams::ContinuationPar; linearAlgo = nothing, kwargs...)
	_linearAlgo = isnothing(linearAlgo) ?  BorderingBLS(_contParams.newtonOptions.linsolver) : linearAlgo
	return continuation(prob, orbitguess, par, lens, _contParams, _linearAlgo; kwargs...)
end

####################################################################################################
"""
$(SIGNATURES)

Perform automatic branch switching from a Hopf bifurcation point labelled `ind_bif` in the list of the bifurcated points on a previously computed branch `br::ContResult`. It first computes a Hopf normal form.

# Arguments

- `F, dF, d2F, d3F`: function `(x,p) -> F(x,p)` and its differentials `(x,p,dx) -> d1F(x,p,dx)`, `(x,p,dx1,dx2) -> d2F(x,p,dx1,dx2)`... These are used to compute the Hopf normal form.
- `br` branch result from a call to `continuation`
- `ind_hopf` index of the bifurcation point in `br`
- `contParams` parameters for the call to `continuation`
- `prob` problem used to specify the way the periodic orbit is computed. It can be [`PeriodicOrbitTrapProblem`](@ref), [`ShootingProblem`](@ref) or [`PoincareShootingProblem`](@ref) .

# Optional arguments

- `linearPO` linear algorithm used for the computation of periodic orbits when `prob` is [`PeriodicOrbitTrapProblem`](@ref))
- `Jᵗ` is the jacobian adjoint, used for computation of the eigen-elements of the jacobian adjoint, needed to compute the spectral projector for the Hopf normal form. For matrix-free methods, `transpose` is not readily available and the user must provide a dedicated method. In the case of Matrix / Sparse based jacobian, `Jᵗ` should not be passed as it is computed internally more efficiently, i.e. it avoid recomputing the jacobian as it would be if you pass `Jᵗ = (x, p) -> transpose(dF(x, p))`
- `δ = 1e-8` used for finite differences
- `δp` used to specify a particular guess for the parameter on the bifurcated branch which is otherwise determined by `contParams.ds`. This allows to use a step larger than `contParams.dsmax`.
- `ampfactor = 1` factor which alter the amplitude of the bifurcated solution. Useful to magnify the bifurcated solution when the bifurcated branch is very steep.
- `usedeflation = true` whether to use nonlinear deflation (see [Deflated problems](@ref)) to help finding the guess on the bifurcated branch
- `updateSectionEveryStep = 0` updates the section every `updateSectionEveryStep` step during continuation
- `linearPO` specify the way the jacobian is computed.
- all `kwargs` from [`continuation`](@ref)

A modified version of `prob` is passed to `plotSolution` and `finaliseSolution`.

!!! note "Linear solver"
    You have to be careful about the options `contParams.newtonOptions.linsolver`. In the case of Matrix-Free solver, you have to pass the right number of unknowns `N * M + 1`. Note that the options for the preconditioner are not accessible yet.

!!! tip "Jacobian transpose"
    The adjoint of the jacobian `J` is computed internally when `Jᵗ = nothing` by using `transpose(J)` which works fine when `J` is an `AbstractArray`. In this case, do not pass the jacobian adjoint like `Jᵗ = (x, p) -> transpose(d_xF(x, p))` otherwise the jacobian would be computed twice!

!!! warning "Hessian"
    The hessian of `F`, when `d2F` is not `nothing`, is computed with Finite differences.
"""
function continuation(F, dF, d2F, d3F, br::AbstractBranchResult, ind_bif::Int, _contParams::ContinuationPar, prob::AbstractPeriodicOrbitProblem ; Jᵗ = nothing, δ = 1e-8, δp = nothing, ampfactor = 1, usedeflation = false, nev = _contParams.nev, updateSectionEveryStep = 0, kwargs...)
	# compute the normal form of the branch point
	verbose = get(kwargs, :verbosity, 0) > 1 ? true : false
	verbose && (println("--> Considering bifurcation point:"); _show(stdout, br.specialpoint[ind_bif], ind_bif))

	cb = get(kwargs, :callbackN, cbDefault)

	hopfpt = hopfNormalForm(F, dF, d2F, d3F, br, ind_bif ; Jᵗ = Jᵗ, δ = δ, nev = nev, verbose = verbose)

	# compute predictor for point on new branch
	ds = isnothing(δp) ? _contParams.ds : δp
	Ty = typeof(ds)
	pred = predictor(hopfpt, ds; verbose = verbose, ampfactor = Ty(ampfactor))

	verbose && printstyled(color = :green, "#"^61*
			"\n--> Start branching from Hopf bif. point to periodic orbits.
			 \n--> Bifurcation type: ", hopfpt.type,
			"\n----> newp = ", pred.p, ", δp = ", pred.p - br.specialpoint[ind_bif].param,
			"\n----> amplitude = ", pred.amp,
			"\n----> period = ", abs(2pi/pred.ω), "\n")

	# we compute a phase so that the constraint equation
	# < u(0) − u_hopf, ψ > is satisfied, i.e. equal to zero.
	ζr = real.(hopfpt.ζ); ζi = imag.(hopfpt.ζ)
	# this phase is for POTrap problem constraint to be satisfied
	ϕ = atan(dot(ζr, ζr), dot(ζi, ζr))
	verbose && printstyled(color = :green,"----> phase ϕ = ", ϕ/pi, "⋅π\n")

	M = getM(prob)
	orbitguess_a = [pred.orbit(t - ϕ) for t in LinRange(0, 2pi, M+1)[1:M]]

	# build the variable to hold the functional for computing PO based on finite differences
	probPO, orbitguess = problemForBS(prob, F, dF, br.params, hopfpt, ζr, orbitguess_a, abs(2pi/pred.ω))

	if _contParams.newtonOptions.linsolver isa GMRESIterativeSolvers
		_contParams = @set _contParams.newtonOptions.linsolver.N = length(orbitguess)
	end

	# pass the problem to the plotting and printSolution functions
	_plotsol = get(kwargs, :plotSolution, nothing)
	_plotsol2 = isnothing(_plotsol) ? (x,p; k...) -> nothing : (x,p; k...) -> _plotsol(x, (prob = probPO, p = p); k...)

	if :printSolution in keys(kwargs)
		_printsol = get(kwargs, :printSolution, nothing)
 		kwargs = @set kwargs[:printSolution] = (x, p; k...) -> _printsol(x, (prob = probPO, p = p); k...)
	end
	# perform continuation
	branch, u, τ = continuation(probPO, orbitguess, setParam(br, pred.p), br.lens, _contParams; kwargs..., plotSolution = _plotsol2, updateSectionEveryStep = updateSectionEveryStep)

	return Branch(branch, hopfpt), u, τ

end

ext: jl


text:
@enum OperationType begin
    Insert
    Setindex
    Push
    Delete
end

struct Operation
    type::OperationType
    stmt
end

mutable struct SlotRecord
    slotnames::Vector{Symbol}
    newslotmap::Dict{Int, Symbol}
    changemap::Dict{Symbol, Int}
    counter::Int
end

function SlotRecord(ci::CodeInfo)
    slotnames = copy(ci.slotnames)
    newslotmap = Dict{Int, Symbol}()
    changemap = Dict{Symbol, Int}()
    SlotRecord(slotnames, newslotmap, changemap, 0)
end

struct NewSlotNumber
    id::Int
end

@active NewSlotNumber(x) begin
    if x isa NewSlotNumber
        Some(x.id)
    else
        nothing
    end
end

function Base.setindex!(rc::SlotRecord, slot::Symbol, v::Int)
    rc.slotnames[v] = slot
    return SlotNumber(v)
end

function Base.insert!(rc::SlotRecord, v::Int, slot::Symbol)
    insert!(rc.slotnames, v, slot)
    for k in v+1:length(rc.slotnames)
        name = rc.slotnames[k]
        rc.changemap[name] = get(rc.changemap, name, 0) + 1
    end

    id = rc.counter += 1
    rc.newslotmap[id] = slot
    return NewSlotNumber(id)
end

# let's not support this since it requires
# us to check if the slot number is used or not
Base.delete!(::SlotRecord, ::Int) = error("delete slot is not supported")

mutable struct StmtRecord
    counter::Int
    newssa::Dict{Int, Vector{NewSSAValue}}
    data::Dict{Int, Vector{Operation}}
end

StmtRecord() = StmtRecord(0, Dict{Int, Vector{NewSSAValue}}(), Dict{Int, Vector{Operation}}())

Base.getindex(rc::StmtRecord, v) = rc.data[v]
Base.haskey(rc::StmtRecord, v) = haskey(rc.data, v)
Base.iterate(rc::StmtRecord) = iterate(rc.data)
Base.iterate(rc::StmtRecord, st) = iterate(rc.data, st)

function Base.setindex!(rc::StmtRecord, stmt, v)
    push!(
        get!(rc.data, v, Operation[]),
        Operation(Setindex, stmt)
    )
    return stmt
end

function insert_before!(rc::StmtRecord, v, stmt)
    pushfirst!(
        get!(rc.data, v, Operation[]),
        Operation(Insert, stmt)
    )

    rc.counter += 1
    var = NewSSAValue(rc.counter) # length(code) + 1
    pushfirst!(get!(rc.newssa, v, NewSSAValue[]), var)
    return var
end

function insert_after!(rc::StmtRecord, v, stmt)
    push!(
        get!(rc.data, v, Operation[]),
        Operation(Insert, stmt)
    )

    rc.counter += 1
    var = NewSSAValue(rc.counter) # length(code) + 1
    push!(get!(rc.newssa, v, NewSSAValue[]), var)
    return var
end

function Base.delete!(rc::StmtRecord, v)
    push!(
        get!(rc.data, v, Operation[]),
        Operation(Delete, nothing)
    )
    return rc
end

mutable struct NewCodeInfo
    src::CodeInfo
    pc::Int # current SSAValue
    slots::SlotRecord
    stmts::StmtRecord
end

function NewCodeInfo(ci::CodeInfo)
    NewCodeInfo(ci, 1, SlotRecord(ci), StmtRecord())
end

Base.length(ci::NewCodeInfo) = length(ci.src.code)
Base.getindex(ci::NewCodeInfo, idx::Int) = ci.src.code[idx]
Base.eltype(::NewCodeInfo) = Tuple{Int, Any}

function Base.iterate(ci::NewCodeInfo, st::Int=1)
    if st > length(ci)
        ci.pc = 1
        return
    end

    ci.pc = st
    return (st, ci[st]), st + 1
end

function Base.setindex!(ci::NewCodeInfo, stmt, v::Int)
    ci.stmts[v] = stmt
    return stmt
end

function Base.insert!(ci::NewCodeInfo, v::Int, stmt)
    return insert_before!(ci.stmts, v, stmt)
end

function Base.delete!(ci::NewCodeInfo, v::Int)
    delete!(ci.stmts, v)
    return ci
end

function Base.push!(ci::NewCodeInfo, stmt)
    return insert_after!(ci.stmts, ci.pc, stmt)
end

function _emit_code_with_changemap(ci::NewCodeInfo)
    changemap = fill(0, length(ci))
    code, codelocs = [], Int32[]

    for (v, stmt) in enumerate(ci.src.code)
        loc = ci.src.codelocs[v]

        if !haskey(ci.stmts, v)
            push!(code, stmt)
            push!(codelocs, loc)
            continue
        end

        need_push = true
        for op in ci.stmts[v]
            if op.type == Insert
                push!(code, op.stmt)
                push!(codelocs, loc)
                changemap[v] += 1
            elseif op.type == Setindex # replace the old one
                push!(code, op.stmt)
                push!(codelocs, loc)
                need_push = false
            elseif op.type == Push
                push!(code, op.stmt)
                push!(codelocs, loc)
            elseif op.type == Delete
                changemap[v] -= 1
                need_push = false
                continue
            else
                error("unknown operation")
            end
        end

        if need_push
            push!(code, stmt)
            push!(codelocs, loc)
        end
    end
    return code, codelocs, changemap
end

function emit_code(ci::NewCodeInfo)
    code, codelocs, changemap = _emit_code_with_changemap(ci)
    Core.Compiler.renumber_ir_elements!(code, changemap)
    # NOTE: this must after renumber_ir_elements!
    # since now the changemap is accumulated
    replace_new_ssavalue!(code, ci, changemap)
    return code, codelocs
end

function emit_slot_changemap(ci::NewCodeInfo)
    changemap = fill(0, length(ci.src.slotnames))

    for (old, name) in enumerate(ci.src.slotnames)
        if haskey(ci.slots.changemap, name)
            changemap[old] = ci.slots.changemap[name]
        end
    end

    newslotmap = Dict{Int, Int}()
    for (id, slot) in ci.slots.newslotmap
        newslotmap[id] = findfirst(isequal(slot), ci.slots.slotnames)
    end
    return newslotmap, changemap
end

function finish(ci::NewCodeInfo; inline::Bool=true)
    code, codelocs = emit_code(ci)
    newslotmap, changemap = emit_slot_changemap(ci)
    update_slots!(code, newslotmap, changemap)

    new_ci = copy(ci.src)
    new_ci.code = code
    new_ci.codelocs = codelocs
    new_ci.slotnames = ci.slots.slotnames
    new_ci.slotflags = [0x00 for _ in ci.slots.slotnames]
    new_ci.inferred = false # only supports untyped CodeInfo
    new_ci.inlineable = inline
    new_ci.ssavaluetypes = length(code)
    return new_ci
end

function update_slots!(code::Vector, newslotmap, changemap)
    for (v, stmt) in enumerate(code)
        code[v] = update_slots(stmt, newslotmap, changemap)
    end
    return code
end

function update_slots(e, newslotmap, changemap)
    @match e begin
        SlotNumber(id) => SlotNumber(id + changemap[id])
        NewSlotNumber(id) => SlotNumber(newslotmap[id])
        ReturnNode(x) => ReturnNode(update_slots(x, newslotmap, changemap))
        GotoIfNot(cond, dest) => GotoIfNot(update_slots(cond, newslotmap, changemap), dest)
        NewvarNode(SlotNumber(id)) => NewvarNode(SlotNumber(id+slotmap[id]))
        Expr(head, args...) => Expr(head, map(x->update_slots(x, newslotmap, changemap), e.args)...)
        _ => e
    end
end

function replace_new_ssavalue(e, newssamap)
    @match e begin
        NewSSAValue(id) => SSAValue(newssamap[id])
        GotoIfNot(NewSSAValue(id), dest) => GotoIfNot(SSAValue(newssamap[id]), dest)
        ReturnNode(NewSSAValue(id)) => ReturnNode(SSAValue(newssamap[id]))
        Expr(head, args...) => Expr(head, map(x->replace_new_ssavalue(x, newssamap), args)...)
        _ => e
    end
end

function newssamap(ci::NewCodeInfo, changemap::Vector{Int})
    d = Dict{Int, Int}()
    for v in 1:length(ci.src.code)
        ssa = v + changemap[v]
        haskey(ci.stmts.newssa, v) || continue
        newssavalues = ci.stmts.newssa[v]
        for (k, new) in enumerate(newssavalues)
            d[new.id] = ssa - length(newssavalues) + k - 1
        end
    end
    return d
end

function replace_new_ssavalue!(code::Vector, ci::NewCodeInfo, changemap::Vector{Int})
    map = newssamap(ci, changemap)
    for (v, stmt) in enumerate(code)
        code[v] = replace_new_ssavalue(stmt, map)
    end
    return code
end

function Base.map(f, ci::NewCodeInfo)
    for (v, stmt) in ci
        ci[v] = f(stmt)
    end
    return finish(ci)
end


text:
using Test, Krylov, LinearOperators, LinearAlgebra, SparseArrays, Printf

include("test_diom.jl")
include("test_cgs.jl")
include("test_dqgmres.jl")
include("gen_lsq.jl")
include("check_min_norm.jl")
include("test_aux.jl")
include("test_cg.jl")
include("test_cg_lanczos.jl")
include("test_minres.jl")
include("test_symmlq.jl")
include("test_cgls.jl")
include("test_crls.jl")
include("test_cgne.jl")
include("test_crmr.jl")
include("test_variants.jl")
include("test_lslq.jl")
include("test_lsqr.jl")
include("test_lsmr.jl")
include("test_craig.jl")
include("test_craigmr.jl")
include("test_cr.jl")

include("test_alloc.jl")


text:
using Documenter, NeuralVerification

makedocs(sitename = "NeuralVerification.jl",
         pages = ["index.md", "problem.md", "solvers.md", "functions.md", "existing_implementations.md"])

deploydocs(
    repo = "github.com/sisl/NeuralVerification.jl.git",
)


text:
# ------------ Attribute alias and APIs ------------ #

const Attribute = MLIR.API.MlirAttribute

get_context(attr::Attribute) = MLIR.API.mlirAttributeGetContext(attr)
get_type(attr::Attribute) = MLIR.API.mlirAttributeGetType(attr)
is_null(attr::Attribute) = MLIR.API.mlirAttributeIsNull(attr)
Base.:(==)(attr1::Attribute, attr2::Attribute) = MLIR.API.mlirAttributeEqual(attr1, attr2)
function parse_attribute(ctx::Context, attr::String)
    ref = StringRef(attr)
    MLIR.API.mlirAttributeParseGet(ctx, ref)
end
dump(attr::Attribute) = MLIR.API.mlirAttributeDump(attr)

# Constructor.
Attribute(ctx::Context, attr::String) = parse_attribute(ctx, attr)

@doc(
"""
const Attribute = MLIR.API.MlirAttribute
""", Attribute)

# ------------ Named attribute alias and APIs ------------ #

const NamedAttribute = MLIR.API.MlirNamedAttribute

create_named_attribute(name, attr::Attribute) = MLIR.API.mlirNamedAttributeGet(name, attr)

# Constructor.
function NamedAttribute(ctx::Context, name::String, attr::Attribute)
    id = MLIR.IR.Identifier(ctx, name) # owned by MLIR
    create_named_attribute(id, attr)
end

@doc(
"""
const NamedAttribute = MLIR.API.MlirNamedAttribute
""", NamedAttribute)


text:
# This file was generated by the Julia Swagger Code Generator
# Do not modify this file directly. Modify the swagger specification instead.

mutable struct ExpressRouteCircuitRoutesTable <: SwaggerModel
    network::Any # spec type: Union{ Nothing, String } # spec name: network
    nextHop::Any # spec type: Union{ Nothing, String } # spec name: nextHop
    locPrf::Any # spec type: Union{ Nothing, String } # spec name: locPrf
    weight::Any # spec type: Union{ Nothing, Int32 } # spec name: weight
    path::Any # spec type: Union{ Nothing, String } # spec name: path

    function ExpressRouteCircuitRoutesTable(;network=nothing, nextHop=nothing, locPrf=nothing, weight=nothing, path=nothing)
        o = new()
        validate_property(ExpressRouteCircuitRoutesTable, Symbol("network"), network)
        setfield!(o, Symbol("network"), network)
        validate_property(ExpressRouteCircuitRoutesTable, Symbol("nextHop"), nextHop)
        setfield!(o, Symbol("nextHop"), nextHop)
        validate_property(ExpressRouteCircuitRoutesTable, Symbol("locPrf"), locPrf)
        setfield!(o, Symbol("locPrf"), locPrf)
        validate_property(ExpressRouteCircuitRoutesTable, Symbol("weight"), weight)
        setfield!(o, Symbol("weight"), weight)
        validate_property(ExpressRouteCircuitRoutesTable, Symbol("path"), path)
        setfield!(o, Symbol("path"), path)
        o
    end
end # type ExpressRouteCircuitRoutesTable

const _property_map_ExpressRouteCircuitRoutesTable = Dict{Symbol,Symbol}(Symbol("network")=>Symbol("network"), Symbol("nextHop")=>Symbol("nextHop"), Symbol("locPrf")=>Symbol("locPrf"), Symbol("weight")=>Symbol("weight"), Symbol("path")=>Symbol("path"))
const _property_types_ExpressRouteCircuitRoutesTable = Dict{Symbol,String}(Symbol("network")=>"String", Symbol("nextHop")=>"String", Symbol("locPrf")=>"String", Symbol("weight")=>"Int32", Symbol("path")=>"String")
Base.propertynames(::Type{ ExpressRouteCircuitRoutesTable }) = collect(keys(_property_map_ExpressRouteCircuitRoutesTable))
Swagger.property_type(::Type{ ExpressRouteCircuitRoutesTable }, name::Symbol) = Union{Nothing,eval(Meta.parse(_property_types_ExpressRouteCircuitRoutesTable[name]))}
Swagger.field_name(::Type{ ExpressRouteCircuitRoutesTable }, property_name::Symbol) =  _property_map_ExpressRouteCircuitRoutesTable[property_name]

function check_required(o::ExpressRouteCircuitRoutesTable)
    true
end

function validate_property(::Type{ ExpressRouteCircuitRoutesTable }, name::Symbol, val)
end


text:
using tele
using NCDatasets, Revise, Statistics, PyPlot, PyCall
using CFTime, DSP, FFTW
using DataFrames, NaNMath
#analyzing data from https://psl.noaa.gov/data/gridded/data.noaa.ersst.v5.html

#FILE PATHS
data_dir = "/home/brynn/Documents/JP/860/tsa_data/"
ersst_p = data_dir * "sst.mnmean.nc"
# cesm_p2 = data_dir * "tos_Omon_CCSM4_past1000_r1i1p1_135001-185012.nc"
cesm_p2 = data_dir * "tos_Omon_CCSM4_historical_r1i1p1_185001-200512.nc"

#READ NETCDFS TO BOUNDING BOX
bbox = [71, 55, 290, 350]
(late, lone, timee, datae) = read_netcdf(ersst_p, "sst", bbox)
(latc, lonc, timec, datac) = read_netcdf(cesm_p2, "tos", bbox)

#FILTER CHARACTERISTICS
#sampling = monthly
#cutoff all signals below 3 year periodicity
responsetype = Lowpass(0.1*10^(-7),fs = 3.80265176 * 10^(-7))
designmethod = Butterworth(4)

#CALCULATE ANOMALY AND FILTER
"""
all this permutedims! nonsense is because my anomaly code assumes lon x lon x time
and filter assumes time x lon x lat
"""

#ersst
dataep = zeros(length(timee), length(lone), length(late))
permutedims!(dataep, datae, [3,1,2])
filte = filt(digitalfilter(responsetype, designmethod), dataep)
permutedims!(datae, filte, [2,3,1])
anome = anomaly(datae, timee)
permutedims!(filte, anome, [3,1,2])
#filte holds filtered anomalized data

#cesm
timec = reinterpret(typeof(timee[begin]), timec) #convert from DatetimeNoLeap to DateTime
datac = datac .- 273.15 #convert to celsius

datacp = zeros(length(timec), length(lonc), length(latc))
permutedims!(datacp, datac, [3,1,2])
filtc = filt(digitalfilter(responsetype, designmethod), datacp)
permutedims!(datac, filtc, [2,3,1])
anomc = anomaly(datac, timec)
permutedims!(filtc, anomc, [3,1,2])
#filtc holds filtered, anomalized data

anomc = anomaly(datac, timec)
anomcp = zeros(length(timec), length(lonc), length(latc))
permutedims!(anomcp, anomc, [3,1,2])
filtc = filt(digitalfilter(responsetype, designmethod), anomcp)

#COMPUTE EOFS
eofs = pyimport("eofs")

#cesm
solver = eofs.standard.Eof(filtc)
savetype = ".pdf"
plot_eofs(solver, lonc, latc, timec, "images/eofc"*savetype)

#ersst
solver = eofs.standard.Eof(filte)
plot_eofs(solver, lone, late, timee, "images/eofe"*savetype)

#CORRELATION PLOT
#at every grid point, correlate the mean time series with the grid point's timeseries
tavgc = spatial_mean(anomc)
tavgc_filt = filt(digitalfilter(responsetype, designmethod), tavgc)
tavge = spatial_mean(anome)
tavge_filt = filt(digitalfilter(responsetype, designmethod), tavge)

xcc = zeros(length(timec)*2 - 1, length(lonc), length(latc))
xce = zeros(length(timee)*2 - 1, length(lone), length(late))
regc = zeros(length(lonc), length(latc))
rege = zeros(length(lone), length(late))

for i in 1:length(lonc)
    for j in 1:length(latc)
        ts = filtc[:, i, j]
        xcc[:, i, j] = xcorr(tavgc_filt, ts)
        regc[i, j] = normalequation(1:length(ts), ts)
    end
end

for i in 1:length(lone)
    for j in 1:length(late)
        ts = filte[:, i, j]
        xce[:, i, j] = xcorr(tavge_filt, ts)
        rege[i, j] = normalequation(1:length(ts), ts)
    end
end

plot_corr(xcc, latc, lonc, timec, "images/xcc"*savetype)
plot_corr(xce, late, lone, timee, "images/xce"*savetype)

plot_reg(regc, latc, lonc, "images/regc"*savetype)
plot_reg(rege, late, lone, "images/rege"*savetype)

#extreme value anlaysis
quant = quantile(tavgc, [0.25, 0.5, 0.75])
thresh1 = (quant[3] - quant[1])* 1 + quant[3]
thresh2 = quant[1] - (quant[3] - quant[1])* 1
ind1 = findall(>(thresh1), tavgc)
ind2 = findall(<(thresh2), tavgc)
extreme1 = temp_mean(anomc[:,:, ind1])
extreme2 = temp_mean(anomc[:,:, ind2])

#Pot extreme values
contplot(latc,lonc,extreme2, "images/ext_anomc_low"*savetype)
contplot(latc, lonc, extreme1, "images/ext_anomc_high"*savetype)

#extreme value anlaysis
quant = quantile(tavge, [0.25, 0.5, 0.75])
thresh1 = (quant[3] - quant[1])* 1 + quant[3]
thresh2 = quant[1] - (quant[3] - quant[1])* 1
ind1 = findall(>(thresh1), tavge)
ind2 = findall(<(thresh2), tavge)
extreme1 = temp_mean(anome[:,:, ind1])
extreme2 = temp_mean(anome[:,:, ind2])

#Pot extreme values
contplot(late,lone,extreme2, "images/ext_anome_low"*savetype)
contplot(late, lone, extreme1, "images/ext_anome_high"*savetype)


text:
### A Pluto.jl notebook ###
# v0.17.2

using Markdown
using InteractiveUtils

# ╔═╡ 22c33fe6-4efd-11ec-351c-25a9d73ec50a
begin
	let
		using Pkg
		Pkg.activate(mktempdir())
		Pkg.Registry.update()
		Pkg.add("PlutoUI")
		Pkg.add("CairoMakie")
		Pkg.add("BenchmarkTools.jl")
		Pkg.add(url="https://github.com/Dale-Black/DistanceTransforms.jl")
	end

	using PlutoUI
	using CairoMakie
	using BenchmarkTools
	using DistanceTransforms
end

# ╔═╡ 0095c4d9-e9d0-4769-83ff-fd7f4346be8f
md"""
## Import packages
First, let's import the most up-to-date version of DistanceTransforms.jl, which can be found on the main/master branch of the [GitHub repository](https://github.com/Dale-Black/DistanceTransforms.jl).

Because we are using the unregistered version (most recent) we will need to `Pkg.add` this explicitly, without using Pluto's built-in package manager. Be aware, this can take a long time, especially if this is the first time being downloaded. Future work on this package will focus on improving this.

To help with the formatting of this documentation we will also add [PlutoUI.jl](https://github.com/JuliaPluto/PlutoUI.jl).

Lastly, to visualize results and time the functions were going to add [Makie.jl](https://github.com/JuliaPlots/Makie.jl) and [BenchmarkTools.jl](https://github.com/JuliaCI/BenchmarkTools.jl).
"""

# ╔═╡ d70e8f0c-9052-47a4-bac7-af8362afda2c
md"""
## Quick start
Distance transforms are an important part of many computer vision-related tasks. Let's create some sample data and see how DistanceTransforms.jl can be used to apply efficient distance transform operations on arrays in Julia.
"""

# ╔═╡ 6c01894f-442a-497c-9519-dfc18b63be06
md"""
The quintessential distance transform operation in DistanceTransforms.jl is just a wrapper function combining the `feature_transform` and `distance_transform` functions from the excellent [ImageMorphology](https://github.com/JuliaImages/ImageMorphology.jl) package.

To utilize this in DistanceTransforms.jl all one must do is call `euclidean(x)`
"""

# ╔═╡ 023529de-8f6b-4339-89a6-23d8b8fb881f
array1 = [
	1 0 0 1
	1 1 1 1
	0 0 1 1
]

# ╔═╡ cc17f05a-6e29-4a81-a45f-2dcfdedf2e88
euclidean(array1)

# ╔═╡ 0ad7423a-f2ff-417a-bfd1-e400a994e88c
md"""
As you can see, each element in the array is either a zero, which remains zero, or is one that then gets replaced by the Euclidean distance to the nearest zero.

We can see this easily using Makie.jl.
"""

# ╔═╡ fde6da34-cd4d-4d97-aeb5-5a745cdffa13

# ╔═╡ Cell order:
# ╟─0095c4d9-e9d0-4769-83ff-fd7f4346be8f
# ╠═22c33fe6-4efd-11ec-351c-25a9d73ec50a
# ╟─d70e8f0c-9052-47a4-bac7-af8362afda2c
# ╟─6c01894f-442a-497c-9519-dfc18b63be06
# ╠═023529de-8f6b-4339-89a6-23d8b8fb881f
# ╠═cc17f05a-6e29-4a81-a45f-2dcfdedf2e88
# ╟─0ad7423a-f2ff-417a-bfd1-e400a994e88c
# ╠═fde6da34-cd4d-4d97-aeb5-5a745cdffa13


text:
"""
    loop_task1(indx::Int)

Interpolate all variables for one record
"""
function loop_task1(indx::Int)
    task=YAML.load(open("task.yml"))
    M=load(task["Specs"]["file"])
    MetaFile=loop_task1(indx,M)
end

"""
    loop_task1(indx,M)

Loop over a subset of model output files (`filList[indx]`), apply
`MatrixInterp` (`M`) as a postprocessing step, and write the result
to file (one subfolder for each variable)
"""
function loop_task1(indx,M)
   task=YAML.load(open("task.yml"))
   dirIn=task["InputDir"][1]
   filIn=task["InputFile"][1]
   dirOut=task["OutputDir"]
   !isdir(dirOut) ? mkdir(dirOut) : nothing
   MTRX=M["MTRX"]
   msk=M[task["Specs"]["mask"]]
   siz=Tuple(task["OutputSize"])

   tmp1=readdir(dirIn)
   tmp1=filter(x -> occursin(filIn,x),tmp1)
   filList=filter(x -> occursin(".data",x),tmp1)
   maximum(indx)>length(filList) ? error("missing files: "*filIn*"*") : nothing
   filList=filList[indx]

   !isa(filList,Array) ? filList=[filList] : nothing
   nf=length(filList)
   MetaFile=filList[1]
   MetaFile=dirIn*MetaFile[1:end-5]*".meta"
   MetaFile=MetaFileRead(MetaFile)

   nv=length(MetaFile["fldList"])
   nd=MetaFile["nDims"]
   dims=Int.(MetaFile["dimList"][:,1])
   prec=MetaFile["dataprec"]

   for ff=1:nf
      fil=dirIn*filList[ff]
      #println(fil)
      fid = open(fil)
      for vv=1:nv
         tmp=Array{Float32,2}(undef,(90,1170))
         nd==3 ? tmp=Array{Float32,3}(undef,(90,1170,50)) : nothing
         read!(fid,tmp)
         tmp = hton.(tmp)
         !ismissing(msk) ? tmp=tmp.*msk : nothing
         tmp=MatrixInterp(tmp,MTRX,siz)
         tmp=Float32.(tmp)
         #
         filOut=dirOut*strip(MetaFile["fldList"][vv])*"/"
         !isdir(filOut) ? mkdir(filOut) : nothing
         filOut=filOut*strip(MetaFile["fldList"][vv])
         filOut=filOut*fil[length(dirIn*filIn)+1:end]
         #println(filOut)
         #
         nd==3 ? recl=720*360*50*4 : recl=720*360*4
         f =  FortranFile(filOut,"w",access="direct",recl=recl,convert="big-endian")
         write(f,rec=1,tmp)
         close(f)
         #to re-read file:
         #f =  FortranFile(filOut,"r",access="direct",recl=recl,convert="big-endian");
         #tmp1=read(f,rec=1,(Float32,(720,360))); close(f);
      end
      close(fid)
   end

   return MetaFile
end

"""
    loop_task2(indx)

Loop over a subset of model output files (`filList[indx]`), convert Rirr*
to Satellite Chl as a postprocessing step, and write the result
to file (one subfolder for each variable)
"""
function loop_task2(indx)

   task=YAML.load(open("task.yml"))

   dirIn=task["InputDir"][1]
   filIn=task["InputFile"][1]
   dirOut=task["OutputDir"]
   filOut=task["OutputFile"]
   !isdir(dirOut) ? mkdir(dirOut) : nothing

   siz=Tuple(task["OutputSize"])
   prec=Float32

   C=task["Specs"]["PolyCoeffs"]
   a0=C[1]; a1=C[2]; a2=C[3]; a3=C[4]; a4=C[5];

   tmp1=readdir(dirIn*filIn)
   tmp1=filter(x -> occursin(filIn,x),tmp1)
   filList=filter(x -> occursin(".data",x),tmp1)
   maximum(indx)>length(filList) ? error("missing files: "*filIn*"*") : nothing
   filList=filList[indx]

   !isa(filList,Array) ? filList=[filList] : nothing
   nf=length(filList)

   wv_cci=[412, 443, 490, 510, 555, 670];
   wv_drwn3=[400,425,450,475,500,525,550,575,600,625,650,675,700];
   jj=Array{Int64,1}(undef,6)
   ww=Array{Float64,1}(undef,6)
   for ii=1:6
      tmp=wv_cci[ii].-wv_drwn3
      kk=maximum(findall(tmp.>=0))
      jj[ii]=kk
      ww[ii]=tmp[kk]/(wv_drwn3[kk+1]-wv_drwn3[kk])
   end

   #C=[0.2424    -2.7423 +1.8017 +0.0015 -1.2280] #OC3M-547 (MODIS)
   #C=[0.3272    -2.9940 +2.7218 -1.2259 -0.5683] #OC4 (SeaWifs, CCI)
   #a0=C[1]; a1=C[2]; a2=C[3]; a3=C[4]; a4=C[5];

   for ff=1:nf
      fil=dirIn*filIn*"/"*filList[ff]

      #1) read Rirr
      Rirr=Array{Float32,3}(undef,(siz[1],siz[2],13))
      for vv=1:13
         fil0=replace(fil,task["InputFile"][1] => task["InputFile"][vv])
         tmp=Array{Float32,2}(undef,siz)
         fid = open(fil0)
         read!(fid,tmp)
         Rirr[:,:,vv] = hton.(tmp)
         close(fid)
      end

      #2) convert to Rrs & interpolate in frequency space
      tmp=Rirr/3
      Rrs0=(0.52*tmp)./(1.0 .-1.7*tmp)
      Rrs=Array{Float32,3}(undef,(siz[1],siz[2],6))
      for vv=1:6
         tmp0=Rrs0[:,:,jj[vv]]
         tmp1=Rrs0[:,:,jj[vv]+1]
         Rrs[:,:,vv]=tmp0.*(1-ww[vv])+tmp1.*ww[vv]
      end

      #3) compute "satellite Chl" in model
      RRSB=max.(Rrs[:,:,2],Rrs[:,:,3])
      RRSG=Rrs[:,:,5]
      X = log10.(RRSB./RRSG) #ratio of blue to green
      chld=exp10.(a0.+a1*X+a2*X.^2+a3*X.^3+a4*X.^4);

      #println(fil)
      recl=720*360*4
      filOut=dirOut*"chld/"
      !isdir(filOut) ? mkdir(filOut) : nothing
      filOut=filOut*replace(filList[ff],task["InputFile"][1] => task["OutputFile"])
      f =  FortranFile(filOut,"w",access="direct",recl=recl,convert="big-endian")
      write(f,rec=1,Float32.(chld))
      close(f)
      #to re-read file:
      #f =  FortranFile(filOut,"r",access="direct",recl=recl,convert="big-endian");
      #tmp1=read(f,rec=1,(Float32,(720,360))); close(f);
   end

   return filList
end

"""
    loop_task3(indx)

Loop over a subset of model output files (`filList[indx]`), compute
export as a postprocessing step, and write the result
to file (one subfolder for each variable)
"""
function loop_task3(indx)

   task=YAML.load(open("task.yml"))

   dirIn=task["InputDir"][1]
   filIn=task["InputFile"][1]
   dirOut=task["OutputDir"]
   filOut=task["OutputFile"]
   !isdir(dirOut) ? mkdir(dirOut) : nothing

   siz=Tuple(task["InputSize"])
   prec=Float32
   recl=prod(siz)*4

   tmp1=readdir(dirIn*filIn)
   tmp1=filter(x -> occursin(filIn,x),tmp1)
   filList=filter(x -> occursin(".data",x),tmp1)
   maximum(indx)>length(filList) ? error("missing files: "*filIn*"*") : nothing
   filList=filList[indx]

   !isa(filList,Array) ? filList=[filList] : nothing
   nf=length(filList)

   Wsink=task["Specs"]["Wsink"]
   #convert from m/s to m/day
   Wsink=86400.0*Wsink

   for ff=1:nf
      #1) read biomass
      fil=dirIn*filIn*"/"*filList[ff]
      tmp=Array{Float32,3}(undef,siz)
      fid = open(fil)
      read!(fid,tmp)
      tmp = hton.(tmp)
      close(fid)

      #2) compute export rate
      xprt=similar(tmp);
      for kk=1:siz[3]-1
         tmp0=Wsink*max.(tmp[:,:,kk],0.0)
         tmp1=Wsink*max.(tmp[:,:,kk],0.0)
         tmp2=isfinite.(tmp[:,:,kk+1])
         tmp1[.!tmp2].=NaN
         xprt[:,:,kk]=tmp1
      end
      xprt[:,:,siz[3]].=NaN

      #println(fil)
      filOut=dirOut*task["OutputFile"]*"/"
      !isdir(filOut) ? mkdir(filOut) : nothing
      filOut=filOut*replace(filList[ff],task["InputFile"][1] => task["OutputFile"])
      f =  FortranFile(filOut,"w",access="direct",recl=recl,convert="big-endian")
      write(f,rec=1,Float32.(xprt))
      close(f)
      #to re-read file:
      #f =  FortranFile(filOut,"r",access="direct",recl=recl,convert="big-endian");
      #tmp1=read(f,rec=1,(Float32,(720,360,50))); close(f);
   end

   return filList
end

"""
    loop_task4(indx)

Loop over model output files and time average them
as a postprocessing step, and write the result
to file (one subfolder for each variable)
"""
function loop_task4(indx)

   task=YAML.load(open("task.yml"))

   dirIn=task["InputDir"]
   filIn=task["InputFile"][1]
   dirOut=task["OutputDir"]
   !isdir(dirOut) ? mkdir(dirOut) : nothing

   siz=Tuple(task["InputSize"])
   prec=Float32

   tmp1=readdir(dirIn*filIn)
   tmp1=filter(x -> occursin(filIn,x),tmp1)
   filList=filter(x -> occursin(".data",x),tmp1)

   #maximum(indx)>length(filList) ? error("missing files: "*filIn*"*") : nothing
   #filList=filList[indx]

   !isa(filList,Array) ? filList=[filList] : nothing
   nf=length(filList)

   YearStart=task["Specs"]["YearStart"]
   YearEnd=task["Specs"]["YearEnd"]
   AverageType=task["Specs"]["AverageType"]

   nrec=missing
   AverageType=="daily" ? nrec=366 : nothing
   AverageType=="monthly" ? nrec=12 : nothing

   dayList=collect(1:nrec)
   dayList=dayList[indx]
   tmpList=Array{Array,1}(undef,nrec)

   for ff=dayList

      #identify all relevant records
      tmp=fill(false,length(filList))
      for ii=1:length(filList)
         if AverageType=="daily"
            tt=DateTime(1992,1,1,12)+Dates.Day(ii-1)
            dd=Day(Dates.Day(1)+tt-DateTime(year(tt),1,1,12))
            year(tt)>=YearStart && year(tt)<=YearEnd && dd.value==ff ? tmp[ii]=true : nothing
         end
         if AverageType=="monthly"
            tt=DateTime(1992,1,15)+Dates.Month(ii-1)
            year(tt)>=YearStart && year(tt)<=YearEnd && Month(tt).value==ff ? tmp[ii]=true : nothing
         end
      end
      recList=findall(tmp)
      tmpList[ff]=recList
      isempty(recList) ? error("no record found to average") : nothing

      println(length(recList))

      #compute time average
      ave = fill(0.0,siz)
      fac = inv(length(recList))
      for ii=recList
         fil=dirIn*filIn*"/"*filList[ii]
         #println(fil)
         tmp=Array{prec,length(siz)}(undef,siz)
         fid = open(fil)
         read!(fid,tmp)
         tmp = hton.(tmp)
         close(fid)
         ave=ave+fac*tmp
      end

      #output to file
      recl=prod(siz)*4
      filOut=dirOut*filIn*"Clim/"
      !isdir(filOut) ? mkdir(filOut) : nothing
      gg=recList[1]
      filOut=filOut*replace(filList[gg],filIn => filIn*"Clim")
      #println(filOut)
      f =  FortranFile(filOut,"w",access="direct",recl=recl,convert="big-endian")
      write(f,rec=1,prec.(ave))
      close(f)
      #to re-read file:
      #f =  FortranFile(filOut,"r",access="direct",recl=recl,convert="big-endian");
      #tmp1=read(f,rec=1,(Float32,(720,360,50))); close(f);

   end

   return tmpList
end

max_issues_repo_name: darwinproject/CBIOMES-Processing-legacy

max_forks_repo_name: darwinproject/CBIOMES-Processing-legacy


text:
"""
    predict_class_prior(X, α)

Predicts the latent class probabilities,
given covariates `X`
and logit regression coefficients `α`.

# Arguments
- `X`: A matrix of covariates.
- `α`: A matrix of logit regression coefficients.

# Return Values
- `prob`: A matrix of latent class probabilities.
- `max_prob_idx`: A matrix of the most likely latent class for each observation.
"""
function predict_class_prior(X, α)
    tmp = exp.(LogitGating(α, X))
    return (prob = tmp, max_prob_idx = [findmax(tmp[i,:])[2] for i in 1:size(tmp)[1]])
end

"""
    predict_class_posterior(Y, X, α, model;
        exact_Y = true, exposure_past = nothing)

Predicts the latent class probabilities,
given observations `Y`, covariates `X`,
logit regression coefficients `α` and a specified `model` of expert functions.

# Arguments
- `Y`: A matrix of responses.
- `X`: A matrix of covariates.
- `α`: A matrix of logit regression coefficients.
- `model`: A matrix specifying the expert functions.

# Optional Arguments
- `exact_Y`: `true` or `false` (default), indicating if `Y` is observed exactly or with censoring and truncation.
- `exposure_past`: A vector indicating the time exposure (past) of each observation. If nothing is supplied, it is set to 1.0 by default.

# Return Values
- `prob`: A matrix of latent class probabilities.
- `max_prob_idx`: A matrix of the most likely latent class for each observation.
"""
function predict_class_posterior(Y, X, α, model; exact_Y = true, exposure_past = nothing)
    if exact_Y == true
        Y = _exact_to_full(Y)
    end

    if isnothing(exposure_past)
        exposure_past = fill(1.0, size(X)[1])
    end

    model_exp = exposurize_model(model, exposure = exposure_past)

    gate = LogitGating(α, X)
    ll_np_list = loglik_np(Y, gate, model_exp)
    z_e_obs = EM_E_z_obs(ll_np_list.gate_expert_ll_comp, ll_np_list.gate_expert_ll)
    return (prob = z_e_obs, max_prob_idx = [findmax(z_e_obs[i,:])[2] for i in 1:size(z_e_obs)[1]])
end

"""
    predict_mean_prior(X, α, model;
        exposure_future = nothing)

Predicts the mean values of response,
given covariates `X`,
logit regression coefficients `α` and a specified `model` of expert functions.

# Arguments
- `X`: A matrix of covariates.
- `α`: A matrix of logit regression coefficients.
- `model`: A matrix specifying the expert functions.

# Optional Arguments
- `exposure_future`: A vector indicating the time exposure (future) of each observation. If nothing is supplied, it is set to 1.0 by default.

# Return Values
- A matrix of predicted mean values of response, based on prior probabilities.
"""
function predict_mean_prior(X, α, model; exposure_future = nothing)
    if isnothing(exposure_future)
        exposure_future = fill(1.0, size(X)[1])
    end

    model_exp = exposurize_model(model, exposure = exposure_future)

    weights = predict_class_prior(X, α).prob
    means = mean.(model_exp)

    result = fill(NaN, size(X)[1], size(model)[1])

    for i in 1:size(X)[1]
        result[i,:] = means[:,:,i] * weights[i,:]
    end

    return result
end

"""
    predict_mean_posterior(Y, X, α, model;
        exact_Y = true, exposure_past = nothing, exposure_future = nothing)

Predicts the mean values of response,
given observations `Y`, covariates `X`,
logit regression coefficients `α` and a specified `model` of expert functions.

# Arguments
- `Y`: A matrix of responses.
- `X`: A matrix of covariates.
- `α`: A matrix of logit regression coefficients.
- `model`: A matrix specifying the expert functions.

# Optional Arguments
- `exact_Y`: `true` or `false` (default), indicating if `Y` is observed exactly or with censoring and truncation.
- `exposure_past`: A vector indicating the time exposure (past) of each observation. If nothing is supplied, it is set to 1.0 by default.
- `exposure_future`: A vector indicating the time exposure (future) of each observation. If nothing is supplied, it is set to 1.0 by default.

# Return Values
- A matrix of predicted mean values of response, based on posterior probabilities.
"""
function predict_mean_posterior(Y, X, α, model; exact_Y = true, exposure_past = nothing, exposure_future = nothing)
    # if exact_Y == true
    #     Y = _exact_to_full(Y)
    # end

    if isnothing(exposure_past)
        exposure_past = fill(1.0, size(X)[1])
    end

    if isnothing(exposure_future)
        exposure_future = fill(1.0, size(X)[1])
    end

    model_exp = exposurize_model(model, exposure = exposure_future)

    weights = predict_class_posterior(Y, X, α, model, exact_Y = exact_Y, exposure_past = exposure_past).prob
    means = mean.(model_exp)

    result = fill(NaN, size(X)[1], size(model)[1])

    for i in 1:size(X)[1]
        result[i,:] = means[:,:,i] * weights[i,:]
    end

    return result
end

"""
    predict_var_prior(X, α, model;
        exposure_future = nothing)

Predicts the variance of response,
given covariates `X`,
logit regression coefficients `α` and a specified `model` of expert functions.

# Arguments
- `X`: A matrix of covariates.
- `α`: A matrix of logit regression coefficients.
- `model`: A matrix specifying the expert functions.

# Optional Arguments
- `exposure_future`: A vector indicating the time exposure of each observation. If nothing is supplied, it is set to 1.0 by default.

# Return Values
- A matrix of predicted variance of response, based on prior probabilities.
"""
function predict_var_prior(X, α, model; exposure_future = nothing)
    if isnothing(exposure_future)
        exposure_future = fill(1.0, size(X)[1])
    end

    model_exp = exposurize_model(model, exposure = exposure_future)

    weights = predict_class_prior(X, α).prob

    c_mean = mean.(model_exp)
    g_mean = predict_mean_prior(X, α, model, exposure_future = exposure_future)
    c_var = var.(model_exp)

    var_c_mean = fill(NaN, size(X)[1], size(model)[1])
    mean_c_var = fill(NaN, size(X)[1], size(model)[1])

    for i in 1:size(X)[1]
        var_c_mean[i,:] = (c_mean[:,:,i] .- g_mean[i,:] ).^2 * weights[i,:]
        mean_c_var[i,:] = c_var[:,:,i] * weights[i,:]
    end

    return var_c_mean + mean_c_var
end

"""
    predict_var_posterior(Y, X, α, model;
        exact_Y = true, exposure_past = nothing, exposure_future = nothing)

Predicts the variance of response,
given observations `Y`, covariates `X`,
logit regression coefficients `α` and a specified `model` of expert functions.

# Arguments
- `Y`: A matrix of responses.
- `X`: A matrix of covariates.
- `α`: A matrix of logit regression coefficients.
- `model`: A matrix specifying the expert functions.

# Optional Arguments
- `exact_Y`: `true` or `false` (default), indicating if `Y` is observed exactly or with censoring and truncation.
- `exposure_past`: A vector indicating the time exposure (past) of each observation. If nothing is supplied, it is set to 1.0 by default.
- `exposure_future`: A vector indicating the time exposure (future) of each observation. If nothing is supplied, it is set to 1.0 by default.

# Return Values
- A matrix of predicted variance of response, based on posterior probabilities.
"""
function predict_var_posterior(Y, X, α, model; exact_Y = true, exposure_past = nothing, exposure_future = nothing)
    # if exact_Y == true
    #     Y = _exact_to_full(Y)
    # end

    if isnothing(exposure_past)
        exposure_past = fill(1.0, size(X)[1])
    end

    if isnothing(exposure_future)
        exposure_future = fill(1.0, size(X)[1])
    end

    model_exp = exposurize_model(model, exposure = exposure_future)

    weights = predict_class_posterior(Y, X, α, model, exact_Y = exact_Y, exposure_past = exposure_past).prob

    c_mean = mean.(model_exp)
    g_mean = predict_mean_posterior(Y, X, α, model, exact_Y = exact_Y, exposure_past = exposure_past, exposure_future = exposure_future)
    c_var = var.(model_exp)

    var_c_mean = fill(NaN, size(X)[1], size(model)[1])
    mean_c_var = fill(NaN, size(X)[1], size(model)[1])

    for i in 1:size(X)[1]
        var_c_mean[i,:] = (c_mean[:,:,i] .- g_mean[i,:] ).^2 * weights[i,:]
        mean_c_var[i,:] = c_var[:,:,i] * weights[i,:]
    end

    return var_c_mean + mean_c_var
end

"""
    predict_limit_prior(X, α, model, limit;
        exposure_future = nothing)

Predicts the limit expected value (LEV) of response,
given covariates `X`,
logit regression coefficients `α` and a specified `model` of expert functions.

# Arguments
- `X`: A matrix of covariates.
- `α`: A matrix of logit regression coefficients.
- `model`: A matrix specifying the expert functions.
- `limit`: A matrix specifying the cutoff point.

# Optional Arguments
- `exposure_future`: A vector indicating the time exposure (future) of each observation. If nothing is supplied, it is set to 1.0 by default.

# Return Values
- A matrix of predicted limit expected value of response, based on prior probabilities.
"""
function predict_limit_prior(X, α, model, limit; exposure_future = nothing)
    if isnothing(exposure_future)
        exposure_future = fill(1.0, size(X)[1])
    end

    model_exp = exposurize_model(model, exposure = exposure_future)

    weights = predict_class_prior(X, α).prob

    result = fill(NaN, size(X)[1], size(model)[1])

    for i in 1:size(X)[1]
        means = vcat([hcat([lev(model_exp[d,j,i], limit[i,d]) for j in 1:size(model_exp)[2]]...) for d in 1:size(model_exp)[1]]...)
        result[i,:] = means * weights[i,:]
    end

    return result
end

"""
    predict_limit_posterior(Y, X, α, model, limit;
        exact_Y = true, exposure_past = nothing, exposure_future = nothing)

Predicts the limit expected value (LEV) of response,
given observations `Y`, covariates `X`,
logit regression coefficients `α` and a specified `model` of expert functions.

# Arguments
- `Y`: A matrix of responses.
- `X`: A matrix of covariates.
- `α`: A matrix of logit regression coefficients.
- `model`: A matrix specifying the expert functions.
- `limit`: A vector specifying the cutoff point.

# Optional Arguments
- `exact_Y`: `true` or `false` (default), indicating if `Y` is observed exactly or with censoring and truncation.
- `exposure_past`: A vector indicating the time exposure (past) of each observation. If nothing is supplied, it is set to 1.0 by default.
- `exposure_future`: A vector indicating the time exposure (future) of each observation. If nothing is supplied, it is set to 1.0 by default.

# Return Values
- A matrix of predicted limit expected value of response, based on posterior probabilities.
"""
function predict_limit_posterior(Y, X, α, model, limit; exact_Y = true, exposure_past = nothing, exposure_future = nothing)
    if isnothing(exposure_past)
        exposure_past = fill(1.0, size(X)[1])
    end

    if isnothing(exposure_future)
        exposure_future = fill(1.0, size(X)[1])
    end

    model_exp = exposurize_model(model, exposure = exposure_future)

    weights = predict_class_posterior(Y, X, α, model, exact_Y = exact_Y, exposure_past = exposure_past).prob

    result = fill(NaN, size(X)[1], size(model)[1])

    for i in 1:size(X)[1]
        means = vcat([hcat([lev(model_exp[d,j,i], limit[i,d]) for j in 1:size(model_exp)[2]]...) for d in 1:size(model_exp)[1]]...)
        result[i,:] = means * weights[i,:]
    end

    return result
end

"""
    predict_excess_prior(X, α, model, limit;
        exposure_future = nothing)

Predicts the excess expectation of response,
given covariates `X`,
logit regression coefficients `α` and a specified `model` of expert functions.

# Arguments
- `X`: A matrix of covariates.
- `α`: A matrix of logit regression coefficients.
- `model`: A matrix specifying the expert functions.
- `limit`: A vector specifying the cutoff point.

# Optional Arguments
- `exposure_future`: A vector indicating the time exposure (future) of each observation. If nothing is supplied, it is set to 1.0 by default.

# Return Values
- A matrix of predicted excess expectation of response, based on prior probabilities.
"""
function predict_excess_prior(X, α, model, limit; exposure_future = nothing)
    if isnothing(exposure_future)
        exposure_future = fill(1.0, size(X)[1])
    end

    model_exp = exposurize_model(model, exposure = exposure_future)

    weights = predict_class_prior(X, α).prob

    result = fill(NaN, size(X)[1], size(model)[1])

    for i in 1:size(X)[1]
        means = vcat([hcat([excess(model_exp[d,j,i], limit[i,d]) for j in 1:size(model_exp)[2]]...) for d in 1:size(model_exp)[1]]...)
        result[i,:] = means * weights[i,:]
    end

    return result
end

"""
    predict_excess_posterior(Y, X, α, model, limit;
        exact_Y = true, exposure_past = nothing, exposure_future = nothing)

Predicts the excess expectation of response,
given observations `Y`, covariates `X`,
logit regression coefficients `α` and a specified `model` of expert functions.

# Arguments
- `Y`: A matrix of responses.
- `X`: A matrix of covariates.
- `α`: A matrix of logit regression coefficients.
- `model`: A matrix specifying the expert functions.
- `limit`: A vector specifying the cutoff point.

# Optional Arguments
- `exact_Y`: `true` or `false` (default), indicating if `Y` is observed exactly or with censoring and truncation.
- `exposure_past`: A vector indicating the time exposure (past) of each observation. If nothing is supplied, it is set to 1.0 by default.
- `exposure_future`: A vector indicating the time exposure (future) of each observation. If nothing is supplied, it is set to 1.0 by default.

# Return Values
- A matrix of predicted excess expectation of response, based on posterior probabilities.
"""
function predict_excess_posterior(Y, X, α, model, limit; exact_Y = true, exposure_past = nothing, exposure_future = nothing)
    if isnothing(exposure_past)
        exposure_past = fill(1.0, size(X)[1])
    end

    if isnothing(exposure_future)
        exposure_future = fill(1.0, size(X)[1])
    end

    model_exp = exposurize_model(model, exposure = exposure_future)

    weights = predict_class_posterior(Y, X, α, model, exact_Y = exact_Y, exposure_past = exposure_past).prob

    result = fill(NaN, size(X)[1], size(model)[1])

    for i in 1:size(X)[1]
        means = vcat([hcat([excess(model_exp[d,j,i], limit[i,d]) for j in 1:size(model_exp)[2]]...) for d in 1:size(model_exp)[1]]...)
        result[i,:] = means * weights[i,:]
    end

    return result
end

# solve a quantile of a mixture model
# Bisection method seems to give the most stable results
function _solve_continuous_mix_quantile(weights, experts, p)
    p0 = sum(weights .* exp.(expert_ll.(experts, 0.0, 0.0, 0.0, Inf)))
    if p <= p0
        return 0.0
    else
        # init_guess = minimum([maximum(quantile.(experts, 0.90)) 500])
        # init_guess = maximum([maximum(quantile.(experts, 0.90)) 1000])
        init_guess = maximum([maximum(quantile.(experts, p)) 1000])
        VaR = try
            # Roots.find_zero(y -> sum(weights .* exp.(expert_ll.(experts, 0.0, 0.0, y, Inf))) - p, init_guess, Roots.Order2())
            Roots.find_zero(y -> sum(weights .* exp.(expert_ll.(experts, 0.0, 0.0, y, Inf))) - p, (0.0, init_guess+500))
        catch;
            NaN
        end
        return VaR
    end
end

function _calc_continuous_CTE(weights, experts, p, VaR)
    m = sum(vec(weights) .* vec(mean.(experts)))
    lim_ev = sum(vec(weights) .* vec(lev.(experts, VaR))) # [lev(model[k], VaR) for k in 1:length(model)]
    return VaR + (m-lim_ev)/(1-p)
end

# calculate CTE based on VaR and p
# function _calc_continuous_mix_CTE(weights, experts, p, means, VaR)
#     return
# end

"""
    predict_VaRCTE_prior(X, α, model, p;
        exposure_future = nothing)

Predicts the `p`-th value-at-risk (VaR) and conditional tail expectation (CTE) of response,
given covariates `X`,
logit regression coefficients `α` and a specified `model` of expert functions.

# Arguments
- `X`: A matrix of covariates.
- `α`: A matrix of logit regression coefficients.
- `model`: A matrix specifying the expert functions.
- `p`: A matrix of probabilities.

# Optional Arguments
- `exposure_future`: A vector indicating the time exposure (future) of each observation. If nothing is supplied, it is set to 1.0 by default.

# Return Values
- `VaR`: A matrix of predicted VaR of response, based on prior probabilities.
- `CTE`: A matrix of predicted CTE of response, based on prior probabilities.
"""
function predict_VaRCTE_prior(X, α, model, p; exposure_future = nothing)

    if isnothing(exposure_future)
        exposure_future = fill(1.0, size(X)[1])
    end

    model_exp = exposurize_model(model, exposure = exposure_future)

    weights = predict_class_prior(X, α).prob

    VaR = fill(NaN, size(X)[1], size(model)[1])
    CTE = fill(NaN, size(X)[1], size(model)[1])
    for i in 1:size(X)[1]
        for k in 1:size(model)[1]
            VaR[i,k] = _solve_continuous_mix_quantile(weights[i,:], model_exp[k,:,i], p[i,k])
            CTE[i,k] = _calc_continuous_CTE(weights[i,:], model_exp[k,:,i], p[i,k], VaR[i,k])
        end
    end
    # return VaR

    # VaR = vcat([hcat([_solve_continuous_mix_quantile(weights[i,:], model[k,:], p) for k in 1:size(model)[1] ]...) for i in 1:size(X)[1]]...)
    # CTE = vcat([hcat([_calc_continuous_CTE(weights[i,:], model[k,:], p, VaR[i,k]) for k in 1:size(model)[1] ]...) for i in 1:size(X)[1]]...)
    return (VaR = VaR, CTE = CTE)
end

"""
    predict_VaRCTE_posterior(Y, X, α, model, p;
        exact_Y = true, exposure_past = nothing, exposure_future = nothing)

Predicts the `p`-th value-at-risk (VaR) and conditional tail expectation (CTE) of response,
given observations `Y`, covariates `X`,
logit regression coefficients `α` and a specified `model` of expert functions.

# Arguments
- `X`: A matrix of covariates.
- `α`: A matrix of logit regression coefficients.
- `model`: A matrix specifying the expert functions.
- `p`: A matrix of probabilities.

# Optional Arguments
- `exact_Y`: `true` or `false` (default), indicating if `Y` is observed exactly or with censoring and truncation.
- `exposure_past`: A vector indicating the time exposure (past) of each observation. If nothing is supplied, it is set to 1.0 by default.
- `exposure_future`: A vector indicating the time exposure (future) of each observation. If nothing is supplied, it is set to 1.0 by default.

# Return Values
- `VaR`: A matrix of predicted VaR of response, based on posterior probabilities.
- `CTE`: A matrix of predicted CTE of response, based on posterior probabilities.
"""
function predict_VaRCTE_posterior(Y, X, α, model, p; exact_Y = true, exposure_past = nothing, exposure_future = nothing)
    if isnothing(exposure_past)
        exposure_past = fill(1.0, size(X)[1])
    end

    if isnothing(exposure_future)
        exposure_future = fill(1.0, size(X)[1])
    end

    model_exp = exposurize_model(model, exposure = exposure_future)

    weights = predict_class_posterior(Y, X, α, model, exact_Y = exact_Y, exposure_past = exposure_past).prob

    VaR = fill(NaN, size(X)[1], size(model)[1])
    CTE = fill(NaN, size(X)[1], size(model)[1])
    for i in 1:size(X)[1]
        for k in 1:size(model)[1]
            VaR[i,k] = _solve_continuous_mix_quantile(weights[i,:], model_exp[k,:,i], p[i,k])
            CTE[i,k] = _calc_continuous_CTE(weights[i,:], model_exp[k,:,i], p[i,k], VaR[i,k])
        end
    end
    # return VaR

    # VaR = vcat([hcat([_solve_continuous_mix_quantile(weights[i,:], model[k,:], p) for k in 1:size(model)[1] ]...) for i in 1:size(X)[1]]...)
    # CTE = vcat([hcat([_calc_continuous_CTE(weights[i,:], model[k,:], p, VaR[i,k]) for k in 1:size(model)[1] ]...) for i in 1:size(X)[1]]...)
    return (VaR = VaR, CTE = CTE)
end

max_stars_repo_name: sparktseung/JuliaLRMoE

max_issues_repo_name: sparktseung/JuliaLRMoE

max_forks_repo_name: sparktseung/JuliaLRMoE


text:
const DimOrDimTuple = Union{Dimension,Tuple{Vararg{<:Dimension}}}
const IntOrIntTuple = Union{Int,Tuple{Vararg{<:Int}}}

struct Ag end
struct DisAg end

"""
    aggregate(method, object, scale; filename, progress, skipmissing)

Aggregate a Raster, or all arrays in a RasterStack or RasterSeries, by `scale` using
`method`.

# Arguments

- `method`: a function such as `mean` or `sum` that can combine the
  value of multiple cells to generate the aggregated cell, or a [`Locus`]($DDlocusdocs)
  like `Start()` or `Center()` that specifies where to sample from in the interval.
- `object`: Object to aggregate, like `AbstractRasterSeries`, `AbstractStack`,
  `AbstractRaster` or `Dimension`.
- `scale`: the aggregation factor, which can be an integer, a tuple of integers
  for each dimension, or any `Dimension`, `Selector` or `Int` combination you can
  usually use in `getindex`. Using a `Selector` will determine the scale by the
  distance from the start of the index.

When the aggregation `scale` of is larger than the array axis, the length of the axis is used.

# Keywords

- `progress`: show a progress bar.
- `skipmissingval`: if `true`, any `missingval` will be skipped during aggregation, so that
    only areas of all missing values will be aggregated to `missingval`. If `false`, any
    aggegrated area containing a `missingval` will be assigned `missingval`.
- `filename`: a filename to write to directly, useful for large files.
- `suffix`: a string or value to append to the filename.
    A tuple of `suffix` will be applied to stack layers. `keys(st)` are the default.

# Example

```jldoctest
using Rasters, Statistics, Plots
using Rasters: Center
st = read(RasterStack(WorldClim{Climate}; month=1))
ag = aggregate(Center(), st, (Y(20), X(20)); skipmissingval=true, progress=false)
plot(ag)
savefig("build/aggregate_example.png")
# output

```

![aggregate](aggregate_example.png)

Note: currently it is faster to aggregate over memory-backed arrays.
Use [`read`](@ref) on `src` before use where required.
"""
function aggregate end
function aggregate(method, series::AbstractRasterSeries, scale, args...;
    progress=true, kw...
)
    f(A) = aggregate(method, A, scale, args...; progress=false, kw...)
    if progress
        ProgressMeter.@showprogress "Aggregating series..." map(f, series)
    else
        map(f, series)
    end
end
function aggregate(method, stack::AbstractRasterStack, scale;
    keys=keys(stack), filename=nothing, suffix=keys, progress=true, kw...
)
    f(A, suffix) = aggregate(method, A, scale; filename, suffix, kw...)

    layers = if progress
        ProgressMeter.@showprogress "Aggregating stack..." map(f, values(stack), Tuple(suffix))
    else
        map(f, values(stack), Tuple(suffix))
    end

    return DD.rebuild_from_arrays(stack, layers)
end
function aggregate(method, src::AbstractRaster, scale;
    suffix=nothing, filename=nothing, kw...
)
    dst = alloc_ag(method, src, scale; filename, suffix, kw...)
    aggregate!(method, dst, src, scale; kw...)
end
aggregate(method, d::Dimension, scale) = rebuild(d, aggregate(method, lookup(d), scale))
function aggregate(method, lookup::LookupArray, scale)
    intscale = _scale2int(Ag(), lookup, scale)
    intscale == 1 && return lookup
    start, stop = _endpoints(method, lookup, intscale)
    newlookup = lookup[start:scale:stop]
    if lookup isa AbstractSampled
        sp = aggregate(method, span(lookup), scale)
        return rebuild(newlookup; span=sp)
    else
        return newlookup
    end
end
aggregate(method, span::Span, scale) = span
aggregate(method, span::Regular, scale) = Regular(val(span) * scale)

"""
    aggregate!(method, dst::AbstractRaster, src::AbstractRaster, scale; skipmissingval=false)

Aggregate array `src` to array `dst` by `scale`, using `method`.

# Arguments

- `method`: a function such as `mean` or `sum` that can combine the
  value of multiple cells to generate the aggregated cell, or a [`Locus`]($DDlocusdocs)
  like `Start()` or `Center()` that species where to sample from in the interval.
- `scale`: the aggregation factor, which can be an integer, a tuple of integers
  for each dimension, or any `Dimension`, `Selector` or `Int` combination you can
  usually use in `getindex`. Using a `Selector` will determine the scale by the
  distance from the start of the index in the `src` array.

When the aggregation `scale` of is larger than the array axis, the length of the axis is used.

# Keywords

- `progress`: show a progress bar.
- `skipmissingval`: if `true`, any `missingval` will be skipped during aggregation, so that
    only areas of all missing values will be aggregated to `missingval`. If `false`, any
    aggegrated area containing a `missingval` will be assigned `missingval`.

Note: currently it is _much_ faster to aggregate over memory-backed arrays.
Use [`read`](@ref) on `src` before use where required.
"""
function aggregate!(locus::Locus, dst::AbstractRaster, src, scale; kw...)
    aggregate!((locus,), dst, src, scale)
end
function aggregate!(loci::Tuple{Locus,Vararg}, dst::AbstractRaster, src, scale; kw...)
    intscale = _scale2int(Ag(), dims(src), scale)
    offsets = _agoffset.(loci, intscale)
    broadcast!(dst, CartesianIndices(dst)) do I
        val = src[(upsample.(Tuple(I), intscale) .+ offsets)...]
        val === missingval(src) ? missingval(dst) : val
    end
end
# Function/functor methods
function aggregate!(f, dst::AbstractRaster, src, scale; skipmissingval=false)
    intscale = _scale2int(Ag(), dims(src), scale)
    broadcast!(dst, CartesianIndices(dst)) do I
        upper = upsample.(Tuple(I), intscale)
        lower = upper .+ intscale .- 1
        block = if isdisk(src)
            src[map(:, upper, lower)...]
        else
            view(src, map(:, upper, lower)...)
        end
        if skipmissingval
            # All missing values return a missing value
            if all(map(x -> x === missingval(src), block))
                _missingval_or_missing(dst)
            else
                # Skip missing values
                f((x for x in block if x !== missingval(src)))
            end
        else
            # Any missing values return a missing value
            if any(map(x -> x === missingval(src), block))
                _missingval_or_missing(dst)
            else
                f(block)
            end
        end
    end
end

"""
    disaggregate(method, object, scale; filename, progress, keys)

Disaggregate array, or all arrays in a stack or series, by some scale.

# Arguments

- `method`: a function such as `mean` or `sum` that can combine the
  value of multiple cells to generate the aggregated cell, or a [`Locus`]($DDlocusdocs)
  like `Start()` or `Center()` that species where to sample from in the interval.
- `object`: Object to aggregate, like `AbstractRasterSeries`, `AbstractStack`,
  `AbstractRaster` or a `Dimension`.
- `scale`: the aggregation factor, which can be an integer, a tuple of integers
  for each dimension, or any `Dimension`, `Selector` or `Int` combination you can
  usually use in `getindex`. Using a `Selector` will determine the scale by the
  distance from the start of the index.

# Keywords

- `progress`: show a progress bar.

Note: currently it is faster to aggregate over memory-backed arrays.
Use [`read`](@ref) on `src` before use where required.

"""
function disaggregate end
function disaggregate(method, series::AbstractRasterSeries, scale; progress=true, kw...)
    f = i -> disaggregate(method, series[i], scale; progress=false, kw...)
    return if progress
        ProgressMeter.@showprogress "Disaggregating series..." map(f, 1:length(series))
    else
        map(f, 1:length(series))
    end
end
function disaggregate(method, stack::AbstractRasterStack, scale;
    keys=keys(stack), suffix=keys, filename=nothing, progress=true
)
    f(A, suffix) = disaggregate(method, A, scale; filename, suffix)

    layers = if progress
        ProgressMeter.@showprogress "Disaggregating stack..." map(f, values(stack), Tuple(suffix))
    else
        map(f, values(stack), Tuple(suffix))
    end
    return DD.rebuild_from_arrays(stack, layers)
end
function disaggregate(method, src::AbstractRaster, scale;
    suffix=nothing, filename=nothing, kw...
)
    dst = alloc_disag(method, src, scale; filename, suffix, kw...)
    disaggregate!(method, dst, src, scale)
end
function disaggregate(locus::Locus, dim::Dimension, scale)
    rebuild(dim, disaggregate(locus, lookup(dim), scale))
end
function disaggregate(locus, lookup::LookupArray, scale)
    intscale = _scale2int(DisAg(), lookup, scale)
    intscale == 1 && return lookup

    len = length(lookup) * intscale
    step_ = step(lookup) / intscale
    start = lookup[1] - _agoffset(locus, intscale) * step_
    stop = start + (len - 1)  * step_
    index = LinRange(start, stop, len)
    if lookup isa AbstractSampled
        sp = disaggregate(locus, span(lookup), intscale)
        return rebuild(lookup; data=index, span=sp)
    else
        return rebuild(lookup; data=index)
    end
end

disaggregate(method, span::Span, scale) = span
disaggregate(method, span::Regular, scale) = Regular(val(span) / scale)

"""
    disaggregate!(method, dst::AbstractRaster, src::AbstractRaster, filename, scale)

Disaggregate array `src` to array `dst` by some scale, using `method`.

- `method`: a function such as `mean` or `sum` that can combine the
  value of multiple cells to generate the aggregated cell, or a [`Locus`]($DDlocusdocs)
  like `Start()` or `Center()` that species where to sample from in the interval.
- `scale`: the aggregation factor, which can be an integer, a tuple of integers
  for each dimension, or any `Dimension`, `Selector` or `Int` combination you can
  usually use in `getindex`. Using a `Selector` will determine the scale by the
  distance from the start of the index in the `src` array.

Note: currently it is faster to aggregate over memory-backed arrays.
Use [`read`](@ref) on `src` before use where required.
"""
function disaggregate!(locus::Locus, dst::AbstractRaster, src, scale)
    disaggregate!((locus,), dst, src, scale)
end
function disaggregate!(loci::Tuple{Locus,Vararg}, dst::AbstractRaster, src, scale)
    intscale = _scale2int(DisAg(), dims(src), scale)
    broadcast!(dst, CartesianIndices(dst)) do I
        val = src[(downsample.(Tuple(I), intscale))...]
        val === missingval(src) ? missingval(dst) : val
    end
end

# Allocate an array of the correct size to aggregate `A` by `scale`
alloc_ag(method, A::AbstractRaster, scale; kw...) = alloc_ag((method,), A, scale; kw...)
function alloc_ag(method::Tuple, A::AbstractRaster, scale;
    filename=nothing, suffix=nothing, skipmissingval=nothing
)
    intscale = _scale2int(Ag(), dims(A), scale)
    # Aggregate the dimensions
    dims_ = aggregate.(method, dims(A), intscale)
    # Dim aggregation determines the array size
    sze = map(length, dims_)
    agT = ag_eltype(method, A)
    if missingval(A) isa Nothing
        T = agT
        mv = nothing
    else
        T = promote_type(agT, typeof(missingval(A)))
        mv = convert(T, missingval(A))
    end
    return create(filename, T, dims_; name=name(A), suffix, missingval=mv)
end

# Allocate an array of the correct size to disaggregate `A` by `scale`
function alloc_disag(method, A::AbstractRaster, scale; kw...)
    alloc_disag((method,), A, scale; kw...)
end
function alloc_disag(method::Tuple, A::AbstractRaster, scale;
    filename=nothing, suffix=nothing
)
    intscale = _scale2int(DisAg(), dims(A), scale)
    dims_ = disaggregate.(method, dims(A), intscale)
    # Dim aggregation determines the array size
    sze = map(length, dims_)
    T = ag_eltype(method, A)
    mv = convert(T, missingval(A))
    return create(filename, T, dims_; name=name(A), suffix, missingval=mv)
end

# Handle how methods like `mean` can change the type
ag_eltype(method::Tuple{<:Locus,Vararg}, A) = eltype(A)
function ag_eltype(method::Tuple{<:Any}, A)
    method_returntype = typeof(method[1](zero(eltype(A))))
    promote_type(eltype(A), method_returntype)
end

# Convert indicies from the aggregated array to the larger original array.
upsample(index::Int, scale::Int) = (index - 1) * scale + 1
upsample(index::Int, scale::Colon) = index

# Convert indicies from the original array to the aggregated array.
downsample(index::Int, scale::Int) = (index - 1) ÷ scale + 1
downsample(index::Int, scale::Colon) = index

# Convert scale or tuple of scale to integer using dims2indices
function _scale2int(x, dims::DimTuple, scale::Tuple)
    map((d, s) -> _scale2int(x, d, s), dims, DD.dims2indices(dims, scale))
end
_scale2int(x, dims::DimTuple, scale::Int) = map(d -> _scale2int(x, d, scale), dims)
_scale2int(x, dims::DimTuple, scale::Colon) = map(d -> _scale2int(x, d, scale), dims)
_scale2int(x, dim::Dimension, scale::Int) = _scale2int(x, lookup(dim), scale)
_scale2int(::Ag, l::LookupArray, scale::Int) = scale > length(l) ? length(l) : scale
_scale2int(::DisAg, l::LookupArray, scale::Int) = scale
_scale2int(x, dim::Dimension, scale::Colon) = 1

_agoffset(locus::Locus, l::LookupArray, scale) = _agoffset(locus, scale)
_agoffset(method, l::LookupArray, scale) = _agoffset(locus(l), scale)
_agoffset(locus::Start, scale) = 0
_agoffset(locus::End, scale) = scale - 1
_agoffset(locus::Center, scale) = scale ÷ 2

function _endpoints(method, l::LookupArray, scale)
    start = firstindex(l) + _agoffset(method, l, scale)
    stop = (length(l) ÷ scale) * scale
    return start, stop
end


text:
using Thrift
import Thrift.process, Thrift.meta

# include the generated hello module
include("gen-jl/hello/hello.jl");

# create a client instance with our choice of protocol and transport
clnt_transport = TSocket(19999)
proto = TBinaryProtocol(clnt_transport)
clnt = hello.SayHelloClient(proto)

# open a connection
open(clnt_transport)
# invoke service and print the result
println(hello.hello_to(clnt, "Julia"))
# close connection
close(clnt_transport)


text:
# Aiyagari Model Test
# Michele Fornino
#---------------------------------------------------#
#			Ramsey-Cass-Koopmans MODEL 				#
#---------------------------------------------------#

"""
	RCK_params{F<:Real, I<:Integer}

	Composite type containing all parameters of the standard Ramsey-Cass-Koopmans model.

"""
@with_kw struct RCK_params{F<:Real, I<:Integer} @deftype F

	# Parameters
	σ = 1.0
	n = 0.0
	α = 0.36
	δ = log(1 + 0.1)
	ρ = log(1 + 0.04)
	# σz = 0.2
	# θz = 0.1
	zd = 1.0

	# Derived Parameters
	γ = 1.0 / σ

	# Grids Parameters
	Nk::I = 2000
	Nz::I = 1
	# zGrid_lower_quantile = 1e-2
	# zGrid_upper_quantile = 1 - 1e-2
	kGrid_lower = 0.01
	# kGrid_upper = 20.
	kGrid_upper = 1.5 * (zd * α / (n + δ + ρ)) ^ (1 / (1 - α))

	# Assertions
	@assert σ > 0.0
	@assert n >= 0.0
	@assert δ > 0.0
	@assert ρ > 0.0
	@assert (α > 0.0) & (α < 1.0)
	# @assert θz > 0
	# @assert σz > 0
	@assert zd > 0

	# Check values for grids make sense
	@assert Nk > 1
	@assert Nz > 0
	# @assert zGrid_upper_quantile > zGrid_lower_quantile
	# @assert zGrid_lower_quantile > 0.0
	# @assert zGrid_upper_quantile < 1.0
	@assert kGrid_upper > kGrid_lower
	@assert kGrid_lower >= 0.

	# Enforce definitions (prevents accidental creation of object with wrong derived parameters)
	@assert γ == 1 / σ

end

"""
	function HamiltonJacobiBellmanEquations.HJBEquation(params::RCK_params)

Overloaded constructor for input of type RCK_params.

"""
function HamiltonJacobiBellmanEquations.HJBEquation(params::RCK_params)

	# Use UnPack.jl to localize all members of params
	@unpack_RCK_params params

	# Utility Function and derivatives.
	function u(c::T)::T where T<:Real
		if σ == 1.0
			return log(c)
		else
			return (c ^ (1 - γ) - 1) / (1 - γ)
		end
	end

	function u′(c::T)::T where T<:Real
		return c ^ -γ
	end

	function u′⁻¹(x::T)::T where T<:Real
		return x ^ (- 1 / γ)
	end

	function f(k::T, z::T)::T where T<:Real
		return z * (k ^ α)
	end

	# function u(c::T)::T where T<:Real
	# 	return - 1 / γ * exp(-γ * c)
	# end

	# function u′(c::T)::T where T<:Real
	# 	return exp( - γ * c)
	# end

	# function u′⁻¹(x::T)::T where T<:Real
	# 	return - 1 / γ * log(x)
	# end

	# Returns Function
	function flow_returns(state::AbstractVector{T}, c::Union{T, AbstractVector{T}})::T where T<:Real
		return u(c[1])
	end

	# Policy Function in term of marginal value function
	function optimal_controls(state::AbstractVector{T}, ∂V∂k::Union{T, AbstractVector{T}})::T where T<:Real
		return u′⁻¹(∂V∂k[1])
	end

	# State equations
	function state_equations(state::AbstractVector{T}, c::Union{T, AbstractVector{T}})::T where T<:Real
		return f(state[1], state[2]) - (n + δ) * state[1] - c[1]
	end

	# State constraint
	function state_constraints(state::AbstractVector{T})::T where T<:Real
		return u′(f(state[1], state[2]) - (n + δ) * state[1])
	end

	# Tuple of endogenous state grids
	kGrid = Vector(range(kGrid_lower, kGrid_upper, length = Nk))
	state_grids_endog = (kGrid, )

	# Exogenous shock
	exog_shock = ContinuousTimeMarkovChain(params)

	# Allocate hjb object
	return HJBEquation(flow_returns,
					   state_equations,
					   optimal_controls,
					   state_constraints,
					   state_grids_endog,
					   exog_shock,
					   ρ)

end

"""
	function MarkovChains.ContinuousTimeMarkovChain(params::RCK_params)

Overloaded constructor for input of type RCK_params.

"""
function MarkovChains.ContinuousTimeMarkovChain(params::RCK_params)

	@unpack_RCK_params params

	# Construct exog_shock ContinuousTimeMarkovChain object (if problem is stochastic).
	if Nz > 1

		# Obtain discretize diffusion process object as a MarkovChains.ContinuousTimeMarkovChain.
		μ = z -> -θz .* (z .- zd)
		σ = z -> sqrt(2 * θz * σz^2 / zd) .* sqrt.(z)
		dp = ItoDiffusionProcess(μ, σ)

		# Theoretical benchmark for CIR stationary distribution, used for zGrid.
		shape = zd^2 / σz^2
		rate = zd / shape
		Fz_theory = Gamma(shape, rate)
		zGrid_lower = quantile(Fz_theory, zGrid_lower_quantile)
		zGrid_upper = quantile(Fz_theory, zGrid_upper_quantile)

		# Allocate Grid
		zGrid = Vector(range(zGrid_lower, zGrid_upper, length = Nz))

		# Exogenous Shock: ContinuousTimeMarkovChain object
		exog_shock = ContinuousTimeMarkovChain(dp, zGrid)

	elseif Nz == 1

		exog_shock = ContinuousTimeMarkovChain(reshape([0], 1, 1), [zd])

	end

	return exog_shock

end


text:
# This file is a part of JuliaFEM.
# License is MIT: see https://github.com/JuliaFEM/JuliaFEM.jl/blob/master/LICENSE.md

"""
Problem u(X) = u₀ in Γ(d)
"""
mutable struct Dirichlet <: BoundaryProblem
    formulation :: Symbol
    variational :: Bool
    dual_basis :: Bool
    order :: Int
end

function Dirichlet()
    Dirichlet(:incremental, false, false, 1)
end

""" Return dual basis transformation matrix Ae. """
function get_dualbasis(element::Element, time::Float64, order=1)
    nnodes = length(element)
    De = zeros(nnodes, nnodes)
    Me = zeros(nnodes, nnodes)
    for ip in get_integration_points(element, order)
        detJ = element(ip, time, Val{:detJ})
        w = ip.weight*detJ
        N = element(ip, time)
        De += w*Matrix(Diagonal(vec(N)))
        Me += w*N'*N
    end
    return De, Me, De*inv(Me)
end

function get_formulation_type(problem::Problem{Dirichlet})
    return problem.properties.formulation
end

function assemble!(problem::Problem{Dirichlet}, time::Float64=0.0;
                   auto_initialize=true)
    # FIXME: boilerplate
    if !isempty(problem.assembly)
        @warn("Assemble problem $(problem.name): problem.assembly is not empty and assembling, are you sure you know what are you doing?")
    end
    if isempty(problem.elements)
        @warn("Assemble problem $(problem.name): problem.elements is empty, no elements in problem?")
    else
        first_element = first(problem.elements)
        unknown_field_name = get_unknown_field_name(problem)
        if !haskey(first_element, unknown_field_name)
            @warn("Assemble problem $(problem.name): seems that problem is uninitialized.")
            if auto_initialize
                @info("Initializing problem $(problem.name) at time $time automatically.")
                initialize!(problem, time)
            end
        end
    end

    if hasmethod(assemble_prehook!, Tuple{typeof(problem), Float64})
        assemble_prehook!(problem, time)
    end

    if problem.properties.variational
        for element in get_elements(problem)
            assemble!(problem.assembly, problem, element, time)
        end
    else # nodal collocation
        field_vals = Dict{Int64, Float64}()
        field_name = get_parent_field_name(problem)
        field_dim = get_unknown_field_dimension(problem)
        for element in get_elements(problem)
            gdofs = get_gdofs(problem, element)
            for i=1:field_dim
                haskey(element, field_name*" $i") || continue
                ldofs = gdofs[i:field_dim:end]
                xis = get_reference_coordinates(element)
                vals = Float64[]
                for xi in xis
                    g = element(field_name*" $i", xi, time)
                    # u = u_prev + Δu ⇒ Δu = u - u_prev
                    if haskey(element, field_name)
                        g_prev = element(field_name, xi, time)
                        g -= g_prev[i]
                    end
                    push!(vals, g)
                end
                for (dof, g) in zip(ldofs, vals)
                    field_vals[dof] = g
                end
            end
        end
        for (k, v) in field_vals
            FEMBase.add!(problem.assembly.C1, k, k, 1.0)
            FEMBase.add!(problem.assembly.C2, k, k, 1.0)
            FEMBase.add!(problem.assembly.g, k, 1, v)
        end
    end

    if hasmethod(assemble_posthook!, Tuple{typeof(problem), Float64})
        assemble_posthook!(problem, time)
    end
end

function assemble!(assembly::Assembly, problem::Problem{Dirichlet},
                   element::Element, time::Float64)

    # get dimension and name of PARENT field
    nnodes = length(element)
    field_dim = get_unknown_field_dimension(problem)
    field_name = get_parent_field_name(problem)
    gdofs = get_gdofs(problem, element)
    props = problem.properties

    if problem.properties.dual_basis
        De, Me, Ae = get_dualbasis(element, time)
    else
        Ae = I
        De = zeros(nnodes, nnodes)
        for ip in get_integration_points(element, props.order)
            N = element(ip, time)
            detJ = element(ip, time, Val{:detJ})
            De += ip.weight*N'*N*detJ
        end
    end

    # left hand side
    for i=1:field_dim
        ldofs = gdofs[i:field_dim:end]
        if haskey(element, field_name*" $i")
            add!(assembly.C1, ldofs, ldofs, De)
            add!(assembly.C2, ldofs, ldofs, De)
        end
    end

    # right hand side
    for ip in get_integration_points(element, props.order)
        detJ = element(ip, time, Val{:detJ})
        w = ip.weight*detJ
        N = element(ip, time)

        for i=1:field_dim
            ldofs = gdofs[i:field_dim:end]
            if haskey(element, field_name*" $i")
                g = element(field_name*" $i", ip, time)
                # u = u_prev + Δu ⇒ Δu = u - u_prev
                if haskey(element, field_name)
                    g_prev = element(field_name, ip, time)
                    g -= g_prev[i]
                end
                add!(assembly.g, ldofs, w*g*Ae*N')
            end
        end

    end

end

function postprocess!(problem::Problem{Dirichlet}, time::Float64, ::Type{Val{Symbol("reaction force")}})
    la = problem("lambda", time)
    rf = Dict(nid => -lai for (nid, lai) in la)
    update!(problem, "reaction force", time => rf)
end


text:
function cycle!(θ, mdl, obs)
    # action
    P = AnimalBehavior.observ(mdl, obs.s; θ...)
    # update
    AnimalBehavior.evol!(mdl, obs.s, obs.a, obs.r; θ...)
    return P
end

# Sample for a single session
function sample(mdl::Tm, data::StructVector, args...; kwargs...) where Tm <: AbstractMCMC.AbstractModel
    sample(Random.default_rng(), mdl, data, args...; kwargs...)
end

function sample(rng::AbstractRNG, mdl::Tm, data::StructVector, args...; kwargs...) where Tm <: AbstractMCMC.AbstractModel
    @model model(A) = begin
        θ = @submodel mdl
        θ = check_tuple_types(θ)

        P = [cycle!(θ, mdl, obs) for obs in data]
        A ~ arraydist(P)
        return
    end

    return sample(rng, model(data.a), args...; kwargs...)
end

# Sample for multiple sessions
function sample(mdl::Tm, data::Array{Ts}, args...; kwargs...) where {Tm <: AbstractMCMC.AbstractModel, Ts <: StructVector}
    sample(Random.default_rng(), mdl, data, args...; kwargs...)
end

function sample(rng::AbstractRNG, mdl::Tm, data::Array{Ts}, args...; kwargs...) where {Tm <: AbstractMCMC.AbstractModel, Ts <: StructVector}
    nsess = length(data)
    @model model(A) = begin
        θ = @submodel mdl
        θ = check_tuple_types(θ)

        for sess in 1:nsess
            sessdat = data[sess]
            θ_init = deepcopy(θ)
            P = [cycle!(θ_init, mdl, obs) for obs in sessdat]
            A[sess] ~ arraydist(P)
        end
        return
    end
    actions = [sessdat.a for sessdat in data]
    return sample(rng, model(actions), args...; kwargs...)
end

max_stars_repo_name: sqwayer/AnimalBehavior

max_issues_repo_name: sqwayer/AnimalBehavior


text:
module Functions

import Base.Broadcast: broadcast_shape
import Base.Iterators: flatten, isdone, repeated, Stateful
import LinearAlgebra: Diagonal, I
import OrderedCollections: OrderedDict

import TensorOperations.@tensor

import ..Transforms: TransformChain
import ..Elements: Element

# Base exports
export ast, callable, compile, isconstant, iselconstant, restype, typetree
export CompiledArray, CompiledSparseArray, Compiled

# Type exports
export Shape, Index, Indices
export Add, ApplyTransform, ApplyTransformGrad
export Constant, Contract, GetIndex
export Inflate, InsertAxis, Inv
export Monomials, Neg, Outer, Product, Reshape, Tupl, Zeros

# Constructor exports
export dimtrans, fulltrans, rawpoint, point
export grad, elemindex, rootcoords, outer

# Some typedefs
Shape = Tuple{Vararg{Int}}

include("Functions/bases.jl")

# More typedefs
Index = Union{ArrayEvaluable{Int}, Colon}
Indices = Vector{Index}

include("Functions/utils.jl")
include("Functions/definitions.jl")
include("Functions/constructors.jl")

end # module

max_stars_repo_name: TheBB/Jutils

max_issues_repo_name: TheBB/Jutils

max_forks_repo_name: TheBB/Jutils


text:
using DashBootstrapComponents

email_input = dbc_row(
    [
        dbc_label("Email", html_for = "example-email-row", width = 2),
        dbc_col(
            dbc_input(
                type = "email",
                id = "example-email-row",
                placeholder = "Enter email",
            ),
            width = 10,
        ),
    ],
    className = "mb-3",
);

password_input = dbc_row(
    [
        dbc_label("Password", html_for = "example-password-row", width = 2),
        dbc_col(
            dbc_input(
                type = "password",
                id = "example-password-row",
                placeholder = "Enter password",
            ),
            width = 10,
        ),
    ],
    className = "mb-3",
);

radios_input = dbc_row(
    [
        dbc_label("Radios", html_for = "example-radios-row", width = 2),
        dbc_col(
            dbc_radioitems(
                id = "example-radios-row",
                options = [
                    Dict("label" => "First radio", "value" => 1),
                    Dict("label" => "Second radio", "value" => 2),
                    Dict(
                        "label" => "Third disabled radio",
                        "value" => 3,
                        "disabled" => true,
                    ),
                ],
            ),
            width = 10,
        ),
    ],
    className = "mb-3",
);

form = dbc_form([email_input, password_input, radios_input]);

max_stars_repo_name: glsdown/dash-bootstrap-components

max_issues_repo_name: glsdown/dash-bootstrap-components

max_forks_repo_name: glsdown/dash-bootstrap-components


text:
module AcousticWaveLab
    using Requires
    include("FDMtx/FDMtx.jl")
    include("plot/plotting.jl")
    include("imaging/imaging.jl")
    include("sourceLocation/sourceLocation.jl")
    @require PyPlot include("plot/plotting.jl")
end

max_stars_repo_name: WneleiGao/AcousticWaveLab

max_issues_repo_name: WneleiGao/AcousticWaveLab

max_forks_repo_name: WneleiGao/AcousticWaveLab


text:
include("metrics/metrics.jl")
include("evaluation.jl")
include("launch_experiment.jl")
include("training.jl")
include("benchmark_solving.jl")