text:
""" Data structure for (directed) wiring diagrams, aka string diagrams.

A (directed) wiring diagram consists of a collection of boxes with input and
output ports connected by wires. A box can be atomic (possessing no internal
structure) or can itself be a wiring diagram. Thus, wiring diagrams can be
nested recursively. Wiring diagrams are closely related to what the CS
literature calls "directed graphs with ports" or more simply "port graphs". The
main difference is that a wiring diagram has an "outer box": a wiring diagram
has its own ports that can be connected to the ports of its boxes.

This module provides a generic data structure for wiring diagrams. Arbitrary
data can be attached to the boxes, ports, and wires of a wiring diagram. The
diagrams are "abstract" in the sense that they cannot be directly rendered as
raster or vector graphics. However, they form a useful intermediate
representation that can be serialized to and from GraphML or translated into
Graphviz or other declarative diagram languages.
"""
module DirectedWiringDiagrams
export AbstractBox, Box, WiringDiagram, Wire, Port, PortKind,
  InputPort, OutputPort, input_ports, output_ports, input_id, output_id,
  outer_ids, boxes, box_ids, nboxes, nwires, box, wires, has_wire, graph,
  add_box!, add_boxes!, add_wire!, add_wires!, rem_box!, rem_boxes!, rem_wire!,
  rem_wires!, port_value, validate_ports, is_permuted_equal,
  all_neighbors, neighbors, outneighbors, inneighbors, in_wires, out_wires,
  singleton_diagram, induced_subdiagram, encapsulated_subdiagram,
  substitute, encapsulate

using Compat
using AutoHashEquals
using LightGraphs, MetaGraphs
import LightGraphs: all_neighbors, neighbors, outneighbors, inneighbors

# Data types
############

""" Kind of port: input or output.
"""
@enum PortKind InputPort OutputPort

""" A port on a box to which wires can be connected.
"""
@auto_hash_equals struct Port
  box::Int
  kind::PortKind
  port::Int
end
set_box(port::Port, box::Int) = Port(box, port.kind, port.port)

function Base.isless(p1::Port, p2::Port)::Bool
  # Lexicographic order.
  p1.box < p2.box ||
    (p1.box == p2.box &&
      (p1.kind < p2.kind || (p1.kind == p2.kind && p1.port < p2.port)))
end

""" A wire connecting one port to another.
"""
@auto_hash_equals struct Wire{Value}
  value::Value
  source::Port
  target::Port
end

Wire(value, src::Tuple{Int,PortKind,Int}, tgt::Tuple{Int,PortKind,Int}) =
  Wire(value, Port(src[1],src[2],src[3]), Port(tgt[1],tgt[2],tgt[3]))
Wire(value, src::Tuple{Int,Int}, tgt::Tuple{Int,Int}) =
  Wire(value, Port(src[1],OutputPort,src[2]), Port(tgt[1],InputPort,tgt[2]))
Wire(value, pair::Pair) = Wire(value, first(pair), last(pair))

Wire(src::Port, tgt::Port) = Wire(nothing, src, tgt)
Wire(src::Tuple, tgt::Tuple) = Wire(nothing, src, tgt)
Wire(pair::Pair) = Wire(nothing, first(pair), last(pair))

function Base.show(io::IO, wire::Wire)
  skip_kind = wire.source.kind == OutputPort && wire.target.kind == InputPort
  show_port = (io::IO, port::Port) -> begin
    if skip_kind
      print(io, "($(port.box),$(port.port))")
    else
      print(io, "($(port.box),$(string(port.kind)),$(port.port)")
    end
  end
  print(io, "Wire(")
  if !isnothing(wire.value)
    show(io, wire.value)
    print(io, ", ")
  end
  show_port(io, wire.source)
  print(io, " => ")
  show_port(io, wire.target)
  print(io, ")")
end

function Base.isless(w1::Wire, w2::Wire)::Bool
  # Lexicographic order.
  isless(w1.source, w2.source) ||
    (w1.source == w2.source && isless(w1.target, w2.target))
end

""" Internal wiring diagram data corresponding to `Port`. Do not use directly.
"""
@auto_hash_equals struct PortData
  kind::PortKind
  port::Int
end
to_port_data(port::Port) = PortData(port.kind, port.port)
from_port_data(port::PortData, v::Int) = Port(v, port.kind, port.port)

""" Internal wiring diagram data corresponding to `Wire`. Do not use directly.
"""
@auto_hash_equals struct WireData{Value}
  value::Value
  source::PortData
  target::PortData
end
function to_wire_data(wire::Wire)
  WireData(wire.value, to_port_data(wire.source), to_port_data(wire.target))
end
function from_wire_data(wire::WireData, edge::Edge)
  Wire(wire.value,
       from_port_data(wire.source, src(edge)),
       from_port_data(wire.target, dst(edge)))
end

""" Base type for any box (node) in a wiring diagram.

This type represents an arbitrary black box with inputs and outputs.
"""
abstract type AbstractBox end

input_ports(box::AbstractBox)::Vector = box.input_ports
output_ports(box::AbstractBox)::Vector = box.output_ports

""" An atomic box in a wiring diagram.

These boxes have no internal structure.
"""
@auto_hash_equals struct Box{Value} <: AbstractBox
  value::Value
  input_ports::Vector
  output_ports::Vector
end

Box(inputs::Vector, outputs::Vector) = Box(nothing, inputs, outputs)

function Base.show(io::IO, box::Box)
  print(io, "Box(")
  if !isnothing(box.value)
    show(io, box.value)
    print(io, ", ")
  end
  print(io, "[")
  join(io, [sprint(show, port) for port in box.input_ports], ",")
  print(io, "], [")
  join(io, [sprint(show, port) for port in box.output_ports], ",")
  print(io, "])")
end

""" A directed wiring diagram, aka a string diagram.

The wiring diagram is implemented using the following internal data structures.
A LightGraphs `DiGraph` stores the "skeleton" of the diagram: a simple directed
graph with the boxes as vertices and with an edge between two vertices iff there
is at least one wire between the corresponding boxes. There are two special
vertices, accessible via `input_id` and `output_id`, representing the input and
output ports, respectively.

The `DiGraph` is wrapped inside a `MetaDiGraph` to attach properties to the
vertices and edges. For each edge, an edge property stores the list of wires
between the source and target boxes.
"""
mutable struct WiringDiagram{Theory} <: AbstractBox
  graph::MetaDiGraph
  value::Any
  input_ports::Vector
  output_ports::Vector

  function WiringDiagram{T}(value, inputs::Vector, outputs::Vector) where T
    graph = MetaDiGraph()
    diagram = new{T}(graph, value, inputs, outputs)
    add_vertices!(graph, 2)
    return diagram
  end
  function WiringDiagram(d::WiringDiagram{T}) where T
    # Copy constructor for shallow copy.
    new{T}(copy(d.graph), d.value, d.input_ports, d.output_ports)
  end
end

function WiringDiagram{T}(inputs::Vector, outputs::Vector) where T
  WiringDiagram{T}(nothing, inputs, outputs)
end
WiringDiagram(args...) = WiringDiagram{Any}(args...)

input_id(::WiringDiagram) = 1
output_id(::WiringDiagram) = 2
outer_ids(::WiringDiagram) = (1,2)

""" Check equality of wiring diagrams.

Warning: This method checks exact equality of the underlying graph
representation, not mathematical equality which involves graph isomorphism.

See also: `is_permuted_equal`
"""
function Base.:(==)(d1::WiringDiagram, d2::WiringDiagram)
  (input_ports(d1) == input_ports(d2) && output_ports(d1) == output_ports(d2) &&
   d1.value == d2.value && graph(d1) == graph(d2) &&
   boxes(d1) == boxes(d2) && sort!(wires(d1)) == sort!(wires(d2)))
end

""" Check equality of wiring diagram under permutation of boxes.

When the boxes in the first diagram `d1` are permuted according to `σ`,
does it become identical to the second diagram `d2`?
"""
function is_permuted_equal(d1::WiringDiagram, d2::WiringDiagram, σ::Vector{Int})
  @assert nboxes(d1) == length(σ) && nboxes(d2) == length(σ)
  d1_ids, d2_ids = box_ids(d1), box_ids(d2)
  box_map = Dict{Int,Int}(d1_ids[σ[i]] => d2_ids[i] for i in eachindex(σ))
  is_induced_equal(d1, d2, box_map)
end
function is_induced_equal(d1::WiringDiagram, d2::WiringDiagram, box_map::Dict{Int,Int})
  box_map[input_id(d1)] = input_id(d2)
  box_map[output_id(d1)] = output_id(d2)
  map_wire = wire::Wire -> Wire(wire.value,
    set_box(wire.source, box_map[wire.source.box]),
    set_box(wire.target, box_map[wire.target.box]))
  (input_ports(d1) == input_ports(d2) && output_ports(d1) == output_ports(d2) &&
   all(box(d1,v) == box(d2,box_map[v]) for v in box_ids(d1)) &&
   sort!(map(map_wire, wires(d1))) == sort!(wires(d2)))
end

Base.copy(diagram::WiringDiagram) = WiringDiagram(diagram)

function Base.show(io::IO, diagram::WiringDiagram{T}) where T
  sshowcompact = x -> sprint(show, x, context=:compact => true)
  print(io, "WiringDiagram")
  if T != Any
    print(io, "{$T}")
  end
  print(io, "(")
  if !isnothing(diagram.value)
    show(io, diagram.value)
    print(io, ", ")
  end
  print(io, "[")
  join(io, map(sshowcompact, input_ports(diagram)), ",")
  print(io, "], [")
  join(io, map(sshowcompact, output_ports(diagram)), ",")
  print(io, "], ")
  if get(io, :compact, false)
    print(io, "{$(nboxes(diagram)) boxes}, {$(nwires(diagram)) wires}")
  else
    print(io, "\n[ $(input_id(diagram)) => {inputs},\n  ")
    print(io, "$(output_id(diagram)) => {outputs},\n  ")
    join(io, [ "$v => $(sshowcompact(box(diagram, v)))"
               for v in box_ids(diagram) ], ",\n  ")
    print(io, " ],\n[ ")
    join(io, map(sshowcompact, wires(diagram)), ",\n  ")
    print(io, " ]")
  end
  print(io, ")")
end

# Imperative interface
######################

# Basic accessors.

box(f::WiringDiagram, v::Int) = get_prop(f.graph, v, :box)
boxes(f::WiringDiagram) = AbstractBox[ box(f,v) for v in box_ids(f) ]
nboxes(f::WiringDiagram) = nv(graph(f)) - 2

function box_ids(f::WiringDiagram)
  Int[ v for v in 1:nv(graph(f)) if !(v in outer_ids(f)) ]
end

function wires(f::WiringDiagram, edge::Edge)
  if has_edge(f.graph, edge)
    Wire[ from_wire_data(data, edge) for data in get_prop(f.graph,edge,:wires) ]
  else
    Wire[]
  end
end
wires(f::WiringDiagram, src::Int, tgt::Int) = wires(f, Edge(src,tgt))
wires(f::WiringDiagram) = vcat((wires(f,e) for e in edges(graph(f)))...)
nwires(f::WiringDiagram) =
  sum(Int[ length(get_prop(f.graph,e,:wires)) for e in edges(graph(f)) ])

function has_wire(f::WiringDiagram, src::Int, tgt::Int)
  has_edge(graph(f), Edge(src, tgt))
end
function has_wire(f::WiringDiagram, wire::Wire)
  wire in wires(f, wire.source.box, wire.target.box)
end
has_wire(f::WiringDiagram, pair::Pair) = has_wire(f, Wire(pair))

function input_ports(f::WiringDiagram, v::Int)
  if v == input_id(f)
    error("Input vertex does not have input ports within wiring diagram")
  elseif v == output_id(f)
    output_ports(f)
  else
    input_ports(box(f, v))
  end
end

function output_ports(f::WiringDiagram, v::Int)
  if v == input_id(f)
    input_ports(f)
  elseif v == output_id(f)
    error("Output vertex does not have output ports within wiring diagram")
  else
    output_ports(box(f, v))
  end
end

function port_value(f::WiringDiagram, port::Port)
  get_ports = port.kind == InputPort ? input_ports : output_ports
  get_ports(f, port.box)[port.port]
end

# Graph mutation.

function add_box!(f::WiringDiagram, box::AbstractBox)
  @assert add_vertex!(f.graph, :box, box)
  return nv(f.graph)
end

function add_boxes!(f::WiringDiagram, boxes)
  for box in boxes
    add_box!(f, box)
  end
end

function rem_box!(f::WiringDiagram, v::Int)
  @assert !(v in outer_ids(f))
  rem_vertex!(f.graph, v)
end

function rem_boxes!(f::WiringDiagram, vs)
  # Remove boxes in descending order of vertex ID to maintain ID stability.
  for v in sort!(collect(vs), rev=true)
    rem_box!(f, v)
  end
