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

using MLStyle
using StaticArrays
using Reexport
import Tables

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

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

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

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

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

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

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

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

unquote(x::QuoteNode) = x.value

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

# - Tables.jl interface

Tables.istable(sat::StructACSetTable) = true

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

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

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

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

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

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

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

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

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

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

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

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

# Mapping

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

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

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

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

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

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

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

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

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

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

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

end


module Runner
    #
    # This model actually runs all the calculations over a collection of households and stores the results in dataframes.
    # Presently it also contains code to summarise the output dataframes. FIXME output needs to be in own module.
    #
    using Parameters: @with_kw
    using DataFrames: DataFrame, DataFrameRow, Not, select!
    using CSV

    using BudgetConstraints: BudgetConstraint

    using ScottishTaxBenefitModel

    using .Definitions
    using .Utils
    using .STBParameters
    using .STBIncomes
    using .STBOutput

    using .RunSettings:
        Settings

    using .BenefitGenerosity:
        adjust_disability_eligibility!,
        initialise

    using .ModelHousehold:
        Household,
        Person,
        get_benefit_units

    using .Results:
        BenefitUnitResult,
        HouseholdResult,
        IndividualResult

    using .FRSHouseholdGetter:
        get_household,
        initialise

    using .Uprating: load_prices

    using .SingleHouseholdCalculations: do_one_calc

    export
        do_one_run!,
        summarise_inc_frame

    # fixme move the output stuff (mostly) to an Output.jl module

    function do_one_run!(
        settings :: Settings,
        params   :: Vector{TaxBenefitSystem{T}} ) :: NamedTuple where T # fixme simpler way of declaring this?
        num_systems = size( params )[1]
        println("start of do_one_run; using $(settings.means_tested_routing) routing")
        load_prices( settings, false )
        for p in 1:num_systems
            println("sys $p")
            println(params[p].it)
        end
        if settings.num_households == 0
            println( "getting households" )
            @time settings.num_households, settings.num_people, nhh2 =
                FRSHouseholdGetter.initialise( settings )
            BenefitGenerosity.initialise( MODEL_DATA_DIR*"/disability/" )
        end

        # vary generosity of disability benefits
        for sysno in 1:num_systems
            adjust_disability_eligibility!( params[sysno].nmt_bens )
        end

        frames :: NamedTuple = initialise_frames( T, settings, num_systems )
        frame_starts = FrameStarts(0,0,0,0)
        println( "starting run " )
        @time for hno in 1:settings.num_households
            hh = FRSHouseholdGetter.get_household( hno )
            if hno % 100 == 0
                println( "on household hno $hno hid=$(hh.hid) year=$(hh.interview_year)")
            end
            for sysno in 1:num_systems
                res = do_one_calc( hh, params[sysno], settings )
                frame_starts = add_to_frames!( frames, hh, res,  sysno, frame_starts, num_systems )
            end
        end #household loop
        println( "dumping frames" )
        dump_frames( settings, frames )
        return frames
    end # do one run

end # module


module POMDPGym
    using POMDPs
    using POMDPModels
    import POMDPModelTools:render
    using POMDPModelTools
    using POMDPPolicies
    using PyCall
    import PyCall: hasproperty
    using Parameters
    using Random
    using ImageCore
    using ImageShow
    using Images
    using Distributions
    import Cairo
    using Compose
    using StaticArrays
    using LinearAlgebra
    using GridInterpolations

    include("pycalls.jl")

    export GymPOMDP, reset!, step!, finished, render, close, init_mujoco_render
    include("gym_pomdp.jl")

    export torgb, preproc_atari_frame, AtariPOMDP
    include("atari_helpers.jl")

    export GridWorldMDP, GWPos, LavaWorldMDP, PendulumMDP, PendulumPOMDP, InvertedPendulumPOMDP, InvertedPendulumMDP, ContinuousBanditMDP, random_lava, ContinuumWorldMDP, Circle, Vec2, Vec4, LagrangeConstrainedPOMDP, RewardModPOMDP, isfailure,
    CollisionAvoidanceMDP, OptimalCollisionAvoidancePolicy, EpisodicSafetyGym,
    CartPoleMDP
    include("extra/gridworld.jl")
    include("extra/lavaworld.jl")
    include("extra/pendulum.jl")
    include("extra/continuous_bandit.jl")
    include("extra/continuumworld.jl")
    include("extra/lagrange_constraints.jl")
    include("extra/reward_mod.jl")
    include("extra/collision_avoidance.jl")
    include("extra/episodic_safety_gym.jl")
    include("extra/cartpole.jl")

    export AdversarialMDP, AdversarialPOMDP, AdditiveAdversarialMDP, AdditiveAdversarialPOMDP, disturbances, disturbanceindex
    include("extra/adversarialmdp.jl")
end # module


#!/usr/bin/env julia
using CSV, DataFrames, PGFPlotsX
it_covid19_pop = CSV.File("data/it_covid19_pop.csv") |> DataFrame
plotdicts = [
    Dict(:num => :CasesFemale, :den => :PopFemale, :lab => "incidens kvinnor"),
    Dict(:num => :CasesMale, :den => :PopMale, :lab => "incidens män"),
    Dict(:num => :DeathsFemale, :den => :PopFemale, :lab => "mortalitet kvinnor"),
    Dict(:num => :DeathsMale, :den => :PopMale, :lab => "mortalitet män")
]
push!(PGFPlotsX.CUSTOM_PREAMBLE, raw"\usepackage{fontspec}")
push!(PGFPlotsX.CUSTOM_PREAMBLE, raw"\setmainfont{Asana-Math}")
p = @pgf Axis({ymode="log", xlabel="Ålder", title="COVID-19 Italien till 2020-03-19",
        legend_pos="outer north east", xmajorgrids, ymajorgrids, yminorgrids})
ages = it_covid19_pop[:AgeGrp10Start]
for pd in plotdicts
    rate = it_covid19_pop[pd[:num]] ./ it_covid19_pop[pd[:den]]
    @pgf push!(p, PlotInc(Table([ages, rate])), LegendEntry(pd[:lab]))
end
pgfsave("../../images/it_covid19.svg", p)


@testset "right" begin
    _test_binary_op(TimeDag.right, (x, y) -> y)
end

@testset "left" begin
    _test_binary_op(TimeDag.left, (x, y) -> x)
end

@testset "align" begin
    @test _eval(TimeDag.align(n4, n1)) == b1
    @test _eval(TimeDag.align(n1, n4)) == Block([
        DateTime(2000, 1, 1) => 1,
        DateTime(2000, 1, 2) => 2,
        DateTime(2000, 1, 3) => 3,
        DateTime(2000, 1, 4) => 4,
        DateTime(2000, 1, 5) => 4,
        DateTime(2000, 1, 6) => 4,
        DateTime(2000, 1, 7) => 4,
    ])
end

@testset "coalign" begin
    @testset "unary" begin
        @test coalign(n4) === n4
        @test coalign(n4; alignment=INTERSECT) === n4
        @test coalign(n4; alignment=LEFT) === n4
        @test coalign(n4; alignment=UNION) === n4
    end

    @testset "binary" begin
        for alignment in (INTERSECT, LEFT, UNION)
            na, nb = coalign(n1, n2; alignment)
            @test _eval(na) == _eval(left(n1, n2, alignment))
            @test _eval(nb) == _eval(right(n1, n2, alignment))
        end
    end

    @testset "ternary" begin
        for alignment in (INTERSECT, LEFT, UNION)
            na, nb, nc = coalign(n1, n2, n3; alignment)

            l(x, y) = left(x, y, alignment)
            r(x, y) = right(x, y, alignment)

            @test _eval(na) == _eval(l(n1, l(n2, n3)))
            @test _eval(nb) == _eval(r(n1, l(n2, n3)))
            @test _eval(nc) == _eval(r(n1, r(n2, n3)))
        end
    end

    @testset "ordering" begin
        # We don't want to generate new different nodes if we coalign the same nodes in a
        # different order, when alignment is order-independent.
        nodes = [n1, n2, n3]
        for alignment in (INTERSECT, UNION)
            for indices in permutations([1, 2, 3])
                new_nodes = coalign(nodes[indices]...; alignment)
                @test new_nodes === coalign(nodes...; alignment)[indices]
            end
        end

        # For the case of LEFT, the first node must stay in place, but we should be able to
        # permute the order of the others and not change the answer.
        for indices in permutations([1, 2, 3])
            new_nodes = coalign(n4, nodes[indices]...; alignment=LEFT)
            @test new_nodes === coalign(n4, nodes...; alignment=LEFT)[[1; 1 .+ indices]]
        end
    end
end

@testset "first_knot" begin
    @test _eval(first_knot(n4)) == b4[1:1]
    @test _eval(first_knot(n_boolean)) == b_boolean[1:1]
    @test _eval(empty_node(Float64)) == Block{Float64}()
end

@testset "active_count" begin
    @test _eval(active_count(n1)) == Block([b1.times[1]], [1])
    @test _eval(active_count(n1, n2)) == Block([b1.times[1], b2.times[1]], [1, 2])
    @test _eval(active_count(n1, n2, n3)) == Block([b1.times[1], b2.times[1]], [2, 3])
    @test _eval(active_count(n1, n2, n3, n4)) == Block([b1.times[1], b2.times[1]], [3, 4])

    # Testing for optimisations.
    @test active_count(n1, n2) === active_count(n2, n1)
    @test active_count(n1, n2, n3) === active_count(n2, n1, n3)
    @test active_count(n1, n2, n3) === active_count(n3, n1, n2)
    @test active_count(n1, n2, n3) === active_count(n3, n2, n1)
