text:
# TODO
# we need to have nquat directly inside the model

################################################################################
# z indices
################################################################################

"""
	Returns the indices of q2 in z and Δq2 in Δz. Important to note that for the indexing
	in z, we want to set quat = false. For indexing in Δz, we set quat = true, so that nquat =
	number of quaternions.
"""
function index_q2(model::ContactModel, env::Environment; quat::Bool = false)
	nq = model.dim.q
	nquat = quat ? model.dim.quat : 0
	iq2 = Vector(1:nq - nquat)
	return iq2
end

"""
	Returns the indices of γ1 in z and Δγ1 in Δz. Important to note that for the indexing
	in z, we want to set quat = false. For indexing in Δz, we set quat = true, so that nquat =
	number of quaternions.
"""
function index_γ1(model::ContactModel, env::Environment; quat::Bool = false)
	nq = model.dim.q
	nc = model.dim.c
	nquat = quat ? model.dim.quat : 0
	off = (nq - nquat)
	iγ1 = Vector(off .+ (1:nc))
	return iγ1
end

"""
	Returns the indices of b1 in z and Δb1 in Δz. Important to note that for the indexing
	in z, we want to set quat = false. For indexing in Δz, we set quat = true, so that nquat =
	number of quaternions.
"""
function index_b1(model::ContactModel, env::Environment; quat::Bool = false)
	nq = model.dim.q
	nc = model.dim.c
	nquat = quat ? model.dim.quat : 0
	nb = nc * friction_dim(env)
	off = (nq - nquat) + nc
	ib1 = Vector(off .+ (1:nb))
	return ib1
end

"""
	Returns the indices of ψ1 in z and Δψ1 in Δz. Important to note that for the indexing
	in z, we want to set quat = false. For indexing in Δz, we set quat = true, so that nquat =
	number of quaternions.
"""
function index_ψ1(model::ContactModel, env::Environment; quat::Bool = false)
	nq = model.dim.q
	nc = model.dim.c
	nquat = quat ? model.dim.quat : 0
	nb = nc * friction_dim(env)
	off = (nq - nquat) + nc + nb
	iψ1 = Vector(off .+ (1:nc))
	return iψ1
end

"""
	Returns the indices of s1 in z and Δs1 in Δz. Important to note that for the indexing
	in z, we want to set quat = false. For indexing in Δz, we set quat = true, so that nquat =
	number of quaternions.
"""
function index_s1(model::ContactModel, env::Environment; quat::Bool = false)
	nq = model.dim.q
	nc = model.dim.c
	nquat = quat ? model.dim.quat : 0
	nb = nc * friction_dim(env)
	off = (nq - nquat) + nc + nb + nc
	is1 = Vector(off .+ (1:nc))
	return is1
end

"""
	Returns the indices of η1 in z and Δη1 in Δz. Important to note that for the indexing
	in z, we want to set quat = false. For indexing in Δz, we set quat = true, so that nquat =
	number of quaternions.
"""
function index_η1(model::ContactModel, env::Environment; quat::Bool = false)
	nq = model.dim.q
	nc = model.dim.c
	nquat = quat ? model.dim.quat : 0
	nb = nc * friction_dim(env)
	off = (nq - nquat) + nc + nb + nc + nc
	iη1 = Vector(off .+ (1:nb))
	return iη1
end

"""
	Returns the indices of s2 in z and Δs2 in Δz. Important to note that for the indexing
	in z, we want to set quat = false. For indexing in Δz, we set quat = true, so that nquat =
	number of quaternions.
"""
function index_s2(model::ContactModel, env::Environment; quat::Bool = false)
	nq = model.dim.q
	nc = model.dim.c
	nquat = quat ? model.dim.quat : 0
	nb = nc * friction_dim(env)
	off = (nq - nquat) + nc + nb + nc + nc + nb
	is2 = Vector(off .+ (1:nc))
	return is2
end

################################################################################
# θ indices
################################################################################

"""
	Returns the indices of q0 in θ.
"""
function index_q0(model::ContactModel)
	nq = model.dim.q
	iq0 = Vector(1:nq)
	return iq0
end

"""
	Returns the indices of q1 in θ.
"""
function index_q1(model::ContactModel)
	nq = model.dim.q
	off = nq
	iq1 = Vector(off .+ (1:nq))
	return iq1
end

"""
	Returns the indices of u1 in θ.
"""
function index_u1(model::ContactModel)
	nq = model.dim.q
	nu = model.dim.u
	off = nq + nq
	iu1 = Vector(off .+ (1:nu))
	return iu1
end

"""
	Returns the indices of w1 in θ.
"""
function index_w1(model::ContactModel)
	nq = model.dim.q
	nu = model.dim.u
	nw = model.dim.w
	off = nq + nq + nu
	iw1 = Vector(off .+ (1:nw))
	return iw1
end

"""
	Returns the indices of μ in θ.
"""
function index_μ(model::ContactModel)
	nq = model.dim.q
	nu = model.dim.u
	nw = model.dim.w
	off = nq + nq + nu + nw
	iμ = Vector(off .+ (1:1))
	return iμ
end

"""
	Returns the indices of h in θ.
"""
function index_h(model::ContactModel)
	nq = model.dim.q
	nu = model.dim.u
	nw = model.dim.w
	off = nq + nq + nu + nw + 1
	ih = Vector(off .+ (1:1))
	return ih
end

################################################################################
# Residual indices
################################################################################

"""
	Returns the indices of the dynamics equation in the residual r.
"""
function index_dyn(model::ContactModel, env::Environment; quat::Bool = false)
	nq = model.dim.q
	nquat = quat ? model.dim.quat : 0
	idyn = Vector(1:nq - nquat)
	return idyn
end

"""
	Returns the indices of the impact equation in the residual r.
"""
function index_imp(model::ContactModel, env::Environment; quat::Bool = false)
	nq = model.dim.q
	nc = model.dim.c
	nquat = quat ? model.dim.quat : 0
	off = (nq - nquat)
	iimp = Vector(off .+ (1:nc))
	return iimp
end

"""
	Returns the indices of the maximum dissipation principle equation in the residual r.
"""
function index_mdp(model::ContactModel, env::Environment; quat::Bool = false)
	nq = model.dim.q
	nc = model.dim.c
	nquat = quat ? model.dim.quat : 0
	nb = nc * friction_dim(env)
	off = (nq - nquat) + nc
	imdp = Vector(off .+ (1:nb))
	return imdp
end

"""
	Returns the indices of the friction cone equation in the residual r.
"""
function index_fri(model::ContactModel, env::Environment; quat::Bool = false)
	nq = model.dim.q
	nc = model.dim.c
	nquat = quat ? model.dim.quat : 0
	nb = nc * friction_dim(env)
	off = (nq - nquat) + nc + nb
	ifri = Vector(off .+ (1:nc))
	return ifri
end

"""
	Returns the indices of the bilinear impact equation in the residual r.
"""
function index_bimp(model::ContactModel, env::Environment; quat::Bool = false)
	nq = model.dim.q
	nc = model.dim.c
	nquat = quat ? model.dim.quat : 0
	nb = nc * friction_dim(env)
	off = (nq - nquat) + nc + nb + nc
	ibimp = Vector(off .+ (1:nc))
	return ibimp
end

"""
	Returns the indices of the bilinear maximum dissipation principle equation in the residual r.
"""
function index_bmdp(model::ContactModel, env::Environment; quat::Bool = false)
	nq = model.dim.q
	nc = model.dim.c
	nquat = quat ? model.dim.quat : 0
	nb = nc * friction_dim(env)
	off = (nq - nquat) + nc + nb + nc + nc
	ibmdp = Vector(off .+ (1:nb))
	return ibmdp
end