end

function add_wire!(f::WiringDiagram, wire::Wire)
  # Check compatibility of port types.
  validate_ports(port_value(f, wire.source), port_value(f, wire.target))

  # Add edge and edge properties.
  edge = Edge(wire.source.box, wire.target.box)
  if !has_edge(f.graph, edge)
    add_edge!(f.graph, src(edge), dst(edge), :wires, WireData[])
  end
  push!(get_prop(f.graph, edge, :wires), to_wire_data(wire))
end
add_wire!(f::WiringDiagram, pair::Pair) = add_wire!(f, Wire(pair))

function add_wires!(f::WiringDiagram, wires)
  for wire in wires
    add_wire!(f, wire)
  end
end

function rem_wire!(f::WiringDiagram, wire::Wire)
  edge = Edge(wire.source.box, wire.target.box)
  edge_data = to_wire_data(wire)
  wires = get_prop(f.graph, edge, :wires)
  deleteat!(wires, findlast(d -> d == edge_data, wires))
  if isempty(wires)
    rem_edge!(f.graph, edge)
  end
end
rem_wire!(f::WiringDiagram, pair::Pair) = rem_wire!(f, Wire(pair))

function rem_wires!(f::WiringDiagram, wires)
  for wire in wires
    rem_wire!(f, wire)
  end
end

function rem_wires!(f::WiringDiagram, src::Int, tgt::Int)
  rem_edge!(f.graph, Edge(src, tgt))
end

""" Check compatibility of source and target ports.

The default implementation is a no-op.
"""
function validate_ports(source_port, target_port) end

# Graph properties.

""" Retrieve the underlying LightGraphs graph.

Do not mutate it! All mutations should pass through the `WiringDiagram` methods:
`add_box!`, `rem_box!`, etc.
"""
graph(diagram::WiringDiagram) = diagram.graph.graph

# Convenience methods delegated to LightGraphs.
all_neighbors(d::WiringDiagram, v::Int) = all_neighbors(graph(d), v)
neighbors(d::WiringDiagram, v::Int) = neighbors(graph(d), v)
outneighbors(d::WiringDiagram, v::Int) = outneighbors(graph(d), v)
inneighbors(d::WiringDiagram, v::Int) = inneighbors(graph(d), v)

""" Get all wires coming into or out of the box.
"""
function wires(d::WiringDiagram, v::Int)
  result = wires(d, v, v)
  for u in inneighbors(d, v)
    if u != v
      append!(result, wires(d, u, v))
    end
  end
  for u in outneighbors(d, v)
    if u != v
      append!(result, wires(d, v, u))
    end
  end
  result
end

""" Get all wires coming into the box.
"""
function in_wires(d::WiringDiagram, v::Int)
  result = Wire[]
  for u in inneighbors(d, v)
    append!(result, wires(d, u, v))
  end
  result
end

""" Get all wires coming into the port.
"""
function in_wires(d::WiringDiagram, port::Port)
  result = Wire[]
  for v in inneighbors(d, port.box)
    for wire in wires(d, v, port.box)
      if wire.target == port
        push!(result, wire)
      end
    end
  end
  result
end
function in_wires(d::WiringDiagram, v::Int, port::Int)
  in_wires(d, Port(v, InputPort, port))
end

""" Get all wires coming out of the box.
"""
function out_wires(d::WiringDiagram, v::Int)
  result = Wire[]
  for u in outneighbors(d, v)
    append!(result, wires(d, v, u))
  end
  result
end

""" Get all wires coming out of the port.
"""
function out_wires(d::WiringDiagram, port::Port)
  result = Wire[]
  for v in outneighbors(d, port.box)
    for wire in wires(d, port.box, v)
      if wire.source == port
        push!(result, wire)
      end
    end
  end
  result
end
function out_wires(d::WiringDiagram, v::Int, port::Int)
  out_wires(d, Port(v, OutputPort, port))
end

# Other constructors.

""" A wiring diagram with a single box connected to its outer ports.
"""
function singleton_diagram(T::Type, box::AbstractBox)
  inputs, outputs = input_ports(box), output_ports(box)
  d = WiringDiagram{T}(inputs, outputs)
  v = add_box!(d, box)
  add_wires!(d, ((input_id(d),i) => (v,i) for i in eachindex(inputs)))
  add_wires!(d, ((v,i) => (output_id(d),i) for i in eachindex(outputs)))
  return d
end
singleton_diagram(box::AbstractBox) = singleton_diagram(Any, box)

""" The wiring diagram induced by a subset of its boxes.

See also `encapsulated_subdiagram`.
"""
function induced_subdiagram(d::WiringDiagram{T}, vs::Vector{Int}) where T
  sub = WiringDiagram{T}(input_ports(d), output_ports(d))
  vmap = Dict(input_id(d) => input_id(sub), output_id(d) => output_id(sub))
  for v in vs
    vmap[v] = add_box!(sub, box(d, v))
  end
  for wire in wires(d)
    src, tgt = wire.source, wire.target
    if haskey(vmap, src.box) && haskey(vmap, tgt.box)
      add_wire!(sub,
        Wire(set_box(src, vmap[src.box]), set_box(tgt, vmap[tgt.box])))
    end
  end
  return sub
end

# Substitution
##############

""" Substitute wiring diagrams for boxes.

Performs one or more substitutions. When performing multiple substitutions, the
substitutions are simultaneous.

This operation implements the operadic composition of wiring diagrams
(`ocompose`).
"""
function substitute(d::WiringDiagram; kw...)
  substitute(d, filter(v -> box(d,v) isa WiringDiagram, box_ids(d)); kw...)
end
function substitute(d::WiringDiagram, v::Int; kw...)
  substitute(d, v, box(d,v)::WiringDiagram; kw...)
end
function substitute(d::WiringDiagram, vs::Vector{Int}; kw...)
  substitute(d, vs, WiringDiagram[ box(d,v) for v in vs ]; kw...)
end
function substitute(d::WiringDiagram, v::Int, sub::WiringDiagram; kw...)
  substitute(d, [v], [sub]; kw...)
end

function substitute(d::WiringDiagram{T}, vs::Vector{Int}, subs::Vector{<:WiringDiagram};
                    merge_wire_values=default_merge_wire_values) where T
  # In outline, the algorithm is:
  #
  # 1. Create an empty wiring diagram.
  # 2. Add *all* boxes of original diagram and the diagrams to be substituted
  #    (in the appropriate order).
  # 3. Add *all* wires of original diagram and the diagrams to be substituted.
  # 4. Remove the boxes that were substituted (hence also removing extraneous
  #    wires from step 3).
  #
  # This may seem convoluted, but it is the simplest way I know to handle the
  # problem of *instantaneous wires*. Some authors ban instantaneous wires, but
  # we need them to represent identities, braidings, etc.
  @assert length(vs) == length(subs)
  result = WiringDiagram{T}(d.value, input_ports(d), output_ports(d))

  # Add boxes by interleaving, in the correct order, the non-substituted boxes
  # of the original diagram and the internal boxes of the substituted diagrams.
  # At the very end, add the substituted boxes too.
  vmap = Dict(input_id(d) => input_id(result), output_id(d) => output_id(result))
  sub_maps = Dict(zip(vs, ((sub, Dict{Int,Int}()) for sub in subs)))
  for v in box_ids(d)
    if haskey(sub_maps, v)
      sub, sub_map = sub_maps[v]
      for u in box_ids(sub)
        sub_map[u] = add_box!(result, box(sub, u))
      end
    else
      vmap[v] = add_box!(result, box(d, v))
    end
  end
  for v in vs
    vmap[v] = add_box!(result, box(d, v))
  end

  # Add the wires of the original diagram, then add the internal wires of the
  # substituted diagrams.
  for wire in wires(d)
    add_wire!(result, Wire(wire.value,
      set_box(wire.source, vmap[wire.source.box]),
      set_box(wire.target, vmap[wire.target.box])))
  end
  for v in vs
    substitute_wires!(result, vmap[v], sub_maps[v]...;
      merge_wire_values=merge_wire_values)
  end

  # Finally, remove the substituted boxes. Because they were added last, this
  # will not change the IDs of the other boxes.
  rem_boxes!(result, (vmap[v] for v in vs))
  result
end

""" Substitute wires of sub-diagram into containing wiring diagram.
"""
function substitute_wires!(d::WiringDiagram, v::Int,
                           sub::WiringDiagram, sub_map::Dict{Int,Int};
                           merge_wire_values=default_merge_wire_values)
  for wire in wires(sub)
    src = get(sub_map, wire.source.box, 0)
    tgt = get(sub_map, wire.target.box, 0)
    # Special case: wire from input port to output port.
    if wire.source.box == input_id(sub) && wire.target.box == output_id(sub)
      for in_wire in in_wires(d, v, wire.source.port)
        for out_wire in out_wires(d, v, wire.target.port)
          add_wire!(d, Wire(
            merge_wire_values(in_wire.value, wire.value, out_wire.value),
            in_wire.source, out_wire.target))
        end
      end
    # Special case: wire from input port to internal box.
    elseif wire.source.box == input_id(sub)
      for in_wire in in_wires(d, v, wire.source.port)
        add_wire!(d, Wire(
          merge_wire_values(in_wire.value, wire.value, nothing),
          in_wire.source, set_box(wire.target, tgt)))
      end
    # Special case: wire from internal box to output port.
    elseif wire.target.box == output_id(sub)
      for out_wire in out_wires(d, v, wire.target.port)
        add_wire!(d, Wire(
          merge_wire_values(nothing, wire.value, out_wire.value),
          set_box(wire.source, src), out_wire.target))
      end
    # Default case: wire between two internal boxes.
    else
      add_wire!(d, Wire(
        merge_wire_values(nothing, wire.value, nothing),
        set_box(wire.source, src), set_box(wire.target, tgt)))
    end
  end
end

default_merge_wire_values(::Any, middle::Any, ::Any) = middle

# Encapsulation
###############

""" Encapsulate multiple boxes within a single sub-diagram.

This operation is a (one-sided) inverse to subsitution (see `substitute`).
"""
function encapsulate(d::WiringDiagram, vs::Vector{Int}; value=nothing, kw...)
  encapsulate(d, [vs]; values=[value], kw...)
end

function encapsulate(d::WiringDiagram{T}, vss::Vector{Vector{Int}};
    discard_boxes::Bool=false, make_box=Box, values=nothing) where T
  if isempty(vss); return d end
  if any(isempty(vs) for vs in vss)
    error("Cannot encapsulate an empty set of boxes")
  end
  if !allunique(reduce(vcat, vss))
    error("Cannot encapsulate overlapping sets of boxes")
  end
  if isnothing(values)
    values = repeat([nothing], length(vss))
  end
  result = WiringDiagram{T}(d.value, input_ports(d), output_ports(d))

  # Add boxes, both encapsulated and non-encapsulated, to new wiring diagram.
  encapsulated_representatives = Dict(
    minimum(vs) => (vs, value) for (vs, value) in zip(vss, values))
  all_encapsulated = Set(v for vs in vss for v in vs)
  vmap = Dict(input_id(d) => input_id(result), output_id(d) => output_id(result))
  port_map = Dict{Port,Port}()
  for v in box_ids(d)
    if haskey(encapsulated_representatives, v)
      vs, value = encapsulated_representatives[v]
      sub, sub_map = encapsulated_subdiagram(d, vs;
        discard_boxes=discard_boxes, make_box=make_box, value=value)
      subv = add_box!(result, sub)
      merge!(port_map, Dict(
        port => from_port_data(data, subv) for (port, data) in sub_map))
    elseif v ∉ all_encapsulated
      vmap[v] = add_box!(result, box(d, v))
    end
  end

  # Add wires to new wiring diagram.
  for wire in wires(d)
    src, tgt = wire.source, wire.target
    new_src = if haskey(vmap, src.box); set_box(src, vmap[src.box])
      elseif haskey(port_map, src); port_map[src]
      else; continue end
    new_tgt = if haskey(vmap, tgt.box); set_box(tgt, vmap[tgt.box])
      elseif haskey(port_map, tgt); port_map[tgt]
      else; continue end
    add_wire!(result, Wire(new_src, new_tgt))
  end
  result
end

