text:
using ComplexMixtures
using PDBTools
using Plots
using LaTeXStrings
using EasyFit

function fig() # to take care of globals only

# Plot defaults
plot_font = "Computer Modern"
default(
    fontfamily=plot_font,
    linewidth=2.5,
    framestyle=:box,
    label=nothing,
    grid=false,
    palette=:tab10
)
scalefontsizes()
#scalefontsizes(1.3)

# Load pdb file of the system
system = readPDB("../simulation/equilibrated.pdb")

# Select the atoms corresponding to glycerol and water
popc = select(system,"resname POPC")
water = select(system,"water")
ethanol = select(system,"resname ETOH")

# Set the complete membrane as a single solute
solute = Selection(popc,nmols=1)

# Load previously computed data
mddf_water_POPC = load("./mddf_water_POPC.json")
mddf_ethanol_POPC = load("./mddf_ethanol_POPC.json")

# Set a solvent structure, uses to retrieve the atomic Contributions
# of ethanol to the MDDF
solvent = Selection(ethanol,natomspermol=9)

l = @layout [ grid(2,1)  b{0.3w} ;
              c          d{0.5w} ;
              grid(2,1) grid(2,1) ]
plot(layout=l)

# Water-POPC MDDF and KB integral. Here we use `movavg` from `EasyFit` to soften
# the noise
x = mddf_water_POPC.d # distances
y = movavg(mddf_water_POPC.mddf,n=10).x
plot!(x,y,label="Water",
    #xlabel=L"\textrm{Distance / \AA}",
    ylabel="MDDF",subplot=1
)

y = movavg(mddf_water_POPC.kb,n=10).x/1000
plot!(x,y,xlabel=L"\textrm{Distance / \AA}",ylabel=L"\textrm{KB~/~L~mol^{-1}}",subplot=2)
plot!(xlim=(0,10),subplot=2)

# Ethanol-POPC MDDF and KB integral
x = mddf_ethanol_POPC.d
y = movavg(mddf_ethanol_POPC.mddf,n=10).x
plot!(x,y,label="Ethanol")
plot!(xlim=(0,10),xticks=:none,subplot=1)

y = movavg(mddf_ethanol_POPC.kb,n=10).x/1000
plot!(x,y,xlabel=L"\textrm{Distance / \AA}",ylabel=L"\textrm{KB~/~L~mol^{-1}}",subplot=2)
plot!(xlim=(0,10),subplot=2)

# Contributions of the ethanol groups
groups = [
    (select(ethanol,"name O1 or name HO1"),"Hydroxyl"),
    (select(ethanol,"not name O1 and not name HO1"),"Aliphatic chain"),
]
x = mddf_ethanol_POPC.d
y = movavg(mddf_ethanol_POPC.mddf,n=10).x
plot!(x,y,label="Total",subplot=3)
for group in groups
    group_contrib = contrib(solvent,mddf_ethanol_POPC.solvent_atom,group[1])
    y = movavg(group_contrib,n=10).x
    plot!(x,y,label=group[2],subplot=3)
end
plot!(
    xlim=(1,8),
    xlabel=L"\textrm{Distance / \AA}",
    ylabel="MDDF", subplot=3
)

# Contributions of POPC groups to ethanol-POPC distribution. Bellow is the list
# of the atom types of each POPC group. We need to list those to select specifically
# the atoms of each chemical group of interest.
groups = [
  (["N","C12","H12A","H12B","C13","H13A","H13B","H13C","C14",
    "H14A","H14B","H14C","C15","H15A","H15B","H15C","C11","H11A","H11B"],"Choline"),
  ([ "P","O13","O14","O12" ],"Phosphate"),
  (["O11","C1","HA","HB","C2","HS","O21","C3","HX","HY","O31"],"Glycerol"),
  (["O22","C21","H2R","H2S","C22","C23","H3R","H3S","C24","H4R","H4S",
    "C25","H5R","H5S","C26","H6R","H6S","C27","H7R","H7S","C28","H8R","H8S",
    "C29","H91","C210","H101","C211","H11R","H11S","C212","H12R","H12S",
    "C213","H13R","H13S","C214","H14R","H14S","C215","H15R","H15S",
    "C216","H16R","H16S","C217","H17R","H17S","C218","H18R","H18S","H18T"],"Oleoyl"),
  (["C31","O32","C32","H2X","H2Y","C33","H3X","H3Y","C34","H4X","H4Y",
    "C35","H5X","H5Y","C36","H6X","H6Y","C37","H7X","H7Y","C38","H8X",
    "H8Y","C39","H9X","H9Y","C310","H10X","H10Y","C311","H11X","H11Y",
    "C312","H12X","H12Y","C313","H13X","H13Y","C314","H14X","H14Y","C315",
    "H15X","H15Y","C316","H16X","H16Y","H16Z"],"Palmitoyl")
]
x = mddf_ethanol_POPC.d
y = movavg(mddf_ethanol_POPC.mddf,n=10).x
plot!(x,y,label="Total",subplot=4)
for group in groups
    # Retrieve the contributions of the atoms of this group
    group_contrib = contrib(solute,mddf_ethanol_POPC.solute_atom,group[1])
    y = movavg(group_contrib,n=10).x
    plot!(x,y,label=group[2],subplot=4)
end
plot!(
    xlim=(1.3,5),
    ylim=(0,1.8),
    xlabel=L"\textrm{Distance / \AA}",
    ylabel="MDDF",
    subplot=4
)

# the same for POPC-water
x = mddf_water_POPC.d
y = movavg(mddf_water_POPC.mddf,n=10).x
plot!(x,y,label="Total",subplot=5)
for group in groups
    group_contrib = contrib(solute,mddf_water_POPC.solute_atom,group[1])
    y = movavg(group_contrib,n=10).x
    plot!(x,y,label=group[2],subplot=5)
end
plot!(
    xlim=(1.3,5),
    ylim=(0,1.8),
    xlabel=L"\textrm{Distance / \AA}",
    ylabel="MDDF",
    subplot=5
)

# Map of interactions of ethanol with Oleoyl groups. Now we split the
# Oleoyl chain into its components, ordered along the chain.
oleoyl_groups = [
  (["O22","C21"],"CO"),
  (["H2R","H2S","C22"],L"\textrm{CH_2}"),
  (["C23","H3R","H3S"],L"\textrm{CH_2}"),
  (["C24","H4R","H4S"],L"\textrm{CH_2}"),
  (["C25","H5R","H5S"],L"\textrm{CH_2}"),
  (["C26","H6R","H6S"],L"\textrm{CH_2}"),
  (["C27","H7R","H7S"],L"\textrm{CH_2}"),
  (["C28","H8R","H8S"],L"\textrm{CH_2}"),
  (["C29","H91"],L"\textrm{CH}"),
  (["C210","H101"],L"\textrm{CH}"),
  (["C211","H11R","H11S"],L"\textrm{CH_2}"),
  (["C212","H12R","H12S"],L"\textrm{CH_2}"),
  (["C213","H13R","H13S"],L"\textrm{CH_2}"),
  (["C214","H14R","H14S"],L"\textrm{CH_2}"),
  (["C215","H15R","H15S"],L"\textrm{CH_2}"),
  (["C216","H16R","H16S"],L"\textrm{CH_2}"),
  (["C217","H17R","H17S"],L"\textrm{CH_2}"),
  (["C218","H18R","H18S","H18T"],L"\textrm{CH_3}"),
]

gcontrib = zeros(length(mddf_ethanol_POPC.d),length(oleoyl_groups))
for (ig, group) in pairs(oleoyl_groups)
    gcontrib[:,ig] .= movavg(contrib(solute,mddf_ethanol_POPC.solute_atom,group[1]),n=10).x
end
labels = [ oleoyl_groups[i][2] for i in 1:length(oleoyl_groups) ]
idmin = findfirst( d -> d > 1.5, mddf_ethanol_POPC.d)
idmax = findfirst( d -> d > 3.0, mddf_ethanol_POPC.d)

contourf!(
    1:length(oleoyl_groups),
    mddf_ethanol_POPC.d[idmin:idmax],
    gcontrib[idmin:idmax,:],
    color=cgrad(:tempo),linewidth=1,linecolor=:black,
    colorbar=:none,levels=10,
    ylabel=L"\textrm{Distance/\AA}",xrotation=60,
    xticks=(1:length(oleoyl_groups),labels),
    subplot=6
)

# And with Palmitoyl groups
palmitoyl_groups = [
  (["C31","O32"],"CO"),
  (["C32","H2X","H2Y"],L"\textrm{CH_2}"),
  (["C33","H3X","H3Y"],L"\textrm{CH_2}"),
  (["C34","H4X","H4Y"],L"\textrm{CH_2}"),
  (["C35","H5X","H5Y"],L"\textrm{CH_2}"),
  (["C36","H6X","H6Y"],L"\textrm{CH_2}"),
  (["C37","H7X","H7Y"],L"\textrm{CH_2}"),
  (["C38","H8X","H8Y"],L"\textrm{CH_2}"),
  (["C39","H9X","H9Y"],L"\textrm{CH_2}"),
  (["C310","H10X","H10Y"],L"\textrm{CH_2}"),
  (["C311","H11X","H11Y"],L"\textrm{CH_2}"),
  (["C312","H12X","H12Y"],L"\textrm{CH_2}"),
  (["C313","H13X","H13Y"],L"\textrm{CH_2}"),
  (["C314","H14X","H14Y"],L"\textrm{CH_2}"),
  (["C315","H15X","H15Y"],L"\textrm{CH_2}"),
  (["C316","H16X","H16Y","H16Z"],L"\textrm{CH_3}")
]
gcontrib = zeros(length(mddf_ethanol_POPC.d),length(palmitoyl_groups))
for (ig, group) in pairs(palmitoyl_groups)
    gcontrib[:,ig] .= movavg(contrib(solute,mddf_ethanol_POPC.solute_atom,group[1]),n=10).x
end
labels = [ palmitoyl_groups[i][2] for i in 1:length(palmitoyl_groups) ]
idmin = findfirst( d -> d > 1.5, mddf_ethanol_POPC.d)
idmax = findfirst( d -> d > 3.0, mddf_ethanol_POPC.d)

# This creates a contour plot, with some options to make the plot look nicer.
contourf!(
    1:length(palmitoyl_groups),
    mddf_ethanol_POPC.d[idmin:idmax],
    gcontrib[idmin:idmax,:],
    color=cgrad(:tempo),linewidth=1,linecolor=:black,
    colorbar=:none,levels=10,
    xlabel="Group",ylabel=L"\textrm{Distance/\AA}",xrotation=60,
    xticks=(1:length(palmitoyl_groups),labels),subplot=7
)

annotate!( 10, 2.7, text("Ethanol-Oleoyl", :left, 12, plot_font), subplot=6)
annotate!(  8, 2.7, text("Ethanol-Palmitoyl", :left, 12, plot_font), subplot=7)

# And now the same for water-POPC

gcontrib = zeros(length(mddf_water_POPC.d),length(oleoyl_groups))
for (ig, group) in pairs(oleoyl_groups)
    gcontrib[:,ig] .= movavg(contrib(solute,mddf_water_POPC.solute_atom,group[1]),n=10).x
end
labels = [ oleoyl_groups[i][2] for i in 1:length(oleoyl_groups) ]
idmin = findfirst( d -> d > 1.5, mddf_water_POPC.d)
idmax = findfirst( d -> d > 3.0, mddf_water_POPC.d)

contourf!(
    1:length(oleoyl_groups),
    mddf_water_POPC.d[idmin:idmax],
    gcontrib[idmin:idmax,:],
    color=cgrad(:tempo),linewidth=1,linecolor=:black,
    colorbar=:none,levels=10,
    ylabel=L"\textrm{Distance/\AA}",xrotation=60,
    xticks=(1:length(oleoyl_groups),labels),subplot=8
)

gcontrib = zeros(length(mddf_water_POPC.d),length(palmitoyl_groups))
for (ig, group) in pairs(palmitoyl_groups)
    gcontrib[:,ig] .= movavg(contrib(solute,mddf_water_POPC.solute_atom,group[1]),n=10).x
end
labels = [ palmitoyl_groups[i][2] for i in 1:length(palmitoyl_groups) ]
idmin = findfirst( d -> d > 1.5, mddf_water_POPC.d)
idmax = findfirst( d -> d > 3.0, mddf_water_POPC.d)

contourf!(
    1:length(palmitoyl_groups),
    mddf_water_POPC.d[idmin:idmax],
    gcontrib[idmin:idmax,:],
    color=cgrad(:tempo),linewidth=1,linecolor=:black,
    colorbar=:none,levels=10,
    xlabel="Group",ylabel=L"\textrm{Distance/\AA}",xrotation=60,
    xticks=(1:length(palmitoyl_groups),labels),subplot=9
)

annotate!( 10, 2.7, text("Water-Oleoyl", :left, 12, plot_font), subplot=8)
annotate!(  8, 2.7, text("Water-Palmitoyl", :left, 12, plot_font), subplot=9)

for (c,lab) in [((-21.8, 16.4),"A)"),
                ((-21.8, 14.1),"B)"),
                (( -2.4, 15.2),"C)"),
                ((-21.8,  9.4),"D)"),
                (( -2.4,  9.4),"E)"),
                ((-21.8,  6.0),"F)"),
                (( -2.4,  6.0),"G)")]
    annotate!(
        c[1],c[2],
        text(lab,plot_font,12),
        subplot=9,
    )
end

