text:
#-------------------------------------------------------------------
#* EMSO Model Library (EML) Copyright (C) 2004 - 2007 ALSOC.
#*
#* This LIBRARY is free software; you can distribute it and/or modify
#* it under the therms of the ALSOC FREE LICENSE as available at
#* http://www.enq.ufrgs.br/alsoc.
#*
#* EMSO Copyright (C) 2004 - 2007 ALSOC, original code
#* from http://www.rps.eng.br Copyright (C) 2002-2004.
#* All rights reserved.
#*
#* EMSO is distributed under the therms of the ALSOC LICENSE as
#* available at http://www.enq.ufrgs.br/alsoc.
#*--------------------------------------------------------------------
#* Author: Gerson Balbueno Bicca
#* $Id: HairpinIncr.mso								$
#*------------------------------------------------------------------
type HairpinIncr
	HairpinIncr()=begin
		PP=outers.PP
		NComp=outers.NComp
		new(
			DanaPlugin(Dict{Symbol,Any}(
				:Brief=>"External Physical Properties",
				:Type=>"PP"
			)),
			DanaInteger(Dict{Symbol,Any}(
				:Brief=>"Number of Components"
			)),
			constant(Dict{Symbol,Any}(
				:Brief=>"Pi Number",
				:Default=>3.14159265,
				:Symbol=>"\\pi"
			)),
			DanaInteger(Dict{Symbol,Any}(
				:Brief=>"Number of zones",
				:Default=>2
			)),
			DanaInteger(Dict{Symbol,Any}(
				:Brief=>"Number of incremental points",
				:Default=>3
			)),
			length(Dict{Symbol,Any}(
				:Brief=>"Outside Diameter of Inner Pipe",
				:Lower=>1e-6
			)),
			length(Dict{Symbol,Any}(
				:Brief=>"Inside Diameter of Inner Pipe",
				:Lower=>1e-10
			)),
			length(Dict{Symbol,Any}(
				:Brief=>"Inside Diameter of Outer pipe",
				:Lower=>1e-10
			)),
			length(Dict{Symbol,Any}(
				:Brief=>"Effective Tube Length of one segment of Pipe",
				:Lower=>0.1,
				:Symbol=>"L_{pipe}"
			)),
			conductivity(Dict{Symbol,Any}(
				:Brief=>"Tube Wall Material Thermal Conductivity",
				:Default=>1.0,
				:Symbol=>"K_{wall}"
			)),
			positive(Dict{Symbol,Any}(
				:Brief=>"Inside Fouling Resistance",
				:Unit=>"m^2*K/kW",
				:Default=>1e-6,
				:Lower=>0
			)),
			positive(Dict{Symbol,Any}(
				:Brief=>"Outside Fouling Resistance",
				:Unit=>"m^2*K/kW",
				:Default=>1e-6,
				:Lower=>0
			)),
			length(Dict{Symbol,Any}(
				:Brief=>"Inner Side Outlet Nozzle Diameter",
				:Default=>0.036,
				:Lower=>1e-6
			)),
			length(Dict{Symbol,Any}(
				:Brief=>"Inner Side Inlet Nozzle Diameter",
				:Default=>0.036,
				:Lower=>1e-6
			)),
			length(Dict{Symbol,Any}(
				:Brief=>"Outer Side Outlet Nozzle Diameter",
				:Default=>0.036,
				:Lower=>1e-6
			)),
			length(Dict{Symbol,Any}(
				:Brief=>"Outer Side Inlet Nozzle Diameter",
				:Default=>0.036,
				:Lower=>1e-6
			)),
			positive(Dict{Symbol,Any}(
				:Brief=>"Inner Side Inlet Nozzle Pressure Loss Coeff",
				:Default=>1.1
			)),
			positive(Dict{Symbol,Any}(
				:Brief=>"Inner Side Outlet Nozzle Pressure Loss Coeff",
				:Default=>0.7
			)),
			positive(Dict{Symbol,Any}(
				:Brief=>"Outer Side Inlet Nozzle Pressure Loss Coeff",
				:Default=>1.1
			)),
			positive(Dict{Symbol,Any}(
				:Brief=>"Outer Side Outlet Nozzle Pressure Loss Coeff",
				:Default=>0.7
			)),
			Summary_Hairpin(Dict{Symbol,Any}(
				:Brief=>"Results for The Whole Heat Exchanger"
			)),
			UpperPipe_basic(Dict{Symbol,Any}(
				:Brief=>"Upper Pipe Results"
			)),
			LowerPipe_basic(Dict{Symbol,Any}(
				:Brief=>"Lower Pipe Results"
			)),
			[
				:(Summary.A = LowerPipe.Details.A+UpperPipe.Details.A),
				:(Summary.Qtotal = LowerPipe.Details.Qtotal+UpperPipe.Details.Qtotal),
				:(Summary.Inner.Pdrop = Summary.Inner.Pdnozzle_in+Summary.Inner.Pdnozzle_out+LowerPipe.Inner.PressureDrop.Pd_fric(1) +UpperPipe.Inner.PressureDrop.Pd_fric(Npoints)),
				:(Summary.Outer.Pdrop = Summary.Outer.Pdnozzle_in+Summary.Outer.Pdnozzle_out+LowerPipe.Outer.PressureDrop.Pd_fric(Npoints) +UpperPipe.Outer.PressureDrop.Pd_fric(1)),
				:(UpperPipe.Outer.PressureDrop.Plocal([1:N]) = UpperPipe.Outer.PressureDrop.Plocal(Npoints) - UpperPipe.Outer.PressureDrop.Pd_fric([1:N])),
				:(UpperPipe.Inner.PressureDrop.Plocal([1:N]+1) = UpperPipe.Inner.PressureDrop.Plocal(1) - UpperPipe.Inner.PressureDrop.Pd_fric([1:N]+1)),
				:(LowerPipe.Inner.PressureDrop.Plocal([1:N]) = LowerPipe.Inner.PressureDrop.Plocal(Npoints) - LowerPipe.Inner.PressureDrop.Pd_fric([1:N])),
				:(LowerPipe.Outer.PressureDrop.Plocal([1:N]+1) = LowerPipe.Outer.PressureDrop.Plocal(1) - LowerPipe.Outer.PressureDrop.Pd_fric([1:N]+1)),
				:(Summary.Inner.Vnozzle_in = UpperPipe.Inner.Properties.Inlet.Fw/(UpperPipe.Inner.Properties.Inlet.rho*(0.25*Pi*Dinozzle_Inner^2))),
				:(Summary.Inner.Vnozzle_out = LowerPipe.Inner.Properties.Outlet.Fw/(LowerPipe.Inner.Properties.Outlet.rho*(0.25*Pi*Donozzle_Inner^2))),
				:(Summary.Outer.Vnozzle_in = LowerPipe.Outer.Properties.Inlet.Fw/(LowerPipe.Outer.Properties.Inlet.rho*(0.25*Pi*Dinozzle_Outer^2))),
				:(Summary.Outer.Vnozzle_out = UpperPipe.Outer.Properties.Outlet.Fw/(UpperPipe.Outer.Properties.Outlet.rho*(0.25*Pi*Donozzle_Outer^2))),
				:(Summary.Inner.Pdnozzle_in = 0.5*InnerKinlet*UpperPipe.Inner.Properties.Inlet.rho*Summary.Inner.Vnozzle_in^2),
				:(Summary.Inner.Pdnozzle_out = 0.5*InnerKoutlet*LowerPipe.Inner.Properties.Outlet.rho*Summary.Inner.Vnozzle_out^2),
				:(Summary.Outer.Pdnozzle_in = 0.5*OuterKinlet*LowerPipe.Outer.Properties.Inlet.rho*Summary.Outer.Vnozzle_in^2),
				:(Summary.Outer.Pdnozzle_out = 0.5*OuterKoutlet*UpperPipe.Outer.Properties.Outlet.rho*Summary.Outer.Vnozzle_out^2),
				:(Summary.Inner.RVsquare_in = UpperPipe.Inner.Properties.Inlet.rho*(Summary.Inner.Vnozzle_in)^2),
				:(Summary.Inner.RVsquare_out = LowerPipe.Inner.Properties.Outlet.rho*(Summary.Inner.Vnozzle_out)^2),
				:(Summary.Outer.RVsquare_in = LowerPipe.Outer.Properties.Inlet.rho*(Summary.Outer.Vnozzle_in)^2),
				:(Summary.Outer.RVsquare_out = UpperPipe.Outer.Properties.Outlet.rho*(Summary.Outer.Vnozzle_out)^2),
				:(Summary.Outer.hcoeff = sum(UpperPipe.Outer.HeatTransfer.hcoeff+LowerPipe.Outer.HeatTransfer.hcoeff)/(2*N)),
				:(Summary.Inner.hcoeff = sum(UpperPipe.Inner.HeatTransfer.hcoeff+LowerPipe.Inner.HeatTransfer.hcoeff)/(2*N)),
			],
			[
				"Total Exchange Surface Area","Total Duty","Total Pressure Drop Inner Side","Total Pressure Drop Outer Side","Outer Pipe Local Pressure","Inner Pipe Local Pressure","Inner Pipe Local Pressure","Outer Pipe Local Pressure","Velocity Inner Side Inlet Nozzle","Velocity Inner Side Outlet Nozzle","Velocity Outer Side Inlet Nozzle","Velocity Outer Side Outlet Nozzle","Pressure Drop Inner Side Inlet Nozzle","Pressure Drop Inner Side Outlet Nozzle","Pressure Drop Outer Side Inlet Nozzle","Pressure Drop Outer Side Outlet Nozzle","Inner Side Inlet Nozzle rho-V^2","Inner Side Outlet Nozzle rho-V^2","Outer Side Inlet Nozzle rho-V^2","Outer Side Outlet Nozzle rho-V^2","Average Film Coefficient Outer Side","Average Film Coefficient Inner Side",
			],
			[:PP,:NComp,:Pi,:N,:Npoints,:DoInner,:DiInner,:DiOuter,:Lpipe,:Kwall,:Rfi,:Rfo,:Donozzle_Inner,:Dinozzle_Inner,:Donozzle_Outer,:Dinozzle_Outer,:InnerKinlet,:InnerKoutlet,:OuterKinlet,:OuterKoutlet,],
			[:Summary,:UpperPipe,:LowerPipe,]
		)
	end
	PP::DanaPlugin
	NComp::DanaInteger
	Pi::constant
	N::DanaInteger
	Npoints::DanaInteger
	DoInner::length
	DiInner::length
	DiOuter::length
	Lpipe::length
	Kwall::conductivity
	Rfi::positive
	Rfo::positive
	Donozzle_Inner::length
	Dinozzle_Inner::length
	Donozzle_Outer::length
	Dinozzle_Outer::length
	InnerKinlet::positive
	InnerKoutlet::positive
	OuterKinlet::positive
	OuterKoutlet::positive
	Summary::Summary_Hairpin
	UpperPipe::UpperPipe_basic
	LowerPipe::LowerPipe_basic
	equations::Array{Expr,1}
	equationNames::Array{String,1}
	parameters::Array{Symbol,1}
	variables::Array{Symbol,1}
	attributes::Dict{Symbol,Any}
