"""
    ResultsStruct

This struct gets updated in the GHPGHX.jl module and is used to record the results/outputs.
"""
Base.@kwdef mutable struct ResultsStruct
    # These get modified by the function below
    # Some of these values are required inputs, those that do not have "= value"
    X_Now::Array{Float64, 1}
    FX_Now::Array{Float64, 1}
    N_Bores::Array{Int64, 1}
    N_Bores_Final::Int64 = 0
    FX_Final::Float64 = 0.0
    Length_Boreholes::Float64 = 0.0
    StorageVolume::Float64 = 0.0
    Depth_DST::Float64 = 0.0
    Mdot_Borehole::Float64 = 0.0
    Radius_Field::Float64 = 0.0
    N_Circuits::Int64 = 0
    L_Circuit::Float64 = 0.0
    Mass_Circuit::Float64 = 0.0

    # These get reset each sizing iteration
    Tmax_GHX::Float64 = -9.99e20
    Tmin_GHX::Float64 = 9.99e20
    Power_GHXPump::Float64 = 0.0
    Qfluid_GHXPump::Float64 = 0.0
    Power_WSHP_H::Float64 = 0.0
    Power_WSHP_C::Float64 = 0.0
    LoadMet_WSHP_H::Float64 = 0.0
    LoadMet_WSHP_C::Float64 = 0.0
    Q_Rejected_Total::Float64 = 0.0
    Q_Absorbed_Total::Float64 = 0.0
    Q_GHX_Net::Float64 = 0.0
    Q_GHX_In::Float64 = 0.0
    Q_GHX_Top::Float64 = 0.0
    Q_GHX_Sides::Float64 = 0.0
    Q_GHX_Bottom::Float64 = 0.0
    Q_GHX_Stored::Float64 = 0.0
    Qf_GHX_Stored::Float64 = 0.0

    # Array results
    total_hours::Int64  # Required argument to be passed by InputsStruct
    P_GHXPump_Hourly::Array{Float64, 1} = zeros(total_hours)
    P_WSHPh_Hourly::Array{Float64, 1} = zeros(total_hours)
    P_WSHPc_Hourly::Array{Float64, 1} = zeros(total_hours)
    Qh_Hourly::Array{Float64, 1} = zeros(total_hours)
    Qc_Hourly::Array{Float64, 1} = zeros(total_hours)
    EWT::Array{Float64, 1} = zeros(total_hours)
end

"""
    init_sizing!(r::ResultsStruct, p::InputsStruct, size_iter::Int64)

Performs the initial sizing for the GHPGHX.jl `size_borefield(p)` function.
"""
function init_sizing!(r::ResultsStruct, p::InputsStruct, size_iter::Int64)
    #Changes each sizing iteration
    r.N_Bores[size_iter] = floor(r.X_Now[size_iter] / p.Depth_Bores) + 1
    r.Length_Boreholes = r.X_Now[size_iter] / r.N_Bores[size_iter]
    if p.SpacingType == 1
        r.StorageVolume = r.Length_Boreholes * r.N_Bores[size_iter] * pi * (p.BoreSpacing * 0.525)^2
    else
        r.StorageVolume = p.BoreSpacing^2 * r.Length_Boreholes * r.N_Bores[size_iter]
    end
    r.Depth_DST = 1.5 * r.Length_Boreholes
    r.Mdot_Borehole = p.Mdot_GHXPump / r.N_Bores[size_iter] * p.N_Series
    r.Radius_Field = (r.StorageVolume / (pi * r.Length_Boreholes))^0.5
    r.N_Circuits = r.N_Bores[size_iter] / p.N_Series
    r.L_Circuit = 2.0 * p.N_Series * r.Length_Boreholes + (p.N_Series - 1) / p.N_Series * r.Radius_Field
    r.Mass_Circuit = r.L_Circuit * pi * p.Ri_Pipe^2 * p.Rho_GHXFluid
    nothing
end

"""
    get_ghpghx_results_for_reopt(r::ResultsStruct, p::InputsStruct)

Extract the required results for input to REopt /job: average hourly dispatch performance and other summary performance metrics.
"""
function get_ghpghx_results_for_reopt(r::ResultsStruct, p::InputsStruct)
    results_dict = Dict{Any,Any}()
    results_dict["number_of_boreholes"] = r.N_Bores_Final
    results_dict["length_boreholes_ft"] = round(r.Length_Boreholes * p.METER_TO_FEET, digits=1)
    results_dict["yearly_heating_heatpump_electric_consumption_series_kw"] = zeros(8760)
    results_dict["yearly_cooling_heatpump_electric_consumption_series_kw"] = zeros(8760)
    results_dict["yearly_ghx_pump_electric_consumption_series_kw"] = zeros(8760)
    # Get average electric consumption
    for yr in 1:p.simulation_years
        results_dict["yearly_heating_heatpump_electric_consumption_series_kw"] += round.(r.P_WSHPh_Hourly[(yr-1)*8760+1:yr*8760] / p.simulation_years, digits=3)
        results_dict["yearly_cooling_heatpump_electric_consumption_series_kw"] += round.(r.P_WSHPc_Hourly[(yr-1)*8760+1:yr*8760] / p.simulation_years, digits=3)
        results_dict["yearly_ghx_pump_electric_consumption_series_kw"] += round.(r.P_GHXPump_Hourly[(yr-1)*8760+1:yr*8760] / p.simulation_years, digits=3)
    end
    results_dict["yearly_total_electric_consumption_series_kw"] =
        results_dict["yearly_heating_heatpump_electric_consumption_series_kw"] +
        results_dict["yearly_cooling_heatpump_electric_consumption_series_kw"] +
        results_dict["yearly_ghx_pump_electric_consumption_series_kw"]
    results_dict["yearly_total_electric_consumption_kwh"] = round(sum(results_dict["yearly_total_electric_consumption_series_kw"]), digits=1)
    results_dict["peak_heating_heatpump_thermal_ton"] = round(p.PeakTons_WSHP_H, digits=3)
    results_dict["peak_cooling_heatpump_thermal_ton"] = round(p.PeakTons_WSHP_C, digits=3)
    results_dict["peak_combined_heatpump_thermal_ton"] = round(p.PeakTons_WSHP_GHX, digits=3)
    results_dict["max_eft_f"] = round(maximum(r.EWT) * 1.8 + 32.0)
    results_dict["min_eft_f"] = round(minimum(r.EWT) * 1.8 + 32.0)
    # Calculate average COP for heating and cooling; estimate allocation of pump power to heating and cooling by thermal energy served
    heating_thermal_kwh = sum(p.HeatingThermalLoadKW)
    cooling_thermal_kwh = sum(p.CoolingThermalLoadKW)
    heating_thermal_frac = heating_thermal_kwh / (heating_thermal_kwh + cooling_thermal_kwh)
    cooling_thermal_frac = cooling_thermal_kwh / (heating_thermal_kwh + cooling_thermal_kwh)
    ghx_pump_electric_kwh = sum(results_dict["yearly_ghx_pump_electric_consumption_series_kw"])
    heating_ghx_pump_electric_kwh = heating_thermal_frac * ghx_pump_electric_kwh
    cooling_ghx_pump_electric_kwh = cooling_thermal_frac * ghx_pump_electric_kwh
    heating_heatpump_electric_kwh = sum(results_dict["yearly_heating_heatpump_electric_consumption_series_kw"])
    cooling_heatpump_electric_kwh = sum(results_dict["yearly_cooling_heatpump_electric_consumption_series_kw"])
    results_dict["heating_cop_avg"] = round(heating_thermal_kwh / (heating_heatpump_electric_kwh + heating_ghx_pump_electric_kwh), digits=3)
    results_dict["cooling_cop_avg"] = round(cooling_thermal_kwh / (cooling_heatpump_electric_kwh + cooling_ghx_pump_electric_kwh), digits=3)
    results_dict["solved_eft_error_f"] = round(r.FX_Final * 1.8, digits=3)
    return results_dict
end


"""
nullpert.jl

Null Perturbation
"""

immutable NullPerturbation <: Perturbation
  ncomp::Int # number of components
end

#
# Thermodynamic functions
#
kinetic(null::NullPerturbation, ref::ReferenceSystem)::Float64 = 0
internal(null::NullPerturbation, ref::ReferenceSystem)::Float64 = 0
entropy(null::NullPerturbation, ref::ReferenceSystem, T::Float64)::Float64 = 0
helmholtz(null::NullPerturbation, ref::ReferenceSystem)::Float64 = 0

function pairpotential(null::NullPerturbation)::Array{Function,2}
  N::Int = null.ncomp
  ret = Array{Function,2}(N,N)

  for i in 1:N, j in 1:N
    i > j && continue
    u(r)::Float64 = 0
    ret[i,j] = ret[j,i] = u
  end

  return ret
end


"""

```
mddf_frame_self!(iframe::Int, framedata::FrameData, options::Options, RNG, R::Result)
```

Computes the MDDF for a single frame, modifies the data in the `R` (`Result`) structure.

"""
function mddf_frame_self!(
    iframe::Int,
    framedata::FrameData,
    options::Options,
    RNG,
    R::Result,
)

    # Simplify code by assigning some shortened names
    trajectory = framedata.trajectory
    volume_frame = framedata.volume_frame
    rdf_count_random_frame = framedata.rdf_count_random_frame
    md_count_random_frame = framedata.md_count_random_frame
    dc = framedata.dc
    dmin_mol = framedata.dmin_mol
    dref_mol = framedata.dref_mol
    x_solvent_random = framedata.x_solvent_random
    lc_solvent = framedata.lc_solvent
    solute = trajectory.solute
    solvent = trajectory.solvent
    x_solute = trajectory.x_solute
    x_solvent = trajectory.x_solvent

    # Reset counters for this frame
    reset!(volume_frame)
    @. rdf_count_random_frame = 0.0
    @. md_count_random_frame = 0.0

    # get pbc sides in this frame
    sides = getsides(trajectory, iframe)

    volume_frame.total = sides[1] * sides[2] * sides[3]
    R.volume.total = R.volume.total + volume_frame.total

    R.density.solute = R.density.solute + (solute.nmols / volume_frame.total)
    R.density.solvent = R.density.solvent + (solvent.nmols / volume_frame.total)

    # Add the box side information to the box structure, in this frame
    box = Box(options.lcell, sides, R.cutoff)

    # Will wrap everything relative to the reference atom of the first molecule
    # and move everything such that that center is in the origin. This is important
    # to simplify the computation of cell indexes, as the minimum coordinates are
    # automatically -side/2 at each direction
    solute_center = x_solute[R.irefatom]
    wrap!(x_solvent, sides, solute_center)
    center_to_origin!(x_solvent, solute_center)

    # Initialize linked cells
    initcells!(x_solvent, box, lc_solvent)

    # Check if the cutoff is not too large considering the periodic cell size
    if R.cutoff > sides[1] / 2.0 || R.cutoff > sides[2] / 2.0 || R.cutoff > sides[3] / 2.0
        error(
            """
        in MDDF: cutoff or dbulk > periodic_dimension/2 in frame: $iframe
                 max(cutoff,dbulk) = $(R.cutoff)
                 sides read from file = $(sides)
                 If sides are zero it is likely that the PBC information is not available
                 in the trajectory file.
      """,
        )
    end

    n_solvent_in_bulk = 0.0
    local n_dmin_in_bulk
    for isolvent = 1:solvent.nmols-1
        # We need to do this one solute molecule at a time to avoid exploding the memory
        # requirements
        x_this_solute = viewmol(isolvent, x_solvent, solvent)

        # Compute all distances between solute and solvent atoms which are smaller than the
        # cutoff (this is the most computationally expensive part), the distances are returned
        # in the dc structure
        cutoffdistances_self!(
            R.cutoff,
            x_this_solute,
            x_solvent,
            lc_solvent,
            box,
            dc,
            solvent,
            isolvent,
        )

        # For each solute molecule, update the counters (this is highly suboptimal, because
        # within updatecounters there are loops over solvent molecules, in such a way that
        # this will loop with cost nsolute*nsolvent. However, I cannot see an easy solution
        # at this point with acceptable memory requirements
        n_dmin_in_bulk, n_dref_in_bulk =
            updatecounters!(R, solute, solvent, dc, dmin_mol, dref_mol)
        n_solvent_in_bulk += n_dref_in_bulk
    end
    # The normalization below is tricky. The number that comes out from updatecounters is the
    # sum, for every solvent molecule (minus one) of the distances that were not found to be in
    # the solute domain. The total number of distances is n^2, because the sum inside updatecounter
    # is made for all molecules (the same function is used for cross-distribution), but we called updatecounters
    # only (n-1) times, so the actual sum of the distances considered is n(n-1). From this set
    # the number of distances returned in n_dref_in_bulk is r=n(n-1)-nd/2, where nd is the
    # number of distances in the domain. Thus, the number of distances in the domain, considering
    # symmetric terms, is nd=2n(n-1)-2r. The average number of distances in the domain, per
    # molecule, is thus nd/n=2(n-1)-2r/n. Finally, the number of solvent molecules in
    # bulk, for each molecule, is the total number of other molecules, (n-1), minus the
    # number of molecules in the domain, that is (n-1)-nd/n=(n-1)-2(n-1)+2r/n, which
    # finally simplifies to 2r/n-(n-1), which is the equation below.
    n_solvent_in_bulk = 2 * n_solvent_in_bulk / solvent.nmols - (solvent.nmols - 1)

    #
    # Computing the random-solvent distribution to compute the random minimum-distance count
    #
    for i = 1:options.n_random_samples
        # generate random solvent box, and store it in x_solvent_random
        for j = 1:solvent.nmols
            # Choose randomly one molecule from the bulk, if there are actually bulk molecules
            if n_dmin_in_bulk != 0
                jmol =
                    dmin_mol[random(RNG, solvent.nmols-n_dmin_in_bulk+1:solvent.nmols)].jmol
            else
                jmol = random(RNG, 1:solvent.nmols)
            end
            # Indexes of this molecule in the x_solvent array
            x_ref = viewmol(jmol, x_solvent, solvent)
            # Indexes of the random molecule in random array
            x_rnd = viewmol(j, x_solvent_random, solvent)
            # Generate new random coordinates (translation and rotation) for this molecule
            random_move!(x_ref, R.irefatom, sides, x_rnd, RNG)
        end

        # wrap random solvent coordinates to origin
        wrap!(x_solvent_random, sides)

        # Initialize linked cells
        initcells!(x_solvent_random, box, lc_solvent)

        # Choose randomly one solute molecule to be the solute in this sample
        i_rand_mol = random(RNG, 1:solvent.nmols)
        x_this_solute = viewmol(i_rand_mol, x_solvent, solvent)

        # Compute all distances between solute and solvent atoms which are smaller than the
        # cutoff (this is the most computationally expensive part), the distances are returned
        # in the dc structure (here we do not use the self function)
        cutoffdistances!(R.cutoff, x_this_solute, x_solvent_random, lc_solvent, box, dc)

        # Update the counters and get the number of solvent molecules in bulk
        updatecounters!(
            R,
            rdf_count_random_frame,
            md_count_random_frame,
            solvent,
            dc,
            dmin_mol,
            dref_mol,
        )

    end # random solvent sampling

    # Update global counters with the data of this frame
    update_counters_frame!(
        R,
        rdf_count_random_frame,
        md_count_random_frame,
        volume_frame,
        solute,
        solvent,
        n_solvent_in_bulk,
    )

    nothing
