text:
module CSetDataStructures
export @acset_type, @abstract_acset_type, @declare_schema, FreeSchema,
  StructACSet, StructCSet, ACSetTableType

using MLStyle
using StaticArrays
using Reexport
import Tables

@reexport using ..ACSetInterface
using ..IndexUtils
using ...Theories, ...Present, ...Syntax
using ...Theories: FreeSchema, SchemaDesc, SchemaDescType, CSetSchemaDescType,
  SchemaDescTypeType, ob_num, codom_num, attr, attrtype
import ...Present: Presentation

# StructACSet Struct Generation
###############################

abstract type StructACSet{S<:SchemaDescType,Ts<:Tuple,Idxed,UniqueIdxed} <: ACSet end

const StructCSet = StructACSet{S,Tuple{},Idxed,UniqueIdxed} where
  {S<:CSetSchemaDescType,Idxed,UniqueIdxed}

q(s::Symbol) = Expr(:quote,s)
q(s::GATExpr) = q(nameof(s))

""" Creates a quoted named tuple used for `StructACSet`s
"""
function pi_type(dom::Vector, F::Function)
  :(NamedTuple{($(q.(dom)...),),Tuple{$(F.(dom)...)}})
end

""" Creates a quoted element of a named tuple used for `StructACSet`s
"""
function pi_type_elt(dom::Vector, f::Function)
  if length(dom) > 0
    Expr(:tuple, Expr(:parameters, [Expr(:kw, nameof(d), f(d)) for d in dom]...))
  else # workaround for Julia 1.0
    :(NamedTuple())
  end
end

""" Create the struct declaration for a `StructACSet` from a Presentation
"""
function struct_acset(name::Symbol, parent, p::Presentation{Schema}; index=[], unique_index=[])
  obs = p.generators[:Ob]
  homs = p.generators[:Hom]
  attr_types = p.generators[:AttrType]
  Ts = nameof.(attr_types)
  attrs = p.generators[:Attr]
  idxed = (;[x => x ∈ index for x in [nameof.(homs);nameof.(attrs)]]...)
  unique_idxed = (;[x => x ∈ unique_index for x in [nameof.(homs);nameof.(attrs)]]...)
  indexed_homs = filter(f -> idxed[nameof(f)], homs)
  unique_indexed_homs = filter(f -> unique_idxed[nameof(f)], homs)
  indexed_attrs = filter(a -> idxed[nameof(a)], attrs)
  unique_indexed_attrs = filter(a -> unique_idxed[nameof(a)], attrs)
  parameterized_type, new_call = if length(attr_types) > 0
    (:($name{$(nameof.(attr_types)...)}), :(new{$(nameof.(attr_types)...)}))
  else
    name, :new
  end
  schema_type = SchemaDescTypeType(p)
  obs_t = :($(GlobalRef(StaticArrays, :MVector)){$(length(obs)), Int})
  quote
    struct $parameterized_type <: $parent{$schema_type, Tuple{$(Ts...)}, $idxed, $unique_idxed}
      obs::$obs_t
      homs::$(pi_type(homs, _ -> :(Vector{Int})))
      attrs::$(pi_type(attrs, a -> :(Vector{$(nameof(codom(a)))})))
      hom_indices::$(pi_type(indexed_homs, _ -> :(Vector{Vector{Int}})))
      hom_unique_indices::$(pi_type(unique_indexed_homs, _ -> :(Vector{Int})))
      attr_indices::$(pi_type(indexed_attrs, a -> :(Dict{$(nameof(codom(a))),Vector{Int}})))
      attr_unique_indices::$(pi_type(unique_indexed_attrs, a -> :(Dict{$(nameof(codom(a))), Int})))
      function $parameterized_type() where {$(nameof.(attr_types)...)}
        $new_call(
          $obs_t(zeros(Int, $(length(obs)))),
          $(pi_type_elt(homs, _ -> :(Int[]))),
          $(pi_type_elt(attrs, a -> :($(nameof(codom(a)))[]))),
          $(pi_type_elt(indexed_homs, _ -> :(Vector{Int}[]))),
          $(pi_type_elt(unique_indexed_homs, _ -> Int[])),
          $(pi_type_elt(indexed_attrs, a -> :(Dict{$(nameof(codom(a))),Vector{Int}}()))),
          $(pi_type_elt(unique_indexed_attrs, a -> :(Dict{$(nameof(codom(a))),Int}())))
        )
      end
    end
  end
end

unquote(x::QuoteNode) = x.value

macro acset_type(head)
  head, parent = @match head begin
    Expr(:(<:), h, p) => (h,p)
    _ => (head, GlobalRef(CSetDataStructures, :StructACSet))
  end
  name, schema, idx_args = @match head begin
    Expr(:call, name, schema, idx_args...) => (name, schema, idx_args)
    _ => error("Unsupported head for @acset_type")
  end
  abstract_name = Symbol("Abstract" * string(name))
  # We assign this to an anonymous variable so that there is something
  # to attach documentation to.
  # TODO: this is a hack and there should be a better way
  quote
    tmp = $(esc(:eval))($(GlobalRef(CSetDataStructures, :struct_acset))(
      $(Expr(:quote, name)), $(Expr(:quote, parent)), $(esc(schema)), $(idx_args...)))
  end
end

macro abstract_acset_type(head)
  type, parent = @match head begin
    Expr(:(<:), h, p) => (h,p)
    _ => (head, GlobalRef(CSetDataStructures, :StructACSet))
  end
  esc(quote
    abstract type $type{S,Ts,Idxed,UniqueIdxed} <: $parent{S,Ts,Idxed,UniqueIdxed} end
  end)
end

# StructACSet DataFrames
########################

struct StructACSetFrame{Ob, S, Ts, Idxed, UniqueIdxed, Attrs} <:
    StructACSet{S, Ts, Idxed, UniqueIdxed}
  obs::MVector{1,Int}
  homs::NamedTuple{(),Tuple{}}
  attrs::Attrs
  hom_indices::NamedTuple{(),Tuple{}}
  hom_unique_indices::NamedTuple{(),Tuple{}}
  attr_indices::NamedTuple{(),Tuple{}}
  attr_unique_indices::NamedTuple{(),Tuple{}}
  function StructACSetFrame{Ob, S, Ts, Idxed, UniqueIdxed, Attrs}() where
      {Ob,S,Ts,Idxed,UniqueIdxed,Attrs}
    new{Ob,S,Ts,Idxed,UniqueIdxed,Attrs}(
      MVector(0),
      (;),
      Attrs([[] for a in attr(S)]),
      (;),
      (;),
      (;),
      (;)
    )
  end
end

function ACSetTableDesc(::Type{S}, ob::Symbol) where {S}
  s = SchemaDesc(S)
  attrs = filter(s.attrs) do attr
    s.doms[attr] == ob
  end
  doms = filter(s.doms) do (f,i)
    f ∈ attrs
  end
  codoms = filter(s.codoms) do (f,i)
    f ∈ attrs
  end
  s′ = SchemaDesc([ob], Symbol[], s.attrtypes, attrs, doms, codoms)
  SchemaDescTypeType(s′)
end

function ACSetTableDataType(::Type{<:StructACSet{S,Ts}}, ob::Symbol) where {S,Ts}
  S′ = ACSetTableDesc(S,ob)
  rowtype = NamedTuple{
    attr(S′),
    Tuple{[Vector{Ts.parameters[codom_num(S′,a)]} for a in attr(S′)]...}
  }
  StructACSetFrame{(ob,), S′, Ts,
                   NamedTuple(a => false for a in attr(S′)),
                   NamedTuple(a => false for a in attr(S′)),
                   rowtype}
end

function ACSetTableUnionAll(::Type{<:StructACSet{S}}, ob::Symbol) where {S}
  S′ = ACSetTableDesc(S,ob)
  type_vars = [ TypeVar(at) for at in attrtype(S′) ]
  rowtype = NamedTuple{
    attr(S′),
    Tuple{[Vector{type_vars[codom_num(S′,a)]} for a in attr(S′)]...}
  }
  T = StructACSetFrame{(ob,), S′, Tuple{type_vars...},
                       NamedTuple(a => false for a in attr(S′)),
                       NamedTuple(a => false for a in attr(S′)),
                       rowtype}
  foldr(UnionAll, type_vars, init=T)
end

function ACSetTableType(X::Type, ob::Symbol; union_all::Bool=false)
  (union_all ? ACSetTableUnionAll : ACSetTableDataType)(X, ob)
end

# StructACSet Operations
########################

Presentation(::Type{<:StructACSet{S}}) where S = Presentation(S)
Presentation(::StructACSet{S}) where S = Presentation(S)

# Accessors
###########

function Base.:(==)(x1::T, x2::T) where T <: StructACSet
  # The indices hold redundant information, so need not be compared.
  unref(x) = x[]
  unref.(values(x1.obs)) == unref.(values(x2.obs)) && x1.homs == x2.homs && x1.attrs == x2.attrs
end

ACSetInterface.acset_schema(::StructACSet{S}) where {S} = SchemaDesc(S)

@inline ACSetInterface.nparts(acs::StructACSet, ob::Symbol) = _nparts(acs, Val{ob})

@generated function _nparts(acs::StructACSet{S}, ::Type{Val{ob}}) where {S,ob}
  :(acs.obs[$(ob_num(S,ob))])
end

@inline ACSetInterface.has_part(acs::StructACSet, ob::Symbol) = _has_part(acs, Val{ob})

outgoing(acs::StructACSet, ob::Symbol) = _outgoing(acs, Val{ob})

@generated function _outgoing(::StructACSet{S}, ::Type{Val{ob}}) where {S,ob}
  s = SchemaDesc(S)
  Tuple(filter(f -> s.doms[f] == ob, vcat(s.homs,s.attrs)))
end

@generated function _has_part(acs::StructACSet{S}, ::Type{Val{ob}}) where
  {obs, S <: SchemaDescType{obs}, ob}
  ob ∈ obs
end

@inline ACSetInterface.has_subpart(acs::StructACSet, f::Symbol) = _has_subpart(acs, Val{f})

@generated function _has_subpart(acs::StructACSet{S}, ::Type{Val{f}}) where
  {f, obs, homs, attrtypes, attrs, S <: SchemaDescType{obs,homs,attrtypes,attrs}}
  f ∈ homs || f ∈ attrs
end

@inline ACSetInterface.subpart(acs::StructACSet, f::Symbol) = _subpart(acs,Val{f})

@generated function _subpart(acs::StructACSet{S}, ::Type{Val{f}}) where {S,f}
  s = SchemaDesc(S)
  if f ∈ s.homs
    :(acs.homs.$f)
  elseif f ∈ s.attrs
    :(acs.attrs.$f)
  else
    throw(ArgumentError("$(repr(f)) not in $(s.homs) or $(s.attrs)"))
  end
end

@inline Base.getindex(acs::StructACSet, args...) = ACSetInterface.subpart(acs, args...)

@inline ACSetInterface.incident(acs::StructACSet, part, f::Symbol; copy::Bool=false) =
  _incident(acs, part, Val{f}; copy=copy)

broadcast_findall(xs, array::AbstractArray) =
  broadcast(x -> findall(y -> x == y, array), xs)

function get_attr_index(idx::AbstractDict, k)
  get(idx, k, [])
end
function get_attr_index(idx::AbstractDict, k::AbstractArray)
  get_attr_index.(Ref(idx), k)
end

"""
We keep the main body of the code generating out of the @generated function
so that the code-generating function only needs to be compiled once.
"""
function incident_body(s::SchemaDesc, idxed::AbstractDict{Symbol,Bool},
                       unique_idxed::AbstractDict{Symbol,Bool}, f::Symbol)
  if f ∈ s.homs
    if idxed[f]
      quote
        indices = $(GlobalRef(ACSetInterface,:view_or_slice))(acs.hom_indices.$f, part)
        copy ? Base.copy.(indices) : indices
      end
    elseif unique_idxed[f]
      quote
        indices = $(GlobalRef(ACSetInterface,:view_or_slice))(acs.hom_unique_indices.$f, part)
        copy ? Base.copy.(indices) : indices
      end
    else
      :(broadcast_findall(part, acs.homs.$f))
    end
  elseif f ∈ s.attrs
    if idxed[f]
      quote
        indices = get_attr_index(acs.attr_indices.$f, part)
        copy ? Base.copy.(indices) : indices
      end
    elseif unique_idxed[f]
      quote
        get.(Ref(acs.attr_unique_indices.$f), part, Ref(0))
      end
    else
      :(broadcast_findall(part, acs.attrs.$f))
    end
  else
    throw(ArgumentError("$(repr(f)) not in $(s.homs)"))
  end
end

@generated function _incident(acs::StructACSet{S,Ts,Idxed,UniqueIdxed},
                              part, ::Type{Val{f}}; copy::Bool=false) where
  {S,Ts,Idxed,UniqueIdxed,f}
  incident_body(SchemaDesc(S),pairs(Idxed),pairs(UniqueIdxed),f)
end

# Mutators
##########

@inline ACSetInterface.add_parts!(acs::StructACSet, ob::Symbol, n::Int) = _add_parts!(acs, Val{ob}, n)

function add_parts_body(s::SchemaDesc, idxed::AbstractDict,
                        unique_idxed::AbstractDict, ob::Symbol)
  code = quote
    m = acs.obs[$(ob_num(s, ob))]
    nparts = m + n
    newparts = (m+1):m+n
    acs.obs[$(ob_num(s, ob))] = nparts
  end
  for f in s.homs
    if s.doms[f] == ob
      push!(code.args, quote
              resize!(acs.homs.$f, nparts)
              acs.homs.$f[newparts] .= 0
            end)
    end
    if s.codoms[f] == ob && idxed[f]
      push!(code.args, quote
            resize!(acs.hom_indices.$f, nparts)
            for i in newparts
              acs.hom_indices.$f[i] = Int[]
            end
            end)
    elseif s.codoms[f] == ob && unique_idxed[f]
      push!(code.args, quote
            resize!(acs.hom_unique_indices.$f, nparts)
            for i in newparts
              acs.hom_unique_indices.$f[i] = 0
            end
            end)
    end
    code
  end
  for a in s.attrs
    if s.doms[a] == ob
      push!(code.args,:(resize!(acs.attrs.$a, nparts)))
    end
  end
  push!(code.args, :(newparts))
  code
end

""" This generates the _add_parts! methods for a specific object of a `StructACSet`.
"""
@generated function _add_parts!(acs::StructACSet{S,Ts,Idxed,UniqueIdxed},
                                ::Type{Val{ob}}, n::Int) where
  {S, Ts, Idxed, UniqueIdxed, ob}
  add_parts_body(SchemaDesc(S),pairs(Idxed),pairs(UniqueIdxed),ob)
end