end


using Test, RecursiveArrayTools, StaticArrays

using DiffEqBase: NAN_CHECK

@test !NAN_CHECK(3.0+4.0im)
@test NAN_CHECK(NaN)

u1 = ones(3)
@test !NAN_CHECK(u1)
u1′ = copy(u1)
u1′[2] = NaN
@test NAN_CHECK(u1′)

u2 = [SA[1.0 1.0; 1.0 1.0] for i = 1:3]
@test !NAN_CHECK(u2)
u2′ = copy(u2)
u2′[2] = SA[1.0 NaN; 1.0 1.0]
@test NAN_CHECK(u2′)

u3 = VectorOfArray([ones(5), ones(5)])
@test !NAN_CHECK(u3)
u3′ = recursivecopy(u3)
u3′[2][3] = NaN
@test NAN_CHECK(u3′)

u4 = ArrayPartition(u1, u2, u3)
@test !NAN_CHECK(u4)
u4_1 = ArrayPartition(u1′, u2, u3)
@test NAN_CHECK(u4_1)
u4_2 = ArrayPartition(u1, u2′, u3)
@test NAN_CHECK(u4_2)
u4_3 = ArrayPartition(u1, u2, u3′)
@test NAN_CHECK(u4_3)

@test !NAN_CHECK(ArrayPartition(u4, u4))
@test NAN_CHECK(ArrayPartition(u4, u4_1))
@test NAN_CHECK(ArrayPartition(u4, u4_2))
@test NAN_CHECK(ArrayPartition(u4, u4_3))


abstract type AbstractJaxTracer{T,N} <: AbstractJaxArray{T,N} end
#  o::PyObject
#  dims::NTuple{N,Int}
#end

const AbstractTracerVector{T} = AbstractJaxTracer{T,1}
const AbstractTracerMatrix{T} = AbstractJaxTracer{T,2}

Base.size(a::AbstractJaxTracer) = a.dims
Base.length(a::AbstractJaxTracer) = prod(a.dims)
Base.eltype(a::AbstractJaxTracer{T}) where {T} = T
Base.axes(a::AbstractJaxTracer) = map(Base.OneTo, size(a))

#wrap python properties
aval(a::AbstractJaxTracer) = a.o.aval

PyCall.PyObject(a::AbstractJaxTracer) = a.o

function Base.show(io::IO, ::MIME"text/plain", a::TA) where {TA<:AbstractJaxTracer}
  println(io, "$(_szstr(a.dims)) $TA:")
  print(io, " ", a.o)
end

function Base.show(io::IO, a::TA) where {TA<:AbstractJaxTracer}
  print(io, "$(_szstr(a.dims)) $TA")
end

for name in (:TracedArray, :ParallelTracerArray, :JVPTracerArray, :BatchTracerArray)
  @eval begin
    struct $name{T,N} <: AbstractJaxTracer{T,N}
      o::PyObject
      dims::NTuple{N,Int}
    end

    function $name(parr::PyObject)
      aval = parr.aval
      T = np_to_jl_type(aval.dtype)
      return $name{T,length(aval.shape)}(parr, reverse(aval.shape))
    end
    Base.convert(::Type{<:$name}, o::PyObject) = $name(o)
  end
end


module OpenSpielWrapper

using ..AlphaZero
using OpenSpiel

struct Spec{G} <: GI.AbstractGameSpec
  spiel_game :: G # loaded with OpenSpiel.load_game
  suppress_warnings :: Bool
  heuristic_value :: Function
  vectorize_state :: Function
  symmetries :: Function
  action_string :: Function
  parse_action :: Function
  read_state :: Function
end

mutable struct Env{G, S, P} <: GI.AbstractGameEnv
  spec :: Spec{G}
  state :: S
  curplayer :: P
end

#####
##### Default functions and constructors
#####

default_heurisic_value(env) = 0.0
# not handling terminal states, not flip board for two-player games
default_vectorize_state(spec, state) = convert(Array{Float32}, observation_tensor(state))
default_symmetries(spec, state) = GI.symmetries(spec, state)
default_action_string(spec, action) = GI.action_string(spec, action)
default_parse_action(spec, str::String) = GI.parse_action(spec, str)
default_read_state(spec) = GI.read_state(spec)

"""
    Spec(spiel_game; kwargs...) <: AbstractGameSpec

Wrap an OpenSpiel game object (`CxxWrap.StdLib.SharedPtrAllocated{OpenSpiel.Game}`)
into an AlphaZero game spec.

# Keyword arguments

The following optional functions from `GameInterface` are not present in
OpenSpiel.jl and can be provided as keyword arguments:

  - `heuristic_value`
  - `vectorize_state`
  - `symmetries`
  - `action_string`
  - `parse_action`
  - `read_state`

You can silence warnings by setting `suppress_warnings=true`.
"""

function Spec(
  spiel_game;
  suppress_warnings = false,
  heuristic_value = default_heurisic_value,
  vectorize_state = default_vectorize_state,
  symmetries = default_symmetries,
  action_string = default_action_string,
  parse_action = default_parse_action,
  read_state = default_read_state
)
  nplayers = num_players(spiel_game)
  @assert nplayers == 1 || nplayers == 2 "
  AlphaZero.jl only supports games with one or two players."
  if heuristic_value == default_heurisic_value && !suppress_warnings
    @warn "The `default_heuristic_value` function was not provided and so " *
      "algorithms such as MinMax may not work with this environment."
  end
  return Spec(
    spiel_game,
    suppress_warnings,
    heuristic_value,
    vectorize_state,
    symmetries,
    action_string,
    parse_action,
    read_state
  )
end

#####
##### GameInterface API
#####

function GI.init(spec::Spec)
  state = new_initial_state(spec.spiel_game)
  curplayer = current_player(state)
  return Env(spec, state, curplayer)
end

GI.spec(env::Env) = env.spec

# Queries on specifications

GI.two_players(spec::Spec) = num_players(spec.spiel_game) == 2

# TODO kinda ugly, may be game relevant
GI.actions(spec::Spec) = collect(0:num_distinct_actions(spec.spiel_game)-1)

GI.vectorize_state(spec::Spec, state) = spec.vectorize_state(spec, state)

# Operations on environments

GI.current_state(env::Env) = deepcopy(env.state)

function GI.set_state!(env::Env, state)
  env.state = deepcopy(state)
  env.curplayer = current_player(state)
end

GI.game_terminated(env::Env) = is_terminal(env.state)

function GI.white_playing(env::Env)
  wp = env.curplayer==0
  return is_terminal(env.state) && GI.two_players(env.spec) ? !wp : wp
end

# not all games implement legal_actions_mask
function GI.actions_mask(env::Env)
  try
    return convert(Array{Bool}, legal_actions_mask(env.state))
  catch e
    @error "AlphaZero.jl only works with OpenSpiel environments that implement " *
      "the `legal_actions_mask` function."
    rethrow(e)
  end
end

#TODO is_chance_node, is_simultaneous_node
function GI.play!(env::Env, action)
  apply_action(env.state, action)
  if !is_terminal(env.state)
    env.curplayer = current_player(env.state)
  end
end

GI.white_reward(env::Env) = returns(env.state)[1]

# available_actions dispatched
GI.available_actions(env::Env) = legal_actions(env.state)

GI.heuristic_value(env::Env) = env.spec.heuristic_value(env)

# Symmetries

GI.symmetries(spec, state) = spec.symmetries(spec.spiel_game, state)

# Interactive utilities

function GI.render(env::Env)
  show(env.state)
  print("\n\n") # GameInterface expects the rendering to end with an empty line
end

# action_string, parse action require Env instead of Spec for OpenSpiel functions
GI.action_string(env::Env, action) = action_to_string(env.state, action)
GI.action_string(spec::Spec, action) = spec.action_string(spec, action)

GI.parse_action(env::Env, str::String) = string_to_action(env.state, str)
GI.parse_action(spec::Spec, str::String) = spec.parse_action(spec, str)

GI.read_state(spec::Spec) = spec.read_state(spec.spiel_game)

end # module


function process_all_retained!(gen_fn::Unfold{T,U}, params::Tuple, argdiffs::Tuple,
                               choices_or_selection, prev_length::Int, new_length::Int,
                               retained_and_targeted::Set{Int}, state) where {T,U}

    len_diff = argdiffs[1]
    init_state_diff = argdiffs[2]
    param_diffs = argdiffs[3:end] # a tuple of diffs

    if any(diff != NoChange() for diff in param_diffs)

        # visit every retained kernel application
        state_diff = init_state_diff
        for key=1:min(prev_length,new_length)
            state_diff = process_retained!(gen_fn, params, choices_or_selection,
                key, (NoChange(), state_diff, param_diffs...), state)
        end

    else
        # every parameter diff is NoChange()

        # visit only certain retained kernel applications
        to_visit::Vector{Int} = sort(collect(retained_and_targeted))
        key = 0
        state_diff = init_state_diff
        if state_diff != NoChange()
            key = 1
            visit = true
            while visit && key <= min(prev_length, new_length)
                state_diff = process_retained!(gen_fn, params, choices_or_selection,
                    key, (NoChange(), state_diff, param_diffs...), state)
                key += 1
                visit = (state_diff != NoChange())
            end
        end
        for i=1:length(to_visit)
            if key > to_visit[i]
                # we have already visited it
                continue
            end
            key = to_visit[i]
            visit = true
            while visit && key <= min(prev_length, new_length)
                state_diff = process_retained!(gen_fn, params, choices_or_selection,
                    key, (NoChange(), state_diff, param_diffs...), state)
                key += 1
                visit = (state_diff != NoChange())
            end
        end
    end
