# Clique types

# this is a developmental type, will be standardized after conclusion of #1010
# TODO resolve type instability
const MsgRelativeType = Vector{NamedTuple{(:variables, :likelihood), Tuple{Vector{Symbol},DFG.AbstractRelative}}}

const MsgPriorType = Dict{Symbol, MsgPrior{<:ManifoldKernelDensity}}

"""
    $TYPEDEF

Internal development types used during consolidation.  Stores relative and prior information making up a joint likelihood
message passed upward on the Bayes tree.
"""
mutable struct _MsgJointLikelihood
  relatives::IIF.MsgRelativeType
  priors::IIF.MsgPriorType
end

"""
  $(TYPEDEF)
Belief message for message passing on the tree.  This should be considered an incomplete joint probility.

Notes:
- belief -> Dictionary of [`TreeBelief`](@ref)
- variableOrder -> Ordered variable id list of the seperators in cliqueLikelihood
- cliqueLikelihood -> marginal distribution (<: `SamplableBelief`) over clique seperators.
- Older names include: productFactor, Fnew, MsgPrior, LikelihoodMessage

DevNotes:
- Used by both nonparametric and parametric.
- Objective for parametric case: `MvNormal(μ=[:x0;:x2;:l5], Σ=[+ * *; * + *; * * +])`.
- Part of the consolidation effort, see #459.
- Better conditioning for joint structure in the works using deconvolution, see #579, #635.
  - TODO confirm why <: Singleton.

$(TYPEDFIELDS)
"""
mutable struct LikelihoodMessage{T <: MessageType} <: AbstractPrior
  sender::NamedTuple{(:id,:step),Tuple{Int,Int}}
  status::CliqStatus
  belief::Dict{Symbol, TreeBelief} # will eventually be deprecated
  variableOrder::Vector{Symbol}
  cliqueLikelihood::Union{Nothing,SamplableBelief}  # TODO drop the Union
  msgType::T
  hasPriors::Bool
  # this is different from belief[].inferdim, as the total available infer dims remaining during down msgs -- see #910
  childSolvDims::Dict{Int, Float64}
  # calc differential factors for joint in the child clique
  jointmsg::_MsgJointLikelihood
  # diffJoints::Vector{NamedTuple{(:variables, :likelihood), Tuple{Vector{Symbol},DFG.AbstractRelative}}}
end

"""
    $TYPEDEF

Cache messages being passed on the tree, one container per clique.

Notes
- See model 2 (?) on IIF #674
"""
mutable struct MessageBuffer
  # up receive message buffer (multiple children, multiple messages)
  upRx::Dict{Int, LikelihoodMessage}
  # down receive message buffer (one parent)
  downRx::Union{Nothing, LikelihoodMessage}
  # RESERVED up outgoing message buffer (one parent)
  upTx::Union{Nothing, LikelihoodMessage}
  # RESERVED down outgoing message buffer (multiple children but one message)
  downTx::Union{Nothing, LikelihoodMessage}
end
MessageBuffer() = MessageBuffer(Dict{Int, LikelihoodMessage}(), nothing, nothing, nothing)

##==============================================================================
## BayesTreeNodeData
##==============================================================================

"""
$(TYPEDEF)

Data structure for each clique in the Bayes (Junction) tree.
"""
mutable struct BayesTreeNodeData
  status::CliqStatus
  frontalIDs::Vector{Symbol}
  separatorIDs::Vector{Symbol}
  inmsgIDs::Vector{Symbol} # Int
  potIDs::Vector{Symbol} # Int # this is likely redundant TODO -- remove
  potentials::Vector{Symbol}
  partialpotential::Vector{Bool}

  dwnPotentials::Vector{Symbol}
  dwnPartialPotential::Vector{Bool}

  cliqAssocMat::Array{Bool,2}
  cliqMsgMat::Array{Bool,2}
  directvarIDs::Vector{Symbol}
  directFrtlMsgIDs::Vector{Symbol}
  msgskipIDs::Vector{Symbol}
  itervarIDs::Vector{Symbol}
  directPriorMsgIDs::Vector{Symbol}
  debug
  debugDwn

  allmarginalized::Bool
  initialized::Symbol
  upsolved::Bool
  downsolved::Bool
  isCliqReused::Bool             # holdover

  # JT Local messages saved for cache and debugging, see IIF #675
  messages::MessageBuffer
end

## Packed types for serialization

mutable struct PackedBayesTreeNodeData
  frontalIDs::Vector{Symbol}
  separatorIDs::Vector{Symbol}
  inmsgIDs::Vector{Symbol} # Int
  potIDs::Vector{Symbol} # Int # this is likely redundant TODO -- remove
  potentials::Vector{Symbol}
  partialpotential::Vector{Bool}
  dwnPotentials::Vector{Symbol}
  dwnPartialPotential::Vector{Bool}
  cliqAssocMat::Array{Bool,2}
  cliqMsgMat::Array{Bool,2}
  directvarIDs::Vector{Symbol} # Int
  directFrtlMsgIDs::Vector{Symbol} # Int
  msgskipIDs::Vector{Symbol} # Int
  itervarIDs::Vector{Symbol} # Int
  directPriorMsgIDs::Vector{Symbol} # Int
end

## Full Clique Types

struct CliqueId{T}
  value::T
end

"""
    $(TYPEDEF)
Structure to store clique data
DEV NOTES: To replace TreeClique completely
    $(FIELDS)
"""
mutable struct TreeClique
  "Interger id unique within a tree with userId, robotId, sessionId"
  id::CliqueId{Int64} # not to be confused with the underlying index used by LightGraphs.jl, see issue #540
  "Data as `BayesTreeNodeData`"
  data::BayesTreeNodeData
  "Drawing attributes"
  attributes::Dict{String, Any}
  #solveInProgress #on a clique level a "solve in progress" might be very handy
end

#


# hash
hash(c::AbstractGray, hx::UInt) = hash(gray(c), hx)
hash(c::TransparentGray, hx::UInt) = hash(alpha(c), hash(gray(c), hx))
hash(c::AbstractRGB, hx::UInt) = hash(blue(c), hash(green(c), hash(red(c), hx)))
hash(c::TransparentRGB, hx::UInt) = hash(alpha(c), hash(blue(c), hash(green(c), hash(red(c), hx))))

Base.adjoint(c::Colorant) = c

# gamut{min,max}
gamutmax(::Type{T}) where {T<:HSV} = (360,1,1)
gamutmin(::Type{T}) where {T<:HSV} = (0,0,0)
gamutmax(::Type{T}) where {T<:HSL} = (360,1,1)
gamutmin(::Type{T}) where {T<:HSL} = (0,0,0)
gamutmax(::Type{T}) where {T<:Lab} = (100,128,128)
gamutmin(::Type{T}) where {T<:Lab} = (0,-127,-127)
gamutmax(::Type{T}) where {T<:LCHab} = (100,1,360)
gamutmin(::Type{T}) where {T<:LCHab} = (0,0,0)
gamutmax(::Type{T}) where {T<:YIQ} = (1,0.5226,0.5226)
gamutmin(::Type{T}) where {T<:YIQ} = (0,-0.5957,-0.5957)

gamutmax(::Type{T}) where {T<:AbstractGray} = (1,)
gamutmax(::Type{T}) where {T<:TransparentGray} = (1,1)
gamutmax(::Type{T}) where {T<:AbstractRGB} = (1,1,1)
gamutmax(::Type{T}) where {T<:TransparentRGB} = (1,1,1,1)
gamutmin(::Type{T}) where {T<:AbstractGray} = (0,)
gamutmin(::Type{T}) where {T<:TransparentGray} = (0,0)
gamutmin(::Type{T}) where {T<:AbstractRGB} = (0,0,0)
gamutmin(::Type{T}) where {T<:TransparentRGB} = (0,0,0,0)

# rand
for t in [Float16,Float32,Float64,N0f8,N0f16,N0f32]
    @eval _rand(::Type{T}) where {T<:Union{AbstractRGB{$t},AbstractGray{$t}}} =
          mapc(x->rand(eltype(T)), base_colorant_type(T)())
end

function _rand(::Type{T}) where T<:Colorant
    Gmax = gamutmax(T)
    Gmin = gamutmin(T)
    Mi = eltype(T) <: FixedPoint ? 1.0/typemax(eltype(T)) : 1.0
    A = rand(eltype(T), length(T))
    for j in eachindex(Gmax)
        A[j] = A[j] * (Mi * (Gmax[j]-Gmin[j])) + Gmin[j]
    end
    T(A...)
end

if Base.VERSION < v"1.6.0-DEV.1083"
    reinterpretc(::Type{T}, A::Array{C}) where {T<:Real,C<:AbstractGray} = reinterpret(T, A)
    reinterpretc(::Type{T}, A::Array{C}) where {T<:Real,C<:Colorant} = reshape(reinterpret(T, A), (sizeof(C)÷sizeof(T), size(A)...))
else
    reinterpretc(::Type{T}, A::Array{C}) where {T<:Real,C<:Colorant} = reinterpret(reshape, T, A)
end

function _rand(::Type{T}, sz::Dims) where T<:Colorant
    Mi = eltype(T) <: FixedPoint ? 1.0/typemax(eltype(T)) : 1.0
    A = Array{T}(undef, sz)
    Tr = eltype(T) <: FixedPoint ? FixedPointNumbers.rawtype(eltype(T)) : eltype(T)
    nchannels = sizeof(T)÷sizeof(eltype(T))
    Au = Random.UnsafeView(convert(Ptr{Tr}, pointer(A)), length(A)*nchannels)
    rand!(Au)

    Gmax = gamutmax(T)
    Gmin = gamutmin(T)
    s = (Gmax .- Gmin) .* Mi
    if !all(iszero, Gmin) || !all(==(1), s)
        Ar = reinterpretc(eltype(T), A)
        if T<:AbstractGray
            s1, offset1 = s[1], Gmin[1]
            for I in CartesianIndices(A)
                Ar[I] = Ar[I] * s1 + offset1
            end
        else
            for I in CartesianIndices(A)
                for j in eachindex(s)
                    Ar[j,I] = Ar[j,I] * s[j] + Gmin[j]
                end
            end
        end
    end
    return A
end

rand(::Type{T}, sz::Dims...) where {T<:Colorant} = _rand(ccolor(T, base_colorant_type(T){Float64}), sz...)

rand(::Type{Gray24}) = Gray24(rand(N0f8))
rand(::Type{Gray24}, sz::Dims) = Gray24.(rand(N0f8,sz))
rand(::Type{AGray32}) = AGray32(rand(N0f8),rand(N0f8))
rand(::Type{AGray32}, sz::Dims) = AGray32.(rand(N0f8,sz),rand(N0f8,sz))

# broadcast

# Mapping a function over color channels
"""
    mapc(f, rgb) -> rgbf
    mapc(f, rgb1, rgb2) -> rgbf

`mapc` applies the function `f` to each color channel of the input
color(s), returning an output color in the same colorspace.

# Examples:

    julia> mapc(x->clamp(x,0,1), RGB(-0.2,0.3,1.2))
    RGB{Float64}(0.0,0.3,1.0)

    julia> mapc(max, RGB(0.1,0.8,0.3), RGB(0.5,0.5,0.5))
    RGB{Float64}(0.5,0.8,0.5)

    julia> mapc(+, RGB(0.1,0.8,0.3), RGB(0.5,0.5,0.5))
    RGB{Float64}(0.6,1.3,0.8)
"""
@inline mapc(f, c::C) where {C<:ColorantN{1}} =
    base_colorant_type(C)(f(comp1(c)))
@inline mapc(f, c::C) where {C<:ColorantN{2}} =
    base_colorant_type(C)(f(comp1(c)), f(comp2(c)))
@inline mapc(f, c::C) where {C<:ColorantN{3}} =
    base_colorant_type(C)(f(comp1(c)), f(comp2(c)), f(comp3(c)))
@inline mapc(f, c::C) where {C<:ColorantN{4}} =
    base_colorant_type(C)(f(comp1(c)), f(comp2(c)), f(comp3(c)), f(comp4(c)))
@inline mapc(f, c::C) where {C<:ColorantN{5}} =
    base_colorant_type(C)(f(comp1(c)), f(comp2(c)), f(comp3(c)), f(comp4(c)), f(comp5(c)))

@inline mapc(f, x::Number) = f(x)

mapc(f::F, x, y) where F = _mapc(_same_colorspace(x,y), f, x, y)
_mapc(::Type{C}, f, x::ColorantN{1}, y::ColorantN{1}) where C =
    C(f(comp1(x), comp1(y)))
_mapc(::Type{C}, f, x::ColorantN{2}, y::ColorantN{2}) where C =
    C(f(comp1(x), comp1(y)), f(comp2(x), comp2(y)))
_mapc(::Type{C}, f, x::ColorantN{3}, y::ColorantN{3}) where C =
    C(f(comp1(x), comp1(y)), f(comp2(x), comp2(y)), f(comp3(x), comp3(y)))
_mapc(::Type{C}, f, x::ColorantN{4}, y::ColorantN{4}) where C =
    C(f(comp1(x), comp1(y)), f(comp2(x), comp2(y)), f(comp3(x), comp3(y)), f(comp4(x), comp4(y)))
_mapc(::Type{C}, f, x::ColorantN{5}, y::ColorantN{5}) where C =
    C(f(comp1(x), comp1(y)), f(comp2(x), comp2(y)), f(comp3(x), comp3(y)), f(comp4(x), comp4(y)), f(comp5(x), comp5(y)))