annotate!(
    5,2,
    text("Ethanol contributions",plot_font,10),
    subplot=3
)
annotate!(
    3.8,0.3,
    text("PolyACR contributions\nto Ethanol MDDF",plot_font,10),
    subplot=4
)
annotate!(
    3.8,0.4,
    text("PolyACR contributions\nto Water MDDF",plot_font,10),
    subplot=5
)

plot!(
    leftmargin=10Plots.Measures.mm,
    size=(1000,800)
)

savefig("./popc.png")
savefig("./figure7.pdf")

end; fig()


text:
include("testtypevers.jl")
include("testutilapi.jl")
include("testwellknown.jl")
include("testcodec.jl")

ext: jl


text:
################################################################################
# This set of functions prepares OpenMendel to evaluate the likelihood
# of a collection of pedigrees by the Elston-Stewart algorithm.
################################################################################
#
# Required OpenMendel packages and modules.
#
# using DataStructures
# using GeneralUtilities
# using GeneticUtilities
# using ReadData

mutable struct Instruction
  instructions :: Int
  instructstart :: Vector{Int}
  instructfinish :: Vector{Int}
  operation :: Vector{Int}
  site :: Vector{Int}
  rankarray :: Vector{Int}
  space :: Vector{Int}
  extra :: Array{Vector{Int}, 1}
end # mutable struct Instruction

export Instruction
export allele_consolidation, genotype_elimination!
export orchestrate_likelihood, model_recombination_fractions

penetrance_and_prior_array = 1
transmission_array         = 2
pure_add                   = 3
multiply_and_add           = 4
pure_multiply              = 5
quit_processing_pedigree   = 6

export penetrance_and_prior_array, transmission_array, pure_add
export multiply_and_add, pure_multiply, quit_processing_pedigree

"""
Orchestrate allele lumping, genotype consolidation,
and creation of the Elston-Stewart instructions.
"""
function orchestrate_likelihood(pedigree::Pedigree, person::Person,
  nuclear_family::NuclearFamily, locus::Locus,
  keyword::Dict{AbstractString, Any})

  pedigrees = pedigree.pedigrees
  populations = person.populations
  model_loci = locus.model_loci
  eliminate_genotypes = keyword["eliminate_genotypes"]
  lump_alleles = keyword["lump_alleles"]
  product_mode = keyword["product_mode"]
  complexity_threshold = keyword["complexity_threshold"]
  #
  # Define the frequency of the lumped allele at each model locus.
  #
##
#   for i = 1:person.people
#     for l = 1:locus.loci
#       println(person.name[i],"  ",locus.name[l],"  ",person.genotype[i, l])
#     end
#   end
  if lump_alleles
    lumped_frequency = zeros(pedigrees, populations, model_loci)
    for l = 1:model_loci
      loc = locus.model_locus[l]
##
      if loc != locus.mutable_locus
##println(" Allele lumping at locus = ",locus.name[loc])
        lumped_frequency[:, :, l] =
          allele_consolidation(pedigree, person, locus, loc)
      end
    end
  else
    lumped_frequency = zeros(0, 0, 0)
  end
  #
  # Perform genotype elimination over the model loci.
  #
  if eliminate_genotypes
    for l = 1:model_loci
      loc = locus.model_locus[l]
##
      if loc != locus.mutable_locus
##println(" Genotype elimination at locus = ",locus.name[loc])
        genotype_elimination!(pedigree, person, nuclear_family, locus, loc)
      end
    end
  end
  #
  # Find an upper bound on the number of instructions.
  #
  bound = person.people
  if product_mode
    bound = bound * model_loci
  end
  bound = 6 * bound + pedigrees
  #
  # Allocate the components of the instruction vector and the first
  # and last instruction associated with each pedigree.
  #
  elston_stewart_count = zeros(pedigrees)
  instruction_start = zeros(Int, pedigrees)
  instruction_finish = zeros(Int, pedigrees)
  operation = zeros(Int, bound)
  site = zeros(Int, bound)
  rankarray = zeros(Int, bound)
  space = zeros(Int, bound)
  extra = Array{Array{Int, 1}}(undef, bound)
  for i = 1:bound
    extra[i] = zeros(Int, 0)
  end
  #
  # Find a good summation sequence.
  #
  instructions = prepare_elston_stewart(pedigree, person, locus,
    elston_stewart_count, instruction_start, instruction_finish, operation,
    site, rankarray, space, extra, bound, product_mode, complexity_threshold)
  #
  # Return the instructions for the Elston-Stewart algorithm.
  #
  instruction = Instruction(instructions, instruction_start,
    instruction_finish, operation, site, rankarray, space, extra)
  locus.lumped_frequency = lumped_frequency
  return (instruction, elston_stewart_count)
end  # function orchestrate_likelihood

"""
Consolidate alleles in each pedigree at locus number loc.
The lumped allele is numbered alleles + 1,
where alleles counts the number of alleles at the locus.
The frequency of the lumped allele and the lumping status
of each pedigree is returned along with the altered
genotype set for each person.
Warning: Allele consolidation should precede genotype elimination.
"""
function allele_consolidation(pedigree::Pedigree, person::Person,
  locus::Locus, loc::Int)
  #
  # Initialize variables.
  #
  xlinked = locus.xlinked[loc]
  pedigrees = pedigree.pedigrees
  (populations, alleles) = size(locus.frequency[loc])
  maximum_genotypes = alleles^2
  lumped_allele = alleles + 1
  lumped = falses(pedigrees)
  lumped_frequency = ones(pedigrees, populations)
  #
  # Find the alleles observed or partially observed in
  # the current pedigree.
  #
  for ped = 1:pedigrees
    observed = BitSet()
    for i = pedigree.pedstart[ped]:pedigree.twin_finish[ped]
      if xlinked && person.male[i]
        if length(person.genotype[i, loc]) < alleles
          for (gm, gp) in person.genotype[i, loc]
            push!(observed, gm)
          end
        end
      else
        if length(person.genotype[i, loc]) < maximum_genotypes
          for (gm, gp) in person.genotype[i, loc]
            push!(observed, gm)
            push!(observed, gp)
          end
        end
      end
    end
    #
    # Find the frequency of the lumped allele.
    #
    if length(observed) < alleles - 1
      lumped[ped] = true
      for a in observed
        f = locus.frequency[loc][:, a]
        lumped_frequency[ped, :] = lumped_frequency[ped, :] - f
      end
    else
      lumped_frequency[ped, :] .= 0.0
    end
    #
    # Construct a new genotype set for each person, replacing
    # each unobserved allele by the lumped allele.
    #
    for i = pedigree.pedstart[ped]:pedigree.twin_finish[ped]
      new_genotype_set = Set{Tuple{Int, Int}}()
      for (gm, gp) in person.genotype[i, loc]
        if in(gm, observed)
          hm = gm
        else
          hm = lumped_allele
        end
        if in(gp, observed)
          hp = gp
        else
          hp = lumped_allele
        end
        push!(new_genotype_set, (hm, hp))
      end
      person.genotype[i, loc] = deepcopy(new_genotype_set)
    end
  end
  #
  # Recount the number of homozygous genotypes per person and return
  # the lumped genotypes and lumped allele frequencies.
  #
  count_homozygotes!(person, loc)
  return lumped_frequency
end # function allele_consolidation

"""
Perform genotype elimination at locus number loc.
Genotypes for person i are stored in the set person.genotype[i, loc].
"""
function genotype_elimination!(pedigree::Pedigree, person::Person,
  nuclear_family::NuclearFamily, locus::Locus, loc::Int)

  people = person.people
  pedigrees = pedigree.pedigrees
  families = nuclear_family.families
  xlinked = locus.xlinked[loc]
  stable_pedigree = falses(pedigrees)
  stable_family = falses(families)
  saved_genotype = Array{Set{Tuple{Int, Int}}}(undef, people)
  #
  # Begin elimination for locus loc.
  #
  progress = true
  while progress
  #
  # Loop over all nuclear families.
  #
    for fam = 1:families
      stable_family[fam] = true
      ped = nuclear_family.pedigree[fam]
      if stable_pedigree[ped]; continue; end
      #
      # Initialize the saved sets of genotypes to be empty.
      #
      mom = nuclear_family.mother[fam]
      dad = nuclear_family.father[fam]
      saved_genotype[mom] = Set{Tuple{Int, Int}}()
      saved_genotype[dad] = Set{Tuple{Int, Int}}()
      for sib in nuclear_family.sib[fam]
        saved_genotype[sib] = Set{Tuple{Int, Int}}()
      end
      #
      # Loop over all combinations of maternal and paternal genotypes.
      #
      for (m1, m2) in person.genotype[mom, loc]
        for (d1, d2) in person.genotype[dad, loc]
          family_match = true
          #
          # Loop over all siblings.
          #
          for sib in nuclear_family.sib[fam]
            #
            # Loop over all genotypes of the current sib.
            #
            sib_match = false
            for (s1, s2) in person.genotype[sib, loc]
              #
              # Check for compatibility between the sib genotype and
              # the parental genotypes.
              #
              maternal_match = (s1 == m1 || s1 == m2)
              if !xlinked || !person.male[sib]
                paternal_match = (s2 == d1 || s2 == d2)
                sib_match = maternal_match && paternal_match
              else
                sib_match = maternal_match
              end
              #
              # As soon as a compatible sib genotype is found, stop looking.
              #
              if sib_match; break; end
            end  # sib genotype loop
            #
            # If no genotype of the current sib matches, exit.
            #
            family_match = family_match && sib_match
            if !family_match; break; end
          end  # sib loop
          #
          # If all sibs provide a match, save all matching parental and
          # sibling genotypes.
          #
          if !family_match; continue; end
          push!(saved_genotype[mom], (m1, m2))
          push!(saved_genotype[dad], (d1, d2))
          for sib in nuclear_family.sib[fam]
            for (s1, s2) in person.genotype[sib, loc]
              maternal_match = (s1 == m1 || s1 == m2)
              if !xlinked || !person.male[sib]
                paternal_match = (s2 == d1 || s2 == d2)
              else
                paternal_match = true
              end
              if maternal_match && paternal_match
                push!(saved_genotype[sib], (s1, s2))
              end
            end
          end  # sib loop
        end  # dad loop
      end  # mom loop
      #
      # Replace the current genotype sets in the family by the saved
      # genotype sets.
      #
      if length(saved_genotype[mom]) < length(person.genotype[mom, loc])
        stable_family[fam] = false
        person.genotype[mom, loc] = intersect(saved_genotype[mom],
          person.genotype[mom, loc])
      end
      if length(saved_genotype[dad]) < length(person.genotype[dad, loc])
        stable_family[fam] = false
        person.genotype[dad, loc] = intersect(saved_genotype[dad],
          person.genotype[dad, loc])
      end
      for sib in nuclear_family.sib[fam]
        if length(saved_genotype[sib]) < length(person.genotype[sib, loc])
          stable_family[fam] = false
          person.genotype[sib, loc] = intersect(saved_genotype[sib],
            person.genotype[sib, loc])
        end
      end  # sib loop
    end  # family loop
    #
    # Genotype elimination continues until genotypes stabilize
    # across all pedigrees.
    #
    fill!(stable_pedigree, true)
    for fam = 1:families
      ped = nuclear_family.pedigree[fam]
      stable_pedigree[ped] = stable_pedigree[ped] && stable_family[fam]
    end
    progress = !all(stable_pedigree)
  end # progress loop
  #
  # Recount the number of homozygous genotypes per person.
  #
  count_homozygotes!(person, loc)
end # function genotype_elimination!

"""
Specify the steps of the Elston-Stewart algorithm for each pedigree.
"""
function prepare_elston_stewart(pedigree::Pedigree, person::Person,
  locus::Locus, elston_stewart_count::Vector{Float64},
  instruction_start::Vector{Int}, instruction_finish::Vector{Int},
  operation::Vector{Int}, site::Vector{Int}, rankarray::Vector{Int},
  space::Vector{Int}, extra::Vector{Vector{Int}}, bound::Int,
  product_mode::Bool, complexity_threshold::Float64)

  pedigrees = pedigree.pedigrees
  model_loci = locus.model_loci
  free_recombination = locus.free_recombination
  instructions = 0
  #
  # Record the homozygote counts at the model loci.
  #
  homozygotes = zeros(Int, person.people, model_loci)
  for l = 1:model_loci
    homozygotes[:, l] = person.homozygotes[:, locus.model_locus[l]]
  end
  #
  # Create the instructions pedigree by pedigree.
  #
  for ped = 1:pedigrees
    ped_start = pedigree.pedstart[ped]
    ped_finish = pedigree.pedfinish[ped]
    ped_size = ped_finish - ped_start + 1
    pedigree_product_mode = product_mode
    #
    # Set up the adjacency graph of the current pedigree. From the adjacency
    # graph, determine a summation sequence by the greedy algorithm.
    # Note that greedy changes the set array graph.
    #
    (nodes, graph, weight) = adjacency_graph(pedigree, person, locus,
      homozygotes, pedigree_product_mode, ped, model_loci, free_recombination)
    (lowest_cost_sum, permutation) = greedy(nodes, graph, weight, ped_size)
    #
    # When likelihood evaluation can proceed in product mode, check if
    # non-product mode will do better.
    #
    if product_mode && !free_recombination
      product_mode_cost = lowest_cost_sum
      pedigree_product_mode = false
      (nodes, graph, weight) = adjacency_graph(pedigree, person, locus,
        homozygotes, pedigree_product_mode, ped, model_loci, free_recombination)
      (lowest_cost_sum, permutation) = greedy(nodes, graph, weight, ped_size)
      #
      # When product mode is superior, revert to it.
      #
      if lowest_cost_sum > product_mode_cost
        pedigree_product_mode = true
        (nodes, graph, weight) = adjacency_graph(pedigree, person, locus,
          homozygotes, pedigree_product_mode, ped, model_loci,
          free_recombination)
        (lowest_cost_sum, permutation) = greedy(nodes, graph, weight, ped_size)
      end
    end
    #
    # Decide whether the current pedigree is too demanding for likelihood
    # computation. If not, set up the index sets corresponding to the
    # various arrays encountered in likelihood evaluation.
    #
    if lowest_cost_sum <= complexity_threshold
      (work_arrays, index_set) = create_index_set(pedigree, person, locus,
        nodes, ped, pedigree_product_mode)
      #
      # Prepare the instructions for the Elston-Stewart algorithm.
      #
      instruction_start[ped] = instructions + 1
      elston_stewart_count[ped] = 0.0
      instructions = prepare_instructions(pedigree, person, locus,
        elston_stewart_count, index_set, weight, permutation, extra,
        operation, rankarray, space, site, work_arrays, instructions, nodes,
        bound, ped, pedigree_product_mode)
      instruction_finish[ped] = instructions
    #
    # If the current pedigree is too complex, set the Elston-Stewart count
    # equal to the lowest cost sum.
    #
    else
      instructions = instructions + 1
      instruction_start[ped] = instructions
      elston_stewart_count[ped] = lowest_cost_sum
      instruction_finish[ped] = instructions
    end
  end
  return instructions