end


###
### LongSequence specific specializations of src/biosequence/transformations.jl
###

"""
    resize!(seq, size, [force::Bool])

Resize a biological sequence `seq`, to a given `size`. Does not resize the underlying data
array unless the new size does not fit. If `force`, always resize underlying data array.
"""
function Base.resize!(seq::LongSequence{A}, size::Integer, force::Bool=false) where {A}
    if size < 0
        throw(ArgumentError("size must be non-negative"))
    else
        if force | (seq_data_len(A, size) > seq_data_len(A, length(seq)))
            resize!(seq.data, seq_data_len(A, size))
        end
        seq.len = size
        return seq
    end
end

function Base.filter!(f, seq::LongSequence)
    writeindex = bitindex(seq, 0)
    chunk = zero(UInt64)
    bps = bits_per_symbol(seq)
    for i in eachindex(seq)
        encoded_symbol = extract_encoded_element(bitindex(seq, i), encoded_data(seq))
        symbol = decode(Alphabet(seq), encoded_symbol)
        if f(symbol)
            writeindex += bps
            chunk |= encoded_symbol << offset(writeindex)
            # If the offset is now zero, the chunk is full
            if iszero(offset(writeindex + bps))
                seq.data[index(writeindex)] = chunk
                chunk = zero(UInt64)
            end
        end
    end
    # If it's not zero, we need to write the first chunk
    if !iszero(offset(writeindex + bps))
        seq.data[index(writeindex)] = chunk
    end
    resize!(seq, div((writeindex + bps).val, bps))
end

function Base.map!(f, seq::LongSequence)
    for i in 1:lastindex(seq)
        unsafe_setindex!(seq, f(inbounds_getindex(seq, i)), i)
    end
    return seq
end

"""
    reverse!(seq::LongSequence)

Reverse a biological sequence `seq` in place.
"""
Base.reverse!(seq::LongSequence{<:Alphabet}) = _reverse!(seq, BitsPerSymbol(seq))

"""
    reverse(seq::LongSequence)

Create reversed copy of a biological sequence.
"""
Base.reverse(seq::LongSequence{<:Alphabet}) = _reverse(seq, BitsPerSymbol(seq))

# Fast path for non-inplace reversion
@inline function _reverse(seq::LongSequence{A}, B::BT) where {A <: Alphabet,
    BT <: Union{BitsPerSymbol{2}, BitsPerSymbol{4}, BitsPerSymbol{8}}}
    cp = LongSequence{A}(unsigned(length(seq)))
    reverse_data_copy!(identity, cp.data, seq.data, seq_data_len(seq) % UInt, B)
    return zero_offset!(cp)
end

_reverse(seq::LongSequence{<:Alphabet}, ::BitsPerSymbol) = reverse!(copy(seq))

# Generic fallback
function _reverse!(seq::LongSequence{<:Alphabet}, ::BitsPerSymbol)
    i, j = 1, lastindex(seq)
    @inbounds while i < j
        seq[i], seq[j] = seq[j], seq[i]
        i += 1
        j -= 1
    end
    return seq
end

@inline function _reverse!(seq::LongSequence{<:Alphabet}, B::BT) where {
    BT <: Union{BitsPerSymbol{2}, BitsPerSymbol{4}, BitsPerSymbol{8}}}
    # We need to account for the fact that the seq may not use all its stored data
    reverse_data!(identity, seq.data, seq_data_len(seq) % UInt, B)
    return zero_offset!(seq)
end

# Reversion of chunk bits may have left-shifted data in chunks, this function right shifts
# all chunks by up to 63 bits.
# This is written so it SIMD parallelizes - careful with changes
@inline function zero_offset!(seq::LongSequence{A}) where A <: Alphabet
    offs = (64 - offset(bitindex(seq, length(seq)) + bits_per_symbol(A()))) % UInt
    zero_offset!(seq, offs)
end

@inline function zero_offset!(seq::LongSequence{A}, offs::UInt) where A <: Alphabet
    rshift = offs
    lshift = 64 - rshift
    len = length(seq.data)
    @inbounds if !iszero(lshift)
        this = seq.data[1]
        for i in 1:len-1
            next = seq.data[i+1]
            seq.data[i] = (this >>> (unsigned(rshift) & 63)) | (next << (unsigned(lshift) & 63))
            this = next
        end
        seq.data[len] >>>= (unsigned(rshift) & 63)
    end
    return seq
end

# Reverse chunks in data vector and each symbol within a chunk. Chunks may have nonzero
# offset after use, so use zero_offset!
@inline function reverse_data!(pred, data::Vector{UInt64}, len::UInt, B::BT) where {
    BT <: Union{BitsPerSymbol{2}, BitsPerSymbol{4}, BitsPerSymbol{8}}}
    @inbounds @simd ivdep for i in 1:len >>> 1
        data[i], data[len-i+1] = pred(reversebits(data[len-i+1], B)), pred(reversebits(data[i], B))
    end
    @inbounds if isodd(len)
        data[len >>> 1 + 1] = pred(reversebits(data[len >>> 1 + 1], B))
    end
end

@inline function reverse_data_copy!(pred, dst::Vector{UInt64}, src::Vector{UInt64}, len::UInt,
    B::BT) where {BT <: Union{BitsPerSymbol{2}, BitsPerSymbol{4}, BitsPerSymbol{8}}}
    @inbounds @simd for i in eachindex(dst)
        dst[i] = pred(reversebits(src[len - i + 1], B))
    end
end

"""
    complement!(seq)

Make a complement sequence of `seq` in place.
"""
function complement!(seq::LongSequence{A}) where {A<:NucleicAcidAlphabet}
    seqdata = seq.data
    @inbounds for i in eachindex(seqdata)
        seqdata[i] = complement_bitpar(seqdata[i], Alphabet(seq))
    end
    return seq
end

function complement!(s::LongSubSeq{A}) where {A <: NucleicAcidAlphabet}
    bps = bits_per_symbol(A())
    bi = firstbitindex(s)
    i = 1
    stop = lastbitindex(s) + bps
    @inbounds while (!iszero(offset(bi)) & (bi < stop))
        s[i] = complement(s[i])
        bi += bps
        i += 1
    end
    @inbounds for j in index(bi):index(stop)-1
        s.data[j] = complement_bitpar(s.data[j], Alphabet(s))
        bi += 64
        i += symbols_per_data_element(s)
    end
    @inbounds while bi < stop
        s[i] = complement(s[i])
        bi += bps
        i += 1
    end
    return s
end

function reverse_complement!(seq::LongSequence{<:NucleicAcidAlphabet})
    pred = x -> complement_bitpar(x, Alphabet(seq))
    reverse_data!(pred, seq.data, seq_data_len(seq) % UInt, BitsPerSymbol(seq))
    return zero_offset!(seq)
end

function reverse_complement(seq::LongSequence{<:NucleicAcidAlphabet})
    cp = typeof(seq)(unsigned(length(seq)))
    pred = x -> complement_bitpar(x, Alphabet(seq))
    reverse_data_copy!(pred, cp.data, seq.data, seq_data_len(seq) % UInt, BitsPerSymbol(seq))
    return zero_offset!(cp)
end

function Random.shuffle!(seq::LongSequence)
    # Fisher-Yates shuffle
    @inbounds for i in 1:lastindex(seq) - 1
        j = rand(i:lastindex(seq))
        seq[i], seq[j] = seq[j], seq[i]
    end
    return seq
end


module Compatibility_test

using SpikeSorting

if VERSION > v"0.7-"
    using Test
else
    using Base.Test
end

num_channels=1;
(buf,nums)=output_buffer(num_channels);

function make_voltage(num_channels)

    v=rand(1:1000, 1000, num_channels);
    count = 1
    for i=525:550
        v[i,1] = v[i,1] + count
        count = count + 10
    end
    for i=551:775
        v[i,1]+=1
        count -= 10
    end
    v
end

v=make_voltage(num_channels)

detects=subtypes(SpikeSorting.Detect)
aligns=subtypes(SpikeSorting.Align)
features=subtypes(SpikeSorting.Feature)
reduces=subtypes(SpikeSorting.Reduction)
clusters=subtypes(SpikeSorting.Cluster)
thresholds=subtypes(SpikeSorting.Threshold)

for d in detects
    for a in aligns
        for f in features
            for r in reduces
                for c in clusters
                    for t in thresholds

                        s1=create_multi(d(),c(),a(),f(),r(),t(),num_channels)

                        cal!(s1,v,buf,nums,0)
                        cal!(s1,v,buf,nums,1)
                        cal!(s1,v,buf,nums,2)
                        onlinesort!(s1,v,buf,nums)

                            @test s1[1].sigend == v[end-74:end,1]

                    end
                end
            end
        end
    end
end

end


"""
    SubstitutionModel

Abstract type for substitution models,
using a continous time Markov model on a phylogeny.
Adapted from [SubstitutionModels.jl](https://github.com/BioJulia/SubstitutionModels.jl/)
in BioJulia.

For variable rates, see [`RateVariationAcrossSites`](@ref)

For sub types, see [`NucleicAcidSubstitutionModel`](@ref), [`TraitSubstitutionModel`](@ref)

All these models are supposed to have fields `rate` and `eigeninfo`.
"""
abstract type SubstitutionModel end #ideally, we'd like this to be SubstitutionModels.SubstitionModel
const SM = SubstitutionModel
const Qmatrix = StaticArrays.SMatrix{4, 4, Float64}
const Pmatrix = StaticArrays.MMatrix{4, 4, Float64}
const Bmatrix = StaticArrays.MMatrix{2, 2, Float64}

"""
    TraitSubstitutionModel

For subtypes, see [`BinaryTraitSubstitutionModel`](@ref),
[`EqualRatesSubstitutionModel`](@ref),
[`TwoBinaryTraitSubstitutionModel`](@ref)
"""
abstract type TraitSubstitutionModel{T} <: SubstitutionModel end #this accepts labels
const TSM = TraitSubstitutionModel{T} where T #T is type of labels

"""
    NucleicAcidSubstitutionModel

Adapted from [SubstitutionModels.jl](https://github.com/BioJulia/SubstitutionModels.jl/)
in BioJulia. The same [`Q`](@ref) and [`P`](@ref) function names are used for the
transition rates and probabilities.

For subtypes, see [`JC69`](@ref), [`HKY85`](@ref)
"""
abstract type NucleicAcidSubstitutionModel <: SubstitutionModel end
const NASM = NucleicAcidSubstitutionModel
function Base.show(io::IO, obj::SM)
    error("show not defined for $(typeof(obj)).")
end

"""
    nparams(model)

Number of parameters for a given trait evolution model
(length of field `model.rate`).
"""
nparams(obj::SM) = error("nparams not defined for $(typeof(obj)).")

"""
    setrates!(model, rates)

update rates then call [`seteigeninfo!`](@ref) to update a model's eigeninfo
"""
function setrates!(obj::SM, rates::AbstractVector)
    obj.rate[:] = rates
    seteigeninfo!(obj)
end

"""
    seteigeninfo!(obj)

Calculate eigenvalue & eigenfector information for a substitution model (SM) object
(as needed to calculate transition rate matrices) and store this info within the object.
"""
function seteigeninfo!(obj::SM)
    error("seteigeninfo! not finalized yet for generic substitution models")
    # eig_vals, eig_vecs = eigen(Q(obj)) #this assumes that every SM has an eigeninfo and has
    # obj.eigeninfo[:] = eig_vals, eig_vecs, inv(eig_vecs)
    # then make P! use seteigeninfo! by for a generic SM
end

"""
    getlabels(model)

State labels of a substitution model.
"""
function getlabels(obj::SM)
    error("Model must be of type TraitSubstitutionModel or NucleicAcidSubstitutionModel. Got $(typeof(obj))")
end

"""
    nstates(model)

Number of character states for a given evolution model.
"""
nstates(obj::SM) = error("nstates not defined for $(typeof(obj)).")