end

"""
Process all new applications.
"""
function process_all_new!(gen_fn::Unfold{T,U}, params::Tuple, choices_or_selection,
                          prev_len::Int, new_len::Int, state) where {T,U}
    for key=prev_len+1:new_len
        process_new!(gen_fn, params, choices_or_selection, key, state)
    end
end


# ------------------------------------------------------------------
# Licensed under the ISC License. See LICENSE in the project root.
# ------------------------------------------------------------------

"""
    AbstractProblem

A generic problem in geostatistics.
"""
abstract type AbstractProblem end

#------------------
# IMPLEMENTATIONS
#------------------
include("problems/estimation_problem.jl")
include("problems/simulation_problem.jl")
include("problems/learning_problem.jl")


## code for figure gen
include("src/simulation/PaperSimulation.jl")
include("src/exampleanalysis/PaperAnalysis.jl")
include("src/figures/PaperFigures.jl")
using Plots, StatsPlots, Dates, OceanWaveSpectralFitting
theme(:vibrant)
fontsize = 10
default(xtickfontsize = fontsize, ytickfontsize = fontsize, colorbar_tickfontsize = fontsize, legendfontsize = fontsize, guidefontsize = fontsize)
pyplot(dpi = 300)
foldername = "fig/"
isdir(foldername) || mkdir(foldername)

## Figure 1 - example S(ω,ϕ)
pyplot(size = (1000,300))
plot(PaperFigures.Figure1(), colorbar = false)
savefig(foldername * "figure1_nocolorbar.png")
plot(PaperFigures.ColorFig1(PaperFigures.Figure1()))
savefig(foldername * "figure1_colorbar.png")
println("Finished figure 1")

## Figure 2 - S(ω,ϕ) vs f(ω)
pyplot(size = (1000,500))
plot(plot(PaperFigures.Figure2a()), plot(PaperFigures.Figure2b()), layout = (1,2))
savefig(foldername * "figure2.png")
println("Finished figure 2")

## Figure 3 - buoy data summary
pyplot(size = (1000,1000))
plot(PaperFigures.Figure3())
savefig(foldername * "figure3.png")
println("Finished figure 3")

## Figure 4 - S(ω,ϕ), f(ω), D(ω,ϕ)
pyplot(size = (1000,250))
plot(PaperFigures.Figure4())
savefig(foldername * "figure4.png")
println("Finished figure 4")

## Figure 5 - Vary parameters
pyplot(size = (1000,600))
plot(PaperFigures.Figure5(),legend=false)
savefig(foldername * "figure5.png")
println("Finished figure 5")

## Figure 6 - Scenario S(ω,ϕ)
pyplot(size = (1000,300))
plot(PaperFigures.Figure6())
savefig(foldername * "figure6_nocolorbar.png")
plot(PaperFigures.ColorFig6(PaperFigures.Figure6()))
savefig(foldername * "figure6_colorbar.png")
println("Finished figure 6")

## Figure 7 - Marginal parameter estimates (simulation)
pyplot(size = (1000,800))
plot(PaperFigures.Figure7(truncate=true))
savefig(foldername * "figure7.png")
plot(PaperFigures.Figure7(truncate=false))
savefig(foldername * "figure7untruncated.png")
println("Finished figure 7")

## Figure 8 - Directional parameter estimates (simulation)
pyplot(size = (1000,1000))
plot(PaperFigures.Figure8(truncate=true))
savefig(foldername * "figure8.png")
plot(PaperFigures.Figure8(truncate=false))
savefig(foldername * "figure8untruncated.png")
println("Finished figure 8")

## Figure 9 - Heatmap of ̂R(ω,t)
pyplot(size = (1000,250))
plot(PaperFigures.Figure9())
savefig(foldername * "figure9.png")
println("Finished figure 9")

## Figure 10 - Spectrogram with cutoffs
pyplot(size = (1000,250))
plot(PaperFigures.Figure10())
savefig(foldername * "figure10.png")
println("Finished figure 10")

## Figure 11 - Parameter estimates (buoy data)
pyplot(size = (1000,750))
plot(PaperFigures.Figure11())
savefig(foldername * "figure11.png")
println("Finished figure 11")


using Test, TestPackage3

@static if VERSION > v"1.3"
  @test hello3("Julia") == "Hello, Julia"
elseif VERSION > v"1.2"
  @test domath3(2.0) ≈ 7.0
else
  multiply3(2.0) == 8.0
end


#=
    ---------------------------------------------------------------------------
    Implements a few collections that the main mcmc sampler from the file
    `mcmc.jl` uses to keep track of the current state and history of where it
    was. The following structures are defined:
    - AccptTracker   : for tracking historical acceptance rate
    - ParamHistory   : for storing the parameter chain
    - ActionTracker  : keeps track of what to do on a given iteration
    - Workspace      : main workspace for `mcmc` function from `mcmc.jl`
    - ParamUpdtDefn  : defines a single parameter update step
    - GibbsDefn      : defines an entire gibbs sweep of parameter updates
    ---------------------------------------------------------------------------
=#

import Base: last, getindex, length, display, eltype

#===============================================================================
                        Workspace for the MCMC chain