_same_colorspace(x::Colorant, y::Colorant) = _same_colorspace(base_colorant_type(x),
                                                              base_colorant_type(y))
_same_colorspace(::Type{C}, ::Type{C}) where {C<:Colorant} = C
@noinline _same_colorspace(::Type{C1}, ::Type{C2}) where {C1<:Colorant,C2<:Colorant} =
    throw(ArgumentError("$C1 and $C2 are from different colorspaces"))

"""
    reducec(op, v0, c)

Reduce across color channels of `c` with the binary operator
`op`. `v0` is the neutral element used to initiate the reduction. For
grayscale,

    reducec(op, v0, c::Gray) = op(v0, comp1(c))

whereas for RGB

    reducec(op, v0, c::RGB) = op(comp3(c), op(comp2(c), op(v0, comp1(c))))

If `c` has an alpha channel, it is always the last one to be folded into the reduction.
"""
@inline reducec(op, v0, c::ColorantN{1}) = op(v0, comp1(c))
@inline reducec(op, v0, c::ColorantN{2}) = op(comp2(c), op(v0, comp1(c)))
@inline reducec(op, v0, c::ColorantN{3}) = op(comp3(c), op(comp2(c), op(v0, comp1(c))))
@inline reducec(op, v0, c::ColorantN{4}) = op(comp4(c), op(comp3(c), op(comp2(c), op(v0, comp1(c)))))
@inline reducec(op, v0, c::ColorantN{5}) = op(comp5(c), op(comp4(c), op(comp3(c), op(comp2(c), op(v0, comp1(c))))))

@inline reducec(op, v0, x::Number) = op(v0, x)

"""
    mapreducec(f, op, v0, c)

Reduce across color channels of `c` with the binary operator `op`,
first applying `f` to each channel. `v0` is the neutral element used
to initiate the reduction. For grayscale,

    mapreducec(f, op, v0, c::Gray) = op(v0, f(comp1(c)))

whereas for RGB

    mapreducec(f, op, v0, c::RGB) = op(f(comp3(c)), op(f(comp2(c)), op(v0, f(comp1(c)))))

If `c` has an alpha channel, it is always the last one to be folded into the reduction.
"""
@inline mapreducec(f, op, v0, c::ColorantN{1}) = op(v0, f(comp1(c)))
@inline mapreducec(f, op, v0, c::ColorantN{2}) = op(f(comp2(c)), op(v0, f(comp1(c))))
@inline mapreducec(f, op, v0, c::ColorantN{3}) =
    op(f(comp3(c)), op(f(comp2(c)), op(v0, f(comp1(c)))))
@inline mapreducec(f, op, v0, c::ColorantN{4}) =
    op(f(comp4(c)), op(f(comp3(c)), op(f(comp2(c)), op(v0, f(comp1(c))))))
@inline mapreducec(f, op, v0, c::ColorantN{5}) =
    op(f(comp5(c)), op(f(comp4(c)), op(f(comp3(c)), op(f(comp2(c)), op(v0, f(comp1(c)))))))

@inline mapreducec(f, op, v0, x::Number) = op(v0, f(x))


# # A Stochastic Splicing Simulator

# ## Overview
# Slicing is a process by which subsequences of preRNA (called introns, by definition) are removed.
# The remaining subsequences (called exons) are rejoined together.
# The resulting RNA sequence with introns removed is called mature RNA (mRNA).

# The splicing process is often stochastic.  Given a preRNA sequence, there is a
# distribution over corresponding mRNA sequences.
# As is almost always the case, the stochasticity results from incomplete information.
# Sources of uncertainty:
# - Structure: RNA modifications / bulged nucleotides make pairing of rNA and snRNPs variable.
# - Mutation: Single nucleotide variation (Q: At what point do mutations happen?)
# - Spliceosome composition: Variability in RNA-Binding Protein concentrations of cell

# ## Model

# At a high level, this is an abstract biophysical model that simulates splicing using a discrete-time simulation.

# __Trans-elements__ are small nucleus RNA, proteins and their combination called small- nuclear
# ribonuclear proteins (snRNP).  Trans elements perform splicing by binding to the RNA
# in a sytematic fashion.

using BioSequences, Random

"A Trans-Element / Factor: Proten or snRNP that can bind to primanry RNA transcript"
abstract type TransElement end

# There are many properties of a trans-element that may affect its role in splicing.
# The simplest representation associates a transelement only with a name:

"Symbolic TransElement -- represents only its name"
struct OpaqueTransElement <: TransElement
  name::Symbol
end

# A __spliceosome__ is a complex of transelements.
# We model a spliceosome as a partial function from trans-elements to real values representing their
# relative concentrations:

"Collection of transelements at different concentrations"
struct Spliceosome{T <: TransElement, R <: Real}
  elems::Dict{T, R}   # Maps from transelement to its concentration in spliceosome
  # FIXME: Ensure concentratios are normalized
end

"Is transelement `te` in the splicesosome `sos`?"
insos(te, sos::Spliceosome) = te in keys(sos.elems)
# Spliceosome(elems::Dict{T, R}) where {T, R} = Spliceosome{T, R}(elems)

"concentration of transelement `te` in spliceosome `so`"
concentration(sos, te) = sos.elems[te]

"Number of trans-elements in splicesosome"
ntranselems(sos::Spliceosome) = length(sos.elems)

# The __primary transcript__ is a sequence of RNA to be spliced by the spliceosome
const RNASequence = BioSequence{RNAAlphabet{4}}

# Trans-elements can bind to sequences in primary transcript.
# A __bound sequence__ or transcript is a primary transcript with any number of trans-elements bound to it.

"Position -- Natural number 1:length(seq)"
const Pos = Int

"Primary transcript with tran-elements bound to it"
struct BoundSeq{T <: TransElement}
  seq::RNASequence             # Primary transcriipt
  binds::Vector{Tuple{Pos, T}} # position -> transelement bound at that position
end

BoundSeq(seq::RNASequence, ::Type{T} = TransElement) where T =
  BoundSeq(seq, Tuple{Pos, T}[])

# ### Interactions
# Splicing literature makes a distinction between the core splicing machinery
# and __splicing regulary elements__.
# The core splicing machinery is carried out by the __major splicesosome__, which
# consists of exactly 5 rnSNAs and around 150 proteins, which combine to make
# 5 rnSNPs: U1, U2, U4, U5, U6.
# The core splicing mechanism is an ordered, stepwise assembly of discrete snRNP particles
# on the pre-mRNA primary transcript.  First U1 binds to the 5' splice site, U2 binds
# to a branch site, and so on until a split occurs.

# Splicing regulartory elements (SREs) are patterns on RNA that inhibit or promote splicing machinery.
# They are believed to be the main mechanism for alternative splcing.

# We will model both the core machinery and SREs within the same framework.
# We model Interactions as tuples of the form $(p, e)$, where $p$ is a __pattern__ (described below)
# and $e$ is a set of __effects__:

struct Interaction{P, TE, F}
  pattern::P
  te::TE
  loc::Pos
  factor::F
end

# #### Patterns

# A pattern (or motif) is a property of an RNA sequence.
# Formally a pattern is a set of strings, i.e., a language.
# For instance, a g-run is the set of strings $\{G, GG, GGG, GGGG, GGGGG, \dots\}$.
# A pattern $p$ __matches__ a string $s$ if any element of $p$ is a subsequence of $s$.

# We assume that the extent to which a pattern has occured is graded.
# A __pattern weight function__ maps patterns to a non-negative real number, denoting the degree to which it has matched.
# p(seq) = 1 indicates a perfect match and p(seq) = 0 indicates no match at all.

# #### EFfects

# An effect $E \subseteq \textrm{TransElements} \times \textrm{Location} \times \mathbb{R}$ is a ternary relation:
# a set of tuples of the form $(t, l, f)$, where $t$ is a transelement, $l$ is an integer valued location, and $f$ is a factor.
# Informally, the semantics of an interaction $(p, \{(t_1, l_1, f_1), (t_2, l_2, f_2), \dots\})$.
# are that if a pattern $p$ is found, then the probability that transelement $t_i$ will bind to location $l_i$ is increated (or decreased) multiplicatively by factor $f_i$.

"Returns a protein extractor that returns the singleton `te` iff its in the splicesosome"
only(te::T) where T = sos -> insos(te, sos) ? [te] : T[]

"Initial probabilities for (TransElement, Position)"
initT(sos::Spliceosome{T, R}, seq) where {T, R} = Dict{Tuple{T, Pos}, R}()

"Absolute position where pattern matched"
matchpos(m::BioSequences.RE.RegexMatch) = m.captured[1]

"Simulate one binding step"
function step(rng, bseq, sos, intrs)
  T = initT(sos, bseq.seq)

  # For each interaction
  for intr in intrs
    # and for each match of the interaction's pattern
    for x in eachmatch(intr.pattern, bseq.seq)
      # For each
      abspos = matchpos(x) + intr.loc
      for protein in intr.te(sos)
        if (protein, abspos) ∉ keys(T)
          T[protein, abspos] = 1.0
        end
        T[protein, abspos] *= intr.factor
      end
    end
  end
  bseq = bind(rng, T, bseq, sos)
  maybesplice(bseq)
end

"Simulate a binding process"
function bind(rng, T, bseq, sos)
  @show T
  # Select transelement with probability proportional to concentration
  # bseq

  # How
end

"Simulate core splicing machinery -- checks for presence of u1 and u2"
maybesplice(bseq) = bseq

# The splicosome is then simulated by simply repeating this process until some criterion of convergence is me
"Simulate the splicing"
function splice(rng, bseq::BoundSeq, sos, intrs; n = 1000)
  for i = 1:n
    bseq = step(rng, bseq, sos, intrs)
    bseq = maybesplice(bseq)
  end
  bseq
end

splice(bseq, sos, intrs) = splice(Random.GLOBAL_RNG, bseq, sos, intrs)
splice(rng, rnaseq::RNASequence, sos, intrs) = splice(rng, BoundSeq(rnaseq), sos, intrs)

# Model
# A splicing round occurs when U1 and U2AF proteins reach a global minimum

# # Measurement (Sequencing)

# Unfortunately, we do not have exact data, and hence we must model the stochasticity in the sequencing process.
# The following model of sequencing assumes that from a given `rna`, `n` (possibly differing, due to alternative splicing) mRNAs
# are produced, and the sequencing measures a single sunsequence (of random length) from each

"Sequence RNA process"
function sequence(rng, seq, n)
  sequences = RNASequence[]

  splice_start(rng) = uniform(rng, 1:length(seq))

  splce_length(rng) = uniform(rng, 1:100)

  for i = 1:n
    mRNA = simsplice(rng, seq, n)
    lb = splice_start(rng)
    ub = lb + splice_length(rng)
    push!(sequences, mRNA[lb:ub]) # FIXME: Might go out of bounds
  end
  sequences
end


module Perturbo
export ElecHam, bands, hamiltonian!

include("./types.jl")
include("./WignerSeitzCell.jl")
include("./load_data.jl")
include("./electronic_structure.jl")

end # module


export iauDtf2d
"""
Encode date and time fields into 2-part Julian Date (or in the case
of UTC a quasi-JD form that includes special provision for leap
seconds).

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

Status:  support function.

Given:
   scale     char[]  time scale ID (Note 1)
   iy,im,id  int     year, month, day in Gregorian calendar (Note 2)
   ihr,imn   int     hour, minute
   sec       double  seconds

Returned:
   d1,d2     double  2-part Julian Date (Notes 3,4)

Returned (function value):
            int     status: +3 = both of next two
                              +2 = time is after end of day (Note 5)
                              +1 = dubious year (Note 6)
                              0 = OK
                              -1 = bad year
                              -2 = bad month
                              -3 = bad day
                              -4 = bad hour
                              -5 = bad minute
                              -6 = bad second (<0)

Notes:

   1. scale identifies the time scale.  Only the value "UTC" (in upper
   case) is significant, and enables handling of leap seconds (see
   Note 4).

   2. For calendar conventions and limitations, see iauCal2jd.

   3. The sum of the results, d1+d2, is Julian Date, where normally d1
   is the Julian Day Number and d2 is the fraction of a day.  In the
   case of UTC, where the use of JD is problematical, special
   conventions apply:  see the next note.

   4. JD cannot unambiguously represent UTC during a leap second unless
   special measures are taken.  The SOFA internal convention is that
   the quasi-JD day represents UTC days whether the length is 86399,
   86400 or 86401 SI seconds.  In the 1960-1972 era there were
   smaller jumps (in either direction) each time the linear UTC(TAI)
   expression was changed, and these "mini-leaps" are also included
   in the SOFA convention.

   5. The warning status "time is after end of day" usually means that
   the sec argument is greater than 60.0.  However, in a day ending
   in a leap second the limit changes to 61.0 (or 59.0 in the case
   of a negative leap second).

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

   7. Only in the case of continuous and regular time scales (TAI, TT,
   TCG, TCB and TDB) is the result d1+d2 a Julian Date, strictly
   speaking.  In the other cases (UT1 and UTC) the result must be
   used with circumspection;  in particular the difference between
   two such results cannot be interpreted as a precise time
   interval.

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

This revision:  2013 July 26

SOFA release 2018-01-30

Copyright (C) 2018 IAU SOFA Board.  See notes at end.
"""
# int iauDtf2d(const char *scale, int iy, int im, int id,
#              int ihr, int imn, double sec, double *d1, double *d2)
function iauDtf2d(scale::String, iy::Real, im::Real, id::Real, ihr::Real, imn::Real, sec::Real)
   ref_d1 = Ref{Float64}(0.0)
   ref_d2 = Ref{Float64}(0.0)

   status = ccall((:iauDtf2d, libsofa_c), Cint,
                  (Cstring, Cint, Cint, Cint, Cint, Cint, Cdouble, Ref{Cdouble}, Ref{Cdouble}),
                  scale, convert(Int32, iy), convert(Int32, im), convert(Int32, id),
                  convert(Int32, ihr), convert(Int32, imn), convert(Float64, sec), ref_d1, ref_d2)

   if status != 0
      @warn "Non-zero return code form iauDtf2d"
   end

   return status, ref_d1[], ref_d2[]