""" Create an encapsulating box for a set of boxes in a wiring diagram.

To a first approximation, the union of input ports of the given boxes will
become the inputs ports of the encapsulating box and likewise for the output
ports. However, when copies or merges occur, as in a cartesian or cocartesian
category, a simplification procedure may reduce the number of ports on the
encapsulating box.

Specifically:

1. Each input port of an encapsulated box will have at most one incoming wire
from the encapsulating outer box, and each output port of an encapsulated box
will have at most one outgoing wire to the encapsulating outer box.

2. A set of ports connected to the same outside (non-encapsulated) ports will be
simplified into a single port of the encapsulating box.

See also `induced_subdiagram`.
"""
function encapsulated_subdiagram(d::WiringDiagram{T}, vs::Vector{Int};
    discard_boxes::Bool=false, make_box=Box, value=nothing) where T
  # Add encapsulated box to new wiring diagram.
  inputs, outputs = [], []
  result = discard_boxes ? nothing : WiringDiagram{T}(value, inputs, outputs)
  vmap = if discard_boxes
    Dict(v => nothing for v in vs)
  else
    Dict(v => add_box!(result, box(d, v)) for v in vs)
  end

  # Process wires into, or out of, encapsulated boxes.
  port_map = Dict{Port,PortData}()
  inner_port_map = Dict{Tuple{Vector{Port},Any},Int}()
  for v in vs
    # Add input ports to encapsulating box and corresponding wire.
    for (port, value) in enumerate(input_ports(d, v))
      tgt = Port(v, InputPort, port)
      srcs = sort!([ wire.source for wire in in_wires(d, tgt)
                     if !haskey(vmap, wire.source.box) ])
      if isempty(srcs) continue end
      src = get!(inner_port_map, (srcs, value)) do
        push!(inputs, value)
        port_data = port_map[tgt] = PortData(InputPort, length(inputs))
        port_data.port
      end
      if discard_boxes; continue end
      add_wire!(result,
        Wire(Port(input_id(result), OutputPort, src), set_box(tgt, vmap[v])))
    end

    # Add output ports to encapsulating box and corresponding wire.
    for (port, value) in enumerate(output_ports(d, v))
      src = Port(v, OutputPort, port)
      tgts = sort([ wire.target for wire in out_wires(d, src)
                    if !haskey(vmap, wire.target.box) ])
      if isempty(tgts) continue end
      tgt = get!(inner_port_map, (tgts, value)) do
        push!(outputs, value)
        port_data = port_map[src] = PortData(OutputPort, length(outputs))
        port_data.port
      end
      if discard_boxes; continue end
      add_wire!(result,
        Wire(set_box(src, vmap[v]), Port(output_id(result), InputPort, tgt)))
    end

    # Add wires between two encapsulated boxes.
    if discard_boxes; continue end
    for wire in out_wires(d, v)
      src, tgt = wire.source, wire.target
      if haskey(vmap, src.box) && haskey(vmap, tgt.box) # Clause #1 always true.
        add_wire!(result,
          Wire(set_box(src, vmap[src.box]), set_box(tgt, vmap[tgt.box])))
      end
    end
  end

  # Yield input and output port lists with the tightest possible types.
  inputs, outputs = [ x for x in inputs ], [ x for x in outputs ]
  if discard_boxes
    result = make_box(value, inputs, outputs)
  else
    result.input_ports, result.output_ports = inputs, outputs
  end
  (result, port_map)
end

end


text:
################################################################################
# Damper Force
################################################################################

function damper_force(joint::Translational{T},
        xa::AbstractVector, va::AbstractVector, qa::Quaternion, ωa::AbstractVector,
        xb::AbstractVector, vb::AbstractVector, qb::Quaternion, ωb::AbstractVector,
        timestep;
        unitary::Bool=false) where T

    damper = unitary ? 1.0 : joint.damper
    Aᵀ = zerodimstaticadjoint(nullspace_mask(joint))
    input = damper * Aᵀ * -minimal_velocities(joint, xa, va, qa, ωa, xb, vb, qb, ωb, timestep) # in the a frame

    return input
end

function damper_force(relative::Symbol, joint::Translational{T},
    xa::AbstractVector, va::AbstractVector, qa::Quaternion, ωa::AbstractVector,
    xb::AbstractVector, vb::AbstractVector, qb::Quaternion, ωb::AbstractVector,
    timestep;
    unitary::Bool=false) where T

    input = damper_force(joint, xa, va, qa, ωa, xb, vb, qb, ωb, timestep, unitary=unitary) # in the a frame
    inputa = impulse_transform(relative, joint, xa, qa, xb, qb) * input

    return inputa
end

damper_impulses(relative::Symbol, joint::Translational, pbody::Node, cbody::Node, timestep;
    unitary::Bool=false) =
    timestep * damper_force(relative, joint,
        current_configuration_velocity(pbody.state)...,
        current_configuration_velocity(cbody.state)...,
        timestep;
        unitary=unitary)
damper_impulses(relative::Symbol, joint::Translational{T,3}, pbody::Node, cbody::Node, timestep; unitary::Bool=false) where T = szeros(T, 6)

################################################################################
# Damper Jacobian
################################################################################

function damper_force_jacobian_configuration(jacobian_relative::Symbol,
        joint::Translational{T},
        xa::AbstractVector, va::AbstractVector, qa::Quaternion, ωa::AbstractVector,
        xb::AbstractVector, vb::AbstractVector, qb::Quaternion, ωb::AbstractVector,
        timestep;
        unitary::Bool=false) where T

    damper = unitary ? 1.0 : joint.damper
    Aᵀ = zerodimstaticadjoint(nullspace_mask(joint))
    ∇input = damper * Aᵀ * -minimal_velocities_jacobian_configuration(jacobian_relative, joint, xa, va, qa, ωa, xb, vb, qb, ωb, timestep) # in the a frame

    return ∇input
end

function damper_force_jacobian_velocity(jacobian_relative::Symbol,
    joint::Translational{T},
    xa::AbstractVector, va::AbstractVector, qa::Quaternion, ωa::AbstractVector,
    xb::AbstractVector, vb::AbstractVector, qb::Quaternion, ωb::AbstractVector,
    timestep;
    unitary::Bool=false) where T

    damper = unitary ? 1.0 : joint.damper
    Aᵀ = zerodimstaticadjoint(nullspace_mask(joint))
    ∇input = damper * Aᵀ * -minimal_velocities_jacobian_velocity(jacobian_relative, joint, xa, va, qa, ωa, xb, vb, qb, ωb, timestep) # in the a frame

    return ∇input
end

function damper_jacobian_configuration(relative::Symbol, jacobian_relative::Symbol,
        joint::Translational{T},
        xa::AbstractVector, va::AbstractVector, qa::Quaternion, ωa::AbstractVector,
        xb::AbstractVector, vb::AbstractVector, qb::Quaternion, ωb::AbstractVector,
        timestep::T;
        unitary::Bool=false) where T

    input = damper_force(joint, xa, va, qa, ωa, xb, vb, qb, ωb, timestep, unitary=unitary)

    ∇xq = impulse_transform(relative, joint, xa, qa, xb, qb) *
        damper_force_jacobian_configuration(jacobian_relative, joint, xa, va, qa, ωa, xb, vb, qb, ωb, timestep, unitary=unitary)
    ∇xq += impulse_transform_jacobian(relative, jacobian_relative, joint, xa, qa, xb, qb, input)

    return timestep * ∇xq
end

function damper_jacobian_configuration(relative::Symbol, jacobian::Symbol,
    joint::Translational,
    pbody::Node, cbody::Node,
    timestep::T;
    attjac::Bool = true) where T

    return damper_jacobian_configuration(relative, jacobian,
        joint,
        current_configuration_velocity(pbody.state)...,
        current_configuration_velocity(cbody.state)...,
        timestep; unitary=false)
end

function damper_jacobian_velocity(relative::Symbol, jacobian_relative::Symbol,
        joint::Translational{T},
        xa::AbstractVector, va::AbstractVector, qa::Quaternion, ωa::AbstractVector,
        xb::AbstractVector, vb::AbstractVector, qb::Quaternion, ωb::AbstractVector,
        timestep::T;
        unitary::Bool=false) where T

    ∇xq = impulse_transform(relative, joint, xa, qa, xb, qb) *
        damper_force_jacobian_velocity(jacobian_relative, joint, xa, va, qa, ωa, xb, vb, qb, ωb, timestep, unitary=unitary)

    return timestep * ∇xq
end

function damper_jacobian_velocity(relative::Symbol, jacobian::Symbol,
    joint::Translational,
    pbody::Node, cbody::Node,
    timestep::T, attjac::Bool=true) where T

    return damper_jacobian_velocity(relative, jacobian,
        joint,
        current_configuration_velocity(pbody.state)...,
        current_configuration_velocity(cbody.state)...,
        timestep; unitary=false)
end

damper_jacobian_configuration(relative::Symbol, jacobian::Symbol, joint::Translational{T,3}, pbody::Node, cbody::Node, timestep::T; attjac::Bool = true, unitary::Bool=false) where T = attjac ? szeros(T, 6, 6) : szeros(T, 6, 7)
damper_jacobian_velocity(relative::Symbol, jacobian::Symbol, joint::Translational{T,3}, pbody::Node, cbody::Node, timestep::T, unitary::Bool=false) where T = szeros(T, 6, 6)

ext: jl


text:
############################################################################################
"""
$(TYPEDEF)

Contains several useful information for sampled parameter.

# Fields
$(TYPEDFIELDS)
"""
struct PrintedParameter{A<:Tuple, B<:Tuple}
    "Unique tagged parameter of all kernels."
    tagged::A
    "Parameter names based on their dimension"
    printed::B
    function PrintedParameter(
        tagged::A,
        printed::B
    ) where {A<:Tuple, B<:Tuple}
        return new{A,B}(tagged, printed)
    end
end

############################################################################################
"""
$(TYPEDEF)

Contains several useful information for constructing sampler.

# Fields
$(TYPEDFIELDS)
"""
struct SamplingInfo{A<:PrintedParameter, U<:BaytesCore.UpdateBool, B<:BaytesCore.UpdateBool}
    "Parameter settings for printing."
    printedparam::A
    "Total number of sampling iterations."
    iterations::Int64
    "Burnin iterations."
    burnin::Int64
    "Number of consecutive samples taken for diagnostics output."
    thinning::Int64
    "Number of algorithms used while sampling."
    Nalgorithms::Int64
    "Number of chains used while sampling."
    Nchains::Int64
    "Boolean if sampler need to be updated after a proposal run. This is the case for, e.g. PMCMC., or for adaptive tempering."
    captured::U
    "Boolean if temperature is adapted for target function."
    tempered::B
    function SamplingInfo(
        printedparam::A,
        iterations::Int64,
        burnin::Int64,
        thinning::Int64,
        Nalgorithms::Int64,
        Nchains::Int64,
        captured::U,
        tempered::B
    ) where {A<:PrintedParameter, U<:BaytesCore.UpdateBool, B<:BaytesCore.UpdateBool}
        return new{A, U, B}(printedparam, iterations, burnin, thinning, Nalgorithms, Nchains, captured, tempered)
    end
end

############################################################################################
"""
$(SIGNATURES)
Check if we can capture results from last proposal step. Typically only possible if a single MCMC or SMC step applied.

# Examples
```julia
```

"""
function update(datatune::DataTune, temperingadaption::B, smc::SMCConstructor) where {B<:BaytesCore.UpdateBool}
    #!NOTE: If a single SMC constructor is used, we can capture previous results and do not need to update sampler before new iteration.
    #!NOTE: SMC only updates parameter if UpdateTrue(). Log-target/gradients will always be updated if jitter step applied.
    return BaytesCore.UpdateFalse()
end
function update(datatune::DataTune{<:B}, temperingadaption::BaytesCore.UpdateFalse, mcmc::MCMCConstructor) where {B<:Batch}
    #!NOTE: If a MCMC constructor is used and no tempering is applied, we can capture previous results and do not need to update sampler before new iteration.
    #!NOTE: This only holds for Batch data in MCMC case
    #!NOTE: MCMC updates log target evaluation and eventual gradients if UpdateTrue()
    return BaytesCore.UpdateFalse()
end
function update(datatune::DataTune, temperingadaption, args...)
    #!NOTE: In all other cases, we have to update kernels for proposal before new propagation
    return BaytesCore.UpdateTrue()
end

############################################################################################
"""
$(SIGNATURES)
Obtain maximum number of iterations based on datatune.
This will always be user input, except if rolling/expanding data is used, in which case iterations are capped.

# Examples
```julia
```

"""
function maxiterations(datatune::DataTune, iterations::Integer)
    return iterations
end
function maxiterations(
    datatune::DataTune{<:D}, iterations::Integer
) where {D<:Union{Expanding,Rolling}}
    # Compute maximum iterations based on data stepsize
    datatune_max = maximum(datatune.config.size) - datatune.structure.index.current
    # Compare to Sampling iteration specification
    if datatune_max >= iterations
        println("Max iterations in Expanding/Rolling datatune set to ", datatune_max, ". Iterations in SampleDefault are ", iterations, ". Sampling iterations set to ", iterations)
        return iterations
    elseif datatune_max < iterations
        println("Max iterations in Expanding/Rolling datatune set to ", datatune_max, ". Iterations in SampleDefault are ", iterations, ". Sampling iterations set to ", datatune_max)
        return datatune_max
    end