end # function prepare_elston_stewart

"""
Define the adjacency graph connecting the indices of summation.
"""
function adjacency_graph(pedigree::Pedigree, person::Person, locus::Locus,
  homozygotes::Matrix{Int}, pedigree_product_mode::Bool, ped::Int,
  model_loci::Int, free_recombination::Bool)

  ped_start = pedigree.pedstart[ped]
  ped_finish = pedigree.pedfinish[ped]
  ped_size = ped_finish - ped_start + 1
  offset = ped_start - 1
  #
  # In product mode, nodes of the adjacency graph are person-locus pairs.
  # In non-product mode, nodes are simply people.
  #
  if pedigree_product_mode
    nodes = ped_size * model_loci
  else
    nodes = ped_size
  end
  #
  # Each node has a set of neighbors graph[node]. The weight of a node equals
  # the number of its associated genotypes or multi-locus genotypes.
  #
  weight = ones(Int, nodes)
  graph = empties(nodes)
  #
  # In product mode, loop over all person-locus pairs. Number nodes by
  # putting them in a matrix with loci on rows and people on columns.
  # A node is numbered by the its column-major position in the matrix.
  #
  if pedigree_product_mode
    for i = ped_start:ped_finish
      mom = person.mother[i]
      dad = person.father[i]
      for j = 1:model_loci
        jloc = locus.model_locus[j]
        node = (i - ped_start) * model_loci + j
        weight[node] = length(person.genotype[i, jloc])
        for k = 1:model_loci
          kloc = locus.model_locus[k]
          self_connected = false
          if mom != 0
            if !loci_split(homozygotes[mom, :], j, k, free_recombination)
              if j != k; self_connected = true; end
              push!(graph[node], (mom - ped_start) * model_loci + k)
            end
          end
          if dad != 0
            if !loci_split(homozygotes[dad, :], j, k, free_recombination)
              if j != k; self_connected = true; end
              push!(graph[node], (dad - ped_start) * model_loci + k)
            end
          end
          if length(person.children[i]) > 0 &&
            !loci_split(homozygotes[i, :], j, k, free_recombination)
            if j != k; self_connected = true; end
            for kid in person.children[i]
              push!(graph[node], (kid - ped_start) * model_loci + k)
            end
          end
          if self_connected
            push!(graph[node], (i - ped_start) * model_loci + k)
          end
        end
      end
    end
  else
  #
  # In non-product mode, two people are neighbors if and only if one
  # is the parent of the other. If this is true, then they occur as
  # indices in the same transmission array. To insure that nodes are
  # numbered starting with 1, subtract an offset from each person number.
  #
    for node = 1:nodes
      i = node + offset
      for kid in person.children[i]
        push!(graph[node], kid - offset)
      end
      if person.mother[i] != 0
        push!(graph[node], person.mother[i] - offset)
        push!(graph[node], person.father[i] - offset)
      end
      for l = 1:model_loci
        loc = locus.model_locus[l]
        weight[node] = weight[node] * length(person.genotype[i, loc])
      end
    end
  end
  return (nodes, graph, weight)
end # function adjacency_graph

"""
Decide if two model loci are split within person i
in the sense of being separated by an obligate heterozygous locus.
Under free recombination, the likelihood factors over different loci.
"""
function loci_split(homozygotes, locus_1::Int, locus_2::Int,
                    free_recombination::Bool)

  if free_recombination
    split_loci = locus_1 != locus_2
  else
    split_loci = false
    for loc = min(locus_1, locus_2) + 1:max(locus_1, locus_2) - 1
      if homozygotes[loc] == 0
        return true
      end
    end
  end
  return split_loci
end # function loci_split

"""
Select a good summation sequence by the greedy method
and store it in a permutation. Observe that the cost
of removing a node from the adjacency graph is a simple surrogate
for the number of arithmetic operations involved in summing over
the genotypes possible for the node.
"""
function greedy(nodes::Int, graph::Vector{BitSet},
                weight::Vector{Int}, ped_size::Int)

  log_cost = zeros(nodes)
  log_weight = zeros(nodes)
  #
  # Initialize the permutation and the log weight attached to each node.
  #
  permutation = collect(1:nodes)
  for node = 1:nodes
    log_weight[node] = log(weight[node])
  end
  lowest_cost_sum = 0.0
  #
  # Compute the log cost of eliminating a node.
  #
  for node = 1:nodes
    log_sum = log_weight[node]
    for element in graph[node]
      log_sum = log_sum + log_weight[element]
    end
    log_cost[node] = log_sum
  end
  #
  # Find the lowest cost node to eliminate. Note that the variable last
  # limits the search to nodes representing the current locus when we are
  # in product mode. Nodes are therefore considered in clumps starting
  # with the leftmost locus and ending with the rightmost locus.
  #
  for i = 1:nodes
    lowest_cost = log_cost[permutation[i]]
    lowest_cost_index = i
    last = (div(i - 1, ped_size) + 1) * ped_size
    for j = i + 1:last
      competing_cost = log_cost[permutation[j]]
      if competing_cost < lowest_cost
        (lowest_cost, lowest_cost_index) = (competing_cost, j)
      end
    end
    low = exp(lowest_cost)
    lowest_cost_sum = lowest_cost_sum + min(low, 1e20)
    #
    # Insert the lowest cost node into the permutation.
    #
    (permutation[i], permutation[lowest_cost_index]) =
      (permutation[lowest_cost_index], permutation[i])
    best = permutation[i]
    #
    # Expand the neighborhood of each neighbor of the lowest cost node and
    # Delete the lowest cost node and neighbor from each revised neighborhood.
    #
    for neighbor in graph[best]
      graph[neighbor] = union(graph[neighbor], graph[best])
      delete!(graph[neighbor], best)
      delete!(graph[neighbor], neighbor)
      #
      # Recompute the log cost of removing each neighbor.
      #
      log_sum = log_weight[neighbor]
      for element in graph[neighbor]
        log_sum = log_sum + log_weight[element]
      end
      log_cost[neighbor] = log_sum
    end
  end
  return (lowest_cost_sum, permutation)
end  # function greedy

"""
Create the index sets for all initial arrays in likelihood computation.
"""
function create_index_set(pedigree::Pedigree, person::Person,
  locus::Locus, nodes::Int, ped::Int, pedigree_product_mode::Bool)
  #
  # Allocate the index sets and a variable n tracking the current index set.
  #
  model_loci = locus.model_loci
  free_recombination = locus.free_recombination
  ped_start = pedigree.pedstart[ped]
  ped_finish = pedigree.pedfinish[ped]
  offset = ped_start - 1
  index_set = empties(6 * nodes)
  n = 0
  #
  # In product mode, loop over all person-locus pairs. Prepare for the
  # penetrance-prior arrays first.
  #
  if pedigree_product_mode
    for i = ped_start:ped_finish
      for l = 1:model_loci
        n = n + 1
        push!(index_set[n], (i - ped_start) * model_loci + l)
      end
    end
    #
    # Next prepare for the transmission arrays by looping over all children.
    # Transmission arrays will be considered parent by parent.
    #
    for i = ped_start:ped_finish
      if person.mother[i] != 0
        for parent = 1:2
          if parent == 1
            j = person.mother[i]
          else
            j = person.father[i]
          end
          #
          # Find a first and last locus in a block of loci not split by
          # an obligate heterozygous locus in the parent.
          #
          (first, last) = (1, 0)
          while last < model_loci
          #
          # Under free recombination, each transmission arrays completely
          # factors.
          #
            if free_recombination
              last = first
            else
              last = model_loci
              for l = first + 1:model_loci
                loc = locus.model_locus[l]
                if person.homozygotes[j, loc] == 0
                  last = l
                  break
                end
              end
            end
            #
            # Insert the parent-locus pairs into the index set before the
            # child-locus pairs.
            #
            n = n + 1
            spread = last - first + 1
            for l = first:last
              push!(index_set[n], (j - ped_start) * model_loci + l)
              push!(index_set[n], (i - ped_start) * model_loci + l)
            end
            #
            # Exit if the last locus has been hit. Otherwise, reset the
            # first locus to the current last locus and find another block
            # of loci.
            #
            if last == model_loci; break; end
            if free_recombination; last = first + 1; end
            first = last
          end
        end
      end
    end
  else
    #
    # In nonproduct mode, prepare for the penetrance-prior arrays first.
    #
    for i = ped_start:ped_finish
      n = n + 1
      push!(index_set[n], i - offset)
    end
    #
    # Next prepare for the transmission arrays by looping over all children.
    #
    for i = ped_start:ped_finish
      if person.mother[i] != 0
        n = n + 1
        push!(index_set[n], person.mother[i] - offset)
        push!(index_set[n], i - offset)
        n = n + 1
        push!(index_set[n], person.father[i] - offset)
        push!(index_set[n], i - offset)
      end
    end
  end
  #
  # Return the initial index sets.
  #
  return (n, index_set)
end # function create_index_set

"""
Prepare the instructions for performing the Elston-Stewart algorithm.
"""
function prepare_instructions(pedigree::Pedigree, person::Person, locus::Locus,
  elston_stewart_count::Vector{Float64}, index_set::Vector{BitSet},
  weight::Vector{Int}, permutation::Vector{Int}, extra::Vector{Vector{Int}},
  operation::Vector{Int}, rankarray::Vector{Int}, space::Vector{Int},
  site::Vector{Int}, work_arrays::Int, instructions::Int, nodes::Int,
  bound::Int, ped::Int, pedigree_product_mode::Bool)

  model_loci = locus.model_loci
  ped_start = pedigree.pedstart[ped]
  offset = ped_start - 1
  #
  # Let n be the index of the current instruction. Prepare a set
  # to hold which arrays are involved in the current operation.
  #
  n = instructions
  active_array = trues(6 * nodes)
  involved = BitSet()
  union_set = BitSet()
  #
  # Give the instructions for creating the initial arrays.
  #
  for i = 1:work_arrays
    n = n + 1
    if i <= nodes
      operation[n] = penetrance_and_prior_array
      extra[n] = zeros(Int, 3)
      #
      # Load the starting and finishing loci and the person involved
      # in a penetrance and prior array.
      #
      if pedigree_product_mode
        j = first(index_set[i]) - 1
        extra[n][1] = mod(j, model_loci) + 1
        extra[n][2] = extra[n][1]
        extra[n][3] = div(j, model_loci) + ped_start
      else
        extra[n][1] = 1
        extra[n][2] = model_loci
        extra[n][3] = first(index_set[i]) + offset
      end
    else
      operation[n] = transmission_array
      extra[n] = zeros(Int, 4)
      #
      # Load the starting and finishing loci and the parent and child
      # involved in a transmission array.
      #
      if pedigree_product_mode
        j = first(index_set[i]) - 1
        k = last(index_set[i]) - 1
        extra[n][1] = mod(j, model_loci) + 1
        extra[n][2] = mod(k, model_loci) + 1
        extra[n][3] = div(j, model_loci) + ped_start
        extra[n][4] = div(k, model_loci) + ped_start
      else
        extra[n][1] = 1
        extra[n][2] = model_loci
        extra[n][3] = first(index_set[i]) + offset
        extra[n][4] = last(index_set[i]) + offset
      end
    end
    #
    # Load the site, the rank, and the size of the current array.
    #
    site[n] = i
    rankarray[n] = length(index_set[i])
    space[n] = product_weights(index_set[i], weight)
  end
  #
  # Give the instructions for creating and destroying arrays in computing
  # the likelihood of the current pedigree.
  #
  for node = 1:nodes
    i = permutation[node]
    #
    # Search through the index sets and find those containing the node i
    # to be eliminated.
    #
    for iteration = 1:100000
      involved = BitSet()
      for k = 1:work_arrays
        if active_array[k] && i in index_set[k]
          push!(involved, k)
        end
      end
      #
      # Exit the loop if no index containing i remains.
      #
      if isempty(involved); break; end
      #
      # Update the instruction index.
      #
      n = n + 1
      #
      # For a pure add, record the array to be summed.
      #
      if length(involved) == 1
        operation[n] = pure_add
        j = first(involved)
        site[n] = j
        rankarray[n] = length(index_set[j])
        space[n] = product_weights(index_set[j], weight)
        elston_stewart_count[ped] = elston_stewart_count[ped] + space[n]
        active_array[j] = false
        break
      #
      # For a pure multiply or a multiply and add,
      # choose a pair of arrays to combine.
      #
      else
        best_1 = 1
        best_2 = 2
        pivot = i
        #
        # For a pure multiply, find the best pair of arrays
        # to multiply by a greedy algorithm.
        #
        if length(involved) > 2
          pivot = 0
          best_size = Inf
          for j = 1:length(involved)
            k = select_set_element(involved, j)
            for l = j + 1:length(involved)
              m = select_set_element(involved, l)
              union_set = union(index_set[k], index_set[m])
              criterion = product_weights(union_set, weight)
              if criterion < best_size
                best_1 = j
                best_2 = l
                best_size = criterion
              end
            end
          end
        end
        #
        # Form the union of the index sets corresponding to
        # the only or best pair of arrays.
        #
        j = select_set_element(involved, best_1)
        k = select_set_element(involved, best_2)
        union_set = union(index_set[j], index_set[k])
        #
        # Combine nodes and insert the instructions for a pure multiply or
        # a multiply and add.
        #
        (rankarray[n], extra[n]) = combine_nodes(index_set[j], index_set[k],
          union_set, weight, j, k, n, pivot)
        #
        # Delete node i in a multiply and add. Also record the current
        # operation.
        #
        space[n] = product_weights(union_set, weight)
        if length(involved) == 2
          elston_stewart_count[ped] = elston_stewart_count[ped] + 2 * space[n]
          delete!(union_set, i)
          operation[n] = multiply_and_add
          space[n] = product_weights(union_set, weight)
        else
          elston_stewart_count[ped] = elston_stewart_count[ped] + space[n]
          operation[n] = pure_multiply
        end
        #
        # Activate the array created and determine its location and space
        # requirement.
        #
        work_arrays = work_arrays + 1
        index_set[work_arrays] = deepcopy(union_set)
        active_array[work_arrays] = true
        site[n] = work_arrays
        #
        # Deactivate the two arrays destroyed.
        #
        active_array[j] = false
        active_array[k] = false
      end
    end
  end
  #
  # Give the instruction for finishing the current pedigree and update the
  # number of instructions
  #
  n = n + 1
  operation[n] = quit_processing_pedigree
  return n