end
export HairpinIncr
function set(in::HairpinIncr)
	#"Pi Number"
	Pi = 3.14159265
	 #"Number of incremental points"
	Npoints = N+1

end
function setEquationFlow(in::HairpinIncr)
	addEquation(1)
	addEquation(2)
	addEquation(3)
	addEquation(4)
	# FIXME: NOZZLE PRESSURE DROP MUST BE ADDED
	addEquation(5)
	# FIXME: NOZZLE PRESSURE DROP MUST BE ADDED
	addEquation(6)
	# FIXME: NOZZLE PRESSURE DROP MUST BE ADDED
	addEquation(7)
	# FIXME: NOZZLE PRESSURE DROP MUST BE ADDED
	addEquation(8)
	addEquation(9)
	addEquation(10)
	addEquation(11)
	addEquation(12)
	addEquation(13)
	addEquation(14)
	addEquation(15)
	addEquation(16)
	addEquation(17)
	addEquation(18)
	addEquation(19)
	addEquation(20)
	addEquation(21)
	addEquation(22)
end
function atributes(in::HairpinIncr,_::Dict{Symbol,Any})
	fields::Dict{Symbol,Any}=Dict{Symbol,Any}()
	fields[:Pallete]=true
	fields[:Icon]="icon/hairpin"
	fields[:Brief]="Incremental Hairpin Heat Exchanger. "
	fields[:Info]="Incremental approach for Hairpin heat exchanger.
OBS: LEFT = 0 e N = N, RIGTH= L e N=1
"
	drive!(fields,_)
	return fields
end
HairpinIncr(_::Dict{Symbol,Any})=begin
	newModel=HairpinIncr()
	newModel.attributes=atributes(newModel,_)
	newModel
end


text:
# site_counting.jl
# ================
#
# Site counting framework for biological sequences.
#
# This file is a part of BioJulia.
# License is MIT: https://github.com/BioJulia/BioSequences.jl/blob/master/LICENSE.md

# Site types
# ----------

abstract type Site end
abstract type Position <: Site end

"""
A `Certain` site describes a site where both of two aligned sites are not an
ambiguity symbol or a gap.
"""
struct Certain <: Position end

"""
An `Gap` site describes a site where either of two aligned sites are a
gap symbol '-'.
"""
struct Gap <: Position end