"""
	Returns the indices of the bilinear friction equation in the residual r.
"""
function index_bfri(model::ContactModel, env::Environment; quat::Bool = false)
	nq = model.dim.q
	nc = model.dim.c
	nquat = quat ? model.dim.quat : 0
	nb = nc * friction_dim(env)
	off = (nq - nquat) + nc + nb + nc + nc + nb
	ibfri = Vector(off .+ (1:nc))
	return ibfri
end

################################################################################
# Aggregated Indices
################################################################################

"""
	Returns the indices of z and Δz, where 3 groups are formed.

	w1 = (q2)
	w2 = (γ1, b1, ψ1)
	w3 = (s1, η1, s2)

	Δw1 = (Δq2)
	Δw2 = (Δγ1, Δb1, Δψ1)
	Δw3 = (Δs1, Δη1, Δs2)

	Note that Δq2 does not have the same size as q2 when optimizing over
	non-euclidean spaces.
"""
function linearization_var_index(model::ContactModel, env::Environment; quat::Bool = false)
	iq2 = index_q2(model, env, quat = quat)
	iγ1 = index_γ1(model, env, quat = quat)
	ib1 = index_b1(model, env, quat = quat)
	iψ1 = index_ψ1(model, env, quat = quat)
	is1 = index_s1(model, env, quat = quat)
	iη1 = index_η1(model, env, quat = quat)
	is2 = index_s2(model, env, quat = quat)
	iw1 = iq2
	iw2 = [iγ1; ib1; iψ1]
	iw3 = [is1; iη1; is2]
	return iw1, iw2, iw3
end

"""
	Returns the indices of the residual r, where 3 groups are formed.
	r = (rdyn, rrst, rbil)
"""
function linearization_term_index(model::ContactModel, env::Environment; quat::Bool = false)
	# dyn = [dyn]
	# rst = [s1  - ..., ≡ ialt
	#        η1  - ...,
	#        s2  - ...,]
	# bil = [γ1 .* s1 .- κ;
	#        b1 .* η1 .- κ;
	#        ψ1 .* s2 .- κ]
	idyn = index_dyn(model, env, quat = quat)
	iimp = index_imp(model, env, quat = quat)
	imdp = index_mdp(model, env, quat = quat)
	ifri = index_fri(model, env, quat = quat)
	ibimp = index_bimp(model, env, quat = quat)
	ibmdp = index_bmdp(model, env, quat = quat)
	ibfri = index_bfri(model, env, quat = quat)

	irst = [iimp; imdp; ifri]
	ibil = [ibimp; ibmdp; ibfri]
	ialt = iimp
	return idyn, irst, ibil, ialt
end

"""
	Returns the positive orthant indices in z or Δz.
"""
function index_ort(model::ContactModel, env::Environment{<:World,LinearizedCone}; quat::Bool = false)
	iγ1 = index_γ1(model, env, quat = quat)
	ib1 = index_b1(model, env, quat = quat)
	iψ1 = index_ψ1(model, env, quat = quat)
	is1 = index_s1(model, env, quat = quat)
	iη1 = index_η1(model, env, quat = quat)
	is2 = index_s2(model, env, quat = quat)
	return [[iγ1; ib1; iψ1], [is1; iη1; is2]]
end

"""
	Returns the positive orthant indices in z or Δz.
"""
function index_ort(model::ContactModel, env::Environment{<:World,NonlinearCone}; quat::Bool = false)
	iγ1 = index_γ1(model, env, quat = quat)
	is1 = index_s1(model, env, quat = quat)
	return [iγ1, is1]
end

"""
	Returns the second order cone indices in z or Δz.
"""
index_soc(model::ContactModel, env::Environment{<:World,LinearizedCone}; quat::Bool = false) = [Vector{Int}[], Vector{Int}[]]

"""
	Returns the second order cone indices in z or Δz.
"""
function index_soc(model::ContactModel, env::Environment{<:World,NonlinearCone}; quat::Bool = false)
	nc = model.dim.c
	nf = friction_dim(env)

	ib1 = index_b1(model, env, quat = quat) # primal cones: vector part
	iψ1 = index_ψ1(model, env, quat = quat) # primal cones: scalar part
	iη1 = index_η1(model, env, quat = quat) # dual cones: vector part
	is2 = index_s2(model, env, quat = quat) # dual cones: scalar part

	pr_idx = [[iψ1[i]; iη1[(i - 1) * nf .+ (1:nf)]] for i = 1:nc]
	du_idx = [[is2[i]; ib1[(i - 1) * nf .+ (1:nf)]] for i = 1:nc]
	[pr_idx, du_idx]
end

function num_var(model::ContactModel, env::Environment; quat::Bool = false)
	nq = model.dim.q
	nc = model.dim.c
	nquat = quat ? model.dim.quat : 0
	nb = nc * friction_dim(env)
	(nq - nquat) + nc + nb + nc + nc + nb + nc
end

function num_data(model::ContactModel)
	nq = model.dim.q
	nu = model.dim.u
	nw = model.dim.w
	nq + nq + nu + nw + 1 + 1
end

function num_bilinear(model::ContactModel, env::Environment)
	nc = model.dim.c
	nb = nc * friction_dim(env)
	nc + nb + nc
end

################################################################################
# Packing and Unpacking Variables
################################################################################

function unpack_θ(model::ContactModel, θ)
	iq0 = index_q0(model)
	iq1 = index_q1(model)
	iu1 = index_u1(model)
	iw1 = index_w1(model)
	iμ  = index_μ(model)
	ih  = index_h(model)

	q0 = θ[iq0]
	q1 = θ[iq1]
	u1 = θ[iu1]
	w1 = θ[iw1]
	μ = θ[iμ]
	h = θ[ih]
	return q0, q1, u1, w1, μ, h
end

function pack_θ(model::ContactModel, q0, q1, u1, w1, μ, h)
	return [q0; q1; u1; w1; μ; h]
end

function unpack_z(model::ContactModel, env::Environment, z)
	iq2 = index_q2(model, env, quat = false)
	iγ1 = index_γ1(model, env, quat = false)
	ib1 = index_b1(model, env, quat = false)
	iψ1 = index_ψ1(model, env, quat = false)
	is1 = index_s1(model, env, quat = false)
	iη1 = index_η1(model, env, quat = false)
	is2 = index_s2(model, env, quat = false)

	# system variables
	q2 = z[iq2]
	γ1 = z[iγ1]
	b1 = z[ib1]
	ψ1 = z[iψ1]
	s1 = z[is1]
	η1 = z[iη1]
	s2 = z[is2]
	return q2, γ1, b1, ψ1, s1, η1, s2
end

function pack_z(model::ContactModel, env::Environment{<:World,LinearizedCone}, q2, γ1, b1, ψ1, η1)
	s1 = ϕ_func(model, env, q2)
	s2 = model.μ_world * γ1 .- E_func(model, env) * b1
	return pack_z(model, env, q2, γ1, b1, ψ1, s1, η1, s2)
end

function pack_z(model::ContactModel, env::Environment{<:World,NonlinearCone}, q2, γ1, b1, ψ1, η1)
	s1 = ϕ_func(model, env, q2)
	s2 = model.μ_world .* γ1
	return pack_z(model, env, q2, γ1, b1, ψ1, s1, η1, s2)
end

function pack_z(model::ContactModel, env::Environment, q2, γ1, b1, ψ1, s1, η1, s2)
	return [q2; γ1; b1; ψ1; s1; η1; s2]
end

################################################################################
# Aggregated indices
################################################################################

function index_equr(model::ContactModel, env::Environment; quat::Bool = false)
	idyn = index_dyn(model, env, quat = quat)
	iimp = index_imp(model, env, quat = quat)
	imdp = index_mdp(model, env, quat = quat)
	ifri = index_fri(model, env, quat = quat)
	iequ = [idyn; iimp; imdp; ifri]
	return iequ