end


# General idea here: An ILPOptimizer is a lightweight type used to create the a full model.
# An `ILPHolder` is the struct that holds most of the metadata about ILP construction.
# An `ILPOptimizer` is used to construct an `ILPHolder`.

# ILP Optimizer
abstract type ILPType end

struct Static <: ILPType end
struct Synchronous <: ILPType end
struct Asynchronous <: ILPType end

struct ILPOptimizer{T, U <: ILPType} <: AbstractOptimizer{T}
    # PMM to DRAM ratio - depending on `T`
    x::T
end

# Hijack `_getratio` to forward to the correct field.
_getratio(I::ILPOptimizer) = I.x

function (O::ILPOptimizer{Rational{Int64}})(data, backend::nGraph.Backend)
    bounds = Profiler.allocation_bounds(data)

    top = max(
        bounds.upper_bound / (getratio(O) + 1) / 1E6,
    )

    x = fill(round(Int, top), size(nodes(data)))
    println("Trying to use $(maximum(x)) MB of memory")
    return ILPHolder(O, x, backend; defrag = !iszero(_numerator(O)))
end

function (O::ILPOptimizer{Int64})(data, backend::nGraph.Backend)
    x = fill(round(Int, getlimit(O) / 1E6), size(nodes(data)))
    println("Trying to use $(maximum(x)) MB of memory")
    return ILPHolder(O, x, backend)
end

Static(x::T) where {T <: OPTIM_TYPES} = ILPOptimizer{T, Static}(x)
Synchronous(x::T) where {T <: OPTIM_TYPES} = ILPOptimizer{T, Synchronous}(x)
Asynchronous(x::T) where {T <: OPTIM_TYPES} = ILPOptimizer{T, Asynchronous}(x)

#####
##### TensorMeta
#####

struct TensorMeta
    graph::MetaGraph

    # Nodes using this tensor
    users::Vector{XNode}

    # Look-up a node wrapper, get the node that serves as a reference for this
    reference_map::Dict{XNode, XNode}

    # Flag to indicate that this is a fixed tensor - can decrease variable generation.
    isfixed::Bool
end

#####
##### ILPHolder
#####

mutable struct ILPHolder{T <: ILPType}
    dram_limits::Vector{Int}

    descriptors::Dict{XTensor{XNode}, TensorMeta}
    async_move_vars::Dict{XNode, Vector{JuMP.VariableRef}}

    # Need to key this with string names instead of XNode because of recompiling the ngraph
    # function and reporifling.
    #
    # New XNodes get created that don't hash the same as the old ones.
    node_to_limit_index::Dict{String, Int}

    # Bandwidths
    read_bandwidth::Int64
    write_bandwidth::Int64
    read_bandwidth_async::Int64
    write_bandwidth_async::Int64

    # Flag to determine if we need to defrag
    defrag::Bool
end

# Build an ILP Holder that will constain most of the metadata for IL: Construction.
function ILPHolder(O::ILPOptimizer{<:Any, T}, limits, backend; defrag = true) where {T}
    return ILPHolder{T}(
        limits,
        Dict{XTensor{XNode}, TensorMeta}(),
        Dict{XNode, Vector{JuMP.VariableRef}}(),
        Dict{String, Int}(),
        _bw_remote_local_sync(backend),
        _bw_local_remote_sync(backend),
        _bw_remote_local_async(backend),
        _bw_local_remote_async(backend),
        defrag,
    )
end

name(::Type{Static}) = "static"
name(::Type{Synchronous}) = "synchronous"
name(::Type{Asynchronous}) = "asynchronous"
name(::ILPOptimizer{<:Any, T}) where {T} = name(T)

# Struct to be passed around since all these items are generally used together anyways.
mutable struct Frame{T <: ILPHolder}
    modeltype::T
    model::JuMP.Model
    profile_data::FunctionData

    # Function parameters that are in the "local" memory and thus should be counted for the
    # total memory consumption.
    local_args::Vector{XTensor{XNode}}
end

Frame(modeltype, model::JuMP.Model, profile_data::FunctionData) =
    Frame(modeltype, model, profile_data, XTensor{XNode}[])

limit(F::Frame) = limit(F.modeltype)
JuMP.optimize!(F::Frame) = optimize!(F.model)

#####
##### Model Types
#####

# Implementations
include("tensor_graphs.jl")
include("formulation.jl")
include("defrag.jl")
include("configure.jl")
include("inspect.jl")
include("compare.jl")


using CpuId
using Libdl

## this lets us load CpuId only once

if cpufeature(:AVX2)
    libsuffix = :avx2
    println("AVX2 support detected, vml_avx2 selected")
else
    libsuffix = :avx
    println("AVX2 support missing, vml_avx selected")
end

if Sys.iswindows()
    rtlib = :mkl_rt
    corelib = :mkl_core
    lib = Symbol(:mkl_vml_, libsuffix)
else
    rtlib = :libmkl_rt
    corelib = :libmkl_core
    lib = Symbol(:libmkl_vml_, libsuffix)
end

depsjl_path = joinpath(@__DIR__, "deps.jl")
open(depsjl_path, "w") do depsjl_file
    println(depsjl_file, strip("""
        ## This file was autogenerated by build.jl.
        ## Do not edit.
        import Libdl

        const lib = :$lib
        const rtlib = :$rtlib
        const corelib = :$corelib
        """))
end


using ArgParse
using Dates
using Distributed
@everywhere include("./mining.jl")
@everywhere using Distributed
@everywhere using Dates
@everywhere using .Mining

ccall(:jl_exit_on_sigint, Cvoid, (Cint,), 0)

function parse_cmd()
    s = ArgParseSettings()

    @add_arg_table! s begin
            "data_path"
                help = "data folder relative to data directory"
                required = true
            "--sort_heuristic"
                default = "plain" # alternative is "uncovered_nodes" but this doesn't work as well
                arg_type = String
            "--surprise_threshold"
                default = 10.
                arg_type = Float64
            "--size_threshold"
                default = 10.
                arg_type = Float64
            "--stop_threshold"
                default = Inf
                arg_type = Float64
            "--stop_after_n_structures"
                default = Inf
                arg_type = Float64
            "--overwrite", "-o"
                default = false
                arg_type = Bool
            "--single_process", "-s"
                default = false
                arg_type = Bool
            "--debug", "-d"
                default = false
                arg_type = Bool
    end
    return parse_args(s)
end

@everywhere function build_model_for_dataset(dataset::String, surprise_threshold::Float64, size_threshold::Float64, stop_threshold::Float64, stop_after_n_structures::Float64, sort_heuristic::String; overwrite::Bool=false, debug::Bool=false)
    println("$(now()) - Starting to compute model for $(dataset) with size threshold $(size_threshold), stop threshold $(stop_threshold), and at most $(stop_after_n_structures) structures...")
    flush(stdout)
    Mining.compute_anytime_model(dataset, surprise_threshold, size_threshold, stop_threshold, stop_after_n_structures, sort_heuristic; overwrite=overwrite, debug=debug)
end

function main()
    parsed_args = parse_cmd()
    dsbase = parsed_args["data_path"]
    datasets = [dsbase * "/" * x for x in readdir(dsbase) if (endswith(x, ".txt") || endswith(x, ".csv")) && !occursin("nodemap",x)]
    params = [parsed_args["surprise_threshold"], parsed_args["size_threshold"], parsed_args["stop_threshold"], parsed_args["stop_after_n_structures"], parsed_args["sort_heuristic"]]

    if parsed_args["single_process"]
        for dataset in datasets
            build_model_for_dataset(dataset, params...; overwrite=parsed_args["overwrite"], debug=parsed_args["debug"])
        end
    else
        pmap(x->build_model_for_dataset(x, params...; overwrite=parsed_args["overwrite"], debug=parsed_args["debug"]), datasets)
    end
end

main()


import Base: getindex
using CUDA: CuArray, CuPtr
using Knet.LibKnet8: @knet8
using Knet.KnetArrays: checkbetween

## Indexing with Int array: used in nll.

function getindex(x::CuArray{T}, i::AbstractArray{I}) where {T,I<:Real}
    y = similar(x, size(i))
    if isempty(y); return y; end
    i = Array{Int32}(i)
    checkbetween(i, 1, length(x))
    unsafe_getindex!(x,y,CuArray{Int32}(i))
    return y
end

for F in (32,64); T=Symbol("Float$F"); @eval begin

## Indexing with CuArray{Int32}: low level, only Int32 supported, no bounds checking

    function unsafe_getindex!(x::CuArray{$T}, y::CuArray{$T}, i::CuArray{Int32})
        @knet8($("getents_$F"),(Cint,CuPtr{Int},CuPtr{$T},CuPtr{$T}), length(i), i, x, y)
        return y
    end

end; end


function Geoid(arg0::NormalizedSphericalHarmonicsProvider, arg1::ReferenceEllipsoid)
    return Geoid((NormalizedSphericalHarmonicsProvider, ReferenceEllipsoid), arg0, arg1)
end

function get_body_frame(obj::Geoid)
    return jcall(obj, "getBodyFrame", Frame, ())
end

function get_ellipsoid(obj::Geoid)
    return jcall(obj, "getEllipsoid", ReferenceEllipsoid, ())
end

function get_intersection_point(obj::Geoid, arg0::FieldLine, arg1::FieldVector3D, arg2::Frame, arg3::FieldAbsoluteDate)
    return jcall(obj, "getIntersectionPoint", FieldGeodeticPoint, (FieldLine, FieldVector3D, Frame, FieldAbsoluteDate), arg0, arg1, arg2, arg3)
end

function get_intersection_point(obj::Geoid, arg0::Line, arg1::Vector3D, arg2::Frame, arg3::AbsoluteDate)
    return jcall(obj, "getIntersectionPoint", GeodeticPoint, (Line, Vector3D, Frame, AbsoluteDate), arg0, arg1, arg2, arg3)
end

function get_undulation(obj::Geoid, arg0::jdouble, arg1::jdouble, arg2::AbsoluteDate)
    return jcall(obj, "getUndulation", jdouble, (jdouble, jdouble, AbsoluteDate), arg0, arg1, arg2)
end

function project_to_ground(obj::Geoid, arg0::TimeStampedPVCoordinates, arg1::Frame)
    return jcall(obj, "projectToGround", TimeStampedPVCoordinates, (TimeStampedPVCoordinates, Frame), arg0, arg1)
end

function project_to_ground(obj::Geoid, arg0::Vector3D, arg1::AbsoluteDate, arg2::Frame)
    return jcall(obj, "projectToGround", Vector3D, (Vector3D, AbsoluteDate, Frame), arg0, arg1, arg2)
end

function transform(obj::Geoid, arg0::FieldGeodeticPoint)
    return jcall(obj, "transform", FieldVector3D, (FieldGeodeticPoint,), arg0)
end

function transform(obj::Geoid, arg0::FieldVector3D, arg1::Frame, arg2::FieldAbsoluteDate)
    return jcall(obj, "transform", FieldGeodeticPoint, (FieldVector3D, Frame, FieldAbsoluteDate), arg0, arg1, arg2)
end

function transform(obj::Geoid, arg0::GeodeticPoint)
    return jcall(obj, "transform", Vector3D, (GeodeticPoint,), arg0)
end

function transform(obj::Geoid, arg0::Vector3D, arg1::Frame, arg2::AbsoluteDate)
    return jcall(obj, "transform", GeodeticPoint, (Vector3D, Frame, AbsoluteDate), arg0, arg1, arg2)
end


using KernelAbstractions, Test

@kernel function mul2(A)
    I = @index(Global, Cartesian)
    A[I] = 2 * A[I]
end

@kernel function add3(A, B, C)
    I = @index(Global, Cartesian)
    A[I] = B[I] + C[I]
end

@kernel function addi(A, C, i)
    I = @index(Global, Cartesian)
    A[I] = i + C[I]
end

function test_typed_kernel_dynamic(backend, ArrayT)
    A = ArrayT(ones(1024, 1024))
    kernel = mul2(backend())
    res = if backend == CPU
        @ka_code_typed kernel(A, ndrange=size(A), workgroupsize=16)
    else
        @ka_code_typed kernel(A, ndrange=size(A), workgroupsize=(32, 32))
    end
    @test isa(res, Pair{Core.CodeInfo, DataType})
    @test isa(res[1].code, Array{Any,1})
end

function test_typed_kernel_dynamic_no_info(backend, ArrayT)
    A = ArrayT(ones(1024, 1024))
    B = similar(A)
    C = similar(A)
    kernel = add3(backend())
    res = @ka_code_typed kernel(A, B, C, ndrange=size(A))
    @test isa(res, Pair{Core.CodeInfo, DataType})
    @test isa(res[1].code, Array{Any,1})
end

function test_typed_kernel_static(backend, ArrayT)
    A = ArrayT(ones(1024, 1024))
    kernel = if backend == CPU
        mul2(backend(), 16)
    else
        mul2(backend(), (32, 32))
    end
    res = @ka_code_typed kernel(A, ndrange=size(A))
    @test isa(res, Pair{Core.CodeInfo, DataType})
    @test isa(res[1].code, Array{Any,1})