"""
An `Ambiguous` site describes a site where either of two aligned sites are an
ambiguity symbol.
"""
struct Ambiguous <: Position end

"""
A `Match` site describes a site where two aligned nucleotides are the
same biological symbol.
"""
struct Match <: Position end

"""
A `Mismatch` site describes a site where two aligned nucleotides are not the
same biological symbol.
"""
struct Mismatch <: Position end

# Bitparallel counting
# --------------------
include("count_sites_bitpar.jl")

# Now we overload some internal methods for each S <: Site.
# So as different types of site can be counted in bit-parallel
# manner.

for A in (DNAAlphabet, RNAAlphabet)
    @eval begin

        # Gaps
        @inline function bp_chunk_count(::Type{Gap}, ::Type{$A{4}}, x::UInt64)
            return count_zero_nibbles(x)
        end

        @inline function bp_chunk_count(::Type{Gap}, ::Type{$A{4}}, a::UInt64, b::UInt64)
            # Count the gaps in a, count the gaps in b, subtract the number of shared gaps.
            return count_zero_nibbles(a) + count_zero_nibbles(b) - count_zero_nibbles(a | b)
        end

        # Certain
        @inline function bp_chunk_count(::Type{Certain}, ::Type{$A{4}}, x::UInt64)
            x = enumerate_nibbles(x)
            x = x ⊻ 0x1111111111111111
            return count_zero_nibbles(x)
        end

        @inline function bp_chunk_count(::Type{Certain}, ::Type{$A{4}}, a::UInt64, b::UInt64)
            x = enumerate_nibbles(a) ⊻ 0x1111111111111111
            y = enumerate_nibbles(b) ⊻ 0x1111111111111111
            return count_zero_nibbles(x | y)
        end

        # Ambiguous
        @inline function bp_chunk_count(::Type{Ambiguous}, ::Type{$A{4}}, x::UInt64)
            return count_nonzero_nibbles(enumerate_nibbles(x) & 0xEEEEEEEEEEEEEEEE)
        end

        @inline function bp_chunk_count(::Type{Ambiguous}, ::Type{$A{4}}, a::UInt64, b::UInt64)
            return count_nonzero_nibbles((enumerate_nibbles(a) | enumerate_nibbles(b)) & 0xEEEEEEEEEEEEEEEE)
        end

        # Match
        @inline function bp_chunk_count(::Type{Match}, ::Type{$A{4}}, a::UInt64, b::UInt64)
            return count_zero_nibbles(a ⊻ b)
        end

        @inline function bp_chunk_count(::Type{Match}, ::Type{$A{2}}, a::UInt64, b::UInt64)
            return count_zero_bitpairs(a ⊻ b)
        end

        # Mismatch
        @inline function bp_chunk_count(::Type{Mismatch}, ::Type{$A{4}}, a::UInt64, b::UInt64)
            return count_nonzero_nibbles(a ⊻ b)
        end

        @inline function bp_chunk_count(::Type{Mismatch}, ::Type{$A{2}}, a::UInt64, b::UInt64)
            return count_nonzero_bitpairs(a ⊻ b)
        end
    end
end

for s in (Match, Gap)
    @eval bp_correct_emptyspace(::Type{$s}, ::Type{A}) where {A<:NucAlphs} = true
end

# Specific Base.count methods
# ---------------------------

# In a general case, go to the bit-parallel counting method.
@inline Base.count(::Type{P}, a::BioSequence{A}, b::BioSequence{A}) where {P<:Position,A<:NucAlphs} = bitpar_counter(P, a, b)

# Some specific edge cases...
for A in (DNAAlphabet, RNAAlphabet)

    # Specific count methods for some edge cases regarding counting sites in
    # 2 bit encoded DNA and RNA sequences.
    seqtype = BioSequence{A{2}}
    @eval begin
        @inline function Base.count(::Type{Certain}, a::$seqtype, b::$seqtype)
            return min(length(a), length(b))
        end
        @inline Base.count(::Type{Gap}, a::$seqtype, b::$seqtype) = 0
        @inline Base.count(::Type{Ambiguous}, a::$seqtype, b::$seqtype) = 0
    end
end

# Specific Base.count sliding window methods
# ------------------------------------------

function Base.count(::Type{P}, a::BioSequence{A}, b::BioSequence{A}, width::Int, step::Int) where {P<:Position,A<:NucAlphs}
    len = min(length(a), length(b))
    ritr = StepRange(width, step, len)
    width -= 1
    results = Vector{IntervalValue{Int,bp_counter_type(P, A)}}(undef, length(ritr))
    r = 1
    @inbounds for i in ritr
        idx = (i - width):i
        results[r] = IntervalValue(first(idx), last(idx), count(P, a[idx], b[idx]))
        r += 1
    end
    return results
end

# Specific count_pairwise methods
# -------------------------------

@inline diag_val(::Type{Int}) = Int(0)

function count_pairwise(::Type{P}, seqs::Vararg{BioSequence{A},N}) where {P<:Position,A<:NucAlphs,N}
    @assert N >= 2 "At least two sequences are required."
    counts = Matrix{bp_counter_type(P, A)}(undef, (N, N))
    for i in 1:N
        counts[i,i] = diag_val(eltype(counts))
        for j in (i+1):N
            counts[i,j] = counts[j,i] = count(P, seqs[i], seqs[j])
        end
    end
    return counts
end

# General, non-specific Base.count methods
# ----------------------------------------

"""
    Base.count{S<:Site}(::Type{S}, a::BioSequence, b::BioSequence)

Count the number of sites of type `S`, between two sequences a and b.
"""
@inline function Base.count(::Type{S}, a::BioSequence, b::BioSequence) where {S<:Site}
    seqs = promote(a, b)
    return count(S, seqs...)
end

"""
    Base.count{S<:Site}(::Type{S}, a::BioSequence, b::BioSequence, width::Int, step::Int)

Count the number of sites of type `S`, between two sequences a and b, using a
sliding window of a given `width`, and `step`.
The `width` and `step` must be provided as number of base pairs.
"""
function Base.count(::Type{S}, a::BioSequence, b::BioSequence, width::Int, step::Int) where {S<:Site}
    seqs = promote(a, b)
    return count(S, seqs..., width, step)
end

"""
    count_pairwise{S<:Site,N}(::Type{S}, seqs::Vararg{BioSequence,N})

Count the number of sites of type `S`, between each possible pair of the sequences
provided (`seqs`).

This returns a symmetric matrix contining the output for each pairwise
operation.
"""
function count_pairwise(::Type{S}, seqs::Vararg{BioSequence,N}) where {S<:Site,N}
    seqs = promote(seqs...)
    return count_pairwise(S, seqs...)
end

ext: jl


text:
include("common.jl")

