text:
# Copyright (c) 2017: Miles Lubin and contributors
# Copyright (c) 2017: Google Inc.
#
# Use of this source code is governed by an MIT-style license that can be found
# in the LICENSE.md file or at https://opensource.org/licenses/MIT.

"""
    abstract type AbstractBridge <: MOI.Bridges.AbstractType

Subtype of [`MOI.Bridges.AbstractBridge`](@ref) for constraint
bridges.
"""
abstract type AbstractBridge <: MOI.Bridges.AbstractBridge end

"""
    MOI.supports_constraint(
        BT::Type{<:AbstractBridge},
        F::Type{<:MOI.AbstractFunction},
        S::Type{<:MOI.AbstractSet},
    )::Bool

Return a `Bool` indicating whether the bridges of type `BT` support bridging
`F`-in-`S` constraints.

## Implementation notes

 * This method depends only on the type of the inputs, not the runtime values.
 * There is a default fallback, so you need only implement this method for
   constraint types that the bridge implements.
"""
function MOI.supports_constraint(
    ::Type{<:AbstractBridge},
    ::Type{<:MOI.AbstractFunction},
    ::Type{<:MOI.AbstractSet},
)
    return false
end

"""
    concrete_bridge_type(
        BT::Type{<:AbstractBridge},
        F::Type{<:MOI.AbstractFunction},
        S::Type{<:MOI.AbstractSet}
    )::Type

Return the concrete type of the bridge supporting `F`-in-`S` constraints.

This function can only be called if `MOI.supports_constraint(BT, F, S)` is
`true`.

## Example

The [`SplitIntervalBridge`](@ref) bridges a [`MOI.VariableIndex`](@ref)-in-[`MOI.Interval`](@ref)
constraint into a [`MOI.VariableIndex`](@ref)-in-[`MOI.GreaterThan`](@ref) and a
[`MOI.VariableIndex`](@ref)-in-[`MOI.LessThan`](@ref) constraint.

```jldoctest; setup=:(using MathOptInterface; const MOI = MathOptInterface)
julia> MOI.Bridges.Constraint.concrete_bridge_type(
           MOI.Bridges.Constraint.SplitIntervalBridge{Float64},
           MOI.VariableIndex,
           MOI.Interval{Float64},
       )
MathOptInterface.Bridges.Constraint.SplitIntervalBridge{Float64, MathOptInterface.VariableIndex, MathOptInterface.Interval{Float64}, MathOptInterface.GreaterThan{Float64}, MathOptInterface.LessThan{Float64}}
```
"""
function concrete_bridge_type(
    ::Type{BT},
    ::Type{<:MOI.AbstractFunction},
    ::Type{<:MOI.AbstractSet},
) where {BT}
    return BT
end

function concrete_bridge_type(
    b::MOI.Bridges.AbstractBridgeOptimizer,
    F::Type{<:MOI.AbstractFunction},
    S::Type{<:MOI.AbstractSet},
)
    return concrete_bridge_type(MOI.Bridges.bridge_type(b, F, S), F, S)
end

"""
    bridge_constraint(
        BT::Type{<:AbstractBridge},
        model::MOI.ModelLike,
        func::AbstractFunction,
        set::MOI.AbstractSet,
    )::BT

Bridge the constraint `func`-in-`set` using bridge `BT` to `model` and returns
a bridge object of type `BT`.

## Implementation notes

 * The bridge type `BT` should be a concrete type, that is, all the type
   parameters of the bridge must be set.
"""
function bridge_constraint end


text:
struct GradConformity <: Conformity end
const H1Conformity = GradConformity

function get_face_own_nodes(reffe::GenericLagrangianRefFE{GradConformity},conf::GradConformity)
  p = get_polytope(reffe)
  orders = get_orders(reffe)
  nodes, face_own_nodes = compute_nodes(p,orders)
  face_own_nodes
end

function get_own_nodes_permutations(reffe::GenericLagrangianRefFE{GradConformity},conf::GradConformity)
  p = get_polytope(reffe)
  face_own_nodes = get_face_own_nodes(reffe)
  dofs = get_dof_basis(reffe)
  interior_nodes = dofs.nodes[face_own_nodes[end]]
  compute_own_nodes_permutations(p,interior_nodes)
end

function get_face_own_nodes_permutations(reffe::GenericLagrangianRefFE{GradConformity},conf::GradConformity)
  own_nodes_permutations = get_own_nodes_permutations(reffe)
  reffaces = reffe.reffe.metadata
  _reffaces = vcat(reffaces...)
  face_own_nodes_permutations = map(get_own_nodes_permutations,_reffaces)

  _compute_face_own_nodes_permutations(
    num_nodes(reffe),
    get_face_own_nodes(reffe),
    face_own_nodes_permutations,
    own_nodes_permutations)
end

function _compute_face_own_nodes_permutations(
  nnodes,
  face_own_nodes,
  face_own_nodes_permutations,
  own_nodes_permutations)

  if length(face_own_nodes_permutations) == 0
    # Vertex degenerated case
    return [own_nodes_permutations,]
  else
    if nnodes == length(face_own_nodes[end])
      # 0-order degenerated case
      _face_own_nodes_permutations = map( (x) -> fill(Int[],length(x)) , face_own_nodes_permutations )
    else
      # Standard case
      _face_own_nodes_permutations = copy(face_own_nodes_permutations)
    end
    push!(_face_own_nodes_permutations,own_nodes_permutations)
    return _face_own_nodes_permutations
  end
end

# API particular to LagrangianRefFE{GradConformity}

"""
    ReferenceFE{N}(reffe::GenericLagrangianRefFE{GradConformity},iface::Integer) where N
"""
function ReferenceFE{N}(reffe::GenericLagrangianRefFE{GradConformity},iface::Integer) where N
  reffaces = reffe.reffe.metadata
  reffaces[N+1][iface]
end

function ReferenceFE{D}(reffe::GenericLagrangianRefFE{GradConformity,D},iface::Integer) where D
  @assert iface==1 "Only one D-face"
  reffe
end

"""
    get_reffaces(
      ::Type{ReferenceFE{d}},
      reffe::GenericLagrangianRefFE{GradConformity}) where d -> Vector{GenericLagrangianRefFE{GradConformity,M,d}}
"""
function get_reffaces(::Type{ReferenceFE{d}},reffe::GenericLagrangianRefFE{GradConformity}) where d
  ftype_to_reffe, _ = _compute_reffes_and_face_types(reffe,Val{d}())
  ftype_to_reffe
  [reffe for reffe in ftype_to_reffe]
end

"""
    get_face_type(reffe::GenericLagrangianRefFE{GradConformity}, d::Integer) -> Vector{Int}
"""
function get_face_type(reffe::GenericLagrangianRefFE{GradConformity}, d::Integer)
  _, iface_to_ftype = _compute_reffes_and_face_types(reffe,Val{d}())
  iface_to_ftype
end

function _compute_reffes_and_face_types(reffe::GenericLagrangianRefFE{GradConformity},::Val{d}) where d
  p = get_polytope(reffe)
  iface_to_reffe = [ ReferenceFE{d}(reffe,iface) for iface in 1:num_faces(p,d) ]
  _find_unique_with_indices(iface_to_reffe)
end

"""
    is_first_order(reffe::GenericLagrangianRefFE{GradConformity}) -> Bool
"""
function is_first_order(reffe::GenericLagrangianRefFE{GradConformity})
  p = get_polytope(reffe)
  r = true
  r = r && num_vertices(p) == num_nodes(reffe)
  r = r && get_vertex_node(reffe) == collect(1:num_nodes(reffe))
  r
end

"""
    is_affine(reffe::GenericLagrangianRefFE{GradConformity}) -> Bool

Query if the `reffe` leads to an afine map
(true only for first order spaces on top of simplices)
"""
function is_affine(reffe::GenericLagrangianRefFE{GradConformity})
  p = get_polytope(reffe)
  is_first_order(reffe) && is_simplex(p)
end

"""
    is_P(reffe::GenericLagrangianRefFE{GradConformity})
"""
function is_P(reffe::GenericLagrangianRefFE{GradConformity})
  is_simplex(get_polytope(reffe))
end

"""
   is_Q(reffe::GenericLagrangianRefFE{GradConformity})
"""
function is_Q(reffe::GenericLagrangianRefFE{GradConformity})
  monomials = get_prebasis(reffe)
  n = length(get_exponents(monomials))
  is_n_cube(get_polytope(reffe)) && (prod(get_orders(reffe).+1) == n)
end

"""
   is_S(reffe::GenericLagrangianRefFE{GradConformity})
"""
function is_S(reffe::GenericLagrangianRefFE{GradConformity})
  is_n_cube(get_polytope(reffe)) && ! is_Q(reffe)
end

function to_dict(reffe::GenericLagrangianRefFE{GradConformity})
  p = get_polytope(reffe)
  b = get_prebasis(reffe)
  dict = Dict{Symbol,Any}()
  dict[:orders] = collect(get_orders(reffe))
  dict[:extrusion] = Array(TensorValues.get_array(get_extrusion(p)))
  if is_S(reffe)
    dict[:space] = "serendipity"
  else
    dict[:space] = "default"
  end
  dict[:value] = string(get_value_type(b))
  dict
end

function from_dict(::Type{<:LagrangianRefFE},dict::Dict{Symbol,Any})
  orders = Tuple(dict[:orders])
  extrusion = Tuple(dict[:extrusion])
  if dict[:value] == "Float64"
    value = Float64
  else
    @notimplemented
  end
  space = dict[:space]
  p = Polytope(extrusion...)
  if space == "default"
    reffe = LagrangianRefFE(value,p,orders)
  elseif space == "serendipity"
    reffe = SerendipityRefFE(value,p,orders)
  else
    @unreachable "unknown space type"
  end
  reffe
end

# Construction of LagrangianRefFE from Polytopes

"""
    LagrangianRefFE(::Type{T},p::Polytope,orders) where T
    LagrangianRefFE(::Type{T},p::Polytope,order::Int) where T

Builds a `LagrangianRefFE` object on top of the given polytope. `T` is the type of
the value of the approximation space (e.g., `T=Float64` for scalar-valued problems,
`T=VectorValue{N,Float64}` for vector-valued problems with `N` components). The arguments `order` or `orders`
are for the polynomial order of the resulting space, which allows isotropic or anisotropic orders respectively
(provided that the cell topology allows the given anisotropic order). The argument `orders` should be an
indexable collection of `D` integers (e.g., a tuple or a vector), being `D` the number of space dimensions.

In order to be able to use this function, the type of the provided polytope `p` has to implement the
following additional methods. They have been implemented for `ExtrusionPolytope` in the library. They
need to be implemented for new polytope types in order to build Lagangian reference elements on top of them.

- [`compute_monomial_basis(::Type{T},p::Polytope,orders) where T`](@ref)
- [`compute_own_nodes(p::Polytope,orders)`](@ref)
- [`compute_face_orders(p::Polytope,face::Polytope,iface::Int,orders)`](@ref)

The following methods are also used in the construction of the `LagrangianRefFE` object. A default implementation
of them is available in terms of the three previous methods. However, the user can also implement them for
new polytope types increasing customization possibilities.

- [`compute_nodes(p::Polytope,orders)`](@ref)
- [`compute_own_nodes_permutations(p::Polytope, interior_nodes)`](@ref)
- [`compute_lagrangian_reffaces(::Type{T},p::Polytope,orders) where T`](@ref)
"""
function LagrangianRefFE(::Type{T},p::Polytope{D},orders;space::Symbol=_default_space(p)) where {T,D}
  if space == :P && is_n_cube(p)
    return _PDiscRefFE(T,p,orders)
  elseif space == :S && is_n_cube(p)
    SerendipityRefFE(T,p,orders)
  else
    if any(orders.==0) && !all(orders.==0)
      cont = map(i -> i == 0 ? DISC : CONT,orders)
      return _cd_lagrangian_ref_fe(T,p,orders,cont)
    else
      return _lagrangian_ref_fe(T,p,orders)
    end
  end
end

function _default_space(p)
  if is_n_cube(p)
    :Q
  else
    :P
  end
end

function _lagrangian_ref_fe(::Type{T},p::Polytope{D},orders) where {T,D}

  prebasis = compute_monomial_basis(T,p,orders)
  nodes, face_own_nodes = compute_nodes(p,orders)
  dofs = LagrangianDofBasis(T,nodes)
  reffaces = compute_lagrangian_reffaces(T,p,orders)

  nnodes = length(dofs.nodes)
  ndofs = length(dofs.dof_to_node)
  metadata = reffaces
  _reffaces = vcat(reffaces...)
  face_nodes = _generate_face_nodes(nnodes,face_own_nodes,p,_reffaces)
  face_own_dofs = _generate_face_own_dofs(face_own_nodes, dofs.node_and_comp_to_dof)
  face_dofs = _generate_face_dofs(ndofs,face_own_dofs,p,_reffaces)

  if all(orders .== 0 ) && D>0
    conf = L2Conformity()
  else
    conf = GradConformity()
  end

  reffe = GenericRefFE(
    ndofs,
    p,
    prebasis,
    dofs,
    conf,
    metadata,
    face_dofs)

  GenericLagrangianRefFE(reffe,face_nodes)

end

function MonomialBasis(::Type{T},p::Polytope,orders) where T
  compute_monomial_basis(T,p,orders)
end

function LagrangianDofBasis(::Type{T},p::Polytope,orders) where T
  nodes, _ = compute_nodes(p,orders)
  LagrangianDofBasis(T,nodes)
end

# Helpers for LagrangianRefFE constructor

function _generate_face_nodes(nnodes,face_to_own_nodes,polytope,reffaces)

    face_to_num_fnodes = map(num_nodes,reffaces)
    push!(face_to_num_fnodes,nnodes)

    face_to_lface_to_own_fnodes = map(get_face_own_nodes,reffaces)
    push!(face_to_lface_to_own_fnodes,face_to_own_nodes)

    face_to_lface_to_face = get_faces(polytope)

  _generate_face_nodes_aux(
    nnodes,
    face_to_own_nodes,
    face_to_num_fnodes,
    face_to_lface_to_own_fnodes,
    face_to_lface_to_face)
end

function _generate_face_dofs(ndofs,face_to_own_dofs,polytope,reffaces)

    face_to_num_fdofs = map(num_dofs,reffaces)
    push!(face_to_num_fdofs,ndofs)

    face_to_lface_to_own_fdofs = map(get_face_own_dofs,reffaces)
    push!(face_to_lface_to_own_fdofs,face_to_own_dofs)

    face_to_lface_to_face = get_faces(polytope)

  _generate_face_nodes_aux(
    ndofs,
    face_to_own_dofs,
    face_to_num_fdofs,
    face_to_lface_to_own_fdofs,
    face_to_lface_to_face)