end # function prepare_instructions

"""
Consolidate the nodes for the multiplicand and product arrays into blocks.
Generate the instructions for a pure multiply or a multiply and add.
a and b are the index sets for the multiplicand arrays,
and a_union_b is the index set for the product array.
n is the current instruction index, and pivot is the node to
be eliminated in a multiply and add. a_location and b_location give
the locations of the multiplicand arrays in the evolving sequence of arrays.
"""
function combine_nodes(a::BitSet, b::BitSet, a_union_b::BitSet,
  weight::Vector{Int}, a_location::Int, b_location::Int, nodes::Int, pivot::Int)
  #
  # Initialize current positions in the original index sets and
  # the block index sets. Also initialize the position of the pivot
  # and the type indicator for block construction.
  #
  extent = zeros(Int, length(a_union_b))
  a_set_position = 1
  b_set_position = 1
  a_block_position = 0
  b_block_position = 0
  union_block_position = 0
  pivot_block_position = 0
  previous_type = -2
  a_stride = 0
  b_stride = 0
  #
  # Create and initialize arrays.
  #
  c = length(a_union_b)
  a_block = zeros(Int, c)
  b_block = zeros(Int, c)
  a_increment = zeros(Int, c)
  b_increment = zeros(Int, c)
  #
  # March through the union of the two original index sets, defining
  # blocks as we go. If the a set has been exhausted, flag its current
  # node as impossible; likewise for the b set.
  #
  while a_set_position <= length(a) || b_set_position <= length(b)
    if a_set_position > length(a)
      current_a_node = nodes + 1
    else
      current_a_node = select_set_element(a, a_set_position)
    end
    if b_set_position > length(b)
      current_b_node = nodes + 1
    else
      current_b_node = select_set_element(b, b_set_position)
    end
    #
    # A new block involving indices from the a set must be started.
    #
    if current_a_node < current_b_node
      new_type = 1
      current_node = current_a_node
      a_set_position = a_set_position + 1
    #
    # A new block involving indices from the b set must be started.
    #
    elseif current_a_node > current_b_node
      new_type = -1
      current_node = current_b_node
      b_set_position = b_set_position + 1
    #
    # The current block continues.
    #
    else
      new_type = 0
      current_node = current_a_node
      a_set_position = a_set_position + 1
      b_set_position = b_set_position + 1
    end
    #
    # A block boundary is hit whenever the pivot node for a multiply and
    # add is encountered.
    #
    if current_node == pivot
      previous_type = -2
      pivot_block_position = union_block_position + 1
    end
#
# Update the product of the weights associated with a continuing block.
#
    if new_type == previous_type
      extent[union_block_position] = extent[union_block_position] *
        weight[current_node]
    #
    # For a new block, record its position, initial weight, and its membership
    # in the a and/or b block index sets.
    #
    else
      union_block_position = union_block_position + 1
      extent[union_block_position] = weight[current_node]
      if new_type == 1 || new_type == 0
        a_block_position = a_block_position + 1
        a_block[a_block_position] = union_block_position
      end
      if new_type == -1 || new_type == 0
        b_block_position = b_block_position + 1
        b_block[b_block_position] = union_block_position
      end
      #
      # Record whether the current block involves a indices, b indices, or
      # both. The pivot marks the beginning of a new block.
      #
      previous_type = new_type
      if current_node == pivot; previous_type = -2; end
    end
  end
  #
  # Insert the rank of the product array into the current instruction.
  #
  rankarray = union_block_position
  a_blocks = a_block_position
  b_blocks = b_block_position
  #
  # Allocate the extra array for the current instruction. Record the
  # locations of the multiplicand arrays. Also load the pivot range, which
  # determines whether a pure multiply or a multiply and add is performed.
  #
  extra = zeros(Int, 5 + 3 * rankarray)
  extra[1] = a_location
  extra[2] = b_location
  if pivot_block_position == 0
    pivot_range = 0
  else
    pivot_range = extent[pivot_block_position] - 1
  end
  extra[3] = pivot_range
  #
  # Calculate the increments for the multiplicand arrays. Also calculate
  # their strides for a multiply and sum.
  #
  a_block_position = 1
  b_block_position = 1
  a_prod = 1
  b_prod = 1
  for union_block_position = 1:rankarray
    if union_block_position == pivot_block_position
      a_stride = a_prod
      b_stride = b_prod
    end
    if a_block_position <= a_blocks && a_block[a_block_position] == union_block_position
      a_prod = a_prod * extent[union_block_position]
      a_increment[union_block_position] = 1
      a_block_position = a_block_position + 1
    else
      a_increment[union_block_position] = 1 - a_prod
    end
    if b_block_position <= b_blocks && b_block[b_block_position] == union_block_position
      b_prod = b_prod * extent[union_block_position]
      b_increment[union_block_position] = 1
      b_block_position = b_block_position + 1
    else
      b_increment[union_block_position] = 1 - b_prod
    end
  end
  #
  # Load the strides into the current instruction.
  #
  extra[4] = a_stride
  extra[5] = b_stride
  #
  # Adjust the increment of each index pivot to correct for the add
  # operation in a multiply and add.
  #
  for i = 1:pivot_block_position - 1
    a_increment[i] = a_increment[i] - a_stride * pivot_range
    b_increment[i] = b_increment[i] - b_stride * pivot_range
  end
  #
  # Load the extents and increments into the current instruction.
  #
  true_rank = rankarray
  if pivot_block_position != 0; true_rank = rankarray - 1; end
  j = 5
  for i = 1:rankarray
    if i != pivot_block_position
      j = j + 1
      extra[j] = extent[i]
      extra[j + true_rank] = a_increment[i]
      extra[j + 2 * true_rank] = b_increment[i]
    end
  end
  rankarray = true_rank
  return (rankarray, extra)
end # function combine_nodes

"""
Find the product of the weights of the elements of the integer set S.
"""
function product_weights(s::BitSet, weight::Vector{Int})

  p = 1
  for element in s
    p = p * weight[element]
  end
  return p
end # function product_weights

"""
Compute recombination fractions between adjacent model loci.
"""
function model_recombination_fractions(locus::Locus,
  keyword::Dict{AbstractString, Any})

  n = locus.model_loci - 1
  theta = zeros(2, n)
  choice = keyword["genetic_map_function"]
  for l = 1:n
    j = locus.model_locus[l]
    k = locus.model_locus[l + 1]
    d = abs(locus.morgans[1, k] - locus.morgans[1, j])
    theta[1, l] = map_function(d, choice)
    d = abs(locus.morgans[2, k] - locus.morgans[2, j])
    theta[2, l] = map_function(d, choice)
  end
  return theta
end # function model_recombination_fractions

max_stars_repo_name: OpenMendel/MendelBase

max_issues_repo_name: OpenMendel/MendelBase

max_forks_repo_name: OpenMendel/MendelBase


text:
# Adapted from C++ GMG Solver from Harald Köstler
# "Multigrid HowTo: A simple Multigrid solver in C++ in less than 200 lines of code"
#  https://www10.cs.fau.de/publications/reports/TechRep_2008-03.pdf

export PoissonGMG,solve!,boundary_coeff,treatboundary,fres,residual,restrict_residual,interpolate_correct,smooth,maxlevels
using SparseArrays
using LinearAlgebra
using BenchmarkTools
using LoopVectorization

struct PoissonGMG{SMOOTHER}
    Sol::Vector{Array{Float64,2}}
    Solb::Vector{Array{Float64,2}}
    RHS::Vector{Array{Float64,2}}
    invh2::Vector{Float64}
    smoothers::Vector{SMOOTHER}
    nlevels::Int
    nprae::Int
    npost::Int
    ncoarse::Int
    bc::NTuple{2,NTuple{2,BoundaryCondition}}
end

function maxlevels(n)
    try
       Int(log(n)/log(2)-1)
    catch
        error("nx=ny must be a power of two")
    end
end

function PoissonGMG(l2D,::Type{SMOOTHER};nlevels=-1,nprae=2,npost=1,ncoarse=20) where {SMOOTHER}
    Lx,Ly=l2D.Lx,l2D.Ly
    nx,ny=l2D.nx,l2D.ny
    nx!=ny && error("nx must be equal to ny")
    Lx!=Ly && error("Lx must be equal to Ly")
    bc=l2D.bc
    nlevels>maxlevels(nx) && error("nlevels too large")
    (nlevels==-1) && (nlevels=maxlevels(nx))
    # @show nlevels

    # add 1 ghost layer in each direction
    sizeadd=2
    Sol=Vector{Array{Float64,2}}(undef,nlevels)
    Solb=Vector{Array{Float64,2}}(undef,nlevels)
    RHS=Vector{Array{Float64,2}}(undef,nlevels)
    invh2=Vector{Float64}(undef,nlevels)
    smoothers=Vector{SMOOTHER}(undef,nlevels)
    nrows,ncols=nx,ny
    h2=1.0
    L=Lx
    for i=1:nlevels
        nlevels
        Sol[i]=zeros(nrows+sizeadd,ncols+sizeadd)
        Solb[i]=zeros(nrows+sizeadd,ncols+sizeadd)
        RHS[i]=zeros(nrows+sizeadd,ncols+sizeadd)
        h=L/(nrows)
        h2=h*h
        ih2=1/h2
        invh2[i]=ih2
        smoothers[i]=SMOOTHER(invh2[i],nrows,ncols,bc,(i==1))

        nrows=div(nrows,2)
        ncols=div(ncols,2)

    end
    PoissonGMG{SMOOTHER}(Sol,Solb,RHS,invh2,smoothers,nlevels,nprae,npost,ncoarse,bc)
end

function solve!(Xxy,Bxy,g::PoissonGMG)
    RHS,Sol,Solb,invh2=g.RHS,g.Sol,g.Solb,g.invh2
    nlevels=size(RHS,1)
    r1,s1=RHS[1],Sol[1]
    nrows,ncols=size(s1)

    @inbounds for j=2:ncols-1
        for i=2:nrows-1
            r1[i,j]=Bxy[i-1,j-1]
            s1[i,j]=Xxy[i-1,j-1]
        end
    end

    res_1=residual(1,g)
    residuals=img_solver(nlevels,res_1,g)
    (nx,ny)=size(Xxy)

    @inbounds for j=2:ncols-1
        for i=2:nrows-1
            Xxy[i-1,j-1]=s1[i,j]
        end
    end
    residuals
end

function img_solver(nlevels,res_1,g::PoissonGMG)
    res=res_1
    res_old=0.
    #no . of pre -, post -, and coarse smoothing steps
    nprae=g.nprae
    npost=g.npost
    ncoarse=g.ncoarse
    lev=1
    max_iters=20

    residuals=[res_1]
    for i=1:max_iters
        res_old=res
        VCycle(lev,g)
        res=residual(lev,g)
        push!(residuals,res)
        (res<1.e-9) && return residuals
        # @show lev,i,res,res_1/res,res/res_old
    end
    residuals
end

function smooth(lev,g::PoissonGMG)
    treatboundary(lev,g.Sol[lev],g)
    smooth(g.Sol[lev],g.RHS[lev],g.smoothers[lev])
    treatboundary(lev,g.Sol[lev],g)

end