function setup_plots()
    # Setup environment for making automated plots
    ENV["GKS_ENCODING"] = "utf8"
    ENV["GKSwstype"]    = "100"
    ENV["PLOTS_TEST"]   = "true"

    gr()
    default(size=tuple(Int.(ceil.(0.75 .* [600, 400]))...),
            guidefontsize=10, grid=false)
end

function dump_table_convergence(fn::AbstractString; kwargs...)
    open(fn, "w") do io
        dump_table_convergence(io; kwargs...)
    end
end
function dump_table_convergence(; kwargs...)
    dump_table_convergence(stdout; kwargs...)
end
function dump_table_convergence(io::IO;
                                selected=["Fe2CrGa"], allscfkeys=nothing,
                                tol=1e-10, ncthresh=100, labels=Dict{String,String}(),
                                precons=nothing,
                               )
    separator_after = String[]
    for (i, case) in enumerate(selected)
        if i+1 <= length(selected) && selected[i+1] == :separator
            push!(separator_after, case)
        end
    end
    selected = filter!(a -> a != :separator, selected)

    if isnothing(allscfkeys)
        allscfkeys = union(Set.(default_scfkeys(case) for case in selected)...)
    end
    allscfkeys = sort([allscfkeys...], by=parse_scfkey_label)
    adaptkeys  = filter(startswith("adaptive"), allscfkeys)
    fixedkeys  = filter(startswith("standard"), allscfkeys)
    allkeys    = vcat(fixedkeys, adaptkeys)
    n_adaptive = length(adaptkeys)
    n_standard = length(fixedkeys)

    if length(adaptkeys) == 1
        adaptheader = "adaptive"
    else
        adaptheader = "adaptive damping"
    end

    println(io, raw"\begin{tabular}{ll|*{" * string(n_standard) * raw"}{c}|*{"
            * string(n_adaptive) *  raw"}{c}}")
    println(io, raw"System&Precond." *
                raw"&\multicolumn{" * string(n_standard) * raw"}{c|}{fixed damping $\alpha$}" *
                raw"&\multicolumn{" * string(n_adaptive) * raw"}{c}{" * adaptheader * raw"}" *
                raw"\\\\")

    print(io, "&")
    for scfkey in fixedkeys
      _, alpha = parse_scfkey_label(scfkey)
      print(io, "& $alpha ")
    end
    if length(adaptkeys) == 1
        print(io, "&damping")
    else
        for scfkey in adaptkeys
          _, alpha = parse_scfkey_label(scfkey)
          print(io, raw"& $\atmin = " * string(alpha) * raw"$ ")
        end
    end
    println(io, "\\\\")
    println(io, raw"\hline")

    for case in selected
        directory = joinpath(@__DIR__, case)
        collectkeys = default_scfkeys(case)
        mintotals = collect_minimal_total_energy(directory, collectkeys)

        min_n_diags_fixed = typemax(Int)
        n_diags = Vector{Union{Missing, Int}}()
        for scfkey in allkeys
            if isempty(mintotals)
                push!(n_diags, typemax(Int))
                continue
            end

            method, _ = parse_scfkey_label(scfkey)
            jsonfile = joinpath(directory, "$scfkey.scfres.json")
            if !isfile(jsonfile)
                push!(n_diags, missing)
                continue
            end

            data = open(JSON.parse, jsonfile)
            energies = data["energies"]
            mintrial  = minimum(energies[max(end-4, 1):end])
            bestmatch = sortperm(abs.(mintrial .- mintotals))[1]
            values = abs.(energies .- mintotals[bestmatch])
            n_conv = findfirst(values .< tol)
            if isnothing(n_conv) && values[end] < 3tol
                n_conv = length(values)
            end
            if case == "Fe2CrGa" && scfkey == "standard_09"
                # Dirty hack, because I manually verified this to be (barely) not converged
                n_conv = nothing
            end
            if isnothing(n_conv)
                push!(n_diags, typemax(Int))
            else
                # unit_cell_volume = abs(det(hcat(data["lattice"]...)))
                # dvol = unit_cell_volume / prod(data["fft_size"])
                # println(case, "   ", scfkey, "   ", data["residuals"][1] * sqrt(dvol),
                #               "   ", data["residuals"][n_conv] * sqrt(dvol),
                #               "   ", data["residuals"][n_conv]/data["residuals"][1])

                n_diag = sum(data["ndiag"][1:n_conv])
                push!(n_diags, n_diag)
                if method == "fixed"
                    min_n_diags_fixed = min(min_n_diags_fixed, n_diag)
                end
            end
        end

        if case in keys(labels)
            print(io, labels[case], raw"  &  ")
        else
            print(io, raw"\ce{", case, raw"}  &  ")
        end
        if case in keys(precons)
            print(io, precons[case], raw"  &  ")
        else
            print(io, raw"--  &  ")
        end

        for (scfkey, n_diag) in zip(allkeys, n_diags)
            method, _ = parse_scfkey_label(scfkey)
            endchar = scfkey == allkeys[end] ? raw" \\\\" : " & "

            if ismissing(n_diag)
                print(io, "--", endchar)
            elseif n_diag > ncthresh
                print(io, raw"$\times$", endchar)
            elseif method == "fixed" && n_diag == min_n_diags_fixed
                print(io, raw"\textbf{", n_diag, "}", endchar)
            else
                print(io, n_diag, endchar)
            end
        end
        println(io)
        if case in separator_after
            println(io, raw"\hline")
        end
    end
    println(io, raw"\end{tabular}")
end

function dump_latex_tables()
    file = "convergence_table.tex"
    allscfkeys = ["adaptive"]
    for i in 1:9
        push!(allscfkeys, "standard_0$i")
    end
    push!(allscfkeys, "standard_10")
    selected = ["Al_nodiis_Kerker", "Al_nodiis", "Al_Kerker", "Al", "AlVac_Kerker", "AlVac",
                :separator,
                "GaAs",
                :separator,
                "CoFeMnGa", "Fe2CrGa", "Fe2MnAl", "FeNiF6", "Mn2RuGa", "Mn3Si", "Mn3Si_AFM",
                :separator,
                "Cr19", "WFe"
               ]
    labels = Dict(
        "Al_nodiis"        => (raw"\ce{Al8} supercell"),
        "Al_nodiis_Kerker" => (raw"\ce{Al8} supercell"),
        "Al"               => (raw"\ce{Al40} supercell"),
        "Al_Kerker"        => (raw"\ce{Al40} supercell"),
        "GaAs"             => (raw"\ce{Ga20As20} supercell"),
        "AlVac_Kerker"     => (raw"\ce{Al40} surface"),
        "AlVac"            => (raw"\ce{Al40} surface"),
        "Mn3Si_AFM"        => (raw"\ce{Mn3Si}$^\text{AFM}$"),
        "Cr19"             => (raw"\ce{Cr19} defect"),
        "WFe"              => (raw"\ce{Fe28W8} multilayer\cite{Marks2021}"),
    )
    precons = Dict(
        "Al_nodiis"        => raw"None$^\dagger$",
        "Al_nodiis_Kerker" => raw"Kerker$^\dagger$",
        "Al_Kerker"        => "Kerker",
        "Al"               => "None",
        "GaAs"             => "None",
        "AlVac"            => "None",
        "AlVac_Kerker"     => "Kerker",
        "CoFeMnGa"         => "Kerker",
        "Fe2CrGa"          => "Kerker",
        "Fe2MnAl"          => "Kerker",
        "FeNiF6"           => "Kerker",
        "Mn2RuGa"          => "Kerker",
        "Mn3Si"            => "Kerker",
        "Mn3Si_AFM"        => "Kerker",
        "Cr19"             => "Kerker",
        "WFe"              => "Kerker",
    )
    dump_table_convergence(file; selected, labels, precons, allscfkeys)