end

function index_ortr(model::ContactModel, env::Environment{<:World,LinearizedCone}; quat::Bool = false)
	ibimp = index_bimp(model, env, quat = quat)
	ibmdp = index_bmdp(model, env, quat = quat)
	ibfri = index_bfri(model, env, quat = quat)
	iortr = [ibimp; ibmdp; ibfri]
	return iortr
end

function index_ortr(model::ContactModel, env::Environment{<:World,NonlinearCone}; quat::Bool = false)
	ibimp = index_bimp(model, env, quat = quat)
	iortr = ibimp
	return iortr
end

function index_socr(model::ContactModel, env::Environment{<:World,LinearizedCone}; quat::Bool = false)
	isocr = Vector{Int64}()
	return isocr
end

function index_socr(model::ContactModel, env::Environment{<:World,NonlinearCone}; quat::Bool = false)
	ibmdp = index_bmdp(model, env, quat = quat)
	ibfri = index_bfri(model, env, quat = quat)
	isocr = [ibmdp; ibfri]
	return isocr
end

################################################################################
# Optimization Space
################################################################################

abstract type OptimizationSpace
end

mutable struct OptimizationSpace13 <: OptimizationSpace
	# Set the residual to 0
	# r(z) = 0
	# z <- z + Δz

	# Dimensions
	nz::Int # dimension of the optimization variable z
	nΔ::Int # dimension of the optimization variable Δz and of the residual r
	ny::Int # dimension of the bilinear vaiables
	nquat::Int # number of quaternions

	# Variables
	dynz::Vector{Int} # indices of the variables associated with the dynamics constraints in z
	dynΔ::Vector{Int} # indices of the variables associated with the dynamics constraints in Δz
	ortz::Vector{Vector{Int}} # indices of the variables associated with the positive ORThant constraints in z
	ortΔ::Vector{Vector{Int}} # indices of the variables associated with the positive ORThant constraints in Δz
	socz::Vector{Vector{Vector{Int}}} # indices of the variables associated with the Second Order Cone constraints in z
	socΔ::Vector{Vector{Vector{Int}}} # indices of the variables associated with the Second Order Cone constraints in Δz

	# Residual
	equr::Vector{Int} # indices of the residual associated with the EQUality constraints in r
	ortr::Vector{Int} # indices of the residual associated with the positive ORThant constraints in r
	socr::Vector{Int} # indices of the residual associated with the Second Order Cone constraints in r
	dyn::Vector{Int} # indices of the residual associated with the dynamics constraints in r
	rst::Vector{Int} # indices of the residual associated with the remaining constraints in r
	bil::Vector{Int} # indices of the residual associated with the bilinear constraints in r
	alt::Vector{Int} # indices of the residual associated with the altitude constraints in r
end

function OptimizationSpace13(model::ContactModel, env::Environment)
	# Dimensions
	nz = num_var(model, env, quat = false)
	nΔ = num_var(model, env, quat = true)
	ny = num_bilinear(model, env)
	nquat = model.dim.quat

	# Variables
	dynz = index_q2(model, env, quat = false)
	dynΔ = index_q2(model, env, quat = true)
	ortz = index_ort(model, env, quat = true)
	ortΔ = index_ort(model, env, quat = false)
	socz = index_soc(model, env, quat = true)
	socΔ = index_soc(model, env, quat = false)

	# Residual
	# dyn = [dyn]
	# rst = [s1  - ..., ≡ ialt
	#        η1  - ...,
	#        s2  - ...,]
	# bil = [γ1 .* s1 .- κ;
	#        b1 .* η1 .- κ;
	#        ψ1 .* s2 .- κ]

	equr = index_equr(model, env, quat = true)
	ortr = index_ortr(model, env, quat = true)
	socr = index_socr(model, env, quat = true)

	dyn = index_dyn(model, env, quat = true)
	imp = index_imp(model, env, quat = true)
	mdp = index_mdp(model, env, quat = true)
	fri = index_fri(model, env, quat = true)
	bimp = index_bimp(model, env, quat = true)
	bmdp = index_bmdp(model, env, quat = true)
	bfri = index_bfri(model, env, quat = true)

	rst = [imp; mdp; fri]
	bil = [bimp; bmdp; bfri]
	alt = imp

	s = OptimizationSpace13(
		nz, nΔ, ny, nquat,
		dynz, dynΔ, ortz, ortΔ, socz, socΔ,
		equr, ortr, socr, dyn, rst, bil, alt)
	return s
end

function OptimizationSpace13()
	v1 = Vector{Int}()
	v2 = Vector{Vector{Int}}()
	v3 = Vector{Vector{Vector{Int}}}()

	s = OptimizationSpace13(
		0, 0, 0, 0,
		v1, v1, v2, v2, v3, v3,
		v1, v1, v2, v1, v1, v1, v1)
	return s
end


text:
function add_load_model!(pm::GenericPowerModel, n::Int)
    bus_loads = Dict([(i, []) for i in PowerModels.ids(pm, n,:bus)])
     for (i,load) in PowerModels.ref(pm, n,:load)
        push!(bus_loads[load["load_bus"]], i)
     end
     pm.ref[:nw][n][:bus_loads] = bus_loads

end
add_load_model!(pm::GenericPowerModel) = add_load_model!(pm::GenericPowerModel, pm.cnw)

function add_power_factor!(pm::GenericPowerModel, n::Int)
    load_angle = Dict([(i, []) for i in PowerModels.ids(pm, n,:load)])
     for (i,load) in PowerModels.ref(pm, n,:load)
        push!(load_angle[i], atan(load["qref"],load["pref"]))
     end
     pm.ref[:nw][n][:load_angle] = load_angle

end
add_power_factor!(pm::GenericPowerModel) = add_power_factor!(pm::GenericPowerModel, pm.cnw)

ext: jl

max_issues_repo_name: frederikgeth/PowerModelsReliability

max_forks_repo_name: frederikgeth/PowerModelsReliability


text:
using MPI
using Test
using CLIMA.GenericCallbacks
using CLIMA.ODESolvers: AbstractODESolver

mutable struct PseudoSolver <: AbstractODESolver
    t::UInt64
    step::Int
    PseudoSolver() = new(time_ns(), 0)
end
gettime(ps::PseudoSolver) = ps.t

function call_cbs(cbs, init = false)
    foreach(cbs) do cb
        cb(init)
    end
end

function main()
    MPI.Init()
    mpicomm = MPI.COMM_WORLD
    ps = PseudoSolver()

    wtcb_initialized = false
    wtcb_calls = 0
    wtcb = GenericCallbacks.EveryXWallTimeSeconds(1, mpicomm) do (init = false)
        if init
            wtcb_initialized = true
        else
            wtcb_calls = wtcb_calls + 1
        end
    end
    stcb_initialized = false
    stcb_calls = 0
    stcb = GenericCallbacks.EveryXSimulationTime(0.5, ps) do (init = false)
        if init
            stcb_initialized = true
        else
            stcb_calls = stcb_calls + 1
        end
    end
    sscb_initialized = false
    sscb_calls = 0
    sscb = GenericCallbacks.EveryXSimulationSteps(10) do (init = false)
        if init
            sscb_initialized = true
        else
            sscb_calls = sscb_calls + 1
        end
    end
    callbacks = (wtcb, stcb, sscb)

    @testset "Generic Callbacks" begin
        call_cbs(callbacks, true)
        @test wtcb_initialized
        @test stcb_initialized
        @test sscb_initialized
        call_cbs(callbacks)
        @test wtcb_calls == 0
        @test stcb_calls == 0
        @test sscb_calls == 0
        sleep(0.5)
        ps.t = time_ns()
        call_cbs(callbacks)
        @test wtcb_calls == 0
        @test stcb_calls == 1
        @test sscb_calls == 0
        sleep(0.5)
        ps.t = time_ns()
        call_cbs(callbacks)
        @test wtcb_calls == 1
        @test stcb_calls == 2
        @test sscb_calls == 0
        for i in 1:7
            call_cbs(callbacks)
        end
        @test wtcb_calls == 1
        @test stcb_calls == 2
        @test sscb_calls == 1
    end