"""
For example, this is 4 for a `NucleicAcidSubstitutionModel`.

```jldoctest
julia> nstates(JC69([0.03], false))
4

julia> nstates(HKY85([.5], [0.25, 0.25, 0.25, 0.25]))
4
```
"""
function nstates(obj::NASM)
    return 4::Int
end

function getlabels(obj::TSM)
    return obj.label
end

"""
for a given [`NucleicAcidSubstitutionModel`](@ref), labels are symbols
from [BioSymbols](https://github.com/BioJulia/BioSymbols.jl).
For now, only ACGTs are allowed. (When fitting data, any ambiguity code
in the data would be treated as missing value).

# examples

```jldoctest
julia> getlabels(JC69([0.03], false))
4-element Array{BioSymbols.DNA,1}:
 DNA_A
 DNA_C
 DNA_G
 DNA_T

julia> getlabels(HKY85([.5], repeat([0.25], 4)))
4-element Array{BioSymbols.DNA,1}:
 DNA_A
 DNA_C
 DNA_G
 DNA_T

```
"""
function getlabels(::NASM)
    return [BioSymbols.DNA_A, BioSymbols.DNA_C, BioSymbols.DNA_G, BioSymbols.DNA_T]
end

"""
    Q(model)

Substitution rate matrix for a given substitution model:
Q[i,j] is the rate of transitioning from state i to state j.
"""
Q(obj::SM) = error("rate matrix Q not defined for $(typeof(obj)).")

"""
    showQ(IO, model)

Print the Q matrix to the screen, with trait states as labels on rows and columns.
adapted from prettyprint function by mcreel, found 2017/10 at
https://discourse.julialang.org/t/display-of-arrays-with-row-and-column-names/1961/6
"""
function showQ(io::IO, obj::SM)
    M = Q(obj)
    pad = max(8,maximum(length(getlabels(obj))+1))
    for i = 1:size(M,2) # print the header
        print(io, lpad(getlabels(obj)[i],(i==1 ? 2*pad : pad), " "))
    end
    print(io, "\n")
    for i = 1:size(M,1) # print one row per state
        if getlabels(obj) != ""
            print(io, lpad(getlabels(obj)[i],pad," "))
        end
        for j = 1:size(M,2)
            if j == i
                print(io, lpad("*",pad," "))
            else
                fmt = "%$(pad).4f"
                @eval(@printf($io,$fmt,$(M[i,j])))
            end
        end
        print(io, "\n")
    end
end

"""
    P(obj, t)

Probability transition matrix for a [`TraitSubstitutionModel`](@ref), of the form

    P[1,1] ... P[1,k]
       .          .
       .          .
    P[k,1] ... P[k,k]

where P[i,j] is the probability of ending in state j after time t,
given that the process started in state i.
see also: [`P!`](@ref).

HKY example:
```jldoctest
julia> m1 = HKY85([0.5], [0.20, 0.30, 0.30, 0.20])
HKY85 Substitution Model base frequencies: [0.2, 0.3, 0.3, 0.2]
relative rate version with transition/tranversion ratio kappa = 0.5,
 scaled so that there is one substitution per unit time
rate matrix Q:
               A       C       G       T
       A       *  0.4839  0.2419  0.3226
       C  0.3226       *  0.4839  0.1613
       G  0.1613  0.4839       *  0.3226
       T  0.3226  0.2419  0.4839       *

julia> PhyloNetworks.P(m1, 3.0)
4×4 StaticArrays.MArray{Tuple{4,4},Float64,2,16} with indices SOneTo(4)×SOneTo(4):
 0.217509  0.198417  0.190967  0.198417
 0.297625  0.312992  0.297625  0.28645
 0.28645   0.297625  0.312992  0.297625
 0.198417  0.190967  0.198417  0.217509
```

Juke-Cantor example:
```jldoctest
julia> m1 = JC69([1.]);

julia> PhyloNetworks.P(m1, 0.2)
4×4 StaticArrays.MArray{Tuple{4,4},Float64,2,16} with indices SOneTo(4)×SOneTo(4):
 0.824446   0.0585179  0.0585179  0.0585179
 0.0585179  0.824446   0.0585179  0.0585179
 0.0585179  0.0585179  0.824446   0.0585179
 0.0585179  0.0585179  0.0585179  0.824446
```
"""
@inline function P(obj::SM, t::Float64)
    t >= 0.0 || error("substitution model: >=0 branch lengths are needed")
    k = nstates(obj)
    Pmat = MMatrix{k,k}(repeat([0.0], k*k))
    return P!(Pmat, obj, t)
end

"""
    P!(Pmat::AbstractMatrix, obj::SM, t::Float64)

Fill in the input matrix `Pmat` with the transition rates
to go from each state to another in time `t`, according to rates in `Q`.
see also: [`P`](@ref).

```jldoctest
julia> m1 = BinaryTraitSubstitutionModel([1.0,2.0], ["low","high"])
Binary Trait Substitution Model:
rate low→high α=1.0
rate high→low β=2.0

julia> PhyloNetworks.P!(Matrix{Float64}(undef,2,2), m1, 3.0) # fills an uninitialized 2x2 matrix of floats
2×2 Array{Float64,2}:
 0.666708  0.333292
 0.666584  0.333416

julia> m2 = JC69([1.]);

julia> PhyloNetworks.P!(Matrix{Float64}(undef,4,4), m2, 0.2)
4×4 Array{Float64,2}:
 0.824446   0.0585179  0.0585179  0.0585179
 0.0585179  0.824446   0.0585179  0.0585179
 0.0585179  0.0585179  0.824446   0.0585179
 0.0585179  0.0585179  0.0585179  0.824446

```
"""
@inline function P!(Pmat::AbstractMatrix, obj::SM, t::Float64)
    Pmat[:] = exp(Q(obj) * t)
    return Pmat
end

"""
    BinaryTraitSubstitutionModel(α, β [, label])

Model for binary traits, that is, with 2 states. Default labels are "0" and "1".
α is the rate of transition from "0" to "1", and β from "1" to "0".
"""
struct BinaryTraitSubstitutionModel{T} <: TraitSubstitutionModel{T} # fixit: back to mutable struct?
    rate::Vector{Float64}
    label::Vector{T} # most often: T = String, but could be BioSymbols.DNA
    eigeninfo::Vector{Float64}
    function BinaryTraitSubstitutionModel{T}(rate::Vector{Float64}, label::Vector{T}, eigeninfo::Vector{Float64}) where T
        #Warning: this constructor should not be used directly. Use it with the constructors below.
        @assert length(rate) == 2 "binary state: need 2 rates"
        rate[1] >= 0. || error("parameter α must be non-negative")
        rate[2] >= 0. || error("parameter β must be non-negative")
        @assert length(label) == 2 "need 2 labels exactly"
        new(rate, label, eigeninfo)
    end
end
const BTSM = BinaryTraitSubstitutionModel{T} where T
function BinaryTraitSubstitutionModel(r::AbstractVector, label::AbstractVector)
    obj = BinaryTraitSubstitutionModel{eltype(label)}(r::AbstractVector, label::AbstractVector, zeros(3)) # Vector{Float64}(undef,3) for julia v1.0
    seteigeninfo!(obj)
    return obj
end
BinaryTraitSubstitutionModel(α::Float64, β::Float64, label::AbstractVector) = BinaryTraitSubstitutionModel([α,β], label)
BinaryTraitSubstitutionModel(α::Float64, β::Float64) = BinaryTraitSubstitutionModel(α, β, ["0", "1"])

"""
for a [`BinaryTraitSubstitutionModel`]: store eigenvalue (q_01+q_10) and stationary distribution
"""
function seteigeninfo!(obj::BTSM)
    ab = obj.rate[1] + obj.rate[2] #eigenvalue = -(a+b)
    ab > 0. || error("α+β must be positive")
    p0 = obj.rate[2]/ab # asymptotic frequency of state "0"
    p1 = obj.rate[1]/ab # asymptotic frequency of state "1"
    obj.eigeninfo[1] = ab
    obj.eigeninfo[2] = p0
    obj.eigeninfo[3] = p1
end

"""
for a `BinaryTraitSubstitutionModel`, this is 2:

```jldoctest
julia> m1 = BinaryTraitSubstitutionModel([1.0,2.0], ["low","high"])
Binary Trait Substitution Model:
rate low→high α=1.0
rate high→low β=2.0

julia> nstates(m1)
2
```
"""
function nstates(::BTSM)
    return 2::Int
end

nparams(::BTSM) = 2::Int

"""
For a BinaryTraitSubstitutionModel, the rate matrix Q is of the form:

    -α  α
     β -β
"""
@inline function Q(obj::BTSM)
    return Bmatrix(-obj.rate[1], obj.rate[2], obj.rate[1], -obj.rate[2])
end

function Base.show(io::IO, obj::BTSM)
    str = "Binary Trait Substitution Model:\n"
    str *= "rate $(obj.label[1])→$(obj.label[2]) α=$(round(obj.rate[1], digits=5))\n"
    str *= "rate $(obj.label[2])→$(obj.label[1]) β=$(round(obj.rate[2], digits=5))\n"
    print(io, str)
end

function P!(Pmat::AbstractMatrix, obj::BTSM, t::Float64)
    e1 = exp(-obj.eigeninfo[1]*t)
    a0= obj.eigeninfo[2]*e1
    a1= obj.eigeninfo[3]*e1
    Pmat[1,1] = obj.eigeninfo[2]+a1
    Pmat[2,1] = obj.eigeninfo[2]-a0
    Pmat[1,2] = obj.eigeninfo[3]-a1
    Pmat[2,2] = obj.eigeninfo[3]+a0
    return Pmat
end

"""
    TwoBinaryTraitSubstitutionModel(rate [, label])

[`TraitSubstitutionModel`](@ref) for two binary traits, possibly correlated.
Default labels are "x0", "x1" for trait 1, and "y0", "y1" for trait 2.
If provided, `label` should be a vector of size 4, listing labels for
trait 1 first then labels for trait 2.
`rate` should be a vector of substitution rates of size 8.
rate[1],...,rate[4] describe rates of changes in trait 1.
rate[5],...,rate[8] describe rates of changes in trait 2.
In the transition matrix, trait combinations are listed in the following order:
x0-y0, x0-y1, x1-y0, x1-y1.

# example

```julia
model = TwoBinaryTraitSubstitutionModel([2.0,1.2,1.1,2.2,1.0,3.1,2.0,1.1],
        ["carnivory", "noncarnivory", "wet", "dry"]);
model
using PhyloPlots
plot(model) # to visualize states and rates
```
"""
struct TwoBinaryTraitSubstitutionModel <: TraitSubstitutionModel{String}
    rate::Vector{Float64}
    label::Vector{String}
    function TwoBinaryTraitSubstitutionModel(α, label)
        all( x -> x >= 0., α) || error("rates must be non-negative")
        @assert length(α)==8 "need 8 rates"
        @assert length(label)==4 "need 4 labels for all combinations of 2 binary traits"
        new(α, [string(label[1], "-", label[3]), # warning: original type of 'label' lost here
                string(label[1], "-", label[4]),
                string(label[2], "-", label[3]),
                string(label[2], "-", label[4])])
    end
end
const TBTSM = TwoBinaryTraitSubstitutionModel
TwoBinaryTraitSubstitutionModel(α::Vector{Float64}) = TwoBinaryTraitSubstitutionModel(α, ["x0", "x1", "y0", "y1"])

nstates(::TBTSM) = 4::Int
nparams(::TBTSM) = 8::Int

function Q(obj::TBTSM)
    M = fill(0.0,(4,4))
    a = obj.rate
    M[1,3] = a[1]
    M[3,1] = a[2]
    M[2,4] = a[3]
    M[4,2] = a[4]
    M[1,2] = a[5]
    M[2,1] = a[6]
    M[3,4] = a[7]
    M[4,3] = a[8]
    M[1,1] = -M[1,2] - M[1,3]
    M[2,2] = -M[2,1] - M[2,4]
    M[3,3] = -M[3,4] - M[3,1]
    M[4,4] = -M[4,3] - M[4,2]
    return M
end

function Base.show(io::IO, obj::TBTSM)
    print(io, "Substitution model for 2 binary traits, with rate matrix:\n")
    showQ(io, obj)
end

"""
    EqualRatesSubstitutionModel(numberStates, α, labels)

[`TraitSubstitutionModel`](@ref) for traits with any number of states
and equal substitution rates α between all states.
Default labels are "1","2",...

# example

```jldoctest
julia> m1 = EqualRatesSubstitutionModel(2, [1.0], ["low","high"])
Equal Rates Substitution Model with k=2,
all rates equal to α=1.0.
rate matrix Q:
             low    high
     low       *  1.0000
    high  1.0000       *
```
"""
struct EqualRatesSubstitutionModel{T} <: TraitSubstitutionModel{T}
    k::Int
    rate::Vector{Float64}
    label::Vector{T}
    eigeninfo::Vector{Float64}
    function EqualRatesSubstitutionModel{T}(k::Int, rate::Vector{Float64},
        label::Vector{T}, eigeninfo::Vector{Float64}) where T
        k >= 2 || error("parameter k must be greater than or equal to 2")
        @assert length(rate)==1 "rate must be a vector of length 1"
        rate[1] > 0 || error("parameter α (rate) must be positive")
        @assert length(label)==k "label vector of incorrect length"
        new(k, rate, label, eigeninfo)
    end