end

function test_typed_kernel_no_optimize(backend, ArrayT)
    A = ArrayT(ones(1024, 1024))
    B = similar(A)
    C = similar(A)
    kernel = if backend == CPU
        add3(backend(), 16)
    else
        add3(backend(), (32, 32))
    end
    res = @ka_code_typed optimize=false kernel(A, ndrange=size(A))
    @test isa(res, Pair{Core.CodeInfo,Core.TypeofBottom})
    res_opt = @ka_code_typed kernel(A, ndrange=size(A))
    @test size(res[1].code) < size(res_opt[1].code)
end

function test_expr_kernel(backend, ArrayT)
    A = ArrayT(ones(1024, 1024))
    C = similar(A)
    kernel = if backend == CPU
        addi(backend())
    else
        addi(backend(), (32, 32))
    end
    res = @ka_code_typed kernel(A, C, 1+2, ndrange=size(A))
    @test isa(res, Pair{Core.CodeInfo, DataType})
    @test isa(res[1].code, Array{Any,1})
end

function reflection_testsuite(backend, ArrayT)
    test_typed_kernel_dynamic(backend, ArrayT)
    test_typed_kernel_dynamic_no_info(backend, ArrayT)
    test_typed_kernel_static(backend, ArrayT)
    test_typed_kernel_no_optimize(backend, ArrayT)
    test_expr_kernel(backend, ArrayT)
end


############################################################################################
#Include
include("stepnumber.jl")
include("kernel.jl")
include("config.jl")
include("diagnostics.jl")
include("postprocessing.jl")

############################################################################################
#export


# Defines structus and functions relating to display.
# Defines `render` for all game objects.

function blendAlphaColors(x::SDL2.Color, y::SDL2.Color)
    xAlphaPercent = x.a / 255
    yAlphaPercent = y.a / 255
    r = round(x.r*xAlphaPercent + (y.r - Int32(x.r) * xAlphaPercent) * yAlphaPercent)
    g = round(x.g*xAlphaPercent + (y.g - Int32(x.g) * xAlphaPercent) * yAlphaPercent)
    b = round(x.b*xAlphaPercent + (y.b - Int32(x.b) * xAlphaPercent) * yAlphaPercent)
    a = round(x.a + (1-xAlphaPercent)*yAlphaPercent * 255)
    SDL2.Color(r,g,b,a)
end

""" Coordinate system to represent the screen, not the Game World.
Importantly, these are "upside-down" from WorldCoords: (0,0) is top-left, and
higher numbers go down and to the right.
"""
abstract type ScreenCoords end

""" ScreenPixelCoords are absolute space on screen, in actual pixels. """
struct ScreenPixelCoords <: ScreenCoords end

"""
    ScreenPixelPos(1200, 1150)
Absolute position on screen, in pixels. (0,0) is top-left of screen.
"""
struct ScreenPixelPos <: AbstractPos{ScreenPixelCoords}  # 0,0 == top-left
    x::Int
    y::Int
end
ScreenPixelPos(x::Number, y::Number) = ScreenPixelPos(convert.(Int, floor.((x,y)))...)

# TODO: CONSIDER SWITCHING TO DOWN-SCALING INSTEAD of up-scaling, so you can get
# more precise sizes (1.5 "pixels" -> 3 pixels). Or consider using Floats for UIPixelCoords.
""" UIPixelCoords are space on screen, in un-dpi-scaled "pixels". """
struct UIPixelCoords <: ScreenCoords end

"""
    UIPixelPos(400, 450)
Position on screen, in un-dpi-scaled "pixels". (0,0) is top-left of screen.

These should be used whenever placing anything on the screen, since they are
indepdent of resolution-scaling.
"""
struct UIPixelPos <: AbstractPos{UIPixelCoords}  # 0,0 == top-left
    x::Int
    y::Int
end
UIPixelPos(x::Number, y::Number) = UIPixelPos(convert.(Int, floor.((x,y)))...)

+(a::UIPixelPos, b::UIPixelPos) = UIPixelPos(a.x+b.x, a.y+b.y)
# TODO: 1.0?
+(a::UIPixelPos, x::Number) = UIPixelPos(a.x+x, a.y+x)

""" Absolute size on screen, in pixels. Use with ScreenPixelPos. """
struct ScreenPixelDims <: AbstractDims{ScreenPixelCoords}
    w::Int32  # Int32 to match SDL
    h::Int32
end
""" Size on screen, in un-dpi-scaled "pixels". Use with UIPixelPos. """
struct UIPixelDims <: AbstractDims{UIPixelCoords}
    w::Int
    h::Int
end

mutable struct Camera
    pos::WorldPos
    # Note, these are in WorldPos size.
    w::Threads.Atomic{Float32}   # Note: These are Atomics, since they can be modified by the
    h::Threads.Atomic{Float32}   # windowEventWatcher callback, which can run in another thread!
end
Camera() = Camera(WorldPos(0,0),100,100)

screenCenter() = UIPixelPos(winWidth[]/2, winHeight[]/2)
screenCenterX() = winWidth[]/2
screenCenterY() = winHeight[]/2
screenOffsetFromCenter(x::Int,y::Int) = UIPixelPos(screenCenterX()+x,screenCenterY()+y)

dpiScale() = winWidth_highDPI[] / winWidth[];
worldScale(c::Camera) = dpiScale() * (winWidth[] / cam.w[]);
function toScreenPos(p::WorldPos, c::Camera)
    scale = worldScale(c)
    ScreenPixelPos(
        round(winWidth_highDPI[]/2.0 + scale*(p.x-c.pos.x)), round(winHeight_highDPI[]/2.0 - scale*(p.y-cam.pos.y)))
end
function toScreenPos(p::UIPixelPos, c::Camera)
    scale = dpiScale()
    ScreenPixelPos(round(scale*p.x), round(scale*p.y))
end
toScreenPos(p::ScreenPixelPos, c::Camera) = p
function toWorldPos(p::ScreenPixelPos, c::Camera)
    scale = worldScale(c)
    WorldPos((p.x - c.w[]/2.)/scale, -(p.y - c.h[]/2.)/scale)
end
function toWorldPos(p::UIPixelPos, c::Camera)
    toWorldPos(toScreenPos(p, cam), cam)
end
toWorldPos(p::WorldPos, c::Camera) = p
function toUIPixelPos(p::ScreenPixelPos, c::Camera)
    scale = dpiScale()
    UIPixelPos(round(p.x/scale), round(p.y/scale))
end
function toUIPixelPos(p::WorldPos, c::Camera)
    toUIPixelPos(toScreenPos(p, cam), cam)
end
toUIPixelPos(p::UIPixelPos, c::Camera) = p
function toScreenPixelDims(dims::UIPixelDims,c::Camera)
    scale = dpiScale()
    ScreenPixelDims(round(scale*dims.w), round(scale*dims.h))
end
function toScreenPixelDims(dims::WorldDims,c::Camera)
    scale = worldScale(c)
    ScreenPixelDims(round(scale*dims.w), round(scale*dims.h))
end
toScreenPixelDims(dims::ScreenPixelDims,c::Camera) = dims
function toUIPixelDims(dims::ScreenPixelDims,c::Camera)
    scale = dpiScale()
    UIPixelDims(round(dims.w/scale), round(dims.h/scale))
end
function toUIPixelDims(d::WorldDims, c::Camera)
    toUIPixelDims(toScreenDims(d, cam), cam)
end
toUIPixelDims(dims::UIPixelDims,c::Camera) = dims
function toWorldDims(dims::ScreenPixelDims,c::Camera)
    scale = worldScale(c)
    WorldDims(dims.w/scale, dims.h/scale)
end
function toWorldDims(d::UIPixelDims, c::Camera)
    toWorldDims(toScreenDims(d, cam), cam)
end
toWorldDims(dims::WorldDims,c::Camera) = dims

SetRenderDrawColor(renderer::Ptr{SDL2.Renderer}, c::SDL2.Color) = SDL2.SetRenderDrawColor(
    renderer, Int64(c.r), Int64(c.g), Int64(c.b), Int64(c.a))

# Convenience functions to allow `WorldPos(x,y)...` to become `x,y`
Base.iterate(p::AbstractPos, i=1) = if i==1 return (p.x,2) elseif i==2 return (p.y,3) else nothing end
Base.iterate(p::AbstractDims, i=1) = if i==1 return (p.w,2) elseif i==2 return (p.h,3) else nothing end
# To allow .+ (but doesn't work)
#Base.length(::Union{AbstractPos,AbstractDims}) = 2

topLeftPos(center::P, dims::D) where P<:AbstractPos{Coord} where D<:AbstractDims{Coord} where Coord<:ScreenCoords = P(center.x - dims.w/2.0, center.y - dims.h/2.0)  # positive is down
topLeftPos(center::P, dims::D) where P<:AbstractPos{Coord} where D<:AbstractDims{Coord} where Coord<:WorldCoords = P(center.x - dims.w/2.0, center.y + dims.h/2.)  # positive is up
rectOrigin(center::P, dims::D) where P<:AbstractPos{C} where D<:AbstractDims{C} where {C} = P(center.x - dims.w/2.0, center.y - dims.h/2.0)  # always minus

#makeSDL2RectFromOrigin(p::ScreenPixelPos, d::ScreenPixelDims) = SDL2.Rect(p..., d...)
#makeSDL2RectFromCenter(p::ScreenPixelPos, d::ScreenPixelDims) = SDL2.Rect(rectOrigin(p,d)..., d...)
#makeSDL2RectFromOrigin(p,d) = makeSDL2RectFromCenter(toScreenPos(p), toScreenPixelDims(d))
#makeSDL2RectFromCenter(p,d) =
function renderRectCentered(cam, renderer, center::AbstractPos{C}, dims::AbstractDims{C}, color; outlineColor=nothing) where C
    origin = topLeftPos(center, dims)
    renderRectFromOrigin(cam, renderer, origin, dims, color; outlineColor=outlineColor)
end
function renderRectFromOrigin(cam, renderer, origin::AbstractPos{C}, dims::AbstractDims{C}, color; outlineColor=nothing) where C
    screenPos = toScreenPos(origin, cam)
    rect = SDL2.Rect(screenPos.x, screenPos.y, toScreenPixelDims(dims, cam)...)
    if color != nothing
        SetRenderDrawColor(renderer, color)
        SDL2.RenderFillRect(renderer, Ref(rect) )
    end
    if outlineColor != nothing
        SetRenderDrawColor(renderer, outlineColor)
        SDL2.RenderDrawRect(renderer, Ref(rect) )
    end
end

function renderProgressBar(percent, cam::Camera, renderer, center::AbstractPos{C}, dims::D, color, bgColor, boxColor) where D<:AbstractDims{C} where C
    # bg
    renderRectCentered(cam, renderer, center, dims, bgColor)
    # progress
    origin = topLeftPos(center, dims)
    renderRectFromOrigin(cam, renderer, origin, D(round(percent * dims.w), dims.h), color)
    # outline
    renderRectCentered(cam, renderer, center, dims, nothing; outlineColor=boxColor)
end
function renderUnit(o::UnitTypes, playerColor, cam::Camera, renderer, dims::WorldDims, color)
    # First render the player color, then the unit color (for transparency)
    renderRectCentered(cam, renderer, o.pos, dims, playerColor)
    renderRectCentered(cam, renderer, o.pos, dims, color)

    # render health bar
    healthBarPos = WorldPos(o.pos.x, o.pos.y + dims.h/2 + healthBarRenderOffset)
    healthBarDims = WorldDims(healthBarRenderWidth, healthBarRenderHeight)
    renderProgressBar(health_percent(o), cam, renderer, healthBarPos,
          healthBarDims, healthBarColor, kBuildOpsBgColor,
          if toScreenPixelDims(healthBarDims, cam).h <= 2 # pixels
              nothing
          else
              healthBarOutlineColor
          end)
end
function render(o::Collector, playerColor, cam::Camera, renderer)
    unitColor = blendAlphaColors(playerColor, kCollectorColor)
    dims = WorldDims(collectorRenderWidth, collectorRenderWidth)
    tl = topLeftPos(o.pos, dims)
    for y in [tl.y, tl.y - dims.h]
        for x in tl.x+1 : 5 : tl.x-1 + dims.w
            nudge = Vector2D(rand(0.0:0.01:2.0,2)...)
            renderRectCentered(cam, renderer, WorldPos(x,y)+nudge, WorldDims(2,10), unitColor)
        end
    end
    for x in [tl.x, tl.x + dims.w]
        for y in tl.y-1 : -5 : tl.y+1 - dims.h
            nudge = Vector2D(rand(0.0:0.01:2.0,2)...)
            renderRectCentered(cam, renderer, WorldPos(x,y)+nudge, WorldDims(10,2), unitColor)
        end
    end
    renderUnit(o, playerColor, cam, renderer, dims, kCollectorColor)
end
function render(o::Fighter, playerColor, cam::Camera, renderer)
    dims = WorldDims(unitRenderWidth, unitRenderWidth)
    renderUnit(o, playerColor, cam, renderer, dims, kFighterColor)
end

abstract type AbstractButton end
mutable struct MenuButton <: AbstractButton
    enabled::Bool
    pos::UIPixelPos
    dims::UIPixelDims
    text::String
    callBack
end
mutable struct KeyButton <: AbstractButton
    enabled::Bool
    pos::UIPixelPos
    dims::UIPixelDims
    text::String
    callBack
end