end

############################################################################################
"""
$(SIGNATURES)
Obtain all parameter where output diagnostics can be printed. Separate step to showparam in case parameter is increasing over time and cannot be printed

# Examples
```julia
```

"""
function printedparam(datatune::DataTune, sym, algorithm...)
    return sym
end
function printedparam(
    datatune::DataTune{<:D}, sym, smc::SMCConstructor
) where {D<:Expanding}
    #!TODO: If SMC is applied on Expanding data sequence, latent data is expanding over time - no output diagnostics for this yet.
    return smc.kernel isa SMC2Constructor ? Tuple(setdiff(sym, smc.kernel.propagation.sym)) : sym
end

"""
$(SIGNATURES)
Return all unique targetted parameter, and parameter where diagnostics will be printed. Separate to get_sym.

# Examples
```julia
```

"""
function showparam(model::ModelWrapper, datatune::DataTune, constructor...)
    # Get all tracked symbols in constructor
    sym = BaytesCore.to_Tuple.(map(BaytesCore.get_sym, constructor))
    unique_sym = unique(collect(Iterators.flatten(sym)))
    # Check if all tracked symbols are not fixed in model
    ArgCheck.@argcheck all(haskey(model.val, symbol) for symbol in unique_sym)
    # Check if one of the parameter has increasing dimension (SMC2)
    printed_sym = printedparam(datatune, unique_sym, constructor...)
    # Check if any of the remaining parameter is fixed, and we can thus avoid printing
    fixed = [model.info.constraint[sym] isa ModelWrappers.Fixed for sym in printed_sym]
    # Return unique symbols and uniqute symbols that are not fixed
    return Tuple(unique_sym), Tuple(printed_sym[.!fixed])
end

############################################################################################
function update(datatune::DataTune, data)
    return DataTune(datatune.structure, ArrayConfig(data), datatune.miss)
end
function update(datatune::Vector{<:DataTune}, data)
    return DataTune(datatune[begin].structure, ArrayConfig(data), datatune[begin].miss)
end

############################################################################################
"""
$(SIGNATURES)
Infer types of all Diagnostics container.

# Examples
```julia
```

"""
function infer(
    _rng::Random.AbstractRNG,
    diagnostics::Type{AbstractDiagnostics},
    algorithmsᵛ::AbstractVector,
    model::ModelWrapper,
    data::D,
) where {D}
    diagtypes = map(
        algorithm -> infer(_rng, diagnostics, algorithm, model, data), algorithmsᵛ[begin]
    )
    return diagtypes
end

############################################################################################
"""
$(SIGNATURES)
Return buffer vector of type `diagtypes`. Comes from function `infer`.

# Examples
```julia
```

"""
function diagnosticsbuffer(
    diagtypes::D, iterations::Integer, Nchains::Integer, args...
) where {D}
    return [
        map(diag -> Vector{diag}(undef, iterations), diagtypes) for _ in Base.OneTo(Nchains)
    ]
end
function diagnosticsbuffer(
    diagtypes::D, iterations::Integer, Nchains::Integer, smc::SMC
) where {D}
    return Vector{diagtypes}(undef, iterations)
end

############################################################################################
#export
export SamplingInfo, update, infer


text:
#=
Copyright (c) 2015, Intel Corporation

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:

    * Redistributions of source code must retain the above copyright notice,
      this list of conditions and the following disclaimer.
    * Redistributions in binary form must reproduce the above copyright
      notice, this list of conditions and the following disclaimer in the
      documentation and/or other materials provided with the distribution.
    * Neither the name of Intel Corporation nor the names of its contributors
      may be used to endorse or promote products derived from this software
      without specific prior written permission.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE
FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
=#

include("dss.jl")

function speye_int32(m::Integer)
  rowval = [Int32(x) for x in [1:m;]]
  colptr = [Int32(x) for x in [rowval; m + 1]]
  nzval = ones(Float64, m)
  return SparseMatrixCSC(m, m, colptr, rowval, nzval)
end

function canonicalize_ipm_input(A, b, p, lo, hi)
  m = size(A, 1)
  n = size(A, 2)
  A = full(A)

  nlo = 0
  for i = 1:length(lo)
    if lo[i] > 0
      nlo += 1
    end
  end
  nhi = 0
  for i = 1:length(hi)
    if hi[i] < 1e300
      nhi += 1
    end
  end

  A = [A; zeros(nlo + nhi, n)]
  b = [b; zeros(nlo + nhi)]

  nlo = 0
  for i = 1:length(lo)
    if lo[i] > 0
      nlo += 1
      A[m + nlo, i] = 1
      b[m + nlo] = lo[i]
    end
  end

  nhi = 0
  for i = 1:length(hi)
    if hi[i] < 1e300
      nhi += 1
      A[m + nlo + nhi, i] = 1
      b[m + nlo + nhi] = hi[i]
    end
  end

  new_n = length(b)
  A = [A [zeros(m, nlo + nhi); -eye(nlo) zeros(nlo, nhi); zeros(nhi, nlo) eye(nhi)]]
  A = SparseMatrixCSC{Float64, Int32}(sparse(A))
  p = [p; zeros(nlo + nhi)]

  A, b, p
end

function load_ipm_input(name)
  A = full(matrix_market_read(string(name, ".mtx")))'
  b = matrix_market_read(string(name, "_b.mtx"))
  p = matrix_market_read(string(name, "_c.mtx"))

  n = string(name, "_lo.mtx")
  if isfile(n)
    lo = matrix_market_read(string(name, "_lo.mtx"))
  else
    lo = zeros(0)
  end
  n = string(name, "_hi.mtx")
  if isfile(n)
    hi = matrix_market_read(string(name, "_hi.mtx"))
  else
    hi = zeros(0)
  end

  canonicalize_ipm_input(A, b, p, lo, hi)
end

type Factor <: Factorization{Float64}
  p::Vector{Int}
  function Factor(p::Vector{Int})
    new(p)
  end
end

function cholfact_int32(B :: SparseMatrixCSC{Float64, Int32})
    R = Factor(dss_analyze(B))
    dss_factor(R.p, B)
    R
end

import Base.LinAlg: (\)

function (\)(L::Factor, b::Vector)
    y = zeros(b)
    dss_solve!(L.p, b, y)
    y
end

# Compared to the Jongsoo's original ipm_ref(), the following changes are made
# to make it work around OptFramework issues:
#(1) No default parameter
#(2) No print
#(3) cholfact_int32(B) instead of cholfact(SparseMatrixCSC{Float64, Int64}(B))

# primal-dual interior-point method for problem
# min p'*x s.t. Ax = b, x >= 0

function ipm_ref(A, b, p) # A: constraint coefficients, b: constraint rhs, p: objective
  t0 = time()
  (m,n) = size(A)

# set initial point, based on largest element in (A,b,p)
  bigM = maximum(A)
  bigM = maximum([norm(b, Inf) norm(p, Inf) bigM])

  # Do not know why, Julia cannot automatically figure out the types of x, s and y,
  # for which we miss some patterns. Here manually annotate the types.
  x::Vector{Float64} = 100*bigM*ones(n)
  s::Vector{Float64} = x
  y::Vector{Float64} = zeros(m)

  bc = 1 + maximum([norm(b) norm(p)])

  relResidual = NaN
  iter=1

  blas1_time = 0.
  spgemm_time = 0.
  fact_time = 0.
  trslv_time = 0.
  spmv_time = 0.

  D = speye_int32(n)
  Rp = zeros(size(A, 1))

  # We need Rc to live into the loop so that we can create a temporary based on it.
  # TODO: remove this statement after a solution is found.
  Rc = zeros(x)

  for iter=1:200
    set_matrix_property(Dict(
        :Rc  => SA_HAS_DEDICATED_MEMORY,
      )
    )

    # compute residuals
    spmv_time -= time()
    Rd = A'*y + s - p
    Rp = A*x - b
    spmv_time += time()

    blas1_time -= time()
    Rc = x.*s
    mu = mean(Rc)
    relResidual = norm([Rd; Rp; Rc])/bc
    blas1_time += time()

    if (relResidual <= 1e-7 && mu <= 1e-7) break; end

    blas1_time -= time()
    Rc = Rc - min(0.1, 100*mu)*mu

    # set up the scaling matrix, and form the coefficient matrix for
    # the linear system
    d = min(5.e+15, x./s)
    blas1_time += time()

    spgemm_time -= time()
    D.nzval = d
    B = A*D*A'
    spgemm_time += time()

    # use the form of the Cholesky routine "cholinc" that's best
    # suited to interior-point methods
    fact_time -= time()
    R = cholfact_int32(B)
    fact_time += time()

    # set up the right-hand side
    blas1_time -= time()
    t1 = x.*Rd - Rc;
    blas1_time += time()

    spmv_time -= time()
    t2 = -(Rp + A*(t1./s));
    spmv_time += time()

    # solve it and recover the other step components
    trslv_time -= time()
    dy = R\t2
    trslv_time += time()

    spmv_time -= time()
    temp = A'*dy
    spmv_time += time()

    blas1_time -= time()
    dx = (x.*temp + t1)./s
    ds = -(s.*dx + Rc)./x

    tau = max(.9995, 1 - mu)
    ap = -1/minimum([dx./x; -1])
    ad = -1/minimum([ds./s; -1])
    ap = tau*ap
    ad = tau*ad
    x = x + ap*dx
    s = s + ad*ds
    y = y + ad*dy
    blas1_time += time()
  end

  blas1_time -= time()
  f = p'*x
  blas1_time += time()

  x, time() - t0, spgemm_time, fact_time, blas1_time, trslv_time, spmv_time, iter, relResidual, f[1][1]
end

max_stars_repo_name: IntelLabs/Sparso

max_issues_repo_name: IntelLabs/Sparso

max_forks_repo_name: IntelLabs/Sparso


text:
using SMC
using PyPlot
include("../aide.jl")
include("smc_samplers.jl")
include("hmm_serialization.jl")

function add_axis_labels()
    fontsize=14
    ax = plt[:gca]()
    ax[:set_ylabel]("state", labelpad=-15, fontsize=fontsize)
    ax[:set_xlabel]("time", labelpad=-15, fontsize=fontsize)
end

function format_axes()
    add_axis_labels()
    plt[:tight_layout](pad=0)
end

function add_upper_right_label(label::String)
    plt[:text](0.86, 0.69, label, fontsize=34, transform=plt[:gcf]()[:transFigure])
end

function render_marginals(hmm::HiddenMarkovModel, observations::Vector{Int},
                          sampler::Sampler, num_replicates::Int, label="")
    num_steps = length(observations)
    marginals = zeros(hmm.num_states, num_steps)
    for i=1:num_replicates
        states, _ = simulate(sampler)
        for (t, state) in enumerate(states)
            marginals[state, t] += 1.0
        end
    end
    marginals = marginals / num_replicates
    render_hmm_posterior_marginals!(hmm, marginals, false, [1, 50], [1, 20])
    add_upper_right_label(label)
    format_axes()
end

# where to read data from
DATA_DIR = "data"

# where to write plots to
PLOT_DIR = "plots"
if !Base.Filesystem.isdir(PLOT_DIR)
    Base.Filesystem.mkdir(PLOT_DIR)
end

# load HMM, ground truth states, and observations
hmm = load_hmm("$DATA_DIR/hmm.json")
num_states = num_states(hmm)
num_obs = num_observations(hmm)
ground_truth_states = load_states("$DATA_DIR/states.json")
observations = load_observations("$DATA_DIR/observations.json")
num_steps = length(ground_truth_states)
if length(observations) != num_steps
    error("Dimension mismatch")
end

# fix random seed for Monte Carlo estimates of marginals
srand(1)

const NUM_MARGINAL_REPLICATES = 1000

# plot the ground truth states
println("Plotting ground truth states...")
width = 3
height = 1.1*width*num_states/num_steps
plt[:figure](figsize=(width,height))
render_hmm_states!(hmm, ground_truth_states, false, [1, 50], [1, 20])
format_axes()
plt[:savefig]("$PLOT_DIR/ground_truth_states.pdf")

# exact posterior marginals
println("Plotting exact posterior marginals...")
exact_sampler = HMMExactSampler(hmm, observations)
plt[:figure](figsize=(width,height))
render_marginals(hmm, observations, exact_sampler, NUM_MARGINAL_REPLICATES)
plt[:savefig]("$PLOT_DIR/exact_posterior_marginals.pdf")

# prior proposal, 1 particle
println("Plotting prior proposal, 1 particle, marginals...")
sampler = make_prior_proposal_smc_sampler(hmm, observations, 1)
plt[:figure](figsize=(width,height))
render_marginals(hmm, observations, sampler, NUM_MARGINAL_REPLICATES, "A")
plt[:savefig]("$PLOT_DIR/smc_prior_1p_marginals.pdf")

