using Celeste: Model
using Base.Test

"""
test log_prob.jl and log_prob_util.jl
"""

#####################
## Helper functions #
#####################

function test_that_star_truth_is_most_likely_log_prob()

    # init ground truth star
    images, ea, body = SampleData.true_star_init()

    # turn list of catalog entries a list of LatentStateParams
    source_states = [Model.catalog_entry_to_latent_state_params(body[s])
                     for s in 1:length(body)]

    # create logpdf function handle
    star_logpdf, star_logprior =
        Model.make_star_logpdf(ea.images, ea.S, ea.N,
                               source_states,
                               ea.patches, ea.active_sources,
                               ea.psf_K, ea.num_allowed_sd)

    # extract the star-specific parameters from source_states[1] for the
    # logpdf function
    star_state = Model.extract_star_state(source_states[1])
    println(star_state)
    best_ll = star_logpdf(star_state)
    println("Best ll ", best_ll)

    # perterb lnr
    lnr = star_state[1]
    for bad_lnr in [lnr*.1, lnr*.5, 1.5*lnr, 2.0*lnr]
        bad_state = deepcopy(star_state)
        bad_state[1] = bad_lnr
        @test best_ll > star_logpdf(bad_state)
    end

    # perturb ra, dec
    ra, dec = star_state[end-1], star_state[end]
    for bad_ra in [.5*ra, .8*ra, 1.1*ra, 1.5*ra]
        bad_state = deepcopy(star_state)
        bad_state[end-1] = bad_ra
        @test best_ll > star_logpdf(bad_state)
    end

    # perturb colors
    for scale in [.5, 2.2]

      # generate a random direction in param space
      r = randn(length(star_state))
      r = r / sqrt(sum(r.*r))

      bad_state = star_state + scale * r
      @test best_ll > star_logpdf(bad_state)
    end
end

function test_that_gal_truth_is_most_likely_log_prob()
    # init ground truth star
    images, ea, body = gen_sample_galaxy_dataset()

    # turn list of catalog entries a list of LatentStateParams
    source_states = [Model.catalog_entry_to_latent_state_params(body[s])
                     for s in 1:length(body)]
    println("  active source params: ", body[1])
    println("  corresponding ls    : ", source_states[1])

    # create logpdf function handle
    gal_logpdf, gal_logprior =
        Model.make_galaxy_logpdf(ea.images, ea.S, ea.N,
                                 source_states,
                                 ea.patches, ea.active_sources,
                                 ea.psf_K, ea.num_allowed_sd)

    # extract the star-specific parameters from source_states[1] for the
    # logpdf function
    gal_state = Model.extract_galaxy_state(source_states[1])
    println(gal_state)
    best_ll = gal_logpdf(gal_state)
    println("Best ll ", best_ll)

    # unpack true gal parameters
    lnr, col, u, shape = gal_state[1], gal_state[2:5],
                         gal_state[6:7], gal_state[8:11]

    # perterb lnr
    for bad_lnr in [lnr*.1, lnr*.5, 1.5*lnr, 2.0*lnr]
        bad_state = deepcopy(gal_state)
        bad_state[1] = bad_lnr
        @test best_ll > gal_logpdf(bad_state)
    end

    # perturb location
    ra, dec = u[1], u[2]
    for bad_ra in [.5*ra, .8*ra, 1.1*ra, 1.5*ra]
        bad_state = deepcopy(gal_state)
        bad_state[6] = bad_ra
        @test best_ll > gal_logpdf(bad_state)
    end

    # perturb colors
    for scale in [.5, 2.2]
      # generate a random direction in param space
      r = randn(length(gal_state))
      r = r / sqrt(sum(r.*r))

      bad_state = gal_state + scale * r
      @test best_ll > gal_logpdf(bad_state)
    end

    # perturb galaxy shape parameters
    gdev, gab, gangle, gscale = shape
    for bad_shape_scale in [.5, .8, 1.2, 1.3]
        # generate a random direction in param space
        r = randn(length(shape))
        r = r / sqrt(sum(r.*r))
        bad_shape = shape + bad_shape_scale * r

        bad_state = deepcopy(gal_state)
        bad_state[8:11] = bad_shape
        @test best_ll > gal_logpdf(bad_state)
    end
end

function test_color_flux_transform()
    # fluxes --- positive poisson rates
    fluxes = [1., 200., 300., 20., 5.]
    lnr, colors = Model.fluxes_to_colors(fluxes)

    # transform back to colors
    fluxes_back = Model.colors_to_fluxes(lnr, colors)
    for i in 1:length(fluxes)
        @test isapprox(fluxes[i], fluxes_back[i])
    end
end

function test_sigmoid_logit()
    as = [-10., -1., -.001, .001, 1., 10.]
    for a in as

        # test that sigmoid is pushing between 0 and 1
        sig_a = Model.sigmoid(a)
        @test (sig_a >= 0) && (sig_a <= 1)

        # test that logit is equal to inverse sigmoid
        a_prime = Model.logit(sig_a)
        @test isapprox(a, a_prime)

    end
end

function test_gal_shape_constrain()
    con_gal_shape = [.1, .1, .1, .1]
    unc_gal_shape = Model.unconstrain_gal_shape(con_gal_shape)
    back_gal_shape = Model.constrain_gal_shape(unc_gal_shape)

    # test equality of inverse map
    @test length(con_gal_shape) == length(back_gal_shape)
    for i in 1:length(unc_gal_shape)
        @test isapprox(con_gal_shape[i], back_gal_shape[i])
    end
end

####################################
test_sigmoid_logit()
test_gal_shape_constrain()
test_color_flux_transform()
test_that_star_truth_is_most_likely_log_prob()
test_that_gal_truth_is_most_likely_log_prob()


using MathOptInterface
const MOI = MathOptInterface
const CI = MOI.ConstraintIndex
const VI = MOI.VariableIndex

const MOIU = MOI.Utilities

"""
    ADMMIterations()

The number of ADMM iterations completed during the solve.
"""
struct ADMMIterations <: MOI.AbstractModelAttribute
end

MOI.is_set_by_optimize(::ADMMIterations) = true

mutable struct MOISolution
    ret_val::Int
    raw_status::String
    primal::Vector{Float64}
    dual::Vector{Float64}
    slack::Vector{Float64}
    objective_value::Float64
    dual_objective_value::Float64
    objective_constant::Float64
    solve_time_sec::Float64
    iterations::Int
end
MOISolution() = MOISolution(0, # SCS_UNFINISHED
                            "", Float64[], Float64[], Float64[], NaN, NaN, NaN,
                            0.0, 0)

# Used to build the data with allocate-load during `copy_to`.
# When `optimize!` is called, a the data is passed to SCS
# using `SCS_solve` and the `ModelData` struct is discarded
mutable struct ModelData
    m::Int # Number of rows/constraints
    n::Int # Number of cols/variables
    I::Vector{Int} # List of rows
    J::Vector{Int} # List of cols
    V::Vector{Float64} # List of coefficients
    b::Vector{Float64} # constants
    objective_constant::Float64 # The objective is min c'x + objective_constant
    c::Vector{Float64}
end

# This is tied to SCS's internal representation
mutable struct ConeData
    f::Int # number of linear equality constraints
    l::Int # length of LP cone
    q::Int # length of SOC cone
    qa::Vector{Int} # array of second-order cone constraints
    s::Int # length of SD cone
    sa::Vector{Int} # array of semi-definite constraints
    ep::Int # number of primal exponential cone triples
    ed::Int # number of dual exponential cone triples
    p::Vector{Float64} # array of power cone params
    nrows::Dict{Int, Int} # The number of rows of each vector sets, this is used by `constrrows` to recover the number of rows used by a constraint when getting `ConstraintPrimal` or `ConstraintDual`
    function ConeData()
        new(0, 0, 0, Int[], 0, Int[], 0, 0, Float64[], Dict{Int, Int}())
    end
end

mutable struct Optimizer <: MOI.AbstractOptimizer
    cone::ConeData
    maxsense::Bool
    data::Union{Nothing, ModelData} # only non-Void between MOI.copy_to and MOI.optimize!
    sol::MOISolution
    silent::Bool
    options::Dict{Symbol, Any}
    function Optimizer(; kwargs...)
        optimizer = new(ConeData(), false, nothing, MOISolution(), false,
                        Dict{Symbol, Any}())
        for (key, value) in kwargs
            MOI.set(optimizer, MOI.RawParameter(key), value)
        end
        return optimizer
    end

end

MOI.get(::Optimizer, ::MOI.SolverName) = "SCS"

function MOI.set(optimizer::Optimizer, param::MOI.RawParameter, value)
    optimizer.options[param.name] = value
end
function MOI.get(optimizer::Optimizer, param::MOI.RawParameter)
    return optimizer.options[param.name]
end

MOI.supports(::Optimizer, ::MOI.Silent) = true
function MOI.set(optimizer::Optimizer, ::MOI.Silent, value::Bool)
    optimizer.silent = value
end
MOI.get(optimizer::Optimizer, ::MOI.Silent) = optimizer.silent

function MOI.is_empty(optimizer::Optimizer)
    !optimizer.maxsense && optimizer.data === nothing
end
function MOI.empty!(optimizer::Optimizer)
    optimizer.maxsense = false
    optimizer.data = nothing # It should already be nothing except if an error is thrown inside copy_to
    optimizer.sol.ret_val = 0
end

MOIU.supports_allocate_load(::Optimizer, copy_names::Bool) = !copy_names

function MOI.supports(::Optimizer,
                      ::Union{MOI.ObjectiveSense,
                              MOI.ObjectiveFunction{MOI.SingleVariable},
                              MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}})
    return true
end

function MOI.supports(::Optimizer, ::MOI.VariablePrimalStart,
                      ::Type{MOI.VariableIndex})
    return true
end

function MOI.supports(::Optimizer,
                      ::Union{MOI.ConstraintPrimalStart,
                              MOI.ConstraintDualStart},
                      ::Type{<:MOI.ConstraintIndex})
    return true
end

function MOI.supports_constraint(
    ::Optimizer,
    ::Type{<:MOI.VectorAffineFunction{Float64}},
    ::Type{<:Union{MOI.Zeros, MOI.Nonnegatives, MOI.SecondOrderCone,
                   MOI.ExponentialCone, MOI.DualExponentialCone,
                   MOI.PositiveSemidefiniteConeTriangle,
                   MOI.PowerCone, MOI.DualPowerCone}})
    return true
end

function MOI.copy_to(dest::Optimizer, src::MOI.ModelLike; kws...)
    return MOIU.automatic_copy_to(dest, src; kws...)
end

using SparseArrays

# Computes cone dimensions
function constroffset(cone::ConeData,
                      ci::CI{<:MOI.AbstractFunction, MOI.Zeros})
    return ci.value
end
#_allocate_constraint: Allocate indices for the constraint `f`-in-`s`
# using information in `cone` and then update `cone`
function _allocate_constraint(cone::ConeData, f, s::MOI.Zeros)
    ci = cone.f
    cone.f += MOI.dimension(s)
    return ci
end
function constroffset(cone::ConeData,
                      ci::CI{<:MOI.AbstractFunction, MOI.Nonnegatives})
    return cone.f + ci.value
end
function _allocate_constraint(cone::ConeData, f, s::MOI.Nonnegatives)
    ci = cone.l
    cone.l += MOI.dimension(s)
    return ci
end
function constroffset(cone::ConeData,
                      ci::CI{<:MOI.AbstractFunction, MOI.SecondOrderCone})
    return cone.f + cone.l + ci.value
end
function _allocate_constraint(cone::ConeData, f, s::MOI.SecondOrderCone)
    push!(cone.qa, s.dimension)
    ci = cone.q
    cone.q += MOI.dimension(s)
    return ci
end
function constroffset(cone::ConeData,
                      ci::CI{<:MOI.AbstractFunction,
                             MOI.PositiveSemidefiniteConeTriangle})
    return cone.f + cone.l + cone.q + ci.value
end
function _allocate_constraint(cone::ConeData, f,
                              s::MOI.PositiveSemidefiniteConeTriangle)
    push!(cone.sa, s.side_dimension)
    ci = cone.s
    cone.s += MOI.dimension(s)
    return ci
end
function constroffset(cone::ConeData,
                      ci::CI{<:MOI.AbstractFunction, MOI.ExponentialCone})
    return cone.f + cone.l + cone.q + cone.s + ci.value
end
function _allocate_constraint(cone::ConeData, f, s::MOI.ExponentialCone)
    ci = 3cone.ep
    cone.ep += 1
    return ci
end
function constroffset(cone::ConeData,
                      ci::CI{<:MOI.AbstractFunction, MOI.DualExponentialCone})
    return cone.f + cone.l + cone.q + cone.s + 3cone.ep + ci.value
end
function _allocate_constraint(cone::ConeData, f, s::MOI.DualExponentialCone)
    ci = 3cone.ed
    cone.ed += 1
    return ci
end
function constroffset(cone::ConeData,
                      ci::CI{<:MOI.AbstractFunction, <:MOI.PowerCone})
    return cone.f + cone.l + cone.q + cone.s + 3cone.ep + 3cone.ed + ci.value
end
function _allocate_constraint(cone::ConeData, f, s::MOI.PowerCone)
    ci = length(cone.p)
    push!(cone.p, s.exponent)
    return ci
end
function constroffset(cone::ConeData,
                      ci::CI{<:MOI.AbstractFunction, <:MOI.DualPowerCone})
    return cone.f + cone.l + cone.q + cone.s + 3cone.ep + 3cone.ed + ci.value