end
const ERSM = EqualRatesSubstitutionModel{T} where T
function EqualRatesSubstitutionModel(k::Int, rate::Vector{Float64}, label::AbstractVector)
    obj = EqualRatesSubstitutionModel{eltype(label)}(k::Int, rate::Vector{Float64}, label::AbstractVector, zeros(1)) # Vector{Float64}(undef,3) for julia v1.0
    seteigeninfo!(obj)
    return obj
end
EqualRatesSubstitutionModel(k::Int, α::Float64, label::AbstractVector) = EqualRatesSubstitutionModel(k,[α],label)
EqualRatesSubstitutionModel(k::Int, α::Float64) = EqualRatesSubstitutionModel(k, [α], string.(1:k))

"""
for a [`EqualRatesSubstitutionModel`]: store lambda = k/(k-1), where k is the number of states
"""
function seteigeninfo!(obj::ERSM)
    q = (obj.k/(obj.k-1.0))
    obj.eigeninfo[1] = q
end

function nstates(obj::ERSM)
    return obj.k
end
nparams(::ERSM) = 1::Int

function Base.show(io::IO, obj::ERSM)
    str = "Equal Rates Substitution Model with k=$(obj.k),\n"
    str *= "all rates equal to α=$(round(obj.rate[1], digits=5)).\n"
    str *= "rate matrix Q:\n"
    print(io, str)
    showQ(io, obj)
end
function Q(obj::ERSM)
    #this might be wrong, doesnt match ERSM proof
    α = obj.rate[1]
    M = fill(α, (obj.k,obj.k))
    d = -(obj.k-1) * α
    for i in 1:obj.k
        M[i,i] = d
    end
    return M
end

"""
    randomTrait(model, t, start)
    randomTrait!(end, model, t, start)

Simulate traits along one edge of length t.
`start` must be a vector of integers, each representing the starting value of one trait.
The bang version (ending with !) uses the vector `end` to store the simulated values.

# Examples
```jldoctest
julia> m1 = BinaryTraitSubstitutionModel(1.0, 2.0)
Binary Trait Substitution Model:
rate 0→1 α=1.0
rate 1→0 β=2.0

julia> using Random; Random.seed!(12345);

julia> randomTrait(m1, 0.2, [1,2,1,2,2])
5-element Array{Int64,1}:
 1
 2
 1
 1
 2
```
"""
function randomTrait(obj::TSM, t::Float64, start::AbstractVector{Int})
    res = Vector{Int}(undef, length(start))
    randomTrait!(res, obj, t, start)
end

@doc (@doc randomTrait) randomTrait!
function randomTrait!(endTrait::AbstractVector{Int}, obj::TSM, t::Float64, start::AbstractVector{Int})
    Pt = P(obj, t)
    k = size(Pt, 1) # number of states
    w = [aweights(Pt[i,:]) for i in 1:k]
    for i in 1:length(start)
        endTrait[i] =sample(1:k, w[start[i]])
    end
    return endTrait
end

"""
    randomTrait(model, net; ntraits=1, keepInternal=true, checkPreorder=true)

Simulate evolution of discrete traits on a rooted evolutionary network based on
the supplied evolutionary model. Trait sampling is uniform at the root.

optional arguments:

- `ntraits`: number of traits to be simulated (default: 1 trait).
- `keepInternal`: if true, export character states at all nodes, including
  internal nodes. if false, export character states at tips only.

output:

- matrix of character states with one row per trait, one column per node;
  these states are *indices* in `model.label`, not the trait labels themselves.
- vector of node labels (for tips) or node numbers (for internal nodes)
  in the same order as columns in the character state matrix

# examples

```jldoctest
julia> m1 = BinaryTraitSubstitutionModel(1.0, 2.0, ["low","high"]);

julia> net = readTopology("(((A:4.0,(B:1.0)#H1:1.1::0.9):0.5,(C:0.6,#H1:1.0::0.1):1.0):3.0,D:5.0);");

julia> using Random; Random.seed!(1234);

julia> trait, lab = randomTrait(m1, net)
([1 2 … 1 1], ["-2", "D", "-3", "-6", "C", "-4", "H1", "B", "A"])

julia> trait
1×9 Array{Int64,2}:
 1  2  1  1  2  2  1  1  1

julia> lab
9-element Array{String,1}:
 "-2"
 "D"
 "-3"
 "-6"
 "C"
 "-4"
 "H1"
 "B"
 "A"
```
"""
function randomTrait(obj::TSM, net::HybridNetwork;
    ntraits=1::Int, keepInternal=true::Bool, checkPreorder=true::Bool)
    net.isRooted || error("net needs to be rooted for preorder recursion")
    if(checkPreorder)
        preorder!(net)
    end
    nnodes = net.numNodes
    M = Matrix{Int}(undef, ntraits, nnodes) # M[i,j]= trait i for node j
    randomTrait!(M,obj,net)
    if !keepInternal
        M = getTipSubmatrix(M, net, indexation=:cols) # subset columns only. rows=traits
        nodeLabels = [n.name for n in net.nodes_changed if n.leaf]
    else
        nodeLabels = [n.name == "" ? string(n.number) : n.name for n in net.nodes_changed]
    end
    return M, nodeLabels
end

function randomTrait!(M::Matrix, obj::TSM, net::HybridNetwork)
    recursionPreOrder!(net.nodes_changed, M, # updates M in place
            updateRootRandomTrait!,
            updateTreeRandomTrait!,
            updateHybridRandomTrait!,
            obj)
end

function updateRootRandomTrait!(V::AbstractArray, i::Int, obj)
    sample!(1:nstates(obj), view(V, :, i)) # uniform at the root
    return
end

function updateTreeRandomTrait!(V::Matrix,
    i::Int,parentIndex::Int,edge::Edge,
    obj)
    randomTrait!(view(V, :, i), obj, edge.length, view(V, :, parentIndex))
end

function updateHybridRandomTrait!(V::Matrix,
        i::Int, parentIndex1::Int, parentIndex2::Int,
        edge1::Edge, edge2::Edge, obj)
    randomTrait!(view(V, :, i), obj, edge1.length, view(V, :, parentIndex1))
    tmp = randomTrait(obj, edge2.length, view(V, :, parentIndex2))
    for j in 1:size(V,1) # loop over traits
        if V[j,i] == tmp[j] # both parents of the hybrid node have the same trait
            continue # skip the rest: go to next trait
        end
        if rand() > edge1.gamma
            V[j,i] = tmp[j] # switch to inherit trait of parent 2
        end
    end
end

"""
    JC69(rate, relative)

Jukes Cantor (1969) nucleic acid substitution model, which has a single rate parameter.
`rate` corresponds to the absolute diagonal elements, that is, the rate of change
(to any of the other 2 states). Individual rates are `rate`/3.
If `relative` is true (default), the transition matrix [`Q`](@ref) is normalized
to an average of 1 transition per unit of time: in which case `rate` is set to 1.0.

# examples

```jldoctest
julia> m1 = JC69([0.25], false)
Jukes and Cantor 69 Substitution Model,
absolute rate version
off-diagonal rates equal to 0.25/3.
rate matrix Q:
               A       C       G       T
       A       *  0.0833  0.0833  0.0833
       C  0.0833       *  0.0833  0.0833
       G  0.0833  0.0833       *  0.0833
       T  0.0833  0.0833  0.0833       *

julia> nstates(m1)
4

julia> nparams(m1)
1

julia> m2 = JC69([0.5])
Jukes and Cantor 69 Substitution Model,
relative rate version
off-diagonal rates equal to 1/3
rate matrix Q:
               A       C       G       T
       A       *  0.3333  0.3333  0.3333
       C  0.3333       *  0.3333  0.3333
       G  0.3333  0.3333       *  0.3333
       T  0.3333  0.3333  0.3333       *

julia> nparams(m2)
0
```
"""
struct JC69 <: NucleicAcidSubstitutionModel
    rate::Vector{Float64}
    relative::Bool
    eigeninfo::Vector{Float64} #TODO change to MVector, see if its faster
    function JC69(rate::Vector{Float64}, relative::Bool, eigeninfo::Vector{Float64})
        #Warning: this constructor should not be used directly. Use the constructor below
        #   which will call this.
        0 <= length(rate) <= 1 || error("rate not a valid length for a JC69 model")
        all(rate .> 0.0) || error("All elements of rate must be positive for a JC69 model")
        new(rate, relative, eigeninfo)
    end
end
function JC69(rate::AbstractVector, relative=true::Bool)
    obj = JC69(rate, relative, zeros(1))
    seteigeninfo!(obj)
    return obj
end
JC69(rate::Float64, relative=true::Bool) = JC69([rate], relative)

"""
for [`JC69`](@ref): store lambda = 4/3 (if relative) or rate * 4/3 (absolute).
"""
function seteigeninfo!(obj::JC69)
    if obj.relative
        lambda = (4.0/3.0) #eigen value of the Q matrix
    else
        lambda = (4.0/3.0)*obj.rate[1] #eigen value of the Q matrix
    end
    obj.eigeninfo[1] = lambda
end

function Base.show(io::IO, obj::JC69)
    str = "Jukes and Cantor 69 Substitution Model,\n"
    if obj.relative == true
        str *= "relative rate version\n"
        str *= "off-diagonal rates equal to 1/3\n"
    else
        str *= "absolute rate version\n"
        str *= "off-diagonal rates equal to $(round(obj.rate[1], digits=5))/3.\n"
    end
    str *= "rate matrix Q:\n"
    print(io, str)
    showQ(io, obj)
end

"""
    HKY85(rate, pi, relative)

A nucleic acid substitution model based on Hasegawa et al. 1985 substitution model.
`rate` should be a vector of 1 or 2 rates, and `pi` a vector of 4 probabilities summing to 1.

If `relative` is false, the 2 rates represent the transition rate and the transversion rate,
α and β. If `relative` is true (default), only the first rate is used and represents the
transition/transversion ratio: κ=α/β. The rate transition matrix Q is normalized to have
1 change / unit of time on average, i.e. the absolute version of Q is divided by
`2(piT*piC + piA*piG)α + 2(piY*piR)β`.

`nparams` returns 1 or 2.
In other words: the stationary distribution is not counted in the number of parameters
(and `fitdiscrete` does not optimize the pi values at the moment).

# examples

```jldoctest
julia> m1 = HKY85([.5], [0.20, 0.30, 0.30, 0.20])
HKY85 Substitution Model base frequencies: [0.2, 0.3, 0.3, 0.2]
relative rate version with transition/tranversion ratio kappa = 0.5,
 scaled so that there is one substitution per unit time
rate matrix Q:
               A       C       G       T
       A       *  0.4839  0.2419  0.3226
       C  0.3226       *  0.4839  0.1613
       G  0.1613  0.4839       *  0.3226
       T  0.3226  0.2419  0.4839       *

julia> nstates(m1)
4

julia> m2 = HKY85([0.5, 0.5], [0.20, 0.30, 0.30, 0.20], false)
HKY85 Substitution Model base frequencies: [0.2, 0.3, 0.3, 0.2]
absolute rate version with transition/transversion ratio kappa = a/b = 1.0
 with rates a = 0.5 and b = 0.5
rate matrix Q:
               A       C       G       T
       A       *  0.1500  0.1500  0.1000
       C  0.1000       *  0.1500  0.1000
       G  0.1000  0.1500       *  0.1000
       T  0.1000  0.1500  0.1500       *

```
"""
struct HKY85 <: NucleicAcidSubstitutionModel
    rate::Vector{Float64}
    pi::Vector{Float64}
    relative::Bool
    eigeninfo::Vector{Float64}
    function HKY85(rate::Vector{Float64}, pi::Vector{Float64}, relative::Bool, eigeninfo::Vector{Float64})
        #Warning: this constructor should not be used directly. Use the constructor below,
        #   which will call this.
        all(rate .> 0.) || error("All elements of rate must be positive")
        1 <= length(rate) <= 2 || error("rate has invalid length for HKY85 model")
        if relative length(rate) == 1 || error("the relative version of HKY85 takes a single rate")
        else        length(rate) == 2 || error("the absolute version of HKY85 takes 2 rates")
        end
        length(pi) == 4 || error("pi must be of length 4")
        all(0. .< pi.< 1.) || error("All base proportions must be between 0 and 1")
        sum(pi) == 1. || error("Base proportions must sum to 1")
        new(rate, pi, relative, eigeninfo)
    end
end
function HKY85(rate::AbstractVector, pi::Vector{Float64}, relative=true::Bool)
    obj = HKY85(rate, pi, relative, zeros(5))
    seteigeninfo!(obj)
    return obj
end
HKY85(rate::Float64, pi::Vector{Float64}) = HKY85([rate], pi, true)

"""
for [`HKY85`](@ref): store piR, piY, the 2 non-zero eigenvalues and a scaling factor
"""
function seteigeninfo!(obj::HKY85)
    piA = obj.pi[1]; piC = obj.pi[2]; piG = obj.pi[3]; piT = obj.pi[4]
    piR = piA + piG
    piY = piT + piC
    obj.eigeninfo[4] = piR
    obj.eigeninfo[5] = piY

    a = obj.rate[1]
    if obj.relative
        lambda = 2*a*(piT*piC + piA*piG) + 2*(piY*piR) #for sub/year interpretation
        b = 1.0/lambda #scaling factor so that branch lengths can be interpreted as substitutions
        obj.eigeninfo[2] = -((piR * a*b) + (piY * b))
        obj.eigeninfo[3] = -((piY * a*b) + (piR * b))
    else
        b = obj.rate[2]
        obj.eigeninfo[2] = -((piR * a) + (piY * b))
        obj.eigeninfo[3] = -((piY * a) + (piR * b))
    end
    obj.eigeninfo[1] = -b
end