mutable struct CheckboxButton
    toggled::Bool
    button::MenuButton
end

import Base.run
run(b::AbstractButton) = b.callBack()
function run(b::CheckboxButton)
    b.toggled = !b.toggled
    b.button.callBack(b.toggled)
end

# pointwise subtraction with bounds checking (floors to 0)
-(a::SDL2.Color, b::Int) = SDL2.Color(a.r-min(b,a.r), a.g-min(b,a.g), a.b-min(b,a.b), a.a-min(b,a.a))
SDL2.Color(1,5,1,1) - 2 == SDL2.Color(0,3,0,0)

function render(b::AbstractButton, cam::Camera, renderer, color, fontSize)
    if (!b.enabled)
         return
    end
    topLeft = topLeftPos(b.pos, b.dims)
    screenPos = toScreenPos(topLeft, cam)
    rect = SDL2.Rect(screenPos..., toScreenPixelDims(b.dims, cam)...)
    x,y = Int[0], Int[0]
    SDL2.GetMouseState(pointer(x), pointer(y))
    if clickedButton == b
        if mouseOnButton(UIPixelPos(x[],y[]),b,cam)
            color = color - 50
        else
            color = color - 30
        end
    else
        if mouseOnButton(UIPixelPos(x[],y[]),b,cam)
            color = color - 10
        end
    end
    SetRenderDrawColor(renderer, color)
    SDL2.RenderFillRect(renderer, Ref(rect) )
    renderText(renderer, cam, b.text, b.pos; fontSize = fontSize)
end

function render(b::MenuButton, cam::Camera, renderer)
    render(b, cam, renderer, kMenuButtonColor, kMenuButtonFontSize)
end
function render(b::KeyButton, cam::Camera, renderer)
    render(b, cam, renderer, kKeySettingButtonColor, kKeyButtonFontSize)
end

function render(b::CheckboxButton, cam::Camera, renderer)
    # Hack: move button text offcenter before rendering to accomodate checkbox
    offsetText = " "
    text_backup = b.button.text
    b.button.text = offsetText * b.button.text
    render(b.button, cam, renderer)
    b.button.text = text_backup

    # Render checkbox
    render_checkbox_square(b.button, 6, SDL2.Color(200,200,200, 255), cam, renderer)

    if b.toggled
        # Inside checkbox "fill"
        render_checkbox_square(b.button, 8, SDL2.Color(100,100,100, 255), cam, renderer)
    end
end

function render_checkbox_square(b::AbstractButton, border, color, cam, renderer)
    checkbox_radius = b.dims.h/2.0 - border  # (checkbox is a square)
    topLeft = topLeftPos(b.pos, b.dims)
    topLeft = topLeft + border
    screenPos = toScreenPos(topLeft, cam)
    screenDims = toScreenPixelDims(UIPixelDims(checkbox_radius*2, checkbox_radius*2), cam)
    rect = SDL2.Rect(screenPos..., screenDims...)
    SetRenderDrawColor(renderer, color)
    SDL2.RenderFillRect(renderer, Ref(rect) )
end

# ---- Text Rendering ----

fonts_cache = Dict()
txt_cache = Dict()
function sizeText(txt, fontName, fontSize)
    sizeText(txt, loadFont(dpiScale(), fontName, fontSize))
end
function sizeText(txt, font::Ptr{SDL2.TTF_Font})
   fw,fh = Cint[1], Cint[1]
   SDL2.TTF_SizeText(font, txt, pointer(fw), pointer(fh))
   return ScreenPixelDims(fw[1],fh[1])
end
function loadFont(scale, fontName, fontSize)
   fontSize = scale*fontSize
   fontKey = (fontName, Cint(round(fontSize)))
   if haskey(fonts_cache, fontKey)
       font = fonts_cache[fontKey]
   else
       font = SDL2.TTF_OpenFont(fontKey...)
       font == C_NULL && throw(ErrorException("Failed to load font '$fontKey'"))
       fonts_cache[fontKey] = font
   end
   return font
end
function createText(renderer, cam, txt, fontName, fontSize)
   font = loadFont(dpiScale(), fontName, fontSize)
   txtKey = (font, txt)
   if haskey(txt_cache, txtKey)
       tex = txt_cache[txtKey]
   else
       text = SDL2.TTF_RenderText_Blended(font, txt, SDL2.Color(20,20,20,255))
       tex = SDL2.CreateTextureFromSurface(renderer,text)
       #SDL2.FreeSurface(text)
       txt_cache[txtKey] = tex
   end

   fw,fh = Cint[1], Cint[1]
   SDL2.TTF_SizeText(font, txt, pointer(fw), pointer(fh))
   fw,fh = fw[1],fh[1]

   return tex, ScreenPixelDims(fw, fh)
end
@enum TextAlign centered leftJustified rightJustified
function renderText(renderer, cam::Camera, txt::String, pos::UIPixelPos
                    ; fontName = defaultFontName,
                     fontSize=defaultFontSize, align::TextAlign = centered)
   tex, fDims = createText(renderer, cam, txt, fontName, fontSize)
   renderTextSurface(renderer, cam, pos, tex, toUIPixelDims(fDims,cam), align)
end

function renderTextSurface(renderer, cam::Camera, pos::AbstractPos{C},
                           tex::Ptr{SDL2.Texture}, dims::AbstractDims{C}, align::TextAlign) where C
   screenPos = toScreenPos(pos, cam)
   screenDims = toScreenPixelDims(dims, cam)
   x,y, fw, fh = screenPos.x, screenPos.y, screenDims.w, screenDims.h
   renderPos = SDL2.Rect(0,0,0,0)
   if align == centered
       renderPos = SDL2.Rect(Int(floor(x-fw/2.0)), Int(floor(y-fh/2.0)), fw,fh)
   elseif align == leftJustified
       renderPos = SDL2.Rect(Int(floor(x)), Int(floor(y-fh/2.0)), fw,fh)
   else # align == rightJustified
       renderPos = SDL2.Rect(Int(floor(x-fw)), Int(floor(y-fh/2.0)), fw,fh)
   end
   SDL2.RenderCopy(renderer, tex, C_NULL, pointer_from_objref(renderPos))
   #SDL2.DestroyTexture(tex)
end

function hcat_render_text(lines, renderer, cam, gap, pos::UIPixelPos;
         fixedWidth=nothing, fontName=defaultFontName, fontSize=defaultFontSize)
    numLines = size(lines)[1]
    if fixedWidth != nothing
        widths = fill(fixedWidth,size(lines))
    else
        widths = [toUIPixelDims(sizeText(line, fontName, fontSize), cam).w for line in lines]
    end
    totalWidth = sum(widths) + gap*(numLines-1)
    runningWidth = 0
    leftMostPos = pos.x - totalWidth/2.0
    text_centers = []
    for i in 1:numLines
        linePos = leftMostPos + runningWidth
        leftPos = UIPixelPos(linePos, pos.y)
        renderText(renderer, cam, lines[i], leftPos; fontName=fontName, fontSize=fontSize, align=leftJustified)
        runningWidth += widths[i] + gap
        push!(text_centers, UIPixelPos(leftPos.x + widths[i]÷2, leftPos.y))
    end
    return text_centers
end

#  ------- Image rendering ---------

function render(t::Ptr{SDL2.Texture}, pos::AbstractPos{C}, cam::Camera, renderer; size::Union{Cvoid, AbstractDims{C}} = nothing) where C
    if (t == C_NULL) return end
    pos = toScreenPos(pos, cam)
    if size != nothing
        size = toScreenPixelDims(size, cam)
        w = size.w
        h = size.h
    else
        w,h,access = Cint[1], Cint[1], Cint[1]
        format = Cuint[1]
        SDL2.QueryTexture( t, format, access, w, h );
        w,h = w[], h[]
    end
    rect = SDL2.Rect(pos.x - w÷2,pos.y - h÷2,w,h)
    SDL2.RenderCopy(renderer, t, C_NULL, pointer_from_objref(rect))
end


"""
Things this needs to abstract:

  abstract type LinearSolver

Preconditioning:
  abstract type AbstractPC
  abstract type AbstractPetscMatPC
  abstract type AbstractPetscMatFreePC
  union of those two

Linear operators:
  abstract type AbstractLO
  abstract type AbstractDenseLO
  abstract type AbstractSparseDirectLO
  abstract type AbstractPetscMatLO
  abstract type AbstractPetscMatFreeLO
  union of those two

functions:
  calcLinearOperator  - calculate Jacobian, or if matfree, calc some stuff
  applyLinearOperator  - does Ax=b
  applyLinearOperatorTranspose  - does A'x = b
  needParallelData
  free

Todo:
  which implements the others (type -> function)
  signatures for functions

"""


"""
`module  JAC.Statistical`
    ... a submodel of JAC that contains all structs and methods to deal with time-dependent statistical tensors.
"""
module Statistical

    using ..ManyElectron

    """
    `struct  Statistical.ResonanceR`  ... defines a type for a resonance state in the continuum with a well-defined bound-ionic core,
             one or several electrons in the continuum, a widths as well as a loss rate due to additional decay processes.

        + isBound            ::Bool                ... True if it just refers to a bound state with no electron in the continuum.
        + ionLevel           ::Level               ... Level of the bound-state core.
        + widths             ::Float64             ... Widths of the resonance state.
        + lossRate           ::Float64             ... Loss rate due to processes that are not considered explicitly in the time evolution.
    """
    struct ResonanceR
        isBound              ::Bool
        ionLevel             ::Level
        widths               ::Float64
        lossRate             ::Float64
    end

    """
    `Statistical.ResonanceR()`  ... constructor for an `empty` instance of Statistical.ResonanceR().
    """
    function ResonanceR()
        ResonanceR( true, Level(), 0., 0. )
    end

    # `Base.show(io::IO, resonance::Statistical.ResonanceR)`  ... prepares a proper printout of the variable resonance::Statistical.ResonanceR.
    function Base.show(io::IO, resonance::Statistical.ResonanceR)
        println(io, "isBound:              $(resonance.isBound)  ")
        println(io, "ionLevel:             $(resonance.ionLevel)  ")
        println(io, "widths:               $(resonance.widths)  ")
        println(io, "lossRate:             $(resonance.lossRate)  ")
    end

    """
    `struct  Statistical.Tensor`  ... defines a type for a statistical tensor of given rank k, projection q and w.r.t. two resonances.

        + k                  ::AngularJ64          ... Rank of the tensor
        + q                  ::AngularM64          ... Projection of the tensor rank.
        + a                  ::ResonanceR          ... Level/resonance a with total symmetry alpha J.
        + b                  ::ResonanceR          ... Level/resonance a with total symmetry alpha' J'.
        + value              ::Complex{Float64}    ... value of the statistical tensor
    """
    struct Tensor
        k                    ::AngularJ64
        q                    ::AngularM64
        a                    ::ResonanceR
        b                    ::ResonanceR
        value                ::Complex{Float64}
    end

    """
    `Statistical.Tensor()`  ... constructor for an `empty` instance of Statistical.Tensor().
    """
    function Tensor()
        Tensor( AngularJ64(0), AngularM64(0), ResonanceR(), ResonanceR(), Complex(0.) )
    end

    # `Base.show(io::IO, tensor::Statistical.Tensor)`  ... prepares a proper printout of the variable tensor::Statistical.Tensor.
    function Base.show(io::IO, tensor::Statistical.Tensor)
        println(io, "k:              $(tensor.k)  ")
        println(io, "q:              $(tensor.q)  ")
        println(io, "a:              $(tensor.a)  ")
        println(io, "b:              $(tensor.b)  ")
        println(io, "value:          $(tensor.value)  ")
    end

end # module


using Distributions
export levitt_adapter, bayesian_adapter, constant_adapter, delta, estimate,
  update!

abstract type Adapter end

"""
    update!(adapter,response,correct)

Updates any internal state for the adapter when the listener responds with
`response` and the correct response is `correct`. Usually not called directly,
but instead called within `response`, when the adapter is passed as the first
argument. May take a while to run.
"""
function update!(adapter,response,correct)
end

"""
    estimate(adapter)

Returns the mean and error of the adapters threshold estimate. May take some
time to run.
"""
function estimate(adapter)
end

"""
    delta(adapter)

Returns the next delta that should be tested to help estimate the threshold.
"""
function delta(adapter)
end

"""
    response([fn],adapter,[key1] => ["resp1"],...;correct=[resp],
             [show_feedback=true],
             [feedback=Dict(true=>"Correct",false=>"Wrong!")]
             keys...)

Record a response in a n-alternative forced choice task and update
an adapter.

The first response recieved is interpreted as the actual response. Subsequent
responses will be recorded, without a delta or correct value set, and appending
"late_" to the specified response string.

# Function Callback

Optionally, upon participant response, `fn` receives two arguments: the
provided response, and the correct response.

# Keyword Arguments

- `correct` the response string corresponding to the correct response
- `show_feedback` (default = true): whether to show feedback to the
  participant after they respond.
- `feedback` the text to display to a participant when they are correct
  (for the true key, defaults to "Correct!") or incorrect (for the false key,
  defaults to "Wrong!").

Any additional keyword arguments are added as column values when the
response is recorded.

"""
function response(track::Adapter,resp::Pair...;keys...)
  response((x,y) -> nothing,track,resp...;keys...)
end