end

function _generate_face_nodes_aux(
  nnodes,
  face_to_own_nodes,
  face_to_num_fnodes,
  face_to_lface_to_own_fnodes,
  face_to_lface_to_face)

  if nnodes == length(face_to_own_nodes[end])
    face_fnode_to_node = fill(Int[],length(face_to_own_nodes))
    face_fnode_to_node[end] = collect(1:nnodes)
    return face_fnode_to_node
  end

  face_fnode_to_node = Vector{Int}[]
  for (face, nfnodes) in enumerate(face_to_num_fnodes)
    fnode_to_node = zeros(Int,nfnodes)
    lface_to_face = face_to_lface_to_face[face]
    lface_to_own_fnodes = face_to_lface_to_own_fnodes[face]
    for (lface, faceto) in enumerate(lface_to_face)
      own_nodes = face_to_own_nodes[faceto]
      own_fnodes = lface_to_own_fnodes[lface]
      fnode_to_node[own_fnodes] = own_nodes
    end
    push!(face_fnode_to_node,fnode_to_node)
  end

  face_fnode_to_node
end

# Constructors taking Int

function LagrangianRefFE(::Type{T},p::Polytope{D},order::Int;space::Symbol=_default_space(p)) where {T,D}
  orders = tfill(order,Val{D}())
  LagrangianRefFE(T,p,orders;space=space)
end

function MonomialBasis(::Type{T},p::Polytope{D},order::Int) where {D,T}
  orders = tfill(order,Val{D}())
  MonomialBasis(T,p,orders)
end

function LagrangianDofBasis(::Type{T},p::Polytope{D},order::Int) where {T,D}
  orders = tfill(order,Val{D}())
  LagrangianDofBasis(T,p,orders)
end

# Queries needed to be implemented for polytopes in order to use them
# for building LagrangianRefFEs in a seamless way

"""
    compute_monomial_basis(::Type{T},p::Polytope,orders) where T -> MonomialBasis

Returns the monomial basis of value type `T` and order per direction described by `orders`
on top of the polytope `p`.
"""
function compute_monomial_basis(::Type{T},p::Polytope,orders) where T
  @abstractmethod
end

"""
    compute_own_nodes(p::Polytope{D},orders) where D -> Vector{Point{D,Float64}}

Returns the coordinates of the nodes owned by the interior of the polytope
associated with a Lagrangian space with the order per direction described by `orders`.
"""
function compute_own_nodes(p::Polytope,orders)
  @abstractmethod
end

"""
    compute_face_orders(p::Polytope,face::Polytope,iface::Int,orders)

Returns a vector or a tuple with the order per direction at the face `face`
of the polytope `p` when restricting the order per direction `orders` to this face.
`iface` is the face id of `face` in the numeration restricted to the face dimension.
"""
function compute_face_orders(p::Polytope,face::Polytope,iface::Int,orders)
  @abstractmethod
end

"""
    compute_nodes(p::Polytope,orders)

When called

    node_coords, face_own_nodes = compute_nodes(p,orders)

Returns `node_coords`, the nodal coordinates of all the Lagrangian nodes associated with the order per direction
`orders`, and `face_own_nodes`, being a vector of vectors indicating which nodes are owned by each of
the faces of the polytope `p`.
"""
function compute_nodes(p::Polytope,orders)
  _compute_nodes(p,orders)
end

"""
    compute_own_nodes_permutations(
      p::Polytope, own_nodes_coordinates) -> Vector{Vector{Int}}

Returns a vector of vectors with the permutations of the nodes owned by the interior of the
polytope.
"""
function compute_own_nodes_permutations(p::Polytope, interior_nodes)
  perms = _compute_node_permutations(p, interior_nodes)
  perms
end

"""
    compute_lagrangian_reffaces(::Type{T},p::Polytope,orders) where T

Returns a tuple of length `D` being the number of space dimensions.
The entry `d+1` of this tuple contains a vector of `LagrangianRefFE`
one for each face of dimension `d` on the boundary of the polytope.
"""
function compute_lagrangian_reffaces(::Type{T},p::Polytope,orders) where T
  _compute_lagrangian_reffaces(T,p,orders)
end

# Default implementations

function _compute_nodes(p,orders)
  if any( orders .== 0)
    _compute_constant_nodes(p,orders)
  elseif all(orders .== 1)
    _compute_linear_nodes(p)
  else
    _compute_high_order_nodes(p,orders)
  end
end

function _compute_constant_nodes(p,orders)
  @assert all( orders .== 0) "If an order is 0 in some direction, it should be 0 also in the others"
  x = compute_own_nodes(p,orders)
  facenodes = [Int[] for i in 1:num_faces(p)]
  push!(facenodes[end],1)
  x, facenodes
end

function _compute_linear_nodes(p)
  x = get_vertex_coordinates(p)
  facenodes = [Int[] for i in 1:num_faces(p)]
  for i in 1:num_vertices(p)
    push!(facenodes[i],i)
  end
  x, facenodes
end

function _compute_high_order_nodes(p::Polytope{D},orders) where D
  nodes = Point{D,Float64}[]
  facenodes = [Int[] for i in 1:num_faces(p)]
  _compute_high_order_nodes_dim_0!(nodes,facenodes,p)
  for d in 1:(num_dims(p)-1)
    _compute_high_order_nodes_dim_d!(nodes,facenodes,p,orders,Val{d}())
  end
  _compute_high_order_nodes_dim_D!(nodes,facenodes,p,orders)
  (nodes, facenodes)
end

function _compute_high_order_nodes_dim_0!(nodes,facenodes,p)
  x = get_vertex_coordinates(p)
  k = 1
  for vertex in 1:num_vertices(p)
    push!(nodes,x[vertex])
    push!(facenodes[vertex],k)
    k += 1
  end
end

@noinline function _compute_high_order_nodes_dim_d!(nodes,facenodes,p,orders,::Val{d}) where d
  x = get_vertex_coordinates(p)
  offset = get_offset(p,d)
  k = length(nodes)+1
  for iface in 1:num_faces(p,d)
    face = Polytope{d}(p,iface)
    face_ref_x = get_vertex_coordinates(face)
    face_prebasis = MonomialBasis(Float64,face,1)
    change = inv(evaluate(face_prebasis,face_ref_x))
    face_shapefuns = change_basis(face_prebasis,change)
    face_vertex_ids = get_faces(p,d,0)[iface]
    face_x = x[face_vertex_ids]
    face_orders = compute_face_orders(p,face,iface,orders)
    face_interior_nodes = compute_own_nodes(face,face_orders)
    face_high_x = evaluate(face_shapefuns,face_interior_nodes)*face_x
    for xi in 1:length(face_high_x)
      push!(nodes,face_high_x[xi])
      push!(facenodes[iface+offset],k)
      k += 1
    end
  end
end

function _compute_high_order_nodes_dim_D!(nodes,facenodes,p,orders)
  k = length(nodes)+1
  p_high_x = compute_own_nodes(p,orders)
  for xi in 1:length(p_high_x)
    push!(nodes,p_high_x[xi])
    push!(facenodes[end],k)
    k += 1
  end
end

_compute_node_permutations(::Polytope{0}, interior_nodes) = [[1]]

function _compute_node_permutations(p, interior_nodes)
  vertex_to_coord = get_vertex_coordinates(p)
  lbasis = MonomialBasis(Float64,p,1)
  change = inv(evaluate(lbasis,vertex_to_coord))
  lshapefuns = change_basis(lbasis,change)
  perms = get_vertex_permutations(p)
  map = evaluate(lshapefuns,interior_nodes)
  pvertex_to_coord = similar(vertex_to_coord)
  node_perms = Vector{Int}[]
  tol = 1.0e-10
  for vertex_to_pvertex in perms
    node_to_pnode = fill(INVALID_PERM,length(interior_nodes))
    pvertex_to_coord[vertex_to_pvertex] = vertex_to_coord
    pinterior_nodes = map*pvertex_to_coord
    for node in 1:length(interior_nodes)
      x = interior_nodes[node]
      pnode = findfirst(i->norm(i-x)<tol,pinterior_nodes)
      if pnode != nothing
         node_to_pnode[node] = pnode
      end
    end
    push!(node_perms,node_to_pnode)
  end
  node_perms
end

_compute_lagrangian_reffaces(::Type{T},p::Polytope{0},orders) where T = ()

function _compute_lagrangian_reffaces(::Type{T},p::Polytope{D},orders) where {T,D}
  reffaces = [ LagrangianRefFE{d}[]  for d in 0:D ]
  p0 = Polytope{0}(p,1)
  reffe0 = LagrangianRefFE(T,p0,())
  for vertex in 1:num_vertices(p)
    push!(reffaces[0+1],reffe0)
  end
  offsets = get_offsets(p)
  for d in 1:(num_dims(p)-1)
    offset = offsets[d+1]
    for iface in 1:num_faces(p,d)
      face = Polytope{d}(p,iface)
      face_orders = compute_face_orders(p,face,iface,orders)
      refface = LagrangianRefFE(T,face,face_orders)
      push!(reffaces[d+1],refface)
    end
  end
  tuple(reffaces...)
end

# Particular implementation for ExtrusionPolytope

function LagrangianRefFE(p::ExtrusionPolytope)
  order = 1
  LagrangianRefFE(Float64,p,order)
end

function compute_monomial_basis(::Type{T},p::ExtrusionPolytope{D},orders) where {D,T}
  extrusion = Tuple(p.extrusion)
  terms = _monomial_terms(extrusion,orders)
  MonomialBasis{D}(T,orders,terms)
end

function compute_own_nodes(p::ExtrusionPolytope{D},orders) where D
  extrusion = Tuple(p.extrusion)
  if all(orders .== 0)
    _interior_nodes_order_0(p)
  else
    _interior_nodes(extrusion,orders)
  end
end

function _interior_nodes_order_0(p)
  x = get_vertex_coordinates(p)
  x0 = sum(x) / length(x)
  [x0,]
end

function compute_face_orders(p::ExtrusionPolytope,face::ExtrusionPolytope{D},iface::Int,orders) where D
  d = num_dims(face)
  offset = get_offset(p,d)
  nface = p.dface.nfaces[iface+offset]
  face_orders = _eliminate_zeros(Val{D}(),nface.extrusion,orders)
  Tuple(face_orders)
end

function _eliminate_zeros(::Val{d},a,o) where d
  b = zero(mutable(Point{d,Int}))
  D = num_components(a)
  k = 1
  for i in 1:D
    m = a[i]
    if (m != 0)
      b[k] = o[i]
      k += 1
    end
  end
  Point(b)
end

function compute_nodes(p::ExtrusionPolytope{D},orders) where D
  _nodes, facenodes = _compute_nodes(p,orders)
  if any( orders .== 0)
    return (_nodes, facenodes)
  end
  terms = _coords_to_terms(_nodes,orders)
  nodes = _terms_to_coords(terms,orders)
  (nodes, facenodes)
end

# Helpers for the ExtrusionPolytope-related implementation

function _monomial_terms(extrusion::NTuple{D,Int},orders) where D
  terms = CartesianIndex{D}[]
  if D == 0
    push!(terms,CartesianIndex(()))
    return terms
  end
  _check_orders(extrusion,orders)
  M = mutable(VectorValue{D,Int})
  term = zero(M)
  _orders = M(orders)
  k = 0
  _add_terms!(terms,term,extrusion,_orders,D,k)
  terms
end

function _interior_nodes(extrusion::NTuple{D,Int},orders) where D
  _check_orders(extrusion,orders)
  terms = CartesianIndex{D}[]
  M = mutable(VectorValue{D,Int})
  term = zero(M)
  _orders = M(orders)
  k = 1
  _add_terms!(terms,term,extrusion,_orders,D,k)
  _terms_to_coords(terms,orders)
end

function _check_orders(extrusion,orders)
  D = length(extrusion)
  @assert length(orders) == D "container of orders not long enough"
  _orders = collect(orders)
  if extrusion[D] == HEX_AXIS
    _orders[D] = 0
  end
  for d in (D-1):-1:1
    if (extrusion[d] == HEX_AXIS || d == 1) && _orders[d+1] == 0
      _orders[d] = 0
    end
  end
  nz = _orders[_orders .!= 0]
  if length(nz) > 1
    @assert all(nz .== nz[1]) "The provided anisotropic order is not compatible with polytope topology"
  end
  nothing
end

function _add_terms!(terms,term,extrusion,orders,dim,k)
  _term = copy(term)
  _orders = copy(orders)
  indexbase = 1
  for i in k:(_orders[dim]-k)
    _term[dim] = i + indexbase
    if dim > 1
      if (extrusion[dim] == TET_AXIS) && i != 0
        _orders .-= 1
      end
      _add_terms!(terms,_term,extrusion,_orders,dim-1,k)
    else
      push!(terms,CartesianIndex(Tuple(_term)))
    end
  end
end

function _coords_to_terms(coords::Vector{<:Point{D}},orders) where D
  indexbase = 1
  terms = CartesianIndex{D}[]
  P = Point{D,Int}
  t = zero(mutable(P))
  for x in coords
    for d in 1:D
      t[d] = round(x[d]*orders[d]) + indexbase
    end
    term = CartesianIndex(Tuple(t))
    push!(terms,term)
  end
  terms
end

function  _terms_to_coords(terms::Vector{CartesianIndex{D}},orders) where D
  P = Point{D,Float64}
  indexbase = 1
  nodes = P[]
  x = zero(mutable(P))
  for t in terms
    for d in 1:D
      x[d] = (t[d] - indexbase) / orders[d]
    end
    node = P(x)
    push!(nodes,node)
  end
  nodes
end

function _extract_nonzeros(mask,values)
  b = Int[]
  for (m,n) in zip(mask,values)
    if (m != 0)
      push!(b, n)
    end
  end
  return Tuple(b)
end

# Precomputed instances

"""
    const VERTEX1 = LagrangianRefFE(Float64,VERTEX,1)
"""
const VERTEX1 = LagrangianRefFE(Float64,VERTEX,1)

"""
    const SEG2 = LagrangianRefFE(Float64,SEGMENT,1)
"""
const SEG2 = LagrangianRefFE(Float64,SEGMENT,1)