end

function purify_maxcond(jsonfile, mintotals)
    energies = open(JSON.parse, jsonfile, "r")["energies"]
    ndiags   = open(JSON.parse, jsonfile, "r")["ndiag"]

    @assert length(mintotals) == 3
    minenes = [abs.(energies .- mintotals[1]),
               abs.(energies .- mintotals[2]),
               abs.(energies .- mintotals[3])]
    ibest = sortperm(minimum.(minenes))[1]
    values = minenes[ibest]
    idx = findfirst(values .< 1e-16)
    if !isnothing(idx)
        values = values[1:idx-1]
        ndiags = ndiags[1:idx-1]
    end

    (; values, cumdiags=cumsum(ndiags))
end

function plot_maxcond(; iteration=80, tol=1e-10, legend=:bottomleft, label=nothing)
    mintotals = [
        -333.3052730117624,
        -333.30261995802,
        -333.30262027079226,
    ]

    errors   = Dict{String, Vector{Float64}}()
    cumdiags = Dict{String, Vector{Int}}()
    for maxcond in (2, 4, 6, 8, 10)
        key = @sprintf "%02i" maxcond
        jsonfile = joinpath(@__DIR__, "Fe2CrGa_maxcond$key/standard_02.scfres.json")
        dlabel = "maxcond = 1e$maxcond"
        pf = purify_maxcond(jsonfile, mintotals)
        errors[dlabel]   = pf.values
        cumdiags[dlabel] = pf.cumdiags
    end

    for key in ("1e2", "1e4", "1e6", "1e8", "1e10")
        jsonfile = joinpath(@__DIR__, "Fe2CrGa_40guess/maxcond_$key.scfres.json")
        dlabel = "maxcond = $key (restart)"
        pf = purify_maxcond(jsonfile, mintotals)
        errors[dlabel]   = pf.values
        cumdiags[dlabel] = 40 .+ pf.cumdiags
    end

    titleargs = ()
    if !isnothing(label)
        titleargs = (title=label, titlelocation=:left, titlefontsize=10)
    end

    ymax = 0
    maxiter = 0
    p = plot(; yaxis=:log, legend, titleargs...)
    for (i, maxcond) in enumerate([2, 4, 6, 8, 10])
        dlabel = "maxcond = 1e$maxcond"
        maxiter = max(maxiter, length(errors[dlabel]))
        ymax    = max(ymax,    maximum(errors[dlabel]))

        plot!(p, cumdiags[dlabel], errors[dlabel];
              label=dlabel, color=i, mark=:x, linestyle=:solid, markerstrokecolor=i,
              lw=1.5, markersize=4)

        rlabel = dlabel * " (restart)"
        if rlabel in keys(errors)
            plot!(p, cumdiags[rlabel], errors[rlabel];
                  label=rlabel, color=i, mark=:+, linestyle=:dot, markerstrokecolor=i,
                  lw=1.5, markersize=4)
        end
    end
    ylims!(p, (tol, 2 * ymax))
    ylabel!(p, "Total energy absolute error")

    xlims!(p, (0, min(iteration, maxiter)))
    xlabel!(p, "Number of Hamiltonian diagonalizations")

    p
end

function default_scfkeys(case)
    directory  = joinpath(@__DIR__, case)
    [key[1:end-12] for key in readdir(directory)
     if (    endswith(key, ".scfres.json")
         && !occursin("adaptive_slowstart", key)
         && !occursin("espresso", key)
         && !occursin("density", key)
         && !occursin("nlsolve", key)
         && !occursin("highprec", key))
    ]
end

function plot_convergence(case; iteration=60, tol=1e-10,
                          scfkeys=default_scfkeys(case),
                          colormap=Dict{Float64,Int}(),
                          legend=:topright, label=nothing,
                          mintotals=nothing, showmins=true, show_αmin=true,
                          show_damping=false)
    # First pass: For DFTK runs collect mimimal total energy
    directory = joinpath(@__DIR__, case)
    collectkeys = default_scfkeys(case)
    if isnothing(mintotals)
        mintotals = collect_minimal_total_energy(directory, collectkeys)
    end
    isempty(mintotals) && return plot()

    errors   = Dict{String, Vector{Float64}}()
    cumdiags = Dict{String, Vector{Int}}()
    dampings = Dict{String, Vector{Float64}}()
    whichmin = Dict{String, Int}()
    for scfkey in sort(scfkeys)
        method, _ = parse_scfkey_label(scfkey)
        method in ("espresso", "highprec") && error(
            "espresso or highprec not supported."
        )

        resultfile = joinpath(directory, scfkey * ".scfres.json")
        isfile(resultfile) || error("Invalid scfkey $scfkey")

        energies  = open(JSON.parse, resultfile, "r")["energies"]
        ndiags    = open(JSON.parse, resultfile, "r")["ndiag"]
        damps     = open(JSON.parse, resultfile, "r")["α"]

        energies === nothing && continue
        isempty(energies) && continue
        # Find best matching minimum
        mintrial  = minimum(energies[max(end-4, 1):end])
        bestmatch = sortperm(abs.(mintrial .- mintotals))[1]
        values = abs.(energies .- mintotals[bestmatch])
        idx = findfirst(values .< 1e-16)
        if !isnothing(idx)
            values = values[1:idx-1]
            ndiags = ndiags[1:idx-1]
            damps  = damps[1:idx-1]
        end
        isempty(values) && continue

        whichmin[scfkey] = bestmatch
        errors[scfkey]   = values
        dampings[scfkey] = damps
        cumdiags[scfkey] = cumsum(ndiags)
    end
    isempty(errors) && return plot()

    titleargs = ()
    if !isnothing(label)
        titleargs = (title=label, titlelocation=:left, titlefontsize=10)
    end

    ymax = 0
    maxiter = 0
    p = plot(; yaxis=:log, legend, titleargs...)
    for (i, scfkey) in enumerate(sort(collect(keys(errors)), by=parse_scfkey_label))
        method, alpha = parse_scfkey_label(scfkey)
        if method == "adaptive" && show_αmin
            label = method * " (α̃ₘᵢₙ = $alpha)"
        elseif method == "adaptive"
            label = "adaptive"
        elseif method == "fixed"
            label = method * " (α = $alpha)"
        else
            error("Unknown method: $method")
        end
        if length(mintotals) > 1 && showmins
            label *= " (min $(whichmin[scfkey]))"
        end
        maxiter = max(maxiter, length(errors[scfkey]))
        ymax    = max(ymax,    maximum(errors[scfkey]))

        if alpha in keys(colormap)
            color = colormap[alpha]
        else
            color = isempty(colormap) ? 1 : 1 + maximum(values(colormap))
            colormap[alpha] = color
        end
        if method == "adaptive"
            mark = :o
        elseif method == "fixed"
            mark = :x
        else
            error("Unknown method")
        end
        linestyle = :solid

        plot!(p, cumdiags[scfkey], errors[scfkey];
              label, color, mark, linestyle, markerstrokecolor=color,
              lw=1.5, markersize=4)
    end
    ylims!(p, (tol, 2 * ymax))
    ylabel!(p, "Total energy absolute error")

    if show_damping
        q = twinx()
        for (i, scfkey) in enumerate(sort(collect(keys(errors)), by=parse_scfkey_label))
            method, alpha = parse_scfkey_label(scfkey)
            method == "adaptive" || continue
            (alpha in keys(colormap)) || continue
            plot!(q, cumdiags[scfkey], abs.(dampings[scfkey]), mark=:+,
                  color=colormap[alpha], label="", linestyle=:dot)
        end
        ylims!(q, (0.0, 1.0))
        ylabel!(q, "Damping parameter")
    end

    xlims!(p, (0, min(iteration, maxiter)))
    xlabel!(p, "Number of Hamiltonian diagonalizations")

    p