end

main()


text:
"""
    $(SIGNATURES)

Test by regression that simulated data are consistent with "true" outcomes.

Regress simulated on true.
Check for 0 intercept, 1 slope.

Note: If data are very far from 0, then intercept can be quite a bit off.

# Arguments
- trueM
    true outcome by state
- simM
    simulated outcome by state
- wtM
    weights (e.g. to downweight rare states or states that cannot be solved precisely)
    scalar weights are expanded. Provide wtM = []
- doShow
    show scatter plot?

# Outputs
dev = max(abs(beta - [0;1]) / se(beta))
    how many std deviations are regression coefficient from 45 degree line?
betaV
    regression coefficients; should be [0;1]
seV
    std errors

Change:
    test hypothesis of 0 intercept and unit slope. Return prob of reject
    GLM cannot do that yet.
"""
function regression_test(trueM :: AbstractArray{F1}, simM :: AbstractArray{F1};
    weights = similar(trueM, 0)) where F1 <: AbstractFloat

    @assert size(trueM) == size(simM)
    n = length(trueM);
    @assert n > 2  "Too few observations: $n"

    m = fit(LinearModel, hcat(ones(F1, n),  vec(trueM)),  vec(simM);
        wts = weights);
    betaV = coef(m);
    seV = stderror(m);
    dev = maximum(abs.(betaV .- [zero(F1), one(F1)]) ./ max.(F1(0.01), seV));
    return dev, betaV, seV
end

# ---------------


text:
# fbp/z-list.jl

include("image_geom.jl")
export image_geom
export MIRT_image_geom

include("sino_geom.jl")
export sino_geom
export MIRT_sino_geom

include("rotate2d.jl")
export rotate2d

include("cuboid_im.jl") # must be after image_geom.jl
export cuboid_im

include("disk-phantom.jl")
include("ellipse_im.jl")
include("ellipsoid_im.jl")
include("ellipse_sino.jl")

include("rect_im.jl")
export rect_im
include("rect_sino.jl")


text:
"""
A trace is a Vector{CartesianIndex{2}} and additionally supports `draw!, Matrix`.
"""
const Trace = Vector{CartesianIndex{2}}

"""
    Blob

Represents a blob. Internally stores a Kalman filter, a counter that is incremented when the blob is not assigned a measurement and a trace of all locations and all seen measurements. If no measurement was seen for a particular time step, `OOB = CartesianIndex(0,0)` is recorded.

This type supports `location, trace, tracem, lifetime, draw!`
"""
mutable struct Blob
    kf
    counter::Float64
    trace::Trace
    tracem::Trace
end

Base.Broadcast.broadcastable(b::Blob) = Ref(b)

Blob() = Blob(KalmanParams(1,1), CartesianIndex(1,1))

"""
    trace(b::Blob)

Get the location trace of a blob. Use `tm = Matrix(t::Trace)` to get an N×2 matrix. Use `tm = replace(Float64.(tm), 0=>NaN)` to create a matrix with NaNs where there were missing measurements, this is useful for plotting since it creates a gap where the missing measurement was.
"""
trace(b::Blob) = b.trace

"""
    tracem(b::Blob)

Get the measurement trace of a blob. Use `skipmissing(t::Trace)` to filter out missing measurements. Use `tm = Matrix(t::Trace)` to get an N×2 matrix. Use `tm = replace(Float64.(tm), 0=>NaN)` to create a matrix with NaNs where there were missing measurements, this is useful for plotting since it creates a gap where the missing measurement was.
"""
tracem(b::Blob) = b.tracem
Base.skipmissing(trace::Trace) = filter(!=(OOB), trace)
Base.Matrix(trace::Trace) = [getindex.(trace,1) getindex.(trace,2)]

struct KalmanParams
    σw::Float64
    σe::Float64
end

Base.Broadcast.broadcastable(b::KalmanParams) = Ref(b)

"""
    BlobTracker{T <: AbstractCorrespondence}

Example `bt = BlobTracker(sizes=3:5, σw=2.0, σe = 10.0)`

# Optional Keyword arguments:
- `params::KalmanParams`: Holds the parameters for the kalman filter `σw,σe`
- `amplitude_th = 0.0001`: blobs must be at least this prominent to be considered
- `kill_counter_th::Int = 10`: after this many steps without an assigned measurement a blob will die
- `sizes::AbstractVector`: vector of numbers determining the size scales at which blobs are detected. See docs for [`Images.blob_LoG`](@ref).
- `preprocessor = ((storage, img)->nothing` a function that processes `img` and stores the result in `storage`
- `correspondence::AbstractCorrespondence = HungarianCorrespondence(p=1.0, dist_th=2)`: Determines how blobs are assigned to measurements
- `mask = nothing`: An optional boolean image that is false where you want to ignore blobs and true where you want to track them.
"""
Base.@kwdef mutable struct BlobTracker{T<:AbstractCorrespondence}
    params::KalmanParams
    amplitude_th = 0.0001
    kill_counter_th::Int = 10
    sizes::AbstractVector
    preprocessor = (storage,img)-> (storage .= Gray.(img))
    correspondence::T = HungarianCorrespondence(p=1.0, dist_th=2)
    mask = nothing
end

"""
    BlobTracker(sizes, σw, σe; kwargs...)

Helper constructor that accepts keyword arguments for the `KalmanParams`

#Arguments:
- `σw`: Dynamics standard deviation
- `σe`: Measurement standard deviation
"""
function BlobTracker(sizes, σw, σe; kwargs...)
    BlobTracker(;sizes = sizes, params = KalmanParams(σw, σe), kwargs...)
end

Base.Broadcast.broadcastable(b::BlobTracker) = Ref(b)

dt = 1
A = @SMatrix [1. 0 dt 0; 0 1 0 dt; 0 0 1 0; 0 0 0 1] #state update matrice
B = @SVector [(dt^2/2), (dt^2/2), dt, dt]
C = @SMatrix [1. 0 0 0; 0 1 0 0]

Re(σe) = diagm([σe^2, σe^2])
Rw(σw) =  [dt^4/4 0 dt^3/2 0;
    0 dt^4/4 0 dt^3/2;
    dt^3/2 0 dt^2 0;
    0 dt^3/2 0 dt^2].*σw^2 + 1e-6I

Re(params::KalmanParams) = Re(params.σe)
Rw(params::KalmanParams) = Rw(params.σw)

"""
    Blob(params::KalmanParams, coord::CartesianIndex)

Spawn a blob using settings from `params` at given coordinates.
"""
function Blob(params::KalmanParams,coord::CartesianIndex)
    kf = KalmanFilter(A,B,C,0,Rw(params),Re(params),MvNormal(10Rw(params)))
    kf.x[1] = coord[1]
    kf.x[2] = coord[2]
    Blob(kf, 0.0, [coord], [coord])
end

"""
    location(b)

Get the location of a blob as CartesianIndex
"""
Distributions.location(b::Blob) = CartesianIndex(round.(Int,(b.kf.x[1], b.kf.x[2])))
Distributions.location(b) = b.location
Distributions.location(b::Tuple) = CartesianIndex(b)
lifetime(b) = findlast(c->c != OOB, b.tracem)

Base.CartesianIndex(blob::Blob) = CartesianIndex((round.(Int,blob.kf.x[1:2])...,))