end
function _allocate_constraint(cone::ConeData, f, s::MOI.DualPowerCone)
    ci = length(cone.p)
    # SCS' convention: dual cones have a negative exponent.
    push!(cone.p, -s.exponent)
    return ci
end
function constroffset(optimizer::Optimizer, ci::CI)
    return constroffset(optimizer.cone, ci::CI)
end
function MOIU.allocate_constraint(optimizer::Optimizer, f::F, s::S) where {F <: MOI.AbstractFunction, S <: MOI.AbstractSet}
    return CI{F, S}(_allocate_constraint(optimizer.cone, f, s))
end

# Vectorized length for matrix dimension n
sympackedlen(n) = div(n*(n+1), 2)
# Matrix dimension for vectorized length n
sympackeddim(n) = div(isqrt(1+8n) - 1, 2)
trimap(i::Integer, j::Integer) = i < j ? trimap(j, i) : div((i-1)*i, 2) + j
trimapL(i::Integer, j::Integer, n::Integer) = i < j ? trimapL(j, i, n) : i + div((2n-j) * (j-1), 2)
function _sympackedto(x, n, mapfrom, mapto)
    @assert length(x) == sympackedlen(n)
    y = similar(x)
    for i in 1:n, j in 1:i
        y[mapto(i, j)] = x[mapfrom(i, j)]
    end
    y
end
sympackedLtoU(x, n=sympackeddim(length(x))) = _sympackedto(x, n, (i, j) -> trimapL(i, j, n), trimap)
sympackedUtoL(x, n=sympackeddim(length(x))) = _sympackedto(x, n, trimap, (i, j) -> trimapL(i, j, n))

function sympackedUtoLidx(x::AbstractVector{<:Integer}, n)
    y = similar(x)
    map = sympackedLtoU(1:sympackedlen(n), n)
    for i in eachindex(y)
        y[i] = map[x[i]]
    end
    y
end

# Scale coefficients depending on rows index
# rows: List of row indices
# coef: List of corresponding coefficients
# d: dimension of set
# rev: if true, we unscale instead (e.g. divide by √2 instead of multiply for PSD cone)
function _scalecoef(rows, coef, d, rev)
    scaling = rev ? 1 / √2 : 1 * √2
    output = copy(coef)
    diagidx = BitSet()
    for i in 1:d
        push!(diagidx, trimap(i, i))
    end
    for i in 1:length(output)
        if !(rows[i] in diagidx)
            output[i] *= scaling
        end
    end
    return output
end
# Unscale the coefficients in `coef` with respective rows in `rows` for a set `s`
scalecoef(rows, coef, s) = _scalecoef(rows, coef, MOI.dimension(s), false)
# Unscale the coefficients in `coef` with respective rows in `rows` for a set of type `S` with dimension `d`
unscalecoef(rows, coef, d) = _scalecoef(rows, coef, sympackeddim(d), true)

output_index(t::MOI.VectorAffineTerm) = t.output_index
variable_index_value(t::MOI.ScalarAffineTerm) = t.variable_index.value
variable_index_value(t::MOI.VectorAffineTerm) = variable_index_value(t.scalar_term)
coefficient(t::MOI.ScalarAffineTerm) = t.coefficient
coefficient(t::MOI.VectorAffineTerm) = coefficient(t.scalar_term)
# constrrows: Recover the number of rows used by each constraint.
# When, the set is available, simply use MOI.dimension
constrrows(s::MOI.AbstractVectorSet) = 1:MOI.dimension(s)
# When only the index is available, use the `optimizer.ncone.nrows` field
constrrows(optimizer::Optimizer, ci::CI{<:MOI.AbstractVectorFunction, <:MOI.AbstractVectorSet}) = 1:optimizer.cone.nrows[constroffset(optimizer, ci)]

orderval(val, s) = val
function orderval(val, s::MOI.PositiveSemidefiniteConeTriangle)
    sympackedUtoL(val, s.side_dimension)
end
orderidx(idx, s) = idx
function orderidx(idx, s::MOI.PositiveSemidefiniteConeTriangle)
    sympackedUtoLidx(idx, s.side_dimension)
end
function MOIU.load_constraint(optimizer::Optimizer, ci::MOI.ConstraintIndex, f::MOI.VectorAffineFunction, s::MOI.AbstractVectorSet)
    A = sparse(output_index.(f.terms), variable_index_value.(f.terms), coefficient.(f.terms))
    # sparse combines duplicates with + but does not remove zeros created so we call dropzeros!
    dropzeros!(A)
    I, J, V = findnz(A)
    offset = constroffset(optimizer, ci)
    rows = constrrows(s)
    optimizer.cone.nrows[offset] = length(rows)
    i = offset .+ rows
    b = f.constants
    if s isa MOI.PositiveSemidefiniteConeTriangle
        # FIXME shouldn't scale before ?
        b = sympackedUtoL(b, s.side_dimension)
        b = scalecoef(rows, b, s)
        V = scalecoef(I, V, s)
        I = sympackedUtoLidx(I, s.side_dimension)
    end
    # The SCS format is b - Ax ∈ cone
    optimizer.data.b[i] = b
    append!(optimizer.data.I, offset .+ I)
    append!(optimizer.data.J, J)
    append!(optimizer.data.V, -V)
end

function MOIU.allocate_variables(optimizer::Optimizer, nvars::Integer)
    optimizer.cone = ConeData()
    VI.(1:nvars)
end

function MOIU.load_variables(optimizer::Optimizer, nvars::Integer)
    cone = optimizer.cone
    m = cone.f + cone.l + cone.q + cone.s + 3cone.ep + 3cone.ed + 3length(cone.p)
    I = Int[]
    J = Int[]
    V = Float64[]
    b = zeros(m)
    c = zeros(nvars)
    optimizer.data = ModelData(m, nvars, I, J, V, b, 0., c)
    # `optimizer.sol` contains the result of the previous optimization.
    # It is used as a warm start if its length is the same, e.g.
    # probably because no variable and/or constraint has been added.
    if length(optimizer.sol.primal) != nvars
        optimizer.sol.primal = zeros(nvars)
    end
    @assert length(optimizer.sol.dual) == length(optimizer.sol.slack)
    if length(optimizer.sol.dual) != m
        optimizer.sol.dual = zeros(m)
        optimizer.sol.slack = zeros(m)
    end
end

function MOIU.allocate(::Optimizer, ::MOI.VariablePrimalStart,
                       ::MOI.VariableIndex, ::Union{Nothing, Float64})
end
function MOIU.allocate(::Optimizer, ::MOI.ConstraintPrimalStart,
                       ::MOI.ConstraintIndex,
                       ::Union{Nothing, AbstractVector{Float64}})
end
function MOIU.allocate(::Optimizer, ::MOI.ConstraintDualStart,
                       ::MOI.ConstraintIndex,
                       ::Union{Nothing, AbstractVector{Float64}})
end
function MOIU.allocate(optimizer::Optimizer, ::MOI.ObjectiveSense, sense::MOI.OptimizationSense)
    optimizer.maxsense = sense == MOI.MAX_SENSE
end
function MOIU.allocate(::Optimizer, ::MOI.ObjectiveFunction,
                       ::MOI.Union{MOI.SingleVariable,
                                   MOI.ScalarAffineFunction{Float64}})
end

function MOIU.load(::Optimizer, ::MOI.VariablePrimalStart,
                   ::MOI.VariableIndex, ::Nothing)
end
function MOIU.load(optimizer::Optimizer, ::MOI.VariablePrimalStart,
                   vi::MOI.VariableIndex, value::Float64)
    optimizer.sol.primal[vi.value] = value
end
function MOIU.load(::Optimizer, ::MOI.ConstraintPrimalStart,
                   ::MOI.ConstraintIndex, ::Nothing)
end
function MOIU.load(optimizer::Optimizer, ::MOI.ConstraintPrimalStart,
                   ci::MOI.ConstraintIndex, value)
    offset = constroffset(optimizer, ci)
    rows = constrrows(optimizer, ci)
    optimizer.sol.slack[offset .+ rows] .= value
end
function MOIU.load(::Optimizer, ::MOI.ConstraintDualStart,
                   ::MOI.ConstraintIndex, ::Nothing)
end
function MOIU.load(optimizer::Optimizer, ::MOI.ConstraintDualStart,
                   ci::MOI.ConstraintIndex, value)
    offset = constroffset(optimizer, ci)
    rows = constrrows(optimizer, ci)
    optimizer.sol.dual[offset .+ rows] .= value
end
function MOIU.load(::Optimizer, ::MOI.ObjectiveSense, ::MOI.OptimizationSense)
end
function MOIU.load(optimizer::Optimizer, ::MOI.ObjectiveFunction,
                   f::MOI.SingleVariable)
    MOIU.load(optimizer,
              MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}(),
              MOI.ScalarAffineFunction{Float64}(f))
end
function MOIU.load(optimizer::Optimizer, ::MOI.ObjectiveFunction,
                   f::MOI.ScalarAffineFunction)
    c0 = Vector(sparsevec(variable_index_value.(f.terms), coefficient.(f.terms),
                          optimizer.data.n))
    optimizer.data.objective_constant = f.constant
    optimizer.data.c = optimizer.maxsense ? -c0 : c0
    return nothing
end

function MOI.optimize!(optimizer::Optimizer)
    cone = optimizer.cone
    m = optimizer.data.m
    n = optimizer.data.n
    A = sparse(optimizer.data.I, optimizer.data.J, optimizer.data.V)
    b = optimizer.data.b
    objective_constant = optimizer.data.objective_constant
    c = optimizer.data.c
    optimizer.data = nothing # Allows GC to free optimizer.data before A is loaded to SCS

    options = optimizer.options
    if optimizer.silent
        options = copy(options)
        options[:verbose] = 0
    end

    linear_solver, options = sanitize_SCS_options(options)

    sol = SCS_solve(linear_solver, m, n, A, b, c,
                    cone.f, cone.l, cone.qa, cone.sa, cone.ep, cone.ed, cone.p,
                    optimizer.sol.primal, optimizer.sol.dual,
                    optimizer.sol.slack; options...)

    ret_val = sol.ret_val
    primal = sol.x
    dual = sol.y
    slack = sol.s
    objective_value = (optimizer.maxsense ? -1 : 1) * sol.info.pobj
    dual_objective_value = (optimizer.maxsense ? -1 : 1) * sol.info.dobj
    solve_time = (sol.info.setupTime + sol.info.solveTime) / 1000
    optimizer.sol = MOISolution(ret_val, raw_status(sol.info), primal, dual,
                                slack, objective_value, dual_objective_value,
                                objective_constant, solve_time, sol.info.iter)
end

function MOI.get(optimizer::Optimizer, ::MOI.SolveTime)
    return optimizer.sol.solve_time_sec
end
function MOI.get(optimizer::Optimizer, ::MOI.RawStatusString)
    return optimizer.sol.raw_status
end

function MOI.get(optimizer::Optimizer, ::ADMMIterations)
    return optimizer.sol.iterations
end

# Implements getter for result value and statuses
# SCS returns one of the following integers:
# -7 SCS_INFEASIBLE_INACCURATE
# -6 SCS_UNBOUNDED_INACCURATE
# -5 SCS_SIGINT
# -4 SCS_FAILED
# -3 SCS_INDETERMINATE
# -2 SCS_INFEASIBLE  : primal infeasible, dual unbounded
# -1 SCS_UNBOUNDED   : primal unbounded, dual infeasible
#  0 SCS_UNFINISHED  : never returned, used as placeholder
#  1 SCS_SOLVED
#  2 SCS_SOLVED_INACCURATE
function MOI.get(optimizer::Optimizer, ::MOI.TerminationStatus)
    s = optimizer.sol.ret_val
    @assert -7 <= s <= 2
    if s == -7
        return MOI.ALMOST_INFEASIBLE
    elseif s == -6
        return MOI.ALMOST_DUAL_INFEASIBLE
    elseif s == 2
        return MOI.ALMOST_OPTIMAL
    elseif s == -5
        return MOI.INTERRUPTED
    elseif s == -4
        return MOI.NUMERICAL_ERROR
    elseif s == -3
        return MOI.SLOW_PROGRESS
    elseif s == -2
        return MOI.INFEASIBLE
    elseif s == -1
        return MOI.DUAL_INFEASIBLE
    elseif s == 1
        return MOI.OPTIMAL
    else
        @assert s == 0
        return MOI.OPTIMIZE_NOT_CALLED
    end
end

function MOI.get(optimizer::Optimizer, attr::MOI.ObjectiveValue)
    MOI.check_result_index_bounds(optimizer, attr)
    value = optimizer.sol.objective_value
    if !MOIU.is_ray(MOI.get(optimizer, MOI.PrimalStatus()))
        value += optimizer.sol.objective_constant
    end
    return value
end
function MOI.get(optimizer::Optimizer, attr::MOI.DualObjectiveValue)
    MOI.check_result_index_bounds(optimizer, attr)
    value = optimizer.sol.dual_objective_value
    if !MOIU.is_ray(MOI.get(optimizer, MOI.DualStatus()))
        value += optimizer.sol.objective_constant
    end
    return value
end

function MOI.get(optimizer::Optimizer, attr::MOI.PrimalStatus)
    if attr.N > MOI.get(optimizer, MOI.ResultCount())
        return MOI.NO_SOLUTION
    end
    s = optimizer.sol.ret_val
    if s in (-3, 1, 2)
        MOI.FEASIBLE_POINT
    elseif s in (-6, -1)
        MOI.INFEASIBILITY_CERTIFICATE
    else
        MOI.INFEASIBLE_POINT
    end
end
function MOI.get(optimizer::Optimizer, attr::MOI.VariablePrimal, vi::VI)
    MOI.check_result_index_bounds(optimizer, attr)
    optimizer.sol.primal[vi.value]