"""
    const QUAD4 = LagrangianRefFE(Float64,QUAD,1)
"""
const QUAD4 = LagrangianRefFE(Float64,QUAD,1)

"""
    const TRI3 = LagrangianRefFE(Float64,TRI,1)
"""
const TRI3 = LagrangianRefFE(Float64,TRI,1)

"""
    const TET4 = LagrangianRefFE(Float64,TET,1)
"""
const TET4 = LagrangianRefFE(Float64,TET,1)

"""
    const HEX8 = LagrangianRefFE(Float64,HEX,1)
"""
const HEX8 = LagrangianRefFE(Float64,HEX,1)

ext: jl


text:
using PlotlyJS, LightGraphs
import GraphPlot  # for spring_layout

# Generate a random layout
G =  LightGraphs.euclidean_graph(200, 2, cutoff=0.125)[1]
# Position nodes
pos_x, pos_y = GraphPlot.spring_layout(G)

# Create plot points
edge_x = []
edge_y = []

for edge in edges(G)
    push!(edge_x, pos_x[src(edge)])
    push!(edge_x, pos_x[dst(edge)])
    push!(edge_y, pos_y[src(edge)])
    push!(edge_y, pos_y[dst(edge)])
end

#  Color node points by the number of connections.
color_map = [size(neighbors(G, node))[1] for node in 1:200]

# Create edges
edges_trace = scatter(
    mode="lines",
    x=edge_x,
    y=edge_y,
    line=attr(
        width=0.5,
        color="#888"
    ),
)

# Create nodes
nodes_trace = scatter(
    x=pos_x,
    y=pos_y,
    mode="markers",
    text = [string("# of connections: ", connection) for connection in color_map],
    marker=attr(
        showscale=true,
        colorscale=colors.imola,
        color=color_map,
        size=10,
        colorbar=attr(
            thickness=15,
            title="Node Connections",
            xanchor="left",
            titleside="right"
      )
    )
)

# Create Plot
plot(
    [edges_trace, nodes_trace],
    Layout(
        hovermode="closest",
        title="Network Graph made with Julia",
        titlefont_size=16,
        showlegend=false,
        showarrow=false,
        xaxis=attr(showgrid=false, zeroline=false, showticklabels=false),
        yaxis=attr(showgrid=false, zeroline=false, showticklabels=false)
    )
)

max_stars_repo_name: Ellon-M/visualizations

max_issues_repo_name: Ellon-M/visualizations

max_forks_repo_name: Ellon-M/visualizations


text:
module Gaussian1d
using DataModel
using StatsBase
export Gaussian1dModel
type Gaussian1dModel <: AbstractDataModel
    sdTh::Float64
    sdX::Float64
    x::Array{Float64,1}
    postvar::Float64
    postmean::Float64
    function Gaussian1dModel(  nsdTh::Float64,       nsdX::Float64,       nx::Array{Float64,1})

      n=length(nx)
      postvar=1.0/(1.0/nsdTh^2+n/nsdX^2)
      postmean=sum(nx[1:n]/(n+ nsdX^2/nsdTh^2))
      new(nsdTh,nsdX,nx,postvar,postmean)

    end

end

function posterior(dm::Gaussian1dModel,nobs::Int64)
  x=dm.x
  sdX=dm.sdX
  sdTh=dm.sdTh
  postvar=1.0/(1.0/sdTh^2+nobs/sdX^2)
  postmean=sum(x[1:nobs]/(nobs+ sdX^2/sdTh^2))
  return (postmean,postvar)
end

function DataModel.getllik(dm::Gaussian1dModel;nobs=length(dm.x)) ## add nobs fercitily

    function logdensity(beta)
      return Float64[-0.5*(beta[1]-dm.postmean)^2/dm.postvar ]
    end
    return logdensity
end

function DataModel.getgrad(dm::Gaussian1dModel,batchsize;nobs=length(dm.x))
    x=dm.x
    sdX=dm.sdX
    sdTh=dm.sdTh

    postvar=1.0/(1.0/sdTh^2+nobs/sdX^2)
    postmean=sum(x[1:nobs]/(nobs+ sdX^2/sdTh^2))

    B2=sum(x[1:nobs])/sdX^2
    function grad_log_posterior(beta)
      return Float64[-beta[1]/postvar+B2]
    end
    function grad_log_posterior_sub(beta)

        chosen_indices=sample(1:nobs,batchsize,replace=false)

        return Float64[-beta[1]/postvar+nobs/batchsize*sum(x[chosen_indices])/sdX^2]
    end
    if nobs==batchsize
      return grad_log_posterior
    else
      return grad_log_posterior_sub
    end
end
function DataModel.getfullgrad(dm::Gaussian1dModel;nobs=length(dm.x))
  return DataModel.getgrad(dm,nobs,nobs=nobs)
end

function artificial_Gaussian1d(nObs,seed::Int64 =1)
    srand(seed)
    sdTh=1.0
    sdX=1.0
    th=randn()
    x=th+sdX*randn(nObs)
    srand()
    Gaussian1dModel(sdTh,sdX,x)
end
end


text:
using .NetworkResilience
using Statistics

ensemblesize = 200
n = 50
ns = 10
nd = 40
q = 0.1
k = 4
alphares = 50
alphamin = 0.5
alphamax = 2.5
I = 1.0
kappa = 5.0

println("Computing cascade durations for a range of gamma values...")
Ds = [0.3:0.1:2.0;]
for D in Ds
	D_alias = Int(10*D)
	fname = "data/durations/q1_10_40_0_t_gamma$D_alias.txt"

	a, s, t, fmax, crits = swingcascadeandtvalpha(ensemblesize, n, ns, nd, q, k, alphares, alphamin, alphamax, I, D, kappa)
	open(fname, "w") do f
		for i=1:length(a)
			beta = a[i]
			gamma = t[i]
			write(f, "$beta $gamma \n")
		end
	end
end

D = 1.0
q = 0.0
qs = [0.0:0.1:1.0;]
println("Computing cascade durations for a range of q values...")
for q in qs
	q_alias = Int(10*q)
	fname = "data/durations/q1_10_40_0_t_q$q_alias.txt"

	a, s, t, fmax, crits = swingcascadeandtvalpha(ensemblesize, n, ns, nd, q, k, alphares, alphamin, alphamax, I, D, kappa)
	open(fname, "w") do f
		for i=1:length(a)
			beta = a[i]
			gamma = t[i]
			write(f, "$beta $gamma \n")
		end
	end
end

q = 0.1
kappa = 0.3
kappas = [1.0:1.0:11.0;]
println("Computing cascade durations for a range of kappa values...")
for kappa in kappas
	kappa_alias = Int(kappa)
	fname = "data/durations/q1_10_40_0_t_kappa$kappa_alias.txt"

	a, s, t, fmax, crits = swingcascadeandtvalpha(ensemblesize, n, ns, nd, q, k, alphares, alphamin, alphamax, I, D, kappa)
	open(fname, "w") do f
		for i=1:length(a)
			beta = a[i]
			gamma = t[i]
			write(f, "$beta $gamma \n")
		end
	end
end

kappa = 5.0
println("Computing cascade durations for a range of network sizes n...")
ns = [45:5:100;]
for n in ns
	ns = Int(n/5)
	nd = Int(4*n/5)
	fname = "data/durations/q1_t_$n.txt"

	a, s, t, fmax, crits = swingcascadeandtvalpha(ensemblesize, n, ns, nd, q, k, alphares, alphamin, alphamax, I, D, kappa)
	open(fname, "w") do f
		for i=1:length(a)
			beta = a[i]
			gamma = t[i]
			write(f, "$beta $gamma \n")
		end
	end

end

max_stars_repo_name: o-smith/GridResilience

max_issues_repo_name: o-smith/GridResilience

max_forks_repo_name: o-smith/GridResilience


text:
module TestLinAlg
using CustomTest
using Base.Test
using ControlSystems
a = [-3 2;1 1]
b = [0  1]'
c = [1 -1]
r = [3]'