function response(callback::Function,adapter::Adapter,responses::Pair...;
                  correct=nothing,show_feedback=true,
                  feedback=Dict(true => "Correct!",false => "Wrong!"),
                  keys...)
  addcolumn(:correct)
  addcolumn(:delta)

  if correct ∉ map(x -> x[2],responses)
    error("The value of `correct` must be "*
          join(map(x -> x[2],responses),", "," or "))
  end
  feedback_visuals = map((pair) -> pair[1] => visual(pair[2]),feedback)

  let responded = false
    begin (event) ->
      for (key,resp) in responses
        if iskeydown(event,key)
          if !responded
            responded = true
            record(resp;correct=correct,delta=delta(adapter),keys...)

            update!(adapter,resp,correct)
            callback(resp,correct)
            if show_feedback
              display(feedback_visuals[resp==correct])
            end

          else
            record("late_"*resp;keys...)
          end
        end
      end
    end
  end
end

mutable struct ConstantStimulus
  stimuli::Vector{Float64}
  index::Int
  correct::Vector{Int}
  response::Vector{Int}
end

"""
    constant_adapter(stimuli)

An adapter that can be used to implement the method of constant stimuli: the
specified sequence of stimulus deltas is presented in order to participants.

Strictly speaking, this is not an adaptive tracking procedure. However,
it can be convienient to have the same programming interface for this method as
for adaptive methods. In this way you can easily select between the method of
constant stimuli or some kind of adaptive procedure.
"""
constant_adapter(stimuli) = ConstantStimulus(stimuli,1,Int[],Int[])
delta(adapter::ConstantStimulus) = adapter.stimuli[adapter.index]
function update!(adapter::ConstantStimulus,response,correct)
  adapte.index += 1
  nothing
  # push!(adapter.correct,correct)
  # push!(adapter.response,response)
end
function estimate(adapter::ConstantStimulus)
  error("The `estimate` method for `ConstantStimulus` is not implemented!")
  # TODO: fit a curve to the responses to estimate the threshold???
end

"""
    levitt_adapter([first_delta=0.1],[down=3],[up=1],
                   [big_reverse=3],[big=0.01],[little=0.005],
                   [min_reversals=7],[min_delta=-Inf],[max_delta=Inf],
                   [mult=false])

An adapter that finds a threshold according to a non-parametric statistical
procedure. This approach makes fewer explicit assumptions than
[`bayesian_adapter`](@ref) but may be slower to converge to a threshold.

This finds a threshold by moving the delta down after three correct responses
and up after one incorrect response (these default up and down counts can be
changed). This is the same approach described in Levitt (1971).

# Keyword Arguments

- `first_delta`: the delta that the first trial should present.
- `up`: how many incorrect responses in a row must occur for the delta
  to move up
- `down`: how many correct responses in a row must occur for the delta
  to move down.
- `big`: the amount delta changes by (up or down) at first
- `big_reverse`: how many reveresals (up to down or down to up) must
  occur before `little` is used instead of `big`
- `little`: the amount delta changes by (up or down) after
  `big_reverse` reversals.
- `min_reversals`: the smallest number of reversals that can
  be used to estimate a threshold.
- `min_delta`: the smallest delta allowed
- `max_delta`: the largest delta allowed
- `mult`: whether the delta change should be additive (false) or
   multiplicative (true).
"""
function levitt_adapter(;first_delta=0.1,down=3,up=1,big_reverse=3,
                        big=0.01,little=0.005,min_reversals=7,
                        drop_reversals=3,min_delta=0,max_delta=Inf,
                        mult=false)
  delta = first_delta
  reversals = Float64[]
  last_direction = 0
  num_correct = 0
  num_incorrect = 0
  operator = (mult ? :mult : :add)
  adapter = Levitt{operator}(delta,up,down,big_reverse,big,little,
                             min_reversals,drop_reversals,min_delta,max_delta,
                             reversals,last_direction,num_correct,num_incorrect)
end

type Levitt{OP} <: Adapter
  delta::Float64

  up::Int
  down::Int
  big_reverse::Int
  big::Float64
  little::Float64
  min_reversals::Int
  drop_reversals::Int

  min_delta::Float64
  max_delta::Float64

  reversals::Array{Float64}
  last_direction::Int
  num_correct::Int
  num_incorrect::Int
end
delta(a::Levitt) = a.delta

function update_reversals(adapter::Levitt,direction::Int)
  if (adapter.last_direction != 0 && direction != adapter.last_direction)
    push!(adapter.reversals,adapter.delta)
  end
  adapter.last_direction = direction
end

bound(x,min_,max_) = min(max_,max(min_,x))

up(adapter::Levitt{:add}) = +
down(adapter::Levitt{:add}) = -
up(adapter::Levitt{:mult}) = *
down(adapter::Levitt{:mult}) = /

function update_delta(adapter,op)
  if length(adapter.reversals) < adapter.big_reverse
    adapter.delta = bound(op(adapter.delta,adapter.big),
                          adapter.min_delta,adapter.max_delta)
  else
    adapter.delta = bound(op(adapter.delta,adapter.little),
                          adapter.min_delta,adapter.max_delta)
  end
end

function update!(adapter::Levitt,response,correct)
  if correct == response
    adapter.num_correct += 1
    adapter.num_incorrect = 0
    if adapter.num_correct >= adapter.down
      adapter.num_correct = 0
      update_reversals(adapter,-1)
      update_delta(adapter,down(adapter))
    end
  else
    adapter.num_incorrect += 1
    adapter.num_correct = 0
    if adapter.num_incorrect >= adapter.up
      adapter.num_incorrect = 0
      update_reversals(adapter,1)
      update_delta(adapter,up(adapter))
    end
  end
  adapter
end

estimate_helper(a::Levitt{:add},rev) = mean(rev),std(rev)
estimate_helper(a::Levitt{:mult},rev) =
  exp(mean(log(rev))),exp(mean(log(rev)))*(1-exp(-std(log(rev))))

function estimate(adapter::Levitt)
  if length(adapter.reversals) < adapter.min_reversals
    NaN,NaN
  else
    if isodd(length(adapter.reversals))
      estimate_helper(adapter,adapter.reversals[adapter.drop_reversals:end])
    else
      estimate_helper(adapter,adapter.reversals[adapter.drop_reversals+1:end])
    end
  end
end

type ImportanceSampler <: Adapter
  miss::Float64
  threshold::Float64
  thresh_d::UnivariateDistribution
  inv_slope_d::UnivariateDistribution
  n_samples::Int

  theta_prior::UnivariateDistribution
  inv_slope_prior::UnivariateDistribution

  delta::Float64
  min_delta::Float64
  max_delta::Float64
  delta_repeat::Int

  resp::Dict{Float64,Int}
  N::Dict{Float64,Int}
  thresh_samples::Vector{Float64}
  theta_samples::Vector{Float64}
  inv_slope_samples::Vector{Float64}
  thresh_weights::Vector{Float64}

  repeat2_thresh::Float64
  repeat3_thresh::Float64
end

# TODO: prior is working, but improve it to make the
# parameters easy to specify???

"""
    bayesian_adapter(;first_delta=0.1,
                     n_samples=1000,miss=0.01,threshold=0.79,
                     min_delta=0,max_delta=1,
                     min_plausible_delta = 0.0001,
                     max_plausible_delta = 0.2,
                     repeat3_thresh=1.0,repeat2_thresh=0.1,
                     thresh_prior=
                     Truncated(LogNormal(log(min_plausible_delta),
                                         log(max_plausible_delta/
                                             min_plausible_delta/2)),
                               min_delta,max_delta),
                          inv_slope_prior=TruncatedNormal(0,0.25,0,Inf),
                          thresh_d=thresh_prior,
                          inv_slope_d=inv_slope_prior)

An adapter that finds a threshold according to a parametric statistical
model. This makes more explicit assumptions than the [`levitt_adapter`](@ref)
but will normally find a threshold faster.

The psychometric curve is estimated from user responses using a bayesian
approach. After estimation, each new delta is selected in a greedy fashion,
picking the response that best minimizes entropy according to this psychometric
function. This is a modified version of the approach described in Kontsevich &
Tyler 1999. Specifically, the approach here uses importance sampling instead of
sampling parmeters over a deterministic, uniform grid. This should increase
measurement efficiency if the priors are chosen well.

This algorithm assumes the following functional form for the psychometric
response as a function of the stimulus difference ``Δ``.

``
f(Δ) = λ/2 + (1-λ) Φ((Δ - θ)⋅σ/√2)
``

In the above ``Φ`` is the cumulative distribution function of a normal
distribution, ``λ`` is the miss-rate parameter, indicating the rate at which
listeners make a mistake, even when the delta is large and easily heard, ``θ``
is the 50%-correct threshold, and ``σ`` is the psychometric slope.

For stability and robustness, this adapter begins by repeating the same delta
multiple times and only begins quickly changing deltas trial-by-trial when the
ratio of estiamted standard deviation to mean is small. This functionality can
be adjusted using `repeat3_thresh` and `repeat2_thresh`, or, if you do not wish
to have any repeats, both values can be set to Inf.

# Keyword Arugments

- **first_delta**: the delta to start measuring with
- **n_samples** the number of samples to use during importance sampling.
  The algorithm for selecting new deltas is O(n²).
- **miss** the expected rate at which listeners will make mistakes
  even for easy to percieve differences.
- **threshold** the %-response threshold to be estimated
- **min_delta** the smallest possible delta
- **max_delta** the largest possible delta
- **min_plausible_delta** the smallest plausible delta, should be > 0.
  Used to define a reasonable value for thresh_prior and inv_slope_prior.
- **max_plausible_delta** the largest plausible delta, should be < max_delta.
  Used to define a reasonable value for thresh_prior and inv_slope_prior.
- **thresh_prior** the prior probability distribution across thresholds.
  This influence the way the delta is adapted. By default this is defined in
  terms of min_plausible_delta and max_plausible_delta.
- **inv_slope_prior** the prior probability distribution across inverse slopes.
  By default this is defined in terms of min_plausible_delta and
  max_plausible_delta.
- **thresh_d** the distribution over-which to draw samples for the threshold
  during importance sampling. This defaults to thresh_prior
- **inv_slope_d** the distribution over-which to draw samples for the inverse slope
  during importance sampling. This defaults to inv_slope_prior.
- **repeat2_thresh** the ratio of sd / mean for theta must suprass
  to repeat each delta twice.
- **repeat3_thresh** the ratio of sd / mean for theta must surpass
  to repeat each delta thrice.
"""
function bayesian_adapter(;first_delta=0.1,
                          n_samples=1000,miss=0.01,threshold=0.79,
                          min_delta=0,max_delta=1,
                          min_plausible_delta = 0.0001,
                          max_plausible_delta = 0.2,
                          repeat3_thresh=1.0,repeat2_thresh=0.1,
                          thresh_prior=
                          Truncated(LogNormal(log(min_plausible_delta),
                                              log(max_plausible_delta/
                                                  min_plausible_delta/2)),
                                    min_delta,max_delta),
                          inv_slope_prior=
                          TruncatedNormal(0,2max_plausible_delta,0,Inf),
                          thresh_d=thresh_prior,
                          inv_slope_d=inv_slope_prior)
  delta = first_delta
  delta_repeat = 0
  resp = Dict{Float64,Int}()
  N = Dict{Float64,Int}()
  samples = ones(0)
  weights = ones(0)
  ImportanceSampler(miss,threshold,thresh_d,inv_slope_d,n_samples,
                    thresh_prior,inv_slope_prior,
                    delta,min_delta,max_delta,delta_repeat,
                    resp,N,samples,samples,samples,weights,
                    repeat3_thresh,repeat2_thresh)
end