function Base.show(io::IO, obj::HKY85)
    str = "HKY85 Substitution Model base frequencies: $(obj.pi)\n"
    if obj.relative
        str *= "relative rate version with transition/tranversion ratio kappa = $(round(obj.rate[1], digits=5)),"
        str *= "\n scaled so that there is one substitution per unit time\n"
    else
        str *= "absolute rate version with transition/transversion ratio kappa = a/b = "
        str *= "$(round(obj.rate[1]/obj.rate[2], digits=5))"
        str *= "\n with rates a = $(round(obj.rate[1], digits=5)) and b = $(round(obj.rate[2], digits=5))\n"
    end
    str *= "rate matrix Q:\n"
    print(io, str)
    showQ(io, obj)
end

"""
for `JC69` model: 0 if relative, 1 if absolute
"""
function nparams(obj::JC69)
    return (obj.relative ? 0 : 1)
end

"""
for `HKY85` model: 1 if relative, 2 if absolute
"""
function nparams(obj::HKY85)
    return (obj.relative ? 1 : 2)
end

@inline function Q(obj::JC69)
    if obj.relative
        Q0 = 1.0/3.0
        #we multiply by 1/3 to make branch lengths interpretable as avg(#substitutions) (pi-weighted avg of diagonal is -1)
    else
        Q0 = (1.0/3.0)*obj.rate[1]
    end

    return Qmatrix(-3*Q0, Q0, Q0, Q0,
                    Q0, -3*Q0, Q0, Q0,
                    Q0, Q0, -3*Q0, Q0,
                    Q0, Q0, Q0, -3*Q0)
end

@inline function Q(obj::HKY85)
    piA = obj.pi[1]; piC = obj.pi[2]; piG = obj.pi[3]; piT = obj.pi[4]
    piR = piA + piG
    piY = piT + piC
    if obj.relative
        k = obj.rate[1]
        lambda = (2*k*(piT*piC + piA*piG) + 2*(piY*piR)) # for sub/time interpretation. see HKY docstring
        Q₁  = piA/lambda
        Q₂  = k * Q₁
        Q₃  = piC/lambda
        Q₄  = k * Q₃
        Q₆  = piG/lambda
        Q₅  = k * Q₆
        Q₇  = piT/lambda
        Q₈  = k * Q₇
        Q₉  = -(Q₃ + Q₅ + Q₇)
        Q₁₀ = -(Q₁ + Q₆ + Q₈)
        Q₁₁ = -(Q₂ + Q₃ + Q₇)
        Q₁₂ = -(Q₁ + Q₄ + Q₆)
    else #GT, AT, CG, and AC are less likely (transversions)
        a = obj.rate[1]; b = obj.rate[2]
        Q₁  = b * piA
        Q₂  = a * piA
        Q₃  = b * piC
        Q₄  = a * piC
        Q₅  = a * piG
        Q₆  = b * piG
        Q₇  = b * piT
        Q₈  = a * piT
        Q₉  = -(Q₃ + Q₅ + Q₇) #AA
        Q₁₀ = -(Q₁ + Q₆ + Q₈) #CC
        Q₁₁ = -(Q₂ + Q₃ + Q₇) #GG
        Q₁₂ = -(Q₁ + Q₄ + Q₆) #TT
    end
    return Qmatrix(Q₉,  Q₁,  Q₂,  Q₁, #ACGT
                    Q₃,  Q₁₀, Q₃,  Q₄,
                    Q₅,  Q₆,  Q₁₁, Q₆,
                    Q₇,  Q₈,  Q₇,  Q₁₂)
end

function P!(Pmat::AbstractMatrix, obj::JC69, t::Float64)
    P0 = 0.25 + 0.75 * exp(-t * obj.eigeninfo[1]) #lambda
    P1 = 0.25 - 0.25 * exp(-t * obj.eigeninfo[1])
    Pmat[:,:] .= P1 # puts P1 on the off-diagonals (in particular)
    for i in 1:4 Pmat[i,i]=P0; end # diagonal
    return Pmat
end

function P!(Pmat::AbstractMatrix, obj::HKY85, t::Float64)
    #returns ACGT
    piA = obj.pi[1]; piC = obj.pi[2]; piG = obj.pi[3]; piT = obj.pi[4]
    piR = obj.eigeninfo[4]
    piY = obj.eigeninfo[5]

    e2 = exp(obj.eigeninfo[1] * t)
    e3 = exp(obj.eigeninfo[2] * t)
    e4 = exp(obj.eigeninfo[3] * t)

    P_AA = piA + ((piA * piY)/piR) * e2 + (piG/piR) * e3
    P_CC = piC + ((piC * piR)/piY) * e2 + (piT/piY) * e4
    P_GG = piG + ((piG * piY)/piR) * e2 + (piA/piR) * e3
    P_TT = piT + ((piT * piR)/piY) * e2 + (piC/piY) * e4

    P_TA_CA = piA * (1 - e2) #double name because TA and CA are equal
    P_GA = piA + (piA*piY/piR)*e2 - (piA/piR)*e3 #P6

    P_AC_GC = piC * (1 - e2) #double name because AC and GC are equal #p7
    P_TC = piC + ((piC * piR)/piY) * e2 - (piC/piY) * e4

    P_AG = piG + ((piG * piY)/piR) * e2 - (piG/piR) * e3 #p9
    P_TG_CG = piG * (1 - e2) #double name because TG and CG are equal #p10

    P_AT_GT = piT * (1 - e2) #double name because AT and GT are equal #P11
    P_CT = piT + ((piT * piR)/piY) * e2 - (piT/piY) * e4

    Pmat[1,1] = P_AA
    Pmat[2,2] = P_CC
    Pmat[3,3] = P_GG
    Pmat[4,4] = P_TT

    Pmat[1,2] = P_TA_CA
    Pmat[1,3] = P_GA
    Pmat[1,4] = P_TA_CA

    Pmat[2,1] = P_AC_GC
    Pmat[2,3] = P_AC_GC
    Pmat[2,4] = P_TC

    Pmat[3,1] = P_AG
    Pmat[3,2] = P_TG_CG
    Pmat[3,4] = P_TG_CG

    Pmat[4,1] = P_AT_GT
    Pmat[4,2] = P_CT
    Pmat[4,3] = P_AT_GT
    return Pmat
end

"""
    RateVariationAcrossSites(α=1.0, ncat=4)

Model for variable substitution rates across sites (or across traits) using
the discrete Gamma model (Yang 1994, Journal of Molecular Evolution).
Turn any SM model to SM + Gamma.
Using this model increases the number of parameters by one.

Because the mean of the desired continuous Gamma distribution is 1, we use
shape α and scale θ=1/α (e.g. rate β=α).
The shape parameter is referred to as alpha here.
The Gamma distribution is discretized into `ncat` categories.
In each category, the category's rate multiplier is a normalized quantile of the gamma distribution.

```jldoctest
julia> rv = RateVariationAcrossSites()
Rate Variation Across Sites using Discretized Gamma Model
alpha: 1.0
categories for Gamma discretization: 4
ratemultiplier: [0.14578, 0.51313, 1.07083, 2.27025]

julia> PhyloNetworks.setalpha!(rv, 2.0)

julia> rv
Rate Variation Across Sites using Discretized Gamma Model
alpha: 2.0
categories for Gamma discretization: 4
ratemultiplier: [0.31907, 0.68336, 1.10898, 1.8886]

julia> RateVariationAcrossSites(2.0, 4)
Rate Variation Across Sites using Discretized Gamma Model
alpha: 2.0
categories for Gamma discretization: 4
ratemultiplier: [0.31907, 0.68336, 1.10898, 1.8886]
```
"""
mutable struct RateVariationAcrossSites
    alpha::Float64
    ncat::Int
    ratemultiplier::Array{Float64}
    function RateVariationAcrossSites(alpha = 1.0::Float64, ncat = 4::Int)
        @assert alpha >= 0.0 "alpha must be >= 0"
        @assert ncat > 0 "ncat must be 1 or greater"
        if ncat == 1
            ratemultiplier = [1.0]
        else
            cuts = Vector((0:(ncat-1))/ncat) + repeat([1/(2ncat)], ncat)
            ratemultiplier = quantile.(Distributions.Gamma(alpha, 1/alpha), cuts)
            ratemultiplier ./= mean(ratemultiplier)
        end
        new(alpha, ncat, ratemultiplier)
    end
end
const RVAS = RateVariationAcrossSites

"""
    setalpha!(obj, alpha)

Set the shape parameter `alpha` in a RateVariationAcrossSites model `obj`,
and update the rate multipliers accordingly.
"""
function setalpha!(obj::RateVariationAcrossSites, alpha::Float64)
    @assert alpha >= 0 "alpha must be a float >= 0"
    obj.alpha = alpha
    if obj.ncat > 1
        rv = obj.ratemultiplier
        cuts = Vector((0:(obj.ncat-1))/obj.ncat) + repeat([1/(2*obj.ncat)], obj.ncat)
        rv[:] = quantile.(Distributions.Gamma(alpha, 1/alpha), cuts)
        rv ./= mean(rv)
    end
    return nothing
end

function Base.show(io::IO, obj::RateVariationAcrossSites)
    str = "Rate Variation Across Sites using Discretized Gamma Model\n"
    str *= "alpha: $(round(obj.alpha, digits=5))\n"
    str *= "categories for Gamma discretization: $(obj.ncat)\n"
    str *= "ratemultiplier: $(round.(obj.ratemultiplier, digits=5))\n"
    print(io, str)
end

function nparams(obj::RateVariationAcrossSites)
    return (obj.ncat == 1 ? 0 : 1)
end

"""
    empiricalDNAfrequencies(DNAdata::AbstractDataFrame, DNAweights,
                            correction=true, useambiguous=true)

Estimate base frequencies in DNA data `DNAdata`, ordered ACGT.

- `DNAdata`: data frame. All columns are used. If the first column
  gives species names, find a way to ignore it before calculating empirical
  frequencies, e.g. `empiricalDNAfrequencies(view(DNAdata, :, 2:ncol(DNAdata)))`.
  Data type must be `BioSymbols.DNA` or `Char` or `String`.
  WARNING: this is checked on the first column only.
- `DNAweights`: vector of weights, to weigh each column in `DNAdata`.
- `correction`: if `true`, add 1 to each count and 4 to the denominator
  for a more stable estimator, similar to Bayes prior of 1/4 and
  the Agresti-Coull interval in binomial estimation.
- `useambiguous`: if `true`, ambiguous bases are used (except gaps and Ns).
  For example, `Y` adds 0.5 weight to `C` and 0.5 weight to `T`.
"""
function empiricalDNAfrequencies(dnaDat::AbstractDataFrame, dnaWeights::Vector,
    correctedestimate=true::Bool, useambiguous=true::Bool)

    # warning: checking first column and first row only
    eltypes(dnaDat)[1] == BioSymbols.DNA || eltypes(dnaDat)[1] == Char ||
      dnaDat[1,1] ∈ string.(BioSymbols.alphabet(DNA)) ||
        error("empiricalDNAfrequencies requires data of type String, Char, or BioSymbols.DNA")

    # initialize counts: keys same as BioSymbols.ACGT (which are ordered)
    prior = correctedestimate ? 1.0 : 0.0
    dnacounts = Dict(DNA_A=>prior, DNA_C=>prior, DNA_G=>prior, DNA_T=>prior)

    convert2dna = eltypes(dnaDat)[1] != BioSymbols.DNA
    for j in 1:ncol(dnaDat) # for each column
        col = dnaDat[!,j]
        wt = dnaWeights[j]
        for nuc in col      # for each row
            if convert2dna
                nuc = convert(DNA, nuc[1]) # if nuc is a string, nuc[1] = 1st character
            end
            if nuc ∈ BioSymbols.ACGT
                dnacounts[nuc] += wt
            elseif nuc == DNA_Gap || nuc == DNA_N || !useambiguous
                continue # to next row
            # else: ambiguity, uses BioSequences definitions
            elseif nuc == DNA_M # A or C
                dnacounts[DNA_A] += 0.5*wt
                dnacounts[DNA_C] += 0.5*wt
            elseif nuc == DNA_R # A or G
                dnacounts[DNA_A] += 0.5*wt
                dnacounts[DNA_G] += 0.5*wt
            elseif nuc == DNA_W # A or T/U
                dnacounts[DNA_A] += 0.5*wt
                dnacounts[DNA_T] += 0.5*wt
            elseif nuc == DNA_S # C or G
                dnacounts[DNA_C] += 0.5*wt
                dnacounts[DNA_G] += 0.5*wt
            elseif nuc == DNA_Y # C or T/U
                dnacounts[DNA_C] += 0.5*wt
                dnacounts[DNA_T] += 0.5*wt
            elseif nuc == DNA_K # G or T/U
                dnacounts[DNA_G] += 0.5*wt
                dnacounts[DNA_T] += 0.5*wt
            elseif nuc == DNA_V # A or C or G
                dnacounts[DNA_A] += wt/3
                dnacounts[DNA_C] += wt/3
                dnacounts[DNA_G] += wt/3
            elseif nuc == DNA_H # A or C or T
                dnacounts[DNA_A] += wt/3
                dnacounts[DNA_C] += wt/3
                dnacounts[DNA_T] += wt/3
            elseif nuc == DNA_D # A or G or T/U
                dnacounts[DNA_A] += wt/3
                dnacounts[DNA_G] += wt/3
                dnacounts[DNA_T] += wt/3
            elseif nuc == DNA_B # C or G or T/U
                dnacounts[DNA_C] += wt/3
                dnacounts[DNA_G] += wt/3
                dnacounts[DNA_T] += wt/3
            end
        end
    end
    totalweight = sum(values(dnacounts))
    res = [dnacounts[key]/totalweight for key in BioSymbols.ACGT] # to control the order
    all(0. .<= res .<= 1.) || error("""weird: empirical base frequency < 0 or > 1""")
    return res