# prior proposal, 10 particles
println("Plotting prior proposal, 10 particles, marginals...")
sampler = make_prior_proposal_smc_sampler(hmm, observations, 10)
plt[:figure](figsize=(width,height))
render_marginals(hmm, observations, sampler, NUM_MARGINAL_REPLICATES, "B")
plt[:savefig]("$PLOT_DIR/smc_prior_10p_marginals.pdf")

# optimal proposal, 100 particles
println("Plotting optimal proposal, 100 particles, marginals...")
sampler = make_optimal_proposal_smc_sampler(hmm, observations, 100)
plt[:figure](figsize=(width,height))
render_marginals(hmm, observations, sampler, NUM_MARGINAL_REPLICATES, "C")
plt[:savefig]("$PLOT_DIR/smc_cond_100p_marginals.pdf")

# optimal proposal, 1000 particles
println("Plotting optimal proposal, 1000 particles, marginals...")
sampler = make_optimal_proposal_smc_sampler(hmm, observations, 1000)
plt[:figure](figsize=(width,height))
render_marginals(hmm, observations, sampler, NUM_MARGINAL_REPLICATES)
plt[:savefig]("$PLOT_DIR/smc_cond_1000p_marginals.pdf")

println("done!")


text:
using YaoBlocks, Test
using YaoBlocks.Optimise
using YaoArrayRegister
using YaoBlocks: check_dumpload

block_A(i, j) = control(i, j => shift(2π / (1 << (i - j + 1))))
block_B(n, i) = chain(n, i == j ? put(i => H) : block_A(j, i) for j = i:n)
qft(n) = chain(block_B(n, i) for i = 1:n)

@testset "is pauli" begin
    @test is_pauli(Val(-im)*X)
    @test is_pauli(X, Z, -Z)
    @test !is_pauli(X, Z, -2Z)
end

@testset "map address" begin
    # chain, put, concentrator
    c2 = map_address(
        chain(5, subroutine(5, put(2, 2 => X), (4, 1)), put(5, 3 => X)),
        AddressInfo(10, [2, 1, 4, 6, 3]),
    )
    @test c2 == chain(10, subroutine(10, put(2, 2 => X), (6, 2)), put(10, 4 => X))

    # control, kron, rot
    c3 = map_address(
        chain(5, control(5, 2, 1 => Rx(0.3)), kron(5, 2 => Y, 3 => X)),
        AddressInfo(10, [2, 1, 4, 6, 3]),
    )
    @test c3 == chain(10, control(10, 1, 2 => Rx(0.3)), kron(10, 1 => Y, 4 => X))

    # repeat, measure
    c4 = map_address(
        chain(5, repeat(5, Y, (2, 3)), Measure(5), Measure(5; locs = (3, 2, 1))),
        AddressInfo(10, [2, 1, 4, 6, 3]),
    )
    @test c4 == chain(
        10,
        repeat(10, Y, (1, 4)),
        Measure(10; locs = (2, 1, 4, 6, 3)),
        Measure(10; locs = (4, 1, 2)),
    )

    # sum, cache, scale
    c5 = map_address(
        2 * put(5, 2 => X) + Daggered(Scale(Val(2), put(5, 3 => X) |> cache)),
        AddressInfo(10, [2, 1, 4, 6, 3]),
    )
    @test c5 == 2 * put(10, 1 => X) + Daggered(Scale(Val(2), put(10, 4 => X) |> cache))

    # qft
    c = qft(4)
    @test mat(subroutine(10, c, (6, 2, 3, 7))) ≈
          mat(map_address(c, AddressInfo(10, [6, 2, 3, 7])))
end

@testset "merge pauli *" begin
    test_merge(a, b) = mat(merge_pauli(a, b)) ≈ (mat(a) * mat(b))

    for a in [X, Y, Z], b in [X, Y, Z]
        @test test_merge(a, b)
    end

    @test mat(simplify(X * Y)) ≈ mat(im * Z)
    @test mat(simplify(im * X * im * Y)) ≈ mat(Val(-im) * Z)
    @test mat(simplify(I2 * im * I2)) ≈ mat(im * I2)

    @test mat(simplify(X * Y * Y)) ≈ mat(X) ≈ mat(X) * mat(Y) * mat(Y)
    @test mat(simplify(X * Y * Z)) ≈ mat(im * I2) ≈ mat(X) * mat(Y) * mat(Z)
end

@testset "eliminate nested" begin
    @test simplify(*(X, *(H))) == *(X, H)
    @test simplify(*(X)) == X

    @test simplify(Add(X, Add(X, X))) == 3X
end

@testset "reduce matrices" begin
    @test mat(*(X, Y)) ≈ mat(X) * mat(Y)
    @test mat(*(X, Y)) ≈ mat(simplify(*(X, Y)))
    @test mat(Add(X, Y)) ≈ mat(X) + mat(Y)
end

@testset "composite strcuture" begin
    g = chain(2, kron(1 => chain(X, Y), 2 => X), control(1, 2 => X))
    @test simplify(g) == prod([control(2, 1, 2 => X), kron(2, 1 => (Val(-im) * Z), 2 => X)])
end

@testset "to basic types" begin
    # primitive
    @test to_basictypes(X) == X

    # chain, put, concentrator
    @test to_basictypes(chain(5, put(5, 3 => X))) == chain(5, put(5, 3 => X))
    @test to_basictypes(subroutine(5, put(2, 2 => X), (4, 1))) == put(5, 1 => X)
    @test to_basictypes(subroutine(5, Measure(2, locs = (2,)), (4, 1))) ==
          Measure(5; locs = (1,))
    @test to_basictypes(subroutine(5, Measure(2), (4, 1))) == Measure(5; locs = (4, 1))
    @test to_basictypes(subroutine(5, X, (1,))) == put(5, 1 => X)
    @test to_basictypes(put(5, 3 => X)) == put(5, 3 => X)

    # control, kron, rot
    @test to_basictypes(control(5, 2, 1 => Rx(0.3))) == control(5, 2, 1 => Rx(0.3))
    @test to_basictypes(kron(5, 2 => Y, 3 => X)) == chain(put(5, 2 => Y), put(5, 3 => X))

    # repeat, measure
    @test to_basictypes(repeat(5, Y, (2, 3))) == chain(put(5, 2 => Y), put(5, 3 => Y))
    @test to_basictypes(Measure(5)) == Measure(5)
    @test to_basictypes(Measure(5; locs = (3, 2, 1))) == Measure(5; locs = (3, 2, 1))
    @test to_basictypes(Measure(5; operator = repeat(3, X, 1:3), locs = (3, 2, 1))) ==
          Measure(5; operator = repeat(3, X, 1:3), locs = (3, 2, 1))

    # sum, cache, scale
    @test to_basictypes(2 * put(5, 2 => X)) == 2 * put(5, 2 => X)
    @test to_basictypes(put(5, 3 => X) |> cache) == put(5, 3 => X)
    @test to_basictypes(Daggered(X)) == Daggered(X)

    # tobasic_recursive
    sub = chain(
        5,
        Daggered(put(5, 4 => Rx(0.5)) * 2),
        put(5, 3 => ConstGate.P0 + ConstGate.P1),
        kron(5, 3 => X, 4 => Y),
        subroutine(5, kron(X, Z), (3, 2)),
        Measure(5, operator = X, locs = 1),
    )
    c = chain(10, repeat(10, H, 1:10), subroutine(10, sub, 6:10))

    sub2 = chain(
        10,
        Daggered(put(10, 9 => Rx(0.5)) * 2),
        put(10, 8 => ConstGate.P0 + ConstGate.P1),
        chain(put(10, 8 => X), put(10, 9 => Y)),
        chain(put(10, 8 => X), put(10, 7 => Z)),
        Measure(10, operator = X, locs = 6),
    )
    c2 = chain(10, chain(10, [put(10, i => H) for i = 1:10]), sub2)
    @test simplify(c, rules = [to_basictypes]) == c2

    # subroutine, chain, put, kron and control
    c = chain(
        6,
        [
            put(6, 3 => X),
            kron(6, 2 => X, 4 => X),
            chain(6, [subroutine(6, control(2, 1, 2 => Y), (6, 1))]),
        ],
    )
    @test zero_state(6) |> c ≈ zero_state(6) |> simplify(c, rules = [to_basictypes])
end

@testset "replace block" begin
    @test eliminate_nested(chain(7, chain(7, control(7, 1, 2 => X)), put(7, 4 => X))) ==
          chain(7, [control(7, 1, 2 => X), put(7, 4 => X)])
    @test replace_block(X => Y, chain(put(2, 2 => X), put(2, 1 => Z), kron(X, Y))) ==
          chain(put(2, 2 => Y), put(2, 1 => Z), kron(Y, Y))
end

include("dumpload.jl")
check_dumpload(qft(5))


text:
## -------------
##

# Julia code to compare the probability bounds
using Distributions, Plots

ϵ = 0.1
σ   = 1

p_exact(N) = 1 - cdf(Normal(), sqrt(N) * ϵ / σ)
p_cheby(N) = σ^2 / (ϵ^2 * N);
p_chern(N) = exp(-ϵ^2 * N / (2*σ^2));

ex_args    = (linewidth=2, color=RGB(1.0, 0.5, 0.0), label="Exact", markershape=:circle)
chb_args   = (linewidth=2, color=RGB(0.2, 0.7, 0.1), label="Chebyshev", linestyle=:dashdot)
chern_args = (linewidth=2, color=RGB(0.2, 0.0, 0.8), label="Chernoff")

Ns = [10^i for i in range(1, 3.8, length=50)]
scatter(Ns, p_exact.(Ns);  ex_args...,
        xaxis=:log10,
        yaxis=:log10, yticks=[1e-15, 1e-10, 1e-5, 1e0],
        legend=:bottomleft)
plot!(p_cheby; chb_args...)
plot!(p_chern; chern_args...)

##
## -------------
##

# Julia code to illustrate the weak law of large numbers
using Distributions, Plots

p = 0.5

# Nset = round.(Int, (10).^range(2,5,length=100))
# x = hcat((rand.(Binomial.(Nset, p), 1000)./Nset)...)
# y = x[1:10:end, :]'

Nset = round.((10).^range(2,5,length=100))
x = zeros(length(Nset), 1000)
for (i,N) in enumerate(Nset)
    x[i,:] = rand(Binomial(N, p), 1000) / N
end
y = x[:, 1:10:end]

scatter(Nset, y; xaxis=:log10, xticks=[1e2, 1e3, 1e4, 1e5],
        markershape=:x, color=:black,
        ylabel="sample average",
        legend=false)

plot!(Nset, p .+ 3*sqrt.(p*(1-p) ./ Nset); color=:red, linewidth=4)
plot!(Nset, p .- 3*sqrt.(p*(1-p) ./ Nset); color=:red, linewidth=4)

##
## -------------
##

# Julia: Plot the PDF of the sum of two Gaussians
using Distributions, Plots

n = 10000;

K = 2 # 6

X = DiscreteUniform(1, 6)
μ = mean(X)   # 3.5
σ = std(X)    # sqrt((6^2-1)/12)

Z = sum(rand(X, n) for _ in 1:K)

mk = K * μ
sk = sqrt(K) * σ

bin_range = (floor(mk - 3sk) - 1/2):(ceil(mk + 3sk) + 1/2)
plot_args = (normalize=true, color=RGB(0, 0.5, 0.8), linewidth=2, legend=false,
             bins=bin_range)

histogram(Z; plot_args..., title="⚁"^K)

##
## -------------
##

# Julia: Visualize convergence in distribution
using Distributions, Plots
N, p = 10, 0.5

B = Binomial(N,p)
Z = Normal(mean(B), std(B)) # N*p, sqrt(N*p*(1-p))

p_b(x) = pdf(B, x)
p_n(x) = pdf(Z, x)

c_b(x) = cdf(B, x)
c_n(x) = cdf(Z, x)

bin_args = (color = RGB(0, 0, 0), linewidth=2, label="Binomial")
norm_args = (color = RGB(0.8, 0, 0), linewidth=6, label="Normal")

ns = 0:N
p1 = plot(ns, p_b.(ns); bin_args..., seriestype=:sticks)
plot!(p_n, 0, N; norm_args...)

p2 = plot(c_b, 0, N; bin_args...)
plot!(c_n; norm_args...)

l = @layout [a b]
plot(p1, p2, layout=l)

##
## -------------
##

# Julia: Poisson to Gaussian: convergence in distribution
using Distributions, Plots

N = 4; # N = 10, N = 50;
λ = 1

P = Poisson(N*λ)
Z = Normal(mean(P), std(P))  # Nλ, √(Nλ)

p_b(x) = pdf(P, x) # evaluate on integers
p_n(x) = pdf(Z, x)

c_b(x) = cdf(P, x)
c_n(x) = cdf(Z, x)