function detect_blobs!(ws::Workspace, bt::BlobTracker, img)
    blobs = blob_LoG!(ws.blob_storage,ws.storage,bt.sizes)
    blobs = filter!(x->x.amplitude > bt.amplitude_th, blobs)
    [b.location for b in blobs]
end

"""
    tune_sizes!(ws::Workspace, bt::BlobTracker, img)

Display a small GUI with a slider to help you tune the `sizes`
"""
function tune_sizes!(ws::Workspace, bt::BlobTracker, img)
    prepare_image!(ws, bt, img)
    @manipulate for k = 1:100
        imadjustintensity(imfilter(ws.storage, Kernel.LoG(k)))
    end
end

function prepare_image!(ws::Workspace, bt::BlobTracker, img)
    apply_mask!(ws,bt,img)
    bt.preprocessor(ws.storage, img)
end

function apply_mask!(ws::Workspace,bt::BlobTracker,img)
    bt.mask === nothing && return img
    return ws.storage .= bt.mask .* img
end

for f in (:detect_blobs, :tune_sizes, :prepare_image, :apply_mask)
    fb = Symbol(f, :!)
    @eval $f(bt::BlobTracker, img) = $fb(Workspace(img,bt), bt, img)
end

function blob_LoG!(img_LoG,img::AbstractArray{T,N}, σscales::Union{AbstractVector,Tuple},
    edges::Tuple{Vararg{Bool}}=(ntuple(d->false, Val(N))..., true), σshape=ntuple(d->1, Val(N))) where {T,N}
    sigmas = sort(σscales)
    colons = ntuple(d->Colon(), Val(N))
    @inbounds for isigma in eachindex(sigmas)
        img_LoG[colons...,isigma] .= (-sigmas[isigma]) .* imfilter!(@view(img_LoG[colons...,isigma]), img, Kernel.LoG(ntuple(i->sigmas[isigma]*σshape[i],Val(N))))
    end
    maxima = Images.findlocalmaxima(img_LoG, 1:ndims(img_LoG), edges)
    [Images.BlobLoG(CartesianIndex(rtailr(x.I)), sigmas[x[ndims(img_LoG)]], img_LoG[x]) for x in maxima]
end

@inline rtailr(x) = reverse(Base.tail(reverse(x)))

# function blob_LoG!(img_LoG,img::AbstractArray{T,N}, σscales::Union{AbstractVector,Tuple},
#     edges::Tuple{Vararg{Bool}}=(true, ntuple(d->false, Val(N))...), σshape=ntuple(d->1, Val(N))) where {T,N}
#     sigmas = sort(σscales)
#     colons = ntuple(d->Colon(), Val(N))
#     @inbounds for isigma in eachindex(sigmas)
#         img_LoG[isigma,colons...] = (-sigmas[isigma]) * imfilter(img, Kernel.LoG(ntuple(i->sigmas[isigma]*σshape[i],Val(N))))
#     end
#     maxima = Images.findlocalmaxima(img_LoG, 1:ndims(img_LoG), edges)
#     [Images.BlobLoG(CartesianIndex(Base.tail(x.I)), sigmas[x[1]], img_LoG[x]) for x in maxima]
# end


text:
"""
    required_halo_size(tendency_term)

Returns the required size of halos for a term appearing
in the tendency for a velocity field or tracer field.

Example
=======

```jldoctest
using Oceananigans.Advection: CenteredFourthOrder
using Oceananigans.Grids: required_halo_size

required_halo_size(CenteredFourthOrder())

# output
2
"""

function required_halo_size end

required_halo_size(nothing) = 1

inflat_halo_size_one_dimension(H, Hx, TX          ) = max(H, Hx)
inflat_halo_size_one_dimension(H, Hx, ::Type{Flat}) = 0

function inflate_halo_size(Hx, Hy, Hz, topology, tendency_terms...)

    oldHx = Hx
    oldHy = Hy
    oldHz = Hz

    for term in tendency_terms
         H = required_halo_size(term)
        Hx = inflat_halo_size_one_dimension(H, Hx, topology[1])
        Hy = inflat_halo_size_one_dimension(H, Hy, topology[2])
        Hz = inflat_halo_size_one_dimension(H, Hz, topology[3])
    end

    if Hx != oldHx || Hy != oldHy || Hz != oldHz
        @warn "Inflating model grid halo size to ($Hx, $Hy, $Hz) and recreating grid. " *
        "The model grid will be different from the input grid. To avoid this warning, " *
        "pass halo=($Hx, $Hy, $Hz) when constructing the grid."
    end

    return Hx, Hy, Hz
end


text:
struct RegressionResult <: RegressionModel
    coef::Vector{Float64}   # Vector of coefficients
    vcov::Matrix{Float64}   # Covariance matrix

    coefnames::Vector       # Name of coefficients
    yname::Symbol           # Name of dependent variable
    formula::FormulaTerm        # Original formula

    nobs::Int64             # Number of observations
    dof_residual::Int64      # degrees of freedoms

    # rss::Float64            # Sum of squared residuals
    # tss::Float64            # Total sum of squares
    # r2::Float64             # R squared
    # adjr2::Float64          # R squared adjusted
    # F::Float64              # F statistics
    # p::Float64              # p value for the F statistics
end

StatsBase.nobs(x::RegressionResult) = x.nobs
StatsBase.coef(x::RegressionResult) = x.coef
StatsBase.coefnames(x::RegressionResult) = x.coefnames
StatsBase.vcov(x::RegressionResult) = x.vcov
StatsBase.dof_residual(x::RegressionResult) = x.dof_residual

function StatsBase.confint(x::RegressionResult)
    scale = quantile(TDist(x.dof_residual), 1 - (1-0.95)/2)
    se = stderror(x)
    hcat(x.coef -  scale * se, x.coef + scale * se)
end

# Display Results
function title(::RegressionResult)
    "Generalized Linear Fixed Effects Model"
end

function top(::RegressionResult)
    error("function top has no general method for AbstractRegressionResult")
end

top(x::RegressionResult) = [
            "Number of obs" sprint(show, nobs(x), context = :compact => true);
            "Degrees of freedom" sprint(show, nobs(x) - dof_residual(x), context = :compact => true);
            ]

function coeftable(x::RegressionResult)
    ctitle = title(x)
    ctop = top(x)
    cc = coef(x)
    se = stderror(x)
    coefnms = coefnames(x)
    conf_int = confint(x)
    # put (intercept) last
    if !isempty(coefnms) && ((coefnms[1] == Symbol("(Intercept)")) || (coefnms[1] == "(Intercept)"))
        newindex = vcat(2:length(cc), 1)
        cc = cc[newindex]
        se = se[newindex]
        conf_int = conf_int[newindex, :]
        coefnms = coefnms[newindex]
    end
    tt = cc ./ se
    CoefTable2(
        hcat(cc, se, tt, ccdf.(Ref(FDist(1, dof_residual(x))), abs2.(tt)), conf_int[:, 1:2]),
        ["Estimate","Std.Error","t value", "Pr(>|t|)", "Lower 95%", "Upper 95%" ],
        ["$(coefnms[i])" for i = 1:length(cc)], 4, ctitle, ctop)
end

function Base.show(io::IO, x::RegressionResult)
    show(io, coeftable(x))
end

## Coeftalble2 is a modified Coeftable allowing for a top String matrix displayed before the coefficients.
## Pull request: https://github.com/JuliaStats/StatsBase.jl/pull/119

struct CoefTable2
    mat::Matrix
    colnms::Vector
    rownms::Vector
    pvalcol::Integer
    title::AbstractString
    top::Matrix{AbstractString}
    function CoefTable2(mat::Matrix,colnms::Vector,rownms::Vector,pvalcol::Int=0,
                        title::AbstractString = "", top::Matrix = Any[])
        nr,nc = size(mat)
        0 <= pvalcol <= nc || error("pvalcol = $pvalcol should be in 0,...,$nc]")
        length(colnms) in [0,nc] || error("colnms should have length 0 or $nc")
        length(rownms) in [0,nr] || error("rownms should have length 0 or $nr")
        length(top) == 0 || (size(top, 2) == 2 || error("top should have 2 columns"))
        new(mat,colnms,rownms,pvalcol, title, top)
    end