function VCycle(lev,g)
    #solve problem on coarsest grid ...
    if lev==g.nlevels

        for i=1:g.ncoarse
            smooth(lev,g)
        end
    else
        #... or recursively do V - cycle
        #
        # do some presmoothing steps
        for i=1:g.nprae
            smooth(lev,g)
        end

        restrict_residual(lev,g)
        # initialize the coarse solution to zero
        fill!(g.Sol[lev+1],0.0);
        VCycle(lev+1,g)
        # interpolate error and correct fine solution
        interpolate_correct(lev+1,g)
        # do some postsmoothing steps
        for i=1:g.npost
            smooth(lev,g)
        end
    end
end

boundary_coeff(bc::BoundaryCondition) = (bc==dirichlet) ? -1.0 : 1.0

function treatboundary(lev,Sol,g)
    sl=Sol
    (nrows,ncols)=size(sl)
    bc=g.bc

    cleft=boundary_coeff(bc[1][1])
    cright=boundary_coeff(bc[1][2])
    cbottom=boundary_coeff(bc[2][1])
    ctop=boundary_coeff(bc[2][2])

    @inbounds for j=2:ncols-1
        sl[1,j]=cleft*sl[2,j]
        sl[nrows,j]=cright*sl[nrows-1,j]
    end
    @inbounds for i=2:nrows-1
        sl[i,1]=cbottom*sl[i,2]
        sl[i,ncols]=ctop*sl[i,ncols-1]
    end
    allneumann = bc[1][1]==bc[1][2]==bc[2][1]==bc[2][2]==neumann

end

# @inline fres(x,y,i,j,ih2)= @inbounds x[i,j]+ih2*(y[i+1,j]+y[i,j+1]+y[i,j-1]+y[i-1,j]-4y[i,j])-0.0001y[i,j]
@inline function fres(x,y,i,j,ih2,nx,ny)
    @inbounds x[i,j]+ih2*(y[i+1,j]+y[i,j+1]+y[i,j-1]+y[i-1,j]-4.0*y[i,j])
end

function residual(lev,g)
    sl=g.Sol[lev]
    rl=g.RHS[lev]
    ih2=g.invh2[lev]
    (nrows,ncols)=size(sl)
    res=0.0

    nx=nrows-1
    ny=ncols-1

    for j=2:ny
        @simd for i=2:nx
            rf=fres(rl,sl,i,j,ih2,nx,ny)
            res+=rf*rf
        end
    end
    sqrt(res)
end

# using LoopVectorization

function restrict_residual(lev,g)
    rl1=g.RHS[lev+1]
    sl1=g.Sol[lev+1]
    rl=g.RHS[lev]
    sl=g.Sol[lev]
    (nrows,ncols)=size(rl1)
    ih2=g.invh2[lev]

    nx=nrows-1
    ny=ncols-1
    fill!(rl1,0.0)
    fill!(sl1,0.0)

    @inline finer(i,j)=fres(rl,sl,i,j,ih2,nx,ny)

    @simd for j=2:ncols-1
        fj=2j-2

        for i=2:nrows-1
            fi=2i-2

             @inbounds rl1[i,j]=0.25*(finer(fi,fj)+finer(fi+1,fj)+finer(fi,fj+1)+finer(fi+1,fj+1))
            #  @inbounds rl1[i,j]=0.25*(finer(fi,fj)+finer(fi+1,fj)+finer(fi,fj+1)+finer(fi+1,fj+1))
        end
    end

end

function interpolate_correct(lev,g::PoissonGMG)
    uf=g.Sol[lev-1]
    uc=g.Sol[lev]
    (nrows,ncols)=size(uc)
    ncm1=ncols-1
    nrm1=nrows-1

    @inbounds for j=2:ncm1
        fj=2j-1
        @fastmath for i=2:nrm1
            fi=2i-1
             v=uc[i,j]
             uf[fi  ,fj  ]+=v
             uf[fi-1,fj  ]+=v
             uf[fi  ,fj-1]+=v
             uf[fi-1,fj-1]+=v
        end
    end

end


text:
cd("$(Pkg.dir("Chipmunk"))/deps")
println("Checking dependencies...")

try
    @unix_only run(`which cmake`)
catch exception
    error("You must install cmake to build this package.")
end

println("Good...")

if !isdir("Chipmunk2D")
    println("Cloning Chipmunk source...")

    cd(Pkg.dir("Chipmunk")*"/deps")
    run(`git clone https://github.com/slembcke/Chipmunk2D.git`)
end

cd("Chipmunk2D")
run(`git checkout Chipmunk-7.0.0`)

run(`git checkout Chipmunk-7.0.0`)

run(`cmake . -DBUILD_DEMOS=OFF`)
run(`make`)

ext = ""
@linux_only ext = "so"
@osx_only ext = "dylib"
@windows_only ext = "dll"

@linux_only begin
    cp("src/libchipmunk.so.7.0.0", "../libchipmunk.so", remove_destination=true)
end
@osx_only begin
    cp("src/libchipmunk.7.0.0.dylib", "../libchipmunk.dylib", remove_destination=true)
end
@windows_only begin
    cp("src/libchipmunk.7.0.0.dll", "../libchipmunk.dll", remove_destination=true)
end

cd("..")

run(`gcc -fPIC -I./Chipmunk2D/include/chipmunk -c $(Pkg.dir("Chipmunk"))/src/c/chipmunkjl.c`)
run(`gcc -I./Chipmunk2D/include/chipmunk -L. -lchipmunk -shared -o ./libchipmunkjl.$ext chipmunkjl.o`)
run(`rm chipmunkjl.o`)

# rm("Chipmunk2D", recursive=true)

println("Successfully built Chipmunk.jl!")


text:
function get_geo_id_place_id(place_id::String; options=Dict{String, String}())

    r = get_oauth("https://api.twitter.com/1.1/geo/id/$(place_id).json", options)

    return r.status == 200 ? Places(JSON.parse(String(r.body))) : error("Twitter API returned $(r.status) status")

end

function get_user_suggestions_slug(slug::String; options=Dict{String, String}())

    r = get_oauth("https://api.twitter.com/1.1/users/suggestions/$(slug).json", options)

    return r.status == 200 ? JSON.parse(String(r.body)) : error("Twitter API returned $(r.status) status")

end

function get_user_suggestions_slug_members(slug::String; options=Dict{String, String}())

    r = get_oauth("https://api.twitter.com/1.1/users/suggestions/$(slug)/members.json", options)

    return r.status == 200 ? Users(JSON.parse(String(r.body))) : error("Twitter API returned $(r.status) status")

end

function get_retweets_id(id::String; options=Dict{String, String}())

    r = get_oauth("https://api.twitter.com/1.1/statuses/retweets/$(id).json", options)

    return r.status == 200 ? Tweets(JSON.parse(String(r.body))) : error("Twitter API returned $(r.status) status")
end

function post_destroy_single_tweet_id(id::String; options=Dict{String, String}())

    r = post_oauth("https://api.twitter.com/1.1/statuses/destroy/$(id).json", options)

    return r.status == 200 ? Tweets(JSON.parse(String(r.body))) : error("Twitter API returned $(r.status) status")

end

function post_status_retweet_id(id::String; options=Dict{String, String}())

    r = post_oauth("https://api.twitter.com/1.1/statuses/retweet/$(id).json", options)

    return r.status == 200 ? Tweets(JSON.parse(String(r.body))) : error("Twitter API returned $(r.status) status")

end


text:
export AdmittanceModel
export get_Y, get_P, get_ports, partial_copy, compatible, canonical_gauge
export apply_transform, cascade, connect, cascade_and_unite, ports_to_indices
export unite_ports, open_ports, open_ports_except, short_ports, short_ports_except

"""
    abstract type AdmittanceModel{T}

An abstract representation of a linear mapping from inputs `x` to outputs `y` of
the form `YΦ = Px`, `y = QᵀΦ`. The type parameter `T` is the type of a port name
object (typically `String`, `Symbol`, or `Int`). Subtypes `U <: AdmittanceModel`
are expected to implement:

    get_Y(am::U)
    get_P(am::U)
    get_ports(am::U)
    partial_copy(am::U; Y, P, ports)
    compatible(AbstractVector{U})
"""
abstract type AdmittanceModel{T} end

"""
    get_Y(pso::PSOModel)
    get_Y(bbox::Blackbox)

Return a vector of admittance matrices.
"""
function get_Y end

"""
    get_P(pso::PSOModel)
    get_P(bbox::Blackbox)

Return an input port matrix.
"""
function get_P end

"""
    get_ports(pso::PSOModel)
    get_ports(bbox::Blackbox)

Return a vector of port identifiers.
"""
function get_ports end

"""
    partial_copy(pso::PSOModel{T, U};
        Y::Union{Vector{V}, Nothing}=nothing,
        P::Union{V, Nothing}=nothing,
        Q::Union{V, Nothing}=nothing,
        ports::Union{AbstractVector{W}, Nothing}=nothing) where {T, U, V, W}

    partial_copy(bbox::Blackbox{T, U};
        Y::Union{Vector{V}, Nothing}=nothing,
        P::Union{V, Nothing}=nothing,
        Q::Union{V, Nothing}=nothing,
        ports::Union{AbstractVector{W}, Nothing}=nothing) where {T, U, V, W}

Create a new model with the same fields except those given as keyword arguments.
"""
function partial_copy end

"""
    compatible(psos::AbstractVector{PSOModel{T, U}}) where {T, U}
    compatible(bboxes::AbstractVector{Blackbox{T, U}}) where {T, U}

Check if the models can be cascaded. Always true for PSOModels and true for Blackboxes
that share the same value of `ω`.
"""
function compatible end

"""
    canonical_gauge(pso::PSOModel)
    canonical_gauge(bbox::Blackbox)

Apply an invertible transformation that takes the model to coordinates in which
`P` is `[I ; 0]` (up to floating point errors). Note this will create a dense model.
"""
function canonical_gauge end

"""
    ==(am1::AdmittanceModel, am2::AdmittanceModel)

Test whether two `AdmittanceModel`s are equal. Two models must be of the same type
to be equal, as well as have the same admittances, port matrices, and port names.
"""
# don't remove space between Base. and ==
function Base. ==(am1::T, am2::T) where {T <: AdmittanceModel}
    return all(getfield(am1, name) == getfield(am2, name) for name in fieldnames(T))
end

"""
    isapprox(am1::AdmittanceModel, am2::AdmittanceModel)

Test whether two `AdmittanceModel`s are approximately equal.
"""
Base.isapprox(am1::AdmittanceModel, am2::AdmittanceModel) = false
function Base.isapprox(am1::T, am2::T) where {T <: AdmittanceModel}
    return all(getfield(am1, name) ≈ getfield(am2, name) for name in fieldnames(T))
end

"""
    apply_transform(am::AdmittanceModel, transform::AbstractMatrix{<:Number})

Apply a linear transformation `transform` to the coordinates of the model.
"""
function apply_transform(am::AdmittanceModel, transform::AbstractMatrix{<:Number})
    Y = [transpose(transform) * m * transform for m in get_Y(am)]
    P = eltype(Y)(transpose(transform) * get_P(am))
    return partial_copy(am, Y=Y, P=P)
end

"""
    cascade(ams::AbstractVector{<:AdmittanceModel})
    cascade(ams::AdmittanceModel...)

Cascade the models into one larger block diagonal model.
"""
function cascade(ams::AbstractVector{<:AdmittanceModel})
    @assert !isempty(ams)
    if length(ams) == 1
        return ams[1]
    end
    @assert compatible(ams)
    Y = [cat(m..., dims=(1,2)) for m in zip([get_Y(am) for am in ams]...)]
    P = cat([get_P(am) for am in ams]..., dims=(1,2))
    ports = vcat([get_ports(am) for am in ams]...)
    return partial_copy(ams[1], Y=Y, P=P, ports=ports)
end
cascade(ams::AdmittanceModel...) = cascade(collect(ams))

"""
    ports_to_indices(am::AdmittanceModel, ports::AbstractVector)
    ports_to_indices(am::AdmittanceModel, ports...)

Find the indices corresponding to given ports.
"""
function ports_to_indices(am::AdmittanceModel, ports::AbstractVector)
    am_ports = get_ports(am)
    return [findfirst(isequal(p), am_ports) for p in ports]
end
ports_to_indices(am::AdmittanceModel, ports...) = ports_to_indices(am, collect(ports))

"""
    unite_ports(am::AdmittanceModel, ports::AbstractVector)
    unite_ports(am::AdmittanceModel, ports...)

Unite the given ports into one port.
"""
function unite_ports(am::AdmittanceModel, ports::AbstractVector)
    if length(ports) <= 1
        return am
    end
    port_inds = ports_to_indices(am, ports)
    keep_inds = filter(!in(port_inds[2:end]), 1:length(get_ports(am))) # keep the first port
    P = get_P(am)
    first_vector = P[:, port_inds[1]]
    constraint_mat = transpose(hcat([first_vector - P[:, i] for i in port_inds[2:end]]...))
    constrained_am = apply_transform(am, nullbasis(constraint_mat))
    return partial_copy(constrained_am, P=get_P(constrained_am)[:, keep_inds],
        ports=get_ports(constrained_am)[keep_inds])
end
unite_ports(am::AdmittanceModel, ports...) = unite_ports(am, collect(ports))

"""
    open_ports(am::AdmittanceModel, ports::AbstractVector)
    open_ports(am::AdmittanceModel, ports...)

Remove the given ports.
"""
function open_ports(am::AdmittanceModel, ports::AbstractVector)
    if length(ports) == 0
        return am
    end
    port_inds = ports_to_indices(am, ports)
    keep_inds = filter(!in(port_inds), 1:length(get_ports(am)))
    return partial_copy(am, P=get_P(am)[:, keep_inds],
        ports=get_ports(am)[keep_inds])