bin_args = (linewidth=2, color=RGB(0.0, 0.0, 0.0), label="Poisson")
norm_args = (linewidth=6, colr=RGB(0.8, 0.0, 0.0), label="Gaussian")

ns = 0:2N
p1 = plot(ns, p_b.(ns); bin_args..., seriestype=:sticks, legend=:topright)
plot!(p_n, 0, 2N; norm_args...)

p2 = plot(c_b, 0, 2N; bin_args..., legend=:bottomright)
plot!(c_n; norm_args...)

l = @layout [a b]
plot(p1, p2, layout=l)

##
## -------------
##

# Julia: Visualize the Central Limit Theorem
using Distributions, Plots

N = 10
x = range(0, N, length=1001)
p = 0.5

B = Binomial(N,p)
Z = Normal(mean(B), std(B))  # N*p, sqrt(N*p*(1-p))

p_b(x) = pdf(B, x)
p_n(x) = pdf(Z, x)

c_b(x) = cdf(B, x)
c_n(x) = cdf(Z, x)

xs = range(5-2.5, 5 + 2.5, length=1001)

p1 = plot(p_n, 5-2.5, 5+2.5; fillrange=0 .* xs, color=RGB(0.6, 0.9, 1.0), legend=false,
     title="Binomial PDF")
plot!(0:N, p_b.(0:N); color=RGB(0,0,0), linewidth=2, seriestype=:sticks, layout=1)

p2 = plot(p_n, 5-2.5, 5+2.5, fillrange=0 .* xs, color=RGB(0.6, 0.9, 1.0), legend=false,
     title = "Gaussian PDF")
plot!(p_n, 0, N; color=RGB(0.8, 0, 0), linewidth=6, layout=2)

l = @layout [a b]
plot(p1, p2, layout=l)

##
## -------------
##

# Julia: Central Limit Theorem from moment generating functions
using Distributions, Plots

N = 2
p = 0.5;

B = Bernoulli(p)
μ, σ = mean(B), std(B)/sqrt(N) # for Xbar
Z = Normal(μ, σ)

m_Bernoulli(s) = mgf(B,s)    # (1 - p + p*exp(s))
m_B(s) = m_Bernoulli(s/N)^N  # (X₁ + ⋯ + Xₙ)/n
m_Z(s) = mgf(Z,s)            # exp(μ * s + σ^2*s^2/2)

bin_args = (linewidth=8, color=RGB(0.1, 0.6, 1),
            yaxis=:log, yticks = [10.0^i for i in -1:4],
            label="Binomial MGF")
norm_args = (linewidth=8, color=RGB(0, 0, 0), yaxis=:log, linestyle=:dot, label="Gaussian MGF")

plot(m_B, -10, 10; bin_args..., legend=:topleft)
plot!(m_Z; norm_args...)

##
## -------------
##

# Julia: Failure of Central Limit Theorem at tails
using Distributions, Plots

λ = 1;

function gamma_pdf(N)
    function(x) # return anonymous function; also x -> ... notation
        (sqrt(N)/λ) * pdf(Gamma(N,λ), (x+sqrt(N))/(λ/sqrt(N)))
    end
end
gaussian_pdf(x) = pdf(Normal(), x)

plot(gamma_pdf(1), -1, 5; linewidth=4, color=RGB(0.8, 0.8, 0.8), label="N=1",
     yaxis=:log,
     legend=:bottomleft)
plot!(gamma_pdf(10);   linewidth=4, color=RGB(0.6, 0.6, 0.6), label="N=10")
plot!(gamma_pdf(100);  linewidth=4, color=RGB(0.4, 0.4, 0.4), label="N=100")
plot!(gamma_pdf(1000); linewidth=4, color=RGB(0.2, 0.2, 0.2), label="N=1000")
plot!(gaussian_pdf;    linewidth=4, color=RGB(0.9, 0.0, 0.0), label="Gaussian",
      linestyle=:dash)


text:
abstract type AbstractAnalysis{T} end

struct Analysis{T, N<:NamedTuple} <: AbstractAnalysis{T}
    f::T
    kwargs::N
    function Analysis(f::T; kwargs...) where {T}
        nt = values(kwargs)
        new{T, typeof(nt)}(f, nt)
    end
end

(an::Analysis)(; kwargs...) = Analysis(an.f; kwargs..., an.kwargs...)
(an::Analysis)(args...; kwargs...) = apply_keywords(an.f, args...; kwargs..., an.kwargs...)

function convert_arguments(P::PlotFunc, f::AbstractAnalysis, args...; kwargs...)
    tmp = f(args...; kwargs...) |> to_tuple
    convert_arguments(P, tmp...)
end

const FunctionOrAnalysis = Union{Function, AbstractAnalysis}

adjust_globally(s::FunctionOrAnalysis, traces) = s


text:
using Mimi
using Test
using DataFrames
using VegaLite
using Electron
using Distributions
using Query
using CSVFiles

import Mimi:
    _spec_for_sim_item, menu_item_list, getdataframe, get_sim_results

# Helper function returns true if VegaLite is verison 3 or above, and false otherwise
function _is_VegaLite_v3()
    return isdefined(VegaLite, :vlplot) ? true : false
end

# Get the example
include("mcs/test-model-2/multi-region-model.jl")
using .MyModel
m = construct_MyModel()

N = 100

sd = @defsim begin
    # Define random variables. The rv() is required to disambiguate an
    # RV definition name = Dist(args...) from application of a distribution
    # to an model parameter. This makes the (less common) naming of an
    # RV slightly more burdensome, but it's only required when defining
    # correlations or sharing an RV across parameters.
    rv(name1) = Normal(1, 0.2)
    rv(name2) = Uniform(0.75, 1.25)
    rv(name3) = LogNormal(20, 4)

    # assign RVs to model Parameters
    grosseconomy.share = Uniform(0.2, 0.8)
    emissions.sigma[:, Region1] *= name2
    emissions.sigma[2020:5:2050, (Region2, Region3)] *= Uniform(0.8, 1.2)

    grosseconomy.depk = [Region1 => Uniform(0.08, 0.14),
            Region2 => Uniform(0.10, 1.50),
            Region3 => Uniform(0.10, 0.20)]

    sampling(LHSData, corrlist=[(:name1, :name2, 0.7), (:name1, :name3, 0.5)])

    # indicate which parameters to save for each model run. Specify
    # a parameter name or [later] some slice of its data, similar to the
    # assignment of RVs, above.
    save(grosseconomy.K, grosseconomy.YGROSS, grosseconomy.share_var, grosseconomy.depk_var,
         emissions.E, emissions.E_Global)
end

si = run(sd, m, N)
results_output_dir = tempdir()
si_disk = run(sd, m, N; results_output_dir = results_output_dir, results_in_memory = false)

## 1. Specs and Menu
pairs = [(:grosseconomy, :K), (:grosseconomy, :YGROSS), (:grosseconomy, :share_var),
         (:grosseconomy, :depk_var), (:emissions, :E), (:emissions, :E_Global)]
for (comp, var) in pairs
    results = get_sim_results(si, comp, var)

    static_spec = _spec_for_sim_item(si, comp, var, results; interactive = false)
    interactive_spec = _spec_for_sim_item(si, comp, var, results)

    name = string(comp, " : ", var)
    @test static_spec["name"] == interactive_spec["name"] == name
end

## 2. Explore
w = explore(si, title="Testing Window")
@test typeof(w) == Electron.Window
close(w)

w = explore(si_disk, results_output_dir = results_output_dir)
@test typeof(w) == Electron.Window
close(w)

@test_throws ErrorException explore(si_disk) #should error, no in-memory results

## 3. Plots

function plot_type_test(p)
    if _is_VegaLite_v3()
        @test typeof(p) == VegaLite.VLSpec
    else
        @test typeof(p) == VegaLite.VLSpec{:plot}
    end
end

# trumpet plot
p = Mimi.plot(si, :emissions, :E_Global)
plot_type_test(p)
p = Mimi.plot(si, :emissions, :E_Global; interactive = true)
plot_type_test(p)

p = Mimi.plot(si_disk, :emissions, :E_Global, results_output_dir = results_output_dir)
plot_type_test(p)
p = Mimi.plot(si_disk, :emissions, :E_Global; interactive = true, results_output_dir = results_output_dir)
plot_type_test(p)

@test_throws ErrorException Mimi.plot(si_disk, :emissions, :E_Global) #should error, no in-memory results

# mulitrumpet plot
p = Mimi.plot(si, :emissions, :E)
plot_type_test(p)
p = Mimi.plot(si, :emissions, :E; interactive = true);
plot_type_test(p)

p = Mimi.plot(si_disk, :emissions, :E, results_output_dir = results_output_dir)
plot_type_test(p)
p = Mimi.plot(si_disk, :emissions, :E; interactive = true, results_output_dir = results_output_dir);
plot_type_test(p)

# histogram plot
p = Mimi.plot(si, :grosseconomy, :share_var)
plot_type_test(p)
p = Mimi.plot(si, :grosseconomy, :share_var; interactive = true); # currently just calls static version
plot_type_test(p)

p = Mimi.plot(si_disk, :grosseconomy, :share_var; results_output_dir = results_output_dir)
plot_type_test(p)
p = Mimi.plot(si_disk, :grosseconomy, :share_var; interactive = true, results_output_dir = results_output_dir); # currently just calls static version
plot_type_test(p)

# multihistogram plot
p = Mimi.plot(si, :grosseconomy, :depk_var)
plot_type_test(p)
p = Mimi.plot(si, :grosseconomy, :depk_var; interactive = true);
plot_type_test(p)

p = Mimi.plot(si_disk, :grosseconomy, :depk_var; results_output_dir = results_output_dir)
plot_type_test(p)
p = Mimi.plot(si_disk, :grosseconomy, :depk_var; interactive = true, results_output_dir = results_output_dir);
plot_type_test(p)


text:
## Example of markup which uses Markdown for text formatting and LaTeX for math formatting

using JuliaWeBWorK

## jmt is used so there is no need to escape backslashes
intro = jmt"""

# LaTeX markup

LaTeX markup can be done two ways (Julian or using LaTeX's slash() or slash[])

Inline: ``\LaTeX`` or \(\LaTeX\)

Display math can be done in two ways

```math
\sin(x)^2 = \frac{1}{2}
```

\[
\sin(x)^2 = \frac{1}{2}
\]

The use of dollar signs for math markup is not supported -- rather dollar signs in the `ww` string macro are for $("interpolation").

# Markdown markup

String formatting: **bold**, *italics*, `code`

itemized lists:

* one
* two

1) One
2) Two

horizontal line

-----

verbatim text

```
verbatim text
```

Quotation

> Four score and seven years ago our fathers brought forth on this continent,
> a new nation, conceived in Liberty, and dedicated to the proposition that
> all men are created equal.

Admonition

!!! note "Admonition title"
    this is an admonition

URL:

<https://www.markdownguide.org/basic-syntax/>

[markdown guide](https://www.markdownguide.org/basic-syntax/)

Images

![some image](https://upload.wikimedia.org/wikipedia/commons/thumb/0/0f/Tangent_to_a_curve.svg/200px-Tangent_to_a_curve.svg.png)

# quirks

* The at sign of a macro is **special** and can't be printed without hacking it in directly. (@ is okay but not @ followed by a name

"""

Page(intro, ())


text:
"""
Classical test functions for sparse polynomial optimization. See e.g. [WSMM06].

[WSMM06] Waki, Hayato, Sunyoung Kim, Masakazu Kojima, and Masakazu Muramatsu.
Sums of squares and semidefinite program relaxations for polynomial optimization problems with structured sparsity.
SIAM Journal on Optimization 17, no. 1 (2006): 218-242.
"""
function chained_singular(n::Int)
    # clique size 3
    @assert mod(n, 4)==0
    @polyvar x[1:n]
    return sum([2*i-1 for i in 1:Int(n/2-1)]) do j
        (x[j] + 10*x[j+1])^2 + 5*(x[j+2] - x[j+3])^2 + (x[j+1] - 2*x[j+2])^4 + 10*(x[j] - x[j+3])^4
    end
end

function broyden_banded(n::Int)
    # clique size 7
    @polyvar x[1:n]
    return sum(1:n) do i
        ( x[i]*(2+5*x[i]^2) + 1 + (1 + x[i])*x[i] - sum( (1+x[j])*x[j] for j = maximum([1, i-5]):minimum([n, i+1]) ) )^2
    end
end

function broyden_tridiagonal(n::Int)
    # clique size 3
    @polyvar x[1:n]
    return (( 3 - 2*x[1])*x[1] -2*x[2] +1 )^2 + sum( ((3 - 2*x[i])*x[i] - x[i-1] - 2*x[i+1] + 1)^2 for i = 2:n-1)
end