end

"""
    stationary(substitutionmodel)

Stationary distribution of a Markov model
"""
stationary(mod::SM) = error("stationary not defined for $(typeof(mod)).")

function stationary(mod::JC69)
    return [0.25,0.25,0.25,0.25]
end

function stationary(mod::HKY85)
    return mod.pi
end

function stationary(mod::ERSM) #for ERSM, uniform = stationary
    return [1/mod.k for i in 1:mod.k]
end

function stationary(mod::BTSM)
    return [mod.eigeninfo[2], mod.eigeninfo[3]]
end


module MathOptInterfaceOSQP

include("modcaches.jl")
using .ModificationCaches

using SparseArrays
using MathOptInterface
using MathOptInterface.Utilities
using LinearAlgebra: rmul!

export Optimizer, OSQPSettings, OSQPModel

const MOI = MathOptInterface
const MOIU = MathOptInterface.Utilities
const CI = MOI.ConstraintIndex
const VI = MOI.VariableIndex

const SparseTriplets = Tuple{Vector{Int}, Vector{Int}, Vector{<:Any}}

const Affine = MOI.ScalarAffineFunction{Float64}
const Quadratic = MOI.ScalarQuadraticFunction{Float64}
const VectorAffine = MOI.VectorAffineFunction{Float64}

const Interval = MOI.Interval{Float64}
const LessThan = MOI.LessThan{Float64}
const GreaterThan = MOI.GreaterThan{Float64}
const EqualTo = MOI.EqualTo{Float64}
const IntervalConvertible = Union{Interval, LessThan, GreaterThan, EqualTo}

const Zeros = MOI.Zeros
const Nonnegatives = MOI.Nonnegatives
const Nonpositives = MOI.Nonpositives
const SupportedVectorSets = Union{Zeros, Nonnegatives, Nonpositives}

import OSQP

lower(::Zeros, i::Int) = 0.0
lower(::Nonnegatives, i::Int) = 0.0
lower(::Nonpositives, i::Int) = -Inf
upper(::Zeros, i::Int) = 0.0
upper(::Nonnegatives, i::Int) = Inf
upper(::Nonpositives, i::Int) = 0.0

# TODO: just use ∈ on 0.7 (allocates on 0.6):
function _contains(haystack, needle)
    for x in haystack
        x == needle && return true
    end
    false
end

mutable struct Optimizer <: MOI.AbstractOptimizer
    inner::OSQP.Model
    hasresults::Bool
    results::OSQP.Results
    is_empty::Bool
    silent::Bool
    settings::Dict{Symbol, Any} # need to store these, because they should be preserved if empty! is called
    sense::MOI.OptimizationSense
    objconstant::Float64
    constrconstant::Vector{Float64}
    modcache::ProblemModificationCache{Float64}
    warmstartcache::WarmStartCache{Float64}
    rowranges::Dict{Int, UnitRange{Int}}

    function Optimizer(; kwargs...)
        inner = OSQP.Model()
        hasresults = false
        results = OSQP.Results()
        is_empty = true
        sense = MOI.MIN_SENSE
        objconstant = 0.
        constrconstant = Float64[]
        modcache = ProblemModificationCache{Float64}()
        warmstartcache = WarmStartCache{Float64}()
        rowranges = Dict{Int, UnitRange{Int}}()
        optimizer = new(inner, hasresults, results, is_empty, false,
                        Dict{Symbol, Any}(:verbose => true), sense, objconstant,
                        constrconstant, modcache, warmstartcache, rowranges)
        for (key, value) in kwargs
            MOI.set(optimizer, MOI.RawParameter(key), value)
        end
        return optimizer
    end
end

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

MOI.supports(::Optimizer, ::MOI.Silent) = true
function MOI.set(optimizer::Optimizer, ::MOI.Silent, value::Bool)
    optimizer.silent = value
    if !MOI.is_empty(optimizer)
        if optimizer.silent
            OSQP.update_settings!(optimizer.inner; :verbose => false)
        else
            OSQP.update_settings!(optimizer.inner; :verbose => optimizer.settings[:verbose])
        end
    end
end
MOI.get(optimizer::Optimizer, ::MOI.Silent) = optimizer.silent

MOI.supports(::Optimizer, ::MOI.TimeLimitSec) = true
function MOI.set(model::Optimizer, ::MOI.TimeLimitSec, limit::Real)
    MOI.set(model, OSQPSettings.TimeLimit(), limit)
    return
end

function MOI.set(model::Optimizer, attr::MOI.TimeLimitSec, ::Nothing)
    delete!(model.settings, :time_limit)
    OSQP.update_settings!(model.inner, time_limit=0.0)
    return
end

function MOI.get(model::Optimizer, ::MOI.TimeLimitSec)
    return get(model.settings, :time_limit, nothing)
end

hasresults(optimizer::Optimizer) = optimizer.hasresults

function MOI.empty!(optimizer::Optimizer)
    optimizer.inner = OSQP.Model()
    optimizer.hasresults = false
    optimizer.results = OSQP.Results()
    optimizer.is_empty = true
    optimizer.sense = MOI.MIN_SENSE # model parameter, so needs to be reset
    optimizer.objconstant = 0.
    optimizer.constrconstant = Float64[]
    optimizer.modcache = ProblemModificationCache{Float64}()
    optimizer.warmstartcache = WarmStartCache{Float64}()
    empty!(optimizer.rowranges)
    optimizer
end

MOI.is_empty(optimizer::Optimizer) = optimizer.is_empty

function MOI.copy_to(dest::Optimizer, src::MOI.ModelLike; copy_names=false)
    copy_names && error("Copying names is not supported.")
    MOI.empty!(dest)
    idxmap = MOIU.IndexMap(dest, src)
    assign_constraint_row_ranges!(dest.rowranges, idxmap, src)
    dest.sense, P, q, dest.objconstant = processobjective(src, idxmap)
    A, l, u, dest.constrconstant = processconstraints(src, idxmap, dest.rowranges)
    settings = copy(dest.settings)
    if dest.silent
        settings[:verbose] = false
    end
    OSQP.setup!(dest.inner; P = P, q = q, A = A, l = l, u = u, settings...)
    dest.modcache = ProblemModificationCache(P, q, A, l, u)
    dest.warmstartcache = WarmStartCache{Float64}(size(A, 2), size(A, 1))
    processprimalstart!(dest.warmstartcache.x, src, idxmap)
    processdualstart!(dest.warmstartcache.y, src, idxmap, dest.rowranges)
    dest.is_empty = false
    idxmap
end

"""
Set up index map from `src` variables and constraints to `dest` variables and constraints.
"""
function MOIU.IndexMap(dest::Optimizer, src::MOI.ModelLike)
    idxmap = MOIU.IndexMap()
    vis_src = MOI.get(src, MOI.ListOfVariableIndices())
    for i in eachindex(vis_src)
        idxmap[vis_src[i]] = VI(i)
    end
    i = 0
    for (F, S) in MOI.get(src, MOI.ListOfConstraints())
        MOI.supports_constraint(dest, F, S) || throw(MOI.UnsupportedConstraint{F, S}())
        cis_src = MOI.get(src, MOI.ListOfConstraintIndices{F, S}())
        for ci in cis_src
            i += 1
            idxmap[ci] = CI{F, S}(i)
        end
    end
    idxmap
end

function assign_constraint_row_ranges!(rowranges::Dict{Int, UnitRange{Int}}, idxmap::MOIU.IndexMap, src::MOI.ModelLike)
    startrow = 1
    for (F, S) in MOI.get(src, MOI.ListOfConstraints())
        cis_src = MOI.get(src, MOI.ListOfConstraintIndices{F, S}())
        for ci_src in cis_src
            set = MOI.get(src, MOI.ConstraintSet(), ci_src)
            ci_dest = idxmap[ci_src]
            endrow = startrow + MOI.dimension(set) - 1
            rowranges[ci_dest.value] = startrow : endrow
            startrow = endrow + 1
        end
    end
end

function constraint_rows(rowranges::Dict{Int, UnitRange{Int}}, ci::CI{<:Any, <:MOI.AbstractScalarSet})
    rowrange = rowranges[ci.value]
    length(rowrange) == 1 || error()
    first(rowrange)
end
constraint_rows(rowranges::Dict{Int, UnitRange{Int}}, ci::CI{<:Any, <:MOI.AbstractVectorSet}) = rowranges[ci.value]
constraint_rows(optimizer::Optimizer, ci::CI) = constraint_rows(optimizer.rowranges, ci)

"""
Return objective sense, as well as matrix `P`, vector `q`, and scalar `c` such that objective function is `1/2 x' P x + q' x + c`.
"""
function processobjective(src::MOI.ModelLike, idxmap)
    sense = MOI.get(src, MOI.ObjectiveSense())
    n = MOI.get(src, MOI.NumberOfVariables())
    q = zeros(n)
    if sense != MOI.FEASIBILITY_SENSE
        function_type = MOI.get(src, MOI.ObjectiveFunctionType())
        if function_type == MOI.ScalarAffineFunction{Float64}
            faffine = MOI.get(src, MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}())
            P = spzeros(n, n)
            processlinearterms!(q, faffine.terms, idxmap)
            c = faffine.constant
        elseif function_type == MOI.ScalarQuadraticFunction{Float64}
            fquadratic = MOI.get(src, MOI.ObjectiveFunction{MOI.ScalarQuadraticFunction{Float64}}())
            I = [Int(idxmap[term.variable_index_1].value) for term in fquadratic.quadratic_terms]
            J = [Int(idxmap[term.variable_index_2].value) for term in fquadratic.quadratic_terms]
            V = [term.coefficient for term in fquadratic.quadratic_terms]
            symmetrize!(I, J, V)
            P = sparse(I, J, V, n, n)
            processlinearterms!(q, fquadratic.affine_terms, idxmap)
            c = fquadratic.constant
        else
            throw(MOI.UnsupportedAttribute(MOI.ObjectiveFunction{function_type}()))
        end
        sense == MOI.MAX_SENSE && (rmul!(P, -1); rmul!(q, -1); c = -c)
    else
        P = spzeros(n, n)
        q = zeros(n)
        c = 0.
    end
    sense, P, q, c
end

function processlinearterms!(q, terms::Vector{<:MOI.ScalarAffineTerm}, idxmapfun::Function = identity)
    # This is currently needed to avoid depwarns. TODO: make this nice again:
    if q isa VectorModificationCache
        q[:] = 0
    else
        q .= 0
    end
    for term in terms
        var = term.variable_index
        coeff = term.coefficient
        q[idxmapfun(var).value] += coeff
    end
end

function processlinearterms!(q, terms::Vector{<:MOI.ScalarAffineTerm}, idxmap::MOIU.IndexMap)
    processlinearterms!(q, terms, var -> idxmap[var])
end

function symmetrize!(I::Vector{Int}, J::Vector{Int}, V::Vector)
    n = length(V)
    (length(I) == length(J) == n) || error()
    for i = 1 : n
        if I[i] != J[i]
            push!(I, J[i])
            push!(J, I[i])
            push!(V, V[i])
        end
    end
end

function processconstraints(src::MOI.ModelLike, idxmap, rowranges::Dict{Int, UnitRange{Int}})
    m = mapreduce(length, +, values(rowranges), init=0)
    l = Vector{Float64}(undef, m)
    u = Vector{Float64}(undef, m)
    constant = Vector{Float64}(undef, m)
    bounds = (l, u)
    I = Int[]
    J = Int[]
    V = Float64[]
    for (F, S) in MOI.get(src, MOI.ListOfConstraints())
        processconstraints!((I, J, V), bounds, constant, src, idxmap, rowranges, F, S)
    end
    l .-= constant
    u .-= constant
    n = MOI.get(src, MOI.NumberOfVariables())
    A = sparse(I, J, V, m, n)
    (A, l, u, constant)
end

function processconstraints!(triplets::SparseTriplets, bounds::Tuple{<:Vector, <:Vector}, constant::Vector{Float64},
        src::MOI.ModelLike, idxmap, rowranges::Dict{Int, UnitRange{Int}},
        F::Type{<:MOI.AbstractFunction}, S::Type{<:MOI.AbstractSet})
    cis_src = MOI.get(src, MOI.ListOfConstraintIndices{F, S}())
    for ci in cis_src
        s = MOI.get(src, MOI.ConstraintSet(), ci)
        f = MOI.get(src, MOI.ConstraintFunction(), ci)
        rows = constraint_rows(rowranges, idxmap[ci])
        processconstant!(constant, rows, f)
        processlinearpart!(triplets, f, rows, idxmap)
        processconstraintset!(bounds, rows, s)
    end
    nothing
end

function processconstant!(c::Vector{Float64}, row::Int, f::Affine)
    c[row] = MOI.constant(f, Float64)
    nothing
end

function processconstant!(c::Vector{Float64}, rows::UnitRange{Int}, f::VectorAffine)
    for (i, row) in enumerate(rows)
        c[row] = f.constants[i]
    end
end

function processlinearpart!(triplets::SparseTriplets, f::MOI.ScalarAffineFunction, row::Int, idxmap)
    (I, J, V) = triplets
    for term in f.terms
        var = term.variable_index
        coeff = term.coefficient
        col = idxmap[var].value
        push!(I, row)
        push!(J, col)
        push!(V, coeff)
    end
end