end
open_ports(am::AdmittanceModel, ports...) = open_ports(am, collect(ports))

"""
    open_ports_except(am::AdmittanceModel, ports::AbstractVector)
    open_ports_except(am::AdmittanceModel, ports...)

Remove all ports except those specified.
"""
function open_ports_except(am::AdmittanceModel, ports::AbstractVector)
    return open_ports(am, filter(!in(ports), get_ports(am)))
end
open_ports_except(am::AdmittanceModel, ports...) = open_ports_except(am, collect(ports))

"""
    short_ports(am::AdmittanceModel, ports::AbstractVector)
    short_ports(am::AdmittanceModel, ports...)

Replace the given ports by short circuits.
"""
function short_ports(am::AdmittanceModel, ports::AbstractVector)
    isempty(ports) && return am
    port_inds = ports_to_indices(am, ports)
    keep_inds = filter(!in(port_inds), 1:length(get_ports(am)))
    constraint_mat = transpose(hcat([get_P(am)[:, i] for i in port_inds]...))
    constrained_am = apply_transform(am, nullbasis(constraint_mat))
    return partial_copy(constrained_am, P = get_P(constrained_am)[:, keep_inds],
        ports = get_ports(constrained_am)[keep_inds])
end
short_ports(am::AdmittanceModel, ports...) = short_ports(am, collect(ports))

"""
    short_ports_except(am::AdmittanceModel, ports::AbstractVector)
    short_ports_except(am::AdmittanceModel, ports...)

Replace all ports with short circuits, except those specified.
"""
function short_ports_except(am::AdmittanceModel, ports::AbstractVector)
    return short_ports(am, filter(!in(ports), get_ports(am)))
end
short_ports_except(am::AdmittanceModel, ports...) =
    short_ports_except(am, collect(ports))

"""
    connect(models::AbstractVector{<:AdmittanceModel})
    connect(models::AdmittanceModel...)

Cascade all models and unite ports with the same name. This results in a
combined admittance model with common ports connected.
"""
function connect(models::AbstractVector{<:AdmittanceModel})
    @assert length(models) >= 1
    if length(models) == 1
        return models[1]
    end
    # number the ports so that the names are all distinct and then cascade
    port_number = 1
    function rename(model)
        ports = [(port_number + i - 1, port) for (i, port) in enumerate(get_ports(model))]
        port_number += length(ports)
        return partial_copy(model, ports=ports)
    end
    model = cascade(map(rename, models))
    # merge all ports with the same name
    original_ports = vcat([get_ports(m) for m in models]...)
    for port in unique(original_ports)
        inds = findall([p[2] == port for p in get_ports(model)])
        model = unite_ports(model, get_ports(model)[inds])
    end
    # remove numbering
    return partial_copy(model, ports=[p[2] for p in get_ports(model)])
end
connect(models::AdmittanceModel...) = connect(collect(models))


text:
"""
    abstract type Dof <: Kernel

Abstract type representing a degree of freedom (DOF), a basis of DOFs, and related objects.
These different cases are distinguished by the return type obtained when evaluating the `Dof`
object on a `Field` object. See function [`evaluate_dof!`](@ref) for more details.

The following functions needs to be overloaded

- [`dof_cache`](@ref)
- [`evaluate_dof!`](@ref)

The following functions can be overloaded optionally

- [`dof_return_type`](@ref)

The interface is tested with

- [`test_dof`](@ref)

In most of the cases it is not strictly needed that types that implement this interface
inherit from `Dof`. However, we recommend to inherit from `Dof`, when possible.

"""
abstract type Dof <: Kernel end

"""
    dof_cache(dof,field)

Returns the cache needed to call `evaluate_dof!(cache,dof,field)`
"""
function dof_cache(dof,field)
  @abstractmethod
end

"""
    evaluate_dof!(cache,dof,field)

Evaluates the dof `dof` with the field `field`. It can return either an scalar value or
an array of scalar values depending the case. The `cache` object is computed with function
[`dof_cache`](@ref).

When a mathematical dof is evaluated on a physical field, a scalar number is returned. If either
the `Dof` object is a basis of DOFs, or the `Field` object is a basis of fields,
or both objects are bases, then the returned object is an array of scalar numbers. The first
dimensions in the resulting array are for the `Dof` object and the last ones for the `Field`
object. E.g, a basis of `nd` DOFs evaluated at physical field returns a vector of `nd` entries.
A basis of `nd` DOFs evaluated at a basis of `nf` fields returns a matrix of size `(nd,nf)`.
"""
function evaluate_dof!(cache,dof,field)
  @abstractmethod
end

"""
    dof_return_type(dof,field)

Returns the type for the value obtained with evaluating `dof` with `field`.

It defaults to

    typeof(evaluate_dof(dof,field))
"""
function dof_return_type(dof,field)
  typeof(evaluate_dof(dof,field))
end

# Testers

"""
    test_dof(dof,field,v,comp::Function=(==))

Test that the `Dof` interface is properly implemented
for object `dof`. It also checks if the object `dof`
when evaluated at the field `field` returns the same
value as `v`. Comparison is made with the `comp` function.
"""
function test_dof(dof,field,v,comp::Function=(==))
  if isa(dof,Dof)
    test_kernel(dof,(field,),v,comp)
  end
  r = evaluate_dof(dof,field)
  @test comp(r,v)
  @test typeof(r) == dof_return_type(dof,field)
end

# Implement Kernel interface

@inline kernel_cache(dof::Dof,field) = dof_cache(dof,field)

@inline apply_kernel!(cache,dof::Dof,field) = evaluate_dof!(cache,dof,field)

@inline kernel_return_type(dof::Dof,field) = dof_return_type(dof,field)

# Some API

"""
    evaluate_dof(dof,field)

Equivalent to

    cache = dof_cache(dof,field)
    evaluate_dof!(cache,dof,field)
"""
function evaluate_dof(dof,field)
  cache = dof_cache(dof,field)
  evaluate_dof!(cache,dof,field)
end

"""
    evaluate(dof::Dof,field)

Equivalent to `evaluate_dof(dof,field)`.
"""
evaluate(dof::Dof,field) = evaluate_dof(dof,field)

# Working with arrays of Dofs

"""
    evaluate_dof_array(dof::AbstractArray,field::AbstractArray)

Evaluates the `Dof` objects in the array `dof` at the `Field` objects
at the array `field` element by element.

The result is numerically equivalent to

    map(evaluate_dof, dof, field)

but it is described with a more memory-friendly lazy type.
"""
function evaluate_dof_array(dof::AbstractArray,field::AbstractArray)
  k = DofEval()
  apply(k,dof,field)
end

function evaluate_dof_array(dof::AbstractArray{<:Dof},field::AbstractArray)
  apply(dof,field)
end

"""
    evaluate(dof::AbstractArray{<:Dof},field::AbstractArray)

Equivalent to `evaluate_dof_array(dof,field)`
"""
function evaluate(dof::AbstractArray{<:Dof},field::AbstractArray)
  evaluate_dof_array(dof,field)
end

struct DofEval <: Kernel end

function kernel_cache(k::DofEval,dof,field)
  dof_cache(dof,field)
end

@inline function apply_kernel!(cache,k::DofEval,dof,field)
  evaluate_dof!(cache,dof,field)
end

function kernel_return_type(k::DofEval,dof,field)
  dof_return_type(dof,field)
end


text:
# When can we convert a RandVar to a Distributions.jl RandVar?
# What inference algorithm if any do we want to use when using mean in model
# Options: 1. user specifies
# 2. No inference algorithm: we do samples on the outside
# Need to take into account conditions on x

"Expectation of `x` from `n` samples"
function samplemean(x::RandVar, n; alg = RejectionSample, kwargs...)
    # @show kwargs
  mean(rand(x, n; alg = alg, kwargs...))
end

samplemeanᵣ(x, n; alg = RejectionSample) = ciid(ω -> samplemean(x(ω), n; alg = alg))

meanstd(x) = (m = mean(x); (mean = m, std = stdm(x, m)))
samplemeanstd(x, n; alg, kwargs...) = meanstd(rand(x, n; alg = alg, kwargs...))
samplemeanstdᵣ(x, n; alg = RejectionSample) = ciid(ω -> samplemeanstd(x(ω), n; alg = alg))

const sampleprob = samplemean
const sampleprobᵣ = samplemeanᵣ

# zt: Default to RejectionSample is problematic

"Probability x is true"
prob(x::RandVar, n, israndvar::Type{Val{false}}) =
  sum((rand(x, alg = RejectionSample) for i = 1:n)) / n
lprob(x::RandVar, n = 1000) = lift(prob)(x, n, Val{false})
prob(x::RandVar, n) = prob(x, n, Val{elemtype(x) <: RandVar})

# Specializations
const unidistattrs = [:succprob, :failprob, :maximum, :minimum, :islowerbounded,
                      :isupperbounded, :isbounded, :std, :median, :mode, :modes,
                      :skewness, :kurtosis, :isplatykurtic, :ismesokurtic,
                      :isleptokurtic, :entropy, :mean]

for func in unidistattrs
  expr =
  quote
    $func(x::RandVar, israndvar::Type{Val{false}}) = Djl.$func(distribution(x))
    $func(x::RandVar, israndvar::Type{Val{true}}) = $(:l *ₛ func)(x)
    $(:l *ₛ func)(x::RandVar) = lift($func)(x, Val{false})
    $func(x::RandVar) = $func(x, Val{elemtype(x) <: RandVar})
  end
  eval(expr)
end

# const bindistattrs = [:entropy, :mgf, :cf, :pdf, :logpdf, :loglikelihood, :cdf,
#                       :logcdf, :ccdf, :logccdf, :quantile, :cquantile,
#                       :invlogcdf, :invlogccdf]

# for func in bindistattrs
#   expr =
#   quote
#     $func(x::PrimRandVar, et::Type{<:Real}) = Djl.$func(distribution(x))
#     $func(x::RandVar, et::Type{<:RandVar}, t) = ciid(mean, x)
#     $func(x::RandVar, t) = $func(x, elemtype(x))
#   end
#   eval(expr)
# end
# Djl.entropy(x::RandVar) = entropy(distribution(x))
# Djl.entropy(x::RandVar, base) = entropy(distribution(x), base)

# Djl.mgf(x::RandVar, t) = mgf(distribution(x), t)
# Djl.cf(x::RandVar, t) = cf(distribution(x), t)

# Djl.pdf(x::RandVar, t) = pdf(distribution(x), t)
# Djl.logpdf(x::RandVar, t) = logpdf(distribution(x), t)
# Djl.loglikelihood(x::RandVar, t) = loglikelihood(distribution(x), t)

# Djl.cdf(x::RandVar, t) = cdf(distribution(x), t)
# Djl.logcdf(x::RandVar, t) = logcdf(distribution(x), t)

# Djl.ccdf(x::RandVar, t) = ccdf(distribution(x), t)
# Djl.logccdf(x::RandVar, t) = logccdf(distribution(x), t)

# Djl.quantile(x::RandVar, t) = quantile(distribution(x), t)
# Djl.cquantile(x::RandVar, t) = cquantile(distribution(x), t)

# Djl.invlogcdf(x::RandVar, t) = invlogcdf(distribution(x), t)
# Djl.invlogccdf(x::RandVar, t) = invlogccdf(distribution(x), t)

# Converions between Distributions and Omega

djldist(::T, params...) where {T <: RandVar} = djltype(T)(params...)

djltype(::Type{<:Normal}) = Djl.Normal
djltype(::Type{<:Beta}) = Djl.Beta
djltype(::Type{<:Bernoulli}) = Djl.Bernoulli
djltype(::Type{<:ReplaceRandVar{Prim}}) where Prim = djltype(Prim)

mayberand(x::RandVar) = rand(x)
mayberand(c) = c

"Convert `rv` into a Distributions.jl `Distribution``"
function distribution(rv::RandVar)
  θs = params(rv)
  θisconst = isconstant.(θs)
  @pre all(θisconst) "All params must be constant to convert to Distributions"
  θsc = mayberand.(θs)
  djldist(rv, θsc...)
end


text:
"""
    chebyexp(A)

Compute the exponential of `A`. The exponential is calculated with a Chebyshev series expansion. For \$|A| ≤ 1\$, we have that

```math
\\exp(A) = J₀(i) + 2 ∑_{k=1}^∞ iᵏ Jₖ(-i) Tₖ(A)
```

where \$Jₖ\$ is the Bessel function of the first kind, and \$Tₖ\$ is a Chebyshev polynomial. If \$|A| > 1\$, scaling and squaring is used.
"""
chebyexp(A::AbstractMatrix{T}) where T<:Number = chebyexp!(Matrix{float(T)}(size(A)),A)

function chebyexp!(R::AbstractMatrix{<:Number},A::AbstractMatrix{<:Number})
    # Run alg
    normA = norm(A,1)
    if normA <= 1
        _chebyexp!(R,A)
    else
        p = ceil(Int,log2(normA))
        _chebyexp!(R,A/2^p)
        Q = similar(R)
        for i = 1:p
            copy!(Q,R)
            A_mul_B!(R,Q,Q)
        end
    end
    return R
end