@inline ACSetInterface.set_subpart!(acs::StructACSet, part::Int, f::Symbol, subpart) =
  _set_subpart!(acs, part, Val{f}, subpart)

function set_subpart_body(s::SchemaDesc, idxed::AbstractDict{Symbol,Bool},
                          unique_idxed::AbstractDict{Symbol,Bool}, f::Symbol)
  if f ∈ s.homs
    if idxed[f]
      quote
        @assert 0 <= subpart <= acs.obs[$(ob_num(s, s.codoms[f]))]
        old = acs.homs.$f[part]
        acs.homs.$f[part] = subpart
        if old > 0
          @assert deletesorted!(acs.hom_indices.$f[old], part)
        end
        if subpart > 0
          insertsorted!(acs.hom_indices.$f[subpart], part)
        end
      end
    elseif unique_idxed[f]
      quote
        @assert 0 <= subpart <= acs.obs[$(ob_num(s, s.codoms[f]))]
        @assert acs.hom_unique_indices.$f[subpart] == 0
        old = acs.homs.$f[part]
        if old > 0
          acs.hom_unique_indices.$f[old] = 0
        end
        acs.homs.$f[part] = subpart
        acs.hom_unique_indices.$f[subpart] = part
      end
    else
      quote
        @assert 0 <= subpart <= acs.obs[$(ob_num(s, s.codoms[f]))]
        acs.homs.$f[part] = subpart
      end
    end
  elseif f ∈ s.attrs
    if idxed[f]
      quote
        if isassigned(acs.attrs.$f, part)
          old = acs.attrs.$f[part]
          unset_attr_index!(acs.attr_indices.$f, old, part)
        end
        acs.attrs.$f[part] = subpart
        set_attr_index!(acs.attr_indices.$f, subpart, part)
      end
    elseif unique_idxed[f]
      quote
        @assert subpart ∉ keys(acs.attr_unique_indices.$f) "subpart not unique"
        if isassigned(acs.attrs.$f, part)
          old = acs.attrs.$f[part]
          delete!(acs.attr_unique_indices.$f, old)
        end
        acs.attrs.$f[part] = subpart
        acs.attr_unique_indices.$f[subpart] = part
      end
    else
      :(acs.attrs.$f[part] = subpart)
    end
  else
    throw(ArgumentError("$(f) not in $(s.homs) or $(s.attrs)"))
  end
end

""" This generates the `_set_subparts!` method for a specific arrow (hom/attr) of a StructACSet
"""
@generated function _set_subpart!(acs::StructACSet{S,Ts,Idxed,UniqueIdxed},
                                  part, ::Type{Val{f}}, subpart) where {S,Ts,Idxed,UniqueIdxed,f}
  set_subpart_body(SchemaDesc(S),pairs(Idxed),pairs(UniqueIdxed),f)
end

@inline Base.setindex!(acs::StructACSet, val, ob, part) = set_subpart!(acs, ob, part, val)

@inline Base.setindex!(acs::StructACSet, val, ob) = set_subpart!(acs, ob, val)

@inline ACSetInterface.rem_part!(acs::StructACSet, type::Symbol, part::Int) =
  _rem_part!(acs, Val{type}, part)

function getassigned(acs::StructACSet, arrows, i)
  assigned_subparts = filter(f -> isassigned(subpart(acs,f),i), arrows)
  Dict(f => subpart(acs,i,f) for f in assigned_subparts)
end

function rem_part_body(s::SchemaDesc, idxed, ob::Symbol)
  in_homs = filter(hom -> s.codoms[hom] == ob, s.homs)
  out_homs = filter(f -> s.doms[f] == ob, s.homs)
  out_attrs = filter(f -> s.doms[f] == ob, s.attrs)
  indexed_out_homs = filter(hom -> s.doms[hom] == ob && idxed[hom], s.homs)
  indexed_attrs = filter(attr -> s.doms[attr] == ob && idxed[attr], s.attrs)
  quote
    last_part = acs.obs[$(ob_num(s, ob))]
    @assert 1 <= part <= last_part
    # Unassign superparts of the part to be removed and also reassign superparts
    # of the last part to this part.
    for hom in $(Tuple(in_homs))
      set_subpart!(acs, incident(acs, part, hom, copy=true), hom, 0)
      set_subpart!(acs, incident(acs, last_part, hom, copy=true), hom, part)
    end
    last_row = getassigned(acs, $([out_homs;out_attrs]), last_part)

    # Clear any morphism and data attribute indices for last part.
    for hom in $(Tuple(indexed_out_homs))
      set_subpart!(acs, last_part, hom, 0)
    end
    for attr in $(Tuple(indexed_attrs))
      if haskey(last_row, attr)
        unset_attr_index!(acs.attr_indices[attr], last_row[attr], last_part)
      end
    end

    # Finally, delete the last part and update subparts of the removed part.
    for f in $(Tuple(out_homs))
      resize!(acs.homs[f], last_part - 1)
    end
    for a in $(Tuple(out_attrs))
      resize!(acs.attrs[a], last_part - 1)
    end
    acs.obs[$(ob_num(s, ob))] -= 1
    if part < last_part
      set_subparts!(acs, part, (;last_row...))
    end
  end
end

@generated function _rem_part!(acs::StructACSet{S,Ts,idxed}, ::Type{Val{ob}},
                               part::Int) where {S,Ts,ob,idxed}
  rem_part_body(SchemaDesc(S),pairs(idxed),ob)
end

function Base.copy(acs::StructACSet)
  new_acs = typeof(acs)()
  ACSetInterface.copy_parts!(new_acs, acs)
  new_acs
end

""" Copy parts from a C-set to a C′-set.

The selected parts must belong to both schemas. All subparts common to the
selected parts, including data attributes, are preserved. Thus, if the selected
parts form a sub-C-set, then the whole sub-C-set is preserved. On the other
hand, if the selected parts do *not* form a sub-C-set, then some copied parts
will have undefined subparts.
"""
@generated function ACSetInterface.copy_parts!(to::StructACSet{S},
                                from::StructACSet{S′}; kw...) where {S, S′}
  s, s′ = SchemaDesc(S), SchemaDesc(S′)
  obs = intersect(s.obs, s′.obs)
  :(copy_parts!(to, from, isempty(kw) ? $(Tuple(obs)) : (; kw...)))
end

ACSetInterface.copy_parts!(to::StructACSet, from::StructACSet, obs::Tuple) =
  copy_parts!(to, from, NamedTuple{obs}((:) for ob in obs))
ACSetInterface.copy_parts!(to::StructACSet, from::StructACSet, parts::NamedTuple) =
  _copy_parts!(to, from, replace_colons(from, parts))

@generated function _copy_parts!(to::StructACSet{S}, from::StructACSet{S′},
                                 parts::NamedTuple{obs}) where {S, S′, obs}
  s, s′ = SchemaDesc(S), SchemaDesc(S′)
  @assert obs ⊆ intersect(s.obs, s′.obs)
  homs = intersect(s.homs, s′.homs)
  homs = filter(homs) do hom
    c, c′, d, d′ = s.doms[hom], s′.doms[hom], s.codoms[hom], s′.codoms[hom]
    c == c′ && d == d′ && c ∈ obs && d ∈ obs
  end
  hom_triples = [ (hom, s.doms[hom], s.codoms[hom]) for hom in homs ]
  in_obs = unique!(map(last, hom_triples))
  quote
    newparts = _copy_parts_only!(to, from, parts)
    partmaps = NamedTuple{$(Tuple(in_obs))}(tuple($(map(in_obs) do type
      :(Dict{Int,Int}(zip(parts.$type, newparts.$type)))
    end...)))
    for (name, dom, codom) in $(Tuple(hom_triples))
      for (p, newp) in zip(parts[dom], newparts[dom])
        q = subpart(from, p, name)
        newq = get(partmaps[codom], q, nothing)
        if !isnothing(newq)
          set_subpart!(to, newp, name, newq)
        end
      end
    end
    newparts
  end
end

""" Copy parts from a C-set to a C′-set, ignoring all non-data subparts.

The selected parts must belong to both schemas. Attributes common to both
schemas are also copied, but no other subparts are copied.

See also: [`copy_parts!`](@ref).
"""
@generated function copy_parts_only!(to::StructACSet{S}, from::StructACSet{S′}; kw...) where
  {S,S′}
  s, s′ = SchemaDesc(S), SchemaDesc(S′)
  obs = intersect(s.obs, s′.obs)
  :(copy_parts_only!(to, from, isempty(kw) ? $(Tuple(obs)) : (; kw...)))
end

copy_parts_only!(to::StructACSet, from::StructACSet, obs::Tuple) =
  copy_parts_only!(to, from, NamedTuple{obs}((:) for ob in obs))
copy_parts_only!(to::StructACSet, from::StructACSet, parts::NamedTuple) =
  _copy_parts_only!(to, from, replace_colons(from, parts))

@generated function _copy_parts_only!(to::StructACSet{S}, from::StructACSet{S′},
    parts::NamedTuple{obs}) where {S, S′, obs}
  s, s′ = SchemaDesc(S), SchemaDesc(S′)
  @assert obs ⊆ intersect(s.obs, s′.obs)
  attrs = intersect(s.attrs, s′.attrs)
  attrs = filter(attrs) do attr
    ob, ob′ = s.doms[attr], s′.doms[attr]
    ob == ob′ && ob ∈ obs
  end
  Expr(:block,
    :(newparts = (; $(map(obs) do ob
        Expr(:kw, ob, :(add_parts!(to, $(QuoteNode(ob)), length(parts.$ob))))
        end...))),
    map(attrs) do attr
      ob = s.doms[attr]
      :(set_subpart!(to, newparts.$ob, $(QuoteNode(attr)),
                     subpart(from, parts.$ob, $(QuoteNode(attr)))))
    end...,
    :newparts)
end

function replace_colons(acs::StructACSet, parts::NamedTuple{types}) where {types}
  NamedTuple{types}(map(types, parts) do type, part
    part == (:) ? (1:nparts(acs, type)) : part
  end)
end

# Printing
##########

function Base.show(io::IO, acs::T) where {S,T<:StructACSet{S}}
  s = SchemaDesc(S)
  if get(io, :compact, false)
    print(io, nameof(T))
    print(io, " {")
    join(io, ("$ob = $(nparts(acs,ob))" for ob in s.obs), ", ")
    print(io, "}")
  else
    print(io, T)
    println(io, ":")
    join(io, Iterators.flatten((
      ("  $ob = $(parts(acs,ob))" for ob in s.obs),
      ("  $f : $(s.doms[f]) → $(s.codoms[f]) = $(subpart(acs,f))"
       for f in s.homs),
      ("  $a : $(s.doms[a]) → $(s.codoms[a]) = $(subpart(acs,a))"
       for a in s.attrs),
    )), "\n")
  end
end

# TODO: implement Tables interface
# Steps:
# - Wrapper data structure

struct StructACSetTable{T<:StructACSet, ob} <: Tables.AbstractColumns
  parent::T
end

StructACSetTable(acs::StructACSet, ob::Symbol) = StructACSetTable{typeof(acs), ob}(acs)

ACSetInterface.tables(acs::StructACSet{<:SchemaDescType{obs}}) where {obs} =
  NamedTuple([ob => StructACSetTable(acs, ob) for ob in obs])

parent(sat::StructACSetTable) = getfield(sat, :parent)

struct StructACSetRow{T<:StructACSet, ob} <: Tables.AbstractRow
  parent::T
  idx::Int
end

parent(row::StructACSetRow) = getfield(row, :parent)
idx(row::StructACSetRow) = getfield(row, :idx)

# - Tables.jl interface

Tables.istable(sat::StructACSetTable) = true

Tables.columnaccess(sat::StructACSetTable) = true
Tables.columns(sat::StructACSetTable) = sat

Tables.getcolumn(sat::StructACSetTable{T,ob}, i::Int) where {T,ob} =
  Base.getproperty(sat, outgoing(parent(sat), ob)[i])

Base.getproperty(sat::StructACSetTable, nm::Symbol) = subpart(parent(sat), :, nm)

Tables.getcolumn(sat::StructACSetTable, nm::Symbol) = Base.getproperty(sat, nm)

Base.propertynames(sat::StructACSetTable{T,ob}) where {T,ob} = outgoing(parent(sat), ob)

Tables.columnnames(sat::StructACSetTable) = Base.propertynames(sat)

Tables.rowaccess(sat::StructACSetTable) = true
Tables.rows(sat::StructACSetTable{T,ob}) where {T,ob} =
  StructACSetRow{T,ob}.(Ref(parent(sat)), parts(parent(sat), ob))

Tables.getcolumn(row::StructACSetRow{T,ob}, i::Int) where {T,ob} =
  Base.getproperty(row, outgoing(parent(row), ob)[i])

Base.getproperty(row::StructACSetRow, nm::Symbol) = subpart(parent(row), idx(row), nm)

Tables.getcolumn(row::StructACSetRow, nm::Symbol) = Base.getproperty(row, nm)

Base.propertynames(row::StructACSetRow{T,ob}) where {T,ob} = outgoing(parent(row), ob)

Tables.columnnames(row::StructACSetRow) = Base.propertynames(row)

# Mapping

function Base.map(acs::ACSet; kwargs...)
  _map(acs, (;kwargs...))
end

function sortunique!(x)
  sort!(x)
  unique!(x)
  x
end

function groupby(f::Function, xs)
  d = Dict{typeof(f(xs[1])),Vector{eltype(xs)}}()
  for x in xs
    k = f(x)
    if k in keys(d)
      push!(d[k],x)
    else
      d[k] = [x]
    end
  end
  d
end