===============================================================================#
"""
    struct MCMCWorkspace{T,S,V}

Implements a few collections that the main mcmc sampler from the file
`mcmc.jl` uses to keep track of the current state and history of where it
was. The following structures are defined:
- `θ_chain` for the parameter chain
- `updates` #TODO
- `pos` for the current position of the sampler
- `mean` for the current sample mean of the sampler
- `cov` for the current sample covariance matrix of sampler
- `updates_hist` #TODO
"""
struct MCMCWorkspace{T,S,V}
    θ_chain::Vector{T}
    updates::S
    pos::Vector{Bool}
    mean::Vector{V}
    cov::Matrix{V}
    updates_hist::Vector

    function MCMCWorkspace(setup::MCMCSetup, schedule, θ::T) where T
        # TODO create an object that pre-allocates the memory based on schedule
        # this is difficult due to potential fusing of kernels which results
        # in a random overall number of updates
        θ_chain = [θ]
        S = typeof(setup.updates)
        pos = list_of_pos_params(setup.updates, θ)
        V = eltype(θ)
        new{T,S,V}(θ_chain, setup.updates, pos, copy(θ), zero(θ*θ'), [])
    end
end

function list_of_pos_params(updates, θ)
    pos = fill(false, length(θ))
    for updt in updates
        if typeof(updt) <: ParamUpdate{MetropolisHastingsUpdt}
            ind_pos = indices_of_pos(updt.updt_coord, updt.t_kernel.pos)
            for (c,p) in ind_pos
                if p
                    pos[c] = true
                end
            end
        end
    end
    pos
end

function indices_of_pos(coords, pos)
    coord_indices = indices(coords)
    @assert length(coord_indices) == length(pos)
    zip(coord_indices, pos)
end

function update!(ws::MCMCWorkspace, acc, θ, i)
    not_imputation = !(typeof(ws.updates[i]) <: Imputation)
    not_imputation && push!(ws.θ_chain, θ)
    not_imputation && update_mean_cov!(ws.mean, ws.cov, θ, length(ws.θ_chain), ws.pos)
    register_accpt!(ws.updates[i], acc)
end

function update_mean_cov!(θ_mean, θ_cov, θ, θ_len, pos)
    ϑ = neutralise_domain(θ, pos)
    prev_mean = copy(θ_mean)
    θ_mean .*= θ_len/(θ_len+1)
    θ_mean .+= (ϑ./(θ_len+1))
    old_sum_sq = (θ_len-1)/θ_len * θ_cov + prev_mean * prev_mean'
    new_sum_sq = old_sum_sq + (ϑ * ϑ')/(θ_len)
    θ_cov .= new_sum_sq - (θ_len+1)/θ_len * (θ_mean * θ_mean')
end

function neutralise_domain(θ, pos)
    ϑ = copy(θ)
    ϑ[pos] = log.(ϑ[pos])
    ϑ
end

function readjust!(ws::MCMCWorkspace, mcmc_iter)
    for i in 1:length(ws.updates)
        readjust!(ws.updates[i], ws.cov, mcmc_iter)
    end
end

function fuse!(ws::MCMCWorkspace, schedule)
    # NOTE for now, the `fuse!` function is quite restricted and will fuse all
    # Metropolis-Hastings transition kernels into a single one
    MH_updates = [u for u in ws.updates
                  if typeof(u) <: ParamUpdate{<:MetropolisHastingsUpdt}]
    order_updates!(MH_updates)
    fusion_updt_coord = fuse_coord(MH_updates)
    fusion_cov = Matrix(view(ws.cov, fusion_updt_coord, fusion_updt_coord))
    fusion_kernel = fuse_kernels(MH_updates, fusion_cov)
    fusion_priors = fuse_priors(MH_updates)
    fusion_aux = fuse_aux(MH_updates)
    fusion_param_updt = ParamUpdate(MetropolisHastingsUpdt(), fusion_updt_coord,
                                    ws.θ_chain[end], fusion_kernel,
                                    fusion_priors, fusion_aux)
    push!(ws.updates, fusion_param_updt)

    MH_updt_idx = [i for (i,u) in enumerate(ws.updates)
                   if typeof(u) <: ParamUpdate{<:MetropolisHastingsUpdt}]

    reschedule!(schedule, MH_updt_idx, length(ws.updates))
end

"""
    getindex(g::MCMCWorkspace, i::Int)

Return `i`th definition of parameter update
"""
getindex(g::MCMCWorkspace, i::Int) = g.updates[i]

"""
    length(g::MCMCWorkspace)

Return the total number of parameter updates in a single Gibbs sweep
"""
length(g::MCMCWorkspace) = length(g.updates)

mutable struct SingleElem{T} val::T end

set!(x::SingleElem{T}, y::T) where T = (x.val = y)

#===============================================================================
                    Workspace for the Diffusion Model
===============================================================================#
"""
    Workspace{ObsScheme,S,TX,TW,R,TP,TZ}

The main container of the `mcmc` function from `mcmc.jl` in which most data
pertinent to sampling is stored
"""
struct Workspace{ObsScheme,S,TX,TW,R,TP,TZ}# ,Q, where Q = eltype(result)
    # Related to imputed path
    Wnr::Wiener{S}         # Wiener, driving law
    XXᵒ::Vector{TX}        # Diffusion proposal paths
    XX::Vector{TX}         # Accepted diffusion paths
    WWᵒ::Vector{TW}        # Driving noise of proposal
    WW::Vector{TW}         # Driving noise of the accepted paths
    Pᵒ::Vector{R}          # Guided proposals parameterised by proposal param
    P::Vector{R}           # Guided proposals parameterised by accepted param
    fpt::Vector            # Additional information about first passage times
    # Related to historically sampled paths
    skip_for_save::Int64            # Thining parameter for saving path
    paths::Vector                   # Storage with historical, accepted paths
    time::Vector{Float64}           # Storage with time axis
    # Related to the starting point
    x0_prior::TP
    z::SingleElem{TZ}
    #recompute_ODEs::Vector{Bool}    # Info on whether to recompute H,Hν,c after resp. param updt

    """
        Workspace(setup::DiffusionSetup{ObsScheme})

    Initialise workspace of the mcmc sampler according to a `setup` variable
    """
    function Workspace(setup::DiffusionSetup{ObsScheme}) where ObsScheme
        x0_prior, Wnr = deepcopy(setup.x0_prior), deepcopy(setup.Wnr)
        XX, WW = deepcopy(setup.XX), deepcopy(setup.WW)
        P, fpt = deepcopy(setup.P), deepcopy(setup.fpt)

        # forcedSolve defines type by the starting point, make sure it matches
        x0_guess = eltype(eltype(XX))(setup.x0_guess)
        TW, TX, S, R = eltype(WW), eltype(XX), valtype(Wnr), eltype(P)
        TP = typeof(x0_prior)
        m = length(P)

        y = copy(x0_guess)
        for i in 1:m
            WW[i] = Bridge.samplepath(P[i].tt, zero(S))
            sample!(WW[i], Wnr)
            WW[i], XX[i] = forcedSolve(EulerMaruyamaBounded(), y, WW[i], P[i])    # this will enforce adherence to domain
            while !checkFpt(ObsScheme(), XX[i], fpt[i])
                sample!(WW[i], Wnr)
                forcedSolve!(EulerMaruyamaBounded(), XX[i], y, WW[i], P[i])    # this will enforce adherence to domain
            end
            y = XX[i].yy[end]
        end
        y = x0_guess
        ll = ( logpdf(x0_prior, y)
               + path_log_likhd(ObsScheme(), XX, P, 1:m, fpt, skipFPT=true)
               + lobslikelihood(P[1], y) )

        XXᵒ, WWᵒ, Pᵒ = deepcopy(XX), deepcopy(WW), deepcopy(P)

        # compute the white noise that generates x0_guess under the initial posterior
        z = inv_start_pt(y, x0_prior, P[1])
        TZ = typeof(z)
        z = SingleElem{TZ}(z)

        skip = setup.skip_for_save
        _time = collect(Iterators.flatten(p.tt[1:skip:end-1] for p in P))

         #check_if_recompute_ODEs(setup)
         (workspace = new{ObsScheme,S,TX,TW,R,TP,TZ}(Wnr, XXᵒ, XX, WWᵒ, WW, Pᵒ,
                                                     P, fpt, skip, [], _time,
                                                     x0_prior, z),
          ll = ll, θ = params(P[1].Target))
    end

    function Workspace(ws::Workspace{ObsScheme,S,TX,TW,R̃,TP,TZ}, P::Vector{R},
                       Pᵒ::Vector{R}) where {ObsScheme,S,TX,TW,R̃,R,TP,TZ}
        new{ObsScheme,S,TX,TW,R,TP,TZ}(ws.Wnr, ws.XXᵒ, ws.XX, ws.WWᵒ, ws.WW, Pᵒ,
                                       P, ws.fpt, ws.skip_for_save, ws.paths,
                                       ws.time, ws.x0_prior, ws.z)
    end
end

eltype(::SamplePath{T}) where T = T
eltype(::Type{SamplePath{T}}) where T = T
solver_type(::Workspace{O,B,ST}) where {O,B,ST} = ST
obs_scheme(::Workspace{O}) where O = O()

next(ws::Workspace, ::Any) = ws
function next(ws::Workspace, updt::Imputation{<:Block})
    XX, P, Pᵒ, bl = ws.XX, ws.P, ws.Pᵒ, updt.blocking
    θ = params(P[1].Target)
    vs = find_end_pts(bl, XX)

    P_new = [GuidPropBridge(Pi, bl.Ls[i], vs[i], bl.Σs[i], bl.change_pts[i],
                            θ, bl.aux_flags[i]) for (i,Pi) in enumerate(P)]
    Pᵒ_new = [GuidPropBridge(Pᵒi, bl.Ls[i], vs[i], bl.Σs[i], bl.change_pts[i],
                             θ, bl.aux_flags[i]) for (i,Pᵒi) in enumerate(Pᵒ)]
    Workspace(ws, P_new, Pᵒ_new)
end

"""
    save_imputed!(ws::Workspace)

Save the entire path spanning all segments in `XX`. Only 1 in every `ws.skip_for_save`
points is saved to reduce storage space.
"""
function save_imputed!(ws::Workspace)
    skip = ws.skip_for_save
    push!(ws.paths, collect(Iterators.flatten(ws.XX[i].yy[1:skip:end-1]
                                               for i in 1:length(ws.XX))))
end

function adaptation_object(setup::DiffusionSetup, ws::Workspace)
    adpt = deepcopy(setup.adaptive_prop)
    if typeof(adpt) <: Adaptation{Val{false}}
        return nothing
    end
    m = length(ws.XX)
    resize!(adpt, m, [length(ws.XX[i]) for i in 1:m])
    adpt
end

function create_workspace(setup::MCMCSetup, schedule::MCMCSchedule, θ)
    MCMCWorkspace(setup, schedule, θ)
end

function create_workspace(setup::T) where {T <: ModelSetup}
    Workspace(setup)
end


function testSMC(model::SMCModel, N::Int64, numTrials::Int64)
  smcio = SMCIO{model.particle, model.pScratch}(N, model.maxn)
  println("Running SMC. N = $N")
  for i=1:numTrials
    @time smc!(model, smcio)
    left = max(length(smcio.logZhats)-10,1)
    right = length(smcio.logZhats)
    println(smcio.logZhats[left:right])
  end
end

function testCSMC(model::SMCModel, N::Int64, numTrials::Int64)
  smcio = SMCIO{model.particle, model.pScratch}(N, model.maxn)
  println("Running CSMC. N = $N")
  smc!(model, smcio)
  v::Vector{model.particle} = Vector{model.particle}(undef, smcio.n)
  for i in 1:smcio.n
    v[i] = model.particle()
  end
  SimpleSMC.pickParticle!(v, smcio)
  for i=1:numTrials
    @time csmc!(model, smcio, v, v)
    left = max(length(smcio.logZhats)-10,1)
    right = length(smcio.logZhats)
    println(smcio.logZhats[left:right])
  end
end

function testCSMCBS(model::SMCModel, lM::F, N::Int64, numTrials::Int64) where F<:Function
  smcio = SMCIO{model.particle, model.pScratch}(N, model.maxn)
  println("Running CSMC-BS. N = $N")
  smc!(model, smcio)
  v::Vector{model.particle} = Vector{model.particle}(undef, smcio.n)
  for i in 1:smcio.n
    v[i] = model.particle()
  end
  SimpleSMC.pickParticle!(v, smcio)
  for i=1:numTrials
    @time csmc!(model, lM, smcio, v, v)
    left = max(length(smcio.logZhats)-10,1)
    right = length(smcio.logZhats)
    println(smcio.logZhats[left:right])
  end
end


# Goal: Make sure some combinations that raise errors do so
using Test
using OrdinaryDiffEq
using LinearAlgebra
using DiffEqProblemLibrary.ODEProblemLibrary: importodeproblems;
importodeproblems();
import DiffEqProblemLibrary.ODEProblemLibrary:
    prob_ode_linear, prob_ode_2Dlinear, prob_ode_lotkavoltera, prob_ode_fitzhughnagumo

import ProbNumDiffEq: remake_prob_with_jac

@testset "Fixed-timestep requires dt" begin
    prob = prob_ode_lotkavoltera
    @test_throws ErrorException solve(prob, EK0(), adaptive=false)
    @test solve(prob, EK0(), adaptive=false, dt=0.05) isa ProbNumDiffEq.ProbODESolution
end

@testset "`dense=true` requires `smooth=true`" begin
    prob = prob_ode_lotkavoltera
    @test_throws ErrorException solve(prob, EK0(smooth=false))
end

@testset "`save_everystep=false` requires `smooth=false`" begin
    prob = prob_ode_lotkavoltera
    @test_throws ErrorException solve(prob, EK0(smooth=true), save_everystep=false)
end

@testset "`dense=false` warns if `smooth=true`" begin
    prob = prob_ode_lotkavoltera
    @test_logs (
        :warn,
        "If you set dense=false for efficiency, you might also want to set smooth=false.",
    ) solve(prob, EK0(smooth=true), dense=false)
end


@testset "Testing Compiler" begin
    # Make sure that type conversion works correctly.
    x = rand(Float32, 2, 2)
    @test OneDNN.tracetype(x) == Mjolnir.Shape{typeof(x)}(size(x))
end


# © 2016 Massachusetts Institute of Technology.  See LICENSE file for details.
abstract Instance
Base.keys(i::Instance) = keys(i._storage)
Base.values(i::Instance) = values(i._storage)
Base.getindex(i::Instance, a::ANode) = i._storage[a]

@deprecate evaluator(x...) instance(x...)

instance(n::ANode...) = instance(collect(ANode, n))
instance(T::Type, n::ANode...) = instance(T, collect(ANode, n))
instance(pool::StoragePool, n::ANode...) = instance(pool, collect(ANode, n))

instance{A<:ANode}(n::Vector{A}) = instance(Array{Float64}, n)
instance{A<:ANode}(T::Type, n::Vector{A}) = instance(NullPool(T), n)

output_node(n::ANode) = mutates(n) ? output_node(n.input[1]) : n

type RefCounter
  count::Int
  RefCounter() = new(0)
end

incr(r::RefCounter) = r.count += 1
decr(r::RefCounter) = (@assert r.count > 0; r.count -= 1; r.count == 0)
iszero(r::RefCounter) = r.count == 0

incr{T}(d::Dict{T,RefCounter}, x::T) = incr(get!(RefCounter, d, x))
decr{T}(d::Dict{T,RefCounter}, x::T) = decr(d[x])
iszero{T}(d::Dict{T,RefCounter}, x::T) = iszero(d[x])

function generate_function(sorted::Vector{ANode}, dict)
  b = :(begin end)

  if length(sorted) >= 20
    n = length(sorted)
    for i in 1:8
      ix = div((i-1)*n,8)+1:div(i*n,8)
      f = generate_function(sorted[ix], dict)
      push!(b.args, :($f()::Void))
    end
  else
    for i in sorted
      if sym(i) == :input || sym(i) == :const
        continue
      elseif sym(i) == :load
        if dict[i] != arg(i)
          push!(b.args, :(copy!($(dict[i]), $(arg(i)))))
        end
      else
        in = map(x->dict[x], i.input)
        if mutates(i)
          push!(b.args, :(Seep.do_forward!($i, $(in...))))
        else
          out = dict[i]
          push!(b.args, :(Seep.do_forward!($i, $out, $(in...))))
        end
      end
    end
  end

  push!(b.args, :nothing)

  @gensym SeepFunction
  @eval @compat $SeepFunction() = $b

  return SeepFunction
end

function instance{A<:ANode}(pool::StoragePool, n::Vector{A})
  sorted = toposort(n)

  T = atype(pool)
  dict = Dict{ANode,T}()
  userdict = Dict{ANode,T}()

  refcnt = Dict{ANode,RefCounter}()
  for a in n
    if sym(a) == :noop continue end
    incr(refcnt, output_node(a))
  end

  for a in sorted
    for b in a.input
      incr(refcnt, output_node(b))
    end
  end

  constants = ANode[]
  c_block = :(begin end)
  for i in sorted
    if sym(i) == :input
      userdict[i] = dict[i] = allocate(pool, i.size...)
    elseif sym(i) == :const
      if isa(arg(i), T)
        dict[i] =  arg(i)
      else
        incr(refcnt, i)
        t = dict[i] = allocate(pool, i.size...)
        push!(c_block.args, :(copy!($t, map(eltype($T), $(arg(i))))))
        push!(constants, i)
      end
    end
  end

  @gensym initialize_constants
  @eval $(initialize_constants)() = $c_block

  for i in sorted
    @assert !iszero(refcnt, output_node(i))

    if mutates(i)
      dict[i] = dict[i.input[1]]
    elseif sym(i) == :load && isa(arg(i), T)
      dict[i] = arg(i)
    elseif sym(i) != :input && sym(i) != :const
      dict[i] = allocate(pool, i.size...)
    end

    for a in i.input
      aa = output_node(a)
      if decr(refcnt, aa) && arg(aa) != dict[aa]
        deallocate(pool, dict[aa])
      end
    end
  end

  for i in n
    if sym(i) == :noop continue end
    userdict[i] = dict[i]

    aa = output_node(i)
    if decr(refcnt, aa) && arg(aa) != dict[aa]
      deallocate(pool, dict[aa])
    end
  end

  for i in constants
    decr(refcnt, i)
    deallocate(pool, dict[i])
  end

  for i in keys(refcnt)
    @assert iszero(refcnt, i)
  end

  local b
  if isempty(constants)
    b = :(begin end)
  else
    b = :(begin
      if $(pool).token[1] != $(QuoteNode(initialize_constants))
        $(pool).token[1] = $(QuoteNode(initialize_constants))
        $(initialize_constants)()
      end
    end)
  end

  fun = generate_function(sorted, dict)
  push!(b.args, :($fun()::Void))

  @gensym SeepInstance
  typeblock = quote
    immutable $SeepInstance <: Instance
      _storage::Dict{ANode,$T}
      _pool::StoragePool
    end
    Base.show(io::IO, ::$SeepInstance) = print(io, "Seep Instance with $($(length(sorted))) nodes.")
    @compat (::$SeepInstance)() = $b
    $SeepInstance($userdict, $pool)
  end

  for (n,a) in userdict
    if haskey(userdict, n) && n.name != ""
      if length(filter(x->x.name == n.name, sorted)) > 1
        warn("There are multiple nodes named $(n.name).")
        continue
      end
      push!(typeblock.args[2].args[3].args, :($(symbol(n.name))::$(typeof(a))))
      push!(typeblock.args[end].args, a)
    end
  end

  eval(typeblock)
end

# make sure run,dict = instance(...) works for a while, even though run and dict are the same thing
#Base.start(::Instance) = (Base.depwarn("replace run,dict=instance(...) with inst=instance(...)", :SeepStart); 1)
Base.start(::Instance) = 1
Base.done(::Instance, i::Int) = i > 2
Base.next(a::Instance, i::Int) = (a, i+1)


# ---------------------------------------------------------------------------- #
#
#   meshF2D.jl
#
#   Type for 2D frame meshes
#   Inherits from the abstract "meshF" type
#
#   intrico
#   Fall 2017
#
#   Max Opgenoord
#
# ---------------------------------------------------------------------------- #

"""
    MeshF2D

MeshF2D type:
Type for 2D meshes consisting of triangles. It holds the nodes and connectivity
information.
"""
struct MeshF2D <: MeshF

    dim::Int64              # Dimension of the problem

    n::Int64                # Number of nodes

    p::Matrix{Float64}      # Nodal locations
    t::Matrix{Int64}        # Triangle - node connectivity
    t2f::Matrix{Int64}      # Triangle - face connectivity
    f::Matrix{Int64}        # Face - node/triangle connectivity
    fb::Matrix{Int64}       # Boundary face info
    nodes::Array{Float64,3} # Nodes on which solution is evaluated -- needed to
                          #  compute Jacobian on face
    n2e::Vector{Vector{Int64}} # Node to edge connectivity

    boundtags::Vector{String} # boundary tags

end

"""
    MeshF( mesh::divido.Mesh2D )

Constructor that generates frame from `divido` mesh structure in `mesh`.
"""
function MeshF( mesh::divido.Mesh2D )

    n2e = genNEconnec2D( mesh.f, mesh.n )

    MeshF2D( 2, mesh.n, mesh.p, mesh.t, mesh.t2f, mesh.f, mesh.fb, mesh.nodes, n2e, mesh.boundtags )

end

"""
    MeshF2D( name::String )

Constructor for frame from mesh written in `name`.
"""
function MeshF2D( name::String )

    if name[end-3:end] == ".su2"
        (p_, t_, bel_, tags_) = divido.readSU2_2D( name )
    elseif name[end-3:end] == ".msh" # BAMG
        (p_, t_, bel_, tags_ ) = divido.readBAMG( name )
    elseif name[end-4:end] == ".mesh" # FEFLOA
        (p_, t_, bel_, tags_ ) = divido.readFEFLOA_2D( name )
    else
        error("MeshF2D: Unknown mesh type")
    end

    porder_ = divido.P1()

    (f_, t2f_, nodes_, ploc_, tloc_, fb_) = divido.genmesh( porder_, p_, t_, bel_ )

    n_ = size( p_, 1 )

    n2e_ = genNEconnec2D( f_, n_ )

    MeshF2D( 2, n_, p_, t_, t2f_, f_, fb_, nodes_, n2e_, tags_ )

end

"""
    genNEconnec2D( e::Matrix{Int64}, n::Int64 )

Generates node to edge connectivity given edge connectivity (`e`).
"""
function genNEconnec2D( e::Matrix{Int64}, n::Int64 )

    n2e = fill( fill(0,0), n )

    for jj in 1:size(e,1)

        ii1 = e[jj,1]
        ii2 = e[jj,2]

        if isempty( n2e[ii1] .== jj )
            append!(n2e[ii1], jj)
        end
        if isempty( n2e[ii2] .== jj )
            append!(n2e[ii2], jj)
        end

    end

    for kk in 1:n
        n2e[kk] = sort( n2e[kk] )
    end

    return n2e

end


#module ModiaBase

using .Instantiation
using .ModelElaboration

using Unitful
using .StructuralTransform
using .Synchronous: sample, Clock, previous, hold, positive, positiveChange, positiveEdge

Var(; args...) = Variable(; args...)

Float(value=nothing; info="", size=nothing, unit=NoUnits, displayUnit=NoUnits,
    min=nothing, max=nothing, start=nothing, fixed::Bool=false, nominal=nothing,
    variability=continuous,
    flow::Bool=false, state::Bool=true) = Variable(variability, Float64, size, value,
    unit, displayUnit, min, max, start, fixed, nominal, info, flow, state, general)

Boolean(value=nothing; info="", size=nothing, unit=NoUnits, displayUnit=NoUnits,
    min=nothing, max=nothing, start=nothing, fixed::Bool=false, nominal=nothing,
    variability=continuous,
    flow::Bool=false, state::Bool=true) = Variable(variability, Bool, size, value,
    unit, displayUnit, min, max, start, fixed, nominal, info, flow, state, general)

Integ(value=nothing; info="", size=nothing, unit=NoUnits, displayUnit=NoUnits,
    min=nothing, max=nothing, start=nothing, fixed::Bool=false, nominal=nothing,
    variability=continuous,
    flow::Bool=false, state::Bool=true) = Variable(variability, Int, size, value,
    unit, displayUnit, min, max, start, fixed, nominal, info, flow, state, general)

Str(value=nothing; info="", size=nothing, unit=NoUnits, displayUnit=NoUnits,
    min=nothing, max=nothing, start=nothing, fixed::Bool=false, nominal=nothing,
    variability=continuous,
    flow::Bool=false, state::Bool=true) = Variable(variability, String, size, value,
    unit, displayUnit, min, max, start, fixed, nominal, info, flow, state, general)

#=
Float(; args...) = Var(T=Float64; args...)
Float0(; args...) = Float(size=(); args...)

Boolean(; args...) = Var(T=Bool; args...)
Boolean0(; args...) = Boolean(size=(); args...)

Integ(; args...) = Var(T=Int; args...)
Integ0(; args...) = Integ(size=(); args...)

Str(; args...) = Var(T=String; args...)
Str0(; args...) = Str(size=(); args...)
=#

Parameter(; args...) = Variable(variability=parameter; args...)
Parameter(value; args...) = Variable(variability=parameter, value=value; args...)

Par(; args...) = Variable(variability=parameter; args...)
Par(value; args...) = Variable(variability=parameter, value=value; args...)

const undefined=nothing

#end


function same_geometry(a::AbstractVTKUnstructuredData, b::AbstractVTKUnstructuredData)
    dim(a) != dim(b) && return false
    num_of_points(a) != num_of_points(b) && return false
    num_of_cells(a) != num_of_cells(b) && return false

    _num_of_points = num_of_points(a)
    b_image_point_inds = zeros(Int, _num_of_points)
    for i in 1:_num_of_points
        for j in 1:_num_of_points
            @views if b.point_coords[:,i] == a.point_coords[:,j]
                b_image_point_inds[i] = j
                break
            end
        end
        if b_image_point_inds[i] == 0
            return false
        end
    end

    function map_cell_connectivity(b_cell_connectivity)
        b_cell_connectivity_mapped = zeros(b_cell_connectivity)
        for i in 1:length(b_cell_connectivity)
            b_cell_connectivity_mapped[i] = b_image_point_inds[b_cell_connectivity[i]]
        end
        b_cell_connectivity_mapped
    end

    _num_of_cells = num_of_cells(a)
    b_cell_connectivity_mapped = map(map_cell_connectivity, b.cell_connectivity)
    b_image_cell_inds = zeros(Int, _num_of_cells)
    for i in 1:_num_of_cells
        for j in 1:_num_of_cells
            if similar_cells(b_cell_connectivity_mapped[i], cell_type(b,i), a.cell_connectivity[j], cell_type(a,j))
                b_image_cell_inds[i] = j
                break
            end
        end
        if b_image_cell_inds[i] == 0
            return false
        end
    end
    return true
end
function same_geometry(a::VTKStructuredData, b::VTKStructuredData)
    return a.point_coords == b.point_coords
end
function same_geometry(a::VTKRectilinearData, b::VTKRectilinearData)
    return extents(a) == extents(b) && all(a.point_coords .== b.point_coords)
end
function same_geometry(a::VTKUniformRectilinearData, b::VTKUniformRectilinearData)
    return a.spacing == b.spacing && extents(a) == extents(b)
end
function same_geometry(a::AbstractStaticVTKData, b::AbstractStaticVTKData)
    return same_geometry(promote(a,b)...)
end
function same_ordered_geometry(a::AbstractVTKUnstructuredData, b::AbstractVTKUnstructuredData)
    dim(a) != dim(b) && return false
    num_of_points(a) != num_of_points(b) && return false
    num_of_cells(a) != num_of_cells(b) && return false

    _num_of_points = num_of_points(a)
    @views for i in 1:_num_of_points
        b.point_coords[:,i] == a.point_coords[:,i] || return false
    end

    _num_of_cells = num_of_cells(a)
    for i in 1:_num_of_cells
        similar_cells(a.cell_connectivity[i], cell_type(a,i), b.cell_connectivity[i], cell_type(b,i)) || return false
    end
    return true
end
function same_ordered_geometry(a::AbstractVTKMultiblockData, b::AbstractVTKMultiblockData)
    length(a) == length(b) || return false
    for i in 1:length(a)
        typeof(a[i]) == typeof(b[i]) && (typeof(a[i]) <: AbstractVTKMultiblockData || typeof(a[i]) <: AbstractVTKUnstructuredData ? same_ordered_geometry(a[i], b[i]) : a[i] == b[i]) || return false
    end
    return true
end

function same_geometry_shape(a::VTKStructuredData, b::VTKStructuredData)
    return size(a.point_coords) == size(b.point_coords)
end
function same_geometry_shape(a::VTKRectilinearData, b::VTKRectilinearData)
    return extents(a) == extents(b)
end
function same_geometry_shape(a::VTKUniformRectilinearData, b::VTKUniformRectilinearData)
    return extents(a) == extents(b)
end
function same_geometry_shape(a::AbstractVTKUnstructuredData, b::AbstractVTKUnstructuredData)
    dim(a) != dim(b) && return false
    num_of_points(a) != num_of_points(b) && return false
    num_of_cells(a) != num_of_cells(b) && return false

    _num_of_cells = num_of_cells(a)
    for i in 1:_num_of_cells
        similar_cells(a.cell_connectivity[i], cell_type(a,i), b.cell_connectivity[i], cell_type(b,i)) || return false
    end
    return true
end
same_geometry_shape(a::AbstractStaticVTKData, b::AbstractStaticVTKData) = false

function same_ordered_geometry_shape(a::AbstractVTKMultiblockData, b::AbstractVTKMultiblockData)
    length(a) == length(b) || return false
    for i in 1:length(a)
        typeof(a[i]) == typeof(b[i]) || return false
        if typeof(a[i]) <: AbstractVTKMultiblockData
            same_ordered_geometry_shape(a[i], b[i]) || return false
        else
            same_geometry_shape(a[i], b[i]) || return false
        end
    end
    return true
end

function same_data_shape(a::T, b::T) where {T<:AbstractVTKSimpleData}
    a_keys = keys(a.point_data)
    b_keys = keys(b.point_data)
    for i in a_keys
        in(i, b_keys) && size(a.point_data[i]) == size(b.point_data[i]) || return false
    end
    a_keys = keys(a.cell_data)
    b_keys = keys(b.cell_data)
    for i in a_keys
        in(i, b_keys) && size(a.cell_data[i]) == size(b.cell_data[i]) || return false
    end
    return true
end
same_data_shape(a::AbstractVTKStructuredData, b::AbstractVTKUnstructuredData) = false
same_data_shape(a::AbstractVTKMultiblockData, b::AbstractVTKSimpleData) = false

function same_data_shape(a::T, b::T) where {T<:AbstractVTKMultiblockData}
    length(a) == length(b)
    for i in 1:length(a)
        same_data_shape(a[i], b[i]) || return false
    end
    return true
end


"""
Callback to stop optimization if the function value is below some threshold.
"""
immutable FunctionStop{T}
    val::T
end

function (cb::FunctionStop)(os::OptimizationState)
    os.value ≤ cb.val
end

function (cb::FunctionStop)(tr::OptimizationTrace)
    tr[end].value ≤ cb.val
end

abstract type StopTolerance end;

"""
Tell the optimization procedure to exit when we hit `|g(x)|≤ tol * |g(x_0)|`.

For absolute, rather than relative, tolerance, set `g0norm = 1.0`.
"""
immutable GradientTolerance{T} <: StopTolerance
    # TODO: The g0norm is a bit superfluous
    tol::T
    g0norm::T
end

"""
Tell the optimization procedure to exit when we hit `f(x) ≤ f_L + tol * (f(x_0) - f_L)`,
where `f_L` is the minimum value of `f`.

For absolute, rather than relative, tolerance, set `f0 = 1.0 + fL`.
"""
immutable FunctionTolerance{T} <: StopTolerance
    # TODO: f0 and fL are superflous
    tol::T
    f0::T
    fL::T
end

"""
If we don't know the value, tell the optimization procedure that we want to find
the minimum function value achieved across a list of solvers,
by running them all to within a given gradient norm tolerance and number of iterations.
"""
immutable UnknownFunctionTolerance{T} <: StopTolerance
    # TODO: Maybe we just run these problems to high tolerance before
    #       starting this process, and then store the approximate minimum value.
    gtol::T
    iterations::Int
end

"Log of results from a given optimization run."
immutable OptimizationRun{T,Tf}
    iterations::Int
    fevals::Int
    gevals::Int
    hevals::Int
    cputime::T
    fval::Tf
    f0::Tf
    gnorm::Tf
    g0norm::Tf
    solvername::String
    problemname::String
    success::Bool
end

function OptimizationRun(iterations, fevals, gevals, hevals, cputime, fval,
                         f0, gnorm, g0norm, solvername, problemname,
                         ft::FunctionTolerance)
    success = fval ≤ ft.fL + ft.tol*(ft.f0 - ft.fL)
    OptimizationRun(iterations, fevals, gevals, hevals, cputime, fval, f0,
                    gnorm, g0norm, solvername, problemname, success)
end

function OptimizationRun(iterations, fevals, gevals, hevals, cputime, fval,
                         f0, gnorm, g0norm, solvername, problemname,
                         gt::GradientTolerance)
    success = gnorm ≤ gt.tol*gt.g0norm
    OptimizationRun(iterations, fevals, gevals, hevals, cputime, fval, f0,
                    gnorm, g0norm, solvername, problemname, success)
end

"Collection of options used create performance measures."
immutable TestSetup{T<:Real, Ts<:AbstractVector, Tn<:AbstractVector, Tl<:Real}
    solvers::Ts
    solvernames::Tn
    stoptype::Symbol
    stoptol::T
    timelog::Int # Number of calls to optimize to record timings. timelog = 0 calls optimize once, but does not record the timing
    maxiter::Int
    timelimit::Tl # Soft time limit for each optimize call
end


export PureStateAnsatz, PureChain

struct PureStateAnsatz{A,IT,OT,Anal,S} <: KetNeuralNetwork
    __ansatz::A
    in_size::S
end

"""
    PureStateAnsatz(ansatz, input_size)

Wraps an arbitrary function into a structure that can be used as a variational
ansatz for a pure state, assuming that the width of the input is `input_size`
(corresponding to the number of lattice sites in the input hilbert space).

For this to work, ansatz must be a function respecting the following:
```
ansatz(rand(input_size)) --> scalar
trainable(ansatz) -> NamedTuple
```

To make `trainable` work with your functions, see the docs of [`functor`](@ref)
"""
function PureStateAnsatz(ansatz, in_size)
    if ansatz isa PureStateAnsatz
        throw("ansatz is already a network!")
    end

    if in_size isa AbstractHilbert
        in_size = nsites(in_size)
    end

    # input type
    in_t = real(eltype(trainable_first(ansatz)))
    # output_type
    out_t = Complex{real(eltype(trainable_first(ansatz)))}
    # is_analytic
    anal = true
    ansatz = (ansatz,)
    S=typeof(in_size)
    return PureStateAnsatz{typeof(ansatz), in_t, out_t, anal,S}(ansatz, in_size)
end

"""
    PureChain(N, args...)

Constructs a FeedForward network (a chain) encoding a pure-state. `N` must be
the number of input sites, while args... should be a series of layers.

Note: when you want to use sum, to output the total sum of states, don't use it
and use instead `sum_autobatch`, defined in this package. `sum_autobatch` behaves
as sum only for vectors, while for batches of data it sum along batches.

ex:
```
PureChain(5, Dense(5,3), Dense(3,2), sum_autobatch)
"""
function PureChain(N, args...)
    ch = Chain(args...)
    return PureStateAnsatz(N, ch)
end

@forward PureStateAnsatz.__m_ansatz Base.length, Base.first, Base.last,
         Base.iterate, Base.lastindex

ansatz(psa::PureStateAnsatz) = first(getfield(psa, :__ansatz))
functor(psa::PureStateAnsatz) = psa.__ansatz,
        a -> PureStateAnsatz(a, input_size(psa))

cache(psa::PureStateAnsatz) =
    cache(psa.__m_ansatz, trainable_first(psa), input_type(psa), input_size(psa))

cache(psa::PureStateAnsatz, batch_sz::Int) =
    cache(psa.__m_ansatz, trainable_first(psa), input_type(psa), input_size(psa), batch_sz)

(c::PureStateAnsatz)(x::Vararg{N,V}) where {N,V} = ansatz(c)(x...)
(c::PureStateAnsatz)(cache::NNCache, σ) = ansatz(c)(cache, σ)

logψ!(out::AbstractArray, net::PureStateAnsatz, c::NNBatchedCache, σ::AStateBatch) =
    logψ!(out, ansatz(net), c, σ)

logψ_and_∇logψ!(∇logψ, out::AbstractMatrix, net::PureStateAnsatz, c::NNBatchedCache, σ) =
    logψ_and_∇logψ!(∇logψ[1], out, ansatz(net), c, σ)

logψ_and_∇logψ!(∇logψ, net::PureStateAnsatz, c::NNCache, σ) =
    logψ_and_∇logψ!(∇logψ[1], ansatz(net), c, σ)

input_type(::PureStateAnsatz{A,IT,OT,AN}) where {A,IT,OT,AN} = IT
input_size(psa::PureStateAnsatz) = getfield(psa, :in_size)
out_type(::PureStateAnsatz{A,IT,OT,AN}) where {A,IT,OT,AN} = OT
is_analytic(::PureStateAnsatz{A,IT,OT,AN}) where {A,IT,OT,AN} = AN

Base.propertynames(psa::PureStateAnsatz) = propertynames(ansatz(psa))
Base.getproperty(psa::PureStateAnsatz, x::Symbol) = begin
    x === :__ansatz   && return getfield(psa, x)
    x === :__m_ansatz && return return ansatz(psa)
    getproperty(ansatz(psa), x)
end
Base.setproperty!(psa::PureStateAnsatz, name::Symbol, x) = begin
    name === :__ansatz   && return setfield!(psa, name, x)
    name === :__m_ansatz && return setfield!(psa, ansatz(psa), x)
    setproperty!(ansatz(psa), name, x)
end

Base.getindex(psa::PureStateAnsatz, x::Int) = begin
    @assert x==1
    return ansatz(psa)
end


using SnoopCompile

SnoopCompile.@snoop "glp_compiles.csv" begin
    using GLPlot;GLPlot.init()
    using Colors, GeometryTypes
    glplot(rand(Float32, 32,32))
    glplot(rand(Float32, 32,32), :surface)
    glplot(rand(Point3f0,32))
    glplot(rand(Point3f0,32), :lines)
    glplot(rand(Point2f0,32), :lines)
    glplot(rand(Point2f0,32))
    glplot(RGBA{Float32}[RGBA{Float32}(rand(), rand(), rand(), rand()) for i=1:512, j=1:512])
    while isopen(GLPlot.viewing_screen)
        yield()
    end
end

using GLPlot
data = SnoopCompile.read("glp_compiles.csv")
blacklist = ["MIME"]
pc = SnoopCompile.format_userimg(data[end:-1:1,2], blacklist=blacklist)
SnoopCompile.write(Pkg.dir("GLPlot", "src", "glp_userimg.jl"), pc)


function Pump(; name, P = 1.0)
    @named oneport = OnePort()
    @unpack p = oneport
    ps = @parameters P=P
    eqs = [
        P ~ p
    ]
    extend(ODESystem(eqs, t, [], ps; name), oneport)
end


"""
    LagrangeDual

Lagrangian dual method for dual decomposition. This `mutable struct` constains:
    - `block_model::BlockModel` object
    - `var_to_index` mapping coupling variable to the index wrt the master problem
    - `masiter::Int` sets the maximum number of iterations
    - `tol::Float64` sets the relative tolerance for termination
"""

mutable struct LagrangeDual <: AbstractLagrangeDual
    block_model::BlockModel
    var_to_index::Dict{Tuple{Int,Any},Int} # maps coupling variable to the index wrt the master problem
    heuristics::Vector{Type}
    subsolve_time::Vector{Dict{Int,Float64}}
    subcomm_time::Vector{Float64}
    subobj_value::Vector{Float64}
    master_time::Vector{Float64}

    function LagrangeDual()
        LD = new()
        LD.block_model = BlockModel()
        LD.var_to_index = Dict()
        LD.heuristics = []
        LD.subsolve_time = []
        LD.subcomm_time = []
        LD.subobj_value = []
        LD.master_time = []

        return LD
    end
end

"""
Wrappers of the functions defined for `BlockModel`
"""

add_block_model!(LD::AbstractLagrangeDual, block_id::Integer, model::JuMP.Model) = add_block_model!(LD.block_model, block_id, model)
num_blocks(LD::AbstractLagrangeDual) = num_blocks(LD.block_model)
block_model(LD::AbstractLagrangeDual, block_id::Integer) = block_model(LD.block_model, block_id)
block_model(LD::AbstractLagrangeDual) = block_model(LD.block_model)
has_block_model(LD::AbstractLagrangeDual, block_id::Integer) = has_block_model(LD.block_model, block_id)
num_coupling_variables(LD::AbstractLagrangeDual) = num_coupling_variables(LD.block_model)
coupling_variables(LD::AbstractLagrangeDual) = coupling_variables(LD.block_model)

function set_coupling_variables!(LD::AbstractLagrangeDual, variables::Vector{CouplingVariableRef})
    set_coupling_variables!(LD.block_model, variables)
    variable_keys = [v.key for v in variables]
    # collect all coupling variables
    all_variable_keys = parallel.allcollect(variable_keys)
    set_variables_by_couple!(LD.block_model, all_variable_keys)
    LD.var_to_index = Dict((v.block_id,v.coupling_id) => i for (i,v) in enumerate(all_variable_keys))
end

dual_objective_value(LD::AbstractLagrangeDual) = dual_objective_value(LD.block_model)
dual_solution(LD::AbstractLagrangeDual) = dual_solution(LD.block_model)
primal_objective_value(LD::AbstractLagrangeDual) = primal_objective_value(LD.block_model)
primal_solution(LD::AbstractLagrangeDual) = primal_solution(LD.block_model)

"""
    run!

This runs the Lagrangian dual method for solving the block model.
"""
function run!(LD::AbstractLagrangeDual, LM::AbstractLagrangeMaster, initial_λ = nothing)

    # We assume that the block models are distributed.
    num_all_blocks = parallel.sum(num_blocks(LD))
    num_all_coupling_variables = parallel.sum(num_coupling_variables(LD))

    # initialize λ if it is nothing
    if isnothing(initial_λ)
        initial_λ = zeros(num_all_coupling_variables)
    end
    @assert length(initial_λ) == num_all_coupling_variables

    # check the validity of LagrangeDual
    if num_all_blocks <= 0 || num_all_coupling_variables == 0
        println("Invalid LagrangeDual structure.")
        return
    end

    function solveLagrangeDual(λ::Array{Float64,1})
        @assert length(λ) == num_all_coupling_variables

        # broadcast λ
        if parallel.is_root()
            parallel.bcast(λ)
        end

        # output
        objvals = Dict{Int,Float64}()
        subgrads = Dict{Int,SparseVector{Float64}}()
        subsolve_time = Dict{Int,Float64}()

        # Adjust block objective function
        for var in coupling_variables(LD)
            adjust_objective_function!(LD, var, λ[index_of_λ(LD, var)])
        end

        for (id,m) in block_model(LD)
            # Initialize subgradients
            subgrads[id] = sparsevec(Dict{Int,Float64}(), length(λ))

            # Solver the Lagrange dual
            stime = time()
            solve_sub_block!(m)
            subsolve_time[id] = time() - stime

            status = JuMP.termination_status(m)
            # @assert status in [MOI.OPTIMAL, MOI.LOCALLY_SOLVED]

            # We may want consider other statuses.
            if status in [MOI.OPTIMAL, MOI.LOCALLY_SOLVED]
                try
                    objvals[id] = -JuMP.dual_objective_value(m)
                catch e
                    objvals[id] = -JuMP.objective_value(m)
                end
            else
                @error "Unexpected solution status: $(status)"
            end
        end

        push!(LD.subsolve_time, subsolve_time)

        # Get subgradients
        for var in coupling_variables(LD)
            # @assert has_block_model(LD, var.key.block_id)
            subgrads[var.key.block_id][index_of_λ(LD, var)] = -JuMP.value(var.ref)
        end

        #run heuristics
        if length(LD.heuristics) > 0
            #get the values of all coupling variables
            opt_coupling_val = Dict{Int,SparseVector{Float64}}()
            coupling_ub = Dict{Int,SparseVector{Float64}}()
            coupling_lb = Dict{Int,SparseVector{Float64}}()
            for (id,m) in block_model(LD)
                opt_coupling_val[id] = sparsevec(Dict{Int,Float64}(), num_all_coupling_variables)
                coupling_ub[id] = sparsevec(Dict{Int,Float64}(), num_all_coupling_variables)
                coupling_lb[id] = sparsevec(Dict{Int,Float64}(), num_all_coupling_variables)
            end
            #get variable values and bounds
            for var in coupling_variables(LD)
                opt_coupling_val[var.key.block_id][index_of_λ(LD, var)] = JuMP.value(var.ref)
                if JuMP.has_lower_bound(var.ref)
                    coupling_lb[var.key.block_id][index_of_λ(LD, var)] = JuMP.lower_bound(var.ref)
                else
                    coupling_lb[var.key.block_id][index_of_λ(LD, var)] = - Inf
                end
                if JuMP.has_upper_bound(var.ref)
                    coupling_ub[var.key.block_id][index_of_λ(LD, var)] = JuMP.upper_bound(var.ref)
                else
                    coupling_ub[var.key.block_id][index_of_λ(LD, var)] = + Inf
                end
            end

            opt_coupling_val_combined = parallel.combine_dict(opt_coupling_val)
            coupling_ub_combined = parallel.combine_dict(coupling_ub)
            coupling_lb_combined = parallel.combine_dict(coupling_lb)

            for htype in LD.heuristics
                run!(htype, LD, opt_coupling_val_combined, coupling_ub_combined, coupling_lb_combined)
            end
        end

        # Reset objective coefficients
        for var in coupling_variables(LD)
            reset_objective_function!(LD, var, λ[index_of_λ(LD, var)])
        end

        parallel.barrier()
        comm_time = time()

        # Collect objvals, subgrads
        objvals_combined = parallel.combine_dict(objvals)
        objvals_vec = Vector{Float64}(undef, length(objvals_combined))
        if parallel.is_root()
            for (k,v) in objvals_combined
                objvals_vec[k] = v
            end
            push!(LD.subobj_value, sum(objvals_vec))
        end

        subgrads_combined = parallel.combine_dict(subgrads)

        if parallel.is_root()
            push!(LD.subcomm_time, time() - comm_time)
            # @printf("Subproblem sommunication time: %6.1f sec.\n", time() - comm_time)
        end

        return objvals_vec, subgrads_combined
    end

    if parallel.is_root()
        load!(LM, num_all_coupling_variables, num_all_blocks, solveLagrangeDual, initial_λ)

        # Add bounding constraints to the Lagrangian master
        add_constraints!(LD, LM)

        # This runs the bundle method.
        run!(LM)

        # Copy master solution time
        LD.master_time = get_times(LM)

        # get dual objective value
        LD.block_model.dual_bound = get_objective(LM)

        # get dual solution
        LD.block_model.dual_solution = get_solution(LM)

        # broadcast we are done.
        parallel.bcast(Float64[])
    else
        λ = parallel.bcast(nothing)
        while length(λ) > 0
            solveLagrangeDual(λ)
            λ = parallel.bcast(nothing)
        end
    end
end

"""
This adjusts the objective function of each Lagrangian subproblem.
"""
function adjust_objective_function!(LD::AbstractLagrangeDual, var::CouplingVariableRef, λ::Float64)
    @assert has_block_model(LD, var.key.block_id)
    affobj = objective_function(LD, var.key.block_id)
    @assert typeof(affobj) == AffExpr
    coef = haskey(affobj.terms, var.ref) ? affobj.terms[var.ref] + λ : λ
    JuMP.set_objective_coefficient(block_model(LD, var.key.block_id), var.ref, coef)
end

"""
This resets the objective function of each Lagrangian subproblem.
"""
function reset_objective_function!(LD::AbstractLagrangeDual, var::CouplingVariableRef, λ::Float64)
    @assert has_block_model(LD, var.key.block_id)
    affobj = objective_function(LD, var.key.block_id)
    @assert typeof(affobj) == AffExpr
    coef = haskey(affobj.terms, var.ref) ? affobj.terms[var.ref] - λ : -λ
    JuMP.set_objective_coefficient(block_model(LD, var.key.block_id), var.ref, coef)
end

"""
This wraps the steps to optimize a block problem.
"""
function solve_sub_block!(model::JuMP.Model)
    JuMP.optimize!(model)
    # FIXME (KK): why do we need to do this?
    # reoptimize!(model)
end

"""
This re-optimizes block models if not solved to local optimality
"""
function reoptimize!(model::JuMP.Model)
    solve_itr = 0
    while !(JuMP.termination_status(model) in [MOI.OPTIMAL, MOI.LOCALLY_SOLVED]) && solve_itr < 10
        JuMP.set_start_value.(all_variables(model), rand())
        JuMP.optimize!(model)
        solve_itr += 1
    end
end

"""
Wrappers of other functions
"""
objective_function(LD::AbstractLagrangeDual, block_id::Integer) = JuMP.objective_function(block_model(LD, block_id), QuadExpr).aff

index_of_λ(LD::AbstractLagrangeDual, var::CouplingVariableKey) = LD.var_to_index[var.block_id,var.coupling_id]
index_of_λ(LD::AbstractLagrangeDual, var::CouplingVariableRef) = index_of_λ(LD, var.key)

function write_times(LD::AbstractLagrangeDual; dir = ".")
    write_file!(LD.subsolve_time, "subsolve_time.txt", dir)
    write_file!(LD.master_time, "subcomm_time.txt", dir)
    write_file!(LD.master_time, "master_time.txt", dir)
end

function write_all(LD::AbstractLagrangeDual; dir = ".")
    write_times(LD, dir = dir)
    write_file!(LD.subobj_value, "subobj_value.txt", dir)
end

function get_solution!(LD::AbstractLagrangeDual, method::BM.AbstractMethod)
    LD.block_model.dual_solution = copy(BM.get_solution(method))
end

include("PrimalHeuristics.jl")