function processlinearpart!(triplets::SparseTriplets, f::MOI.VectorAffineFunction, rows::UnitRange{Int}, idxmap)
    (I, J, V) = triplets
    for term in f.terms
        row = rows[term.output_index]
        var = term.scalar_term.variable_index
        coeff = term.scalar_term.coefficient
        col = idxmap[var].value
        push!(I, row)
        push!(J, col)
        push!(V, coeff)
    end
end

function processconstraintset!(bounds::Tuple{<:Vector, <:Vector}, row::Int, s::IntervalConvertible)
    processconstraintset!(bounds, row, MOI.Interval(s))
end

function processconstraintset!(bounds::Tuple{<:Vector, <:Vector}, row::Int, interval::Interval)
    l, u = bounds
    l[row] = interval.lower
    u[row] = interval.upper
    nothing
end

function processconstraintset!(bounds::Tuple{<:Vector, <:Vector}, rows::UnitRange{Int}, s::S) where {S<:SupportedVectorSets}
    l, u = bounds
    for (i, row) in enumerate(rows)
        l[row] = lower(s, i)
        u[row] = upper(s, i)
    end
end

function processprimalstart!(x, src::MOI.ModelLike, idxmap)
    has_primal_start = false
    for attr in MOI.get(src, MOI.ListOfVariableAttributesSet())
        if attr isa MOI.VariablePrimalStart
            has_primal_start = true
        end
    end
    if has_primal_start
        vis_src = MOI.get(src, MOI.ListOfVariableIndices())
        for vi in vis_src
            value = MOI.get(src, MOI.VariablePrimalStart(), vi)
            if value != nothing
                x[idxmap[vi].value] = value
            end
        end
    end
end

function processdualstart!(y, src::MOI.ModelLike, idxmap, rowranges::Dict{Int, UnitRange{Int}})
    for (F, S) in MOI.get(src, MOI.ListOfConstraints())
        has_dual_start = false
        for attr in MOI.get(src, MOI.ListOfConstraintAttributesSet{F, S}())
            if attr isa MOI.ConstraintDualStart
                has_dual_start = true
            end
        end
        if has_dual_start
            cis_src = MOI.get(src, MOI.ListOfConstraintIndices{F, S}())
            for ci in cis_src
                rows = constraint_rows(rowranges, idxmap[ci])
                dual = MOI.get(src, MOI.ConstraintDualStart(), ci)
                if dual != nothing
                    for (i, row) in enumerate(rows)
                        y[row] = -dual[i] # opposite dual convention
                    end
                end
            end
        end
    end
end

## Standard optimizer attributes:
MOI.get(optimizer::Optimizer, ::MOI.ObjectiveSense) = optimizer.sense
function MOI.get(optimizer::Optimizer, a::MOI.NumberOfVariables)
    OSQP.dimensions(optimizer.inner)[1]
end

function MOI.get(optimizer::Optimizer, a::MOI.ListOfVariableIndices)
    [VI(i) for i = 1 : MOI.get(optimizer, MOI.NumberOfVariables())] # TODO: support for UnitRange would be nice
end

## Solver-specific optimizer attributes:
module OSQPSettings

using MathOptInterface
using OSQP

export OSQPAttribute, isupdatable

abstract type OSQPAttribute <: MathOptInterface.AbstractOptimizerAttribute end

# TODO: just use ∈ on 0.7 (allocates on 0.6):
function _contains(haystack, needle)
    for x in haystack
        x == needle && return true
    end
    false
end

for setting in fieldnames(OSQP.Settings)
    Attribute = Symbol(mapreduce(uppercasefirst, *, split(String(setting), '_'))) # to camelcase
    @eval begin
        export $Attribute
        struct $Attribute <: OSQPAttribute end
        Base.Symbol(::$Attribute) = $(QuoteNode(setting))
        isupdatable(::$Attribute) = $(_contains(OSQP.UPDATABLE_SETTINGS, setting))
    end
end
end # module

using .OSQPSettings

_symbol(param::MOI.RawParameter) = Symbol(param.name)
_symbol(a::OSQPAttribute) = Symbol(a)
OSQPSettings.isupdatable(param::MOI.RawParameter) = _contains(OSQP.UPDATABLE_SETTINGS, _symbol(param))
function MOI.set(optimizer::Optimizer, a::Union{OSQPAttribute, MOI.RawParameter}, value)
    (isupdatable(a) || MOI.is_empty(optimizer)) || throw(MOI.SetAttributeNotAllowed(a))
    setting = _symbol(a)
    optimizer.settings[setting] = value
    if !MOI.is_empty(optimizer)
        OSQP.update_settings!(optimizer.inner; setting => value)
    end
end

function MOI.get(optimizer::Optimizer, a::Union{OSQPAttribute, MOI.RawParameter})
    return optimizer.settings[_symbol(a)]
end

## Optimizer methods:
function MOI.optimize!(optimizer::Optimizer)
    processupdates!(optimizer.inner, optimizer.modcache)
    processupdates!(optimizer.inner, optimizer.warmstartcache)
    OSQP.solve!(optimizer.inner, optimizer.results)
    optimizer.hasresults = true
    # Copy previous solution into warm start cache without setting the dirty bit:
    copyto!(optimizer.warmstartcache.x.data, optimizer.results.x)
    copyto!(optimizer.warmstartcache.y.data, optimizer.results.y)
    nothing
end

## Optimizer attributes:
MOI.get(optimizer::Optimizer, ::MOI.RawSolver) = optimizer.inner
MOI.get(optimizer::Optimizer, ::MOI.ResultCount) = optimizer.hasresults ? 1 : 0

MOI.supports(::Optimizer, ::MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}) = true
MOI.supports(::Optimizer, ::MOI.ObjectiveFunction{Quadratic}) = true
MOI.supports(::Optimizer, ::MOI.ObjectiveSense) = true

function MOI.set(optimizer::Optimizer, a::MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}, obj::MOI.ScalarAffineFunction{Float64})
    MOI.is_empty(optimizer) && throw(MOI.SetAttributeNotAllowed(a))
    optimizer.modcache.P[:] = 0
    processlinearterms!(optimizer.modcache.q, obj.terms)
    optimizer.objconstant = MOI.constant(obj)
    nothing
end

function MOI.set(optimizer::Optimizer, a::MOI.ObjectiveFunction{Quadratic}, obj::Quadratic)
    MOI.is_empty(optimizer) && throw(MOI.SetAttributeNotAllowed(a))
    cache = optimizer.modcache
    cache.P[:] = 0
    for term in obj.quadratic_terms
        row = term.variable_index_1.value
        col = term.variable_index_2.value
        coeff = term.coefficient
        row > col && ((row, col) = (col, row)) # upper triangle only
        cache.P[row, col] += coeff
    end
    processlinearterms!(optimizer.modcache.q, obj.affine_terms)
    optimizer.objconstant = MOI.constant(obj)
    nothing
end

function MOI.get(optimizer::Optimizer, a::MOI.ObjectiveValue)
    MOI.check_result_index_bounds(optimizer, a)
    rawobj = optimizer.results.info.obj_val + optimizer.objconstant
    ifelse(optimizer.sense == MOI.MAX_SENSE, -rawobj, rawobj)
end

error_not_solved() = error("Problem is unsolved.")
function check_has_results(optimizer::Optimizer)
    if !hasresults(optimizer)
        error_not_solved()
    end
end

# Since these aren't explicitly returned by OSQP, I feel like it would be better to have a fallback method compute these:
function MOI.get(optimizer::Optimizer, a::MOI.SolveTime)
    check_has_results(optimizer)
    return optimizer.results.info.run_time
end

function MOI.get(optimizer::Optimizer, ::MOI.RawStatusString)
    return string(optimizer.results.info.status)
end

function MOI.get(optimizer::Optimizer, ::MOI.TerminationStatus)
    hasresults(optimizer) || return MOI.OPTIMIZE_NOT_CALLED
    osqpstatus = optimizer.results.info.status
    if osqpstatus == :Unsolved
        return MOI.OPTIMIZE_NOT_CALLED
    elseif osqpstatus == :Interrupted
        return MOI.INTERRUPTED
    elseif osqpstatus == :Dual_infeasible
        return MOI.DUAL_INFEASIBLE
    elseif osqpstatus == :Primal_infeasible
        return MOI.INFEASIBLE
    elseif osqpstatus == :Max_iter_reached
        return MOI.ITERATION_LIMIT
    elseif osqpstatus == :Solved
        return MOI.OPTIMAL
    elseif osqpstatus == :Solved_inaccurate
        return MOI.ALMOST_OPTIMAL
    elseif osqpstatus == :Primal_infeasible_inaccurate
        return MOI.ALMOST_INFEASIBLE
    else
        @assert osqpstatus == :Non_convex
        return MOI.INVALID_MODEL
    end
end

function MOI.get(optimizer::Optimizer, a::MOI.PrimalStatus)
    if a.N > MOI.get(optimizer, MOI.ResultCount())
        return MOI.NO_SOLUTION
    end
    osqpstatus = optimizer.results.info.status
    if osqpstatus == :Unsolved
        return MOI.NO_SOLUTION
    elseif osqpstatus == :Primal_infeasible
        return MOI.INFEASIBILITY_CERTIFICATE
    elseif osqpstatus == :Solved
        return MOI.FEASIBLE_POINT
    elseif osqpstatus == :Primal_infeasible_inaccurate
        return MOI.NEARLY_INFEASIBILITY_CERTIFICATE
    elseif osqpstatus == :Dual_infeasible
        return MOI.INFEASIBILITY_CERTIFICATE
    else # :Interrupted, :Max_iter_reached, :Solved_inaccurate, :Non_convex (TODO: good idea? use OSQP.SOLUTION_PRESENT?)
        return MOI.NO_SOLUTION
    end
end

function MOI.get(optimizer::Optimizer, a::MOI.DualStatus)
    if a.N > MOI.get(optimizer, MOI.ResultCount())
        return MOI.NO_SOLUTION
    end
    osqpstatus = optimizer.results.info.status
    if osqpstatus == :Unsolved
        return MOI.NO_SOLUTION
    elseif osqpstatus == :Dual_infeasible
        return MOI.INFEASIBILITY_CERTIFICATE
    elseif osqpstatus == :Primal_infeasible
        return MOI.INFEASIBILITY_CERTIFICATE
    elseif osqpstatus == :Primal_infeasible_inaccurate
        return MOI.NEARLY_INFEASIBILITY_CERTIFICATE
    elseif osqpstatus == :Solved
        return MOI.FEASIBLE_POINT
    else # :Interrupted, :Max_iter_reached, :Solved_inaccurate, :Non_convex (TODO: good idea? use OSQP.SOLUTION_PRESENT?)
        return MOI.NO_SOLUTION
    end
end

## Variables:
function MOI.is_valid(optimizer::Optimizer, vi::VI)
    vi.value ∈ 1 : MOI.get(optimizer, MOI.NumberOfVariables())
end

## Variable attributes:
function MOI.get(optimizer::Optimizer, a::MOI.VariablePrimal, vi::VI)
    MOI.check_result_index_bounds(optimizer, a)
    x = ifelse(_contains(OSQP.SOLUTION_PRESENT, optimizer.results.info.status), optimizer.results.x, optimizer.results.dual_inf_cert)
    x[vi.value]
end

function MOI.set(optimizer::Optimizer, a::MOI.VariablePrimalStart, vi::VI, value)
    MOI.is_empty(optimizer) && throw(MOI.SetAttributeNotAllowed(a))
    optimizer.warmstartcache.x[vi.value] = value
end

## Constraints:
function MOI.is_valid(optimizer::Optimizer, ci::CI)
    MOI.is_empty(optimizer) && return false
    ci.value ∈ keys(optimizer.rowranges)
end

function MOI.set(optimizer::Optimizer, a::MOI.ConstraintDualStart, ci::CI, value)
    MOI.is_empty(optimizer) && throw(MOI.SetAttributeNotAllowed(a))
    rows = constraint_rows(optimizer, ci)
    for (i, row) in enumerate(rows)
        optimizer.warmstartcache.y[row] = -value[i] # opposite dual convention
    end
    nothing
end

# function modification:
function MOI.set(optimizer::Optimizer, attr::MOI.ConstraintFunction, ci::CI{Affine, <:IntervalConvertible}, f::Affine)
    MOI.is_valid(optimizer, ci) || throw(MOI.InvalidIndex(ci))
    row = constraint_rows(optimizer, ci)
    optimizer.modcache.A[row, :] = 0
    for term in f.terms
        col = term.variable_index.value
        coeff = term.coefficient
        optimizer.modcache.A[row, col] += coeff
    end
    Δconstant = optimizer.constrconstant[row] - f.constant
    optimizer.constrconstant[row] = f.constant
    optimizer.modcache.l[row] += Δconstant
    optimizer.modcache.u[row] += Δconstant
    nothing
end

function MOI.set(optimizer::Optimizer, attr::MOI.ConstraintFunction, ci::CI{VectorAffine, <:SupportedVectorSets}, f::VectorAffine)
    MOI.is_valid(optimizer, ci) || throw(MOI.InvalidIndex(ci))
    rows = constraint_rows(optimizer, ci)
    for row in rows
        optimizer.modcache.A[row, :] = 0
    end
    for term in f.terms
        row = rows[term.output_index]
        col = term.scalar_term.variable_index.value
        coeff = term.scalar_term.coefficient
        optimizer.modcache.A[row, col] += coeff
    end
    for (i, row) in enumerate(rows)
        Δconstant = optimizer.constrconstant[row] - f.constants[i]
        optimizer.constrconstant[row] = f.constants[i]
        optimizer.modcache.l[row] += Δconstant
        optimizer.modcache.u[row] += Δconstant
    end
end