function chained_wood(n::Int)
    # clique size 2
    @assert mod(n, 4)==0
    @polyvar x[1:n]
    p = 1 + sum(100*(x[j+1] - x[j]^2)^2 + (1 - x[j])^2 + 90*(x[j+3] - x[j+2]^2)^2  + (1 - x[j+2])^2
            + 10*(x[j+1] + x[j+3] - 2)^2 + 0.1*(x[j+1] - x[j+3])^4 for j in [2*i-1 for i in 1:Int(n/2-1)])
    return p
end

function generalized_rosenbrock(n::Int)
    # clique size 2
    @polyvar x[1:n]
    p = 1 + sum( 100*(x[i]-x[i-1]^2)^2 + (1-x[i])^2 for i=2:n)
    return p
end

function sos_lower_bound(p, factory, sparsity::Sparsity)
    model = Model(factory)
    @variable(model, t)
    @objective(model, Max, t)
    @constraint(model, p - t in SOSCone(), sparsity=sparsity)
    optimize!(model)
    println(termination_status(model))
    println(objective_value(model))
    return model
end


text:
"""
    SafeStatesProperty{N<:Real} <: Property

Type that represents a safety property characterized by a set of safe states.
The property is satisfied by a given set of states ``X`` if ``X`` is fully
contained in the set of safe states.

### Fields

- `safe`    -- convex set representing the safe states
- `witness` -- witness point (empty vector if not set)

### Notes

The following formula characterizes whether a set ``X`` satisfies a safety
property characterized by a set of safe states 𝑃:

```math
    X \\models 𝑃 \\iff X ⊆ 𝑃.\\texttt{safe}
```
"""
mutable struct SafeStatesProperty{N<:Real} <: Property
    safe::LazySet
    witness::Vector{N}

    SafeStatesProperty{N}(safe::LazySet) where {N<:Real} = new(safe, N[])
end

# type-less convenience constructor
SafeStatesProperty(safe::LazySet{N}) where {N<:Real} =
    SafeStatesProperty{N}(safe)

"""
    dim(𝑃::SafeStatesProperty)::Int

Return the dimension of a property with safe states.

### Input

- `𝑃` -- safety property with safe states

### Output

The dimension of the safe states.
"""
function dim(𝑃::SafeStatesProperty)::Int
    return dim(𝑃.safe)
end

"""
    check(𝑃::SafeStatesProperty, X::LazySet; witness::Bool=false)

Checks whether a convex set is contained in the set of safe states.

### Input

- `𝑃`       -- safety property with safe states
- `X`       -- convex set
- `witness` -- (optional, default: `false`) flag for returning a counterexample
               if the property is violated

### Output

Let ``Y`` be the safe states represented by 𝑃.
* If `witness` option is deactivated: `true` iff ``X ⊆ Y``
* If `witness` option is activated:
  * `(true, [])` iff ``X ⊆ Y``
  * `(false, v)` iff ``X ⊈ Y`` and ``v ∈ X \\setminus Y``
"""
@inline function check(𝑃::SafeStatesProperty, X::LazySet; witness::Bool=false)
    return ⊆(X, 𝑃.safe, witness)
end


text:
using OrdinaryDiffEq
using ForwardDiff
using DataFrames
using LinearAlgebra
using StatsPlots

function sir_rates(u,p)
  (S,I,R) = u
  (β,c,γ) = p
  N = S+I+R
  infection = β*c*I/N*S
  recovery = γ*I
  [infection,recovery]
end;

sir_transitions = float.([[-1  1  0];
                          [ 0 -1  1]])

δt = 0.1
tmax = 40.0
tspan = (0.0,tmax)
t = 0.0:δt:tmax;

u0 = [990.0,10.0,0.0]; # S,I.R
c0 = zeros(3,3) # covariance matrix
uc0 = vcat(diagm(0=>u0), c0)

p = [0.05,10.0,0.25]; # β,c,γ

rates = sir_rates
transitions = sir_transitions
nrates, nstates = size(transitions);

function ode(du, u, p, t)
    du[1:3] = transitions' * rates(u, p)
end;

prob_ode = ODEProblem(ode, u0 ,tspan, p);

sol_ode = solve(prob_ode, Tsit5(), saveat=t);