end

function dump_plots()
    setup_plots()

    let kwargs = ()
        case  = "Al_nodiis_Kerker"
        label = "(b) Al₈ supercell (Kerker, no Anderson)"
        scfkeys = ["standard_02", "standard_05", "standard_09",
                   "adaptive",    "adaptive_05", "adaptive_09"]
        colormap = Dict{Float64,Int}()
        p = plot_convergence(case; colormap, scfkeys, label, kwargs...)
        ylims!(p, (1e-10, 10))
        savefig(p, case * ".pdf")

        case  = "Al_nodiis"
        label = "(a) Al₈ supercell (no preconditioner, no Anderson)"
        scfkeys = ["standard_01", "standard_02", "standard_03",
                   "adaptive_01", "adaptive_03", "adaptive_05"]
        p = plot_convergence(case; colormap, scfkeys, label, kwargs...)
        ylims!(p, (1e-10, 10))
        savefig(p, case * ".pdf")
    end

    let kwargs = ()
        mintotals = [-333.3052730117624, -333.30261995802, -333.30262027079226]
        case  = "Fe2CrGa"
        label = "(b) Fe₂CrGa (Kerker, Anderson)"
        colormap = Dict{Float64,Int}()
        scfkeys = ["standard_02", "standard_03", "standard_04", "standard_05",
                   "adaptive",    "adaptive_03", "adaptive_04", "adaptive_05"]
        p = plot_convergence(case; colormap, scfkeys, label, showmins=false, mintotals, kwargs...)
        savefig(p, case * ".pdf")

        case  = "Fe2CrGa_nodiis"
        label = "(a) Fe₂CrGa (Kerker, no Anderson)"
        scfkeys = ["standard_005", "standard_010", "standard_015",
                   "adaptive_01", "adaptive", "adaptive_03"]
        p = plot_convergence(case; colormap, scfkeys, label,
                             legend=:bottomleft, mintotals, showmins=false, kwargs...)
        savefig(p, case * ".pdf")
    end

    let kwargs = ()
        colormap = Dict{Float64,Int}()
        case = [("GaAs", "Ga₂₀As₂₀ supercell (no preconditioner, Anderson)"),
                ("AlVac", "Al₄₀ surface (no preconditioner, Anderson)"),
                ("Mn3Si", "Mn₃Si (Kerker, Anderson)")]

        for (case, label) in case
            scfkeys = ["standard_01", "standard_05", "standard_08",
                       "adaptive",    "adaptive_05"]
            case == "AlVac" && filter!(!isequal("adaptive_05"), scfkeys)
            p = plot_convergence(case; colormap, scfkeys, label, showmins=false, kwargs...)
            savefig(p, case * ".pdf")
        end
    end
end

function main()
    dump_plots()
    dump_latex_tables()
end

(abspath(PROGRAM_FILE) == @__FILE__) && main()

max_stars_repo_name: mfherbst/supporting-adaptive-damping

max_issues_repo_name: mfherbst/supporting-adaptive-damping

max_forks_repo_name: mfherbst/supporting-adaptive-damping


text:
""" Parse relations in Julia-esque syntax into undirected wiring diagrams.
"""
module ParseRelations
export @relation, parse_relation_diagram

using Compat
using Match

using ...CategoricalAlgebra.CSets, ...Present
using ...Theories: FreeCategory
using ...WiringDiagrams.UndirectedWiringDiagrams
using ...WiringDiagrams.UndirectedWiringDiagrams: TheoryUWD, TheoryTypedUWD

# Data structures
#################

@present TheoryRelationDiagram <: TheoryUWD begin
  Name::Ob
  name::Hom(Box,Name)
  variable::Hom(Junction,Name)
end

const RelationDiagram = CSetType(
  TheoryRelationDiagram, data=[:Name],
  index=[:box, :junction, :outer_junction], unique_index=[:variable])

RelationDiagram(nports::Int) =
  UndirectedWiringDiagram(RelationDiagram, nports,
                          name=Symbol, variable=Symbol)

@present TheoryTypedRelationDiagram <: TheoryTypedUWD begin
  Name::Ob
  name::Hom(Box,Name)
  variable::Hom(Junction,Name)
end

const TypedRelationDiagram = CSetType(
  TheoryTypedRelationDiagram, data=[:Name, :Type],
  index=[:box, :junction, :outer_junction], unique_index=[:variable])

TypedRelationDiagram(port_types::AbstractVector{Symbol}) =
  UndirectedWiringDiagram(TypedRelationDiagram, port_types,
                          name=Symbol, variable=Symbol)

# Parsing
#########