@generated function _map(acs::AT, fns::NamedTuple{map_over}) where
  {S, Ts, AT<:StructACSet{S,Ts}, map_over}
  s = SchemaDesc(S)
  map_over = Symbol[map_over...]

  mapped_attrs = intersect(s.attrs, map_over)
  mapped_attrtypes = intersect(s.attrtypes, map_over)
  mapped_attrs_from_attrtypes = filter(a -> s.codoms[a] ∈ mapped_attrtypes, s.attrs)
  attrs_accounted_for = sortunique!(Symbol[mapped_attrs; mapped_attrs_from_attrtypes])

  affected_attrtypes = sortunique!(map(a -> s.codoms[a], attrs_accounted_for))
  needed_attrs = sort!(filter(a -> s.codoms[a] ∈ affected_attrtypes, s.attrs))

  unnaccounted_for_attrs = filter(a -> a ∉ attrs_accounted_for, needed_attrs)
  unnaccounted_for_attrs == [] ||
    error("not enough functions provided to fully transform ACSet, need functions for: $(unnaccounted_for_attrs)")

  fn_applications = map(attrs_accounted_for) do a
    qa = q(a)
    if a ∈ mapped_attrs
      :($a = (fns[$qa]).(subpart(acs, $qa)))
    else
      d = s.codoms[a]
      :($a = (fns[$(q(d))]).(subpart(acs, $qa)))
    end
  end

  abc = groupby(a -> s.codoms[a], s.attrs)

  attrtype_instantiations = map(enumerate(s.attrtypes)) do (i,d)
    if d ∈ affected_attrtypes
      :(mapreduce(eltype, typejoin,
                  $(Expr(:tuple, (:(fn_vals[$(q(a))]) for a in abc[d])...))))
    else
      :($(Ts[i]))
    end
  end

  quote
    fn_vals = $(Expr(:tuple, fn_applications...))
    new_acs = $(AT.name.wrapper){$(attrtype_instantiations...)}()
    $(Expr(:block, map(s.obs) do ob
           :(add_parts!(new_acs,$(q(ob)),nparts(acs,$(q(ob)))))
      end...))
    $(Expr(:block, map(s.homs) do hom
           :(set_subpart!(new_acs,$(q(hom)),subpart(acs,$(q(hom)))))
      end...))
    $(Expr(:block, map(s.attrs) do attr
        qa = Expr(:quote, attr)
        if attr ∈ attrs_accounted_for
           :(set_subpart!(new_acs,$qa,fn_vals[$qa]))
        else
           :(set_subpart!(new_acs,$qa,subpart(acs,$qa)))
        end
      end...))
    return new_acs
  end
end

end


text:
using Zygote: ignore

include("dqns/dqns.jl")
include("policy_gradient/policy_gradient.jl")
include("searching/searching.jl")
include("cfr/cfr.jl")

ext: jl


text:
using Probably
using Probably: Bits, apply!, H, ⊗

# (N = 2^n)th root of unity
ω(n) = MathConstants.e^(-2π*im/(2^n))

# The Fourier transform is represented by the matrix W.
# We could use W directly, but it's more interesting
# to build it from basic (one or two qubit) gates.
W(n) = ω(n) .^ [i*j for i = 0:2^n-1, j = 0:2^n-1] ./ √(2^n)

# We will use the same recursive decomposition that powers the FFT.

# W(k+1) = 1|I D(k)|  * |W(k) 0  | * R(k+1)
#         √2|I D(k)|    | 0  W(k)|
# Where
# D(k) = diagm(ω(k+1) .^ (0:2^k-1))
# R(k)[i,j] = 2i == j || 2(i-2^k)+1 == j

# R moves the least significant bit to the front
function R(buf)
  for i = length(buf):-1:2
    swap(buf[i],buf[i-1])
  end
end

# Phase shift gates that make up D
Ds(k) = [[1 0; 0 ω(i)] for i = 1:k]

# Apply |I D(k)| by decomposition
#       |I D(k)|
function D(buf)
  ds = Ds(length(buf))
  for i = 2:length(buf)
    apply!(control(ds[i]), buf[1], buf[i])
  end
  hadamard(buf[1])
end

function QFT(buf::QBuffer)
  if length(buf) == 1
    hadamard(buf[1])
  else
    R(buf)
    QFT(buf[2:end])
    D(buf)
  end
  return buf
end

# Test it

buf = QBuffer(4)
# Set up an interesting state
hadamard(buf[1])
cnot(buf[1], buf[2])
# Cheat and check what the result will be
cheat_fft = W(4)*buf.bits.state.ψ
# Apply the QFT
QFT(buf)
buf.bits.state.ψ ≈ cheat_fft


text:
module TestAdaptivity

using FrameFun, DomainSets, Test

@testset begin "Optimal adaptivity"
    P = FrameFun.ExtensionFramePlatform(WeightedSumPlatform(platform(ChebyshevT(10)^2), (x,y)->1.,
            (x,y)->sqrt(x^2+y^2)), .9 .* UnitDisk())
    f = (x,y) -> cos(pi*(x+y)) + sqrt(x^2+y^2)*sin(1+pi*(x+y))
    p = IncrementalCartesianParameterPath{2}()
    PP = parametrizedplatform(P, HierarchyPath(p,ProductPath(p,p)))

    F1, logbook1, _ = FrameFun.Adaptivity.adaptive_approximation(OptimalStyle(), f, PP; criterion = FNAStyle(),δ=1e-2)
    @test size(F1)==(45,)
    @test length(logbook1) == 8
    @test last(logbook1)[3] < .1
    @test norm(coefficients(F1))>10*2.188

    F, logbook, _ = FrameFun.Adaptivity.adaptive_approximation(OptimalStyle(), f, PP; criterion = FNAStyle(),δ=1e-2,FNAη=10)
    @test size(F)==(190,)
    @test length(logbook) == 12
    @test last(logbook1)[3] < .1
    @test norm(coefficients(F))<10*2.188

    P = FrameFun.ExtensionFramePlatform(WeightedSumPlatform(parametrizedplatform(platform(ChebyshevT(10)^2),CartesianParameterPath((1,1))), (x,y)->1.,
            (x,y)->sqrt(x^2+y^2)), .9 .* UnitDisk())
    PP = parametrizedplatform(P)
    f = (x,y) -> cos(pi*(x+y)) + sqrt(x^2+y^2)*sin(1+pi*(x+y))

    F1 = Fun(f, PP; adaptivestyle = OptimalStyle(), criterion = FNAStyle(),δ=1e-2,FNAη=Inf)
    @test size(F1)==(50,)
    F1 = Fun(f, PP; adaptivestyle = OptimalStyle(), criterion = FNAStyle(),δ=1e-2,FNAη=10)
    @test size(F1)==(200,)
end

end


text:
__precompile__(true)
module RubberDucks

export rubber_ducks_in_earth
"""
Computes the number of rubber ducks that would fill the volume of the Earth.

Assumes a spherical Earth, and liquid rubber ducks (to avoid the packing problem).

Arguments:
- `num_ruber::Real` is the first function argument. [_m^3_]
- `r_earth::Real` is the first function argument. Default: 6378000 [_m_]

Returns:
- `num_ducks::Float64` Number of rubber ducks that could
"""
function rubber_ducks_in_earth(duck_volume::Real, r_earth::Real=6378000)
    # Volume of spherrical Earth
    v_earth =  4.0/3.0*pi*r_earth^3

    # Return number of ducks in Earth
    return v_earth/duck_volume
end

export RUBBER_DUCK_PER_SEC
"""
Number of rubber ducks produced per second in the entire world on average.

Units: _ducks/sec_
"""
const RUBBER_DUCK_PER_SEC = 200.0

export add_one
"""
Adds one to a value

Argments:
- `val:Real` A value

Returns:
- `x::Float64` Value plus one
"""
function add_one(val::Real)
    return val + 1
end

end # End of module RubberDucks (This comment is nice to distinguish from an end)


text:
immutable Frechet <: ContinuousUnivariateDistribution
    α::Float64
    β::Float64

    function Frechet(α::Real, β::Real)
    	α > zero(α) && β > zero(β) || error("Both shape and scale must be positive")
    	new(float64(α), float64(β))
    end

    Frechet(α::Real) = Frechet(α, 1.0)
    Frechet() = new(1.0, 1.0)
end

@distr_support Frechet 0.0 Inf

#### Parameters

shape(d::Frechet) = d.α
scale(d::Frechet) = d.β
params(d::Frechet) = (d.α, d.β)

#### Statistics

mean(d::Frechet) = (α = d.α; α > 1.0 ? d.β * gamma(1.0 - 1.0 / α) : Inf)

median(d::Frechet) = d.β * logtwo^(-1.0 / d.α)

mode(d::Frechet) = (iα = -1.0/d.α; d.β * (1.0 - iα) ^ iα)

function var(d::Frechet)
    if d.α > 2.0
        iα = 1.0 / d.α
        return d.β^2 * (gamma(1.0 - 2.0 * iα) - gamma(1.0 - iα)^2)
    else
        return Inf
    end
end

function skewness(d::Frechet)
    if d.α > 3.0
        iα = 1.0 / d.α
        g1 = gamma(1.0 - iα)
        g2 = gamma(1.0 - 2.0 * iα)
        g3 = gamma(1.0 - 3.0 * iα)
        return (g3 - 3.0 * g2 * g1 + 2 * g1^3) / ((g2 - g1^2)^1.5)
    else
        return Inf
    end
end

function kurtosis(d::Frechet)
    if d.α > 3.0
        iα = 1.0 / d.α
        g1 = gamma(1.0 - iα)
        g2 = gamma(1.0 - 2.0 * iα)
        g3 = gamma(1.0 - 3.0 * iα)
        g4 = gamma(1.0 - 4.0 * iα)
        return (g4 - 4.0 * g3 * g1 + 3 * g2^2) / ((g2 - g1^2)^2) - 6.0
    else
        return Inf
    end
end

function entropy(d::Frechet)
    const γ = 0.57721566490153286060  # γ is the Euler-Mascheroni constant
    1.0 + γ / d.α + γ + log(d.β / d.α)
end

#### Evaluation

function logpdf(d::Frechet, x::Float64)
    (α, β) = params(d)
    if x > 0.0
        z = β / x
        return log(α / β) + (1.0 + α) * log(z) - z^α
    else
        return -Inf
    end
end

pdf(d::Frechet, x::Float64) = exp(logpdf(d, x))

cdf(d::Frechet, x::Float64) = x > 0.0 ? exp(-((d.β / x) ^ d.α)) : 0.0
ccdf(d::Frechet, x::Float64) = x > 0.0 ? -expm1(-((d.β / x) ^ d.α)) : 1.0
logcdf(d::Frechet, x::Float64) = x > 0.0 ? -(d.β / x) ^ d.α : -Inf
logccdf(d::Frechet, x::Float64) = x > 0.0 ? log1mexp(-((d.β / x) ^ d.α)) : 0.0

quantile(d::Frechet, p::Float64) = d.β * (-log(p)) ^ (-1.0 / d.α)
cquantile(d::Frechet, p::Float64) = d.β * (-log1p(-p)) ^ (-1.0 / d.α)
invlogcdf(d::Frechet, lp::Float64) = d.β * (-lp)^(-1.0 / d.α)
invlogccdf(d::Frechet, lp::Float64) = d.β * (-log1mexp(lp))^(-1.0 / d.α)

function gradlogpdf(d::Frechet, x::Float64)
    (α, β) = params(d)
    insupport(Frechet, x) ? -(α + 1.0) / x + α * (β^α) * x^(-α-1.0)  : 0.0
end

## Sampling

rand(d::Frechet) = d.β * randexp() ^ (-1.0 / d.α)


text:
include("colortypes.jl")
include("core.jl")
include("map.jl")
include("overlays.jl")
include("io.jl")
include("readnrrd.jl")
include("readremote.jl")
@osx_only include("readOSX.jl")
include("algorithms.jl")
include("edge.jl")


text:
"""
    VariationalDenseBE(in, out, rank,
                        ensemble_size,
                        σ=identity;
                        bias=true,
                        init=glorot_normal,
                        alpha_init=glorot_normal,
                        gamma_init=glorot_normal)
    VariationalDenseBE(layer, alpha_sampler, gamma_sampler,
                        ensemble_bias, ensemble_act, rank)

Creates a bayesian dense BatchEnsemble layer.
Batch ensemble is a memory efficient alternative for deep ensembles. In deep ensembles,
if the ensemble size is N, N different models are trained, making the time and memory
complexity O(N * complexity of one network).
BatchEnsemble generates weight matrices for each member in the ensemble using a
couple of rank 1 vectors R (alpha), S (gamma), RS' and multiplying the result with
weight matrix W element wise. We also call R and S as fast weights. In the bayesian
version of batch ensemble, instead of having point estimates of the fast weights, we
sample them form a trainable parameterized distribution.

Reference - https://arxiv.org/abs/2005.07186

During both training and testing, we repeat the samples along the batch dimension
N times, where N is the ensemble_size. For example, if each mini batch has 10 samples
and our ensemble size is 4, then the actual input to the layer has 40 samples.
The output of the layer has 40 samples as well, and each 10 samples can be considered
as the output of an esnemble member.

# Fields
- `layer`: The dense layer which transforms the pertubed input to output
- `alpha_sampler`: Sampler for the first Fast weight of size (in_dim, ensemble_size)
- `gamma_sampler`: Sampler for the second Fast weight of size (out_dim, ensemble_size)
- `ensemble_bias`: Bias added to the ensemble output, separate from dense layer bias
- `ensemble_act`: The activation function to be applied on ensemble output
- `rank`: Rank of the fast weights (rank > 1 doesn't work on GPU for now)

# Arguments
- `in::Integer`: Input dimension of features
- `out::Integer`: Output dimension of features
- `rank::Integer`: Rank of the fast weights
- `ensemble_size::Integer`: Number of models in the ensemble
- `σ::F=identity`: Activation of the dense layer, defaults to identity
- `init=glorot_normal`: Initialization function, defaults to glorot_normal
- `bias::Bool=true`: Toggle the usage of bias in the dense layer
- `ensemble_bias::Bool=true`: Toggle the usage of ensemble bias
- `ensemble_act::F=identity`: Activation function for enseble outputs
- `alpha_init=TrainableMvNormal`: Initialization function for the alpha fast weight,
                        defaults to TrainableMvNormal
- `gamma_init=TrainableMvNormal`: Initialization function for the gamma fast weight,
                        defaults to TrainableMvNormal
"""
struct VariationalDenseBE{L,F,M,B}
    layer::L
    alpha_sampler::M
    gamma_sampler::M
    ensemble_bias::B
    ensemble_act::F
    rank
end

function VariationalDenseBE(
    in::Integer,
    out::Integer,
    rank::Integer,
    ensemble_size::Integer,
    σ = identity;
    init = glorot_normal,
    bias = true,
    ensemble_bias = true,
    ensemble_act = identity,
    alpha_init = TrainableMvNormal,
    gamma_init = TrainableMvNormal,
    device = cpu
)

    layer = Flux.Dense(in, out, σ; init = init, bias = bias)
    if rank >= 1
        alpha_shape = (in, ensemble_size, rank)
        gamma_shape = (out, ensemble_size, rank)
    else
        error("Rank must be >= 1.")
    end
    alpha_sampler = alpha_init(alpha_shape, init = init, device = device)
    gamma_sampler = gamma_init(gamma_shape, init = init, device = device)
    gamma = gamma_sampler()
    ensemble_bias = create_bias(gamma, ensemble_bias, out, ensemble_size)

    return VariationalDenseBE(
        layer,
        alpha_sampler,
        gamma_sampler,
        ensemble_bias,
        ensemble_act,
        rank,
    )
end

@functor VariationalDenseBE