const sqrt_2 = sqrt(2)
function sample(a::ImportanceSampler)
  if length(a.thresh_samples) == 0
    theta = rand(a.thresh_d,a.n_samples)
    inv_slope = rand(a.inv_slope_d,a.n_samples)
    q_log_prob = logpdf(a.thresh_d,theta) .+ logpdf(a.inv_slope_d,inv_slope)

    udeltas = collect(keys(a.resp))
    x = exp((udeltas' .- theta)./inv_slope)/sqrt_2
    p = (a.miss/2) + (1-a.miss)*cdf(Normal(),x)
    log_prob = sum(enumerate(udeltas)) do x
      i,delta = x
      logpdf.(Binomial.(a.N[delta],p[:,i]),a.resp[delta])
    end + logpdf(a.theta_prior,theta) + logpdf(a.inv_slope_prior,inv_slope)

    lweights = log_prob - q_log_prob
    weights = exp.(lweights - maximum(lweights))
    weights = weights ./ sum(weights)
    ess = 1 / sum(weights.^2)

    # if ess / length(weights) < 0.1
    #   warn("Effective sampling size of importance samples is low ",
    #        "($(round(ess,1))). Consider adjusting the `thresh_d` and `slope_d`.")
    #   record("poor_adaptor_ess",value=ess)
    # end

    dprime = invlogcdf(Normal(),log((a.threshold-a.miss/2)/(1-a.miss)))
    thresh_off = -log(dprime).*inv_slope
    a.theta_samples = theta
    a.inv_slope_samples = inv_slope
    a.thresh_samples = theta + thresh_off
    a.thresh_weights = weights
  end
  a.thresh_samples, a.thresh_weights
end

function update!(adapter::ImportanceSampler,response,correct)
  adapter.thresh_samples = ones(0)
  adapter.resp[adapter.delta] = get(adapter.resp,adapter.delta,0) +
    (response==correct)
  adapter.N[adapter.delta] = get(adapter.N,adapter.delta,0) + 1

  thresh,weights = sample(adapter)
  μ,σ = estimate(adapter)

  sd_ratio = σ / μ
  if sd_ratio > adapter.repeat3_thresh
    repeats = 2
  elseif sd_ratio > adapter.repeat2_thresh
    repeats = 1
  else
    repeats = 0
  end

  if repeats > adapter.delta_repeat
    adapter.delta_repeat += 1
  else
    adapter.delta_repeat = 0

    # TODO: sample these differences to make sure this is a reasonably
    # sized array
    # TODO: debug selection

    entropies = -sum(1:adapter.n_samples) do i
      diff = adapter.thresh_samples[i] .- adapter.theta_samples
      dprime = exp(diff./adapter.inv_slope_samples)/sqrt_2
      p = ((adapter.miss/2) + (1-adapter.miss)*cdf(Normal(),dprime))
      (p .* log(p) + (1-p) .* log(1-p)) .* weights
    end

    _,i = findmax(entropies)

    adapter.delta = adapter.thresh_samples[i]
  end

  adapter
end

function estimate(adapter::ImportanceSampler)
  thresh,weights = sample(adapter)
  μ = sum(thresh .* weights)
  M = sum(weights .> 0)
  σ = sqrt.(sum(((μ .- thresh) .* weights).^2) ./
           ((M-1)/M)*sum(weights))

  μ,σ
end

delta(adapter::ImportanceSampler) = bound(adapter.delta,adapter.min_delta,
adapter.max_delta)


module MeasurementsWrapper

using JavaCall

const AbsoluteDate = @jimport org.orekit.time.AbsoluteDate
const AbstractMeasurement = @jimport org.orekit.estimation.measurements.AbstractMeasurement
const AngularAzEl = @jimport org.orekit.estimation.measurements.AngularAzEl
const AngularRaDec = @jimport org.orekit.estimation.measurements.AngularRaDec
const Class = @jimport java.lang.Class
const ComparableMeasurement = @jimport org.orekit.estimation.measurements.ComparableMeasurement
const EOPHistory = @jimport org.orekit.frames.EOPHistory
const EstimatedEarthFrameProvider = @jimport org.orekit.estimation.measurements.EstimatedEarthFrameProvider
const EstimatedMeasurement = @jimport org.orekit.estimation.measurements.EstimatedMeasurement
const EstimatedMeasurement_Status = @jimport org.orekit.estimation.measurements.EstimatedMeasurement$Status
const EstimationModifier = @jimport org.orekit.estimation.measurements.EstimationModifier
const EstimationsProvider = @jimport org.orekit.estimation.measurements.EstimationsProvider
const FieldAbsoluteDate = @jimport org.orekit.time.FieldAbsoluteDate
const FieldTransform = @jimport org.orekit.frames.FieldTransform
const FieldVector3D = @jimport org.hipparchus.geometry.euclidean.threed.FieldVector3D
const Frame = @jimport org.orekit.frames.Frame
const GeodeticPoint = @jimport org.orekit.bodies.GeodeticPoint
const GroundStation = @jimport org.orekit.estimation.measurements.GroundStation
const InterSatellitesRange = @jimport org.orekit.estimation.measurements.InterSatellitesRange
const JString = @jimport java.lang.String
const List = @jimport java.util.List
const Map = @jimport java.util.Map
const MultiplexedMeasurement = @jimport org.orekit.estimation.measurements.MultiplexedMeasurement
const Object = @jimport java.lang.Object
const ObservableSatellite = @jimport org.orekit.estimation.measurements.ObservableSatellite
const ObservedMeasurement = @jimport org.orekit.estimation.measurements.ObservedMeasurement
const PV = @jimport org.orekit.estimation.measurements.PV
const ParameterDriver = @jimport org.orekit.utils.ParameterDriver
const Position = @jimport org.orekit.estimation.measurements.Position
const Range = @jimport org.orekit.estimation.measurements.Range
const RangeRate = @jimport org.orekit.estimation.measurements.RangeRate
const RealFieldElement = @jimport org.hipparchus.RealFieldElement
const SpacecraftState = @jimport org.orekit.propagation.SpacecraftState
const StationDisplacement = @jimport org.orekit.models.earth.displacement.StationDisplacement
const Stream = @jimport java.util.stream.Stream
const TimeStamped = @jimport org.orekit.time.TimeStamped
const TimeStampedFieldPVCoordinates = @jimport org.orekit.utils.TimeStampedFieldPVCoordinates
const TimeStampedPVCoordinates = @jimport org.orekit.utils.TimeStampedPVCoordinates
const TopocentricFrame = @jimport org.orekit.frames.TopocentricFrame
const Transform = @jimport org.orekit.frames.Transform
const TurnAroundRange = @jimport org.orekit.estimation.measurements.TurnAroundRange
const UT1Scale = @jimport org.orekit.time.UT1Scale
const Vector3D = @jimport org.hipparchus.geometry.euclidean.threed.Vector3D

export AbstractMeasurement
export AngularAzEl
export AngularRaDec
export ComparableMeasurement
export EstimatedEarthFrameProvider
export EstimatedMeasurement
export EstimationModifier
export EstimationsProvider
export GroundStation
export InterSatellitesRange
export MultiplexedMeasurement
export ObservableSatellite
export ObservedMeasurement
export PV
export Position
export Range
export RangeRate
export TurnAroundRange
export add_modifier
export compare_to
export equals
export estimate
export get_base_frame
export get_base_weight
export get_class
export get_clock_offset_driver
export get_coordinates
export get_correlation_coefficients_matrix
export get_count
export get_covariance_matrix
export get_date
export get_derivatives_drivers
export get_dimension
export get_displacements
export get_east_offset_driver
export get_estimated_earth_frame
export get_estimated_measurement
export get_estimated_measurements
export get_estimated_ut1
export get_estimated_value
export get_iteration
export get_master_station
export get_measurements
export get_modifiers
export get_north_offset_driver
export get_number
export get_observed_measurement
export get_observed_value
export get_offset_geodetic_point
export get_offset_to_inertial
export get_parameter_derivatives
export get_parameters_drivers
export get_participants
export get_polar_drift_x_driver
export get_polar_drift_y_driver
export get_polar_offset_x_driver
export get_polar_offset_y_driver
export get_position
export get_prime_meridian_drift_driver
export get_prime_meridian_offset_driver
export get_propagator_index
export get_reference_frame
export get_satellites
export get_slave_station
export get_state_derivatives
export get_state_size
export get_states
export get_station
export get_status
export get_theoretical_standard_deviation
export get_time_offset
export get_transform
export get_velocity
export get_zenith_offset_driver
export hash_code
export is_enabled
export is_two_way
export modify
export notify
export notify_all
export set_enabled
export set_estimated_value
export set_parameter_derivatives
export set_state_derivatives
export set_status
export signal_time_of_flight
export to_string
export wait

include("FilteringWrapper/FilteringWrapper.jl")
include("GenerationWrapper/GenerationWrapper.jl")
include("GnssWrapper/GnssWrapper.jl")
include("ModifiersWrapper/ModifiersWrapper.jl")
include("abstract_measurement.jl")
include("angular_az_el.jl")
include("angular_ra_dec.jl")
include("comparable_measurement.jl")
include("estimated_earth_frame_provider.jl")
include("estimated_measurement.jl")
include("estimation_modifier.jl")
include("estimations_provider.jl")
include("ground_station.jl")
include("inter_satellites_range.jl")
include("multiplexed_measurement.jl")
include("observable_satellite.jl")
include("observed_measurement.jl")
include("position.jl")
include("pv")
include("range.jl")
include("range_rate.jl")
include("turn_around_range.jl")

end


# Integrations with DFTK.jl

# -----------------------------------------------------------------------------
# Integration with AtomsBase
# -----------------------------------------------------------------------------

# ! all functions in this section are copied directly from experimental DFTK code
# ! this code will be deleted once it is live in DFTK

function parse_system(system::AbstractSystem{D}) where {D}
    if !all(periodicity(system))
        error("DFTK only supports calculations with periodic boundary conditions.")
    end

    # Parse abstract system and return data required to construct model
    mtx = austrip.(hcat(bounding_box(system)...))
    T = eltype(mtx)
    lattice = zeros(T, 3, 3)
    lattice[1:D, 1:D] .= mtx

    # Cache for instantiated pseudopotentials
    # (such that the respective objects are indistinguishable)
    cached_pseudos = Dict{String,Any}()
    atoms = map(system) do atom
        if hasproperty(atom, :potential)
            potential = atom.potential
        elseif hasproperty(atom, :pseudopotential)
            pspkey = atom.pseudopotential
            potential = get!(cached_pseudos, pspkey) do
                ElementPsp(AtomsBase.atomic_symbol(atom); psp = load_psp(pspkey))
            end
        else
            potential = ElementCoulomb(AtomsBase.atomic_symbol(atom))
        end

        coordinate = zeros(T, 3)
        coordinate[1:D] = inv(lattice[1:D, 1:D]) * T.(austrip.(position(atom)))
        potential => Vec3{T}(coordinate)
    end

    oldatoms = oldatoms_from_new(atoms)

    (; lattice, atoms = oldatoms)
end
function oldatoms_from_new(atomic_potentials)
    potentials = first.(atomic_potentials)
    potential_groups = [findall(Ref(pot) .== potentials) for pot in Set(potentials)]
    [first(atomic_potentials[first(group)]) => last.(atomic_potentials[group]) for group in potential_groups]
end
function DFTK.model_LDA(system::AbstractSystem; kwargs...)
    parsed = parse_system(system)
    model_LDA(parsed.lattice, parsed.atoms; kwargs...)
end

# -----------------------------------------------------------------------------
# Integration with InteratomicPotentials
# -----------------------------------------------------------------------------

# ! This integration should ultimately move to the DFTK package itself

@kwdef struct DFTKPotential <: ArbitraryPotential
    Ecut::Real
    kgrid::AbstractVector{<:Integer}
    n_bands::Union{Integer,Nothing} = nothing
    tol::Union{AbstractFloat,Nothing} = nothing
    damping::Union{AbstractFloat,Nothing} = nothing
    mixing::Union{Mixing,Nothing} = nothing
    previous_scfres::Ref{Any} = Ref{Any}()
    potential_energy_cache::Dict{Float64,Float64} = Dict{Float64,Float64}()
end
function DFTKPotential(Ecut::Real, kgrid::AbstractVector{<:Integer}; kwargs...)
    DFTKPotential(; Ecut = Ecut, kgrid = kgrid, kwargs...)
end
function DFTKPotential(Ecut::Unitful.Energy, kgrid::AbstractVector{<:Integer}; kwargs...)
    DFTKPotential(; Ecut = austrip(Ecut), kgrid = kgrid, kwargs...)
end

function calculate_scf(system::AbstractSystem, potential::DFTKPotential)
    model = model_LDA(system)
    basis = PlaneWaveBasis(model; Ecut = potential.Ecut, kgrid = potential.kgrid)

    args = (f => getfield(potential, f) for f ∈ (:n_bands, :tol, :damping, :mixing) if !isnothing(getfield(potential, f)))
    extra_args = isassigned(potential.previous_scfres) ? (ψ = potential.previous_scfres[].ψ, ρ = potential.previous_scfres[].ρ) : (;)
    scfres = self_consistent_field(basis; args..., extra_args...)
    potential.previous_scfres[] = scfres
end

function InteratomicPotentials.potential_energy(system::AbstractSystem, potential::DFTKPotential)
    calculate_scf(system, potential).energies.total
end

function InteratomicPotentials.potential_energy(system::DynamicSystem, potential::DFTKPotential)
    get!(potential.potential_energy_cache, austrip(system.time)) do
        InteratomicPotentials.potential_energy(system.system, potential)
    end
end

function InteratomicPotentials.force(system::AbstractSystem, potential::DFTKPotential)
    # TODO: support multiple species
    compute_forces_cart(calculate_scf(system, potential))[1]
end

function InteratomicPotentials.force(system::DynamicSystem, potential::DFTKPotential)
    scf = calculate_scf(system, potential)
    potential.potential_energy_cache[austrip(system.time)] = scf.energies.total
    # TODO: support multiple species
    compute_forces_cart(scf)[1]
end

# -----------------------------------------------------------------------------
# Experimental functions
# -----------------------------------------------------------------------------

# ! This function is experimental and will eventually be removed

function analyze_convergence(system::AbstractSystem, potential::DFTKPotential, cutoffs::AbstractVector{<:Unitful.Energy})
    energies = zeros(Float64, length(cutoffs))
    for i ∈ 1:length(cutoffs)
        parameters = DFTKPotential(
            Ecut = cutoffs[i],
            kgrid = potential.kgrid,
            n_bands = potential.n_bands,
            tol = potential.tol,
            damping = potential.damping,
            mixing = potential.mixing
        )
        @info "Ecut: $(cutoffs[i])"
        energies[i] = InteratomicPotentials.potential_energy(system, parameters)
    end

    plot(
        title = "DFTK Analysis",
        xlab = "Ecut",
        ylab = "Total Energy",
        legend = false,
        cutoffs,
        energies * ENERGY_UNIT
    )
end


module SemidefiniteOptInterface

using MathOptInterface
const MOI = MathOptInterface
const MOIU = MOI.Utilities

using Compat
using Compat.LinearAlgebra # for diag

abstract type AbstractSDOptimizer <: MOI.AbstractOptimizer end

include("interface.jl")

const SVF = MOI.SingleVariable
const VVF = MOI.VectorOfVariables
const VF  = Union{SVF, VVF}
const SAF{T} = MOI.ScalarAffineFunction{T}
const ASF{T} = Union{SVF, SAF{T}}

const ZS = Union{MOI.EqualTo, MOI.Zeros}
const NS = Union{MOI.GreaterThan, MOI.Nonnegatives}
const PS = Union{MOI.LessThan, MOI.Nonpositives}
const DS = MOI.PositiveSemidefiniteConeTriangle
const SupportedSets = Union{ZS, NS, PS, DS}

const VI = MOI.VariableIndex
const CI{F, S} = MOI.ConstraintIndex{F, S}

mutable struct SOItoMOIBridge{T, SIT <: AbstractSDOptimizer} <: MOI.AbstractOptimizer
    sdoptimizer::SIT
    setconstant::Dict{Int64, T}
    blkconstant::Dict{Int, T}
    objconstant::T
    objsign::Int
    objshift::T
    nconstrs::Int
    nblocks::Int
    blockdims::Vector{Int}
    free::BitSet
    varmap::Vector{Vector{Tuple{Int, Int, Int, T, T}}} # Variable Index vi -> blk, i, j, coef, shift # x = sum coef * block(X, blk)[i, j] + shift
    zeroblock::Dict{CI, Int}
    constrmap::Dict{CI, UnitRange{Int}} # Constraint Index ci -> cs
    double::Vector{CI} # created when there are two cones for same variable
    function SOItoMOIBridge{T}(sdoptimizer::SIT) where {T, SIT}
        new{T, SIT}(sdoptimizer, Dict{Int64, T}(), Dict{Int, T}(),
            zero(T), 1, zero(T), 0, 0,
            Int[],
            BitSet(),
            Vector{Tuple{Int, Int, Int, T}}[],
            Dict{CI, Int}(),
            Dict{CI, UnitRange{Int}}(),
            CI[])
    end
end
varmap(optimizer::SOItoMOIBridge, vi::VI) = optimizer.varmap[vi.value]
function setvarmap!(optimizer::SOItoMOIBridge{T}, vi::VI, v::Tuple{Int, Int, Int, T, T}) where T
    setvarmap!(optimizer, vi, [v])
end
function setvarmap!(optimizer::SOItoMOIBridge{T}, vi::VI, vs::Vector{Tuple{Int, Int, Int, T, T}}) where T
    optimizer.varmap[vi.value] = vs
end

SDOIOptimizer(sdoptimizer::AbstractSDOptimizer, T=Float64) = SOItoMOIBridge{T}(sdoptimizer)

include("load.jl")

function MOI.get(optimizer::SOItoMOIBridge, attr::MOI.SolverName)
    return MOI.get(optimizer.sdoptimizer, attr)
end

function MOI.is_empty(optimizer::SOItoMOIBridge)
    isempty(optimizer.double) &&
    isempty(optimizer.setconstant) &&
    isempty(optimizer.blkconstant) &&
    iszero(optimizer.objconstant) &&
    optimizer.objsign == 1 &&
    iszero(optimizer.objshift) &&
    iszero(optimizer.nconstrs) &&
    iszero(optimizer.nblocks) &&
    isempty(optimizer.blockdims) &&
    isempty(optimizer.free) &&
    isempty(optimizer.varmap) &&
    isempty(optimizer.zeroblock) &&
    isempty(optimizer.constrmap)
end
function MOI.empty!(optimizer::SOItoMOIBridge{T}) where T
    for s in optimizer.double
        MOI.delete(m, s)
    end
    MOI.empty!(optimizer.sdoptimizer)
    optimizer.double = CI[]
    optimizer.setconstant = Dict{Int64, T}()
    optimizer.blkconstant = Dict{Int, T}()
    optimizer.objconstant = zero(T)
    optimizer.objsign = 1
    optimizer.objshift = zero(T)
    optimizer.nconstrs = 0
    optimizer.nblocks = 0
    optimizer.blockdims = Int[]
    optimizer.free = BitSet()
    optimizer.varmap = Vector{Tuple{Int, Int, Int, T}}[]
    optimizer.zeroblock = Dict{CI, Int}()
    optimizer.constrmap = Dict{CI, UnitRange{Int}}()
end

function setconstant!(optimizer::SOItoMOIBridge, ci::CI, s) end
function setconstant!(optimizer::SOItoMOIBridge, ci::CI, s::MOI.AbstractScalarSet)
    optimizer.setconstant[ci.value] = MOIU.getconstant(s)
end
function set_constant(optimizer::SOItoMOIBridge,
                      ci::CI{<:MOI.AbstractScalarFunction,
                             <:MOI.AbstractScalarSet})
    return optimizer.setconstant[ci.value]
end
function set_constant(optimizer::SOItoMOIBridge{T}, ci::CI) where T
    return zeros(T, length(optimizer.constrmap[ci]))
end
function addblkconstant(optimizer::SOItoMOIBridge, ci::CI{<:Any, <:Union{NS, PS}}, x)
    blk = -ci.value
    return x .+ optimizer.blkconstant[blk]
end
addblkconstant(optimizer::SOItoMOIBridge, ci::CI, x) = x

function MOI.supports(
    optimizer::SOItoMOIBridge{T},
    ::Union{MOI.ObjectiveSense,
            MOI.ObjectiveFunction{<:Union{MOI.SingleVariable,
                                          MOI.ScalarAffineFunction{T}}}}) where T
    return true
end

# Zeros and Nonpositives supports could be removed thanks to variable bridges
# * `VectorOfVariables`-in-`Zeros` would return a `VectorAffineFunction` with
#   zero constant and no variable created.
# * `VectorOfVariables`-in-`Nonpositives` would create variables in
#   `Nonnegatives` and return a `VectorAffineFunction` containing `-` the
#    variables.
function MOI.supports_constraint(
    ::SOItoMOIBridge, ::Type{MOI.VectorOfVariables},
    ::Type{<:Union{MOI.Zeros, MOI.Nonnegatives, MOI.Nonpositives,
                   MOI.PositiveSemidefiniteConeTriangle}})
    return true
end
# This support could be remove thanks to variable bridges.
# The VectorizeVariableBridge would redirect to the above case and then the
# resulting function would be shifted by the constant.
function MOI.supports_constraint(
    ::SOItoMOIBridge{T}, ::Type{MOI.SingleVariable},
    ::Type{<:Union{MOI.EqualTo{T}, MOI.GreaterThan{T}, MOI.LessThan{T}}}) where T
    return true
end
function MOI.supports_constraint(
    ::SOItoMOIBridge{T}, ::Type{MOI.ScalarAffineFunction{T}},
    ::Type{MOI.EqualTo{T}}) where T
    return true
end

function MOI.copy_to(dest::SOItoMOIBridge, src::MOI.ModelLike; kws...)
    return MOIU.automatic_copy_to(dest, src; kws...)
end
MOIU.supports_allocate_load(::SOItoMOIBridge, copy_names::Bool) = !copy_names

MOI.optimize!(m::SOItoMOIBridge) = MOI.optimize!(m.sdoptimizer)

# Objective

function MOI.get(m::SOItoMOIBridge, ::MOI.ObjectiveValue)
    m.objshift + m.objsign * getprimalobjectivevalue(m.sdoptimizer) + m.objconstant
end

# Attributes

const SolverStatus = Union{MOI.TerminationStatus, MOI.PrimalStatus, MOI.DualStatus}
MOI.get(m::SOItoMOIBridge, s::SolverStatus) = MOI.get(m.sdoptimizer, s)

MOI.get(m::SOItoMOIBridge, ::MOI.ResultCount) = 1

function _getblock(M, blk::Integer, s::Type{<:Union{NS, ZS}})
    return diag(block(M, blk))
end
function _getblock(M, blk::Integer, s::Type{<:PS})
    return -diag(block(M, blk))
end
# Vectorized length for matrix dimension d
sympackedlen(d::Integer) = (d*(d+1)) >> 1
function _getblock(M::AbstractMatrix{T}, blk::Integer, s::Type{<:DS}) where T
    B = block(M, blk)
    d = Compat.LinearAlgebra.checksquare(B)
    n = sympackedlen(d)
    v = Vector{T}(undef, n)
    k = 0
    for j in 1:d
        for i in 1:j
            k += 1
            v[k] = B[i, j]
        end
    end
    @assert k == n
    return v
end
function getblock(M, blk::Integer, s::Type{<:MOI.AbstractScalarSet})
    vd = _getblock(M, blk, s)
    @assert length(vd) == 1
    return vd[1]
end
function getblock(M, blk::Integer, s::Type{<:MOI.AbstractVectorSet})
    return _getblock(M, blk, s)
end

getvarprimal(m::SOItoMOIBridge, blk::Integer, S) = getblock(getX(m.sdoptimizer), blk, S)
function getvardual(m::SOItoMOIBridge, blk::Integer, S)
    z = getZ(m.sdoptimizer)
    b = getblock(z, blk, S)
    return getblock(getZ(m.sdoptimizer), blk, S)
end

function MOI.get(m::SOItoMOIBridge{T}, ::MOI.VariablePrimal, vi::VI) where T
    X = getX(m.sdoptimizer)
    x = zero(T)
    for (blk, i, j, coef, shift) in varmap(m, vi)
        x += shift
        if blk != 0
            x += block(X, blk)[i, j] * sign(coef)
        end
    end
    return x
end
function MOI.get(m::SOItoMOIBridge, vp::MOI.VariablePrimal, vi::Vector{VI})
    return MOI.get.(m, vp, vi)
end

function _getattribute(m::SOItoMOIBridge, ci::CI{<:ASF}, f)
    cs = m.constrmap[ci]
    @assert length(cs) == 1
    return f(m, first(cs))
end
function _getattribute(m::SOItoMOIBridge, ci::CI{<:VVF}, f)
    return f.(m, m.constrmap[ci])
end

function MOI.get(m::SOItoMOIBridge, a::MOI.ConstraintPrimal,
                 ci::CI{F, S}) where {F, S}
    if ci.value >= 0
        return set_constant(m, ci)
    else
        # Variable Function-in-S with S different from Zeros and EqualTo and not a double variable constraint
        blk = -ci.value
        return addblkconstant(m, ci, getvarprimal(m, blk, S))
    end
end

function MOI.get(m::SOItoMOIBridge, ::MOI.ConstraintDual, ci::CI{<:VF, S}) where S<:SupportedSets
    if ci.value < 0
        return getvardual(m, -ci.value, S)
    else
        dual = _getattribute(m, ci, getdual)
        if haskey(m.zeroblock, ci) # ZS
            return dual + getvardual(m, m.zeroblock[ci], S)
        else # var constraint on unfree constraint
            return dual
        end
    end
end

function getdual(m::SOItoMOIBridge{T}, c::Integer) where T
    if c == 0
        return zero(T)
    else
        return -gety(m.sdoptimizer)[c]
    end
end
function MOI.get(m::SOItoMOIBridge, ::MOI.ConstraintDual, ci::CI)
    return _getattribute(m, ci, getdual)
end

include("sdpa.jl")
include("mock.jl")

end # module


using Raylib
using Raylib: RayCamera3D, rayvector, RayVector3, RayColor, raycolor,
    RAYWHITE, LIGHTGRAY, DARKGRAY, MAROON,
    RED, BLUE, LIME, GOLD, SKYBLUE, BLACK

function main()
    screenWidth = 800
    screenHeight = 450

    Raylib.InitWindow(
        screenWidth, screenHeight,
        "raylib [core] example - 3d camera first person"
    )

    camera = RayCamera3D(
        rayvector(4, 2, 4),
        rayvector(0, 1.8, 0),
        rayvector(0, 1, 0),
        60,
        0,
    )

    heights = Vector{Cfloat}(undef, 20)
    positions = Vector{RayVector3}(undef, 20)
    colors = Vector{RayColor}(undef, 20)

    for i = 1:20
        heights[i] = rand(1:12)
        positions[i] = rayvector(
            rand(-15:15),
            heights[i] / 2,
            rand(-15:15),
        )
        colors[i] = raycolor(
            rand(20:255),
            rand(10:55),
            30,
            255,
        )
    end

    Raylib.SetCameraMode(camera, Int(Raylib.CAMERA_FIRST_PERSON))
    Raylib.SetTargetFPS(60)

    while !Raylib.WindowShouldClose()
        Raylib.UpdateCamera!(camera)

        Raylib.BeginDrawing()
        Raylib.ClearBackground(RAYWHITE)

        Raylib.BeginMode3D(camera)

        Raylib.DrawPlane(
            rayvector(0,0,0),
            rayvector(32, 32),
            LIGHTGRAY)

        Raylib.DrawCube(
            rayvector(-16, 2.5, 0), 1f0, 5f0, 32f0,
            BLUE)
        Raylib.DrawCube(
            rayvector(16, 2.5, 0), 1f0, 5f0, 32f0,
            LIME)
        Raylib.DrawCube(
            rayvector(0, 2.5, 16), 32f0, 5f0, 1f0,
            GOLD)

        for i = 1:20
            Raylib.DrawCube(
                positions[i], 2f0, heights[i], 2f0,
                colors[i])
            Raylib.DrawCubeWires(positions[i], 2f0, heights[i], 2f0,
                                 MAROON)

        end

        Raylib.EndMode3D()

        Raylib.DrawRectangle(10, 10, 220, 70,
                             Raylib.Fade(SKYBLUE, 0.5) )
        Raylib.DrawRectangleLines(10, 10, 220, 70, BLUE)

        Raylib.DrawText("First person camera default controls:",
                        20, 20, 10, BLACK)
        Raylib.DrawText("- Move with keys: W, A, S, D",
                        40, 40, 10, DARKGRAY)
        Raylib.DrawText("- Mouse move to look around",
                        40, 60, 10, DARKGRAY)

        Raylib.EndDrawing()
    end

    Raylib.CloseWindow()
end