""" Parse an undirected wiring diagram from a relation in Julia-esque syntax.

Unlike the [`@program`](@ref) macro for directed wiring diagrams, this macro
fundamentally alters the usual semantics of the Julia language. Function calls
with n arguments are now interpreted as assertions that an n-ary relation holds
at a particular point. For example, the operation of composing of binary
relations R ⊆ X × Y and S ⊆ Y × Z can be represented as an undirected wiring
diagram by the macro call

```julia
@relation (x,z) where (x::X, y::Y, z::Z) begin
  R(x,y)
  S(y,z)
end
```

In general, the context in the `where` clause defines the set of junctions in
the diagram and variable sharing defines the wiring of ports to junctions.
"""
macro relation(exprs...)
  Expr(:call, GlobalRef(ParseRelations, :parse_relation_diagram),
       (QuoteNode(expr) for expr in exprs)...)
end

""" Parse an undirected wiring diagram from a relation expression.

For more information, see the corresponding macro [`@relation`](@ref).
"""
function parse_relation_diagram(expr::Expr)
  @match expr begin
    Expr(:function, [head, body]) => parse_relation_diagram(head, body)
    Expr(:->, [head, body]) => parse_relation_diagram(head, body)
    _ => error("Not a function or lambda expression")
  end
end

function parse_relation_diagram(head::Expr, body::Expr)
  @match head begin
    Expr(:where, [Expr(:tuple, args), Expr(:tuple, context)]) =>
      parse_relation_diagram(context, args, body)
    _ => error("Invalid syntax in declaration of outer ports and context")
  end
end

function parse_relation_diagram(context::AbstractVector, args::AbstractVector,
                                body::Expr)
  all_vars, all_types = parse_context(context)
  outer_vars = parse_variables(args)
  outer_vars ⊆ all_vars || error("One of variables $outer_vars is not declared")
  var_types = if isnothing(all_types) # Untyped case.
    vars -> length(vars)
  else # Typed case.
    var_type_map = Dict{Symbol,Symbol}(zip(all_vars, all_types))
    vars -> getindex.(Ref(var_type_map), vars)
  end

  # Create diagram and add outer ports and junctionsm
  UWD = isnothing(all_types) ? RelationDiagram : TypedRelationDiagram
  d = UWD(var_types(outer_vars))
  add_junctions!(d, var_types(all_vars), variable=all_vars)
  set_junction!(d, ports(d, outer=true),
                incident(d, outer_vars, :variable), outer=true)

  # Add boxes to diagram.
  for expr in Base.remove_linenums!(body).args
    name, vars = @match expr begin
      Expr(:call, [name::Symbol, vars...]) => (name, parse_variables(vars))
      _ => error("Invalid syntax in box definition $expr")
    end
    vars ⊆ all_vars || error("One of variables $vars is not declared")
    box = add_box!(d, var_types(vars), name=name)
    set_junction!(d, ports(d, box), incident(d, vars, :variable))
  end

  return d
end

function parse_context(context)
  vars = map(context) do term
    @match term begin
      Expr(:(::), [var::Symbol, type::Symbol]) => (var => type)
      var::Symbol => var
      _ => error("Invalid syntax in term $expr of context")
    end
  end
  if vars isa AbstractVector{Symbol}
    (vars, nothing)
  elseif vars isa AbstractVector{Pair{Symbol,Symbol}}
    (first.(vars), last.(vars))
  else
    error("Context $context mixes typed and untyped terms")
  end
end

parse_variables(vars) = collect(Symbol, vars)

end


text:
using Compat
import Compat.read

export PcapFileHeader, PcapRec, PcapOffline,
       pcap_get_record

type PcapFileHeader
    magic_number::UInt32
    version_major::UInt16
    version_minor::UInt16
    thiszone::Int32
    sigfigs::UInt32
    snaplen::UInt32
    network::UInt32
    PcapFileHeader() = new(0,0,0,0,0,0,0)
end # type PcapFileHeader

type PcapRec
    ts_sec::UInt32
    ts_usec::UInt32
    incl_len::UInt32
    orig_len::UInt32
    payload::Vector{UInt8}
    PcapRec() = new(0,0,0,0, Vector{UInt8}(0))
end # type PcapRec

type PcapOffline
    filename::AbstractString
    file::IO
    filehdr::PcapFileHeader
    record::PcapRec
    is_big::Bool
    function PcapOffline(fn::AbstractString)
        filename = fn
        file = open(fn, "r+")
        filehdr, is_big = decode_hdr(file)
        record = PcapRec()
        new(filename, file, filehdr, record, is_big)
    end # constructor
end # type PcapOffline

#----------
# decode PCap file format header
#----------
function decode_hdr(file::Any)
    filehdr = PcapFileHeader()
    filehdr.magic_number  = read(file, UInt32)
    big_endian = false
    if filehdr.magic_number == 0xd4c3b2a1
        big_endian = true
    end
    filehdr.version_major = big_endian ? ntoh(read(file, UInt16)) : read(file, UInt16)
    filehdr.version_minor = big_endian ? ntoh(read(file, UInt16)) : read(file, UInt16)
    filehdr.thiszone      = read(file, Int32)
    filehdr.sigfigs       = big_endian ? ntoh(read(file, UInt32)) : read(file, UInt32)
    filehdr.snaplen       = big_endian ? ntoh(read(file, UInt32)) : read(file, UInt32)
    filehdr.network       = big_endian ? ntoh(read(file, UInt32)) : read(file, UInt32)
    return [filehdr, big_endian]
end # function decode_hdr

#----------
# decode next record in PCap file
#----------
function pcap_get_record(s::PcapOffline)
    rec = PcapRec()
    if (!eof(s.file))
        rec.ts_sec   = s.is_big ? ntoh(read(s.file, UInt32)) : read(s.file, UInt32)
        rec.ts_usec  = s.is_big ? ntoh(read(s.file, UInt32)) : read(s.file, UInt32)
        rec.incl_len = s.is_big ? ntoh(read(s.file, UInt32)) : read(s.file, UInt32)
        rec.orig_len = s.is_big ? ntoh(read(s.file, UInt32)) : read(s.file, UInt32)
        rec.payload  = read(s.file, rec.incl_len)
        return rec
    end
    nothing
end # function pcap_get_record


text:
# A simple and practical FASTA file parser
# ========================================

import Automa
import Automa.RegExp: @re_str
const re = Automa.RegExp

# Create a machine of FASTA.
fasta_machine = (function ()
    # First, describe FASTA patterns in regular expression.
    lf          = re"\n"
    newline     = re"\r?" * lf
    identifier  = re"[!-~]*"
    description = re"[!-~][ -~]*"
    letters     = re"[A-Za-z*-]*"
    sequence    = re.cat(letters, re.rep(newline * letters))
    record      = re.cat('>', identifier, re.opt(re" " * description), newline, sequence)
    fasta       = re.rep(record)

    # Second, bind action names to each regular expression.
    lf.actions[:enter]          = [:count_line]
    identifier.actions[:enter]  = [:mark]
    identifier.actions[:exit]   = [:identifier]
    description.actions[:enter] = [:mark]
    description.actions[:exit]  = [:description]
    letters.actions[:enter]     = [:mark]
    letters.actions[:exit]      = [:letters]
    record.actions[:exit]       = [:record]

    # Finally, compile the final FASTA pattern into a state machine.
    return Automa.compile(fasta)
end)()