"""
The forward pass for a VariationalDenseBE layer. The input is initially perturbed
using the first fast weight, then passed through the dense layer, and finally
multiplied by the second fast weight.

# Arguments
- `x`: Input tensors
"""
function (be::VariationalDenseBE)(x)
    layer = be.layer
    alpha = be.alpha_sampler()
    gamma = be.gamma_sampler()
    e_b = be.ensemble_bias
    e_σ = be.ensemble_act
    rank = be.rank

    batch_size = size(x)[end]
    in_size = size(alpha)[1]
    out_size = size(gamma)[1]
    ensemble_size = size(alpha)[2]
    samples_per_model = batch_size ÷ ensemble_size

    # TODO: Do we really need repeat? Or can we use broadcast?
    # Alpha, gamma shapes - [units, ensembles, rank]
    alpha = reshape(alpha, (in_size, ensemble_size * rank))
    gamma = reshape(gamma, (out_size, ensemble_size * rank))
    # Repeat breaks on GPU when input dims > 2
    alpha = repeat(alpha, samples_per_model)
    gamma = repeat(gamma, samples_per_model)
    # Reshape alpha, gamma to [units, batch_size, rank]
    alpha = reshape(alpha, (in_size, batch_size, rank))
    gamma = reshape(gamma, (out_size, batch_size, rank))
    # Reshape inputs to [units, batch_size, 1] for broadcasting
    x = Flux.unsqueeze(x, (ndims(x) + 1))
    # Perturb the inputs
    perturbed_x = x .* alpha
    # Dense layer forward pass
    outputs = layer(perturbed_x) .* gamma
    # Reduce the rank dimension through summing it up
    outputs = sum(outputs, dims = ndims(outputs))
    outputs = reshape(outputs, (out_size, samples_per_model, ensemble_size))
    # Reshape ensemble bias
    e_b = Flux.unsqueeze(e_b, ndims(e_b))

    outputs = e_σ.(outputs .+ e_b)
    outputs = reshape(outputs, (out_size, batch_size))
    return outputs
end


text:
function _partition(x::AbstractArray, p::Int, r::Int)::Int
    @inbounds y = x[r]
    i = p-1
    for j ∈ p:(r-1)
        @inbounds if x[j] ≤ y
            i += 1
            @inbounds tmp  = x[i]
            @inbounds x[i] = x[j]
            @inbounds x[j] = tmp
        end
    end
    @inbounds tmp    = x[i+1]
    @inbounds x[i+1] = x[r]
    @inbounds x[r]   = tmp
    return i+1
end

function quicksort!(x::AbstractArray, p::Int, r::Int)
    if p ≥ r
        return nothing
    end
    q::Int = _partition(x, p, r)
    quicksort!(x, p, q-1)
    quicksort!(x, q+1, r)
end


text:
"""
Adapted from sandreza/Learning/sandbox/gaussian_process.jl
https://github.com/sandreza/Learning/blob/master/sandbox/gaussian_process.jl
Changed handling of kernel functions; changed some variable names;
added log marginal likelihood function.
"""

using LinearAlgebra

"""
GP
# Description
- data structure for typical GPR computations
# Data Structure and Description
    kernel::ℱ, a Kernel object
    x_train::𝒮 , an array of vectors (n-length array of D-length vectors)
    α::𝒮2 , an array
    K::𝒰 , matrix or sparse matrix
    CK::𝒱, cholesky factorization of K
    cache::𝒜, cache vector
"""
struct GP{Kernel, 𝒮, 𝒮2, 𝒰, 𝒱, 𝒜}
    kernel::Kernel
    x_train::𝒮
    α::𝒮2
    K::𝒰
    CK::𝒱
    cache::𝒜
end

"""
GP_multiple
# Description
- data structure for GPR computations where each gridpoint in the prediction has a different predictor
# Data Structure and Description
    GPs, Array of GP objects
    kernel, Kernel object
    x_train, an array of vectors (n-length array of D-length vectors)
    stencil_size, scalar stencil length from 1 to the number of gridpoints
    stencil_ranges, index ranges at which to slice a profile
"""
struct GP_multiple
    GPs::Array{GP}
    kernel::Kernel
    x_train
    stencil_size
    stencil_ranges
end