end

## format numbers in the p-value column
function format_scientific(pv::Number)
    return @sprintf("%.3f", pv)
end

function Base.show(io::IO, ct::CoefTable2)
    mat = ct.mat; nr,nc = size(mat); rownms = ct.rownms; colnms = ct.colnms;
    pvc = ct.pvalcol; title = ct.title;   top = ct.top
    if length(rownms) == 0
        rownms = AbstractString[lpad("[$i]",floor(Integer, log10(nr))+3) for i in 1:nr]
    end
    if length(rownms) > 0
        rnwidth = max(4,maximum([length(nm) for nm in rownms]) + 1)
        else
            # if only intercept, rownms is empty collection, so previous would return error
        rnwidth = 4
    end
    rownms = [rpad(nm,rnwidth) for nm in rownms]
    widths = [length(cn)::Int for cn in colnms]
    str = [sprint(show, mat[i,j]; context=:compact => true) for i in 1:nr, j in 1:nc]
    if pvc != 0                         # format the p-values column
        for i in 1:nr
            str[i, pvc] = format_scientific(mat[i, pvc])
        end
    end
    for j in 1:nc
        for i in 1:nr
            lij = length(str[i, j])
            if lij > widths[j]
                widths[j] = lij
            end
        end
    end
    widths .+= 1
    totalwidth = sum(widths) + rnwidth
    if length(title) > 0
        halfwidth = div(totalwidth - length(title), 2)
        println(io, " " ^ halfwidth * string(title) * " " ^ halfwidth)
    end
    if length(top) > 0
        for i in 1:size(top, 1)
            top[i, 1] = top[i, 1] * ":"
        end
        println(io, "=" ^totalwidth)
        halfwidth = div(totalwidth, 2) - 1
        interwidth = 2 +  mod(totalwidth, 2)
        for i in 1:(div(size(top, 1) - 1, 2)+1)
            print(io, top[2*i-1, 1])
            print(io, lpad(top[2*i-1, 2], halfwidth - length(top[2*i-1, 1])))
            print(io, " " ^interwidth)
            if size(top, 1) >= 2*i
                print(io, top[2*i, 1])
                print(io, lpad(top[2*i, 2], halfwidth - length(top[2*i, 1])))
            end
            println(io)
        end
    end
    println(io,"=" ^totalwidth)
    println(io," " ^ rnwidth *
            join([lpad(string(colnms[i]), widths[i]) for i = 1:nc], ""))
    println(io,"-" ^totalwidth)
    for i in 1:nr
        print(io, rownms[i])
        for j in 1:nc
            print(io, lpad(str[i,j],widths[j]))
        end
        println(io)
    end
    println(io,"=" ^totalwidth)
end


text:
#=
make.jl:
- Julia version: 1.0
- Author: qz
- Date: Apr 22, 2019
=#
using Documenter, APMAE4001

makedocs(sitename="APMAE4001 Documentation")

deploydocs(
    repo = "github.com/singularitti/APMAE4001.jl.git",
)


text:
export EQLoc,
EQMag,
EventChannel,
EventTraceData,
FDSNevq,
FDSNevt,
PhaseCat,
SeisEvent,
SeisHdr,
SeisPha,
SeisSrc,
SourceTime,
distaz!,
gcdist,
get_pha!,
read_qml,
show_phases,
write_qml


text: # api for algorithms might go here


text:
# There should be support for a parallel sort, here, I think.

"return a winner individual from a tournament of size t_size from a population"
function tournament_selection(pop::Array{<:Individual}, t_size::Int)
    inds = shuffle!(collect(1:length(pop)))
    sort(pop[inds[1:t_size]])[end]
end

function tournament_selection(geo::Geography, t_size::Int)
    inds = choose_combatants(geo, t_size)
    sort(geo.deme[inds])[end]
end

"return the best individual from a population"
function max_selection(pop::Array{<:Individual})
    sort(pop)[end]
end

function max_selection(geo::Geography)
    max_selection(geo.deme)
end

"return a random individual from a population"
function random_selection(pop::Array{<:Individual})
    pop[rand(1:length(pop))]
end

function random_selection(geo::Geography)
    rand(geo.deme)
end

## Why is t_size hardcoded here?
selection(pop::Array{<:Individual}) = tournament_selection(pop, 4)

selection(geo::Geography) = tournament_selection(geo, 4)


text:
@testset "Evolve" begin
    model = initialise()

    @test length(collect(households(model))) == model[1].households
    @test length(collect(households(model[1], model))) == model[1].households

    @test length(collect(municipalities(model))) == 1

    @test LimnoSES.active_interventions(model[1], 1) == [WastewaterTreatment()]
end


text:
#=
= Step location move.
=#

#= Auxiliary ==================================================================#

function loc_init(Y, S::State, X::Auxiliary) # locs, widths, counts.
    X.locs[:] = S.locs
    loc_metrics(Y,X)
end

function loc_metrics(Y, X::Auxiliary) # widths, counts.
    sp = 0.
    y = 1 # index of the first uncounted datum.
    for i = eachindex(X.widths)
        s = X.locs[i+1]
        X.widths[i] = s-sp
        n = searchsortedfirst(@view(Y[y:end]), s) - 1
        X.counts[i] = n
        sp,y = s,y+n
    end
end

function loc_unconstrained(C::Constants, X::Auxiliary) # tlocs, texps, dlocs.
    locs = @view(X.locs[2:end-1])
    @. X.tlocs = log(locs/(C.R-locs))
end

function loc_constrained(Y, C::Constants, X::Auxiliary) # locs, widths, counts.
    @. X.locs[2:end-1] = C.R/(1+exp.(-X.tlocs))
    loc_metrics(Y,X)
end

function loc_momenta(C::Constants, X::Auxiliary) # ptlocs.
    X.ptlocs[:] = randn(C.k)
end

#= Move =======================================================================#

function loc_move(Y, C::Constants, S::State, X::Auxiliary, U::Counts)

    ε = .8C.ll + rand()*.4C.ll # base step size ± 20%.
    loc_init(Y,S,X)
    loc_unconstrained(C,X)
    loc_momenta(C,X)
    ep = loc_potential(C,X) # previous energy; also computes the gradient.
    kp = loc_kinetic(X)

    # loc_println(X) # debug.
    # println("ep: $ep, kp: $kp") # debug.

    loc_pstep(.5ε,X)
    for i = 1:C.Ll
        loc_qstep(ε,Y,C,X)
        if i != C.Ll
            loc_potential(C,X)
            loc_pstep(ε,X)
        end
    end
    en = loc_potential(C,X)
    loc_pstep(.5ε,X)
    kn = loc_kinetic(X)

    # loc_println(X) # debug.
    # println("en: $en, kn: $kn") # debug.
    # println("diff: $(ep+kp-en-kn)") # debug.

    p = rand()
    l = exp(ep+kp-en-kn)
    # println((p < l) ? "$p < $l" : "$p ≥ $l") # debug.

    if p < l
        S.locs[:] = X.locs
        S.hts[:] = X.hts # order might have changed.
        U.laccepts += 1
    else
        loc_init(Y,S,X)        # leave the auxiliary parameters
        loc_unconstrained(C,X) # in a consistent state.
    end
    U.lprops += 1

end

function loc_pstep(s, X::Auxiliary)
    @. X.ptlocs -= s * X.∇tlocs
end