function _chebyexp!(R::AbstractMatrix{T},A::AbstractMatrix{<:Number}) where T
    # only valid for |A|<1
    n = checksquare(A)
    if T <: BigFloat
        order, coeff0, coeff = 32, cheby_coeff0_big, cheby_coeff_big
    else
        order, coeff0, coeff = 16, cheby_coeff0, cheby_coeff
    end
    # Initialize required matrices
    fill!(R,zero(T))
    Tₖ₋₂ = similar(R); copy!(Tₖ₋₂,A)
    Tₖ₋₁ = eye(T,n,n); A_mul_B!(T(2),A,Tₖ₋₂,T(-1),Tₖ₋₁)
    # Compute expansion
    for i=1:n; R[i,i] = coeff0; end # R₀ = J₀(i)*I, the 0th order
    R .+= coeff[1].*Tₖ₋₂ .+ coeff[2].*Tₖ₋₁ # 1st and 2nd order
    for k in 3:order # higher orders
        A_mul_B!(T(2),A,Tₖ₋₁,T(-1),Tₖ₋₂) # result Tₖ is stored in Tₖ₋₂
        R .+= coeff[k].*Tₖ₋₂ # coeff[k].*Tₖ
        Tₖ₋₂, Tₖ₋₁ = Tₖ₋₁, Tₖ₋₂ # Tₖ₋₁->Tₖ₋₂ and Tₖ->Tₖ₋₁ for next order
    end
    return R
end

# Bessel function values precomputed in Mathematica

# J₀(i)
const cheby_coeff0_big = big"1.2660658777520083355982446252147175376076703113550"
const cheby_coeff0 = Float64(cheby_coeff0_big)

# 2 * iᵏ * Jₖ(-i)
const cheby_coeff_big = BigFloat[
    big"1.1303182079849700544153920552197266146577992432422"
    big"0.27149533953407656236570513998998184589974213622556"
    big"0.044336849848663804952571495259799231058830698339959"
    big"0.0054742404420937326502761684311864595467579461858053"
    big"0.00054292631191394375036214781030755468476712885351629"
    big"0.000044977322954295146654690328110912699086657650642406"
    big"3.1984364624019905058638729766022957272370458074163E-6"
    big"1.9921248066727957259610643848055890533900933857774E-7"
    big"1.1036771725517344326169960913353241812896390172496E-8"
    big"5.5058960796737472504714204020055270687431547280674E-10"
    big"2.4979566169849825227120109342187675410080716361415E-11"
    big"1.0391522306785700504996346724238478525397128556091E-12"
    big"3.9912633564144015128877204015326949127607826796533E-14"
    big"1.4237580108256571488273680253471752219093588992517E-15"
    big"4.7409261025614961710899305606042914145777617483995E-17"
    big"1.4801800572082975003888571658877975360303747318760E-18"
    big"4.3499194949441698455876297633392992805626063964777E-20"
    big"1.2074289272797528890630463524357806390885570735305E-21"
    big"3.1753567370594449606628945704889798438009317678357E-23"
    big"7.9336719716380401114641564996829844420300175294498E-25"
    big"1.8879484042289160772319706157860669889247560558152E-26"
    big"4.2886738765925870898799133815030885460420950260008E-28"
    big"9.3189852817775768480872792470802866623424437483364E-30"
    big"1.9406469749017397597649278461566813645709017660684E-31"
    big"3.8798022492260012156255855282161124021152712081106E-33"
    big"7.4585028873915195213508204862516351326616201311347E-35"
    big"1.3807477824110645231588753652621331312287399205215E-36"
    big"2.4648623717710962928935138361162240264245603184147E-38"
    big"4.2485421925059913850761703704767643098614220928398E-40"
    big"7.0790011762128954933502123970072670493550456762707E-42"
    big"1.1414867782540890660429322724040802483946870774259E-43"
    big"1.7831510375432838840323081019695093079857962301911E-45"
]
const cheby_coeff = Vector{Float64}(cheby_coeff_big)

max_stars_repo_name: jebej/ChebyExp

max_issues_repo_name: jebej/ChebyExp

max_forks_repo_name: jebej/ChebyExp


text:
using UUIDs

export Dependency, BuildDependency, HostBuildDependency,
    is_host_dependency, is_target_dependency, is_build_dependency, is_runtime_dependency

# Pkg.PackageSpec return different types in different Julia versions so...
const PkgSpec = typeof(Pkg.PackageSpec(name="dummy"))

"""
An `AbstractDependency` is a binary dependency of the JLL package.  Dependencies
are installed to `\${prefix}` in the build environment.

Concrete subtypes of `AbstractDependency` are

* [`Dependency`](@ref): a JLL package that is necessary for to build the package
  and to load the generated JLL package.
* [`BuildDependency`](@ref): a JLL package that is necessary only to build the
  package.  This will not be a dependency of the generated JLL package.
* [`HostBuildDependency`](@ref): similar to `BuildDependency`, but it will
  install the artifact for the host platform, instead of that for the target
  platform.
"""
abstract type AbstractDependency end

"""
    is_host_dependency(dep::AbstractDependency) -> Bool

Return whether `dep` is a dependency of the host platform or not.
"""
is_host_dependency

"""
    is_target_dependency(dep::AbstractDependency) -> Bool

Return whether `dep` is a dependency of the target platform or not.
"""
is_target_dependency(dep::AbstractDependency) = !is_host_dependency(dep)

"""
    is_build_dependency(dep::AbstractDependency) -> Bool

Return whether `dep` is a build-time dependency or not.
"""
is_build_dependency

"""
    is_runtime_dependency(dep::AbstractDependency) -> Bool

Return whether `dep` is a runtime dependency or not.
"""
is_runtime_dependency

"""
    Dependency(dep::Union{PackageSpec,String})

Define a binary dependency that is necessary to build the package and load the
generated JLL package.  The argument can be either a string with the name of the
JLL package or a `Pkg.PackageSpec`.
"""
struct Dependency <: AbstractDependency
    pkg::PkgSpec
    build_version::Union{VersionNumber,Nothing}
    Dependency(pkg::PkgSpec, build_version = nothing) = new(pkg, build_version)
end
Dependency(dep::AbstractString, build_version = nothing) = Dependency(PackageSpec(; name = dep), build_version)
is_host_dependency(::Dependency) = false
is_build_dependency(::Dependency) = true
is_runtime_dependency(::Dependency) = true

"""
    BuildDependency(dep::Union{PackageSpec,String})

Define a binary dependency that is necessary only to build the package.  The
argument can be either a string with the name of the JLL package or a
`Pkg.PackageSpec`.
"""
struct BuildDependency <: AbstractDependency
    pkg::PkgSpec
end
BuildDependency(dep::AbstractString) = BuildDependency(PackageSpec(; name = dep))
is_host_dependency(::BuildDependency) = false
is_build_dependency(::BuildDependency) = true
is_runtime_dependency(::BuildDependency) = false

"""
    HostBuildDependency(dep::Union{PackageSpec,String})

Define a binary dependency that is necessary only to build the package.
Different from the [`BuildDependency`](@ref), the artifact for the host
platform will be installed, instead of that for the target platform.

The argument can be either a string with the name of the JLL package or a
`Pkg.PackageSpec`.
"""
struct HostBuildDependency <: AbstractDependency
    pkg::PkgSpec
end
HostBuildDependency(dep::AbstractString) = HostBuildDependency(PackageSpec(; name = dep))
is_host_dependency(::HostBuildDependency) = true
is_build_dependency(::HostBuildDependency) = true
is_runtime_dependency(::HostBuildDependency) = false

getpkg(d::AbstractDependency) = d.pkg
function getpkg(d::Dependency)
    # A Dependency can have a separate build_version, so replace that here:
    pkg = d.pkg
    if d.build_version !== nothing
        pkg = deepcopy(pkg)
        pkg.version = Pkg.Types.VersionSpec(d.build_version)
    end
    return pkg
end

getname(x::PkgSpec) = x.name
getname(x::AbstractDependency) = getname(getpkg(x))

# Wrapper around `Pkg.Types.registry_resolve!` which keeps the type of the
# dependencies.  TODO: improve this
function registry_resolve!(ctx, dependencies::Vector{<:AbstractDependency})
    resolved_dependencies = Pkg.Types.registry_resolve!(ctx, getpkg.(dependencies))
    for idx in eachindex(dependencies)
        dependencies[idx] = typeof(dependencies[idx])(resolved_dependencies[idx])
    end
    return dependencies
end

# We only want to update the registry once per run
registry_updated = false
function update_registry(ctx = Pkg.Types.Context(), outs = stdout)
    global registry_updated
    if !registry_updated
        Pkg.Registry.update(ctx,
            [Pkg.Types.RegistrySpec(uuid = "23338594-aafe-5451-b93e-139f81909106")];
            io=outs,
        )
        registry_updated = true
    end
end

function resolve_jlls(dependencies::Vector; ctx = Pkg.Types.Context(), outs=stdout)
    if isempty(dependencies)
        return true, Dependency[]
    end

    # Don't clobber caller
    # XXX: Coercion is needed as long as we support old-style dependencies.
    dependencies = deepcopy(coerce_dependency.(dependencies))

    # If all dependencies already have a UUID, return early
    if all(x->getpkg(x).uuid !== nothing, dependencies)
        return true, dependencies
    end

    # Resolve, returning the newly-resolved dependencies
    update_registry(ctx, outs)
    dependencies = registry_resolve!(ctx, dependencies)

    # But first, check to see if anything failed to resolve, and warn about it:
    all_resolved = true
    for dep in getpkg.(dependencies)
        if dep.uuid === nothing
            @warn("Unable to resolve $(getname(dep))")
            all_resolved = false
        end
    end
    return all_resolved, dependencies
end

# Add JSON serialization of dependencies
string_or_nothing(x) = isnothing(x) ? x : string(x)

major(v::VersionNumber) = v.major
minor(v::VersionNumber) = v.minor
patch(v::VersionNumber) = v.patch
major(v::Pkg.Types.VersionBound) = v.t[1]
minor(v::Pkg.Types.VersionBound) = v.t[2]
patch(v::Pkg.Types.VersionBound) = v.t[3]
__version(v::VersionNumber) = v
__version(v::Pkg.Types.VersionSpec) = v.ranges[1].lower
version(d::AbstractDependency) = __version(getpkg(d).version)
version(d::Dependency) = __version(d.pkg.version)

for (type, type_descr) in ((Dependency, "dependency"), (BuildDependency, "builddependency"), (HostBuildDependency, "hostdependency"))
    JSON.lower(d::type) = Dict("type" => type_descr,
                               "name" => d.pkg.name,
                               "uuid" => string_or_nothing(d.pkg.uuid),
                               "version-major" => major(version(d)),
                               "version-minor" => minor(version(d)),
                               "version-patch" => patch(version(d)))
end

# When deserialiasing the JSON file, the dependencies are in the form of
# dictionaries.  This function converts the dictionary back to the appropriate
# AbstractDependency.
function dependencify(d::Dict)
    if d["type"] in ("dependency", "builddependency", "hostdependency")
        uuid = isnothing(d["uuid"]) ? d["uuid"] : UUID(d["uuid"])
        version = VersionNumber(d["version-major"], d["version-minor"], d["version-patch"])
        version = version == v"0" ? nothing : version
        T = if d["type"] == "dependency"
            Dependency
        elseif d["type"] == "builddependency"
            BuildDependency
        elseif d["type"] == "hostdependency"
            HostBuildDependency
        end
        return T(PackageSpec(; name = d["name"], uuid = uuid, version = version))
    else
        error("Cannot convert to dependency")
    end
end

# XXX: compatibility functions.  These are needed until we support old-style
# dependencies.
coerce_dependency(dep::AbstractDependency) = dep
function coerce_dependency(dep)
    @warn "Using PackageSpec or string as dependency is deprecated, use Dependency instead"
    Dependency(dep)
end


text:
using Plots
pyplot()

function lkcontour(model,data,theta0,xx,yy)
	thetahat=mle(model,data)
	ff(x,y)=loglikelihood(model,data,[x,y])
	gg=[ff(x,y)::Float64 for x=xx,y=yy]

	qlevels=[0,.25,.5,.75,0.8,0.9,0.95,0.975,0.99,1]
	levels=quantile(vec(gg),qlevels)

	contour(xx,yy,gg,fill=true,levels=levels,color=ColorGradient(:heat,[0,0.98,1]))
	scatter!([theta0[1]],[theta0[2]],c=:green,leg=false)
	scatter!([thetahat[1]],[thetahat[2]],c=:blue,leg=false)
end


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

An immutable `AbstractVector` for which all elements are the same, represented
compactly and as a bitstype if the element type is a bitstype.
"""
struct ConstVector{T} <: AbstractVector{T}
    val::T
    length::Int64
end
Base.size(A::ConstVector) = (A.length, )
@propagate_inbounds Base.getindex(A::ConstVector, i::Int) = (@boundscheck checkbounds(A, i); A.val)
Base.IndexStyle(::Type{<:ConstVector}) = IndexLinear()
Base.unalias(dest, x::ConstVector) = x


text:
function _bracket(f, x, step, scale, maxiter)
    a, ya = x, f(x)
    b, yb = a + step, f(a + step)
    if yb > ya
        a, b = b, a
        ya, yb = yb, ya
        step = -step
    end
    iter = 0
    while iter < maxiter
        iter += 1
        c, yc = b + step, f(b + step)
        if yc > yb
            return a < c ? (a, c) : (c, a)
        end
        a, ya, b, yb = b, yb, c, yc
        step *= scale
    end
    return nothing
end

"""
    bracket(f::Function, x::T=0.0; kwargs...) where {T<:AbstractFloat}

Create a bracket around a local minimum from the intial point `x`. Return the bracket in a
tuple if successful, and `nothing` otherwise.