end
function MOI.get(optimizer::Optimizer, attr::MOI.VariablePrimal, vi::Vector{VI})
    return MOI.get.(optimizer, attr, vi)
end
function MOI.get(optimizer::Optimizer, attr::MOI.ConstraintPrimal,
                 ci::CI{<:MOI.AbstractFunction, S}) where S <: MOI.AbstractSet
    MOI.check_result_index_bounds(optimizer, attr)
    offset = constroffset(optimizer, ci)
    rows = constrrows(optimizer, ci)
    primal = optimizer.sol.slack[offset .+ rows]
    if S == MOI.PositiveSemidefiniteConeTriangle
        primal = sympackedLtoU(primal)
        primal = unscalecoef(rows, primal, length(rows))
    end
    return primal
end

function MOI.get(optimizer::Optimizer, attr::MOI.DualStatus)
    if attr.N > MOI.get(optimizer, MOI.ResultCount())
        return MOI.NO_SOLUTION
    end
    s = optimizer.sol.ret_val
    if s in (-3, 1, 2)
        MOI.FEASIBLE_POINT
    elseif s in (-7, -2)
        MOI.INFEASIBILITY_CERTIFICATE
    else
        MOI.INFEASIBLE_POINT
    end
end
function MOI.get(optimizer::Optimizer, attr::MOI.ConstraintDual,
                 ci::CI{<:MOI.AbstractFunction, S}) where S <: MOI.AbstractSet
    MOI.check_result_index_bounds(optimizer, attr)
    offset = constroffset(optimizer, ci)
    rows = constrrows(optimizer, ci)
    dual = optimizer.sol.dual[offset .+ rows]
    if S == MOI.PositiveSemidefiniteConeTriangle
        dual = sympackedLtoU(dual)
        dual = unscalecoef(rows, dual, length(rows))
    end
    return dual
end

MOI.get(optimizer::Optimizer, ::MOI.ResultCount) = 1


"""
    IntervalConstantSolver(; θ = 0.5, timestep)

BTPDE solver specialized on intervalwise constant `ScalarGradient`s, e.g [`PGSE`](@ref),
[`DoublePGSE`](@ref). Raises an error for other gradients.
"""
Base.@kwdef struct IntervalConstantSolver{T}
    θ::T = 0.5
    timestep::T
end


using Swizzles.ExtrudedArrays
using Swizzles.NamedArrays
using Swizzles.ValArrays
using Swizzles.ArrayifiedArrays
using Swizzles.Virtuals

using Base.Broadcast: broadcastable, Broadcasted, materialize

using Swizzles