@test_approx_eq care(a, b, c'*c, r) [0.5895174372762604 1.8215747248860816; 1.8215747248860823 8.818839806923107]
@test_approx_eq dare(a, b, c'*c, r) [240.73344504496302 -131.09928700089387; -131.0992870008943 75.93413176750603]

a_2 = [-5 -3; 2 -9]
C_212 = ss(a_2, [1; 2], eye(2), [0; 0])
C_222 = ss(a_2, [1 0; 0 2], eye(2), zeros(2,2))

a_3 = [-0.81  0.47 -0.43  1.58  0.25 -0.40  0.92;
       -0.61 -1.89  0.80 -1.59  2.01  0.98 -0.90;
        0.50 -1.18 -2.12 -1.56 -1.14  0.14 -0.87;
       -1.28  2.12  0.47 -1.16  3.65 -1.23 -1.26;
       -0.24 -0.08  1.58 -3.58 -1.26  1.72 -2.59;
        1.28 -0.96 -1.28 -0.57 -2.43 -2.43 -0.36;
       -0.16  1.53 -0.99  1.47  0.61 -2.22 -3.27]
b_3 =  [ 0.49 -0.80  0.10;
        -0.18  0.70  0.72;
         0.00  0.83  2.58;
         1.42  0.00  0.00;
         0.00  0.22  0.19;
         0.00  0.00 -0.08;
         0.00 -1.15 -1.93]
c_3 = [ 0.00  0.84  0.10  0.30  0.49  1.71  0.00
       -1.79 -0.89 -0.54 -0.60  0.74 -0.19  0.00]
d_3 = [ 1.36  0.96  1.44
       -1.07  0.00 -1.96]
C_732 = ss(a_3,b_3,c_3,d_3);

s = tf("s")
f_C_211 = (s+2)*(s+3)/((s+4)*(s+5))
# biquad passband
omega0 = 52.0; Q = 10
f_C_211_bis = (s/(Q*omega0)) / ((s/omega0)^2 + s/(Q*omega0) + 1 );

C_22tf = [0 tf([3,0],[1,1,10]);tf([1,1],[1,5]) tf(2,[1,6])]

da_2 = [0.2 -0.8; -0.8 0.07]
D_221 = ss(da_2, [1 0; 0 2], [1 0], [0 0], 0.005)
D_222 = ss(da_2, [1 0; 0 2], eye(2), zeros(2,2), 0.005)

@test_approx_eq gram(C_212,:c) [0.042016806722689065 0.09663865546218485
    0.09663865546218488 0.24369747899159663]
@test_approx_eq gram(C_212,:o) [0.09523809523809523 -0.0119047619047619
    -0.011904761904761911 0.059523809523809534]
@test_approx_eq gram(C_222,:c) [0.11764705882352941 -0.029411764705882346
    -0.02941176470588234 0.21568627450980393]
@test_approx_eq gram(C_222,:o) [0.09523809523809523 -0.0119047619047619
    -0.011904761904761911 0.059523809523809534]
@test_approx_eq gram(D_221,:c) [11.895658388619264 -7.526846649259635
    -7.526846649259634 12.517564258299078]
@test_approx_eq gram(D_221,:o) [3.1039794541519283 -1.7910988952221152
    -1.7910988952221152 2.197919733616834]
@test_approx_eq gram(D_222,:c) [11.895658388619264 -7.526846649259635
    -7.526846649259634 12.517564258299078]
@test_approx_eq gram(D_222,:o) [5.301899187768763 -3.2250358337314955
    -3.2250358337314955 4.777830864787395]

@test obsv(C_222) == [1 0; 0 1; -5 -3; 2 -9]
@test ctrb(C_222) == [1 0 -5 -6; 0 2 2 -18]
@test obsv(C_212) == [1 0; 0 1; -5 -3; 2 -9]
@test ctrb(C_212) == [1 -11; 2 -16]
@test_approx_eq norm(C_222) 0.5773502691896258
@test_approx_eq norm(C_212) 0.5345224838248488
@test_approx_eq norm(D_222) 4.940973856125768
@test_approx_eq norm(D_221) 3.4490083195926426
@test norm(ss([1],[2],[3],[4])) == Inf

# Test Hinfinity norm computations
tolHinf = 1e-12
@test_approx_eq_eps norm(C_212,Inf,tol=tolHinf)[1] 0.242535625036333 tolHinf
@test_approx_eq_eps norm(C_222,Inf,tol=tolHinf)[1] 0.242535625036333 tolHinf
ninf, fpeak = norm(C_732,Inf,tol=tolHinf)
@test_approx_eq_eps ninf 4.899135403568278 (10*tolHinf)
@test_approx_eq_eps fpeak 6.112977387441163 1e-6
@test_approx_eq_eps norm(f_C_211,Inf,tol=tolHinf)[1] 1.0 (2*tolHinf)
@test_approx_eq norm(f_C_211_bis,Inf,tol=tolHinf)[2] 52.0

ninf, fpeak = norm(C_22tf,Inf,tol=tolHinf)
@test_approx_eq_eps ninf 3.014974550173459 (10*tolHinf)
@test_approx_eq_eps fpeak 3.162123338668049 1e-8

A = [1  100  10000; .01  1  100; .0001  .01  1]
T, P, B = ControlSystems.balance(A)
# The scaling is BLAS dependent. However, the ratio should be the same on all
# machines. We just need to check that T == res * constant
res_diag = [512, 8, 0.0625]
constant = T[1]/res_diag[1]
@test all(diag(T) == res_diag * constant)
@test P == eye(3)
@test_approx_eq B [1.0 1.5625 1.220703125; 0.64 1.0 0.78125; 0.8192 1.28 1.0]

end


text:
##USED TO CREATE IID NORMAL FIGURE FOR ABC
using RareEventABC
using Distributions
using DataFrames
using ProgressMeter
using Lora
using Plots; pgfplots();
using StatPlots;

##Generate observations
srand(1);
d = 25;
yobs = randn(d) .* 3.0;
a = RareEventABC.IIDNormalDetails(yobs, d);

rmse(x::Vector, y::Float64) = (x .- y) .^ 2 |> mean |> sqrt;
rmse(x::Vector, y::Float64, w::Vector) = (w .* ((x .- y) .^ 2)) |> sum |> a->a/sum(w) |> sqrt;

######
###ABC
######

"
Run a trial of ABC

Arguments:

* `N` is number of acceptances required
* `yobs` is observed data
* `threshold` is threshold to investigate
* `scale_true` is the true parameter value
"
function trialABC(N::Int, yobs::Vector, threshold::Float64; scale_true=3.0)
    prior = Uniform(0.0,10.0)
    output_sample = Float64[]
    nsims = 0
    tic()
    @showprogress for i in 1:N
        while true
            prior_draw = rand(prior)
            next_dist = RareEventABC.getDistance(a, randn(d) .* prior_draw)
            nsims += 1
            if next_dist <= threshold
                push!(output_sample, prior_draw)
                break ##i.e. increment i in for loop
            end ##if
        end ##while
    end ##for
    timeABC = toq()
    error = rmse(output_sample, scale_true)
    time_per_sample = timeABC / N
    return DataFrame(threshold = threshold, time_per_ESS = time_per_sample, RMSE = error, accrate = N/nsims, method="ABC", dim=d, mean = mean(output_sample), sd = std(output_sample))
end ##trialABC

srand(2);
temp = [trialABC(500, yobs, eps) for eps in 9.0:30.0];
abc_efficiency = vcat(temp...)

#######
##RE-ABC (non-adaptive)
#######

##Tuning
thresholds_hd = [5.0, 10.0, 15.0, 20.0, 25.0];
particles_hd  = [300, 40,  16,   10,    10];

srand(1);
schedules_hd = Array{Float64, 1}[doSMC(a, [3.0], t, nparticles=p, quiet=true).schedule for (t,p) in zip(thresholds_hd, particles_hd)];
llest = Float64[doSMC(a, [3.0], s, nparticles=p, quiet=true).llest for (s,p) in zip(schedules_hd, particles_hd), i in 1:100];
var(llest, 2)
fin(x) = x[x .> -Inf]; ##Returns finite elements
[var(fin(llest[i,:])) for i in 1:size(llest)[1]] ##Variance if zeros likelihood estimates excluded

"
Run a trial of RE-ABC within PMMH

Arguments:

* `nits` is number of PMMH iterations required
* `yobs` is observed data
* `schedule` is threshold schedule (non-adaptive algorithm)
* `epsilon` is final threshold (adaptive algorithm)
* `scale_true` is the true parameter value
* `scale_init` is starting MCMC state
* `nparticles` is number of SMC particles
* `sd` is standard deviation of MH proposal distribution
"
function trialREABC(nits::Int, yobs::Vector, schedule::Array{Float64, 1}; scale_true=3.0, scale_init=3.0, nparticles=1000, sd=1.0)
    tic()
    thetas, ll = doPMMH(a, schedule, nits, nparticles, Uniform(0.0,10.0), scale_init, Normal(0.0,sd));
    t = toq()
    return DataFrame(threshold = schedule[end], time_per_ESS = t/Lora.ess(thetas), RMSE = rmse(thetas[:,1], scale_true), accrate = length(unique(thetas)) / nits, method="RE-ABC", dim=d, mean = mean(thetas), sd = std(thetas))
end

##Pilot run
srand(2);
sd_hd = sqrt(2.562) * Float64[trialREABC(200, yobs, sched, scale_init=3.0, nparticles=p, sd=0.5)[1,:sd] for (sched,p) in zip(schedules_hd, particles_hd)]

##Main run
srand(3);
temp = [trialREABC(2000, yobs, sched, scale_init=3.0, nparticles=p, sd=s) for (sched,p,s) in zip(schedules_hd, particles_hd, sd_hd)];
abc_hd_efficiency = vcat(temp...)

##Repeat with "bad" schedules
particles_bad  = [300, 150,  60,   20,    10];
srand(1);
schedules_bad = Array{Float64, 1}[doSMC(a, [6.0], t, nparticles=p, quiet=true).schedule for (t,p) in zip(thresholds_hd, particles_bad)];
llest = @showprogress Float64[doSMC(a, [6.0], s, nparticles=p, quiet=true).llest for (s,p) in zip(schedules_bad, particles_bad), i in 1:100];
var(llest, 2)
fin(x) = x[x .> -Inf]; ##Returns finite elements
[var(fin(llest[i,:])) for i in 1:size(llest)[1]] ##Variance if zeros likelihood estimates excluded

##Pilot run
srand(2);
sd_bad = sqrt(2.562) * Float64[trialREABC(200, yobs, sched, scale_init=3.0, nparticles=p, sd=0.5)[1,:sd] for (sched,p) in zip(schedules_bad, particles_bad)]

##Main run
srand(3);
temp = [trialREABC(2000, yobs, sched, scale_init=3.0, nparticles=p, sd=s) for (sched,p,s) in zip(schedules_bad, particles_bad, sd_bad)];
efficiency_bad = vcat(temp...)
efficiency_bad[:method] = "RE-ABC poor schedule"

#######
##RE-ABC (adaptive)
#######

##Tuning
thresholds_hd = [5.0, 10.0, 15.0, 20.0, 25.0];
particles_ahd  = [260, 40,  10,   10,    10];

srand(1);
llest = Float64[doSMC(a, [3.0], t, nparticles=p, quiet=true).llest for (t,p) in zip(thresholds_hd, particles_ahd), i in 1:100];
var(llest, 2)
fin(x) = x[x .> -Inf]; ##Returns finite elements
[var(fin(llest[i,:])) for i in 1:size(llest)[1]] ##Variance if zeros likelihood estimates excluded

##Adaptive version
function trialREABC(nits::Int, yobs::Vector, epsilon::Float64; scale_true=3.0, scale_init=3.0, nparticles=1000, sd=1.0)
    tic()
    thetas, ll = doPMMH(a, epsilon, nits, nparticles, Uniform(0.0,10.0), scale_init, Normal(0.0,sd));
    t = toq()
    return DataFrame(threshold = epsilon, time_per_ESS = t/Lora.ess(thetas), RMSE = rmse(thetas[:,1], scale_true), accrate = length(unique(thetas)) / nits, method="Adaptive RE-ABC", dim=d, mean = mean(thetas), sd = std(thetas))
end

##Pilot run
srand(2);
sd_ahd = sqrt(2.562) * Float64[trialREABC(200, yobs, t, scale_init=3.0, nparticles=p, sd=0.5)[1,:sd] for (t,p) in zip(thresholds_hd, particles_ahd)]

##The results above are sufficiently similar to the non-adaptive case that we use the non-adaptive tuning for a clearer comparison
srand(3);
temp = [trialREABC(2000, yobs, t, scale_init=3.0, nparticles=p, sd=s) for (t,p,s) in zip(thresholds_hd, particles_hd, sd_hd)];
abc_ahd_efficiency = vcat(temp...)

##Closer comparison of threshold 5 case
thetas_a, ll_a = doPMMH(a, 5.0, 2000, 300, Uniform(0.0,10.0), 3.0, Normal(0.0,0.75));
thetas_n, ll_n = doPMMH(a, schedules_hd[1], 2000, 300, Uniform(0.0,10.0), 3.0, Normal(0.0,0.75));
plot(x=1:2000, hcat(thetas_a, thetas_n), line=:steppost, layout=2)
Lora.ess(thetas_a), Lora.ess(thetas_n)

#######
##LIKELIHOOD-BASED INFERENCE
#######
"
Likelihood based inference using importance sampling

Arguments:

* `nits` is number of iterations to perform
* `yobs` is observed data
* `scale_true` is the true parameter value
"
function likelihoodInf(nits::Int, yobs::Vector; scale_true=3.0)
    prior = Uniform(0.0,10.0)
    output_sample = Array(Float64, nits)
    w = Array(Float64, nits)
    tic()
    @showprogress for i in 1:nits
        sig = output_sample[i] = rand(prior)
        w[i] = pdf(Distributions.MvNormal(zeros(d), sig), yobs)
    end ##for
    timeIS = toq()
    error = rmse(output_sample, scale_true, w)
    return DataFrame(threshold = 0.0, time_per_ESS = timeIS / RareEventABC.ess(w), RMSE = error, accrate = 1.0, method="Exact", dim=d)
end ##trialABC

srand(3);
IS_efficiency = likelihoodInf(10000, yobs)

"
Likelihood based inference using MCMC

Arguments:

* `nits` is number of iterations to perform
* `yobs` is observed data
* `scale_true` is the true parameter value
"
function doMCMC(nits::Int, yobs::Vector; scale_true=3.0, scale_init=3.0, sd=1.0, thin=1)
    tic()
    thetas = Float64[]
    theta = scale_init
    L = pdf(Distributions.MvNormal(zeros(d), theta), yobs)
    1 % thin == 0 && push!(thetas, theta)
    tic()
    @showprogress for i in 2:nits
        theta_prop = theta + sd*randn()
        if 0.0 < theta_prop < 10.0
            L_prop = pdf(Distributions.MvNormal(zeros(d), theta_prop), yobs)
            if rand() < L_prop / L
                theta = theta_prop ##Acceptance
            end
        end
        i % thin == 0 && push!(thetas, theta)
    end
    t = toq()
    return DataFrame(threshold = 0.0, time_per_ESS = t/Lora.ess(thetas), RMSE = rmse(thetas[:,1], scale_true), accrate = length(unique(thetas)) / nits, method="MCMC", dim=d, mean = mean(thetas), sd = std(thetas))
end

srand(3);
MCMC_efficiency = doMCMC(10000, yobs, sd=1.8)

#######
##ABC-MCMC
#######

"
Run ABC-MCMC

Arguments:

* `nits` is number of iterations required
* `yobs` is observed data
* `threshold` is threshold to investigate
* `scale_true` is the true parameter value
* `scale_init` is starting MCMC state
* `sd` is standard deviation of MH proposal distribution
* `thin` is a thinning factor. This many iterations are performed for each one stored.
"
function doABCMCMC(nits::Int, yobs::Vector, threshold::Float64; scale_true=3.0, scale_init=3.0, sd=1.0, thin=1)
    tic()
    thetas = Float64[]
    theta = scale_init
    1 % thin == 0 && push!(thetas, theta)
    @showprogress for i in 2:nits
        theta_prop = theta + sd*randn()
        if 0.0 < theta_prop < 10.0
            next_dist = RareEventABC.getDistance(a, randn(d) .* theta_prop)
            if next_dist <= threshold
                theta = theta_prop ##Acceptance
            end
        end
        i % thin == 0 && push!(thetas, theta)
    end
    t = toq()
    return DataFrame(threshold = threshold, time_per_ESS = t/Lora.ess(thetas), RMSE = rmse(thetas[:,1], scale_true), accrate = length(unique(thetas)) / nits, method="ABC-MCMC", dim=d, mean = mean(thetas), sd = std(thetas))
end

prop_sds = sqrt(2.562) * abc_efficiency[:,:sd].data;

srand(3);
temp = [doABCMCMC(10^5, yobs, t, scale_init=3.0, sd=s) for (t,s) in zip(11.0:30.0, prop_sds[3:end])];
push!(temp, doABCMCMC(10^8, yobs, 10.0, scale_init=3.0, sd=prop_sds[2], thin=10^4));
push!(temp, doABCMCMC(2*10^8, yobs, 9.0, scale_init=3.0, sd=prop_sds[1], thin=10^4));
abc_mcmc_efficiency = vcat(temp...)

#######
##PLOTS
#######

eff = vcat(abc_efficiency, abc_hd_efficiency, abc_ahd_efficiency, MCMC_efficiency, abc_mcmc_efficiency);
##eff = vcat(eff, efficiency_bad);

sort!(eff);
default(guidefont=font(25), tickfont=font(25), legendfont=font(25), size=(500,500));
p1 = scatter(eff, :threshold, :RMSE, group=:method,
        markershape=:auto, line=:path);
savefig(p1, "iid_normal_accuracy.pdf") ##For talk
p1 = scatter!(p1, legend=false);
p2 = scatter(eff, :threshold, :time_per_ESS, group=:method,
        markershape=:auto, line=:path, yaxis=("time / ESS (s)", :log10),
        ylims=(1E-6,10.0));
savefig(p2, "iid_normal_efficiency.pdf") ##For talk
p3 = plot(p1,p2, size=(1000,500));
savefig(p3, "iid_normal.pdf") ##For paper

##For talk, just plot ABC results
##abc_efficiency = eff[eff[:method].=="ABC",:];
p4 = scatter(abc_efficiency,
             :threshold, :RMSE,
             markershape=:auto, line=:path, ylims=(0.0,2.5), legend=false);
p5 = scatter(abc_efficiency,
             :threshold, :time_per_ESS, group=:method,
             markershape=:auto, line=:path, yaxis=("time / ESS (s)", :log10),
             ylims=(1E-6,10.0), legend=false);
p6 = plot(p4,p5, size=(1000,500));
savefig(p6, "iid_normal_abc_only.pdf")


text:
"""
    A = JopTaper(spc, dims, frac[, frac_end; mode=:normal, taper=(:cosine,:cosine)])

The linear operator `A` tapers the edges of all or some subset of the dimensions of an array that
belongs to `spc::JetSpace`.  The dimensions that are tapered are given by `dims`, an `Int` tuple.
`A` will taper the beginning and/or end edges of each specified array dimension. The size of the
beginning and end tapers are determined as a fraction of the length of that array dimension, and
this fraction is set using the tuple `frac`, and (optionally) `frac_end`.

If `frac_end::NTuple{M,Float64}` is not set, then the ith entry of `frac::NTuple{M,Float64}`
determines the length of the center portion of the ith dimension that is not tapered.  If
`frac_end::NTuple{M,NTuple{2,Float64}}` is set, then the length of the begining taper of the ith array
dimension is determiend by `frac[i]`, and the length of the end taper is determined by `frac_end[i]`.

The optional named argument `mode` can be used if the taper is applied to a dimension that
corresponds to FFT ordering where the edges are assumed to be at the center and left ends of the dimension.

The optional named argument `taper` is a tuple specifying what type of taper to use at each end.  Available
tapers are :cosine and :heaviside.

# Examples:

## taper for a space containing 1D arrays
```julia
A = JopTaper(JetSpace(Float64,10), (1,), (.75,))
m = ones(domain(A))
d = A*m
```

## taper for a space containing 1D arrays, and only taper at the end
```julia
A = JopTaper(JetSpace(Float64,10), (1,), (0.0,), (0.25,))
m = ones(domain(A))
d = A*m
```

## taper for a space containing 2D arrays, where both dimensions are tapered
```julia
A = JopTaper(JetSpace(Float64,10,11), (1,2), (0.75,0.5))
m = ones(domain(A))
d = A*m
```

## taper for a space containing 3D arrays, where two of the three dimensions are tapered
```julia
A = JopTaper(JetSpace(Float64,10,11,12), (1,3), (0.75,0.5))
m = ones(domain(A))
d = A*m
```
In the next example, array dimension `1` is tapered at the end, and array dimension
`2` is tapered at the beginning and end.
```julia
A = JopTaper(JetSpace(Float64,10,11,12), (1,3), (0.0,0.25), (0.25,0.25))
m = ones(domain(A))
d = A*m
```
"""
function JopTaper(spc::JetAbstractSpace, dim::NTuple{M,Int}, prcbeg::NTuple{M,Real}, prcend::NTuple{M,Real}; mode=:normal, taper=(:cosine,:cosine)) where {M}
    # specify taper via middle part
    @assert length(prcbeg) == length(prcend) == length(dim)
    prc = [[Float64(prcbeg[i]),Float64(prcend[i])] for i=1:length(prcbeg)]
    mode = isa(mode,Symbol) ? ntuple(i->mode, length(dim)) : mode
    JopLn(dom = spc, rng = spc, df! = JopTaper_df!, s = (dim=dim, prc=prc, mode=mode, taper=taper))
end

function JopTaper(spc::JetAbstractSpace, dim::NTuple{M,Int}, prc::NTuple{M,Real}; mode=:normal, taper=(:cosine,:cosine)) where {M}
    # specify taper via middle part
    @assert length(dim) == length(prc)
    prc_lftrht = ntuple(idim->0.5*(1.0 - prc[idim]), length(dim))
    prc = [[Float64(prc_lftrht[idim]),Float64(prc_lftrht[idim])] for idim=1:length(dim)]
    mode = isa(mode,Symbol) ? ntuple(i->mode, length(dim)) : mode
    JopLn(dom = spc, rng = spc, df! = JopTaper_df!, s = (dim=dim, prc=prc, mode=mode, taper=taper))
end

export JopTaper

function JopTaper_df!(d::AbstractArray, m::AbstractArray; dim, prc, mode, taper, kwargs...)
    d .= m
    for idim = 1:length(dim)
        tpr = buildtaper(dim[idim], prc[idim], mode[idim], size(m, dim[idim]), taper)
        for idx in CartesianIndices(size(d))
            d[idx] *= tpr[idx[dim[idim]]]
        end
    end
    d
end

function buildtaper(dim, prc, mode, n, taper)
    tpr = ones(n)

    nlft = round(Int, prc[1] * n)
    nrht = round(Int, prc[2] * n)

    for i=1:nlft
        if taper[1] == :cosine
            x = (i - 1.0) / (nlft - 1.0)
            tpr[i] = 0.5 * (1.0 + cos(pi * (1.0 + x)))
        else # :heaviside
            tpr[i] = 0.0
        end
    end

    for i=1:nrht
        if taper[2] == :cosine
            x = (i - 1.0) / (nrht - 1.0)
            tpr[n+1-i] = 0.5 * (1.0 + cos(pi * (1.0 + x)))
        else # :heaviside
            tpr[n+1-i] = 0.0
        end
    end

    tpr = mode == :fft ? fftshift(tpr) : tpr
end


text:
"""
    IteratorState{T<:TreeCursor}

The state of a [`TreeIterator`](@ref) object.  These are simple wrappers of [`TreeCursor`](@ref) objects which define
a method for [`next`](@ref).  `TreeIterator`s are in turn simple wrappers of `IteratorState`s.

Each `IteratorState` fully determines the current iteration state and therefore the next state can be obtained
with `next` (`nothing` is returned after the final state is reached).
"""
abstract type IteratorState{T<:TreeCursor} end

# we define this explicitly to avoid run-time dispatch
"""
    instance(::Type{<:IteratorState}, node; kw...)

Create an instance of the given [`IteratorState`](@ref) around node `node`.  This is mostly just a constructor
for `IteratorState` except that if `node` is `nothing` it will return `nothing`.
"""
function instance(::Type{S}, node; kw...) where {S<:IteratorState}
    isnothing(node) ? nothing : S(node; kw...)
end

"""
    initial(::Type{<:IteratorState}, node)

Obtain the initial [`IteratorState`](@ref) of the provided type for the node `node`.
"""
function initial end

"""
    next(s::IteratorState)
    next(f, s::IteratorState)

Obtain the next [`IteratorState`](@ref) after the current one.  If `s` is the final state, this will return
`nothing`.

This provides an alternative iteration protocol which only uses the states directly as opposed to
[`Base.iterate`](@ref) which takes an iterator object and the current state as separate arguments.
"""
function next end

"""
    statetype(::Type{<:TreeIterator})

Gives the type of [`IteratorState`](@ref) which is the state of the provided [`TreeIterator`](@ref).
"""
function statetype end

"""
    TreeIterator{T}

An iterator of a tree that implements the AbstractTrees interface.  Every `TreeIterator` is simply a wrapper
of an [`IteratorState`](@ref) which fully determine the iteration state and implement their own alternative
protocol using [`next`](@ref).

## Constructors
All `TreeIterator`s have a one argument constructor `T(node)` which starts iteration from `node`.
"""
abstract type TreeIterator{T} end

_iterator_eltype(::NodeTypeUnknown) = EltypeUnknown()
_iterator_eltype(::HasNodeType) = HasEltype()

Base.IteratorEltype(::Type{<:TreeIterator{T}}) where {T}  = _iterator_eltype(NodeType(T))

Base.eltype(::Type{<:TreeIterator{T}}) where {T} = nodetype(T)
Base.eltype(ti::TreeIterator) = eltype(typeof(ti))

Base.IteratorSize(::Type{<:TreeIterator}) = SizeUnknown()

function Base.iterate(ti::TreeIterator, s::Union{Nothing,IteratorState}=initial(statetype(ti), ti.root))
    isnothing(s) && return nothing
    (nodevalue(s.cursor), next(s))
end

"""
    nodevalues(itr::TreeIterator)

An iterator which returns the `nodevalue` of each node in the tree, equivalent to
`Iterators.map(nodevalue, itr)`.
"""
nodevalues(itr::TreeIterator) = Iterators.map(nodevalue, itr)

"""
    PreOrderState{T<:TreeCursor} <: IteratorState{T}

The iteration state of a tree iterator which guarantees that parent nodes will be visited *before* their children,
i.e. which descends a tree from root to leaves.

This state implements a [`next`](@ref) method which accepts a filter function as its first argument, allowing
tree branches to be skipped.

See [`PreOrderDFS`](@ref).
"""
struct PreOrderState{T<:TreeCursor} <: IteratorState{T}
    cursor::T

    PreOrderState(csr::TreeCursor) = new{typeof(csr)}(csr)
end

PreOrderState(node) = (PreOrderState ∘ TreeCursor)(node)

initial(::Type{PreOrderState}, node) = PreOrderState(node)

function next(f, s::PreOrderState)
    if f(nodevalue(s.cursor))
        ch = children(s.cursor)
        isempty(ch) || return instance(PreOrderState, first(ch))
    end

    csr = s.cursor
    while !isroot(csr)
        n = nextsibling(csr)
        if isnothing(n)
            csr = parent(csr)
        else
            return PreOrderState(n)
        end
    end
    nothing
end
next(s::PreOrderState) = next(_ -> true, s)

"""
    PreOrderDFS{T,F} <: TreeIterator{T}

Iterator to visit the nodes of a tree, guaranteeing that parents
will be visited before their children.

Optionally takes a filter function that determines whether the iterator
should continue iterating over a node's children (if it has any) or should
consider that node a leaf.

e.g. for the tree
```
Any[Any[1,2],Any[3,4]]
├─ Any[1,2]
|  ├─ 1
|  └─ 2
└─ Any[3,4]
   ├─ 3
   └─ 4
```
we will get `[[[1, 2], [3, 4]], [1, 2], 1, 2, [3, 4], 3, 4]`.
"""
struct PreOrderDFS{T,F} <: TreeIterator{T}
    filter::F
    root::T
end

PreOrderDFS(root) = PreOrderDFS(_ -> true, root)

statetype(itr::PreOrderDFS) = PreOrderState

function Base.iterate(ti::PreOrderDFS, s::Union{Nothing,IteratorState}=initial(statetype(ti), ti.root))
    isnothing(s) && return nothing
    (nodevalue(s.cursor), next(ti.filter, s))
end

"""
    PostOrderState{T<:TreeCursor} <: IteratorState{T}

The state of a tree iterator which guarantees that parents are visited *after* their children, i.e. ascends
a tree from leaves to root.

See [`PostOrderDFS`](@ref).
"""
struct PostOrderState{T<:TreeCursor} <: IteratorState{T}
    cursor::T

    PostOrderState(csr::TreeCursor) = new{typeof(csr)}(csr)
end

PostOrderState(node) = (PostOrderState ∘ TreeCursor)(node)

initial(::Type{PostOrderState}, node) = (PostOrderState ∘ descendleft ∘ TreeCursor)(node)

function next(s::PostOrderState)
    n = nextsibling(s.cursor)
    if isnothing(n)
        instance(PostOrderState, parent(s.cursor))
    else
        # descendleft is not allowed to return nothing
        PostOrderState(descendleft(n))
    end
end

"""
    PostOrderDFS{T} <: TreeIterator{T}

Iterator to visit the nodes of a tree, guaranteeing that children
will be visited before their parents.

e.g. for the tree
```
Any[1,Any[2,3]]
├─ 1
└─ Any[2,3]
   ├─ 2
   └─ 3
```
we will get `[1, 2, 3, [2, 3], [1, [2, 3]]]`.
"""
struct PostOrderDFS{T} <: TreeIterator{T}
    root::T
end

statetype(itr::PostOrderDFS) = PostOrderState

"""
    LeavesState{T<:TreeCursor} <: IteratorState{T}

A [`IteratorState`](@ref) of an iterator which visits the leaves of a tree.

See [`Leaves`](@ref).
"""
struct LeavesState{T<:TreeCursor} <: IteratorState{T}
    cursor::T

    LeavesState(csr::TreeCursor) = new{typeof(csr)}(csr)
end

LeavesState(node) = (LeavesState ∘ TreeCursor)(node)

initial(::Type{LeavesState}, node) = (LeavesState ∘ descendleft ∘ TreeCursor)(node)

function next(s::LeavesState)
    csr = s.cursor
    while !isnothing(csr) && !isroot(csr)
        n = nextsibling(csr)
        if isnothing(n)
            csr = parent(csr)
        else
            # descendleft is not allowed to return nothing
            return LeavesState(descendleft(n))
        end
    end
    nothing
end

"""
    Leaves{T} <: TreeIterator{T}

Iterator to visit the leaves of a tree, e.g. for the tree
```
Any[1,Any[2,3]]
├─ 1
└─ Any[2,3]
   ├─ 2
   └─ 3
```
we will get `[1,2,3]`.
"""
struct Leaves{T} <: TreeIterator{T}
    root::T
end

statetype(itr::Leaves) = LeavesState

"""
    SiblingState{T<:TreeCursor} <: IteratorState{T}

A [`IteratorState`](@ref) of an iterator which visits all of the tree siblings after the current sibling.

See [`Siblings`](@ref).
"""
struct SiblingState{T<:TreeCursor} <: IteratorState{T}
    cursor::T

    SiblingState(csr::TreeCursor) = new{typeof(csr)}(csr)
end

SiblingState(node) = (SiblingState ∘ TreeCursor)(node)

initial(::Type{SiblingState}, node) = SiblingState(node)

next(s::SiblingState) = nextsibling(s.cursor)

"""
    Siblings{T} <: TreeIterator{T}

A [`TreeIterator`](@ref) which visits the siblings of a node after the provided node.
"""
struct Siblings{T} <: TreeIterator{T}
    root::T
end

statetype(itr::Siblings) = SiblingState

function _ascend(getparent, select, node)
    isroot(node) && (select(node); return node)
    p = getparent(node)
    while select(node) && !isroot(node)
        node = p
        p = getparent(node)
    end
    node
end

"""
    ascend(select, node)
    ascend(select, tree, node)

Ascend a tree starting from node `node` and at each node choosing whether or not to terminate by calling
`select(n)` where `n` is the current node (`false` terminates).

For the two argument method, parents are obtained by calling `parent(tree, node)`, otherwise by calling
`parent(node)`.

Note that the parent is computed before `select` is executed, allowing modification to its argument
(as long as the tree structure is not altered).
"""
ascend(select, tree, node) = _ascend(n -> parent(tree, n), select, node)
ascend(select, node) = _ascend(parent, select, node)

"""
    descend(select, tree)

Descends the tree, at each node choosing the child given by select callback
or the current node if 0 is returned.

Requires tree nodes to have [`IndexedChildren`](@ref).
"""
function descend(select, ::IndexedChildren, tree)
    idx = select(tree)
    idx == 0 && return tree
    node = children(tree)[idx]
    while true
        idx = select(node)
        idx == 0 && return node
        node = children(node)[idx]
    end
end
descend(select, node) = descend(select, ChildIndexing(node), node)

"""
    StatelessBFS{T} <: TreeIterator{T}

Iterator to visit the nodes of a tree, all nodes of a level will be visited
before their children

This iterator requires [`getdescendant`](@ref) to be valid for all nodes in the
tree, but the nodes do not necessarily need the [`IndexedChildren`](@ref) trait.

e.g. for the tree
```
Any[1,Any[2,3]]
├─ 1
└─ Any[2,3]
   ├─ 2
   └─ 3
```

we will get `[[1, [2,3]], 1, [2, 3], 2, 3]`.

WARNING: This is \$O(n^2)\$, only use this if you know you need it, as opposed to
a more standard statefull approach.
"""
struct StatelessBFS{T} <: TreeIterator{T}
    root::T
end

Base.iterate(ti::StatelessBFS) = (ti.root, [])

function _descend_left(inds, next_node, level)
    while length(inds) ≠ level
        ch = children(next_node)
        isempty(ch) && break
        push!(inds, 1)
        next_node = first(ch)
    end
    inds
end

function _nextind_or_deadend(node, ind, level)
    current_lvl = active_lvl = length(ind)
    active_inds = ind
    # go up until there is a sibling to the right
    while current_lvl > 0
        active_inds = ind[1:(current_lvl-1)]
        parent = getdescendant(node, active_inds)
        cur_child = ind[current_lvl]
        ch = children(parent)
        ni = nextind(ch, cur_child)
        current_lvl -= 1
        if !isnothing(iterate(ch, ni))
            newinds = [active_inds; ni]
            next_node = ch[ni]
            return _descend_left(newinds, next_node, level)
        end
    end
    nothing
end

function Base.iterate(ti::StatelessBFS, ind)
    org_lvl = active_lvl = length(ind)
    newinds = ind
    while true
        newinds = _nextind_or_deadend(ti.root, newinds, active_lvl)
        if isnothing(newinds)
            active_lvl += 1
            active_lvl > org_lvl + 1 && return nothing
            newinds = _descend_left([], ti.root, active_lvl)
        end
        length(newinds) == active_lvl && break
    end
    (getdescendant(ti.root, newinds), newinds)
end

"""
    MapNode{T,C}

A node in a tree which is returned by [`treemap`](@ref).  It consists of a value which is the result of the function
call and an array of the children, which are also of type `MapNode`.

Every `MapNode` is itself a tree with the [`IndexedChildren`](@ref) trait and therefore supports indexing via
[`getdescendant`](@ref).

Use [`AbstractTrees.nodevalue`](@ref) or `mapnode.value` to obtain the wrapped value.
"""
struct MapNode{T,C}
    value::T
    children::C

    function MapNode(f, node)
        (v, ch) = f(node)
        ch′ = map(c -> MapNode(f, c), ch)
        new{typeof(v),typeof(ch′)}(v, ch′)
    end
end
MapNode(node) = MapNode(n -> (nodevalue(n), children(n)), nodevalue(node))

children(μ::MapNode) = μ.children

childrentype(::Type{MapNode{T,C}}) where {T,C} = C

nodevalue(μ::MapNode) = μ.value

ChildIndexing(::MapNode) = IndexedChildren()

function Base.show(io::IO, μ::MapNode)
    print(io, typeof(μ))
    print(io, "(", μ.value, ")")
end

Base.show(io::IO, ::MIME"text/plain", μ::MapNode) = print_tree(io, μ)

"""
    treemap(f, node)

Apply the function `f` to every node in the tree with root `node`, where `f` is a function that returns `(v, ch)` where
`v` is a new value for the node (i.e. as returned by [`nodevalue`](@ref) and `ch` is the new children of the node.
`f` will be called recursively so that all the children returned by `f` for a parent node will again be called for
each child.  This means that to maintain the structure of a tree but merely map to new values, one should define
`f = node -> (g(node), children(node))` for some function `g` which returns a value for the node.

The nodes of the output tree will all be represented by [`MapNode`](@ref) objects wrapping the returned values.
This is necessary in order to guarantee that the output types can describe any tree topology.

Note that in most common cases tree nodes are of a type which depends on their connectedness and the function
`f` should take this into account.  For example the tree `[1, [2, 3]]` contains integer leaves but two
`Vector` nodes.  Therefore, the `f` in `treemap(f, [1, [2,3]])` must be a function that is valid for either
`Int` or `Vector`.  Alternatively, to only operate on leaves do `map(𝒻, Leaves(itr))`.

It's very easy to write an `f` that makes `treemap` stack-overflow.  To avoid this, ensure that `f` eventually
termiantes, i.e. that sometimes it returns empty `children`.  For example, if `f(n) = (nothing, [0; children(n)])` will
stack-overflow because every node will have at least 1 child.

## Examples
```julia
julia> t = [1, [2, 3]];

julia> f(n) = n isa AbstractArray ? (nothing, children(n)) : (n+1, children(n))

julia> treemap(f, t)
nothing
├─ 2
└─ nothing
   ├─ 3
   └─ 4

julia> g(n) = isempty(children(n)) ? (nodevalue(n), []) : (nodevalue(n), [0; children(n)])
g (generic function with 1 method)

julia> treemap(g, t)
Any[1, [2, 3]]
├─ 0
├─ 1
└─ [2, 3]
   ├─ 0
   ├─ 2
   └─ 3
```
"""
treemap(f, node) = MapNode(f, node)


text:
#!/usr/bin/env bash
    #=
    exec julia --project="$(realpath $(dirname $0))/" "${BASH_SOURCE[0]}" "$@" -e "include(popfirst!(ARGS))" \
    "${BASH_SOURCE[0]}" "$@"
    =#

#=
Adapted from https://github.com/Simon-Hohberg/Viola-Jones/
=#

println("\033[1;34m===>\033[0;38m\033[1;38m\tLoading required libraries (it will take a moment to precompile if it is your first time doing this)...\033[0;38m")

include(joinpath(dirname(dirname(@__FILE__)), "src", "FaceDetection.jl"))

using .FaceDetection
const FD = FaceDetection
using Printf: @printf
using Images: imresize

println("...done")

function main(;
    smart_choose_feats::Bool=false,
	scale::Bool=false,
	scale_to::Tuple=(200, 200)
)
	include("constants.jl")
	include("main_data.jl")

    min_size_img = (19, 19) # default for our test dataset
    if smart_choose_feats
        # For performance reasons restricting feature size
        notify_user("Selecting best feature width and height...")

        max_feature_width, max_feature_height, min_feature_height, min_feature_width, min_size_img = determine_feature_size(pos_training_path, neg_training_path; scale = scale, scale_to = scale_to)

        println("...done.  Maximum feature width selected is $max_feature_width pixels; minimum feature width is $min_feature_width; maximum feature height is $max_feature_height pixels; minimum feature height is $min_feature_height.\n")
    else
        min_feature_height = 8
        max_feature_height = 10
        min_feature_width = 8
        max_feature_width = 10
    end

    # classifiers are haar like features
    classifiers = FD.learn(pos_training_path, neg_training_path, num_classifiers, min_feature_height, max_feature_height, min_feature_width, max_feature_width; scale = scale, scale_to = scale_to)

    FD.notify_user("Testing selected classifiers...")
	num_faces = length(filtered_ls(pos_testing_path))
	num_non_faces = length(filtered_ls(neg_testing_path))

	correct_faces = sum(FD.ensemble_vote_all(pos_testing_path, classifiers, scale=scale, scale_to=scale_to))
	correct_non_faces = num_non_faces - sum(FD.ensemble_vote_all(neg_testing_path, classifiers, scale=scale, scale_to=scale_to))
	correct_faces_percent = (correct_faces / num_faces) * 100
	correct_non_faces_percent = (correct_non_faces / num_non_faces) * 100

    faces_frac = string(correct_faces, "/", num_faces)
    faces_percent = string("(", correct_faces_percent, "% of faces were recognised as faces)")
    non_faces_frac = string(correct_non_faces, "/", num_non_faces)
    non_faces_percent = string("(", correct_non_faces_percent, "% of non-faces were identified as non-faces)")

    println("...done.\n")
    FD.notify_user("Result:\n")

    @printf("%10.9s %10.15s %15s\n", "Faces:", faces_frac, faces_percent)
    @printf("%10.9s %10.15s %15s\n\n", "Non-faces:", non_faces_frac, non_faces_percent)
end

@time main(smart_choose_feats=true, scale=true, scale_to=(20, 20))


text:
function gale_gross(x, y, spectrum_x, spectrum_y, weights_x, weights_y)
	function _f(i, x, sx)
		if i == 0
			prod([(x - s)^2 for s in sx])
		else
			prod([(x - s)^2 / (sx[i] - s)^2 for s in sx if s != sx[i]])
		end
	end

	function _ϕ(x, as)
		prod([(x - s)^2 for s in as])
	end

	function _ϕ(i, x, as)
		prod([(x - s)^2 for s in as if s != as[i + 1]])
	end

	μs = weights_x
	νs = weights_y
	m = length(μs)
	n = length(νs)

	αs = rand(Float64, n + 1) * 2 .- 1
	βs = rand(Float64, m + 1) * 2 .- 1

	f(i, x) = _f(i, x, spectrum_x)
	ϕ(i, x) = _ϕ(i, x, αs)
	ϕ(x) = _ϕ(x, αs)
	g(i, y) = _f(i, y, spectrum_y)
	ψ(i, y) = _ϕ(i, y, βs)
	ψ(y) = _ϕ(y, βs)

	m1 = +f(0, x)ϕ(x) * (g(0, y)ϕ(0, x) + sum([(g(j, y) - νs[j]) * ϕ(j, x) for j = 1:n]))
	m2 = -g(0, y)ψ(y) * (f(0, x)ψ(0, y) + sum([(f(i, x) - μs[i]) * ψ(i, y) for i = 1:m]))
	m3 = -(f(0, x)ϕ(x))^2 + (g(0, y)ψ(y))^2

	m1 + m2 + m3
end


text:
module BinaryTreeRectanglePacking

using AbstractTrees
using Plots

include("Structs.jl")
include("Packing.jl")
include("Plotting.jl")

using .Structs, .Packing, .Plotting

export pack

end # module


text:
using Simplex
using CUDA
using SparseArrays
using LinearAlgebra
using MatrixOptInterface
using Test

const MatOI = MatrixOptInterface
const MOI = MatOI.MOI
# CUDA.allowscalar(false)

include("PhaseOne/Artificial.jl")
include("PhaseOne/Cplex.jl")
include("netlib.jl")


text:
#########################
# Entity-based learning #
#########################

# The state of the network learner: it is defined by estimates for all nodes
# and a mask that specifies which of the estimates are to be update (in case the training
# data is used as well

"""
Entity-based network learning model state. It consists of an `Array` with estimates and a an update
mask in the form of a `BitVector` indicating which observation estimates are to be updated (the
ones that are not updated are considered training/stable observations).
"""
mutable struct NetworkLearnerState{T<:AbstractArray, OD<:LearnBase.ObsDimension}
	ê::T			# estimates
	update::BitVector	# which estimates to update
	obsdim::OD

	function NetworkLearnerState(ê::T, update::BitVector, obsdim::OD) where {T<:AbstractArray, OD<:LearnBase.ObsDimension}
		@assert  nobs(ê,obsdim=obsdim) == length(update) "Number of NetworkLearner estimates should be equal to the length of the update mask."
		new{T,OD}(ê, update, obsdim)
	end
end

Base.show(io::IO, m::NetworkLearnerState) = print(io, "NetworkLearner state: $(sum(m.update))/$(length(m.update)) entities can be updated")

"""
Entity-based network learning model type.
"""
mutable struct NetworkLearnerEnt{S,V,
				    NS<:NetworkLearnerState,
				    R<:Vector{<:AbstractRelationalLearner},
				    C<:AbstractCollectiveInferer,
				    A<:Vector{<:AbstractAdjacency},
				    OD<:LearnBase.ObsDimension} <: AbstractNetworkLearner
	state::NS										# state
	Mr::S											# relational model
	fr_exec::V										# relational model execution function
	RL::R											# relational learner
	Ci::C											# collective inferer
	Adj::A											# adjacency information
	size_out::Int										# expected output dimensionality
	obsdim::OD										# observation dimension
end

# Printers
Base.show(io::IO, m::NetworkLearnerEnt) = begin
	println(io,"NetworkLearner, $(m.size_out) estimates, entity-based")
	println(io,"`- state: $(sum(m.state.update))/$(nobs(m.state.update)) entities can be updated");
	print(io,"`- relational model: "); println(io, m.Mr)
	print(io,"`- relational learners: "); println(io, m.RL)
	print(io,"`- collective inferer: "); println(io, m.Ci)
	print(io,"`- adjacency: "); println(io, m.Adj)
	println(io,"`- observations are: $(m.obsdim == ObsDim.Constant{1} ? "rows" : "columns")");
end

####################
# Training methods #
####################
"""
	fit(::Type{NetworkLearnerEnt}, X, update, Adj, fl_train, fl_exec, fr_train, fr_exec [;kwargs])

Training method for the entity-based network learning framework.

# Arguments
  * `Xo::AbstractMatrix` initial estimates for the entities
  * `update::BitVector` mask that indicates wether estimates can be updated (`true` value) or not (`false` value); false values
  generally can be associated with estimates of training samples
  * `Adj::Vector{AbstractAdjacency}` a vector containing the entity relational structures (adjacency objects)
  * `fr_train` relational model training `function`; can be anything that suports the call `fr_train((Xr,y))` where `y = f_targets(Xo)`
  * `fr_exec` relational model prediction `function`; can be anything that suports the call `fr_exec(Mr,Xr)` where `Mr = fr_train((Xr,y))`
  and `Xr` is a dataset of relational variables generated by the relational learner using the estimates `Xo` and the  adjacency structures.

# Keyword arguments
  * `priors::Vector{Float64}` class priors (if applicable)
  * `learner::Symbol` relational learner (i.e. variable generator); available options `:rn`, `:wrn`, `:bayesrn` and `:cdrn` (default `:wrn`)
  * `inference::Symbol` collective inference method; available options `:rl`, `:ic` and `:gs` (default `:rl`)
  * `normalize::Bool` whether to normalize the relational variables per-entity to the L1 norm (default `true`)
  * `f_targets::Function` function that extracts targets from estimates generated by the local/relational models
  (default `f_targets = x->MLDataPattern.targets(indmax,x)`)
  * `obsdim::Int` observation dimension (default `2`)
  * `tol::Float64` maximum admissible mean estimate error for collective inference convergence (default `1e-6`)
  * `κ::Float64` relaxation labeling starting constant, used if `learner == :rl` (default `1.0`)
  * `α::Float64` relaxation labeling decay constant, used if `learner == :rl` (default `0.99`)
  * `maxiter::Int` maximum number of iterations for collective inference (default `100`)
  * `bratio::Float64` percentage of iterations i.e. `maxiter` used for Gibbs sampling burn-in (default `0.1`)
"""
function fit(::Type{NetworkLearnerEnt}, Xo::AbstractMatrix, update::BitVector, Adj::A where A<:Vector{<:AbstractAdjacency},
	     	fr_train, fr_exec;
		learner::Symbol=:wvrn, inference::Symbol=:rl, normalize::Bool=true, f_targets::Function=x->targets(indmax,x),
		obsdim::Int=2, priors::Vector{Float64}=1/nvars(Xo,ObsDim.Constant(obsdim)).*ones(nvars(Xo,ObsDim.Constant(obsdim))),
		tol::Float64=1e-6, κ::Float64=1.0, α::Float64=0.99, maxiter::Int=100, bratio::Float64=0.1)

	# Parse, transform input arguments
	κ = clamp(κ, 1e-6, 1.0)
	α = clamp(α, 1e-6, 1.0-1e-6)
	tol = clamp(tol, 0.0, Inf)
	maxiter = ifelse(maxiter<=0, 1, maxiter)
	bratio = clamp(bratio, 1e-6, 1.0-1e-6)

	@assert obsdim in [1,2] "Observation dimension can have only two values 1 (row-major) or 2 (column-major)."
	@assert all((priors.>=0.0) .& (priors .<=1.0)) "All priors have to be between 0.0 and 1.0."

	# Parse relational learner argument and generate relational learner type
	if learner == :rn
		Rl = SimpleRN
	elseif learner == :wrn
		Rl = WeightedRN
	elseif learner == :cdrn
		Rl = ClassDistributionRN
	elseif learner == :bayesrn
		Rl = BayesRN
	else
		@print_verbose 1 "Unknown relational learner. Defaulting to :wrn."
		Rl = WeightedRN
	end

	# Parse collective inference argument and generate collective inference objects
	if inference == :rl
		Ci = RelaxationLabelingInferer(maxiter, tol, f_targets, κ, α)
	elseif inference == :ic
		Ci = IterativeClassificationInferer(maxiter, tol, f_targets)
	elseif inference == :gs
		Ci = GibbsSamplingInferer(maxiter, tol, f_targets, ceil(Int, maxiter*bratio))
	else
		@print_verbose 1 "Unknown collective inferer. Defaulting to :rl."
		Ci = RelaxationLabelingInferer(maxiter, tol, f_targets, κ, α)
	end

	fit(NetworkLearnerEnt, Xo, update, Adj, Rl, Ci, fr_train, fr_exec;
     		priors=priors, normalize=normalize, obsdim=ObsDim.Constant{obsdim}())
end

"""
Training method for the network learning framework. This method should not be called directly.
"""
function fit(::Type{NetworkLearnerEnt}, Xo::T, update::BitVector, Adj::A, Rl::R, Ci::C, fr_train::U, fr_exec::U2;
	    	normalize::Bool=true, obsdim::LearnBase.ObsDimension=ObsDim.Constant{2}(),
	    	priors::Vector{Float64}=1/nvars(Xo,obsdim).*ones(nvars(Xo,obsdim)) ) where {
			T<:AbstractMatrix,
			A<:Vector{<:AbstractAdjacency},
			R<:Type{<:AbstractRelationalLearner},
			C<:AbstractCollectiveInferer,
			U, U2
		}

	# Step 0: pre-process input arguments and retrieve sizes
	n = nobs(Xo,obsdim)									# number of entities
	p = nvars(Xo,obsdim)									# number of estimates/entity
	m = length(Adj) * p									# number of relational variables

	@assert p == length(priors) "Found $p entities/estimate, the priors indicate $(length(priors))."

	# Step 1: Get relational variables by training and executing the relational learner
	@print_verbose 2 "Calculating relational variables ..."
	mₜ = .!update										# training mask (entities that are not updated
	Xoₜ = datasubset(Xo, mₜ, obsdim)							#    are considered training or high/certainty samples)
	nₜ = sum(mₜ)										# number of training observations
	yₜ = Ci.tf(Xoₜ)										# get targets
	Adjₜ = [adjacency(adjacency_matrix(Aᵢ)[mₜ,mₜ]) for Aᵢ in Adj]
	RL = [fit(Rl, Aᵢₜ, Xoₜ, yₜ; obsdim=obsdim, priors=priors, normalize=normalize)
       		for Aᵢₜ in Adjₜ]								# Train relational learners

	# Pre-allocate relational variables array
	Xr = matrix_prealloc(n, m, obsdim, 0.0)
	Xrᵢ = matrix_prealloc(nₜ, p, obsdim, 0.0)						# Initialize temporary storage
	Xrₜ = datasubset(Xr, mₜ, obsdim)
	for (i,(RLᵢ,Aᵢₜ)) in enumerate(zip(RL,Adjₜ))

		# Apply relational learner
		transform!(Xrᵢ, RLᵢ, Aᵢₜ, Xoₜ, yₜ)

		# Update relational data output
		_Xrₜ = datasubset(Xrₜ, (i-1)*p+1 : i*p, oppdim(obsdim))
		_Xrₜ[:] = Xrᵢ
	end

	# Step 2 : train relational model
	@print_verbose 2 "Training relational model ..."
	Mr = fr_train((Xrₜ,yₜ))

	# Step 3: Apply collective inference
	@print_verbose 2 "Collective inference ..."
	transform!(Xo, Ci, obsdim, Mr, fr_exec, RL, Adj, 0, Xr, update)

	# Step 3: return network learner
	return NetworkLearnerEnt(NetworkLearnerState(Xo,update,obsdim), Mr, fr_exec, RL, Ci, Adj, p, obsdim)
end

"""
Function that calls collective inference using the information in contained in the entity-based network learner
"""
function infer!(model::T) where T<:NetworkLearnerEnt
	p = nvars(model.state.ê, model.obsdim)							# number of estimates/entity
	m = length(model.Adj) * p								# number of relational variables
	Xr = matrix_prealloc(nobs(model.state.ê, model.obsdim), m, model.obsdim, 0.0)

	transform!(model.state.ê, model.Ci, model.obsdim, model.Mr, model.fr_exec, model.RL, model.Adj, 0, Xr, model.state.update)
end


text:
struct QBool
  bit::Bit
end

QBool(v::Bool = false) = QBool(Bit(v))

bits(x::QBool) = x.bit
repr(::Type{QBool}, bits) = bits[1]

Base.show(io::IO, ::QBool) = print(io, "QBool()")

function hadamard(x::QBool)
  apply!(H, x)
  return x
end

function not(x::QBool)
  apply!(X, x)
  return x
end

function cnot(c::QBool, x::QBool)
  apply!(CX, c, x)
  return x
end

function swap(a::QBool, b::QBool)
  apply!(S, a, b)
  return
end

function epr()
  a, b = QBool(), QBool()
  hadamard(a)
  cnot(a, b)
  return a, b
end


text:
module BloomFadeTrigger

using ..Ahorn, Maple

const placements = Ahorn.PlacementDict(
    "Bloom Fade" => Ahorn.EntityPlacement(
        Maple.BloomFadeTrigger,
        "rectangle"
    )
)

function Ahorn.editingOptions(trigger::Maple.BloomFadeTrigger)
    return Dict{String, Any}(
        "positionMode" => Maple.trigger_position_modes
    )
end

end


text:
@testset "Test function `pluralize_eng`" begin
    PostgresORMUtil.pluralize("analysis","eng")
    PostgresORMUtil.pluralize("analyse","fra")
    PostgresORMUtil.pluralize("bateau","fra")
end


text:
using OpenTelemetrySDK
using OpenTelemetryAPI
using Test
using Random

@testset "OpenTelemetrySDK.jl" begin
    include("trace.jl")
    include("metric.jl")
end


text:
function sqrt2(a,n)
if a
return n>0?2.0:1.0
else
return 1.0
end
end

function napier(a,n)
if a
return n>0?Float64(n):2.0
else
return n>1?n-1.0:1.0
end
end

function _pi(a,n)
if a
return n>0?6.0:3.0
else
return (2.0*n-1.0)^2.0 # exponentiation operator
end
end

function calc(f,expansions)
a=true;b=false
r=0.0
for i=expansions:-1:1
r=f(b,i)/(f(a,i)+r)
end
return f(a,0)+r
end

print("""
sqrt 2 = $(calc(sqrt2, 1000))
e = $(calc(napier, 1000))
 pi = $(calc(_pi, 1000))
""")

max_stars_repo_name: ethansaxenian/RosettaDecode

max_issues_repo_name: ethansaxenian/RosettaDecode

max_forks_repo_name: ethansaxenian/RosettaDecode


text:
"""
        INFO HERE
"""

"""
                Debug data from shock
"""
mutable struct CRShockData
   """   Data structure to analyze a single shocked particle.
         To be used with compile option: LMB_CR_DEBUG_TRACK_PARTICLE
   """
   dt::Vector{Float64}
   Mach::Vector{Float64}
   Shock_Speed::Vector{Float64}
   Shock_Compress::Vector{Float64}
   Shock_Energy_In::Vector{Float64}
   Shock_Energy_Real::Vector{Float64}
   Energy_P::Vector{Float64}
   Energy_e::Vector{Float64}

   function CRShockData(N::Int64=0)

       new(zeros(N), zeros(N), zeros(N), zeros(N),
           zeros(N), zeros(N), zeros(N), zeros(N))

   end
end

"""
                Distribution spectrum in momentum space
"""


text:
# using FlightMechanics
# using FlightMechanics.Models
# using FlightMechanics.Aircrafts

export Linearize, get_A, get_B, get_named_A, get_named_B,
        get_lon_A, get_lon_B, get_lat_A, get_lat_B

struct Linearize
    A :: Array{T, 2} where T<:Number
    B :: Array{T, 2} where T<:Number
    # C :: Array{T, 2} where T<:Number
    # D :: 0  # Array{T, 2} where T<:Number   # no feed forward
end

function Linearize(trimmed_ac::Aircraft, trimmed_state::State, trimmed_aerostate::AeroState,
					trimmed_fcs::FCS, env::Environment, dx=0.001::Float64, du=0.001::Float64, tol=1e-6::Float64)

    grav = get_gravity(env)

    ## Cacluating A_matrix
    A = A_matrix_calc(trimmed_ac, trimmed_state, trimmed_aerostate, trimmed_fcs, env, grav, dx, tol)

    ## Cacluating B_matrix
    B = B_matrix_calc(trimmed_ac, trimmed_state, trimmed_aerostate, trimmed_fcs, env, grav, du, tol)

    Linearize(A, B)
end

get_A(linear_sys::Linearize) = linear_sys.A
get_B(linear_sys::Linearize) = linear_sys.B

state_names = ["u", "v", "w", "p", "q", "r", "psi", "theta", "psi", "xe", "ye", "ze"]
control_names = ["delta_e", "delta_a", "delta_r", "delta_t"]
function get_named_A(linear_sys::Linearize)
    A = copy(linear_sys.A)
    convert(Any, A)
    A = hcat(["u", "v", "w", "p", "q", "r", "psi", "theta", "psi", "xe", "ye", "ze"], A)
    A = vcat(reshape([" ", "u", "v", "w", "p", "q", "r", "psi", "theta", "psi", "xe", "ye", "ze"], 1, 13), A)
    return A
end
function get_named_B(linear_sys::Linearize)
    B = copy(linear_sys.B)
    convert(Any, B)
    B = hcat(["u", "v", "w", "p", "q", "r", "psi", "theta", "psi", "xe", "ye", "ze"], B)
    B = vcat(reshape([" ", "delta_e", "delta_a", "delta_r", "delta_t"], 1, 5), B)
    return B
end

function get_lon_A(linear_sys::Linearize)
    A = copy(linear_sys.A)
    A = [
        A[1,1] A[3,1] A[5,1] A[8,1] A[10,1] A[12,1];
        A[1,3] A[3,3] A[5,3] A[8,3] A[10,3] A[12,3];
        A[1,5] A[3,5] A[5,5] A[8,5] A[10,5] A[12,5];
        A[1,8] A[3,8] A[5,8] A[8,8] A[10,8] A[12,8];
        A[1,10] A[3,10] A[5,10] A[8,10] A[10,10] A[12,10];
        A[1,12] A[3,12] A[5,12] A[8,12] A[10,12] A[12,12]
        ]

    return Array(transpose(A))
end

function get_lon_B(linear_sys::Linearize)
    B = copy(linear_sys.B)
    B = [
        B[1,1] B[1,4];
        B[3,1] B[3,4];
        B[5,1] B[5,4];
        B[8,1] B[8,4];
        B[10,1] B[10,4];
        B[12,1] B[12,4];
        ]

return B
end

function get_lat_A(linear_sys::Linearize)
    A = copy(linear_sys.A)
    A = [
        A[2,2] A[4,2] A[6,2] A[7,2] A[9,2] A[11,2];
        A[2,4] A[4,4] A[6,4] A[7,4] A[9,4] A[11,4];
        A[2,6] A[4,6] A[6,6] A[7,6] A[9,6] A[11,6];
        A[2,7] A[4,7] A[6,7] A[7,7] A[9,7] A[11,7];
        A[2,9] A[4,9] A[6,9] A[7,9] A[9,9] A[11,9];
        A[2,11] A[4,11] A[6,11] A[7,11] A[9,11] A[11,11];
        ]

    return Array(transpose(A))
end

function get_lat_B(linear_sys::Linearize)
    B = copy(linear_sys.B)
    B = [
        B[2,2] B[2,3];
        B[4,2] B[4,3];
        B[6,2] B[6,3];
        B[7,2] B[7,3];
        B[9,2] B[9,3];
        B[11,2] B[11,3]
        ]

return B
end

function A_matrix_calc(trimmed_ac, trimmed_state, trimmed_aerostate, trimmed_fcs, env, grav, dx_, tol)

    x = get_sixdof_euler_fixed_mass_state(trimmed_state)
    n = length(x)

    dx = dx_.*x      # purturbation in all state variables
    for i = 1:n   # set perturbations
        if dx[i] == 0.0
            dx[i] = dx_
            if i == n    # 12th state => ze should be negative
                dx[i] = -dx_
            end
        end
    end

    last = zeros(Float64, n,1)  # single column from A_matrix
    A = zeros(Float64, n,n)

    for j = 1:n
        xt = x
        for i = 1:10
            xt[j] = x[j] + dx[j]
            mass, inertia, forces, moments = recalculate_aircraft(xt, trimmed_ac, trimmed_fcs, grav)
            xd1 = six_dof_euler_fixed_mass(xt, mass, inertia, forces, moments)

            xt[j] = x[j] - dx[j]
            mass, inertia, forces, moments = recalculate_aircraft(xt, trimmed_ac, trimmed_fcs, grav)
            xd2 = six_dof_euler_fixed_mass(xt, mass, inertia, forces, moments)

            A[:, j] = (transpose(xd1 - xd2) ./ (2*dx[j]))
            if false in (max.(abs.(A[:,j] - last) ./ abs.(A[:,j] .+ 1e-12)) .< tol)
                break
            end
            dx[j] = 0.5*dx[j]
            last = A[:,j]
            iteration = i
            if iteration == 10
                println("not converged on A, column $j")
            end
        end
    end

    A_ = A.*2
    return A_
end

function B_matrix_calc(trimmed_ac, trimmed_state, trimmed_aerostate, trimmed_fcs, env, grav, du_, tol)

    x = get_sixdof_euler_fixed_mass_state(trimmed_state)
    n = length(x)
    u = get_controls(trimmed_fcs)
    m = length(u)

    du = du_.*u
    for i = 1:m
        if du[i] == 0.0
            du[i] = du_
        end
    end

    last = zeros(Float64, n,1)  # last column of matrix
    B = zeros(Float64, n,m)

    for j = 1:m
        ut = u
        for i = 1:10
            ut[j] = u[j] + du[j]    # adding du to u
            fcs = fcs_from_u(trimmed_fcs, ut)
            mass, inertia, forces, moments = recalculate_aircraft(x, trimmed_ac, trimmed_fcs, grav)
            xd1 = six_dof_euler_fixed_mass(x, mass, inertia, forces, moments)

            ut[j] = u[j] - du[j]
            fcs = fcs_from_u(trimmed_fcs, ut)
            mass, inertia, forces, moments = recalculate_aircraft(x, trimmed_ac, trimmed_fcs, grav)
            xd2 = six_dof_euler_fixed_mass(x, mass, inertia, forces, moments)
            B[:, j] = (transpose(xd1 - xd2) ./ (2*du[j]))
            if false in (max.(abs.(B[:,j] - last) ./ abs.(B[:,j] .+ 1e-12)) .< tol)
                break
            end
            du[j] = 0.5*du[j]  # reducing du further to improve accuracy of the step
            last = B[:,j]
            iteration = i
            if iteration == 10
                println("not converged on B, column $j")
            end
        end
    end

    B_ = B.*2
    return B_

end

#### utils      => only to be used in this file

"""
returns state::State from the state_array
"""
function state_from_x(x)
    pos = PositionEarth(x[10], x[11], x[12])
    vel = [x[1], x[2], x[3]]
    ang_vel = [x[4], x[5], x[6]]
    att = Attitude(x[7], x[8], x[9])

    return State(pos, att, vel, ang_vel, [0.0, 0.0, 0.0], [0.0, 0.0, 0.0])
end

"""
calculate Aircraft after modifying states
"""
function recalculate_aircraft(xt::Array, ac, fcs, grav)
    state = state_from_x(xt)
    env = Environment(state.position, atmos="ISAF16", wind="NoWind", grav="const")
    aerostate = AeroState(state, env)
    ac = calculate_aircraft(ac, fcs, aerostate, state, grav)
    mass = ac.mass_props.mass
    inertia = ac.mass_props.inertia
    forces = ac.pfm.forces
    moments = ac.pfm.moments

    return mass, inertia, forces, moments
end

"""
construct a control_array
"""
function get_controls(fcs::FCS)
    u = [fcs.stick_longitudinal.value, fcs.stick_lateral.value, fcs.pedals.value, fcs.thtl.value]

    return u
end

"""
returns fcs with setting controls to modified value as per control_array
"""
function fcs_from_u(fcs, u)
    set_stick_lon(fcs, u[1])
    set_stick_lat(fcs, u[2])
    set_pedals(fcs, u[3])
    set_thtl(fcs, u[4])

    return fcs
end


text:
# This file is a part of RRRMC.jl. License is MIT: http://github.com/carlobaldassi/RRRMC.jl/LICENCE.md

# This implements an unsorted set which allows for
# fast insertion-deletion and randomly choosing one
# element.

module ArraySets

using Random
using ExtractMacro
using ..Common

export ArraySet, check_consistency

import Base: push!, delete!, length, empty!,
             iterate, in, show
import Random: rand

mutable struct ArraySet
    N::Int
    v::IVec
    pos::IVec
    t::Int
    ArraySet(N::Integer) = new(N, zeros(Int, N), zeros(Int, N), 0)
end

function check_consistency(aset::ArraySet)
    @extract aset : N v pos t
    @assert 0 ≤ t ≤ N
    c = 0
    for i = 1:N
        pos[i] == 0 && continue
        c += 1
        @assert 1 ≤ pos[i] ≤ t (i, pos[i], t)
        @assert v[pos[i]] == i (v[1:t], pos, pos[i], i)
    end
    @assert c == t
    for i = 1:t
        @assert v[i] ≠ 0
        @assert pos[v[i]] == i
    end
end

length(aset::ArraySet) = aset.t

in(aset::ArraySet, i) = pos[i] ≠ 0

function show(io::IO, aset::ArraySet)
    @extract aset : v t
    print(io, "ArraySet(")
    for i = 1:(t-1)
        print(io, v[i], ',')
    end
    t > 0 && print(io, v[t])
    print(io, ')')
end

function push!(aset::ArraySet, i::Integer)
    @extract aset : v pos t
    #@assert pos[i] == 0
    @inbounds begin
    t += 1
    v[t] = i
    pos[i] = t
    aset.t = t
    end
end
function delete!(aset::ArraySet, i::Integer)
    @extract aset : v pos t
    @inbounds begin
    p = pos[i]
    #@assert 0 < p ≤ t
    v[p] = v[t]
    pos[v[p]] = p
    pos[i] = 0
    aset.t -= 1
    end
end

rand(aset::ArraySet) = rand(Random.GLOBAL_RNG, aset)
function rand(r::AbstractRNG, aset::ArraySet)
    @extract aset : v t
    @inbounds i = v[rand(r, 1:t)]
    return i
end

function empty!(aset::ArraySet)
    @extract aset : pos t
    fill!(pos, 0)
    aset.t = 0
    return aset
end

function iterate(aset::ArraySet, i = 1)
    i == aset.t + 1 && return nothing
    @extract aset : v
    @boundscheck 1 ≤ i ≤ aset.t
    @inbounds x = v[i]
    return (x, i + 1)
end

end # module


text:
module ImplicitEquations

import IntervalArithmetic: Interval, diam, isempty
using RecipesBase

include("predicates.jl")
include("intervals.jl")
include("tupper.jl")
include("asciigraphs.jl")
include("plot_recipe.jl")

export Lt, ≪, Le, ≦, Eq, ⩵, Neq, ≶, ≷, Ge, ≧, Gt, ≫
export screen, I_
export asciigraph

#export GRAPH, OInterval
#export Region, compute, negate_op
#export TRUE, FALSE, MAYBE

end # module


text:
#################################################################################################
#
# demo 9: im2dots with OpenCV.jl
# Photoshop processing effect - overlay of small circles over image
# https://opencv-code.com/tutorials/photo-to-colored-dot-patterns-with-opencv/#more-754
#################################################################################################

# Julia filename with full path
inputfile =  joinpath(Pkg.dir("OpenCV"), "./test/images/julia.png")
outfile = joinpath(Pkg.dir("OpenCV"), "./test/images/dotjulia.png")

cxx"""
void im2dots(const char *inputfile, const char *outfile)
{
    const std::string ifname = inputfile;
    const std::string ofname = outfile;

    cv::Mat src = cv::imread(ifname);
    if (!src.data) {
        exit(0);
    }

    cv::Mat dst = cv::Mat::zeros(src.size(), CV_8UC3);
    cv::Mat cir = cv::Mat::zeros(src.size(), CV_8UC1);
    int bsize = 10;

    for (int i = 0; i < src.rows; i += bsize)
    {
        for (int j = 0; j < src.cols; j += bsize)
        {
            cv::Rect rect = cv::Rect(j, i, bsize, bsize) &
                            cv::Rect(0, 0, src.cols, src.rows);

            cv::Mat sub_dst(dst, rect);
            sub_dst.setTo(cv::mean(src(rect)));

            cv::circle(
                cir,
                cv::Point(j+bsize/2, i+bsize/2),
                bsize/2-1,
                cv::Scalar(200,200,200), -1, 2
            );
        }
    }

    cv::Mat cir_32f;
    cir.convertTo(cir_32f, CV_32F);
    cv::normalize(cir_32f, cir_32f, 0, 1, cv::NORM_MINMAX);

    cv::Mat dst_32f;
    dst.convertTo(dst_32f, CV_32F);

    std::vector<cv::Mat> channels;
    cv::split(dst_32f, channels);
    for (int i = 0; i < channels.size(); ++i)
        channels[i] = channels[i].mul(cir_32f);

    cv::merge(channels, dst_32f);
    dst_32f.convertTo(dst, CV_8U);

    cv::imshow("dst", dst);
    cv::waitKey(0);
    cv::imwrite(ofname, dst);
    cv::destroyAllWindows();
}
"""

@cxx im2dots(pointer(inputfile), pointer(outfile))


text:
# Copyright © 2021 Alexander L. Hayes
# Apache 2.0 License

module TestRequest

using Test
using RelationalDatasets

function load_webkb()
    train, test = RelationalDatasets.load("webkb", "v0.0.4", fold=2)
    return length(train.facts), length(test.facts)
end

@testset "Test Loading webkb-v0.0.4-fold2" begin
    @test load_webkb() == (1344, 568)
end

end # module


text:
export  ViscoelasticityMaxwell

@doc raw"""
    mutable struct ViscoelasticityMaxwell
        ρ::Float64 # density
        E::Float64 # Young's modulus
        ν::Float64 # Poisson's ratio
        η::Float64 # Viscosity parameter
        λ::Float64
        μ::Float64 # Lame constants
        σ0::Float64 # stress at last time step
        σ0_::Float64 # σ0 to be updated in `commitHistory
    end

Maxwell model for viscoelasticity. For details, see [this post](https://kailaix.github.io/PoreFlow.jl/dev/viscoelasticity/)

The updating rule is

$$\sigma^{n+1} = H \epsilon^{n+1} + S \sigma^n  - H\epsilon^n$$

"""
mutable struct ViscoelasticityMaxwell
    ρ::Float64 # density
    E::Float64 # Young's modulus
    ν::Float64 # Poisson's ratio
    η::Float64 # Viscosity parameter
    λ::Float64
    μ::Float64 # Lame constants
    σ0::Array{Float64} # stress at last time step
    σ0_::Array{Float64} # σ0 to be updated in `commitHistory`
    ε0::Array{Float64} # strain at last time step
    ε0_::Array{Float64} # ε0 to be updated in `commitHistory`
end

function ViscoelasticityMaxwell(prop::Dict{String, Any})
    ρ = Float64(prop["rho"])
    E = Float64(prop["E"])
    ν = Float64(prop["nu"])
    η = Float64(prop["eta"])
    λ = E*ν/(1+ν)/(1-2ν)
    μ = E/(2(1+ν))
    σ0 = zeros(3); σ0_ = zeros(3)
    ε0 = zeros(3); ε0_ = zeros(3)
    ViscoelasticityMaxwell(ρ, E, ν, η, λ, μ, σ0, σ0_, ε0, ε0_)
end

function getStress(domain::ViscoelasticityMaxwell,  strain::Array{Float64},  Dstrain::Array{Float64}, Δt::Float64)
    # strain, Dstrain are defined at (1-αf)*Δt in the generalized α solver
    ϵ = strain; ϵ0 = domain.ε0
    σ0 = domain.σ0
    λ, μ, η = domain.λ, domain.μ, domain.η
    S = inv([
        1 + 2/3*μ*Δt/η -1/3*μ*Δt/η 0.0
        -1/3*μ*Δt/η 1+2/3*μ*Δt/η 0.0
        0.0 0.0 1+μ*Δt/η
    ])
    H = S * [2μ+λ λ 0.0
            λ 2μ+λ 0.0
            0.0 0.0 μ]
    σ = H * ϵ + S * σ0 - H * ϵ0
    dΔσdΔε = H
    # the Newton solver will iterate until convergence is reached.
    # Therefore, σ and ϵ are true values at (1-αf)*Δt
    # Strictly speaking, to plot σ and ϵ, you should use values at (1-αf)*Δt
    # Note AD solver is different from this solver.
    domain.σ0_ = σ
    domain.ε0_ = ϵ
    return σ, dΔσdΔε
end

function getTangent(domain::ViscoelasticityMaxwell)
    error("Not implemented")
end

function commitHistory(domain::ViscoelasticityMaxwell)
    domain.σ0 = domain.σ0_
    domain.ε0 = domain.ε0_
end


text:
module BlockArrays
using Base.Cartesian
using LinearAlgebra, ArrayLayouts

# AbstractBlockArray interface exports
export AbstractBlockArray, AbstractBlockMatrix, AbstractBlockVector, AbstractBlockVecOrMat
export Block, getblock, getblock!, setblock!, eachblock
export blockaxes, blocksize, blocklength, blockcheckbounds, BlockBoundsError, BlockIndex
export blocklengths, blocklasts, blockfirsts, blockisequal
export BlockRange, blockedrange, BlockedUnitRange

export BlockArray, BlockMatrix, BlockVector, BlockVecOrMat, mortar
export PseudoBlockArray, PseudoBlockMatrix, PseudoBlockVector, PseudoBlockVecOrMat

export undef_blocks, undef, findblock, findblockindex

import Base: @propagate_inbounds, Array, to_indices, to_index,
            unsafe_indices, first, last, size, length, unsafe_length,
            unsafe_convert,
            getindex, show,
            step,
            broadcast, eltype, convert, similar,
            @_inline_meta, _maybetail, tail, @_propagate_inbounds_meta, reindex,
            RangeIndex, Int, Integer, Number,
            +, -, min, max, *, isless, in, copy, copyto!, axes, @deprecate,
            BroadcastStyle, checkbounds, throw_boundserror
using Base: ReshapedArray, dataids

import Base: (:), IteratorSize, iterate, axes1, strides
import Base.Broadcast: broadcasted, DefaultArrayStyle, AbstractArrayStyle, Broadcasted
import LinearAlgebra: lmul!, rmul!, AbstractTriangular, HermOrSym, AdjOrTrans,
                        StructuredMatrixStyle
import ArrayLayouts: _fill_lmul!, MatMulVecAdd, MatMulMatAdd, MatLmulVec, MatLdivVec,
                        materialize!, MemoryLayout, sublayout, transposelayout, conjlayout,
                        triangularlayout, triangulardata

include("blockindices.jl")
include("blockaxis.jl")
include("abstractblockarray.jl")
include("blockarray.jl")
include("pseudo_blockarray.jl")
include("views.jl")
include("show.jl")
include("blockarrayinterface.jl")
include("blockbroadcast.jl")
include("blocklinalg.jl")

end # module


text:
"""
    Model Data
"""

struct ModelData{T,X,U,W}
    dynamics::Vector{Dynamics{T}}
    jacobian_state::Vector{X}
    jacobian_action::Vector{U}
	jacobian_parameter::Vector{W}
end

function model_data(dynamics::Vector{Dynamics{T}}) where T
	jacobian_state = [zeros(d.num_next_state, d.num_state) for d in dynamics]
    jacobian_action = [zeros(d.num_next_state, d.num_action) for d in dynamics]
	jacobian_parameter = [zeros(d.num_next_state, d.num_parameter) for d in dynamics]
    ModelData(dynamics, jacobian_state, jacobian_action, jacobian_parameter)
end

function reset!(data::ModelData)
    H = length(data.dynamics) + 1
    for t = 1:H-1
        fill!(data.jacobian_state[t], 0.0)
        fill!(data.jacobian_action[t], 0.0)
        fill!(data.jacobian_parameter[t], 0.0)
    end
end