# Keywords
- `step=0.01`: the step size to expand the bracket
- `scale=2.0`: the scale factor for step size at each iteration
- `maxiter::Integer=100`: the maximum number of iterations
"""
function bracket(f::Function, x::T=0.0;
        step=0.01, scale=2.0, maxiter::Integer=100) where {T<:AbstractFloat}
    _bracket(f, x, T(step), T(scale), maxiter)
end


text:
using BlockBandedMatrices, BlockArrays, LinearAlgebra, MatrixFactorizations, Test
import BlockBandedMatrices: blockcolsupport

@testset "BlockBandedMatrix QR/QL" begin
    @testset "Square QR" begin
        N = 5
        A = BlockBandedMatrix{Float64}(undef, 1:N,1:N, (2,1))
        A.data .= randn.()

        F = qr(A)
        @test F isa MatrixFactorizations.QR{Float64,<:BlockSkylineMatrix}
        @test F.factors ≈ qr(Matrix(A)).factors
        @test F.τ ≈ LinearAlgebra.qrfactUnblocked!(Matrix(A)).τ

        Q,R = F
        Q̃,R̃ = qr(Matrix(A))
        @test R ≈ R̃
        @test Q ≈ Q̃

        b = randn(size(A,1))
        @test Q'b ≈ Q̃'b
        @test Q*b ≈ Q̃*b

        @test factorize(A) isa MatrixFactorizations.QR{Float64,<:BlockSkylineMatrix}
        @test F\b ≈ ldiv!(F, copy(b)) ≈ Matrix(A)\b ≈ A\b
    end

    @testset "Thin QR" begin
        N = 5
        A = BlockBandedMatrix{Float64}(undef, 1:N+1,1:N, (2,1))
        A.data .= randn.()

        F = qr(A)
        @test F.factors ≈ qr(Matrix(A)).factors
        @test F.τ ≈ LinearAlgebra.qrfactUnblocked!(Matrix(A)).τ

        Q,R = F
        Q̃,R̃ = qr(Matrix(A))
        @test R ≈ R̃
        @test Q ≈ Q̃

        b = randn(size(A,1))
        @test Q'b ≈ Q̃'b
        @test Q*b ≈ Q̃*b

        @test Q'b ≈ lmul!(Q', copy(b)) ≈ Q̃'*b

        @test F\b ≈ ldiv!(F, copy(b))[1:15] ≈ Matrix(A)\b ≈ A\b
    end

    @testset "Wide QR" begin
        N = 5
        A = BlockBandedMatrix{Float64}(undef, 1:N,1:N+1, (2,1))
        A.data .= randn.()

        F = qr(A)
        @test F.factors ≈ qr(Matrix(A)).factors
        @test F.τ ≈ LinearAlgebra.qrfactUnblocked!(Matrix(A)).τ

        Q,R = F
        Q̃,R̃ = qr(Matrix(A))
        @test R ≈ R̃
        @test Q ≈ Q̃

        b = randn(size(A,1))
        @test Q'b ≈ Q̃'b
        @test Q*b ≈ Q̃*b
        Q'b
        # @test_broken DimensionMismatch ldiv!(F, copy(b))

        @test_broken F\b ≈ ldiv!(F, copy(b)) ≈ Matrix(A)\b ≈ A\b
    end

    @testset "Square QL" begin
        N = 5
        A = BlockBandedMatrix{Float64}(undef, 1:N,1:N, (2,1))
        A.data .= randn.()

        F = ql(A)
        @test F.factors ≈ ql(Matrix(A)).factors
        @test F.τ ≈ MatrixFactorizations.qlfactUnblocked!(Matrix(A)).τ

        Q,L = F
        Q̃,L̃ = ql(Matrix(A))
        @test L ≈ L̃
        @test Q ≈ Q̃

        b = randn(size(A,1))
        @test Q'b ≈ Q̃'b
        @test Q*b ≈ Q̃*b

        b = randn(size(A,1))
        @test F\b ≈ ldiv!(F, copy(b)) ≈ Matrix(A)\b ≈ A\b
    end

    @testset "Thin QL" begin
        N = 5
        A = BlockBandedMatrix{Float64}(undef, 1:N+1,1:N, (2,1))
        A.data .= randn.()

        # F = ql(A)
        # @test F.factors ≈ ql(Matrix(A)).factors
        # @test F.τ ≈ MatrixFactorizations.qlfactUnblocked!(Matrix(A)).τ

        # Q,L = F
        # Q̃,L̃ = ql(Matrix(A))
        # @test L ≈ L̃
        # @test Q ≈ Q̃

        # b = randn(size(A,1))
        # @test Q'b ≈ Q̃'b
        # @test Q*b ≈ Q̃*b

        # b = randn(size(A,1))
        # @test F\b ≈ ldiv!(F, copy(b)) ≈ Matrix(A)\b ≈ A\b
    end

    @testset "Wide QL" begin
        N = 5
        A = BlockBandedMatrix{Float64}(undef, 1:N,1:N+1, (2,1))
        A.data .= randn.()

        @test_throws ArgumentError ql(A)
    end

    @testset "Complex QR/QL" begin
        @testset "Simple" begin
            A=BlockBandedMatrix{Float64}(I, [1,1],[1,1], (0,0))
            Af=qr(A)
            B=BlockBandedMatrix{ComplexF64}(I, [1,1],[1,1], (0,0))
            Bf=qr(B)
            @test Af.factors == Bf.factors
            @test Af.τ == Bf.τ
        end

        @testset "Off-diagonals" begin
            for qrl in [qr,ql]
                for T in [Float64, ComplexF64]
                    A = BlockBandedMatrix{T}(undef, [3,2,1,3],[3,2,1,3], (1,1))
                    A.data .= rand(T, size(A.data))
                    Af = qrl(A)
                    @test Af.factors isa BlockBandedMatrix

                    Am = Matrix(A)
                    Amf = qrl(Am)

                    @test Af.factors ≈ Amf.factors
                end
            end
        end
    end

    @testset "Fast QR \\" begin
        A = [1. 2; 3 4]; A = A + A'
        B = [5. 6; 7 8]
        N = 10_000;
        T = BlockBandedMatrix(mortar(Tridiagonal(fill(Matrix(B'),N-1), fill(zeros(2,2),N), fill(B,N-1))))
        z = 1+2im
        @time F = qr(T - z*I);
        @test blockbandwidths(UpperTriangular(F.factors)) == (0,2)
        @test blockcolsupport(UpperTriangular(F.factors),Block(4)) == Block.(2:4)
        V = view(F.factors,Block.(1:N), Block.(1:N))
        @test blocksize(V) == (N,N)
        @test @allocated(blocksize(V)) ≤ 40
        @test blockbandwidths(V) == (1,2)
        @test @allocated(blockbandwidths(V)) ≤ 40
        R = UpperTriangular(V)
        @test axes(R) == axes(V) == axes(F.factors)
        @test blocksize(R) == (N,N)
        @test @allocated(blocksize(R)) ≤ 40
        @test blockbandwidths(R) == (0,2)
        @test @allocated(blockbandwidths(R)) ≤ 40
        @test blockcolsupport(R,Block(4)) == Block.(2:4)
        @test @allocated(blockcolsupport(R,Block(4))) ≤ 40

        b = [1; zeros(size(T,1)-1)]
        B = [Matrix(I,2,2); zeros(size(T,1)-2,2)]
        @test ((T - z*I)\b)[1] ≈ (F\b)[1] ≈ (F \ B)[1,1] ≈ ((T - z*I)\B)[1,1] ≈ -0.1309123477325813 + 0.28471699370329884im
    end
end

# N = 500;
# A = BlockBandedMatrix{Float64}(undef, (1:N,1:N), (2,1));
# A.data .= randn.();
# @time F = qr(A); # 11s
# b = randn(size(A,1));
# @time F\b; # 0.6s


text:
### Julia OpenStreetMap Package ###
### MIT License                 ###
### Copyright 2014              ###

### Functions for finding highway intersections ###

### Generate a list of intersections ###
function findIntersections(highways::Dict{Int,Highway})
    seen = Set{Int}()
    intersections = Dict{Int,Intersection}()

    for hwy in values(highways)
        n_nodes = length(hwy.nodes)

        for i in 1:n_nodes
            node = hwy.nodes[i]

            if i == 1 || i == n_nodes || in(node, seen)
                get!(Intersection, intersections, node)
            else
                push!(seen, node)
            end
        end
    end

    for (hwy_key, hwy) in highways
        n_nodes = length(hwy.nodes)

        for i in 1:n_nodes
            node = hwy.nodes[i]

            if i == 1 || i == n_nodes || haskey(intersections, node)
                push!(intersections[node].highways, hwy_key)
            end
        end
    end

    return intersections
end

### Generate a new list of highways divided up by intersections
function segmentHighways(nodes, highways::Dict{Int,Highway}, intersections, classes, levels=Set(1:10))
    segments = Segment[]
    inters = Set(keys(intersections))

    for (i, class) in classes
        if in(class, levels) && haskey(highways,i)
            highway = highways[i]
            first = 1
            for j = 2:length(highway.nodes)
                if highway.nodes[first] != highway.nodes[j] && (in(highway.nodes[j], inters) || j == length(highway.nodes))
                    node0 = highway.nodes[first]
                    node1 = highway.nodes[j]
                    route_nodes = highway.nodes[first:j]
                    dist = distance(nodes, route_nodes)
                    s = Segment(node0, node1, route_nodes, dist, class, i, true)
                    push!(segments, s)

                    if !highway.oneway
                        s = Segment(node1, node0, reverse(route_nodes), dist, class, i, true)
                        push!(segments, s)
                    end
                    first = j
                end
            end
        end
    end

    return segments
end

### Generate a list of highways from segments, for plotting purposes
function highwaySegments( segments::Vector{Segment} )
    highways = Dict{Int,Highway}()

    for k = 1:length(segments)
        highways[k] = Highway("", 1, true, "", "", "", "$(segments[k].parent)", segments[k].nodes)
    end

    return highways
end

### Cluster highway intersections into higher-level intersections ###
# Note that there may be multiple intersection clusters containing the same
# streets, due to curved streets. Parameter max_dist controls how far apart an
# intersection must be from an existing cluster to create a new cluster.
function findIntersectionClusters( nodes::Dict{Int,ENU},
                                   intersections_in::Dict{Int,Intersection},
                                   highway_clusters::Vector{HighwaySet};
                                   max_dist=15.0 )
    hwy_cluster_mapping = Dict{Int,Int}()
    for k = 1:length(highway_clusters)
        hwys = [highway_clusters[k].highways...]
        for kk = 1:length(hwys)
            hwy_cluster_mapping[hwys[kk]] = k
        end
    end

    # Deep copy intersections dictionary and replace highways with highway
    # sets where available
    intersections = deepcopy(intersections_in)
    for (node,inter) in intersections
        hwys = [inter.highways...]
        for k = 1:length(hwys)
            if haskey(hwy_cluster_mapping,hwys[k])
                hwys[k] = hwy_cluster_mapping[hwys[k]]
            end
        end
        inter.highways = Set(hwys)
    end

    # Group intersections by number of streets contained
    hwy_counts = Vector{Int}[]
    for (node,inter) in intersections
        hwy_cnt = length(inter.highways)
        if hwy_cnt > length(hwy_counts)
            for k = (length(hwy_counts)+1):hwy_cnt
                push!(hwy_counts,Int[])
            end
        end
        push!(hwy_counts[hwy_cnt], node)
    end

    clusters = Set{Int}[]                   # Array of sets of highway IDs in each cluster
    clusters_nodes = Set{Int}[]             # Array of sets of node IDs in each cluster
    intersection_mapping = Dict{Int,Int}()  # [intersection id => index in `clusters`]

    for kk = 1:(length(hwy_counts)-1)
        # Start with intersections with most highways, as they are the best
        # "seeds" for new clusters because all intersection nodes added to the cluster
        # must have their highways be a subset of the highways already in the cluster.
        # Skip checking intersections with only 1 highway (road ends)
        k = length(hwy_counts)+1-kk

        for inter in hwy_counts[k]
            found = false
            for index = 1:length(clusters)
                if issubset(intersections[inter].highways,clusters[index])
                    # Check distance to cluster centroid
                    c = centroid(nodes,[clusters_nodes[index]...])
                    c_dist = distance(c,nodes[inter])
                    if c_dist < max_dist
                        intersection_mapping[inter] = index
                        clusters_nodes[index] = Set([clusters_nodes[index]...,inter])
                        found = true
                        break
                    end
                end
            end
            if !found
                push!(clusters,intersections[inter].highways)
                push!(clusters_nodes,Set(inter))
                intersection_mapping[inter] = length(clusters)
            end
        end
    end

    # Create new node at centroid of each intersection cluster
    cluster_map = Dict{Int,Int}()   # [Intersection Node ID => Cluster Node ID]
    for k = 1:length(clusters_nodes)
        if length(clusters_nodes[k]) > 1
            n = [clusters_nodes[k]...]
            c = centroid(nodes,n)
            cluster_node_id = addNewNode!(nodes,c)

            for j = 1:length(n)
                cluster_map[n[j]] = cluster_node_id
            end
        end
    end

    return cluster_map
end

### Replace Nodes in Highways Using Node Remapping
function replaceHighwayNodes!( highways::Dict{Int,Highway}, node_map::Dict{Int,Int} )
    for (key,hwy) in highways
        all_equal = true
        for k = 1:length(hwy.nodes)
            if haskey(node_map,hwy.nodes[k])
                hwy.nodes[k] = node_map[hwy.nodes[k]]
            end

            if k > 1 && hwy.nodes[k] != hwy.nodes[k-1]
                all_equal = false
            end
        end

        # If all nodes in hwy are now equal, delete it.
        if all_equal
            delete!(highways,key)
        end
    end
    return nothing
end