@testset "Virtuals" begin
    for x in [@_((1, 2, 3) .+ [1, 2, 3]'),
              @_(0 .+ [1, 2, 3] .+ [1 2 3; 4 5 6; 7 8 9]),
              Swizzle(+)(@_(0 .+ [1, 2, 3] .+ [1 2 3; 4 5 6; 7 8 9])),
              Swizzle(+, 1)(@_(0 .+ [1, 2, 3] .+ [1 2 3; 4 5 6; 7 8 9]))]
        @test kept.(keeps(virtualize(:root, typeof(lift_keeps(x))))) == kept.(keeps(x))
        @test eltype(virtualize(:root, typeof(lift_keeps(x)))) == eltype(x)
    end
end


@testset "Bitmask flags" begin
    test(BitmaskDefinition, bitmask_by_name, :VkQueryPipelineStatisticFlagBits, :(
        @bitmask_flag QueryPipelineStatisticFlag::UInt32 begin
            QUERY_PIPELINE_STATISTIC_INPUT_ASSEMBLY_VERTICES_BIT = 1
            QUERY_PIPELINE_STATISTIC_INPUT_ASSEMBLY_PRIMITIVES_BIT = 2
            QUERY_PIPELINE_STATISTIC_VERTEX_SHADER_INVOCATIONS_BIT = 4
            QUERY_PIPELINE_STATISTIC_GEOMETRY_SHADER_INVOCATIONS_BIT = 8
            QUERY_PIPELINE_STATISTIC_GEOMETRY_SHADER_PRIMITIVES_BIT = 16
            QUERY_PIPELINE_STATISTIC_CLIPPING_INVOCATIONS_BIT = 32
            QUERY_PIPELINE_STATISTIC_CLIPPING_PRIMITIVES_BIT = 64
            QUERY_PIPELINE_STATISTIC_FRAGMENT_SHADER_INVOCATIONS_BIT = 128
            QUERY_PIPELINE_STATISTIC_TESSELLATION_CONTROL_SHADER_PATCHES_BIT = 256
            QUERY_PIPELINE_STATISTIC_TESSELLATION_EVALUATION_SHADER_INVOCATIONS_BIT = 512
            QUERY_PIPELINE_STATISTIC_COMPUTE_SHADER_INVOCATIONS_BIT = 1024
        end
    ))

    test(BitmaskDefinition, bitmask_by_name, :VkSparseMemoryBindFlagBits, :(
        @bitmask_flag SparseMemoryBindFlag::UInt32 begin
            SPARSE_MEMORY_BIND_METADATA_BIT = 1
        end
    ))

    test(BitmaskDefinition, bitmask_by_name, :VkShaderStageFlagBits, :(
        @bitmask_flag ShaderStageFlag::UInt32 begin
            SHADER_STAGE_VERTEX_BIT = 1
            SHADER_STAGE_TESSELLATION_CONTROL_BIT = 2
            SHADER_STAGE_TESSELLATION_EVALUATION_BIT = 4
            SHADER_STAGE_GEOMETRY_BIT = 8
            SHADER_STAGE_FRAGMENT_BIT = 16
            SHADER_STAGE_COMPUTE_BIT = 32
            SHADER_STAGE_RAYGEN_BIT_KHR = 256
            SHADER_STAGE_ANY_HIT_BIT_KHR = 512
            SHADER_STAGE_CLOSEST_HIT_BIT_KHR = 1024
            SHADER_STAGE_MISS_BIT_KHR = 2048
            SHADER_STAGE_INTERSECTION_BIT_KHR = 4096
            SHADER_STAGE_CALLABLE_BIT_KHR = 8192
            SHADER_STAGE_TASK_BIT_NV = 64
            SHADER_STAGE_MESH_BIT_NV = 128
            SHADER_STAGE_SUBPASS_SHADING_BIT_HUAWEI = 16384
            SHADER_STAGE_ALL_GRAPHICS = $(Int(0x0000001f))
            SHADER_STAGE_ALL = $(Int(0x7fffffff))
        end
    ))

    test(BitmaskDefinition, bitmask_by_name, :VkCullModeFlagBits, :(
        @bitmask_flag CullModeFlag::UInt32 begin
            CULL_MODE_FRONT_BIT = 1
            CULL_MODE_BACK_BIT = 2
            CULL_MODE_NONE = 0
            CULL_MODE_FRONT_AND_BACK = 3
        end
    ))
end


abstract type EoSModel end
abstract type LatticeFluidModel <: EoSModel end
abstract type SAFTModel <: LatticeFluidModel end
abstract type GCSAFTModel <: SAFTModel end
abstract type CubicModel <: EoSModel end
abstract type ABCubicModel <: CubicModel end #cubics that have an exact polynomial form to solve Z roots, this excludes CPA
abstract type IdealModel <: EoSModel end
abstract type EmpiricHelmholtzModel <: EoSModel end
abstract type EoSParam end


include("balance_law_interface.jl")

using ..Courant
"""
    calculate_dt(dg, model, Q, Courant_number, direction, t)

For a given model, compute a time step satisying the nondiffusive Courant number
`Courant_number`
"""
function calculate_dt(dg, model, Q, Courant_number, t, direction)
    Δt = one(eltype(Q))
    CFL = courant(nondiffusive_courant, dg, model, Q, Δt, t, direction)
    return Courant_number / CFL
end

function create_conservative_state(balance_law::BalanceLaw, grid)
    topology = grid.topology
    # FIXME: Remove after updating CUDA
    h_vgeo = Array(grid.vgeo)
    FT = eltype(h_vgeo)
    Np = dofs_per_element(grid)
    DA = arraytype(grid)

    weights = view(h_vgeo, :, grid.Mid, :)
    weights = reshape(weights, size(weights, 1), 1, size(weights, 2))

    V = vars_state_conservative(balance_law, FT)
    state_conservative = MPIStateArray{FT, V}(
        topology.mpicomm,
        DA,
        Np,
        number_state_conservative(balance_law, FT),
        length(topology.elems),
        realelems = topology.realelems,
        ghostelems = topology.ghostelems,
        vmaprecv = grid.vmaprecv,
        vmapsend = grid.vmapsend,
        nabrtorank = topology.nabrtorank,
        nabrtovmaprecv = grid.nabrtovmaprecv,
        nabrtovmapsend = grid.nabrtovmapsend,
        weights = weights,
    )
    return state_conservative
end

function create_auxiliary_state(balance_law, grid)
    topology = grid.topology
    Np = dofs_per_element(grid)

    h_vgeo = Array(grid.vgeo)
    FT = eltype(h_vgeo)
    DA = arraytype(grid)

    weights = view(h_vgeo, :, grid.Mid, :)
    weights = reshape(weights, size(weights, 1), 1, size(weights, 2))

    V = vars_state_auxiliary(balance_law, FT)
    state_auxiliary = MPIStateArray{FT, V}(
        topology.mpicomm,
        DA,
        Np,
        number_state_auxiliary(balance_law, FT),
        length(topology.elems),
        realelems = topology.realelems,
        ghostelems = topology.ghostelems,
        vmaprecv = grid.vmaprecv,
        vmapsend = grid.vmapsend,
        nabrtorank = topology.nabrtorank,
        nabrtovmaprecv = grid.nabrtovmaprecv,
        nabrtovmapsend = grid.nabrtovmapsend,
        weights = weights,
    )

    dim = dimensionality(grid)
    polyorder = polynomialorder(grid)
    vgeo = grid.vgeo
    device = array_device(state_auxiliary)
    nrealelem = length(topology.realelems)
    event = Event(device)
    event = kernel_init_state_auxiliary!(device, min(Np, 1024), Np * nrealelem)(
        balance_law,
        Val(dim),
        Val(polyorder),
        state_auxiliary.data,
        vgeo,
        topology.realelems,
        dependencies = (event,),
    )
    event = MPIStateArrays.begin_ghost_exchange!(
        state_auxiliary;
        dependencies = event,
    )
    event = MPIStateArrays.end_ghost_exchange!(
        state_auxiliary;
        dependencies = event,
    )
    wait(device, event)

    return state_auxiliary
end

function create_gradient_state(balance_law, grid)
    topology = grid.topology
    Np = dofs_per_element(grid)

    h_vgeo = Array(grid.vgeo)
    FT = eltype(h_vgeo)
    DA = arraytype(grid)

    weights = view(h_vgeo, :, grid.Mid, :)
    weights = reshape(weights, size(weights, 1), 1, size(weights, 2))

    # TODO: Clean up this MPIStateArray interface...
    V = vars_state_gradient_flux(balance_law, FT)
    state_gradient_flux = MPIStateArray{FT, V}(
        topology.mpicomm,
        DA,
        Np,
        number_state_gradient_flux(balance_law, FT),
        length(topology.elems),
        realelems = topology.realelems,
        ghostelems = topology.ghostelems,
        vmaprecv = grid.vmaprecv,
        vmapsend = grid.vmapsend,
        nabrtorank = topology.nabrtorank,
        nabrtovmaprecv = grid.nabrtovmaprecv,
        nabrtovmapsend = grid.nabrtovmapsend,
        weights = weights,
    )

    return state_gradient_flux
end

function create_higher_order_states(balance_law, grid)
    topology = grid.topology
    Np = dofs_per_element(grid)

    h_vgeo = Array(grid.vgeo)
    FT = eltype(h_vgeo)
    DA = arraytype(grid)

    weights = view(h_vgeo, :, grid.Mid, :)
    weights = reshape(weights, size(weights, 1), 1, size(weights, 2))

    ngradlapstate = num_gradient_laplacian(balance_law, FT)
    # TODO: Clean up this MPIStateArray interface...
    V = vars_gradient_laplacian(balance_law, FT)
    Qhypervisc_grad = MPIStateArray{FT, V}(
        topology.mpicomm,
        DA,
        Np,
        3ngradlapstate,
        length(topology.elems),
        realelems = topology.realelems,
        ghostelems = topology.ghostelems,
        vmaprecv = grid.vmaprecv,
        vmapsend = grid.vmapsend,
        nabrtorank = topology.nabrtorank,
        nabrtovmaprecv = grid.nabrtovmaprecv,
        nabrtovmapsend = grid.nabrtovmapsend,
        weights = weights,
    )

    V = vars_hyperdiffusive(balance_law, FT)
    Qhypervisc_div = MPIStateArray{FT, V}(
        topology.mpicomm,
        DA,
        Np,
        ngradlapstate,
        length(topology.elems),
        realelems = topology.realelems,
        ghostelems = topology.ghostelems,
        vmaprecv = grid.vmaprecv,
        vmapsend = grid.vmapsend,
        nabrtorank = topology.nabrtorank,
        nabrtovmaprecv = grid.nabrtovmaprecv,
        nabrtovmapsend = grid.nabrtovmapsend,
        weights = weights,
    )
    return Qhypervisc_grad, Qhypervisc_div
end


__precompile__()
module Map_removeEQ

include("get_kurtosis.jl")
include("remove_eq.jl")

using .Get_kurtosis, .Remove_eq
using SeisIO, Dates, JLD2, Printf, FileIO
using Suppressor

export map_removeEQ

"""
loaddata(finame, path)
"""
function loaddata(finame, path)
    try
        @suppress_err return FileIO.load(finame, path)
    catch
        return false;
    end
end

"""
    ParallelEQremoval(dlid, InputDict::Dict)
    remove earthquake and save it into jld2 file.

"""
function map_removeEQ(dlid, InputDict::Dict)

    #store data

    finame                 =InputDict["finame"]
    IsKurtosisRemoval      =InputDict["IsKurtosisRemoval"]
    max_edgetaper_duration =InputDict["max_edgetaper_duration"]
    kurtosis_tw_sparse     =InputDict["kurtosis_tw_sparse"]
    kurtosis_timewindow    =InputDict["kurtosis_timewindow"]
    kurtosis_threshold     =InputDict["kurtosis_threshold"]
    IsSTALTARemoval        =InputDict["IsSTALTARemoval"]
    stalta_longtimewindow  =InputDict["stalta_longtimewindow"]
    stalta_threshold       =InputDict["stalta_threshold"]
    stalta_absoluteclip    =InputDict["stalta_absoluteclip"]
    max_wintaper_duration  =InputDict["max_wintaper_duration"]
    removal_shorttimewindow=InputDict["removal_shorttimewindow"]
    overlap                =InputDict["overlap"]
    Iswhiten               =InputDict["Iswhiten"]
    freqmin_whiten         =InputDict["freqmin_whiten"]
    freqmax_whiten         =InputDict["freqmax_whiten"]
    plot_kurtosis_α        =InputDict["plot_kurtosis_α"]
    plot_boxheight         =InputDict["plot_boxheight"]
    plot_span              =InputDict["plot_span"]
    plot_fmt               =InputDict["plot_fmt"]
    fodir                  =InputDict["fodir"]
    foname                 =InputDict["foname"]
    fopath                 =InputDict["fopath"]
    IsSaveFig              =InputDict["IsSaveFig"]

    DLtimestamplist        =InputDict["DLtimestamplist"]
    stationlist            =InputDict["stationlist"]
    NumofTimestamp         =InputDict["NumofTimestamp"]

    tstamp = DLtimestamplist[dlid]

    if mod(dlid, round(0.1*NumofTimestamp)+1) == 0
        println(@sprintf("start process %s", tstamp))
    end

    bt_getkurtosis = 0.0
    bt_removeeq = 0.0

    for st = stationlist
        #S = t[joinpath(tstamp, st)]
        st1 = replace(st, "-"=>"")

        # load raw data
        Stemp = loaddata(finame, joinpath(tstamp, st1))

        # skip if it does not exist
        if Stemp == false
            continue;
        end

        #convert if it's SeisChannel
        if Stemp isa SeisIO.SeisChannel
            Stemp = SeisData(Stemp)
        end

        SremEQ = SeisData()

        #loop by SeisChannel and save into temporal wile
        for Seisid = 1:Stemp.n

            Schan = Stemp[Seisid]

            if !haskey(Schan.misc, "dlerror")
                if iszero(Schan.x)
                    Schan.misc["dlerror"] = true
                else
                    Schan.misc["dlerror"] = false
                end
            end

            #---Transient error on combination between download margin and resampling---#
            # S.t may contain index 0 in t[:, 1]. in that case , this causes the error on later part, so discard that data.
            if any(Schan.t[:,1] .== 0)
                println("data gap index found.")
                println(Schan)
                Schan.misc["dlerror"] = true
            end
            #---------------------------------------------------------------------------#

            if Schan.misc["dlerror"] == 0

                dt = 1/Schan.fs
                tvec = collect(0:length(Schan.x)-1) * dt ./ 60 ./ 60

                #tapering to avoid instrumental edge artifacts
                #SeisIO.taper!(Schan,  t_max = max_edgetaper_duration, α=0.05) this is already done by download margin

                S1 = deepcopy(Schan)

                #set long window length to user input since last window of previous channel will have been adjusted
                S1.misc["eqtimewindow"] = fill(true, length(S1.x))

                if IsKurtosisRemoval
                    # compute kurtosis and detect
                    bt_1 = @elapsed S1 = Get_kurtosis.get_kurtosis(S1, float(kurtosis_timewindow), float(kurtosis_tw_sparse))

                    bt_2 = @elapsed S1 = Remove_eq.detect_eq_kurtosis(S1, tw=float(removal_shorttimewindow), kurtosis_threshold=float(kurtosis_threshold), overlap=float(overlap))

                    btsta_1 = 0

                    if IsSTALTARemoval
                        # detect earthquake and tremors by STA/LTA
                        btsta_1 = @elapsed S1 = Remove_eq.detect_eq_stalta(S1, float(stalta_longtimewindow), float(removal_shorttimewindow),
                                            float(stalta_threshold), float(overlap), float(stalta_absoluteclip),  InputDict, tstamp, st1)
                    end

                    if Iswhiten
                        Remove_eq.s_whiten!(S1, freqmin_whiten, freqmax_whiten)# apply spectral whitening on this channel
                    end

                    bt_3 = @elapsed S1 = Remove_eq.remove_eq(S1, Schan, float(plot_kurtosis_α), float(max_wintaper_duration),
                                    plot_boxheight, trunc(Int, plot_span), plot_fmt, fodir, tstamp, tvec, IsSaveFig)

                    bt_getkurtosis += bt_1
                    bt_removeeq += bt_2 + bt_3 + btsta_1

                    #println([bt_2, bt_3, btsta_1])
                    #if mod(dlid, round(0.1*NumofTimestamp)+1) == 0
                    #    println([bt_1, bt_2, btsta_1, bt_3])
                    #end

                else
                    #only STA/LTA
                    if IsSTALTARemoval

                        data.misc["kurtosis"] = zeros(length(S1.x))

                        bt_2 = @elapsed S1 = Remove_eq.detect_eq_stalta(S1, float(stalta_longtimewindow), float(removal_shorttimewindow),
                                            float(stalta_threshold), float(overlap), InputDict, tstamp, st1)

                        if Iswhiten
                            Remove_eq.s_whiten!(S1, freqmin_whiten, freqmax_whiten) # apply spectral whitening on this channel
                        end

                        bt_3 = @elapsed S1 = Remove_eq.remove_eq(S1, S, float(plot_kurtosis_α), float(max_wintaper_duration),
                                        plot_boxheight, trunc(Int, plot_span), plot_fmt, fodir, tstamp, tvec, IsSaveFig)

                        bt_getkurtosis += 0.0
                        bt_removeeq += bt_2 + bt_3

                    else
                        #no removal process.
                        @warn "Both 'IsKurtosisRemoval' and 'IsKurtosisRemoval' are false. No removal process is executed. Abort."
                        exit(0)
                    end

                end

                if InputDict["IsOutputRemovalFrac"]
                    #output removal fraction on this channel
                    eqidlist = S1.misc["eqtimewindow"][:]
                    numofremoval = sum(x->x==false, eqidlist, dims=1)
                    fractionofremoval = numofremoval[1] / length(eqidlist)

                    y, jd = parse.(Int64, split(InputDict["DLtimestamplist"][dlid], ".")[1:2])
                    tstamp_fname = replace(tstamp, ":" => "-")
                    fname_out = join([tstamp_fname, st1, "removalfraction","dat"], '.')
                    fopath = InputDict["removal_fractionpath"]*"/"*fname_out
                    open(fopath, "w") do io
                       write(io, @sprintf("%f\n", fractionofremoval))
                    end
                end

                if InputDict["dumptraces"]
                    #dump raw trace
                    tstamp_fname = replace(tstamp, ":" => "-")
                    fname_out = join([tstamp_fname, st1, "rawdata","dat"], '.')
                    fopath = InputDict["dumppath"]*"/"*fname_out
                    open(fopath, "w") do io
                        for i in 1:length(Schan.x)
                            write(io, @sprintf("%12.8f\n", Schan.x[i]))
                        end
                    end

                    #dump kurtosis trace
                    mkpath(InputDict["dumppath"])
                    tstamp_fname = replace(tstamp, ":" => "-")
                    fname_out = join([tstamp_fname, st1, "kurtosis","dat"], '.')
                    fopath = InputDict["dumppath"]*"/"*fname_out
                    open(fopath, "w") do io
                        for i in 1:length(S1.misc["kurtosis"])
                            write(io, @sprintf("%12.8f\n", S1.misc["kurtosis"][i]))
                        end
                    end

                    #dump removed trace
                    tstamp_fname = replace(tstamp, ":" => "-")
                    fname_out = join([tstamp_fname, st1, "remdata","dat"], '.')
                    fopath = InputDict["dumppath"]*"/"*fname_out
                    open(fopath, "w") do io
                        for i in 1:length(S1.x)
                            write(io, @sprintf("%12.8f\n", S1.x[i]))
                        end
                    end
                end

                #it's not allowed to save this into binary;
                delete!(S1.misc, "eqtimewindow")
                delete!(S1.misc, "kurtosis")

            else
                #download error found: save as it is.
                S1 = Schan
            end

            SremEQ += S1

        end

        # if some of SeisChannels in Stemp have a data, save temp file
        y, jd = parse.(Int64, split(InputDict["DLtimestamplist"][dlid], ".")[1:2])
        tstamp_fname = replace(tstamp, ":" => "-")
        fname_out = join([tstamp_fname, st1, "FDSNWS","dat"], '.')
        # save as intermediate binary file
        t_write = @elapsed SeisIO.wseis(InputDict["tmppath"]*"/"*fname_out, SremEQ)

    end

    return nothing
end

end


#production and acquisition

function DDCost!(options::zefiroOptions,sizes::zefiroSizes,data::zefiroData,output::zefiroOutput)
    key = (options.postprocess ? "postprocessing" : "evaluating")
    println(key*" decommissioning cost")
    decommissioning!(options,sizes,data,output)
    println(key*" wasteManagement cost")
    wasteManagement!(options,sizes,data,output)
    println(key*" siteClearance cost")
    siteClearance!(options,sizes,data,output)
    println(key*" postMonitoring cost")
    postMonitoring!(options,sizes,data,output)

    nothing
end

function decommissioning!(options::zefiroOptions,sizes::zefiroSizes,data::zefiroData,output::zefiroOutput)

    if options.decommissioning == 0
        if options.postprocess
            nothing
        else
            CDDport = (data.Cddport + data.Nddport .* data.Lddport).*transpose(data.turbineQuantity)
            Cremov = (data.Nremove .* data.Vremove).*transpose(data.turbineQuantity)
            output.DECOM .+= CDDport + Cremov
            output.decommissioning .+= CDDport + Cremov
        end
    else
        error("decommissioning mode not found.")
    end
    nothing

end

function wasteManagement!(options::zefiroOptions,sizes::zefiroSizes,data::zefiroData,output::zefiroOutput)

    if options.wasteManagement == 0
        if options.postprocess
            nothing
        else
            Cwproc = data.Wjproc .* data.Cjproc
            Cwtrans = data.Wjtrans .* data.Cjproc ./ data.Wttrans
            Cland = data.Wnr .* data.Cnr
            Csv = data.Wr .* data.SV
            output.DECOM .+= transpose(Cwproc + Cwtrans + Cland - Csv).*transpose(data.turbineQuantity)
            output.wasteManagement .+= transpose(Cwproc + Cwtrans + Cland - Csv).*transpose(data.turbineQuantity)
        end
    else
        error("wasteManagement mode not found.")
    end
    nothing

end

function siteClearance!(options::zefiroOptions,sizes::zefiroSizes,data::zefiroData,output::zefiroOutput)

    if options.siteClearance == 0
        if options.postprocess
            nothing
        else
            output.DECOM .+= data.Asc.*data.Cscunit.*transpose(data.turbineQuantity)
            output.siteClearance .+= data.Asc.*data.Cscunit.*transpose(data.turbineQuantity)
        end
    else
        error("siteClearance mode not found.")
    end
    nothing

end

function postMonitoring!(options::zefiroOptions,sizes::zefiroSizes,data::zefiroData,output::zefiroOutput)

    if options.postMonitoring == 0
        if options.postprocess
            nothing
        else
            output.DECOM .+= data.postCost
            output.postMonitoring .+= data.postCost
        end
    else
        error("postMonitoring mode not found.")
    end
    nothing

end


module Linda

using Printf

import MathOptInterface
const MOI = MathOptInterface

# package code goes here
include("env.jl")
include("column.jl")
include("oracle.jl")
include("master.jl")
include("colgen.jl")

end # module


"""
    $(TYPEDEF)
An abstract type for all VLBI interfermetry data types. See [EHTVisibilityDatum](@ref) for an example.
"""
abstract type AbstractInterferometryDatum{T} end

abstract type AbstractVisibilityDatum{T} <: AbstractInterferometryDatum{T} end
abstract type AbstractLinearPolDatum{S<:AbstractVisibilityDatum, T} <: AbstractInterferometryDatum{T} end
abstract type AbstractCrossPolDatum{S,T} <: AbstractInterferometryDatum{T} end

abstract type ClosureProducts{T} <: AbstractInterferometryDatum{T} end

abstract type Observation{T} end

using DelimitedFiles
using AstroTime: modified_julian

export uvpositions, stations, getdata, arrayconfig,
       getuv, baselines, rescaleuv!

"""
    $(TYPEDEF)
This defined the abstract type for an array configuration. Namely, baseline
times, SEFD's, bandwidth, observation frequencies, etc.
"""
abstract type ArrayConfiguration end

"""
    $(TYPEDEF)
Stores all the non-visibility data products for an EHT array. This is useful when evaluating
model visibilities.

# $(FIELDS)
"""
struct EHTArrayConfiguration{S,F,T<:AbstractArray} <: ArrayConfiguration
    """
    List of stations in the array
    """
    stations::S
    """
    Observing frequency (Hz)
    """
    frequency::F
    """
    Observing bandwith (Hz)
    """
    bandwidth::F
    """
    A struct array of `ArrayBaselineDatum` holding time, u, v, baselines.
    """
    uvsamples::T
end

"""
    $(SIGNATURES)
Get the u, v positions of the array.
"""
function getuv(ac::ArrayConfiguration)
    return ac.uvsamples.u, ac.uvsamples.v
end

"""
    $(SIGNATURES)
Get the u, v, time, freq of the array as a tuple.
"""
function uvtimefreq(ac::EHTArrayConfiguration)
    u,v = getuv(ac)
    t = ac.uvsamples.time
    ν = ac.frequency
    return u, v, t, fill(ν, length(u))
end

"""
    $(TYPEDEF)
A single datum of an observing array.
"""
struct ArrayBaselineDatum{T}
    time::T
    u::T
    v::T
    baseline::Tuple{Symbol, Symbol}
end

const ArrayQuadrangleDatum = NTuple{4, ArrayBaselineDatum{T}} where {T}
const ArrayTriangleDatum = NTuple{3, ArrayBaselineDatum{T}} where {T}

"""
    uvpositions(datum)
Get the uvp positions of an inferometric datum.
"""
uvpositions(D::AbstractVisibilityDatum) = D.u, D.v

"""
    $(SIGNATURES)
The main data product type in `Comrade` this stores the `data` which can be a StructArray
of any `AbstractInterferometryDatum` type.

# ($FIELDS)
"""
Base.@kwdef struct EHTObservation{F,T<:AbstractInterferometryDatum{F},S<:StructArray{T}, N} <: Observation{F}
    """
    StructArray of data productts
    """
    data::S
    """
    modified julia date of the observation
    """
    mjd::N
    """
    RA of the observation in J2000 (deg)
    """
    ra::F
    """
    DEC of the observation in J2000 (deg)
    """
    dec::F
    """
    bandwidth of the observation (Hz)
    """
    bandwidth::F
    """
    frequency of the observation (Hz)
    """
    frequency::F
    """
    Common source name
    """
    source::Symbol
    """
    Time zone used.
    """
    timetype::Symbol = :UTC
end

Base.getindex(data::EHTObservation, i::Int) = data.data[i]
Base.length(data::EHTObservation) = length(data.data)

"""
    $(SIGNATURES)
Get all the stations in a observation. The result is a vector of symbols.
"""
function stations(d::EHTObservation{T,A}) where {T,A<:AbstractInterferometryDatum}
    bl = getdata(d, :baseline)
    s1 = first.(bl)
    s2 = last.(bl)
    return unique([s1..., s2...])
end

function Base.show(io::IO, d::EHTObservation{F,D}) where {F,D}
    println(io, "EHTObservation{$F,$D, ...}")
    println(io, "  source: ", d.source)
    println(io, "  mjd: ", d.mjd)
    println(io, "  frequency: ", d.frequency)
    println(io, "  bandwidth: ", d.bandwidth)
    println(io, "  stations: ", stations(d))
    println(io, "  nsamples: ", length(d))
end

"""
    getdata(obs::EHTObservation, s::Symbol)
Pass-through function that gets the array of `s` from the EHTObservation. For example
say you want the times of all measurement then

```julia
getdata(obs, :time)
```
"""
getdata(obs::Observation, s::Symbol) = getproperty(obs.data, s)
Base.getindex(obs::Observation, i) = getindex(obs.data, i)

"""
    $(SIGNATURES)
A struct holding the information for a single measured visibility.

# $(FIELDS)

"""
Base.@kwdef struct EHTVisibilityDatum{T<:Number} <: AbstractVisibilityDatum{T}
    """
    real component of the visibility (Jy)
    """
    visr::T
    """
    imaginary component of the visibility (Jy)
    """
    visi::T
    """
    error of the visibility (Jy)
    """
    error::T
    """
    x-direction baseline length in λ
    """
    u::T
    """
    y-direction baseline length in λ
    """
    v::T
    """
    Time of the observation in hours
    """
    time::T
    """
    station baseline codes
    """
    baseline::NTuple{2,Symbol}
end

"""
    $(SIGNATURES)
Return the complex visibility of the visibility datum
"""
@inline function visibility(D::EHTVisibilityDatum{T}) where {T}
        return Complex{T}(D.visr, D.visi)
end

"""
    $(SIGNATURES)
Get the amplitude of a visibility datum
"""
function amplitude(D::EHTVisibilityDatum)
    amp = hypot(D.visr,D.visi)
    return EHTVisibilityAmplitudeDatum(amp, D.error,
                                       D.u, D.v,
                                       D.time,
                                       D.frequency,
                                       D.bandwidth,
                                       D.baseline
                                    )
end

"""
    $(SIGNATURES)
A struct holding the information for a single measured visibility amplitude.

# $(FIELDS)

"""
Base.@kwdef struct EHTVisibilityAmplitudeDatum{T<:Number} <: AbstractVisibilityDatum{T}
    """
    amplitude (Jy)
    """
    amp::T
    """
    error of the visibility amplitude (Jy)
    """
    error::T
    """
    x-direction baseline length in λ
    """
    u::T
    """
    y-direction baseline length in λ
    """
    v::T
    """
    Time of the observation in hours
    """
    time::T
    """
    station baseline codes
    """
    baseline::NTuple{2,Symbol}
end

function checktriangle(D1::EHTVisibilityDatum,
                       D2::EHTVisibilityDatum,
                       D3::EHTVisibilityDatum)
    b1 = D1.baseline
    b2 = D2.baseline
    b3 = D3.baseline
    l = length(unique([b1..., b2..., b3...]))
    @assert l == 3 "For a valid closure phase you need 3 unique stations not $l"
    @assert (D1.time == D2.time == D3.time) "For a valid closure phase the times need to match"

end

@inline function amplitude(D::EHTVisibilityAmplitudeDatum{T}) where {T}
    return D.amp
end

@inline function bispectrum(D1::EHTVisibilityDatum, D2::EHTVisibilityDatum, D3::EHTVisibilityDatum)
    visibility(D1)*visibility(D2)*visibility(D3)
end

"""
    rescaleuv!(data)
rescale the u-v lengths according to scale. This can be useful when you want the spatial scales
to be in 1/μas instead of 1/rad
"""
function rescaleuv!(data::EHTObservation{T,D}, scale) where {T,D<:Union{EHTVisibilityAmplitudeDatum,EHTVisibilityDatum}}
    data.data.u .= data.data.u.*scale
    data.data.v .= data.data.v.*scale
    return data
end

"""
    $(TYPEDEF)

A Datum for a single closure phase.

# $(FIELDS)

"""
Base.@kwdef struct EHTClosurePhaseDatum{T<:Number} <: ClosureProducts{T}
    """
    closure phase (rad)
    """
    phase::T
    """
    error of the closure phase assuming the high-snr limit
    """
    error::T
    """
    u (λ) of first station
    """
    u1::T
    """
    v (λ) of first station
    """
    v1::T
    """
    u (λ) of second station
    """
    u2::T
    """
    v (λ) of second station
    """
    v2::T
    """
    u (λ) of third station
    """
    u3::T
    """
    v (λ) of third station
    """
    v3::T
    """
    Measured time of closure phase in hours
    """
    time::T
    """
    station baselines used
    """
    triangle::NTuple{3,Symbol}
end

function rescaleuv!(data::EHTObservation{T,D}, scale) where {T,D<:EHTClosurePhaseDatum}
    data.data.u1 .= data.data.u1.*scale
    data.data.v1 .= data.data.v1.*scale
    data.data.u2 .= data.data.u2.*scale
    data.data.v2 .= data.data.v2.*scale
    data.data.u3 .= data.data.u3.*scale
    data.data.v3 .= data.data.v3.*scale
    return data
end

function stations(d::EHTObservation{T,A}) where {T,A<:EHTClosurePhaseDatum}
    bl = getdata(d, :triangle)
    return unique(vcat(collect.(bl)...))
end

"""
    $(SIGNATURES)
Computes the closure phase of the three visibility datums.

# Notes
We currently use the high SNR Gaussian error approximation for the closure phase.
In the future we may use the moment matching from Monte Carlo sampling.
"""
function closure_phase(D1::EHTVisibilityDatum,
                      D2::EHTVisibilityDatum,
                      D3::EHTVisibilityDatum)

    checktriangle(D1,D2,D3)

    amp1 = amplitude(D1).amp
    amp2 = amplitude(D2).amp
    amp3 = amplitude(D3).amp
    u1,v1 = uvpositions(D1)
    u2,v2 = uvpositions(D2)
    u3,v3 = uvpositions(D3)

    bis = bispectrum(D1, D2, D3)
    s12 = unique([D1.baseline..., D2.baseline...])
    s123 = unique([s12..., D3.baseline...])
    #Use the Gaussian approximation TODO hook this into Measurements.jl?
    err = sqrt((D1.error/amp1)^2 + (D2.error/amp2)^2 + (D3.error/amp3)^2)
    return EHTClosurePhaseDatum(angle(bis), err,
                                u1, v1, u2, v2, u3, v3,
                                time, s123)
end

function baselines(CP::EHTClosurePhaseDatum)
    tri = CP.triangle
    return ((tri[1],tri[2]), (tri[2], tri[3]), (tri[3], tri[1]))
end

uvpositions(datum::EHTClosurePhaseDatum) = (datum.u1, datum.v1, datum.u2, datum.v2, datum.u3, datum.v3)

"""
    $(TYPEDEF)

A Datum for a single log closure amplitude.

# $(FIELDS)

"""
Base.@kwdef struct EHTLogClosureAmplitudeDatum{T<:Number} <: ClosureProducts{T}
    """
    log-closure amplitude
    """
    amp::T
    """
    log-closure amplitude error in the high-snr limit
    """
    error::T
    """
    u (λ) of first station
    """
    u1::T
    """
    v (λ) of first station
    """
    v1::T
    """
    u (λ) of second station
    """
    u2::T
    """
    v (λ) of second station
    """
    v2::T
    """
    u (λ) of third station
    """
    u3::T
    """
    v (λ) of third station
    """
    v3::T
    """
    u (λ) of fourth station
    """
    u4::T
    """
    v (λ) of fourth station
    """
    v4::T
    """
    Observation time of the quadrangle
    """
    time::T
    """
    station codes for the quadrangle
    """
    quadrangle::NTuple{4,Symbol}
end

function rescaleuv!(data::EHTObservation{T,D}, scale) where {T,D<:EHTLogClosureAmplitudeDatum}
    data.data.u1 .= data.data.u1.*scale
    data.data.v1 .= data.data.v1.*scale
    data.data.u2 .= data.data.u2.*scale
    data.data.v2 .= data.data.v2.*scale
    data.data.u3 .= data.data.u3.*scale
    data.data.v3 .= data.data.v3.*scale
    data.data.u4 .= data.data.u4.*scale
    data.data.v4 .= data.data.v4.*scale
    return data
end

function baselines(CP::EHTLogClosureAmplitudeDatum)
    quad = CP.quadrangle
    return ((quad[1],quad[2]), (quad[3], quad[4]), (quad[1], quad[3]), (quad[2], quad[4]))
end

function stations(d::EHTObservation{T,A}) where {T,A<:EHTLogClosureAmplitudeDatum}
    bl = getdata(d, :quadrangle)
    return unique(vcat(collect.(bl)...))
end

uvpositions(datum::EHTLogClosureAmplitudeDatum) = (datum.u1, datum.v1, datum.u2, datum.v2, datum.u3, datum.v3, datum.u4, datum.v4)

# @doc raw"""
#     $(SIGNATURES)
# Compute the log-closure amplitude of 4 visibility datums.
# This uses the same ordering as eht-imaging i.e.

# ```math
# C_A = \log \frac{|V_{1}| |V_{2}|}{|V_{3}| |V_{4}|}
# ```

# # Notes

# Currently we use the high SNR Gaussian approximation to the thermal error.
# In the future we may use the moment matching approach using Monte Carlo sampling.
# """
# function logclosure_amplitude(D1::EHTVisibilityDatum,
#                               D2::EHTVisibilityDatum,
#                               D3::EHTVisibilityDatum,
#                               D4::EHTVisibilityDatum)
#     check_quadrangle(D1,D2,D3,D4)

#     amp1 = amplitude(D1)
#     amp2 = amplitude(D2)
#     amp3 = amplitude(D3)
#     amp4 = amplitude(D4)

#     u1,v1 = uvpositions(D1)
#     u2,v2 = uvpositions(D2)
#     u3,v3 = uvpositions(D3)
#     u4,v4 = uvpositions(D4)

#     #Construct the quadrangle
#     s12 = unique([D1.baseline..., D2.baseline...])
#     s123 = unique([s12..., D3.baseline...])
#     s1234 = unique([s123..., D4.baseline...])

#     # Now do the error propogation
#     err = hypot(D1.err/map1, D2.err/amp2, D3.err/amp3, D4.err/amp4)

#     #construct the lcamp
#     lcamp = log(amp1*amp2/(amp3*amp4))

#     return EHTLogClosureAmplitudeDatum(lcamp, err,
#                                        u1, v1,
#                                        u2, v2,
#                                        u3, v3,
#                                        u4, v4,
#                                        time, s1234)
# end

"""
    $(SIGNATURES)
Extract the array configuration from a visibility EHT observation.
"""
function arrayconfig(vis::EHTObservation{F,A}) where {F,A<:Union{EHTVisibilityDatum, EHTVisibilityAmplitudeDatum}}
    u = getdata(vis, :u)
    v = getdata(vis, :v)
    st = stations(vis)
    times = getdata(vis, :time)
    bandwidth = vis.bandwidth
    frequency = vis.frequency
    error = getdata(vis, :error)
    baseline = getdata(vis, :baseline)
    uvsamples = StructArray{ArrayBaselineDatum}(time=times,
                                        u=u,
                                        v=v,
                                        baseline=baseline,
                                        error_real=error,
                                        error_imag=error
                                    )
    return EHTArrayConfiguration(st, frequency, bandwidth, uvsamples)
end

include(joinpath(@__DIR__, "io.jl"))
include(joinpath(@__DIR__, "ehtim.jl"))


function born_approximation(output::Ts, input::Ts, iflag::Ti;
         irz::Vector, irx::Vector, location_flag="index", dir_sourceside="NULL", normalization_flag=true, mute_index=0, fidiff::TdParams) where {Ts<:String, Ti<:Int64}

    # clean the space
    rm(output, force=true, recursive=true)

    if iflag == 1
       born_approximation_forward(output, input, irz, irx, dir_sourceside, fidiff;
                                  location_flag=location_flag, normalization_flag=normalization_flag, mute_index=mute_index)

    elseif iflag == 2
       born_approximation_adjoint(output, input, dir_sourceside, fidiff;
                                  normalization_flag=normalization_flag, mute_index=mute_index)

    else
       error("undefined operator behavior")
    end

    return nothing
end

"""
   perform y = ax + by for multiple recordings
"""
function recordings_axpby!(a, dir_x::Ts, b, dir_y::Ts) where {Ts<:String}

    # determine number of files
    file_name = readdir(dir_x)
    nx        = length(file_name)
    if file_name[1] == ".DS_Store" # in case of MacOS
       nx = nx - 1
    end

    file_name = readdir(dir_y)
    ny        = length(file_name)
    if file_name[1] == ".DS_Store" # in case of MacOS
       ny = ny - 1
    end

    nx == ny || error("number of files doesn't match")

    r = 0.0
    for i = 1 : nx
        file_name = join(["recordings_" "$i" ".bin"])

        path_x    = joinpath(dir_x, file_name)
        path_y    = joinpath(dir_y, file_name)

        rec_x     = read_recordings(path_x)
        rec_y     = read_recordings(path_y)

        for i2 = 1 : rec_x.nr
            for i1 = 1 : rec_x.nt
                rec_y.p[i1,i2] = a*rec_x.p[i1,i2] + b*rec_y.p[i1,i2]
            end
        end

        write_recordings(path_y, rec_y)
    end

    return nothing
end

"""
   perform y = ax + by for variables in model space
"""
function image_axpby!(a, path_x::Ts, b, path_y::Ts) where {Ts<:String}

    (hdr1, x) = read_RSdata(path_x)
    (hdr2, y) = read_RSdata(path_y)

    for i = 1 : length(x)
        y[i] = a*x[i] + b*y[i]
    end

    write_RSdata(path_y, hdr2, y)

    return nothing
end

function recordings_norm(dir::Ts) where {Tv<:Number, Ts<:String}

    # determine number of files
    file_name = readdir(dir)
    n         = length(file_name)
    if file_name[1] == ".DS_Store" # in case of MacOS
       n = n - 1
    end

    r = 0.0
    for i = 1 : n
        file_name = join(["recordings_" "$i" ".bin"])
        path      = joinpath(dir, file_name)
        rec       = read_recordings(path)
        tmp       = norm(rec.p)
        r        += tmp * tmp
    end

    return sqrt(r)
end


mutable struct WishartStatXMLElement <: MyXMLElement
    el::XMLOrNothing
    traitLikelihood_el::TraitLikelihoodXMLElement
    rm_el::Union{Nothing, RepeatedMeasuresXMLElement}
    mbd_el::Union{Nothing, MBDXMLElement}

    function WishartStatXMLElement(like_el::TraitLikelihoodXMLElement,
                mbd_el::MBDXMLElement)
        return new(nothing, like_el, nothing, mbd_el)
    end

    function WishartStatXMLElement(like_el::TraitLikelihoodXMLElement,
                rm_el::RepeatedMeasuresXMLElement)
        return new(nothing, like_el, rm_el, nothing)
    end
end

mutable struct PrecisionGibbsOperatorXMLElement <: OperatorXMLElement
    el::XMLOrNothing
    wishart_stats::WishartStatXMLElement
    weight::Float64

    function PrecisionGibbsOperatorXMLElement(
            traitLikelihood_el::TraitLikelihoodXMLElement,
            mbd_el::MBDXMLElement)

        wishart_stats = WishartStatXMLElement(traitLikelihood_el, mbd_el)
        return new(nothing, wishart_stats, 1.0)
    end

    function PrecisionGibbsOperatorXMLElement(
        traitLikelihood_el::TraitLikelihoodXMLElement,
        rm_el::RepeatedMeasuresXMLElement)

        wishart_stats = WishartStatXMLElement(traitLikelihood_el, rm_el)
        return new(nothing, wishart_stats, 1.0)
    end

end

mutable struct CompoundPrecisionOperatorXMLElement <: OperatorXMLElement
    el::XMLOrNothing
    diff_el::PrecisionGibbsOperatorXMLElement
    res_el::PrecisionGibbsOperatorXMLElement
    weight::Float64

    function CompoundPrecisionOperatorXMLElement(
            diff_el::PrecisionGibbsOperatorXMLElement,
            res_el::PrecisionGibbsOperatorXMLElement)

        diff_ws = diff_el.wishart_stats
        res_ws = res_el.wishart_stats
        @assert !isnothing(diff_ws.mbd_el) && isnothing(diff_ws.rm_el)
        @assert isnothing(res_ws.mbd_el) && !isnothing(res_ws.rm_el)

        return new(nothing, diff_el, res_el, 1.0)
    end

    function CompoundPrecisionOperatorXMLElement(
                tl_el::TraitLikelihoodXMLElement,
                mbd_el::MBDXMLElement,
                rm_el::RepeatedMeasuresXMLElement)

        diff_el = PrecisionGibbsOperatorXMLElement(tl_el, mbd_el)
        res_el = PrecisionGibbsOperatorXMLElement(tl_el, rm_el)
        return CompoundPrecisionOperatorXMLElement(diff_el, res_el)
    end

end

function make_xml(cpo_el::CompoundPrecisionOperatorXMLElement)
    el = new_element(bn.COMPOUND_PRECISION_OPERATOR)
    set_attribute(el, bn.WEIGHT, cpo_el.weight)

    do_el = new_child(el, bn.DIFFUSION_OPERATOR)
    make_xml(cpo_el.diff_el)
    add_child(do_el, cpo_el.diff_el.el)

    ro_el = new_child(el, bn.RESIDUAL_OPERATOR)
    make_xml(cpo_el.res_el)
    add_child(ro_el, cpo_el.res_el.el)

    cpo_el.el = el
    return el
end

function make_xml(ws_el::WishartStatXMLElement)
    make_xml(ws_el.traitLikelihood_el)

    if isnothing(ws_el.rm_el) && !isnothing(ws_el.mbd_el)
        el = new_element(bn.WISHART_STATISTICS)
        return finish_Wishart_el(ws_el, el)
    elseif !isnothing(ws_el.rm_el) && isnothing(ws_el.mbd_el)
        make_xml(ws_el.rm_el)
        el = new_element(bn.REPEATED_MEASURES_WISHART_STATISTICS)
        add_ref_el(el, ws_el.rm_el.el)
        return finish_Wishart_el(ws_el, el)
    else
        error("Not implemented")
    end
end

function finish_Wishart_el(ws_el::WishartStatXMLElement, el::XMLElement)
    add_ref_el(el, ws_el.traitLikelihood_el.el)
    set_attribute(el, bn.TRAIT_NAME, ws_el.traitLikelihood_el.attrs[bn.TRAIT_NAME])
    ws_el.el = el
    return el
end

function make_xml(pgo_el::PrecisionGibbsOperatorXMLElement)
    el = new_element(bn.PRECISION_GIBBS_OPERATOR)
    set_attribute(el, bn.WEIGHT, pgo_el.weight)
    ws_el = pgo_el.wishart_stats
    make_xml(ws_el)
    add_child(el, ws_el.el)
    if isnothing(ws_el.rm_el) && !isnothing(ws_el.mbd_el)
        make_xml(ws_el.mbd_el)
        add_ref_el(el, ws_el.mbd_el.precision_prior.el)
    elseif !isnothing(ws_el.rm_el) && isnothing(ws_el.mbd_el)
        make_xml(ws_el.rm_el)
        add_ref_el(el, ws_el.rm_el.precision_prior.el)
    else
        error("Not implemented")
    end
    pgo_el.el = el
    return el
end


###
### COIN-OR Clp API Wrapper
###

module Clp

include("ClpCInterface.jl")
include("ClpSolverInterface.jl")
include("MOIWrapper.jl")

using Clp.ClpMathProgSolverInterface
export ClpSolver

end # module


using Fermi.DIIS

using TensorOperations
using LinearAlgebra
using Formatting
import Base: show

export UHF

abstract type UHFAlgorithm end

struct UHF <: AbstractHFWavefunction
    molecule::Molecule
    energy::Float64
    nocc::Int
    nvir::Int
    orbitals::UHFOrbitals
    e_conv::Float64
    d_conv::Float64
end

function UHF(x...)
    if !any(i-> i isa UHFAlgorithm, x)
        UHF(x..., get_uhf_alg())
    else
        # Print the type of arguments given for a better feedback
        args = "("
        for a in x[1:end-1]
            args *= "$(typeof(a)), "
        end
        args = args[1:end-2]*")"
        throw(FermiException("invalid arguments for UHF method: $args"))
    end
end

"""
    Fermi.HartreeFock.get_uhf_alg()

Returns a singleton type corresponding to a UHF implementation.
"""
function get_uhf_alg(N::Int = Options.get("uhf_alg"))
    try
        return get_uhf_alg(Val(N))
    catch MethodError
        throw(FermiException("implementation number $N not available for UHF."))
    end
end

struct UHFa <: UHFAlgorithm end
include("UHFa.jl")
include("UHFHelper.jl")
# And a number is assigned to the implementation
get_uhf_alg(x::Val{1}) = UHFa()

function string_repr(X::UHF)
    out = ""
    out = out*" ⇒ Fermi Unrestricted Hartree--Fock Wave function\n"
    out = out*" ⋅ Basis:                  $(X.orbitals.basis)\n"
    out = out*" ⋅ Energy:                 $(X.energy)\n"
    out = out*" ⋅ Occ. Spatial Orbitals:  $(X.nocc)\n"
    out = out*" ⋅ Vir. Spatial Orbitals:  $(X.nvir)\n"
    out = out*"Convergence: "
    out = out*"ΔE => $(format("{:1.2e}",abs(X.e_conv)))"
    out = out*" Dᵣₘₛ => $(format("{:1.2e}",abs(X.d_conv)))"
    return out
end

function show(io::IO, ::MIME"text/plain", X::UHF)
    print(string_repr(X))
end


WIDTH = 800
HEIGHT = 600
RATIO = WIDTH/(HEIGHT*1f0)
SIZE = WIDTH * HEIGHT
FOV = 60.0f0
CLIP_NEAR = 0.001f0
CLIP_FAR = 10000.0f0

"""
sets glfw window size + viewport
"""
function rezizeWindow(width,height)
  global WIDTH, HEIGHT, RATIO, SIZE
  WIDTH = width
  HEIGHT = height
  RATIO = WIDTH/(HEIGHT*1f0)
  SIZE = WIDTH * HEIGHT
  GLFW.SetWindowSize(window, WIDTH, HEIGHT)
  glViewport(0, 0, WIDTH, HEIGHT)
end

lastCursorPos = [0,0]
cursorPos_old = [0,0]
cursorPos = [0,0]
mouseMove = [0,0]

#MVP = eye(Mat4x4f)
#MVP *= transform([0.0f0,0.0f0,0.0f0],[0.0f0,0.0f0,0.0f0],Float32[])

keyFB = 0
keyLR = 0
keyUD = 0
speed = false

mouseKeyPressed = false
keyPressed = false
keyPressing = false
keyValue = 0

"""
camera object with holds position, rotation, scaling and various matrices like MVP
"""
type Camera
  moved::Bool

  position::AbstractArray
  rotation::AbstractArray
  scale::AbstractArray

  viewMat::AbstractArray
  projectionMat::AbstractArray
  modelMat::AbstractArray

  translateMat::AbstractArray
  rotationMat::AbstractArray
  scalingMat::AbstractArray

  MVP::AbstractArray

  Camera() = new(true, zerosVector3f,zerosVector3f,onesVector3f,
  eyeMat4x4f,eyeMat4x4f,eyeMat4x4f,eyeMat4x4f,eyeMat4x4f,eyeMat4x4f,eyeMat4x4f)
end

CAMERA = Camera()

"""
gets forward vector of camera direction
"""
forward(camera::Camera) = forward(camera.rotationMat)

"""
gets right vector of camera direction
"""
right(camera::Camera) = right(camera.rotationMat)

"""
gets up vector of camera direction
"""
up(camera::Camera) = up(camera.rotationMat)

"""
sets projection matrix
"""
setProjection(camera::Camera, m::AbstractArray) = (camera.projectionMat = m)

"""
sets view matrix
"""
setView(camera::Camera, m::AbstractArray) = (camera.viewMat = m)

VIEW_KEYS=false

"""
event which catches keyboard inputs.
here keys for wireframe, fullscreen and camera movement are defined
"""
function OnKey(window, key::Number, scancode::Number, action::Number, mods::Number)
  if key == 70 && action == 1 # f
    # ...
  elseif key == 71 && action == 1 # g
    #global wireframe
    #wireframe = !wireframe
    #println("Wireframe: ",wireframe)

  elseif key == 72 && action == 1 # h
    #global fullscreen=!fullscreen

  elseif key == 87 #w
    global keyFB = (action > 0)?1:0
  elseif key == 83 #s
    global keyFB = (action > 0)?-1:0
  elseif key == 65 #a
    global keyLR = (action > 0)?-1:0
  elseif key == 68 #d
    global keyLR = (action > 0)?1:0
  elseif key == 32 #space
    global keyUD = (action > 0)?1:0
  elseif key == 67 || key == 341 # c || lctrl
    global keyUD = (action > 0)?-1:0
  elseif key == 340 #lshift
    global speed = (action > 0)
  end

  if key == 75 && action == 1 #k
    global VIEW_KEYS = !VIEW_KEYS
    println("VIEW_KEYS = $VIEW_KEYS")
  end

  if VIEW_KEYS println("K: ", key, " A: ", action, " C: ", scancode, " M: ", mods) end

  if action == 1 global keyPressed = true end
  global keyPressing = action != 0
  global keyValue = key
  nothing
end

"""
event which catches mouse key inpits and hides/shows cursor when mouse button is pressed
"""
function OnMouseKey(window, key::Number, action::Number, mods::Number)
  global cursorPos_old, cursorPos
  global mouseKeyPressed = action != 0
  if action == 1
    cursorPos_old = cursorPos
    GLFW.SetInputMode(window, GLFW.CURSOR, GLFW.CURSOR_DISABLED)
  end
  if action == 0
    GLFW.SetInputMode(window, GLFW.CURSOR, GLFW.CURSOR_NORMAL)
  end
  nothing
end

"""
event which catches mouse position for camera rotation event
"""
function OnCursorPos(window, x::Number, y::Number)
  global CAMERA, lastCursorPos, cursorPos, mouseMove, mouseKeyPressed

  t = [x/800f0,y/600f0]
  cursorPos = t
  l = lastCursorPos; if l == [] l = t end
  n = l - t
  lastCursorPos = t
  mouseMove = n

  if mouseKeyPressed OnRotate(CAMERA) end
  nothing
end

"""
rotates camera
"""
function rotate(camera::Camera, rotation::AbstractArray)
  camera.rotation += rotation
  camera.moved=true
  #println("rotated")
end

"""
moves camera, adds vector to current position
"""
function move(camera::Camera, position::AbstractArray)
  camera.position += position
  camera.moved=true
  #println("moved")
end

"""
event which calculates cursor position shifts and calls rotate function
"""
function OnRotate(camera::Camera)
  global cursorPos, cursorPos_old
  mx = cursorPos[1] - cursorPos_old[1]
  my = cursorPos[2] - cursorPos_old[2]
  #println(mx)
  cursorPos_old = cursorPos
  rotate(camera, [-mx*2f0,my*2f0,0f0]) #[-mx*2,my*2,0f0] #Vec3f((-mx+0.5f0),(-0.5f0+my),0f0)
end

oldposition = CAMERA.position
shiftposition = [0.0f0,0,0]

"""
sets camera position
"""
function setPosition(camera::Camera, position::AbstractArray)
  camera.position = position
  global oldposition = position
  camera.translateMat = translation(camera.position)
  camera.moved = true
end

"""
event which updates positions shifts (left,right,up,down,forward,back)
key is (left,right,up,down,forward,back)
m is direction value with weight (positive, negative)
"""
function OnMove(camera::Camera, key::Symbol, m::Number)
  global shiftposition
  if key == :FORWARD  move(camera, forward(camera)*(m*0.05f0*(!speed?1f0:10f0)))
  elseif key == :RIGHT  move(camera, right(camera)*(-m*0.05f0*(!speed?1f0:10f0))) #+Vec3f(-right*0.02f0,-up*0.02f0,forward*0.02f0)
  elseif key == :UP  move(camera, up(camera)*(-m*0.05f0*(!speed?1f0:10f0)))
  end

  #dif = camera.position
  #step = trunc.(dif)
  #shiftposition = dif - step
  #step = abs.(step)

  #if step[1] >= 2 camera.position[1] = 0 end
  #if step[2] >= 2 camera.position[2] = 0 end
  #if step[3] >= 2 camera.position[3] = 0 end

  #  dist = sqrt(camera.position[1]^2+camera.position[2]^2+camera.position[3]^2)
  # if dist >= 1 camera.position = [0f0,0,0] end
end

"""
update function where camera translation is update only when camera was moved by input.
here cameras MVP Matrix is updated aswell
"""
function Update(camera::Camera)
  if OnTime(0.01)
    if keyFB != 0 OnMove(camera, :FORWARD, keyFB) end
    if keyLR != 0 OnMove(camera, :RIGHT, keyLR) end
    if keyUD != 0 OnMove(camera, :UP, keyUD) end
  end

  if camera.moved
    camera.translateMat = translation(camera.position)
    camera.rotationMat = computeRotation(camera.rotation)
    camera.viewMat = camera.scalingMat * camera.rotationMat * camera.translateMat
    camera.MVP = camera.modelMat * camera.projectionMat * camera.viewMat
  end
end

"""
event which is called by game loop and calls real update function
this event resets camera moved state
"""
function OnUpdate(camera::Camera)
  #if isFocus return end
  Update(camera)
  r = camera.moved
  camera.moved = false
  r
end


export iauUt1utc
"""
Time scale transformation:  Universal Time, UT1, to Coordinated
Universal Time, UTC.

This function is part of the International Astronomical Union's
SOFA (Standards of Fundamental Astronomy) software collection.

Status:  canonical.

Given:
   ut11,ut12  double   UT1 as a 2-part Julian Date (Note 1)
   dut1       double   Delta UT1: UT1-UTC in seconds (Note 2)

Returned:
   utc1,utc2  double   UTC as a 2-part quasi Julian Date (Notes 3,4)

Returned (function value):
               int      status: +1 = dubious year (Note 5)
                                 0 = OK
                              -1 = unacceptable date

Notes:

   1. ut11+ut12 is Julian Date, apportioned in any convenient way
   between the two arguments, for example where ut11 is the Julian
   Day Number and ut12 is the fraction of a day.  The returned utc1
   and utc2 form an analogous pair, except that a special convention
   is used, to deal with the problem of leap seconds - see Note 3.

   2. Delta UT1 can be obtained from tabulations provided by the
   International Earth Rotation and Reference Systems Service.  The
   value changes abruptly by 1s at a leap second;  however, close to
   a leap second the algorithm used here is tolerant of the "wrong"
   choice of value being made.

   3. JD cannot unambiguously represent UTC during a leap second unless
   special measures are taken.  The convention in the present
   function is that the returned quasi JD day UTC1+UTC2 represents
   UTC days whether the length is 86399, 86400 or 86401 SI seconds.

   4. The function iauD2dtf can be used to transform the UTC quasi-JD
   into calendar date and clock time, including UTC leap second
   handling.

   5. The warning status "dubious year" flags UTCs that predate the
   introduction of the time scale or that are too far in the future
   to be trusted.  See iauDat for further details.

Called:
   iauJd2cal    JD to Gregorian calendar
   iauDat       delta(AT) = TAI-UTC
   iauCal2jd    Gregorian calendar to JD

References:

   McCarthy, D. D., Petit, G. (eds.), IERS Conventions (2003),
   IERS Technical Note No. 32, BKG (2004)

   Explanatory Supplement to the Astronomical Almanac,
   P. Kenneth Seidelmann (ed), University Science Books (1992)

This revision:  2013 June 18

SOFA release 2018-01-30

Copyright (C) 2018 IAU SOFA Board.  See notes at end.
"""

# int iauUt1utc(double ut11, double ut12, double dut1,
#               double *utc1, double *utc2)

function iauUt1utc(ut11::Real, ut12::Real, dta::Real)
   ref_tt1 = Ref{Float64}(0.0)
   ref_tt2 = Ref{Float64}(0.0)

   status = ccall((:iauUt1utc, libsofa_c), Cint,
                   (Cdouble, Cdouble, Cdouble, Ref{Cdouble}, Ref{Cdouble}),
                   convert(Float64, ut11), convert(Float64, ut12), convert(Float64, dta),
                   ref_tt1, ref_tt2)

   return status, ref_tt1[], ref_tt2[]
end


push!(LOAD_PATH, "/Users/czhang/Desktop/Projects/dw_/julialib/")
import DimmWitted
DimmWitted.set_libpath("/Users/czhang/Desktop/Projects/dw_/libdw_julia")

######################################
# The following piece of code creates a
# synthetic data set:
#    - Data type is Cdouble
#    - Modle type is Array{Cdouble}
#
nexp = 100000
nfeat = 1024
examples = Array(Cdouble, nexp, nfeat+1)
for row = 1:nexp
	for col = 1:nfeat
		examples[row, col] = 1
	end
	if rand() > 0.8
		examples[row, nfeat+1] = 0
	else
		examples[row, nfeat+1] = 1
	end
end
model = Cdouble[0 for i = 1:nfeat]

######################################
# Define the loss function and gradient
# function for logistic regression
#
function loss(row::Array{Cdouble,1}, model::Array{Cdouble,1})
	const label = row[length(row)]
	const nfeat = length(model)
	d = 0.0
	for i = 1:nfeat
		d = d + row[i]*model[i]
	end
	return (-label * d + log(exp(d) + 1.0))
end

function grad(row::Array{Cdouble,1}, model::Array{Cdouble,1})
	const label = row[length(row)]
	const nfeat = length(model)
	d = 0.0
	for i = 1:nfeat
		d = d + row[i]*model[i]
	end
	d = exp(-d)
		Z = 0.00001 * (-label + 1.0/(1.0+d))
  	for i = 1:nfeat
  		model[i] = model[i] - row[i] * Z
  	end
	return 1.0
end

######################################
# Create a DimmWitted object using data
# and model. You do not need to specify
# the type, they are infer'ed by the
# open() function, which is parametric.
#
dw = DimmWitted.open(examples, model,
                DimmWitted.MR_PERMACHINE,
                DimmWitted.DR_SHARDING,
                DimmWitted.AC_ROW)

######################################
# Register functions.
#
handle_loss = DimmWitted.register_row(dw, loss)
handle_grad = DimmWitted.register_row(dw, grad)

######################################
# Run 10 epoches.
#
for iepoch = 1:10
	rs = DimmWitted.exec(dw, handle_loss)
	println("LOSS: ", rs/nexp)
	rs = DimmWitted.exec(dw, handle_grad)
end


##############################
## determine function names ##
##############################

const VineCppPath = joinpath(Pkg.dir("Copulas"),
                             "src/VineCopulaCPP/libVineCopulaCPP.so.1.0")

pdfName = :_Z13PairCopulaPDFiPKdPdS1_S1_j
cdfName = :_Z13PairCopulaCDFiPKdPdS1_S1_j
hfunName = :_Z14PairCopulaHfuniPKdPdS1_S1_j
vfunName = :_Z14PairCopulaVfuniPKdPdS1_S1_j
hvinName = :_Z17PairCopulaInvHfuniPKdPdS1_S1_j
vinvName = :_Z17PairCopulaInvVfuniPKdPdS1_S1_j

#######################
## VineCpp interface ##
#######################

## copula ID lookup tables
##########################

## list of copula families - order determines copula IDs!
const vineCpp_cops = [:Indep,
                      :AMH,
                      :AsymFGM,
                      :BB1,
                      :BB6,
                      :BB7,
                      :BB8,
                      :Clayton,
                      :FGM,
                      :Frank,
                      :Gaussian,
                      :Gumbel,
                      :IteratedFGM,
                      :Joe,
                      :PartialFrank,
                      :Plackett,
                      :Tawn1,
                      :Tawn2,
                      :Tawn,
                      :t]

## transform to dictionary with families as keys
const vineCpp_IDs = [vineCpp_cops[ii] => ii-1
                     for ii=1:length(vineCpp_cops)]

function getVineCppId(nam::Symbol)
    return vineCpp_IDs[nam]
end

#######################################
## create parametric VineCpp copulas ##
#######################################

copTypeNames_Cpp =
    [symbol(string(vineCpp_cops[ii], "PC_Cpp"))
     for ii=1:length(vineCpp_cops)]

## define copula types through macro
##----------------------------------

macro defineParamCopCpp(nam)
    esc(quote
        type $(nam) <: Copulas.ParamPC_Cpp
            params::FloatVec
        end
    end)
end

for cop in copTypeNames_Cpp
    eval(macroexpand(:(@defineParamCopCpp $cop)))
end

###########################
## VineCpp ID interfaces ##
###########################

## get VineCpp Id for copula object
##---------------------------------

## dict with copula type names as keys
const type2idDictCpp = [copTypeNames_Cpp[ii] => (ii-1)::Int for ii=1:length(copTypeNames_Cpp)]

function getVineCppId(cop::ParamPC_Cpp)
    return type2idDictCpp[symbol(string(typeof(cop)))]
end

## get family name for ID
##-----------------------

function getCopNam(ii::Int)
    return vineCpp_cops[ii+1]
end

## get copula type for ID
##-----------------------

function getCopType_Cpp(ii::Int)
    return copTypeNames_Cpp[ii+1]
end

#############
## display ##
#############

function displayShortString(cop::ParamPC_Cpp)
    copNam = string(typeof(cop))
    par = params(cop)
    nParams = length(par)
    paramsString = string(par[1])
    for ii=2:nParams
        paramsString = join([paramsString, ", ", string(par[ii])])
    end
    return join([copNam, ": ", paramsString])
end

function displayShortString(cop::IndepPC_Cpp)
    return "IndepPC_Cpp"
end

import Base.Multimedia.display
function display(cop::ParamPC_Cpp)
    msg = displayShortString(cop)
    println("$msg")
end

##############
## equality ##
##############

import Base.==
function ==(cop1::ParamPC_Cpp, cop2::ParamPC_Cpp)
    isequal(typeof(cop1), typeof(cop2)) || return false
    return isequal(params(cop1), params(cop2))
end

#######################
## VineCpp functions ##
#######################

## pdf function
##-------------

function pdf(cop::ParamPC_Cpp, u1::FloatVec, u2::FloatVec)
    fam, params = getFamAndParams(cop)

    nObs = length(u1)
    retVals = zeros(Float64, nObs)

    ccall((:_Z13PairCopulaPDFiPKdPdS1_S1_j, VineCppPath),
          Void,
          (Int, Ptr{Float64}, Ptr{Float64}, Ptr{Float64},
           Ptr{Float64}, Int),
          fam, params, u1, u2, retVals, nObs)

    return retVals
end

function cdf(cop::ParamPC_Cpp, u1::FloatVec, u2::FloatVec)
    fam, params = getFamAndParams(cop)

    nObs = length(u1)
    retVals = zeros(Float64, nObs)

    ccall((:_Z13PairCopulaCDFiPKdPdS1_S1_j, VineCppPath),
          Void,
          (Int, Ptr{Float64}, Ptr{Float64}, Ptr{Float64},
           Ptr{Float64}, Int),
          fam, params, u1, u2, retVals, nObs)

    return retVals
end

function hfun(cop::ParamPC_Cpp, u1::FloatVec, u2::FloatVec)
    fam, params = getFamAndParams(cop)

    nObs = length(u1)
    retVals = zeros(Float64, nObs)

    ccall((:_Z14PairCopulaHfuniPKdPdS1_S1_j, VineCppPath),
          Void,
          (Int, Ptr{Float64}, Ptr{Float64}, Ptr{Float64},
           Ptr{Float64}, Int),
          fam, params, u1, u2, retVals, nObs)

    return retVals
end

function vfun(cop::ParamPC_Cpp, u2::FloatVec, u1::FloatVec)
    # u2|u1
    fam, params = getFamAndParams(cop)

    nObs = length(u1)
    retVals = zeros(Float64, nObs)

    ccall((:_Z14PairCopulaVfuniPKdPdS1_S1_j, VineCppPath),
          Void,
          (Int, Ptr{Float64}, Ptr{Float64}, Ptr{Float64},
           Ptr{Float64}, Int),
          fam, params, u1, u2, retVals, nObs)

    return retVals
end

function hinv(cop::ParamPC_Cpp, u1::FloatVec, u2::FloatVec)
    fam, params = getFamAndParams(cop)

    nObs = length(u1)
    retVals = zeros(Float64, nObs)

    ccall((:_Z17PairCopulaInvHfuniPKdPdS1_S1_j, VineCppPath),
          Void,
          (Int, Ptr{Float64}, Ptr{Float64}, Ptr{Float64},
           Ptr{Float64}, Int),
          fam, params, u1, u2, retVals, nObs)

    return retVals
end

function vinv(cop::ParamPC_Cpp, u2::FloatVec, u1::FloatVec)
    # u2|u1
    fam, params = getFamAndParams(cop)

    nObs = length(u1)
    retVals = zeros(Float64, nObs)

    ccall((:_Z17PairCopulaInvVfuniPKdPdS1_S1_j, VineCppPath),
          Void,
          (Int, Ptr{Float64}, Ptr{Float64}, Ptr{Float64},
           Ptr{Float64}, Int),
          fam, params, u1, u2, retVals, nObs)

    return retVals
end


macro forward_func(ex, fs)
    T, field = ex.args[1], ex.args[2].value

    T = esc(T)
    fs = isexpr(fs, :tuple) ? map(esc, fs.args) : [esc(fs)]
    :($([:($f(x::$T, args...) = (Base.@_inline_meta; $f(x.$field, args...)))
        for f in fs]...);
    nothing)
end

@forward_func SEAssemblage.occ nthings, nplaces, occupancy, richness, nrecords, occurring, occupied, thingnames
@forward_func SEThings.commatrix nthings, nplaces, thingnames, placenames, occupancy, richness, nrecords, occurring, occupied
@forward_func SEAssemblage.site placenames

#--------------------------------------------------------------------------
# Basic summary functions

occupancy(com::AbstractComMatrix) = occupancy(occurrences(com))
richness(com::AbstractComMatrix) = richness(occurrences(com))
occurring(com::AbstractComMatrix, idx...) = occurring(occurrences(com), idx...)
occupied(com::AbstractComMatrix, idx...) = occupied(occurrences(com), idx...)

const nspecies = nthings
nthings(com::AbstractComMatrix) = size(com.occurrences, 1)

const nsites = nplaces
nplaces(com::AbstractComMatrix) = size(com.occurrences, 2)
nplaces(sd::SiteData) = size(coordinates(sd.site), 1)
nplaces(sd::SELocations) = DataFrames.ncol(sd.sitestats)
nplaces(gr::GridData) = size(gr.indices, 1)
nplaces(pd::PointData) = size(pd.coords, 1)

nrecords(com::AbstractComMatrix) = _nnz(occurrences(com))

const getspecies = thingoccurrences
thingoccurrences(com::AbstractComMatrix, idx) = thingoccurrences(occurrences(com), idx)

const getsite = placeoccurrences
placeoccurrences(com::AbstractComMatrix, idx) = placeoccurrences(occurrences(com), idx)

const speciesnames = thingnames
thingnames(com::AbstractComMatrix) = com.speciesnames

const sitenames = placenames
placenames(com::AbstractComMatrix) = com.sitenames
placenames(sd::SiteData) = sitenames(sd.site)
placenames(sd::SELocations) = collect(sd.sitestats[:sites])

sitetotals(com::AbstractComMatrix) = vec(colsum(com.occurrences))
speciestotals(com::AbstractComMatrix) = vec(rowsum(com.occurrences))

size(com::AbstractComMatrix) = size(occurrences(com))
size(com::AbstractComMatrix, dims...) = size(occurrences(com), dims...)

#----------------------------------------------------------------------------------------------
# Mutating and concatenating

#-------------------------------------------------------------------------------------------------
# show methods

summary(com::AbstractComMatrix) = "$(nsites(com))x$(nspecies(com)) $(typeof(com))"

function createsummaryline(vec::AbstractVector{T}) where T<:AbstractString
    linefunc(vec) = mapreduce(x->x*", ", *, vec[1:(end-1)])*vec[end]
    length(vec) == 1 && return vec[1]
    length(vec) < 6 && return linefunc(vec)
    linefunc(vec[1:3])*"..."*linefunc(vec[(end-1):end])
end

function show(io::IO, com::ComMatrix)
    sp = createsummaryline(speciesnames(com))
    si = createsummaryline(sitenames(com))
    println(io, "ComMatrix with $(nspecies(com)) species in $(nsites(com)) sites\n\nSpecies names:\n$(sp)\n\nSite names:\n$(si)")
end

function show(io::IO, com::Assemblage)
    sp = createsummaryline(speciesnames(com))
    si = createsummaryline(sitenames(com))
    println(io, "Assemblage with $(nspecies(com)) species in $(nsites(com)) sites\n\nSpecies names:\n$(sp)\n\nSite names:\n$(si)")
end

function show(io::IO, sd::SiteData)
    println(io, "Spatial data set with $(nsites(sd)) sites\n\nSite names:\n$(createsummarylines(sitenames(sd)))")
end

#TODO also create render functions for Juno

getindex(com::AbstractComMatrix, inds...) = ComMatrix(getindex(com.occurrences, inds...))

setindex!(com::AbstractComMatrix, X, inds...) = setindex!(com.occurrences, X, inds...)

function getindex(site::S, inds) where S<:SELocations
  S(coordinates(site)[inds,:], site.sitestats[inds,:])
end

function getindex(com::SEAssemblage, ind::Symbol)
    if ind in names(com.site.sitestats)
        return com.site.sitestats[ind]
    elseif ind in names(com.occ.traits)
        return com.occ.traits[ind]
    else
        error("No such name in traits or sitestats")
    end
end


# This file is a part of BAT.jl, licensed under the MIT License (MIT).

struct VariateTransformChain{
    VF <: VariateForm,
    ST <: VariateSpace,
    SF <: VariateSpace,
    N,
    NTT<:NTuple{N,VariateTransform{VF}}
} <: VariateTransform{VF,ST,SF}
    transforms::NTT
end

function VariateTransformChain(transforms::NTuple{N,VariateTransform{VF}}) where {N,VF}
    shape = varshape(first(values(transforms)))
    @argcheck all(isequal(shape), map(varshape, values(transforms)))
    ST = typeof(target_space(last(transforms)))
    SF = typeof(source_space(last(transforms)))
    NTT = typeof(transforms)
    VariateTransformChain{VF,ST,SF,N,NTT}(transforms)
end

import Base.∘

# ToDo: Utilize specialized merge operations at ends of chains:
∘(a::VariateTransform, b::VariateTransform) = VariateTransformChain((b, a))
∘(a::VariateTransform, b::VariateTransformChain) = VariateTransformChain((b, a.transforms...))
∘(a::VariateTransformChain, b::VariateTransform) = VariateTransformChain((b.transforms..., a))
∘(a::VariateTransformChain, b::VariateTransformChain) = VariateTransformChain((b.transforms..., a.transforms...))

∘(a::VariateTransform{VF}, b::IdentityVT{VF}) where VF = (@argcheck varshape(a) == varshape(b); a)
∘(a::IdentityVT{VF}, b::VariateTransform{VF}) where VF = (@argcheck varshape(a) == varshape(b); b)
∘(a::IdentityVT{VF}, b::IdentityVT{VF}) where VF = (@argcheck varshape(a) == varshape(b); a)

target_space(trafo::VariateTransformChain) = target_space(last(orig.transforms))
source_space(trafo::VariateTransformChain) = source_space(last(orig.transforms))

function _apply_trafo_chain(v::Any, prev_ladj::Real)
    (v = v, ladj = prev_ladj)
end

function _apply_trafo_chain(v::Any, prev_ladj::Real, t_first::VariateTransform, ts::VariateTransform...)
    r1 = apply_vartrafo_impl(t_first, v, prev_ladj)
    r2 = _apply_trafo_chain(r1.v, r1.ladj, ts...)
    (v = r2.v, ladj = r2.ladj)
end

Base.inv(trafo::VariateTransformChain) = VariateTransformChain(reverse(map(inv, trafo.transforms)))

ValueShapes.varshape(trafo::VariateTransformChain) = varshape(first(values(trafo.transforms)))

apply_vartrafo_impl(trafo::VariateTransformChain, v::Any, prev_ladj::Real) = _apply_trafo_chain(v, prev_ladj, trafo.transforms...)

apply_vartrafo_impl(trafo::InvVT{<:VariateTransformChain}, v::Any, prev_ladj::Real) = apply_vartrafo_impl(inv(trafo.orig), v, prev_ladj)

# ToDo: Better fix for ambiguity:
function apply_vartrafo_impl(trafo::VariateTransformChain{Univariate}, v::AbstractArray{<:Real,0}, prev_ladj::Real)
    r = apply_vartrafo_impl(trafo, v[], prev_ladj)
    (v = fill(r.v), ladj = r.ladj)
end


using PkgBenchmark
using HMMBase

# TODO: Add timestamp to result files
# results = benchmarkpkg(parent(pathof(HMMBase)))
results = benchmarkpkg(pwd(), resultfile = "benchmark.json")
export_markdown("benchmark.md", results)


module CommunalHelperChain

using ..Ahorn, Maple
using Ahorn.CommunalHelper

@mapdef Entity "CommunalHelper/Chain" Chain(
    x::Integer,
    y::Integer,
    extraJoints::Integer=0,
    outline::Bool=true,
)

const placements = Ahorn.PlacementDict(
    "Chain (Communal Helper)" => Ahorn.EntityPlacement(
        Chain,
        "line",
    ),
)

Ahorn.nodeLimits(entity::Chain) = 1, 1

function Ahorn.selection(entity::Chain)
    nodes = get(entity.data, "nodes", ())
    x, y = Ahorn.position(entity)

    res = Ahorn.Rectangle[Ahorn.Rectangle(x - 4, y - 4, 8, 8)]

    for node in nodes
        nx, ny = node

        push!(res, Ahorn.Rectangle(nx - 4, ny - 4, 8, 8))
    end

    return res
end

function renderChain(ctx::Ahorn.Cairo.CairoContext, entity::Chain, color::Ahorn.colorTupleType=ChainColor)
    x, y = Ahorn.position(entity)

    start = (x, y)
    stop = get(entity.data, "nodes", [start])[1]
    nx, ny = Int.(stop)
    control = (start .+ stop) ./ 2 .+ (0, 24)

    curve = Ahorn.SimpleCurve(start, stop, control)
    Ahorn.drawSimpleCurve(ctx, curve, color, thickness=1)
end

Ahorn.renderSelectedAbs(ctx::Ahorn.Cairo.CairoContext, entity::Chain) = renderChain(ctx, entity, Ahorn.colors.selection_selected_fc)

Ahorn.renderAbs(ctx::Ahorn.Cairo.CairoContext, entity::Chain, room::Room) = renderChain(ctx, entity, (0.0, 0.5, 0.5, 1.0))

end