function loc_qstep(s, Y, C::Constants, X::Auxiliary)
    @. X.tlocs += s * X.ptlocs
    if !issorted(X.tlocs)
        sortperm!(X.ix, X.tlocs, initialized=true)
        X.tlocs[:] = X.tlocs[X.ix]
        X.ptlocs[:] = X.ptlocs[X.ix]
        X.hts[2:end] = @view(X.hts[2:end])[X.ix]
        X.lhts[2:end] = @view(X.lhts[2:end])[X.ix]
        X.ix[:] = 1:length(X.ix)
    end
    loc_constrained(Y,C,X)
end

#= Energy =====================================================================#

function loc_potential(C::Constants, X::Auxiliary)
    k,R,α,β = C.k, C.R, C.α, C.β
    hts, widths, counts = X.hts, X.widths, X.counts
    locs = @view(X.locs[2:end-1])

    e = sum(@. (widths+β)*hts - (counts+α)*X.lhts - log(widths)) -
        sum(@. log(locs*(R-locs)))
    X.∇tlocs .= @. begin
        (@view(hts[1:end-1]) - @view(hts[2:end]) -
        1/@view(widths[1:end-1]) + 1/@view(widths[2:end]))*(locs*(R-locs)/R) -
        (1-2locs/R)
    end
    return e
end

loc_kinetic(X::Auxiliary) = .5*sum(X.ptlocs.^2)

#= I/O ========================================================================#

function loc_println(X::Auxiliary)
    println("X: locs = $(string(X.locs)), hts = $(string(X.hts))\n",
            "   widths = $(string(X.widths)), counts = $(X.counts)\n",
            "   tlocs = $(string(X.tlocs)), ∇tlocs = $(string(X.∇tlocs))\n",
            "   ptlocs = $(string(X.ptlocs)), ix = $(X.ix)")
end

max_stars_repo_name: kevdur/rjcp

max_issues_repo_name: kevdur/rjcp

max_forks_repo_name: kevdur/rjcp


text:
#using Reexport, RecipesBase

#@reexport using Plots

"
Y = βX + ϵ
E(Y|H) = E(βX + ϵ|H)
       = E(βX|H) + E(ϵ|H) # linearity of E
       = E(βX) + E(ϵ|H) # βX only evaluated  when H
       = βX + E(ϵ|H)

E(ϵ|u > -γW) = E(ϵ|u/σ > -γW/σ) # standazilation
            = ρψ(-γW/σ)/(1-Ψ(-γW/σ))
            = ρψ(γW/σ)/Ψ(γW/σ)) # pdf of standard normal distribution is symmetry
                                # cdf can be transformed by inverse x-axis
"


text:
export invariant_polytopes
using LinearAlgebra

struct SymPoint{VT}
    point::VT
end

include("conitope.jl")
include("balanced_real_polytope.jl")
include("balanced_complex_polytope.jl")

const PolytopeLike = Union{Conitope, BalancedRealPolytope, BalancedComplexPolytope}

function Polyhedra.convexhull!(p::PolytopeLike, v::SymPoint)
    push!(p.points, v.point)
    if p.model !== nothing
        _update_model(p, v.point)
    end
end
function Base.in(v, brp::PolytopeLike)
    if isempty(brp.points)
        return false
    end
    if brp.model === nothing
        _parametrized_model(brp, v)
    else
        _fix(brp, v)
    end
    MOI.optimize!(brp.model)
    status = MOI.get(brp.model, MOI.TerminationStatus())
    if status == MOI.OPTIMAL
        return true
    elseif status == MOI.INFEASIBLE || status == MOI.INFEASIBLE_OR_UNBOUNDED
        # There is no objective so it cannot be unbounded.
        # Therefore, `MOI.INFEASIBLE_OR_UNBOUNDED` means infeasible.
        return false
    else
        error("Solver returned status $status.")
    end
end
function in_ratio(v, p::PolytopeLike)
    if isempty(p.points)
        # The polytope is the set containing the origin only so `v` needs to be
        # multiplied by 0.0 to be mapped to the polytope
        return 0.0
    end
    if p.model === nothing
        _ratio_model(p, v)
        MOI.set(p.model, MOI.ObjectiveSense(), MOI.MAX_SENSE)
        MOI.set(p.model, MOI.ObjectiveFunction{MOI.SingleVariable}(),
                MOI.SingleVariable(p.t_0))
    else
        _fix(p, v)
    end
    MOI.optimize!(p.model)
    status = MOI.get(p.model, MOI.TerminationStatus())
    if status == MOI.OPTIMAL
        r = MOI.get(p.model, MOI.VariablePrimal(), p.t_0)
        return r
    else
        error("Solver returned status $status.")
    end
end

function root_path(parent, leaf, n)
    path = Int[]
    cur = leaf
    while cur > n
        push!(path, cur)
        cur = parent[cur]
    end
    return cur, reverse(path)
end