"""
model(x_train, y_train; kernel; sparsity_threshold = 0.0, robust = true, entry_threshold = sqrt(eps(1.0)))
# Description
Constructs the posterior distribution for a gp. In other words this does the 'training' automagically.
# Arguments
- `x_train`: (array). training inputs (predictors), must be an array of states.
                      length-n array of D-length vectors, where D is the length of each input n is the number of training points.
- `y_train`: (array). training outputs (prediction), must have the same number as x_train
                      length-n array of D-length vectors.
- `kernel`: (Kernel). Kernel object. See kernels.jl.
                      kernel_function(kernel)(x,x') maps predictor x predictor to real numbers.
# Keyword Arguments
- `z`: (vector). values w.r.t. which to derivate the state (default none).
- `normalize`: (bool). whether to normalize the data during preprocessing and reverse the scaling for postprocessing. Can lead to better performance.
- `hyperparameters`: (array). default = []. hyperparameters that enter into the kernel
- `sparsity_threshold`: (number). default = 0.0. a number between 0 and 1 that determines when to use sparse array format. The default is to never use it
- `robust`: (bool). default = true. This decides whether to uniformly scale the diagonal entries of the Kernel Matrix. This sometimes helps with Cholesky factorizations.
- `entry_threshold`: (number). default = sqrt(eps(1.0)). This decides whether an entry is "significant" or not. For typical machines this number will be about 10^(-8) * largest entry of kernel matrix.
# Return
- GP object
"""
function model(x_train, y_train, kernel, zavg; sparsity_threshold = 0.0, robust = true, entry_threshold = sqrt(eps(1.0)))

    # get k(x,x') function from kernel object
    kernel = kernel_function(kernel; z=zavg)
    # fill kernel matrix with values
    K = compute_kernel_matrix(kernel, x_train)

    # get the maximum entry for scaling and sparsity checking
    mK = maximum(K)

    # make Cholesky factorization work by adding a small amount to the diagonal
    if robust
        K += mK*sqrt(eps(1.0))*I
    end

    # check sparsity
    bools = K .> entry_threshold * mK
    sparsity = sum(bools) / length(bools)
    if sparsity < sparsity_threshold
        sparse_K = similar(K) .* 0
        sparse_K[bools] = sK[bools]
        K = sparse(Symmetric(sparse_K))
        CK = cholesky(K)
    else
        CK = cholesky(K)
    end

    # get prediction weights FIX THIS SO THAT IT ALWAYS WORKS
    y = hcat(y_train...)'
    α = CK \ y # α = K + σ_noise*I

    # construct struct
    return GP(kernel, x_train, α', K, Array(CK), zeros(length(x_train)))
end

function stencil_range(D, stencil_size, i)
    k = Int(floor(stencil_size/2))
    if i-k < 1
        return 1:stencil_size
    elseif i-k+stencil_size-1 > D
        return D-stencil_size+1:D
    else
        start = i-k
        start:start+stencil_size-1
    end
end

stencil(data, stencil_range) = [x[stencil_range] for x in data]

"""
model(𝒟::ProfileData; kernel::Kernel = Kernel(), stencil_size=nothing)
# Description
Create an instance of GP using data from ProfileData object 𝒟.
# Arguments
- 𝒟::ProfileData, Data for training the GP
# Keyword Arguments
- kernel::Kernel,
- stencil_size::Int64
"""
function model(𝒟::ProfileData; kernel::Kernel = Kernel(), stencil_size::Int64=0)
    # create instance of GP using data from ProfileData object
    if stencil_size == 0
        return model(𝒟.x_train, 𝒟.y_train, kernel, 𝒟.zavg);
    end

    D = length(𝒟.x_train[1])
    # create instance of GP using data from ProfileData object
    r = [stencil_range(D,stencil_size,i) for i=1:D]
    GPs = [model(stencil(𝒟.x_train,r[i]), stencil(𝒟.y_train,i), kernel, 𝒟.zavg[r[i]]) for i=1:D]
    return GP_multiple(GPs, kernel, 𝒟.x_train, stencil_size, r);
end

"""
model_output(x, 𝒢::GP)
# Description
- Given state x, GP 𝒢, returns the mean GP prediction
# Arguments
- `x`: single scaled state
- `𝒢`: GP object with which to make the prediction
# Return
- `y`: scaled prediction
"""
# function model_output(x, 𝒢::GP)
#     return 𝒢.α * 𝒢.kernel.([x], 𝒢.x_train)
# end

function model_output(x, 𝒢::GP)
    for i in 1:length(𝒢.cache)
        𝒢.cache[i] = 𝒢.kernel(x, 𝒢.x_train[i])
    end
    return 𝒢.α * 𝒢.cache
end

"""
model_output(x, 𝒢::GP_multiple)
# Description
- Given state x, GP_multiple 𝒢, returns the mean GP prediction
# Arguments
- `x`: single scaled state
- `𝒢`: GP_multiple object with which to make the prediction
# Return
- `y`: scaled prediction
"""
function model_output(x, 𝒢::GP_multiple)
    xs = [x[r] for r in 𝒢.stencil_ranges]
    return [model_output(xs[i], 𝒢.GPs[i])[1] for i in 1:length(xs)]
end

"""
uncertainty(x, 𝒢::GP)
# Description
- Given state x and GP 𝒢, output the variance at a point
# Arguments
- `x`: state
- `𝒢`: GP object
# Return
- `var`: variance
"""
function uncertainty(x, 𝒢::GP)
    tmpv = zeros(size(𝒢.x_train)[1])
    for i in eachindex(𝒢.x_train)
        tmpv[i] = 𝒢.kernel(x, 𝒢.x_train[i])
    end
    # no ldiv for suitesparse
    tmpv2 = 𝒢.CK \ tmpv
    var = k(x, x) .- tmpv'*tmpv2  # var(f*) = k(x*,x*) - tmpv'*tmpv2
    return var
end

"""
compute_kernel_matrix(kernel, x)
# Description
- Computes the kernel matrix for GPR
# Arguments
- `k` : (Kernel) kernel function k(a,b).
- `x` : (array of predictors). x[1] is a vector
# Return
- `sK`: (symmetric matrix). A symmetric matrix with entries sK[i,j] = k(x[i], x[j]). This is only meaningful if k(x,y) = k(y,x) (it should)
"""
function compute_kernel_matrix(k, x)

    K = [k(x[i], x[j]) for i in eachindex(x), j in eachindex(x)]

    if typeof(K[1,1]) <: Number
        sK = Symmetric(K)
    else
        sK = K
    end
    return sK
end

"""
mean_log_marginal_loss(y_train, 𝒢::GP; add_constant=false)
# Description
Computes log marginal loss for each element in the output and averages the results.
Assumes noise-free observations.

log(p(y|X)) = -(1/2) * (y'*α + 2*sum(Diagonal(CK)) + n*log(2*pi))
where n is the number of training points and

# Arguments
- `y_train`: (Array). training outputs (prediction), must have the same number as x_train
- `𝒢`: (GP).
# Keyword Arguments
- `add_constant`: (bool). whether to give the exact value of the loss or leave out an added constant for efficiency.

"""
function mean_log_marginal_loss(y_train, 𝒢::GP; add_constant=false)
    n = length(𝒢.x_train)
    D = length(𝒢.x_train[1])

    ys = hcat(y_train...)' # n x D

    if add_constant
        c = sum([log(𝒢.CK[i,i]) for i in 1:n]) + 0.5*n*log(2*pi)
        total_loss=0.0
        for i in 1:D
            total_loss -= 0.5*ys[:,i]'*𝒢.α[:,i] + c
        end
    else
        total_loss=0.0
        for i in 1:D
            total_loss -= 0.5*ys[:,i]'*𝒢.α[:,i]
        end
    end

    return total_loss / D
end

max_stars_repo_name: adelinehillier/LearnConvection

max_issues_repo_name: adelinehillier/LearnConvection

max_forks_repo_name: adelinehillier/LearnConvection


text:
module LineSearch

using Printf
using Stopping
using NLPModels
using NLPModelsJuMP
using Printf
using LinearAlgebra
using OptimizationProblems

include("LSFunctionMetamod.jl")
include("LineModel2mod.jl")

include("phi_dphi.jl")
include("pred_ared.jl")
include("update_H.jl")

include("armijo_backtracking.jl")
include("one-step-size.jl")
include("shamanskii-line-search.jl")

include("TR/TR_generic_ls.jl")
include("TR/TR-ls-step-computation.jl")
include("TR/TR-Nwt-LS.jl")
include("TR/TR-Sec-LS.jl")
include("TR/TR-SecA-LS.jl")

include("ARC/ARC_generic_ls.jl")
include("ARC/ARC-Nwt-LS.jl")
include("ARC/ARC-Sec-LS.jl")
include("ARC/ARC-SecA-LS.jl")
include("ARC/ARC_direction_computation.jl")

include("Bisection/find_interval.jl")
include("Bisection/Biss-Nwt.jl")
include("Bisection/Biss-Sec.jl")
include("Bisection/Biss-SecA.jl")
include("Bisection/Biss-Cub.jl")

include("other-ls/struct-ls.jl")
include("other-ls/hagerzhang2.jl")
include("other-ls/hagerzhang.jl")

end # module


text:
"""
    Pathfinding.set_target!(agent, target, pathfinder::AStar{D})
Calculate and store the shortest path to move the agent from its current position to
`target` (a position e.g. `(1, 5)` or `(1.3, 5.2)`) using the provided `pathfinder`.

Use this method in conjuction with [`move_along_route!`](@ref).
"""
function set_target!(
    agent::A,
    target::Dims{D},
    pathfinder::AStar{D},
) where {D,A<:AbstractAgent}
    path = find_path(pathfinder, agent.pos, target)
    isnothing(path) && return
    pathfinder.agent_paths[agent.id] = path
end

"""
    Pathfinding.set_best_target!(agent, targets, pathfinder::AStar{D}; kwargs...)
Calculate, store, and return the best path to move the agent from its current position to
a chosen target position taken from `targets` using `pathfinder`.

The `condition = :shortest` keyword retuns the shortest path which is shortest out of the
possible target positions. Alternatively, the `:longest` path may also be requested.

Return the position of the chosen target. Return `nothing` if none of the supplied targets are
reachable.
"""
function set_best_target!(
    agent::A,
    targets,
    pathfinder::AStar{D,P,M,Int64};
    condition::Symbol = :shortest,
) where {A<:AbstractAgent,D,P,M}
    @assert condition ∈ (:shortest, :longest)
    compare = condition == :shortest ? (a, b) -> a < b : (a, b) -> a > b
    best_path = Path{D,Int64}()
    best_target = nothing
    for target in targets
        path = find_path(pathfinder, agent.pos, target)
        isnothing(path) && continue
        if isempty(best_path) || compare(length(path), length(best_path))
            best_path = path
            best_target = target
        end
    end

    isnothing(best_target) && return
    pathfinder.agent_paths[agent.id] = best_path
    return best_target
end

"""
    Pathfinding.move_along_route!(agent, model::ABM{<:GridSpace{D}}, pathfinder::AStar{D})
Move `agent` for one step along the route toward its target set by [`Pathfinding.set_target!`](@ref)

For pathfinding in models with [`GridSpace`](@ref).

If the agent does not have a precalculated path or the path is empty, it remains stationary.
"""
function Agents.move_along_route!(
    agent::A,
    model::ABM{<:GridSpace{D},A},
    pathfinder::AStar{D}
) where {D,A<:AbstractAgent}
    isempty(agent.id, pathfinder) && return

    move_agent!(agent, first(pathfinder.agent_paths[agent.id]), model)
    popfirst!(pathfinder.agent_paths[agent.id])
end

"""
    Pathfinding.nearby_walkable(position, model::ABM{<:GridSpace{D}}, pathfinder::AStar{D}, r = 1)
Return an iterator over all [`nearby_positions`](@ref) within "radius" `r` of the given
`position` (excluding `position`), which are walkable as specified by the given `pathfinder`.
"""
nearby_walkable(position, model::ABM{<:GridSpace{D}}, pathfinder::AStar{D}, r = 1) where {D} =
    Iterators.filter(x -> pathfinder.walkmap[x...] == 1, nearby_positions(position, model, r))

"""
    Pathfinding.random_walkable(model, pathfinder::AStar{D})
Return a random position in the given `model` that is walkable as specified by the given
`pathfinder`.
"""
random_walkable(model::ABM{<:GridSpace{D}}, pathfinder::AStar{D}) where {D} =
    Tuple(rand(model.rng, filter(x -> pathfinder.walkmap[x], CartesianIndices(model.space.s))))


text:
using Test, Gmsh

@testset "test_t1.jl" begin include("test_t1.jl") end


text:
@testset "PlanarLayer" begin
    # logpdf of a flow with a planar layer and two-dimensional inputs
    test_ad(randn(7)) do θ
        layer = PlanarLayer(θ[1:2], θ[3:4], θ[5:5])
        flow = transformed(MvNormal(zeros(2), I), layer)
        return logpdf_forward(flow, θ[6:7])
    end
    test_ad(randn(11)) do θ
        layer = PlanarLayer(θ[1:2], θ[3:4], θ[5:5])
        flow = transformed(MvNormal(zeros(2), I), layer)
        return sum(logpdf_forward(flow, reshape(θ[6:end], 2, :)))
    end

    # logpdf of a flow with the inverse of a planar layer and two-dimensional inputs
    test_ad(randn(7)) do θ
        layer = PlanarLayer(θ[1:2], θ[3:4], θ[5:5])
        flow = transformed(MvNormal(zeros(2), I), inv(layer))
        return logpdf_forward(flow, θ[6:7])
    end
    test_ad(randn(11)) do θ
        layer = PlanarLayer(θ[1:2], θ[3:4], θ[5:5])
        flow = transformed(MvNormal(zeros(2), I), inv(layer))
        return sum(logpdf_forward(flow, reshape(θ[6:end], 2, :)))
    end
end


text:
"""
    Base.read!(io::IO, model::FileFormats.MOF.Model)

Read `io` in the MathOptFormat file format and store the result in `model`.
"""
function Base.read!(io::IO, model::Model)
    if !MOI.is_empty(model)
        error("Cannot read model from file as destination model is not empty.")
    end
    options = get_options(model)
    if options.validate
        validate(io)
    end
    object = JSON.parse(io; dicttype = UnorderedObject)
    file_version = _parse_mof_version(object["version"]::UnorderedObject)
    if file_version.major != VERSION.major || file_version.minor > VERSION.minor
        error(
            "Sorry, the file can't be read because this library supports " *
            "v$(VERSION) of MathOptFormat, but the file you are trying to " *
            "read is v$(file_version).",
        )
    end
    name_map = read_variables(model, object)
    read_objective(model, object, name_map)
    read_constraints(model, object, name_map)
    return
end

function read_variables(model::Model, object::Object)
    name_map = Dict{String,MOI.VariableIndex}()
    for variable::typeof(object) in object["variables"]
        name = get(variable, "name", "")::String
        if isempty(name)
            error(
                "Variable name is `\"\"`. MathOptFormat variable names must " *
                "be unique and non-empty.",
            )
        end
        name_map[name] = index = MOI.add_variable(model)
        MOI.set(model, MOI.VariableName(), index, name)
    end
    return name_map
end

function read_objective(
    model::Model,
    object::Object,
    name_map::Dict{String,MOI.VariableIndex},
)
    obj = object["objective"]::typeof(object)
    sense = read_objective_sense(obj["sense"]::String)
    MOI.set(model, MOI.ObjectiveSense(), sense)
    if sense == MOI.FEASIBILITY_SENSE
        return
    end
    func = function_to_moi(obj["function"]::typeof(object), name_map)
    MOI.set(model, MOI.ObjectiveFunction{typeof(func)}(), func)
    return
end

function _add_constraint(
    model::Model,
    object::Object,
    name_map::Dict{String,MOI.VariableIndex},
)
    index = MOI.add_constraint(
        model,
        function_to_moi(object["function"]::typeof(object), name_map),
        set_to_moi(object["set"]::typeof(object)),
    )
    if haskey(object, "name")
        MOI.set(model, MOI.ConstraintName(), index, object["name"]::String)
    end
    return
end

function read_constraints(
    model::Model,
    object::Object,
    name_map::Dict{String,MOI.VariableIndex},
)
    for constraint::typeof(object) in object["constraints"]::Vector
        _add_constraint(model, constraint, name_map)
    end
end

function read_objective_sense(sense::String)
    if sense == "min"
        return MOI.MIN_SENSE
    elseif sense == "max"
        return MOI.MAX_SENSE
    else
        @assert sense == "feasibility"
        return MOI.FEASIBILITY_SENSE
    end
end

macro head_to_val(f, arg1, args...)
    head = gensym()
    body = Expr(
        :if,
        Expr(:call, :(==), head, string(arg1)),
        Expr(:return, Val(arg1)),
    )
    leaf = body
    for arg in args
        new_expr = Expr(
            :elseif,
            Expr(:call, :(==), head, string(arg)),
            Expr(:return, Val(arg)),
        )
        push!(leaf.args, new_expr)
        leaf = new_expr
    end
    push!(leaf.args, Expr(:return, Val(head)))
    quote
        function $(esc(f))($head::String)
            return $body
        end
    end
end

@head_to_val(
    head_to_function,
    SingleVariable,
    VectorOfVariables,
    ScalarAffineFunction,
    ScalarQuadraticFunction,
    VectorAffineFunction,
    VectorQuadraticFunction,
    ScalarNonlinearFunction,
)

"""
    function_to_moi(
        x::Object, name_map::Dict{String, MOI.VariableIndex}
    )

Convert `x` from an MOF representation into a MOI representation.
"""
function function_to_moi(x::Object, name_map::Dict{String,MOI.VariableIndex})
    val = if haskey(x, "type")
        head_to_function(x["type"]::String)
    else
        # TODO(odow): remove when v0.4 no longer supported.
        head_to_function(x["head"]::String)
    end
    return function_to_moi(val, x, name_map)
end

function function_to_moi(
    ::Val{FunctionSymbol},
    ::Object,
    ::Dict{String,MOI.VariableIndex},
) where {FunctionSymbol}
    return error(
        "Version $(VERSION) of MathOptFormat does not support the function: " *
        "$(FunctionSymbol).",
    )
end

# ========== Non-typed scalar functions ==========

function function_to_moi(
    ::Val{:SingleVariable},
    object::Object,
    name_map::Dict{String,MOI.VariableIndex},
)
    return MOI.SingleVariable(name_map[object["variable"]::String])
end

# ========== Typed scalar functions ==========

# Here, we deal with a special case: ScalarAffineTerm, ScalarQuadraticTerm,
# VectorAffineTerm, and VectorQuadraticTerm do not contain a "type" field
# (because it is unnecessary at the JSON level).

function parse_scalar_affine_term(
    object::Object,
    name_map::Dict{String,MOI.VariableIndex},
)
    return MOI.ScalarAffineTerm(
        object["coefficient"]::Float64,
        name_map[object["variable"]::String],
    )
end

function function_to_moi(
    ::Val{:ScalarAffineFunction},
    object::Object,
    name_map::Dict{String,MOI.VariableIndex},
)
    return MOI.ScalarAffineFunction{Float64}(
        parse_scalar_affine_term.(object["terms"], Ref(name_map)),
        object["constant"]::Float64,
    )
end

function parse_scalar_quadratic_term(
    object::Object,
    name_map::Dict{String,MOI.VariableIndex},
)
    return MOI.ScalarQuadraticTerm(
        object["coefficient"]::Float64,
        name_map[object["variable_1"]::String],
        name_map[object["variable_2"]::String],
    )
end

function function_to_moi(
    ::Val{:ScalarQuadraticFunction},
    object::Object,
    name_map::Dict{String,MOI.VariableIndex},
)
    return MOI.ScalarQuadraticFunction{Float64}(
        parse_scalar_affine_term.(object["affine_terms"], Ref(name_map)),
        parse_scalar_quadratic_term.(object["quadratic_terms"], Ref(name_map)),
        object["constant"]::Float64,
    )
end

# ========== Non-typed vector functions ==========

function function_to_moi(
    ::Val{:VectorOfVariables},
    object::Object,
    name_map::Dict{String,MOI.VariableIndex},
)
    return MOI.VectorOfVariables([
        name_map[variable] for variable::String in object["variables"]
    ])
end

# ========== Typed vector functions ==========

function parse_vector_affine_term(
    object::Object,
    name_map::Dict{String,MOI.VariableIndex},
)
    return MOI.VectorAffineTerm(
        object["output_index"],
        parse_scalar_affine_term(
            object["scalar_term"]::typeof(object),
            name_map,
        ),
    )
end

function function_to_moi(
    ::Val{:VectorAffineFunction},
    object::Object,
    name_map::Dict{String,MOI.VariableIndex},
)
    return MOI.VectorAffineFunction{Float64}(
        parse_vector_affine_term.(object["terms"], Ref(name_map)),
        Float64.(object["constants"]),
    )
end

function parse_vector_quadratic_term(
    object::Object,
    name_map::Dict{String,MOI.VariableIndex},
)
    return MOI.VectorQuadraticTerm(
        object["output_index"],
        parse_scalar_quadratic_term(object["scalar_term"], name_map),
    )
end

function function_to_moi(
    ::Val{:VectorQuadraticFunction},
    object::Object,
    name_map::Dict{String,MOI.VariableIndex},
)
    return MOI.VectorQuadraticFunction{Float64}(
        parse_vector_affine_term.(object["affine_terms"], Ref(name_map)),
        parse_vector_quadratic_term.(object["quadratic_terms"], Ref(name_map)),
        Float64.(object["constants"]),
    )
end

# ========== Default fallback ==========

@head_to_val(
    head_to_set,
    ZeroOne,
    Integer,
    LessThan,
    GreaterThan,
    EqualTo,
    Interval,
    Semiinteger,
    Semicontinuous,
    Zeros,
    Reals,
    Nonnegatives,
    Nonpositives,
    SecondOrderCone,
    RotatedSecondOrderCone,
    GeometricMeanCone,
    NormOneCone,
    NormInfinityCone,
    RelativeEntropyCone,
    NormSpectralCone,
    NormNuclearCone,
    RootDetConeTriangle,
    RootDetConeSquare,
    LogDetConeTriangle,
    LogDetConeSquare,
    PositiveSemidefiniteConeTriangle,
    PositiveSemidefiniteConeSquare,
    ExponentialCone,
    DualExponentialCone,
    PowerCone,
    DualPowerCone,
    SOS1,
    SOS2,
    IndicatorSet,
)

"""
    set_to_moi(x::Object)

Convert `x` from an OrderedDict representation into a MOI representation.
"""
function set_to_moi(x::Object)
    if haskey(x, "type")
        return set_to_moi(head_to_set(x["type"]::String), x)
    else
        # TODO(odow): remove when v0.4 no longer supported.
        return set_to_moi(head_to_set(x["head"]::String), x)
    end
end

# ========== Non-typed scalar sets ==========

function set_to_moi(::Val{:ZeroOne}, ::Object)
    return MOI.ZeroOne()
end

function set_to_moi(::Val{:Integer}, ::Object)
    return MOI.Integer()
end

# ========== Typed scalar sets ==========

function set_to_moi(::Val{:LessThan}, object::Object)
    return MOI.LessThan(object["upper"])
end

function set_to_moi(::Val{:GreaterThan}, object::Object)
    return MOI.GreaterThan(object["lower"])
end

function set_to_moi(::Val{:EqualTo}, object::Object)
    return MOI.EqualTo(object["value"])
end

function set_to_moi(::Val{:Interval}, object::Object)
    return MOI.Interval(object["lower"], object["upper"])
end

function set_to_moi(::Val{:Semiinteger}, object::Object)
    return MOI.Semiinteger(object["lower"], object["upper"])
end

function set_to_moi(::Val{:Semicontinuous}, object::Object)
    return MOI.Semicontinuous(object["lower"], object["upper"])
end

# ========== Non-typed vector sets ==========

function set_to_moi(::Val{:Zeros}, object::Object)
    return MOI.Zeros(object["dimension"])
end

function set_to_moi(::Val{:Reals}, object::Object)
    return MOI.Reals(object["dimension"])
end

function set_to_moi(::Val{:Nonnegatives}, object::Object)
    return MOI.Nonnegatives(object["dimension"])
end

function set_to_moi(::Val{:Nonpositives}, object::Object)
    return MOI.Nonpositives(object["dimension"])
end

function set_to_moi(::Val{:SecondOrderCone}, object::Object)
    return MOI.SecondOrderCone(object["dimension"])
end

function set_to_moi(::Val{:RotatedSecondOrderCone}, object::Object)
    return MOI.RotatedSecondOrderCone(object["dimension"])
end

function set_to_moi(::Val{:GeometricMeanCone}, object::Object)
    return MOI.GeometricMeanCone(object["dimension"])
end

function set_to_moi(::Val{:NormOneCone}, object::Object)
    return MOI.NormOneCone(object["dimension"])
end

function set_to_moi(::Val{:NormInfinityCone}, object::Object)
    return MOI.NormInfinityCone(object["dimension"])
end

function set_to_moi(::Val{:RelativeEntropyCone}, object::Object)
    return MOI.RelativeEntropyCone(object["dimension"])
end

function set_to_moi(::Val{:NormSpectralCone}, object::Object)
    return MOI.NormSpectralCone(object["row_dim"], object["column_dim"])
end

function set_to_moi(::Val{:NormNuclearCone}, object::Object)
    return MOI.NormNuclearCone(object["row_dim"], object["column_dim"])
end

function set_to_moi(::Val{:RootDetConeTriangle}, object::Object)
    return MOI.RootDetConeTriangle(object["side_dimension"])
end

function set_to_moi(::Val{:RootDetConeSquare}, object::Object)
    return MOI.RootDetConeSquare(object["side_dimension"])
end

function set_to_moi(::Val{:LogDetConeTriangle}, object::Object)
    return MOI.LogDetConeTriangle(object["side_dimension"])
end

function set_to_moi(::Val{:LogDetConeSquare}, object::Object)
    return MOI.LogDetConeSquare(object["side_dimension"])
end

function set_to_moi(::Val{:PositiveSemidefiniteConeTriangle}, object::Object)
    return MOI.PositiveSemidefiniteConeTriangle(object["side_dimension"])
end

function set_to_moi(::Val{:PositiveSemidefiniteConeSquare}, object::Object)
    return MOI.PositiveSemidefiniteConeSquare(object["side_dimension"])
end

function set_to_moi(::Val{:ExponentialCone}, ::Object)
    return MOI.ExponentialCone()
end

function set_to_moi(::Val{:DualExponentialCone}, ::Object)
    return MOI.DualExponentialCone()
end

# ========== Typed vector sets ==========

function set_to_moi(::Val{:PowerCone}, object::Object)
    return MOI.PowerCone(object["exponent"]::Float64)
end

function set_to_moi(::Val{:DualPowerCone}, object::Object)
    return MOI.DualPowerCone(object["exponent"]::Float64)
end

function set_to_moi(::Val{:SOS1}, object::Object)
    return MOI.SOS1(Float64.(object["weights"]))
end

function set_to_moi(::Val{:SOS2}, object::Object)
    return MOI.SOS2(Float64.(object["weights"]))
end

function set_to_moi(::Val{:IndicatorSet}, object::Object)
    set = set_to_moi(object["set"]::typeof(object))
    indicator = if object["activate_on"]::String == "one"
        MOI.ACTIVATE_ON_ONE
    else
        @assert object["activate_on"]::String == "zero"
        MOI.ACTIVATE_ON_ZERO
    end
    return MOI.IndicatorSet{indicator}(set)
end


text:
module TestQueryNode

using StructuredQueries
using Base.Test

io = IOBuffer()
disp = TextDisplay(IOBuffer())

type MyData end
iris = MyData()

c = @with iris filter(sepal_length > 5.0)
@test source(c) === (iris,)
@test graph(c) === c.graph
show(io, graph(c))
display(disp, graph(c))

c1 = @with iris begin
    filter(sepal_length > 5.0)
    select(petal_width)
end
c2 = @with iris filter(sepal_length > 5.0), select(petal_width)
@test isequal(c1, c2)
@test source(c1) === (iris,)
@test graph(c1) === c1.graph
show(io, graph(c1))
display(disp, graph(c1))

c1 = @with iris begin
    filter(sepal_length > 5.0)
    groupby(species)
    summarize(avg = mean(petal_width))
end
c2 = @with iris filter(sepal_length > 5.0),
    groupby(species),
    summarize(avg = mean(petal_width))
@test isequal(c1, c2)
@test source(c1) === (iris,)
@test graph(c1) === c1.graph
show(io, graph(c1))
display(disp, graph(c1))

end


text:
using Documenter
using dynamicRisk

push!(LOAD_PATH,"../src/")
makedocs(sitename="dynamicRisk.jl Documentation",
         pages = [
            "Home" => "index.md",
            "Mathematical Background" =>Any[
                "back/DiscreteTime.md",
                "back/Limit.md",
                "back/Continuous.md"
            ],
            "Tutorials" => Any["tut/basic.md",
				    "tut/Tree.md",
				    "tut/Lattice.md",
                    "tut/BSDE.md"
            ],
            "Function Library" => "library.md",
         ],
         format = Documenter.HTML(prettyurls = false)
)
# Documenter can also automatically deploy documentation to gh-pages.
# See "Hosting Documentation" and deploydocs() in the Documenter manual
# for more information.
deploydocs(
    repo = "github.com/rubsc/dynamicRisk.jl.git",
    devbranch = "main"
)


text:
using Test
using BEAST

const B3D = BEAST.SparseND.Banded3D

k0 = [2 2; 3 2]
k1 = [2 3; 3 3]
M = 2
N = 2
bandwidth = maximum(k1 - k0 .+ 1)
data = collect(begin
    k = k0[i,j] + l - 1
    v = (k > k1[i,j]) ? 0.0 : Float64((i-1) + (j-1)*M + (k-1)*M*N)
end for l in 1:bandwidth, i in axes(k0,1), j in axes(k0,2))
A = B3D(k0,data,maximum(k1))

@test size(A) == (2,2,3)
@test eltype(A) == Float64
@test A[1,1,1] == 0
@test A[1,1,2] == M*N
@test A[1,1,2] == M*N
@test A[1,1,3] == 0

slice = A[:,:,2]
@test slice == [4 6; 0 7]

F = Array(A)
for m in axes(F,1)
    for n in axes(F,2)
        for l in axes(F,3)
            @test F[m,n,l] ≈ A[m,n,l]
        end
    end
end

st_data = rand(2,5)
cv1 = BEAST.convolve(F, st_data, 4, 2)
cv2 = BEAST.convolve(A, st_data, 4, 2)
@test cv1 ≈ cv2

cv1 = BEAST.convolve(F, st_data, 4, 1)
cv2 = BEAST.convolve(A, st_data, 4, 1)
@test cv1 ≈ cv2

cv1 = BEAST.convolve(F, st_data, 4, 3)
cv2 = BEAST.convolve(A, st_data, 4, 3)
@test cv1 ≈ cv2


text:
#!/usr/bin/env julia

# Experiment in using Julia (I language I revently discovered)
# to solve Sudoku puzzles, from scratch, without looking at other peoples
# solutions.

puzzle = [
            [5,3,0,0,7,0,0,0,0],
            [6,0,0,1,9,5,0,0,0],
            [0,9,8,0,0,0,0,6,0],
            [8,0,0,0,6,0,0,0,3],
            [4,0,0,8,0,3,0,0,1],
            [7,0,0,0,2,0,0,0,6],
            [0,6,0,0,0,0,2,8,0],
            [0,0,0,4,1,9,0,0,5],
            [0,0,0,0,8,0,0,7,9]];

# println(puzzle);

# println(join(puzzle[1], " "))
println(join(map(row -> join(row, " "), puzzle), ";\n"))

# println(puzzle[:][1])

grid = [
5 3 0 0 7 0 0 0 0;
6 0 0 1 9 5 0 0 0;
0 9 8 0 0 0 0 6 0;
8 0 0 0 6 0 0 0 3;
4 0 0 8 0 3 0 0 1;
7 0 0 0 2 0 0 0 6;
0 6 0 0 0 0 2 8 0;
0 0 0 4 1 9 0 0 5;
0 0 0 0 8 0 0 7 9]

# println(grid[:, 1])
# println(grid[1, :])

const GRID_SIZE = 9
ALL_POSSIBLE = BitSet(1:9)
possibles = fill(ALL_POSSIBLE, (GRID_SIZE, GRID_SIZE))
# println("SetDiff: ", setdiff(possibles[1, 1], BitSet(1:3)))

curry(f, x) = (xs...) -> f(x, xs...)
precurry(f, x) = (xs...) -> f(xs..., x)
# println(precurry(setdiff, BitSet(1))(possibles[1,1]))

const sqcol = const sqrow = [1:3, 4:6, 7:9]
function neighbors(r)
    return reduce(union, [row for row in sqrow if !(r in row)])
end

function mySet(r)
    BitSet([collect(row) for row in sqrow if r in row][1])
end

@assert(mySet(4) == BitSet([4, 5, 6]))
@assert(5 in mySet(4))

function mySquare(r)
    for row in sqrow
        if r in row
            return row
        end
    end
end

@assert(mySquare(4) == 4:6)

@assert(neighbors(8) == [1,2,3,4,5,6])

function bitmask(n, subset)
    return subset .& (1 << n) .>> n
end

function column(n, g = grid)
    # return g[n*9-8:n*9]
    return g[:, n]
end

function row(n, g = grid)
    return g[n, :]
end

function squareFor(r, c, g = grid)
    g[mySquare(r), mySquare(c)]
end

@assert(row(2) == [6, 0, 0, 1, 9, 5, 0, 0, 0])
@assert(column(2) == [3, 0, 9, 0, 0, 0, 6, 0, 0])
@assert(squareFor(8,7) == [2 8 0; 0 0 5; 0 7 9])

function -(a::BitSet, b)
    setdiff(a, b)
end

# + = union for BitSets
import Base.+

function +(a::BitSet, b::BitSet)
    union(a, b)
end

# | = intersect for BitSets
import Base.|

function |(a::BitSet, b::BitSet)
    intersect(a, b)
end

function box(value, rows=3, cols=3)
    repeat([value], rows, cols)
end

notEmpty = x -> !isempty(x)

function findOnePossible(grid=grid, possibles=possibles, verbose=false)
    changed = Dict()
    for n in 1:9
        for r in 1:9
            foundSet = box(BitSet(n), 9, 1) .| possibles[r, 1:9]
            if filter(notEmpty, foundSet) == [BitSet(n)]
                ndx = findfirst(notEmpty, foundSet)
                c = first(ndx.I)
                if n in possibles[r, c]
                    key = "$(r)x$(c)"
                    changed[key] = "Setting $(r)x$(c) to $(n)"
                    setValue(n, r, c, possibles, grid)
                end
            end
            c = r
            foundSet = box(BitSet(n), 9, 1) .| possibles[1:9, c]
            if filter(notEmpty, foundSet) == [BitSet(n)]
                ndx = findfirst(notEmpty, foundSet)
                r = first(ndx.I)
                if n in possibles[r, c]
                    key = "$(r)x$(c)"
                    changed[key] = "Setting $(r)x$(c) to $(n)"
                    setValue(n, r, c, possibles, grid)
                end
            end
        end
    end
    if verbose
        print(changed)
    end
    changed
end

function findAllPossibles(grid=grid, possibles=possibles, verbose=false)
    while !isempty(findOnePossible(grid, possibles))
        if verbose
            printGrid(grid)
        else
            print(".")
        end
    end
    println("Done")
    printGrid(grid)
end

@assert((BitSet([1,2,3]) - BitSet(3))  - BitSet(1) == BitSet(2))
# println(possibles[1, 1] -= BitSet(9))

function eraseSquare(n, r, c, p=possibles)
    p[mySquare(r), mySquare(c)] .-= BitSet(n)
end

function eraseCol(n, r, c, p=possibles)
    p[r, neighbors(c)] .-= BitSet(n)
end

function eraseRow(n, r, c, p=possibles)
    p[neighbors(r), c] .-= BitSet(n)
end

function eraser(n, r, c, p=possibles)
    eraseSquare(n, r, c, p)
    eraseRow(n, r, c, p)
    eraseCol(n, r, c, p)
    p
end

# eraser(5, 1, 1)
# println(eraseSquare(3, 1, 1))
# println(eraseCol(3, 1, 1))
# items = eraseRow(3, 1, 1)[1]
# println(items)
# println(reduce(|, [p for p in map(i -> 2^i, items)]))

# Let BitSet handle converting Sets to Integers for us

# function toInt(items::BitSet{Int64})
#     powerz = map(i -> 2^i, items)
#     reduce(|, [p for p in powerz])
# end

# println(toInt(items))

function setValue(n, r, c, p=possibles, g=grid)
    if n > 0
        g[r, c] = n
        p[r, c] = BitSet()
        eraser(n, r, c, p)
    # else
    #     findSingles(g, [r], [c])
    # elseif length(p[r, c]) == 1
    #     # println([r, c, p[r, c]])
    #     return setValue(first(p[r, c]), r, c, p, g)
    end
end

# Initialize possibles from grid
function init(grid=grid, p=possibles)
    for c in 1:9
        for r in 1:9
            setValue(grid[r, c], r, c, p, grid)
        end
    end
end
# init()

# println(grid[4, 1], possibles[1:9, 1])

# println(length(possibles[8, 1]))

# p0 = possibles
# println(possibles)

# println(first(p0[8, 1]))

# println(squareFor(8,1))

# println(grid[8,1])

function printGrid(grid=grid)
    println("===")
    for r in 1:9
        println(grid[r, 1:3], grid[r, 4:6], grid[r, 7:9])
        if r % 3 == 0
            println("---")
        end
    end
end

function findPatterns(grid=grid, rows=1:9, cols=1:9, p=possibles)
    findSingles(grid, rows, cols, p)
end

function findLine(grid, rows=1:3:9, cols=1:3:9, p=possibles)
end

# Find patterns near and set value for cells with only one possible value
function findSingles(grid=grid, rows=1:9, cols=1:9, p=possibles)
    for r in rows
        for c in cols
            if grid[r, c] == 0 && length(p[r, c]) == 1
                # println([first(p[r, c]), r, c])
                setValue(first(p[r, c]), r, c, p, grid)
                # findPatterns(grid, mySquare(r), cols, p)
                # findPatterns(grid, rows, mySquare(c), p)
                findPatterns(grid, rows, [c], p)
                findPatterns(grid, [r], cols, p)
                findPatterns(grid, mySquare(r), mySquare(c), p)
            end
        end
    end
end

# println("0 Grid: ", grid[1:3, 1:3] .== 0)

# first pass simple solve
function solve(grid, p=possibles)
    println("INITIAL")
    init(grid, p)
    printGrid(grid)
    findPatterns(grid, 1:9, 1:9, p)
    findPatterns(grid, 1:9, 1:9, p)
    println("FINAL")
    printGrid(grid)
    println("0 Grid: ", any(grid[1:9, 1:9] .== 0))
    return grid
end

# println(join(map(row -> join(row, "\n"), grid), "")

# setValue(0, 8, 1)
# guessValue(8, 1) == 2
# println([possibles[r, c] if length(possibles[r, c] == 1) for c in 1:9, r in 1:9 ])

answerMatrix = [
 [5,3,4,6,7,8,9,1,2],
 [6,7,2,1,9,5,3,4,8],
 [1,9,8,3,4,2,5,6,7],
 [8,5,9,7,6,1,4,2,3],
 [4,2,6,8,5,3,7,9,1],
 [7,1,3,9,2,4,8,5,6],
 [9,6,1,5,3,7,2,8,4],
 [2,8,7,4,1,9,6,3,5],
 [3,4,5,2,8,6,1,7,9]]

# println("Answer Matrix")
# println(join(map(row -> join(row, " "), answerMatrix), ";\n"))

answerGrid = [
5 3 4 6 7 8 9 1 2;
6 7 2 1 9 5 3 4 8;
1 9 8 3 4 2 5 6 7;
8 5 9 7 6 1 4 2 3;
4 2 6 8 5 3 7 9 1;
7 1 3 9 2 4 8 5 6;
9 6 1 5 3 7 2 8 4;
2 8 7 4 1 9 6 3 5;
3 4 5 2 8 6 1 7 9]

answered = @time solve(grid) == answerGrid

println("Answered: ", answered)

puzzle2 = [ 5 0 0  0 0 4  1 0 0;
            0 4 6  0 0 9  0 0 0;
            9 0 1  0 6 0  8 3 4;

            7 0 5  0 4 3  0 1 0;
            0 0 4  0 8 0  3 0 0;
            0 1 0  6 5 0  4 0 9;

            1 8 2  0 7 0  9 0 3;
            0 0 0  1 0 0  5 4 0;
            0 0 7  3 0 0  0 0 1]

puzzle1 = grid
grid = puzzle2

# Simple Solve and then find all possible solutions
function solveAll(grid, p=fill(ALL_POSSIBLE, (GRID_SIZE, GRID_SIZE)), verbose=false)
    global possibles = p
    println(grid)
    @time solve(grid, p)
    if verbose
        printGrid(p)
    end
    printGrid(grid)
    @time findAllPossibles(grid, p, verbose)
    if verbose
        printGrid(p)
    end
    printGrid(grid)
    check_solved = isSolved(grid)
    println("Is Solved: $(check_solved)")
end

# Check that every row, column, and square has values 1 to 9.
function isSolved(grid)
    if any(grid .== 0)
        return false
    end
    fullSet = BitSet(1:9)
    for n in 1:9
        if BitSet(grid[n, 1:9]) != fullSet || BitSet(grid[1:9, n]) != fullSet
            return false
        end
    end
    for rowSq in [1:3, 4:6, 7:9]
        for colSq in [1:3, 4:6, 7:9]
            if BitSet(grid[rowSq, colSq]) != fullSet
                return false
            end
        end
    end
    return true
end

solveAll(grid)

grid = [
    0 0 0  9 7 0  0 0 6;
    0 7 0  0 0 6  8 0 1;
    0 0 2  0 8 3  0 0 0;

    0 8 6  0 2 0  0 0 0;
    0 0 1  0 9 0  6 0 0;
    0 0 0  0 3 0  2 4 0;

    0 0 0  7 4 0  1 0 0;
    9 0 4  8 0 0  0 6 0;
    7 0 0  0 6 9  0 0 0
]

solveAll(grid)

# TODO: Enter your own Sudoku


text:
using MRIgeneralizedBloch
using MRIgeneralizedBloch:propagator_linear_inversion_pulse
using StaticArrays
using Test

##
T = 500e-6
T2s = 10e-6

R2slT = precompute_R2sl()
R2sl = R2slT[1](T, π, 1, T2s)

function simulate_inversion_pulse(M0, ωz, B1, T, m0s, R1f, R2f, Rx, R1s, R2sl)
    Δ = π * (-1 + .0001:.0001:1) # * 2.17 / 2

    M  = zeros(6)
    for i in eachindex(Δ)
        Rz = exp(hamiltonian_linear(0, 1, Δ[i], 1, 0, 0, 0, 0, 0, 0)) # spoiler
        Ry = exp(hamiltonian_linear(π / T, B1, ωz, T, m0s, R1f, R2f, Rx, R1s, R2sl))
        M += Rz * (Ry * (Rz * M0))
    end
    M /= length(Δ)
end

u_rot = @SMatrix [cos(π) -sin(π) 0 0 0 0;
                  sin(π)  cos(π) 0 0 0 0;
                      0       0  1 0 0 0;
                      0       0  0 1 0 0;
                      0       0  0 0 1 0;
                      0       0  0 0 0 1]

## Test w/o exchange: in this case we expect equivalence except numerical differences
ωz = 0 # off resonance is implemented in the surrounding free presession part
m0s = 0
R1f = 1e-6
R2f = 1e-6
Rx = 1e-6
R1s = 1e-6

for i = 1:10
    local B1 = 1 + randn() / 4
    u_fp = exp(hamiltonian_linear(0, B1, ωz, T / 2, m0s, R1f, R2f, Rx, R1s, R2sl))
    UE = u_fp * propagator_linear_inversion_pulse(π / T, T, B1, R2sl, undef, undef, undef) * u_rot * u_fp

    M0 = SVector(1 - m0s, 0, 0, 0, m0s, 1)
    MB = simulate_inversion_pulse(M0, ωz, B1, T, m0s, R1f, R2f, Rx, R1s, R2sl)
    ME = UE * M0
    @test ME ≈ MB

    M0 = SVector(0, 1 - m0s, 0, 0, m0s, 1)
    MB = simulate_inversion_pulse(M0, ωz, B1, T, m0s, R1f, R2f, Rx, R1s, R2sl)
    ME = UE * M0
    @test ME ≈ MB

    M0 = SVector(0, 0, 1 - m0s, 0, m0s, 1)
    MB = simulate_inversion_pulse(M0, ωz, B1, T, m0s, R1f, R2f, Rx, R1s, R2sl)
    ME = UE * M0
    @test ME ≈ MB
end

## Test with exchange: in this case our approach is an approximation
m0s = 0.15
R1f = 0.3
R2f = 10
Rx = 30
R1s = 2

rtolmax = 1e-2
for i = 1:10
    local B1 = 1 + randn() / 4
    u_fp = exp(hamiltonian_linear(0, B1, ωz, T / 2, m0s, R1f, R2f, Rx, R1s, R2sl))
    UE = u_fp * propagator_linear_inversion_pulse(π / T, T, B1, R2sl, undef, undef, undef) * u_rot * u_fp

    M0 = SVector(1 - m0s, 0, 0, 0, m0s, 1)
    MB = simulate_inversion_pulse(M0, ωz, B1, T, m0s, R1f, R2f, Rx, R1s, R2sl)
    ME = UE * M0
    @test ME[[1,2,3,5,6]] ≈ MB[[1,2,3,5,6]] rtol = rtolmax

    M0 = SVector(0, 1 - m0s, 0, 0, m0s, 1)
    MB = simulate_inversion_pulse(M0, ωz, B1, T, m0s, R1f, R2f, Rx, R1s, R2sl)
    ME = UE * M0
    @test ME[[1,2,3,5,6]] ≈ MB[[1,2,3,5,6]] rtol = rtolmax

    M0 = SVector(0, 0, 1 - m0s, 0, m0s, 1)
    MB = simulate_inversion_pulse(M0, ωz, B1, T, m0s, R1f, R2f, Rx, R1s, R2sl)
    ME = UE * M0
    @test ME[[1,2,3,5,6]] ≈ MB[[1,2,3,5,6]] rtol = rtolmax
end


text:
#verilog-ranges.jl..
#creates ranges that help with making the julia verilog and the verilog much
#more interpretable.

doc"""
  v

  is an appendage that reverses the order of julia's ranges so that they have
  the cosmetic appearance of a verilog range.  This can be used both in setters
  and getters for wire collections.

  For example:
  `my_wire[5:0v]`  references a wire over indices 0-5
  `my_wire[6:2v]`  references a wire over indices 2-6

  For dereferencing (getting) you can reverse this:
  `my_wire[2:6v]`  references the wire over indices 2-6 with the msb order reversed.

  if you are using a variable, be sure to put the value in paretheses.

  For example:

  `my_wire[bits:(bits-4)v]`
"""
immutable v
  value
end

doc"""
  msb

  is a special keyword which indicates that the value should go to the most
  significant bit.  There is no equivalent "lsb", since in most cases this will
  be a fairly easy-to-reference value.  Plan your code accordingly.

  For example:
  my_wire[16:0v] = Wire(0xF000, 16)
  my_wire[msb:12]              # ==> Wire{3:0v}(0xF)

  my_wire_2[16:0v] = Wire(0x0F00, 16)
  my_wire[(msb-4):(msb-8)]     # ==> Wire{3:0v}(0xF)
"""
immutable msb;
  value::Integer
end

doc"""
  VerilogRange

  is a range object that represents a verilog range.  These are declared by
  using the v ranges, so 5:0v is a VerilogRange from 0 to 5.
"""
immutable VerilogRange <: AbstractUnitRange{Int64}
  start::Int64
  stop::Int64
end

doc"""
  RelativeRange

  is a range object that can take values which can be relative to the most
  significant bit.
"""
immutable RelativeRange
  start::Union{Int64, msb}
  stop::Union{Int64, msb}
end

function Base.:*(i, ::Type{v}); v(i); end
#allows

function Base.colon(i::Integer, vv::v)
  #first possibility, vv.value is actually the type msb.
  if vv.value == msb
    RelativeRange(msb(0), i)
  elseif isa(vv.value, msb)
    RelativeRange(vv.value, i)
  else
    VerilogRange(vv.value, i)
  end
end

function Base.colon(m::Union{Type{msb}, msb}, vv::v)
  left = (m == msb) ? msb(0) : m
  (vv.value == msb) ? RelativeRange(msb(0), left) : RelativeRange(vv.value, left)
end

function Base.:-(::Type{msb}, n::Integer)
  msb(n)
end
function Base.:-(m::msb, n::Integer)
  msb(m.value + n)
end
function Base.:+(m::msb, n::Integer)
  (n > m.value) && throw(ArgumentError("attempt to dereference beyond msb"))
  msb(m.value - n)
end

function parse_msb(rr::RelativeRange, vr::VerilogRange)
  true_start = isa(rr.start, msb) ? vr.stop - rr.start.value : rr.start
  true_stop  = isa(rr.stop, msb)  ? vr.stop - rr.stop.value  : rr.stop
  true_stop:(true_start)v
end

export v, msb

#because verilog often uses zero-indexing, a python-style range() operator
#is helpful for creating for loop generators.

Base.range(i::Integer) = (i-1):0v

#making a VerilogRange a well-defined iterable.
Base.start(v::VerilogRange) = v.start
Base.next(v::VerilogRange, state) = (state, state + ((v.stop >= v.start) ? 1 : - 1))
Base.done(v::VerilogRange, state) = state == v.stop + ((v.stop >= v.start) ? 1 : -1)
Base.eltype(v::VerilogRange) = Int64
Base.length(v::VerilogRange) = ((v.stop > v.start) ? (v.stop - v.start) : (v.start - v.stop)) + 1

Base.show(io::IO, v::VerilogRange) = print(io, v.stop, ":", v.start, "v")

function Base.reverse(r::VerilogRange)
  if r.start < r.stop
    r.stop:-1:r.start
  else
    r.stop:r.start
  end
end


text:
using DistributedControl
using Test

@testset "DistributedControl.jl" begin
    # Write your tests here.
end


text:
# extension of merge to EventTraceData fields
function merge_ext!(S::EventTraceData, Ω::Int64, rest::Array{Int64, 1})
  pha   = PhaseCat()
  az    = getindex(getfield(S, :az), Ω)
  baz   = getindex(getfield(S, :baz), Ω)
  dist  = getindex(getfield(S, :dist), Ω)
  for i in rest
    if az == 0.0
      θ = getindex(getfield(S, :az), i)
      (θ != 0.0) && (az = θ)
    end
    if baz == 0.0
      β = getindex(getfield(S, :baz), i)
      (β != 0.0) && (baz = β)
    end
    if dist == 0.0
      # Δ is already in use, so...
      d = getindex(getfield(S, :dist), i)
      (d != 0.0) && (dist = d)
    end
    merge!(pha, getindex(getfield(S, :pha), i))
  end
  # This guarantees that the phase catalog of Ω overwrites others
  merge!(pha, getindex(getfield(S, :pha), Ω))
  setindex!(getfield(S, :az),     az, Ω)
  setindex!(getfield(S, :baz),   baz, Ω)
  setindex!(getfield(S, :dist), dist, Ω)
  setindex!(getfield(S, :pha),   pha, Ω)
  return nothing
end

# Home of all extended merge! methods
merge(S::EventTraceData; v::Integer=KW.v) = (U = deepcopy(S); merge!(U, v=v); return U)
merge!(S::EventTraceData, U::EventTraceData; v::Integer=KW.v) = ([append!(getfield(S, f), getfield(U, f)) for f in tracefields]; S.n += U.n; merge!(S; v=v))
merge!(S::EventTraceData, C::EventChannel; v::Integer=KW.v) = merge!(S, EventTraceData(C), v=v)

function merge(A::Array{EventTraceData,1}; v::Integer=KW.v)
  L::Int64 = length(A)
  n = sum([A[i].n for i = 1:L])
  T = EventTraceData(n)
  [setfield!(T, f, vcat([getfield(A[i],f) for i = 1:L]...)) for f in SeisIO.datafields]
  merge!(T, v=v)
  return T
end
merge(S::EventTraceData, U::EventTraceData; v::Integer=KW.v) = merge(Array{EventTraceData,1}([S,U]), v=v)
merge(S::EventTraceData, C::EventChannel; v::Integer=KW.v) = merge(S, EventTraceData(C), v=v)
merge(C::EventChannel, S::EventTraceData; v::Integer=KW.v) = merge(EventTraceData(C), S, v=v)
merge(C::EventChannel, D::EventChannel; v::Integer=KW.v) = (S = EventTraceData(C,D); merge!(S, v=v); return S)

+(S::EventTraceData, C::EventChannel) = +(S, EventTraceData(C))
+(C::EventChannel, S::EventTraceData) = +(S, EventTraceData(C))
+(C::EventChannel, D::EventChannel) = +(EventTraceData(C), EventTraceData(D))

# Multiplication
# distributivity: (S1+S2)*S3) == (S1*S3 + S2*S3)
*(S::EventTraceData, U::EventTraceData) = merge(Array{EventTraceData,1}([S,U]))
*(S::EventTraceData, C::EventChannel) = merge(S, EventTraceData(C))
function *(C::EventChannel, D::EventChannel)
  s1 = deepcopy(C)
  s2 = deepcopy(D)
  S = merge(EventTraceData(s1),EventTraceData(s2))
  return S
end


text:
using Documenter
using DocExample
Documenter.makedocs(root = "./",
	source = "src",
	build = "build",
	clean = true,
	doctest = true,
	modules = Module[DocExample],
	#repo = "github.com/alenskorobogatova/DocExample.jl.git",
	highlightsig = true,
	sitename = "Example Documentation",
	expandfirst = [],
	pages = [
		"Text Function" => "textIndex.md",
		"Math Function" => "mathIndex.md",
		]
	)
Documenter.deploydocs(
    repo = "github.com/alenskorobogatova/DocExample.jl.git",
    target = "build",
    deps = nothing,
    make = nothing
)


text:
module TestSBL
using Test
using LinearAlgebra
using CompressedSensing: sparse_data, sbl, SBL, fsbl, rmps, update!, rmp, foba, perturb
using SparseArrays

n, m, k = 32, 48, 3
A, x, b = sparse_data(n = n, m = m, k = k)
σ = 1e-2
y = perturb(b, σ/2)

@testset "Sparse Bayesian Learning" begin
    xsbl = sbl(A, y, σ^2) # need to pass variance or covariance of noise
    tol = σ
    # xsbl = droptol!(sparse(xsbl), tol)
    @test findall(abs.(xsbl) .> tol) == x.nzind
    @test isapprox(A*xsbl, b, atol = σ)

    # fast sparse bayesian learning
    xfsbl = fsbl(A, y, σ^2)
    @test findall(abs.(xfsbl) .> tol) == x.nzind
    @test isapprox(A*xfsbl, b, atol = σ)

    # relevance matching pursuit
    xrmps = rmps(A, y, σ^2)
    @test findall(abs.(xrmps) .> tol) == x.nzind
    @test isapprox(A*xrmps, b, atol = σ)

    # rmp with noise variance optimization
    σ_init = σ
    xrmps, σ²_opt = rmps(A, y, Val(true), σ_init^2)
    @test xrmps isa SparseVector
    @test σ²_opt isa Real
    @test norm(A*xrmps-y) < 5sqrt(σ²_opt) * size(A, 1)

    # rmp with noise variance optimization with prior over σ
    σ_init = σ
    xrmps, σ²_opt = rmps(A, y, Val(true), σ_init^2, 1, σ_init^2)
    @test norm(A*xrmps-y) < 5sqrt(σ²_opt) * size(A, 1)
    @test isapprox(σ²_opt, σ^2, rtol = 2) # noise variance is approximately recovered

    # zero noise limit of relevance matching pursuit
    xrmp = rmp(A, y, σ)
    @test xrmp.nzind == x.nzind
    @test isapprox(A*xrmp, b, atol = σ)

end

# bench = false
# using BenchmarkTools
# if bench
#     @btime sbl($A, $b, $(σ^2))
#     @btime rmps($A, $b, $(σ^2))
#     @btime fsbl($A, $b, $(σ^2))
#     @btime rmp($A, $y, $(σ^2))
#     @btime foba($A, $y, $(σ^2))
# end

end # TestSBL


text:
module TestFinalizers

using Test
using AsyncFinalizers
using AsyncFinalizers.Internal: fallback_finalize

using ..Utils: check_executor, outlined, collect_garbage

function test_simple()
    collect_garbage()
end

function test_once()
    ncalls = Threads.Atomic{Int}(0)
    function async_finalizer(_)
        Threads.atomic_add!(ncalls, 1)
    end
    @testset "Initial collect_garbage" begin
        collect_garbage(; async_finalizer = async_finalizer)
    end
    @testset for trial in 1:5
        collect_garbage()
    end
    @test ncalls[] == 1
end

function test_lock()
    check_executor()
    a = 0
    guard = ReentrantLock()
    outlined() do
        local r = Ref(0)
        AsyncFinalizers.onfinalize(r) do _
            function ()
                lock(guard) do
                    a = 1
                end
            end
        end
    end
    GC.gc()

    for _ in 1:1000
        lock(guard) do
            a == 1
        end && break
        sleep(0.01)
    end

    @test a == 1
end

function test_fallback()
    counter = Ref(0)
    f() = counter[] += 1
    T = Int
    @test_logs (:warn, r"queue not initalized") wait(fallback_finalize(f, T)::Task)
    @test counter[] == 1
end

end  # module


text:
struct Constraint{op,L,R}
    lhs::L
    rhs::R
end

Constraint{op}(lhs::L, rhs::R) where {op, L, R} = Constraint{op,L,R}(lhs, rhs)

function print_op(op)
    (op == ≠) && return "≠"
    (op == ~) && return "="
end

Base.show(io::IO, c::Constraint{op}) where {op} = print(io, c.lhs, " ", print_op(op), " ", c.rhs)

function clauses(c::Constraint{!=})

    # ¬(x[i] ∧ y[i])  -- can't take same value

    [¬(c.lhs[i]) ∨ ¬(c.rhs[i]) for i in 1:min(max(c.lhs), max(c.rhs))]
end

function clauses(c::Constraint{~})
    i = c.rhs  # integer
    [c.lhs[i]]  # that Boolean is positive
end

Base.:≠(x::BoundedInteger, y::BoundedInteger) = Constraint{≠}(x, y)
Base.:~(x::BoundedInteger, y::Int) = Constraint{~}(x, y)

macro integer(name, k)
    quoted = Meta.quot(name)

    :($(esc(name)) = BoundedInteger{$k}($quoted))
end

make_vector(name, T, num) = [T(Symbol(name, subscript(i))) for i in 1:num]


text:
#=##############################################################################################
Filename: plot_lift_moment_distribution.jl
Author: Ryan Anderson
Contact: rymanderson@gmail.com
README: `<: Action` function plots the lift distribution of all lifting surfaces at the specified steps
=###############################################################################################

"""
    plot_lift_moment_distribution(aircraft, parameters, freestream, environment, step_range, stepi, step_symbol) <: Action

# Arguments:

* `aircraft::Aircraft`: `Aircraft` system object
* `parameters<:Parameters`: `Parameters` struct
* `freestream::Freestream`: `Freestream` object
* `environment::Environment` `Environment` object
* `step_range::AbstractArray`: array of step for which the simulation is run
* `stepi::Int`: index of the current step
* `step_symbol::String`: defines the step, e.g. `alpha` or `time`

`parameters <: Parameters` requires the following elements:

* `surfacenames::Vector{String}`: names of each lifting surface to be shown in the legend
* `cfs::Vector{Array{Float64,2}}`: vector with length equal to the number of lifting surfaces, each member containing an array of size (3,ns) of force coefficients cd, cy, cl
* `cms::Vector{Array{Float64,2}}`: vector with length equal to the number of lifting surfaces, each member containing an array of size (3,ns) of moment coefficients cmx, cmy, cmz
* `plot_directory::String`: path to the folder where plots will be saved
* `plot_base_name::String`: first part of saved figure file names
* `plot_extension::String`: extension of saved figure file names
* `plotstepi::Vector{Int}`: which steps at which to plot
"""
function plot_lift_moment_distribution(aircraft, parameters, freestream, environment, step_range, stepi, step_symbol)

    # extract plot indices
    plotstepi = parameters.plotstepi
    basename = parameters.plot_base_name

    # cycle through styles
    styles = ["-", "--", "--^", "--v"]

    # extract main wing span
    b = aircraft.wing_system.system.surfaces[1][end].rtr[2] * 2

    if stepi in plotstepi

        # extract info
        surfacenames = parameters.surfacenames
        cfs = parameters.cfs
        cms = parameters.cms
        lifting_line_rs = aircraft.wing_system.lifting_line_rs
        plot_directory = parameters.plot_directory
        plot_base_name = parameters.plot_base_name
        plot_extension = parameters.plot_extension
        # cl_ylim = parameters.cl_ylim
        # cd_ylim = parameters.cd_ylim
        # cy_ylim = parameters.cy_ylim

        # create axes
        fig_lift_distribution = plt.figure(basename * "_lift_distribution")
        fig_cf_distribution = plt.figure(basename * "_cf_distribution")
        nsurfaces = length(surfacenames)
        nsubplotbase = 300 + 10 * nsurfaces
        if stepi == parameters.plotstepi[1]
            fig_lift_distribution.clear()
            fig_lift_distribution.add_subplot(111, ylabel = L"c_l", xlabel = L"2y/b")
            fig_cf_distribution.clear()
            for isurface = 1:nsurfaces
                fig_cf_distribution.add_subplot(nsubplotbase + 2 + (isurface-1) * nsurfaces, ylabel = L"c_d") # drag ylim = cd_ylim,
                fig_cf_distribution.add_subplot(nsubplotbase + 3 + (isurface-1) * nsurfaces, ylabel = L"c_y", xlabel = L"y [m]") # side force ylim = cy_ylim,
                fig_cf_distribution.add_subplot(nsubplotbase + 1 + (isurface-1) * nsurfaces, ylabel = L"c_l", title = surfacenames[isurface]) # lift ylim = cl_ylim,
                # fig_cf_distribution.suptitle("t = $(step_range[stepi]), ti = $stepi") # add title

                # set axes
                # axs[1].set_ylim(cd_ylim)
                # axs[1].set_ylabel(L"c_d")
                # # axs[1].legend()
                # axs[2].set_ylim(cy_ylim)
                # axs[2].set_ylabel(L"c_y")
                # axs[2].set_xlabel(L"y [m]")
                # axs[2].legend()
                # axs[3].set_ylim(cl_ylim)
                # axs[3].set_ylabel(L"c_l")
            end
        end
        ax_cl = fig_lift_distribution.get_axes()[1]
        axs_cf = fig_cf_distribution.get_axes()

        # get color
        cratio = findfirst((x)->x==stepi,plotstepi) / length(plotstepi)

        # plot
        for (isurface,cf) in enumerate(cfs[stepi])
            rs_plot = get_midpoints(lifting_line_rs[isurface][2,:])
            ax_cl.plot(rs_plot ./ b * 2, cf[3,:], styles[isurface], color = (0.05, 0.85-cratio*0.7, 0.15 + 0.75 * cratio), label="$(surfacenames[isurface]), $step_symbol = $(round(step_range[stepi],digits=3))")
            for icf = 1:3
                axs_cf[(isurface - 1) * nsurfaces + icf].plot(rs_plot, cf[icf,:], color=(0.05, 0.85-cratio*0.7, 0.15 + 0.75 * cratio), label="$step_symbol = $(round(step_range[stepi],digits=3))")
            end
        end

        # save
        if stepi == plotstepi[end] # last step
            ax_cl.legend(loc="upper left", bbox_to_anchor=(1.01,1))
            fig_lift_distribution.set_size_inches(10, 6, forward=true)
            fig_lift_distribution.tight_layout()
            savepath_cl = joinpath(plot_directory, plot_base_name * "_lift_distribution" * plot_extension)
            fig_lift_distribution.savefig(savepath_cl, bbox_inches="tight")
            axs_cf[3].legend(loc="upper left", bbox_to_anchor=(1.01,1)) # set legend
            fig_cf_distribution.set_size_inches(6, 10, forward=true)
            fig_cf_distribution.tight_layout() # clean up white space
            savepath_cf = joinpath(plot_directory, plot_base_name * "_cf_distribution" * plot_extension)
            fig_cf_distribution.savefig(savepath_cf, bbox_inches="tight")
        end

        # cms plot
        fig_moment_distribution = plt.figure(basename * "_moment_distribution")
        nsurfaces = length(surfacenames)
        nsubplotbase = 300 + 10 * nsurfaces
        if stepi == parameters.plotstepi[1]
            fig_moment_distribution.clear()
            for isurface = 1:nsurfaces
                fig_moment_distribution.add_subplot(nsubplotbase + 2 + (isurface-1) * nsurfaces, ylabel = L"c_{mx}") # drag ylim = cd_ylim,
                fig_moment_distribution.add_subplot(nsubplotbase + 1 + (isurface-1) * nsurfaces, ylabel = L"c_{my}", title = surfacenames[isurface]) # lift ylim = cl_ylim,
                fig_moment_distribution.add_subplot(nsubplotbase + 3 + (isurface-1) * nsurfaces, ylabel = L"c_{mz}", xlabel = L"y [m]") # side force ylim = cy_ylim,
                # fig_moment_distribution.suptitle("t = $(step_range[stepi]), ti = $stepi") # add title

                # set axes
                # axs[1].set_ylim(cd_ylim)
                # axs[1].set_ylabel(L"c_d")
                # # axs[1].legend()
                # axs[2].set_ylim(cy_ylim)
                # axs[2].set_ylabel(L"c_y")
                # axs[2].set_xlabel(L"y [m]")
                # axs[2].legend()
                # axs[3].set_ylim(cl_ylim)
                # axs[3].set_ylabel(L"c_l")
            end
        end

        # plot
        axs = fig_moment_distribution.get_axes()
        for (isurface,cm) in enumerate(cms[stepi])
            rs_plot = get_midpoints(lifting_line_rs[isurface][2,:])
            for icf = 1:3
                axs[(isurface - 1) * nsurfaces + icf].plot(rs_plot, cm[icf,:], color=(0.05, 0.85-cratio*0.7, 0.15 + 0.75 * cratio), label="$step_symbol = $(round(step_range[stepi],digits=3))")
            end
        end

        # save
        if stepi == plotstepi[end]
            axs[2].legend(loc="upper left", bbox_to_anchor=(1.01,1)) # set legend
            fig_moment_distribution.set_size_inches(6, 10, forward=true)
            fig_moment_distribution.tight_layout() # clean up white space
            savepath = joinpath(plot_directory, plot_base_name * "_moment_distribution" * plot_extension)
            fig_moment_distribution.savefig(savepath, bbox_inches="tight")
        end
    end

    return false
end

"""
plot_lift_moment_distribution(aircraft, step_range)

Method returns initialized elements required for the `parameters <: Parameters` struct during simulation.

# Arguments:

* `aircraft::Aircraft`: system to be simulated
* `step_range::AbstractArray`: defines each step of the simulation

# Returns:

* `surfacenames::Vector{String}`: names of each lifting surface to be shown in the legend
* `cfs::Vector{String}`: vector with length equal to the number of lifting surfaces, each member containing an array of size (3,ns) of force coefficients
* `plot_directory::String`: path to the folder where plots will be saved
* `plot_base_name::String`: first part of saved figure file names
* `plot_extension::String`: extension of saved figure file names
"""
function plot_lift_moment_distribution(aircraft, step_range)

    nwings = length(aircraft.wing_system.surfaces)
    surfacenames = Vector{String}(undef,nwings)
    cfs = [[zeros(3, length(aircraft.wing_system.surfaces[i])) for i in 1:nwings] for j in 1:length(step_range)]
    cms = deepcopy(cfs)
    plot_directory = ""
    plot_base_name = "default"
    plot_extension = ".pdf"
    plotstepi = deepcopy(step_range)

    return surfacenames, cfs, cms, plot_directory, plot_base_name, plot_extension, plotstepi#, cl_ylim, cd_ylim, cy_ylim
end

max_stars_repo_name: byuflowlab/AircraftSystems

max_issues_repo_name: byuflowlab/AircraftSystems

max_forks_repo_name: byuflowlab/AircraftSystems


text:
module StandardAtmosphere

export Atm, Env, kcas, keas, ktas, mach

import Base: show, *

using SIUnits
using SIUnits.ShortUnits

include("atm.jl")
include("env.jl")

end # module


text:
"""
Tmp

Example module from Julia workflow tip https://docs.julialang.org/en/v1/manual/workflow-tips/

Use this module by running Tst.jl
"""
module Tmp

export sayhello

function sayhello()
    println("Hello")
end
end