# It is useful to visualize the state machine for debugging.
# write("fasta.dot", Automa.machine2dot(fasta_machine))
# run(`dot -Tsvg -o fasta.svg fasta.dot`)

# Bind Julia code to each action name (see the `parse_fasta` function defined below).
fasta_actions = Dict(
    :count_line  => :(linenum += 1),
    :mark        => :(mark = p),
    :identifier  => :(identifier = mark == 0 ? "" : String(data[mark:p-1]); mark = 0),
    :description => :(description = mark == 0 ? "" : String(data[mark:p-1]); mark = 0),
    :letters     => :(mark > 0 && unsafe_write(buffer, pointer(data, mark), p - mark); mark = 0),
    :record      => :(push!(records, FASTARecord(identifier, description, take!(buffer)))))

# Define a type to store a FASTA record.
mutable struct FASTARecord
    identifier::String
    description::String
    sequence::Vector{UInt8}
end

# Generate a parser function from `fasta_machine` and `fasta_actions`.
context = Automa.CodeGenContext(generator=:goto, checkbounds=false)
@eval function parse_fasta(data::Union{String,Vector{UInt8}})
    # Initialize variables you use in the action code.
    records = FASTARecord[]
    mark = 0
    linenum = 1
    identifier = description = ""
    buffer = IOBuffer()

    # Initialize variables used by the state machine.
    $(Automa.generate_init_code(context, fasta_machine))
    p_end = p_eof = lastindex(data)

    # This is the main loop to iterate over the input data.
    $(Automa.generate_exec_code(context, fasta_machine, fasta_actions))

    # Check the last state the machine reached.
    if cs != 0
        error("failed to parse on line ", linenum)
    end

    # Finally, return records accumulated in the action code.
    return records
end

# Run the FASTA parser.
records = parse_fasta("""
>NP_003172.1 brachyury protein isoform 1 [Homo sapiens]
MSSPGTESAGKSLQYRVDHLLSAVENELQAGSEKGDPTERELRVGLEESELWLRFKELTNEMIVTKNGRR
MFPVLKVNVSGLDPNAMYSFLLDFVAADNHRWKYVNGEWVPGGKPEPQAPSCVYIHPDSPNFGAHWMKAP
VSFSKVKLTNKLNGGGQIMLNSLHKYEPRIHIVRVGGPQRMITSHCFPETQFIAVTAYQNEEITALKIKY
NPFAKAFLDAKERSDHKEMMEEPGDSQQPGYSQWGWLLPGTSTLCPPANPHPQFGGALSLPSTHSCDRYP
TLRSHRSSPYPSPYAHRNNSPTYSDNSPACLSMLQSHDNWSSLGMPAHPSMLPVSHNASPPTSSSQYPSL
WSVSNGAVTPGSQAAAVSNGLGAQFFRGSPAHYTPLTHPVSAPSSSGSPLYEGAAAATDIVDSQYDAAAQ
GRLIASWTPVSPPSM
""")


text:
module TestEnsembleGibbs

using Test: @test, @testset
using Distributions
using Ensemble
using Ensemble.Stats: logsumexp
using Statistics

const snr_half = 3.0
const mu = log(3.0)
const sigma = 1.0/3.0
const lambda = 100.0

const ptrue = [log(lambda), mu, log(sigma)]

log_odds(x) = 0.5*(x*x - snr_half*snr_half)

function logpselect(x)
    lo = log_odds(x)

    return lo - logsumexp(0.0, lo)
end

function logpnselect(x)
    lo = log_odds(x)

    return -logsumexp(0.0, lo)
end

function draw(lambda=lambda, mu=mu, sigma=sigma)
    n = rand(Poisson(lambda))
    xs = exp.(mu .+ sigma.*randn(n))
    psel = exp.(logpselect.(xs))

    sel = rand(n) .< psel

    xs[sel], xs[.~(sel)]
end

function lognorm_logpdf(xs, mu, sigma)
    out = zeros(size(xs)...)

    log_inv_sqrt_2pi = -0.5*log(2.0*pi)
    log_sigma = log(sigma)
    sigma2 = sigma*sigma

    for i in eachindex(xs)
        log_x = log(xs[i])
        y = log_x - mu
        out[i] = log_inv_sqrt_2pi - log_sigma - log_x - 0.5*y*y/sigma2
    end

    out
end

function make_lnprob(xsdet)
    ndet = length(xsdet)
    function lnprob(params, xsndet)
        nndet = length(xsndet)
        log_lambda, mu, log_sigma = params
        lambda = exp(log_lambda)
        sigma = exp(log_sigma)

        sum(logpselect.(xsdet)) + sum(logpnselect.(xsndet)) + (ndet + nndet)*log_lambda - lambda + sum(lognorm_logpdf(xsdet, mu, sigma)) + sum(lognorm_logpdf(xsndet, mu, sigma))
    end
end

function gibbsupdate(params, xsndet)
    log_lambda, mu, log_sigma = params
    lambda = exp(log_lambda)
    sigma = exp(log_sigma)

    xsdet, xsndet = draw(lambda, mu, sigma)
    xsndet
end

function testall()
    xsdet, xsndet = draw()

    lnprob = make_lnprob(xsdet)
    ps = zeros(3, 100)
    for i in 1:100
        for j in 1:3
            ps[j,i] = ptrue[j] + 1e-3*randn()
        end
    end
    gs = EnsembleGibbs.gibbses(ps, [nothing for i in 1:100], gibbsupdate)
    lnprobs = EnsembleGibbs.lnprobs(ps, gs, lnprob)

    for i in 1:1000
        ps, gs, lnprobs = EnsembleGibbs.update(ps, gs, lnprobs, lnprob, gibbsupdate)
    end

    pts = zeros(3, 100, 100)
    for i in 1:100
        for j in 1:10
            ps, gs, lnprobs = EnsembleGibbs.update(ps, gs, lnprobs, lnprob, gibbsupdate)
        end
        pts[:,:,i] = ps
    end

    @testset "EnsembleGibbs tests" begin
        @test((mean(pts[1,:,:] .- ptrue[1])/std(pts[1,:,:]) < 3))
        @test((mean(pts[2,:,:] .- ptrue[2])/std(pts[2,:,:]) < 3))
        @test((mean(pts[3,:,:] .- ptrue[3])/std(pts[3,:,:]) < 3))
    end
end

end