"""
Computes invariant balanced polytopes for system `s` using the algorithm of
[GP13] with the spectral maximizing product candidate (s.m.p.) `smp`. The
choice between Complex Balanced Polytope, Real Balanced Polytope or Conitope
using invariance of the positive orthant is done automatically. The solver used
is `optimizer_constructor`. This solver should support second order cone if the Complex
Balanced Polytope is used (e.g. if the leading eigenvalue has an imaginary
part). The maximum depth considered is `max_length`, a point `x` is considered
to be in the interior of the polytope if the maximum `λ` such that `λ * x` is
in the polytope is larger than `1 + tol`.

New s.m.p. candidates are considered with tolerance `new_candidate_tol`. A
maximum of `max_cycles` times the current s.m.p. candidate is prepended to this
new candidate to try to find a candidate with growth rate higher than the
current one, we also stop adding cycles if the length reaches `max_smp_length`.
If `gready`, then cycles continue to be prepended even if the it has a higher growth
rate than the current s.m.p. until the growth rate decreases when prepending a new
cycle.

If `verbose` is 1, the number of leaves is printed when the depth is a multiple
of `log_step_length`.

[GP13] N. Guglielmi and V. Protasov.
*Exact computation of joint spectral characteristics of linear operators*.
Foundations of Computational Mathematics 13.1, **2013**, 37-97.
"""
function invariant_polytopes(
    s::AbstractDiscreteSwitchedSystem, optimizer_constructor,
    smp::AbstractPeriodicSwitching; max_length=10, verbose=2, tol=nothing,
    max_cycles=10, new_candidate_tol=1e-6, gready=false, max_smp_length=50,
    log_step_length=10)
    leaves = Int[]
    sets = PolytopeLike[]
    new_smp = true
    max_lowtol = -Inf
    min_uptol = Inf
    max_low_candidate_tol = -Inf
    min_up_candidate_tol = Inf
    while new_smp
        new_smp = false
        n = length(smp.period)
        # The ith vertex is between transition[i] and transition[i+1]
        transition = [smp.period[end]; smp.period[1:(end-1)]]
        parent = [n; 1:(n-1)]
        A = prod(reverse(SwitchOnSafety.integratorfor.(s, smp.period)))
        # non-Hermitian matrices are not diagonalizable by StaticArrays
        if A isa SMatrix && !ishermitian(A)
            A = Matrix(A)
        end
        eig = eigen(Matrix(A))
        # TODO check whether the multiplicity is one
        λind = argmax(abs.(eig.values))
        λ_raw = eig.values[λind]
        λ = abs(λ_raw)^(1/n)
        @assert λ ≈ smp.growthrate
        v1 = eig.vectors[:, λind]
        u1 = inv(eig.vectors)[λind, :]
        if isreal(λ_raw)
            v1 = real.(v1)
            u1 = real.(u1)
        end
        if all(all(x -> x ≥ 0, SwitchOnSafety.integratorfor(s, t))
               for t in transitions(s))
            if any(x -> x < 0, v1)
                v1 = -v1
            end
            @assert all(x -> x ≥ 0, v1)
            sets = [Conitope(statedim(s, mode), Vector{eltype(v1)}[], optimizer_constructor)
                    for mode in modes(s)]
        else
            if isreal(λ_raw)
                sets = [BalancedRealPolytope(statedim(s, mode),
                                             Vector{eltype(v1)}[],
                                             optimizer_constructor)
                        for mode in modes(s)]
            else
                sets = [BalancedComplexPolytope(statedim(s, mode),
                                                Vector{eltype(v1)}[],
                                                optimizer_constructor)
                        for mode in modes(s)]
            end
        end
        vertices = [v1]
        duals = [u1 / dot(u1, v1)]
        for i in 2:n
            # TODO should maybe also add complex conjugate if it is complex like in CJSR
            push!(vertices,
                  SwitchOnSafety.integratorfor(s, transition[i]) * vertices[i-1] / λ)
        end
        for i in n:-1:1
            j = (i % n) + 1
            u = SwitchOnSafety.integratorfor(s, transition[j])' * vertices[j]
            u /= dot(u, vertices[i])
            push!(duals, u)
        end
        _mode(i) = target(s, transition[i])
        for i in 1:n
            if verbose ≥ 3
                println("v_$i = ", vertices[i])
            end
            Polyhedra.convexhull!(sets[_mode(i)], SymPoint(vertices[i]))
        end
        leaves = collect(1:n)
        for k in 1:max_length
            new_smp && break
            isempty(leaves) && break
            if verbose ≥ 2 || (verbose ≥ 1 && ((k % log_step_length) == 0))
                println("Depth $k: $(map(p -> length(p.points), sets)) points, $(length(leaves)) living leaves...")
            end
            new_leaves = Int[]
            for leaf in leaves
                new_smp && break
                root, path = root_path(parent, leaf, n)
                if verbose ≥ 3
                    println("Path starting starting at root $root: $(map(i -> symbol(s, transition[i]), path)):")
                end
                for t in out_transitions(s, _mode(leaf))
                    new_smp && break
                    if k == 1 && t == transition[(leaf % n) + 1]
                        continue
                    end
                    v = SwitchOnSafety.integratorfor(s, t) * vertices[leaf] / λ
                    if verbose ≥ 3
                        print(symbol(s, t), " : ", v, " : ")
                    end
                    if tol === nothing
                        is_in = v in sets[target(s, t)]
                    else
                        r = in_ratio(v, sets[target(s, t)]) - 1
                        is_in = r ≥ tol
                    end
                    if is_in
                        if verbose ≥ 3
                            printstyled("dead leaf", color=:red)
                            if tol !== nothing
                                println(": $r ≥ $tol")
                            end
                            println()
                        end
                        if tol !== nothing
                            min_uptol = min(min_uptol, r)
                        end
                    else
                        if target(s, t) == target(s, transition[root])
                            candidate_rating = abs(real(dot(duals[root], v))) - 1
                            if candidate_rating > new_candidate_tol
                                min_up_candidate_tol = min(min_up_candidate_tol, candidate_rating)
                                period = [[transition[i] for i in path]; t]
                                if verbose ≥ 2
                                    printstyled("New candidate s.m.p. found", bold=true, color=:blue)
                                    println(" because ⟨v_$root*, v⟩ - 1 = $(dot(duals[root], v) - 1) > $new_candidate_tol:")
                                end
                                cycle = [transition[root+1:n]; transition[1:root]]
                                _smp = periodicswitching(s, period, scaling = smp.growthrate)
                                __smp = _smp
                                better = isbetter(_smp, smp)
                                n_cycles = 0
                                for i in 1:max_cycles
                                    if length(_smp.period) + length(cycle) > max_smp_length
                                        break
                                    end
                                    if better
                                        if gready
                                            if i > 1 && !isbetter(_smp, __smp)
                                                # It does not improve, let's stop
                                                _smp = __smp
                                                break
                                            end
                                        else
                                            break
                                        end
                                    end
                                    if verbose ≥ 3
                                        print(_smp)
                                        if better
                                            print(" is better than existing s.m.p. but it might still be improved")
                                        else
                                            print(" is not better yet")
                                        end
                                        println(", adding a cycle.")
                                    end
                                    period = [cycle; period]
                                    __smp = _smp
                                    _smp = periodicswitching(s, period, scaling = smp.growthrate)
                                    n_cycles += 1
                                    if !better
                                        better = isbetter(_smp, smp)
                                    end
                                end
                                if better
                                    new_smp = true
                                    smp = _smp
                                    if verbose ≥ 1
                                        println(smp)
                                    end
                                else
                                    if verbose ≥ 2
                                        printstyled("aborting", bold=true, color=:red)
                                        print(" with smp candidate $_smp of length $(length(_smp.period)) after prefixing the suffix with $n_cycles times the current s.m.p.")
                                        @assert n_cycles ≤ max_cycles
                                        if n_cycles == max_cycles
                                            println(" Increase the value of the `max_cycles` keyword argument to go further.")
                                        else
                                            println(" Increase the value of the `max_smp_length` keyword argument to go further.")
                                        end
                                    end
                                end
                            else
                                max_low_candidate_tol = max(max_low_candidate_tol, candidate_rating)
                            end
                        end
                        if !new_smp
                            if verbose ≥ 3
                                printstyled("living leaf", color=:green)
                                if tol !== nothing
                                    println(": $r < $tol")
                                end
                                println()
                            end
                            if tol !== nothing
                                max_lowtol = max(max_lowtol, r)
                            end
                            push!(transition, t)
                            push!(parent, leaf)
                            push!(vertices, v)
                            push!(new_leaves, length(parent))
                            Polyhedra.convexhull!(sets[_mode(length(parent))], SymPoint(v))
                        end
                    end
                end
            end
            leaves = new_leaves
        end
    end
    if verbose ≥ 1
        if isempty(leaves)
            print("0 living leaves, $smp is a ")
            printstyled("spectral maximizing product", bold=true)
            println(" (s.m.p.)")
        else
            print("Still $(length(leaves)) living leaves, ")
            printstyled("aborting", bold=true, color=:red)
            println(" at depth $max_length. Increase the `max_length` keyword argument to go further.")
        end
        if tol === nothing
            println("Use the `tol` keyword argument to change the threshold to evaluate dead and living leaves.")
        else
            println("Use `tol` < $max_lowtol to eliminate at least one more living leaf.")
            println("Use `tol` > $min_uptol to keep at least one more dead leaf.")
        end
        if isfinite(max_low_candidate_tol)
            println("Use `new_candidate_tol` < $max_low_candidate_tol to consider at least one more candidate as s.m.p.")
        end
        if isfinite(min_up_candidate_tol)
            println("Use `new_candidate_tol` > $min_up_candidate_tol to eliminate at least one s.m.p. candidate.")
        end
    end
    return smp, isempty(leaves), sets
end


text:
"""
    ContactProcess{T} <: AbstractEpidemicModel

A type representing a contact process on network `g` with arbitrary dynamics.
"""
struct ContactProcess{T} <: AbstractEpidemicModel where T
    g::AbstractSimpleGraph
    dynamics::T
end

function init_rates(cp::ContactProcess, state)
    N = nv(cp.g)
    a = Vector{Float64}(undef, N)
    init_rates!(a, state, cp)
    return a
end

progressbar(::ContactProcess) = true


text:
using PkgBugDemo
using Test


text:
struct CoupledMolecule <: Simulation
    constructor::AbstractMoleculeConstructor
    mol::Molecule
    eig
    μ
    Δ
    H
    H0
    S0
    γ
    Γ
    L
end
Base.show(io::IO, sim::LiteSimulation) = Base.print(io, "LiteSimulation(...)")