df_ode = DataFrame(sol_ode(t)')
df_ode[!,:t] = t;

@df df_ode plot(:t,
    [:x1 :x2 :x3],
    label=["S" "I" "R"],
    xlabel="Time",
    ylabel="Number",
    c=[:blue :red :green])

function lna(du, u, p, t)
        mean_vec = diag(u)
        covar_mx = u[nstates+1:nstates*2,:]
        reaction_rates = rates(mean_vec, p)
        reaction_rates_jac = ForwardDiff.jacobian(
            y -> rates(y, p),
            mean_vec)
        A = transitions' * reaction_rates_jac
        du[1:nstates, :] .= diagm(0 => transitions'*reaction_rates)
        du[nstates + 1:end, :] .= A*covar_mx + covar_mx*A' + transitions' * diagm(0 => reaction_rates) * transitions
end;

prob_lna = ODEProblem(lna, uc0, tspan, p);

sol_lna = solve(prob_lna, Tsit5(),saveat=t);

mean_traj = Array{Float64,2}(undef, nstates, length(t))
covar_traj = Array{Array{Float64,2},1}(undef, length(t))
for j in 1:length(t)
    mean_traj[:,j] = diag(sol_lna[j][1:nstates,1:nstates])
    covar_traj[j] = sol_lna[j][nstates+1:end, 1:nstates]
end;

var_traj = zeros(nstates, length(t))
for (idx, elt) in enumerate(covar_traj)
    var_traj[:, idx] = diag(elt)
end
sd_traj = 1.96 .* sqrt.(var_traj);

plot(t, mean_traj',
     ribbon=sd_traj',
     label = ["S" "I" "R"],
     xlabel = "Time",
     ylabel = "Number")


text:
@testset "RCF" begin
  include("RCF/rcf.jl")
  include("RCF/conductor_sieve.jl")
end


text:
using Test, QSimulator

@testset "basis state utilities" begin
    # check that constructors from AbstractVector work
    basis = TensorProductBasis((2,2))
    @test basis == TensorProductBasis([2,2])
    @test TensorProductBasisState(basis, (0,0)) == TensorProductBasisState(basis, [0,0])

    # simple tensor product expansion
    basis = TensorProductBasis((2,2))
    @test [bs.states for bs in basis_states(basis)] == [(0,0), (0,1), (1,0), (1,1)]
    basis = TensorProductBasis((2,3))
    @test [bs.states for bs in basis_states(basis)] == [(0,0), (0,1), (0,2), (1,0), (1,1), (1,2)]

    # go from state -> index
    basis = TensorProductBasis((2,3,4,5))
    all_states = basis_states(basis)
    @test all(QSimulator.index(state) == i for (i, state) in enumerate(all_states))

    # and from index to state
    @test all(basis[i] == state for (i,state) in enumerate(all_states))

    # check that TensorProductBasisState construction catches errors
    basis = TensorProductBasis((2,2))
    # wrong number of subsystems
    @test_throws MethodError TensorProductBasisState(basis, (0,0,0))
    # state index exceeds dimension
    @test_throws AssertionError TensorProductBasisState(basis, (2,0))
end


text:
module TestCopy

using Test
using ThreadsX

@testset "copy vector" begin
    @testset for basesize in [[nothing]; 1:8], n in 0:10, T in [Int, Any]
        ys = Vector{T}(undef, n)
        if basesize === nothing
            @test ThreadsX.copyto!(ys, 1:n) == 1:n
            @test ThreadsX.copy!(ys, 1:n) == 1:n
        else
            @test ThreadsX.copyto!(ys, 1:n; basesize = basesize) == 1:n
            @test ThreadsX.copy!(ys, 1:n; basesize = basesize) == 1:n
        end
    end
end

@testset "copy matrix" begin
    @testset for basesize in [[nothing]; 1:8], n in 0:10, m in 1:5, T in [Int, Any]
        A = Matrix{T}(undef, n, m)
        B = reshape(1:n*m, n, m)
        if basesize === nothing
            @test ThreadsX.copyto!(A, B) == B
        else
            @test ThreadsX.copyto!(A, B; basesize = basesize) == B
        end

        Bt = transpose(reshape(1:n*m, m, n))
        if basesize === nothing
            @test ThreadsX.copyto!(A, Bt) == Bt
        else
            @test ThreadsX.copyto!(A, Bt; basesize = basesize) == Bt
        end
    end
end

end  # module


text:
# precise values were computed with Maple 2018.2.1
@testset "k = 1: Smallest parameter is equal to 0" begin
  a = Float64[0, 1, 10, 100, 1_000, 10_000]
  z = Float64[0.0001, 0.05, 0.25, 0.5, 0.75, 0.95, 0.999, 1.0, 1.001, 1.5, 2.0, 5.0, 10.0, 100.0]

  @testset for (v, i) = enumerate(a), (u, x) = enumerate(z)
    @test SimpleBirthDeathProcess.log_hypergeometric(i, zero(Float64), log(x)) === zero(Float64)
  end
end

@testset "k = 1: Smallest parameter is equal to 1" begin
  a = Float64[1, 10, 100, 1_000, 10_000]
  z = Float64[0.0001, 0.05, 0.25, 0.5, 0.75, 0.95, 0.999, 1.0, 1.001, 1.5, 2.0, 5.0, 10.0, 100.0]

  @testset for (v, i) = enumerate(a), (u, x) = enumerate(z)
    @test SimpleBirthDeathProcess.log_hypergeometric(i, one(Float64), log(x)) ≈ log1p(z[u])
  end
end

@testset "k = 1: Smallest parameter is equal to 10" begin
  a = Float64[10, 100, 1_000, 10_000]
  z = Float64[0.0001, 0.05, 0.25, 0.5, 0.75, 0.95, 0.999, 1.0, 1.001, 1.5, 2.0, 5.0, 10.0, 100.0]

  F = [.52629570750287642e-3 .91738879566125019e-3 .99103116479562278e-3 .99905090243401476e-3;
       .25827405402099644 .44843337932042183 .48365235659705397 .48747343730492292;
       1.2054650954697372 2.0636024443887799 2.2135642778561085 2.2296368083023555;
       2.2363597073435242 3.7717986038306407 4.0248301950016774 4.0516528801920376;
       3.1400087187564869 5.2295234752131476 5.5577840348697503 5.5923027193844829;
       3.7917650954983671 6.2594003839524659 6.6346492459499567 6.6739111333451279;
       3.9432482803681855 6.4962973784476965 6.8816948938946638 6.9219748306325743;
       3.9463086440638964 6.5010739252600879 6.8866736335117214 6.9269738299247465;
       3.9493677809408321 6.5058481865958856 6.8916498955268224 6.9319703314655938;
       5.3418154461353614 8.6418531256852886 9.1091011971693386 9.1575092622170040;
       6.5259297666347161 10.403699599162072 10.926302430593021 10.980124686221491;
       11.370989154827456 17.178317589397033 17.842706622218652 17.910095816706045;
       16.015972663553841 23.166677252626662 23.897200666625793 23.970771578398150;
       35.440236886017119 45.259632511457571 46.062102918942031 46.142294355447970]

  @testset for (v, i) = enumerate(a), (u, x) = enumerate(z)
    @test SimpleBirthDeathProcess.log_hypergeometric(i, Float64(10), log(x)) ≈ F[u, v]
  end
end

@testset "k = 1: Smallest parameter is equal to 100" begin
  a = Float64[100, 1_000, 10_000]
  z = Float64[0.0001, 0.05, 0.25, 0.5, 0.75, 0.95, 0.999, 1.0, 1.001, 1.5, 2.0, 5.0, 10.0, 100.0]

  F = [.50249368807809934e-2 .90987634287470340e-2 .99014802318167912e-2;
       2.4666443059832288 4.4482851939723456 4.8323139970139239;
       11.524281754588844 20.480803126601523 22.117912345460543;
       21.400865955167309 37.453351842420101 40.218677436589141;
       30.072315426599046 51.949247565387039 55.539683534816658;
       36.333830674236772 62.196534594341785 66.303069760338699;
       37.789959356670749 64.554315994835902 68.772395282427217;
       37.819380072632425 64.601858299203765 68.822160169402496;
       37.848789114135625 64.649377950963507 68.871900302643673;
       51.246835825536797 85.918710136477377 91.037419537357267;
       62.656980998947858 103.47455944362613 109.20340200218063;
       109.46579494823216 171.05660004460201 178.35230210076593;
       154.46757511501702 230.86168145561156 238.89032794264213;
       343.01569283539162 451.70431712929944 460.53193435379431]

  @testset for (v, i) = enumerate(a), (u, x) = enumerate(z)
    @test SimpleBirthDeathProcess.log_hypergeometric(i, Float64(100), log(x)) ≈ F[u, v]
  end
end

@testset "k = 1: Smallest parameter is equal to 1000" begin
  a = Float64[1_000, 10_000]
  z = Float64[0.0001, 0.05, 0.25, 0.5, 0.75, 0.95, 0.999, 1.0, 1.001, 1.5, 2.0, 5.0, 10.0, 100.0]

  F = [.50023136359714628e-1 .90913185893301719e-1;
       24.556062009329795 44.447116566494245;
       114.73804119098180 204.65399437047979;
       213.09154263526358 374.27063596879579;
       299.45745843985605 519.14850966958181;
       361.82793031880269 621.56999830698620;
       376.33310417901014 645.13663820919853;
       376.62618055924502 645.61183835006045;
       376.91914076041489 646.08681216570142;
       510.39623395340794 858.68943294195577;
       624.08549449957836 1034.1852145148492;
       1090.6064193359382 1709.8407880825313;
       1539.2529853619575 2307.8125159100825;
       3419.6984988733087 4516.1511309139319]

  @testset for (v, i) = enumerate(a), (u, x) = enumerate(z)
    @test SimpleBirthDeathProcess.log_hypergeometric(i, Float64(1_000), log(x)) ≈ F[u, v]
  end
end

@testset "k in [1, 0, -1]: Smallest parameter is equal to 10" begin
  a = Float64[10, 100, 1_000, 10_000]
  z = Float64[0.0001, 0.05, 0.25, 0.5, 0.75, 0.95, 0.999, 1.0, 1.001, 1.5, 2.0, 5.0, 10.0, 100.0]

  F1 = [.52629570750287642e-3 .91738879566125019e-3 .99103116479562278e-3 .99905090243401476e-3;
        .25827405402099644 .44843337932042183 .48365235659705397 .48747343730492292;
        1.2054650954697372 2.0636024443887799 2.2135642778561085 2.2296368083023555;
        2.2363597073435242 3.7717986038306407 4.0248301950016774 4.0516528801920376;
        3.1400087187564869 5.2295234752131476 5.5577840348697503 5.5923027193844829;
        3.7917650954983671 6.2594003839524659 6.6346492459499567 6.6739111333451279;
        3.9432482803681855 6.4962973784476965 6.8816948938946638 6.9219748306325743;
        3.9463086440638964 6.5010739252600879 6.8866736335117214 6.9269738299247465;
        3.9493677809408321 6.5058481865958856 6.8916498955268224 6.9319703314655938;
        5.3418154461353614 8.6418531256852886 9.1091011971693386 9.1575092622170040;
        6.5259297666347161 10.403699599162072 10.926302430593021 10.980124686221491;
        11.370989154827456 17.178317589397033 17.842706622218652 17.910095816706045;
        16.015972663553841 23.166677252626662 23.897200666625793 23.970771578398150;
        35.440236886017119 45.259632511457571 46.062102918942031 46.142294355447970]

  F2 = [.44998345678670619e-3 .82496193284742848e-3 .89192008545566585e-3 .89914573134024811e-3;
        .22097348997472164 .40326729413131295 .43528341511859189 .43872605598148542;
        1.0337869442705281 1.8559564900380504 1.9921934559411943 2.0066729818594451;
        1.9223436082513246 3.3925865169201566 3.6223249060077170 3.6464873674046837;
        2.7041931940186855 4.7040669720766040 5.0019785898363071 5.0330721749260146;
        3.2696518295546249 5.6306933961520946 5.9711546994858141 6.0065197217477258;
        3.4012513696979216 5.8438472067500863 6.1934952647312047 6.2297770441467050;
        3.4039106862133272 5.8481450451956800 6.1979761202089027 6.2342761434082884;
        3.4065689624653464 5.8524408281779971 6.2024547458599853 6.2387729946938314;
        4.6190964452624217 7.7744591864540570 8.1981572586687954 8.2417579959973827;
        5.6538481761652827 9.3599027325441433 9.8336363685542776 9.8821118577812427;
        9.9150684195832806 15.456667865588829 16.058396715552669 16.119085841417183;
        14.030398994378236 20.846094369967696 21.507440548989568 21.573694019021120;
        31.404933440248328 40.729711990331924 41.455852237488533 41.528064515050814]

  F3 = [.37645737093725218e-3 .73267652494216932e-3 .79281076158312994e-3 .79924057819886665e-3;
        .18501505641666830 .35816765997400180 .38691529076290343 .38997868300551447;
        .86797500773573354 1.6485764401614848 1.7708258096930680 1.7837091877515229;
        1.6185025569124830 3.0137963828233681 3.2198245324530256 3.2413219045320742;
        2.2818894473571170 4.1791318554875865 4.4461791154097941 4.4738416909885490;
        2.7632381871697886 5.0025634774356329 5.3076666945255646 5.3391283763894788;
        2.8754387937166076 5.1919854674496699 5.5053022916911211 5.5375793250461307;
        2.8777067671217975 5.1958048215043223 5.5092852652820301 5.5415785242996434;
        2.8799738794608238 5.1996223498017387 5.5132662568183175 5.5455757253523662;
        3.9166606851026718 6.9077365078259501 7.2872207908624379 7.3260068052879610;
        4.8050596781442441 8.3168240164637134 8.7409782203730292 8.7840991092543809;
        8.4916436064234296 13.735821088365889 14.274095482354261 14.328075953551293;
        12.083136743044405 18.526335751794288 19.117689283855788 19.176616548826644;
        27.419142471984311 36.200625034610757 36.849610483741025 36.913834764572873]

  @testset for (v, i) = enumerate(a), (u, x) = enumerate(z)
    r1, r2, r3 = SimpleBirthDeathProcess.log_hypergeometric_joint(Float64(10), Float64(2), log(x))
    @test r1 ≈ log1p(z[u]) - log(11) + log(11 + 9 * z[u])
    @test r2 ≈ log1p(9 * z[u] / 10)
    @test r3 === zero(Float64)

    r1, r2, r3 = SimpleBirthDeathProcess.log_hypergeometric_joint(i, Float64(10), log(x))
    @test r1 ≈ F1[u, v]
    @test r2 ≈ F2[u, v]
    @test r3 ≈ F3[u, v]
  end
end

@testset "k in [1, 0, -1]: Smallest parameter is equal to 100" begin
  a = Float64[100, 1_000, 10_000]
  z = Float64[0.0001, 0.05, 0.25, 0.5, 0.75, 0.95, 0.999, 1.0, 1.001, 1.5, 2.0, 5.0, 10.0, 100.0]

  F1 = [.50249368807809934e-2 .90987634287470340e-2 .99014802318167912e-2;
        2.4666443059832288 4.4482851939723456 4.8323139970139239;
        11.524281754588844 20.480803126601523 22.117912345460543;
        21.400865955167309 37.453351842420101 40.218677436589141;
        30.072315426599046 51.949247565387039 55.539683534816658;
        36.333830674236772 62.196534594341785 66.303069760338699;
        37.789959356670749 64.554315994835902 68.772395282427217;
        37.819380072632425 64.601858299203765 68.822160169402496;
        37.848789114135625 64.649377950963507 68.871900302643673;
        51.246835825536797 85.918710136477377 91.037419537357267;
        62.656980998947858 103.47455944362613 109.20340200218063;
        109.46579494823216 171.05660004460201 178.35230210076593;
        154.46757511501702 230.86168145561156 238.89032794264213;
        343.01569283539162 451.70431712929944 460.53193435379431]

  F2 = [.49498146993667119e-2 .90069636852402412e-2 .98024558201806584e-2;
        2.4299167476990252 4.4034209515414675 4.7839863843406208;
        11.355103716989390 20.274470401927723 21.896715841394308;
        21.091177618372048 37.076400822137734 39.816463762401394;
        29.642236173597454 51.426769581826806 54.984254027376182;
        35.818373131251335 61.571266000234244 65.640003267968256;
        37.254821359540792 63.905404776958647 68.084634907931466;
        37.283845094226251 63.952470380968845 68.133902133716824;
        37.312857337844685 63.999513560653828 68.183144853323734;
        50.532684991172307 85.055682936009055 90.127004465471585;
        61.794709682825008 102.43570683587193 108.11132468235296;
        108.02335422838683 169.34144506661048 176.56873162702488;
        152.49755416222608 228.54835469900929 236.50137623187276;
        338.99663037344746 447.18251128663564 455.92656616789762]

  F3 = [.48749450306762533e-2 .89151788823844733e-2 .97034316007880775e-2;
        2.3933109019165650 4.3585637366489039 4.7356588611679045;
        11.186458746974319 20.068165915296292 21.675519685361248;
        20.782415661414754 36.699494779300824 39.414250627167662;
        29.213388876149237 50.904347324146949 54.428825174788873;
        35.304348363141624 60.946059184724153 64.976937493197120;
        36.721160676052845 63.256556576223269 67.396875263636285;
        36.749788321316848 63.303145504300101 67.445644828479381;
        36.778404658235848 63.349712236301587 67.494390134699319;
        49.820388773179653 84.192727819818496 89.216590213380003;
        60.934570748315407 101.39693148129660 107.01924823109579;
        106.58389780095158 167.62637678300494 174.78516210553471;
        150.53105487361483 226.23511704031285 234.11242549309786;
        334.98210607276279 442.66079559878089 451.32119896229363]

  @testset for (v, i) = enumerate(a), (u, x) = enumerate(z)
    r1, r2, r3 = SimpleBirthDeathProcess.log_hypergeometric_joint(Float64(100), Float64(2), log(x))
    @test r1 ≈ log1p(z[u]) - log(101) + log(101 + 99 * z[u])
    @test r2 ≈ log1p(99 * z[u] / 100)
    @test r3 === zero(Float64)

    r1, r2, r3 = SimpleBirthDeathProcess.log_hypergeometric_joint(i, Float64(100), log(x))
    @test r1 ≈ F1[u, v]
    @test r2 ≈ F2[u, v]
    @test r3 ≈ F3[u, v]
  end
end

@testset "k in [1, 0, -1]: Smallest parameter is equal to 1000" begin
  a = Float64[1_000, 10_000]
  z = Float64[0.0001, 0.05, 0.25, 0.5, 0.75, 0.95, 0.999, 1.0, 1.001, 1.5, 2.0, 5.0, 10.0, 100.0]

  F1 = [.50023136359714628e-1 .90913185893301719e-1;
        24.556062009329795 44.447116566494245;
        114.73804119098180 204.65399437047979;
        213.09154263526358 374.27063596879579;
        299.45745843985605 519.14850966958181;
        361.82793031880269 621.56999830698620;
        376.33310417901014 645.13663820919853;
        376.62618055924502 645.61183835006045;
        376.91914076041489 646.08681216570142;
        510.39623395340794 858.68943294195577;
        624.08549449957836 1034.1852145148492;
        1090.6064193359382 1709.8407880825313;
        1539.2529853619575 2307.8125159100825;
        3419.6984988733087 4516.1511309139319]

  F2 = [.49948127291680893e-1 .90821447794235120e-1;
        24.519388970102761 44.402282005465535;
        114.56910189594814 204.44779094691963;
        212.78226912186370 373.89390785783417;
        299.02793075953787 518.62632562316103;
        361.31311418683165 620.94506926038766;
        375.79862751389278 644.48807652132610;
        376.09130731173913 644.96280016119640;
        376.38387111361191 645.43729770395815;
        509.68291314254013 857.82683757123930;
        623.22417730264136 1033.1468512938082;
        1089.1653124992746 1708.1262776258678;
        1537.2845364794923 2305.4999102042296;
        3415.6814416938878 4511.6301327159023]

  F3 = [.49873143246801418e-1 .90729711196776349e-1;
        24.482727992831556 44.357448150824494;
        114.40021543877559 204.24159036290622;
        212.47308744782300 373.51718427133454;
        298.59852522947545 518.10414718390371;
        360.79844012945262 620.32014643098337;
        375.26429734277602 643.83952117550954;
        375.55658064673090 644.31376831685064;
        375.84874813773730 644.78778958919225;
        508.96977627099518 856.96424945705929;
        622.36307161979581 1032.1084958509835;
        1087.7245017908888 1706.4117759012473;
        1535.3164371300011 2303.1873134735524;
        3411.6648351122439 4507.1091436000474]

  @testset for (v, i) = enumerate(a), (u, x) = enumerate(z)
    r1, r2, r3 = SimpleBirthDeathProcess.log_hypergeometric_joint(Float64(1_000), Float64(2), log(x))
    @test r1 ≈ log1p(z[u]) - log(1_001) + log(1_001 + 999 * z[u])
    @test r2 ≈ log1p(999 * z[u] / 1_000)
    @test r3 === zero(Float64)

    r1, r2, r3 = SimpleBirthDeathProcess.log_hypergeometric_joint(i, Float64(1_000), log(x))
    @test r1 ≈ F1[u, v]
    @test r2 ≈ F2[u, v]
    @test r3 ≈ F3[u, v]
  end
end