# set modification:
function MOI.set(optimizer::Optimizer, attr::MOI.ConstraintSet, ci::CI{Affine, S}, s::S) where {S <: IntervalConvertible}
    MOI.is_valid(optimizer, ci) || throw(MOI.InvalidIndex(ci))
    interval = S <: Interval ? s : MOI.Interval(s)
    row = constraint_rows(optimizer, ci)
    constant = optimizer.constrconstant[row]
    optimizer.modcache.l[row] = interval.lower - constant
    optimizer.modcache.u[row] = interval.upper - constant
    nothing
end

function MOI.set(optimizer::Optimizer,  attr::MOI.ConstraintSet, ci::CI{VectorAffine, S}, s::S) where {S <: SupportedVectorSets}
    MOI.is_valid(optimizer, ci) || throw(MOI.InvalidIndex(ci))
    rows = constraint_rows(optimizer, ci)
    for (i, row) in enumerate(rows)
        constant = optimizer.constrconstant[row]
        optimizer.modcache.l[row] = lower(s, i) - constant
        optimizer.modcache.u[row] = upper(s, i) - constant
    end
    nothing
end

# partial function modification:
function MOI.modify(optimizer::Optimizer, ci::CI{Affine, <:IntervalConvertible}, change::MOI.ScalarCoefficientChange)
    MOI.is_valid(optimizer, ci) || throw(MOI.InvalidIndex(ci))
    row = constraint_rows(optimizer, ci)
    optimizer.modcache.A[row, change.variable.value] = change.new_coefficient
    nothing
end

# TODO: MultirowChange?

MOI.supports_constraint(optimizer::Optimizer, ::Type{Affine}, ::Type{<:IntervalConvertible}) = true
MOI.supports_constraint(optimizer::Optimizer, ::Type{VectorAffine}, ::Type{<:SupportedVectorSets}) = true

## Constraint attributes:
function MOI.get(optimizer::Optimizer, a::MOI.ConstraintDual, ci::CI)
    MOI.check_result_index_bounds(optimizer, a)
    y = ifelse(_contains(OSQP.SOLUTION_PRESENT, optimizer.results.info.status), optimizer.results.y, optimizer.results.prim_inf_cert)
    rows = constraint_rows(optimizer, ci)
    -y[rows]
end

# Objective modification
function MOI.modify(optimizer::Optimizer, attr::MOI.ObjectiveFunction, change::MOI.ScalarConstantChange)
    MOI.is_empty(optimizer) && throw(MOI.ModifyObjectiveNotAllowed(change))
    optimizer.objconstant = change.new_constant
end

function MOI.modify(optimizer::Optimizer, attr::MOI.ObjectiveFunction, change::MOI.ScalarCoefficientChange)
    MOI.is_empty(optimizer) && throw(MOI.ModifyObjectiveNotAllowed(change))
    optimizer.modcache.q[change.variable.value] = change.new_coefficient
end

# There is currently no ScalarQuadraticCoefficientChange.

MOIU.@model(OSQPModel, # modelname
    (), # scalarsets
    (MOI.Interval, MOI.LessThan, MOI.GreaterThan, MOI.EqualTo), # typedscalarsets
    (MOI.Zeros, MOI.Nonnegatives, MOI.Nonpositives), # vectorsets
    (), # typedvectorsets
    (), # scalarfunctions
    (MOI.ScalarAffineFunction, MOI.ScalarQuadraticFunction), # typedscalarfunctions
    (), # vectorfunctions
    (MOI.VectorAffineFunction,) # typedvectorfunctions
)

end # module


module SteadyState

using QuantumOptics
using IterativeSolvers, LinearMaps, LinearAlgebra, SparseArrays

const T_blas = Union{Float64,Float32,ComplexF64,ComplexF32}
const DecayRates = Union{Vector{T} where T<:Number, Matrix{T} where T<:Number, Nothing}

isblascompatible(M::Array{T,2}) where {T<:T_blas} = true
isblascompatible(M::AbstractArray{T,2}) where T = false

include("interface.jl")
include("generic_method.jl")

end # module


abstract type Aircraft end

# Mass properties
"""
    get_mass(ac::Aircraft)

Get aircraft mass (kg).
"""
function get_mass(ac::Aircraft) end

"""
    get_inertia_tensor(ac::Aircraft)

Get aircraft inertia tensor (kg·m²).
"""
function get_inertia_tensor(ac::Aircraft) end

"""
    get_xcg_mac(ac::Aircraft)

Get aircraft center of gravity position (units of mean aerodynamic chord).
"""
function get_xcg_mac(ac::Aircraft) end

# Geometric properties
"""
    get_surface(ac:Aircraft)

Get aircraft wing surface, ``S`` (m²).
"""
function get_surface(ac::Aircraft) end

"""
    get_chord(ac::Aircraft)

Get aircraft mean aerodynamic chord, ``c``, (m).
"""
function get_chord(ac::Aircraft) end

"""
    get_wing_span(ac::Aircraft)

Get aircraft wing span, ``b``, (m).
"""
function get_wing_span(ac::Aircraft) end

"""
    get_aspect_ratio(ac::Aircraft)

Get aspect ratio, ``b^2/S``.
"""
function get_aspect_ratio(ac::Aircraft)
    b = get_wing_span(ac)
    S = get_surface(ac)
    return b^2 / S
end

# Forces and moments
# TODO: document once headers are standarized. Right now they are too specific to F16Stevens
function calculate_prop_forces_moments end
function calculate_prop_gyro_effects end
function calculate_pdot end

function calculate_aero_forces_moments end

# Flight Control system
function tgear end


@testset "visualization" begin

    @testset "Peristent plots" begin
        # All tests use the same dir to also catch if one plot overwrites another by default
        testdir = "test_persistent_plots"

        struct DummyGraph
            nv::Int
            np::Int
        end
        NaiveNASlib.nvertices(g::DummyGraph) = g.nv
        NaiveGAflux.nparams(g::DummyGraph) = g.np
        struct PlotTestCand <: AbstractCandidate
            fitness::Real
            graph::DummyGraph
            PlotTestCand(fitness, nv, np) = new(fitness, DummyGraph(nv, np))
        end
        NaiveGAflux.model(c::PlotTestCand, f=identity) = f(c.graph)
        NaiveGAflux.fitness(c::PlotTestCand) = c.fitness

        try
            import MemPool
            @testset "PlotFitness" begin

                p = PlotFitness((args...;kwargs...) -> [], testdir)
                @test !isdir(p.basedir)

                p(PlotTestCand.(1:3, [10, 20, 30], [100, 200, 300]))

                @test p.plt == [1, [3, 2]]
                @test p.best == [3]
                @test p.avg == [2]

                p(PlotTestCand.(2:4, [10, 20, 30], [100, 200, 300]))
                @test p.plt == [1, [3, 2], 2, [4, 3]]
                @test p.best == [3, 4]
                @test p.avg == [2, 3]

                p2 = PlotFitness((args...;kwargs...) -> [args...], testdir)
                @test p2.plt == [[1,2], [3 2; 4 3]]
                @test p2.best == p.best
                @test p2.avg == p.avg
            end

            @testset "ScatterPop" begin

                p = ScatterPop((args...;kwargs...) -> true, testdir)
                @test !isdir(p.basedir)

                @test p(PlotTestCand.(1:3, [10, 20, 30], [100, 200, 300]))
                @test p.data ==  [[10 1 100; 20 2 200; 30 3 300]]

                @test p(PlotTestCand.(2:4, [20, 30, 40], [200, 300, 400]))
                @test p.data ==  [[10 1 100; 20 2 200; 30 3 300], [20 2 200; 30 3 300; 40 4 400]]

                p2 = ScatterPop((args...;kwargs...) -> true, testdir)
                @test p2.data == p.data

                @test p2(PlotTestCand.(3:5, [30, 40, 50], [300, 400, 500]))
                @test p2.data ==  [[10 1 100; 20 2 200; 30 3 300], [20 2 200; 30 3 300; 40 4 400], [30 3 300; 40 4 400; 50 5 500]]
            end

            @testset "ScatterOpt" begin
                NaiveGAflux.opt(c::PlotTestCand) = fitness(c) > 2 ? ADAM(nparams(c) - fitness(c)) : Flux.Optimiser([ShieldedOpt(Descent(nparams(c) - fitness(c)))])

                p = ScatterOpt((args...;kwargs...) -> true, testdir)
                @test !isdir(p.basedir)

                @test p(PlotTestCand.(1:3, [10, 20, 30], [100, 200, 300]))
                @test p.data ==  [[1 99.0 Descent; 2 198.0 Descent; 3 297.0 ADAM]]

                @test p(PlotTestCand.(2:4, [20, 30, 40], [200, 300, 400]))
                @test p.data ==  [[1 99 Descent; 2 198 Descent; 3 297 ADAM], [2 198 Descent; 3 297 ADAM; 4 396 ADAM]]

                p2 = ScatterOpt((args...;kwargs...) -> true, testdir)
                @test p2.data == p.data

                @test p2(PlotTestCand.(3:5, [30, 40, 50], [300, 400, 500]))
                @test p2.data ==  [[1 99 Descent; 2 198 Descent; 3 297 ADAM], [2 198 Descent; 3 297 ADAM; 4 396 ADAM],[3 297 ADAM; 4 396 ADAM; 5 495 ADAM]]

                p3 = ScatterOpt((args...;kwargs...) -> true, testdir)
                @test p3(PlotTestCand.(3:5, [30, 40, 50], [300, 400, 500]))
                # Test that something was added
                @test length(p3.data) == 1 + length(p2.data)
                # What was added is just a copy paste of the last thing inserted to p2 so this lazy check works
                @test p3.data[end] == p2.data[end]

            end

        finally
            rm(testdir, force=true, recursive=true)
            MemPool.cleanup()
        end

        @testset "MultiPlot" begin
            struct MultiPlotTestMock
                nr::Int
                seen::Vector{Int}
                MultiPlotTestMock(nr) = new(nr, Int[])
            end
            NaiveGAflux.plotgen(p::MultiPlotTestMock, gen=length(p.seen)) = gen==0 ? p.nr : p.nr * p.seen[gen]
            function (p::MultiPlotTestMock)(pop)
                push!(p.seen, pop)
                return NaiveGAflux.plotgen(p)
            end

            pdata = []
            spy(x...) = push!(pdata, x...)
            p = MultiPlot(spy, MultiPlotTestMock.(2:4)...)

            @test pdata == 2:4

            @test p(3) == pdata ==[2,3,4,6,9,12]
            @test mapfoldl(pp -> pp.seen, vcat, p.plts) == [3,3,3]
        end

        @testset "CbAll" begin
            cb1 = false
            cb2 = false

            CbAll(pop -> cb1 = pop, pop -> cb2 = pop)(123)

            @test cb1 == 123
            @test cb2 == 123
        end
    end
end


# test_decimals.jl

@testset "Test decimal support" begin
    con = DBInterface.connect(DuckDB.DB)

    results = DBInterface.execute(
        con,
        "SELECT 42.3::DECIMAL(4,1) a, 4923.3::DECIMAL(9,1) b, 421.423::DECIMAL(18,3) c, 129481294.3392::DECIMAL(38,4) d"
    )

    # convert to DataFrame
    df = DataFrame(results)
    @test names(df) == ["a", "b", "c", "d"]
    @test size(df, 1) == 1
    @test df.a == [42.3]
    @test df.b == [4923.3]
    @test df.c == [421.423]
    @test df.d == [129481294.3392]

    DBInterface.close!(con)
end

# test returning decimals in a table function
function my_bind_function(info::DuckDB.BindInfo)
    DuckDB.add_result_column(info, "a", FixedDecimal{Int16, 0})
    DuckDB.add_result_column(info, "b", FixedDecimal{Int32, 1})
    DuckDB.add_result_column(info, "c", FixedDecimal{Int64, 2})
    DuckDB.add_result_column(info, "d", FixedDecimal{Int128, 3})
    return missing
end

mutable struct MyInitStruct
    pos::Int64

    function MyInitStruct()
        return new(0)
    end
end

function my_init_function(info::DuckDB.InitInfo)
    return MyInitStruct()
end

function my_main_function(info::DuckDB.FunctionInfo, output::DuckDB.DataChunk)
    init_info = DuckDB.get_init_info(info, MyInitStruct)

    a_array = DuckDB.get_array(output, 1, Int16)
    b_array = DuckDB.get_array(output, 2, Int32)
    c_array = DuckDB.get_array(output, 3, Int64)
    d_array = DuckDB.get_array(output, 4, Int128)
    count = 0
    multiplier = 1
    for i in 1:(DuckDB.VECTOR_SIZE)
        if init_info.pos >= 3
            break
        end
        a_array[count + 1] = 42 * multiplier
        b_array[count + 1] = 42 * multiplier
        c_array[count + 1] = 42 * multiplier
        d_array[count + 1] = 42 * multiplier
        count += 1
        init_info.pos += 1
        multiplier *= 10
    end

    DuckDB.set_size(output, count)
    return
end

@testset "Test returning decimals from a table functions" begin
    con = DBInterface.connect(DuckDB.DB)

    arguments::Vector{DataType} = Vector()
    DuckDB.create_table_function(con, "my_function", arguments, my_bind_function, my_init_function, my_main_function)
    GC.gc()

    # 3 elements
    results = DBInterface.execute(con, "SELECT * FROM my_function()")
    GC.gc()

    df = DataFrame(results)
    @test names(df) == ["a", "b", "c", "d"]
    @test size(df, 1) == 3
    @test df.a == [42, 420, 4200]
    @test df.b == [4.2, 42, 420]
    @test df.c == [0.42, 4.2, 42]
    @test df.d == [0.042, 0.42, 4.2]
end