text:
## Grid generation function.
function make_grid(zsize, ktot)
    dz = zsize / ktot
    z = range(0.5*dz, step=dz, length=ktot) |> collect
    return z
end

## Settings.
float_type = Float32

settings_grid = Dict{String, Any}(
    "itot" => 256,
    "jtot" => 256,
    "ktot" => 256,

    "xsize" => 3200.,
    "ysize" => 3200.,
    "zsize" => 3200.)

settings_fields = Dict(
    "visc" => 5.,
    "alpha" => 9.81 / 300)

settings_boundary = Dict(
    "mom_bot_type" => "Dirichlet",
    "mom_top_type" => "Neumann",
    "s_bot_type" => "Neumann",
    "s_top_type" => "Neumann")

settings_timeloop = Dict(
    "start_time" => 0.,
    "end_time" => 3600.,
    "save_time" => 900.,
    "check_time" => 2.,
    "dt" => 2.)

settings_multidomain = Dict(
    "enable_nudge" => false)

settings_d01 = Dict(
    "grid" => settings_grid,
    "fields" => settings_fields,
    "boundary" => settings_boundary,
    "timeloop" => settings_timeloop,
    "multidomain" => settings_multidomain)

settings_d01["grid"]["z"] = make_grid(settings_d01["grid"]["zsize"], settings_d01["grid"]["ktot"])

settings = [ settings_d01 ]


text:
using Ratios, Test
using FixedPointNumbers

@testset "SimpleRatio" begin
    r = SimpleRatio(1,2)
    @test convert(Float64, r) == 0.5
    @test convert(Float32, r) == 0.5f0
    @test convert(BigFloat, r) == BigFloat(1)/2

    r2 = SimpleRatio(2,3)
    @test r*r2 == SimpleRatio(2,6) == SimpleRatio(1,3)
    @test r2*3 == 3*r2 == 2
    @test r*false == false*r == 0
    @test r/r2 == SimpleRatio(3,4)
    @test r/2 == SimpleRatio(1,4)
    @test 2/r == 4
    @test 4 == 2/r
    @test r+1 == 1+r == SimpleRatio(3,2)
    @test r-1 == SimpleRatio(-1,2)
    @test 1-r == r
    @test r+r2 == SimpleRatio(7,6)
    @test r-r2 == SimpleRatio(-1,6)
    @test r^2 == SimpleRatio(1,4)
    @test -r == SimpleRatio(-1,2)
    @test 0.2*r ≈ 0.1
    @test r == 0.5
    @test 0.5 == r

    @test_throws OverflowError -SimpleRatio(0x02,0x03)

    @test r + SimpleRatio(0x02,0x03) == SimpleRatio(7,6)

    @test SimpleRatio(11, 10) == 11//10
    @test 1//3 + SimpleRatio(1, 5) == 8//15

    @test isfinite(SimpleRatio(0,0)) == false
    @test isfinite(SimpleRatio(1,0)) == false
    @test isfinite(SimpleRatio(2,1)) == true

    @test SimpleRatio(5,3) * 0.035N0f8 == SimpleRatio{Int}(rationalize((5*0.035N0f8)/3))
    @test SimpleRatio(5,3) * 0.035N4f12 == SimpleRatio{Int}(rationalize((5*0.035N4f12)/3))
    @test SimpleRatio(5,3) * -0.03Q0f7 == SimpleRatio{Int}(rationalize((5.0*(-0.03Q0f7))/3))
    r = @inferred(SimpleRatio(0.75Q0f7))
    @test r == 3//4 && r isa SimpleRatio{Int16}
end

ext: jl


text:
module Move_SolidFileMesh

using Modia3D
using Modia3D.StaticArrays

vmat1 = Modia3D.Material(color="LightBlue" , transparency=0.5)    # material of SolidFileMesh
vmat2 = deepcopy(vmat1)                                           # material of convex decomposition of SolidFileMesh
vmat2.transparency = 0.7

font  = Modia3D.Font(fontFamily="Arial", bold=true, charSize=0.2)
rtext = [0.0, 1.0, 0.0]
filenamePascal = joinpath(Modia3D.path, "objects", "pascal", "pascal.obj")

r1 = [0.0, 0.0, 0.0]
r2 = [3.0, 0.0, 0.0]

@assembly FileMeshes begin
  world = Modia3D.Object3D(visualizeFrame=false)

  solidFileMesh = Modia3D.SolidFileMesh(filenamePascal,0.2)
  fileMesh1 = Modia3D.Object3D(world, Modia3D.Solid(                       solidFileMesh, nothing, vmat1       ); r=r1, fixed=false)
  fileMesh2 = Modia3D.Object3D(world, Modia3D.SolidWithConvexDecomposition(solidFileMesh, nothing, vmat1, vmat2); r=r2, fixed=false)
end

fileMeshes = FileMeshes(sceneOptions=Modia3D.SceneOptions(visualizeFrames=true,defaultFrameLength=0.7,enableContactDetection=false))

# Kinematic simulation
Modia3D.initAnalysis!(fileMeshes)

for time = range(0.0, stop=2.0, length=101)
  r1[3]  = 2*time
  Modia3D.set_r!(fileMeshes.fileMesh1, r1)

  Modia3D.updatePosition!(fileMeshes)

  # Visualize shapes
  Modia3D.visualize!(fileMeshes,time)
end
Modia3D.closeAnalysis!(fileMeshes)

println("... success of Move_SolidFileMesh.jl!")

end


text:
"""
    transition_matrices(m::Union{MDP,POMDP})
    transition_matrices(m; sparse=true)

Construct transition matrices for (PO)MDP m.

The returned object is an associative object (usually a Dict), where the keys are actions. Each value in this object is an AbstractMatrix where the row corresponds to the state index of s and the column corresponds to the state index of s'. The entry in the matrix is the probability of transitioning from state s to state s'.
"""
function transition_matrices(m::Union{MDP,POMDP}; sparse::Bool=true)
    transmats = POMDPModelTools.transition_matrix_a_s_sp(m)
    if !sparse
        transmats = [convert(Matrix, t) for t in transmats]
    end
    mtype = typeof(first(transmats))
    oa = ordered_actions(m)
    return Dict{actiontype(m), mtype}(oa[ai]=>transmats[ai] for ai in 1:length(actions(m)))
end

"""
    reward_vectors(m::Union{MDP, POMDP})

Construct reward vectors for (PO)MDP m.

The returned object is an associative object (usually a Dict), where the keys are actions. Each value in this object is an AbstractVector where the index corresponds to the state index of s and the entry is the reward for that state.
"""
function reward_vectors(m::Union{MDP,POMDP})
    d = Dict{actiontype(m), Vector{Float64}}()
    r = StateActionReward(m)
    for a in actions(m)
        rv = zeros(length(states(m)))
        for s in states(m)
            rv[stateindex(m, s)] = r(s, a)
        end
        d[a] = rv
    end
    return d
end


text:
print("""

Test Kabsch algorithm
=====================
""")

let a = [ 1  1 0
         2  1 0
         1+cos(pi/4) 1-sin(pi/4) 0
         1  0 0 ],
    b = [  1 1 0
         1 2 0
         1+cos(pi/4) 1+sin(pi/4)  0 ], # mean(b,1) 1.2357  1.56904  0.0
    sa = a[1:3, :],
    ma = mean(sa, 1) # mean x, y, z for a[1:3,:]

    # Center b, sa, a
    center!(b)
    sa[:,:] = sa .- ma
    a[:,:]  = a  .- ma # center != 0

    @test_approx_eq_eps mean(b,1) [0.0, 0.0, 0.0] 1e-13
    @test_approx_eq_eps mean(sa,1) [0.0, 0.0, 0.0] 1e-13

    F = kabsch(b, sa) # Reference: b
    R = sa * F
    RR = a * F

    @test_approx_eq R b
    @test_approx_eq RR[1:3,:] b

    @test_approx_eq RR[4,:] .- RR[1,:] [1.0, 0.0, 0.0]
    @test_approx_eq sqrt( (RR[1,1] - RR[4,1])^2 + (RR[1,2] - RR[4,2])^2 ) 1.0

    @test_approx_eq_eps rmsd(R, b) 0.0 1e-13
    @test_approx_eq_eps rmsd(RR[1:3,:], b) 0.0 1e-13
end

let P = [51.65 -1.90 50.07
         50.40 -1.23 50.65
         50.68 -0.04 51.54
         50.22 -0.02 52.85],
    Q = [51.30 -2.99 46.54
         51.09 -1.88 47.58
         52.36 -1.20 48.03
         52.71 -1.18 49.38], # P and Q are from BiomolecularStructures.jl' kabsch tests
    Qdistances = Float64[ sqrt(sumabs2(Q[j,:] .- Q[i,:])) for i in 1:4, j in 1:4 ]

    center!(P)
    center!(Q)

    @test_approx_eq_eps mean(P, 1) zeros(3) 1e-13
    @test_approx_eq_eps mean(Q, 1) zeros(3) 1e-13

    rotationmatrix = kabsch(P, Q)
    rotated = Q * rotationmatrix

    @test_approx_eq rmsd(P, rotated) 0.0030426652601371583

    # Internal distances mustn't change
    @test_approx_eq Float64[ sqrt(sumabs2(Q[j,:] .- Q[i,:])) for i in 1:4, j in 1:4 ] Qdistances # Translated
    @test_approx_eq Float64[ sqrt(sumabs2(rotated[j,:] .- rotated[i,:])) for i in 1:4, j in 1:4 ] Qdistances # Translated  and rotated
end

let P = [-1. 0. 0.
         0. 2. 0.
         0. 1. 0.
         0. 1. 1],
    Q = [0. -1. -1.
         0. -1. 0
         0. 0. 0.
         -1. 0. 0.]

    center!(P)
    center!(Q)
    rotated = Q * kabsch(P, Q)
    @test_approx_eq_eps rmsd(P, rotated) 0.695 0.001
end

print("""

RMSF
----
""")

let A = [ -1. -1. -1.
          -1.  0.  1. ],
    B = [  1.  1.  1.
          -1.  0.  1. ],
    w = [ .25 , .75 ]

    @test_throws ArgumentError mean_coordinates(Matrix{Float64}[A[1,:], B])
    @test_throws ArgumentError mean_coordinates(Matrix{Float64}[A[:,1:2], B])
    @test_throws ArgumentError mean_coordinates(Matrix{Float64}[A])
    @test_throws ArgumentError mean_coordinates(Matrix{Float64}[A,  B'])
    @test_throws ArgumentError mean_coordinates(Matrix{Float64}[A', B'])

    @test mean_coordinates(Matrix{Float64}[A, B]) == [ 0. 0. 0.
                                                      -1. 0. 1. ]

    @test mean_coordinates(Matrix{Float64}[A, B], w) == [ 0.5 0.5 0.5
                                                         -1.  0.  1. ]

    @test mean_coordinates(Matrix{Float64}[A, B], reverse(w)) == [ -0.5 -0.5 -0.5
                                                                   -1.   0.   1. ]

    @test rmsf(Matrix{Float64}[A, B]) == [sqrt(3), 0.]
    @test rmsf(Matrix{Float64}[A, B], w) == [sqrt(0.25*6.75 + 0.75*0.75), 0.]

end

print("""

Superimpose PDBs and RMSD
-------------------------
""")

let hemoglobin = read(joinpath(pwd(), "data", "2hhb.pdb.gz"), PDBFile, group="ATOM", model="1")

    α1 = @residues hemoglobin model "1" chain "A" group "ATOM" residue "*"
    α2 = @residues hemoglobin model "1" chain "C" group "ATOM" residue "*"
    β1 = @residues hemoglobin model "1" chain "B" group "ATOM" residue "*"
    β2 = @residues hemoglobin model "1" chain "D" group "ATOM" residue "*"

    a1, a2, rα = superimpose(α1, α2)

    @test_approx_eq_eps rα 0.230 0.001 # Bio3D RMSD

    b1, b2, rβ = superimpose(β1, β2)

    @test_approx_eq_eps rβ 0.251 0.001 # Bio3D & Chimera's MatchMaker RMSD

    @test length(a1) == length(α1)
    @test length(a2) == length(α2)
    @test length(b1) == length(β1)
    @test length(b2) == length(β2)

    @test a1 != α1
    @test a2 != α2

    @test b1 != β1
    @test b2 != β2

    print("""
    rmsd with Vector{PDBResidue}
    """)

    @test_approx_eq_eps rmsd(α1, α2) 0.230 0.001
    @test_approx_eq_eps rmsd(β1, β2) 0.251 0.001

    @test_approx_eq_eps rmsd(a1, a2, superimposed=true) 0.230 0.001
    @test_approx_eq_eps rmsd(b1, b2, superimposed=true) 0.251 0.001

    ca_a1 = CAmatrix(a1)
    ca_a2 = CAmatrix(a2)
    ca_b1 = CAmatrix(b1)
    ca_b2 = CAmatrix(b2)

    print("""
    Centered
    """)

    @test_approx_eq_eps mean(ca_a1,1) zeros(3) 1e-13
    @test_approx_eq_eps mean(ca_a2,1) zeros(3) 1e-13
    @test_approx_eq_eps mean(ca_b1,1) zeros(3) 1e-13
    @test_approx_eq_eps mean(ca_b2,1) zeros(3) 1e-13

    @test PDB._iscentered(ca_a1)
    @test PDB._iscentered(ca_a2)
    @test PDB._iscentered(ca_b1)
    @test PDB._iscentered(ca_b2)

    print("""
    RMSD
    """)

    @test rα == rmsd(ca_a1, ca_a2)
    @test rβ == rmsd(ca_b1, ca_b2)

    print("""
    coordinatesmatrix, centeredresidues and centeredcoordinates
    """)

    @test_approx_eq coordinatesmatrix(centeredresidues(α1)) coordinatesmatrix(a1)
    @test_approx_eq coordinatesmatrix(centeredresidues(β1)) coordinatesmatrix(b1)

    @test_approx_eq centeredcoordinates(α1) coordinatesmatrix(a1)
    @test_approx_eq centeredcoordinates(β1) coordinatesmatrix(b1)

    print("""
    Superimpose with centered PDBs
    """)

    b12, b22, rβ2 = superimpose(b1, b2)
    @test_approx_eq rβ2 rβ

    print("""
    RMSF
    """)

    @test rmsf(Vector{PDBResidue}[α1, α2]) == rmsf(Vector{PDBResidue}[α1, α2], [.5, .5])
    @test rmsf(Vector{PDBResidue}[β1, β2]) == rmsf(Vector{PDBResidue}[β1, β2], [.5, .5])

end


text:
## load packages
using Plots
using Distributions, Distances, DataFrames, StatsBase

"""
makeCitizens(n::Int64)

Create a population of `n` individuals uniformly distributed in a rectangle.

Returns the spatial positions of the citizens, their Euclidean distances from
each other, and initializes their disease status to being unaffected (0).

This function can be generalized to include other spatial distributions, and
other ways to calculate distance between citizens.
"""
function makeCitizens(n::Int64)
    x = rand(Uniform(),n);
    y = rand(Uniform(),n);
    citizenPos = [x y];
    citizenDists=pairwise(Euclidean(),citizenPos,dims=1)
    citizenStatus = DataFrame(ind=1:n,infected=zeros(Int64,n))
    return citizenStatus,citizenDists,citizenPos
end

"""
makeCitizens(citizenPos::Matrix{Float64})

Create a population of `n` individuals uniformly distributed in a rectangle.

Returns the spatial positions of the citizens, their Euclidean distances from
each other, and initializes their disease status to being unaffected (0).
"""
function makeCitizens(citizenPos::Matrix{Float64})
    n = size(citizenPos,1)
    citizenDists = pairwise(Euclidean(),citizenPos,dims=1)
    citizenStatus = DataFrame(ind=1:n,infected=zeros(Int64,n))
    return citizenStatus,citizenDists,citizenPos
end
"""
infectCitizens!(citizenStatus::DataFrame)

Picks `nInfect` random individuals in `citizenStatus` to be infected.  This
function could be generalized by picking someone in a defined region or
with defined characteristics.
"""
function infectCitizens!(citizenStatus::DataFrame,nInfect::Int64)
    n = size(citizenStatus,1)
    citizenStatus[sample(1:n,nInfect),2] .= 1
end

"""
updateStatus!(status::DataFrame,distances::Matrix{Float64},
             time::Int64,infectionDist::Float64,infectionWindow::Int64,
             infectionProb::Float64)

Updates the infection `status` data frame at a given `time` using information
on `infectionDist` (radius of neighborhood which a person can infect),
`infectionWindow` (how long an infected person is infectious), `infectionProb`
(the probability of infecting a susceptible person in a given time period).

The infection model can be tweaked to make it more realistic.
"""
function updateStatus!(status::DataFrame,distances::Matrix{Float64},
        time::Int64,infectionDist::Float64,
        infectionWindow::Int64,infectionProb::Float64)

    # individuals who are still infectious
    whoCanInfect = status.ind[(status.infected.>0).&
                   ((time.-status.infected).<=infectionWindow)]
    # individuals who are susceptible
    whoSusceptible = status.ind[status.infected.==0]
    # which individuals can infect whom
    transmissionPossible = (distances[whoCanInfect,whoSusceptible].<=infectionDist)
    # calculate probability of infection for each person in a given time period
    prob = 1.0.-(1.0.-infectionProb).^(sum(transmissionPossible,dims=1))
    # randomly decide based on infection probability, who is newly infected
    status.infected[status.ind[whoSusceptible[rand.(Bernoulli.(prob))[1,:]]]] .= time
end

"""
followCitizens!(citizens::DataFrame,distances::Matrix{Float64},ntime::Int64=365,
                        infectionDist::Float64,
                        infectionWindow::Int64,infectionProb::Float64)

Follow the citizens in the data drame `citizens` for `ntime` time periods assuming
`infectionDist`, `infectionWindow` an `infectionProb` as in updateStatus function.
"""
function followCitizens!(citizens::DataFrame,distances::Matrix{Float64},ntime::Int64=365,
                        infectionDist::Float64=0.1,
                        infectionWindow::Int64=1,infectionProb::Float64=0.2)
    # print expected number of infections per infected individual
    r0 = round(size(citizens,1)*pi*infectionDist^2*infectionWindow*infectionProb,digits=2)
    for i in 2:ntime
         updateStatus!(citizens,distances,i,infectionDist,
            infectionWindow,infectionProb)
    end
    return r0
end

###
## convert trues to 1
###
function trueone(x)
    if x==true
        return 0
    else
        return 1
    end
end

###################################################

"""
uninfected(status::DataFrame)

Returns number uninfected at end of follow up period.
"""
function uninfected(status::DataFrame)
    return mean(status.infected.>0)
end

"""
newInfections(status::DataFrame)

Make histogram of new infections.
"""
function newInfections(status::DataFrame)
    histogram(status.infected,bins=range(1,maximum(status.infected),step=1))
end

"""
cumulativeInfections(status::DataFrame)

Plot cumulative number of infection as a proportion of the population.
"""
function cumulativeInfections(status::DataFrame)
    plot(1:365,cumsum(counts(status.infected,1:365))./size(status,1),ylim=(0,1))
end

"""
infectionSpread(status::DataFrame,pos::Matrix{Float64})

Make a animation of the spread of the infection using the `status` data and
spatial positions in `pos`.
"""
function infectionSpread(status::DataFrame,pos::Matrix{Float64},ntime::Int64=365)
    spread = @animate for i in 1:ntime
      scatter(pos[:,1],pos[:,2],
              color=trueone.((status.infected.<=i).&(status.infected.>0)),label="")
    end
    return spread
end


text:
mutable struct UnfoldGenerateState{T,U}
    score::Float64
    noise::Float64
    weight::Float64
    subtraces::Vector{U}
    retval::Vector{T}
    num_nonempty::Int
    state::T
end

function process!(gen_fn::Unfold{T,U}, params::Tuple, choices::ChoiceMap,
                  key::Int, state::UnfoldGenerateState{T,U}) where {T,U}
    local subtrace::U
    local new_state::T
    kernel_args = (key, state.state, params...)
    submap = get_submap(choices, key)
    (subtrace, weight) = generate(gen_fn.kernel, kernel_args, submap)
    state.weight += weight
    state.noise += project(subtrace, EmptySelection())
    state.num_nonempty += (isempty(get_choices(subtrace)) ? 0 : 1)
    state.score += get_score(subtrace)
    state.subtraces[key] = subtrace
    new_state = get_retval(subtrace)
    state.state = new_state
    state.retval[key] = new_state
end

function generate(gen_fn::Unfold{T,U}, args::Tuple, choices::ChoiceMap) where {T,U}
    len = args[1]
    init_state = args[2]
    params = args[3:end]
    state = UnfoldGenerateState{T,U}(0., 0., 0.,
        Vector{U}(undef,len), Vector{T}(undef,len), 0, init_state)
    for key=1:len
        process!(gen_fn, params, choices, key, state)
    end
    trace = VectorTrace{UnfoldType,T,U}(gen_fn,
        PersistentVector{U}(state.subtraces), PersistentVector{T}(state.retval),
        args, state.score, state.noise, len, state.num_nonempty)
    (trace, state.weight)
end


text:
function local_update!(model::Potts, T::Real, Js::AbstractArray)
    rng = model.rng
    nsites = numsites(model)
    nbonds = numbonds(model)
    mbeta = -1.0/T

    @inbounds for site in 1:nsites
        center = model.spins[site]
        new_center = mod1(center+rand(rng, 1:(model.Q-1)), model.Q)
        de = 0.0
        for (n,b) in neighbors(model, site)
            de += ifelse(center == model.spins[n], Js[bondtype(model, b)], 0.0)
            de -= ifelse(new_center == model.spins[n], Js[bondtype(model, b)], 0.0)
        end
        if rand(rng) < exp(mbeta*de)
            model.spins[site] = new_center
        end
    end

    return nothing
end

function SW_update!(model::Potts, T::Real, Js::AbstractArray)
    rng = model.rng
    ps = -expm1.((-1.0/T).*Js)
    nsites = numsites(model)
    nbonds = numbonds(model)
    nbt = numbondtypes(model)
    activated_bonds = zeros(Int,nbt)
    uf = UnionFind(nsites)
    @inbounds for bond in bonds(model)
        s1,s2 = source(bond), target(bond)
        bt = bondtype(bond)
        if model.spins[s1] == model.spins[s2] && rand(rng) < ps[bt]
            activated_bonds[bt] += 1
            unify!(uf, s1,s2)
        end
    end
    nc = clusterize!(uf)
    clustersize = zeros(Int, nc)
    clusterspin = rand(rng,1:model.Q, nc)

    @inbounds for site in 1:nsites
        id = clusterid(uf, site)
        model.spins[site] = clusterspin[id]
        clustersize[id] += 1
    end
    return SWInfo(activated_bonds, clustersize, clusterspin)
end

function Wolff_update!(model::Potts, T::Real, Js::AbstractArray)
    rng = model.rng
    ps = -expm1.((-1.0/T).*Js)
    nsites = numsites(model)

    clustersize = 0
    st = Stack(Deque{Int}())
    center = rand(rng, 1:nsites)
    sp = model.spins[center]
    newsp = mod1(sp+rand(rng, 1:(model.Q-1)), model.Q)
    model.spins[center] = newsp
    push!(st, center)
    @inbounds while !isempty(st)
        clustersize += 1
        s = pop!(st)
        for (n,b) in neighbors(model, s)
            bt = bondtype(model,b)
            if model.spins[n] == sp && rand(rng) < ps[bt]
                model.spins[n] = newsp
                push!(st, n)
            end
        end
    end
    return nothing
end

@gen_convert_parameter(Potts, ("T", 1, 1.0), ("J", numbondtypes, 1.0))


text:
# load simulated data
using NeuralNets
using HDF5, JLD
# reconstruct RunResults
results = load("trial1.jld","results")
include("../src/apma2821v/plotting.jl")
plot_vision_rasters(results)
println("Press [Enter] to continue...")
readline()


text:
using InvertibleNetworks, LinearAlgebra, Test
using PyPlot, Flux, Random, JLD
import Flux.Optimise.update!

# Random seed
Random.seed!(100)

#############################################################################################################
# Utils

# Concatenate states Zi and final output
function cat_states(Z_save, X)
    Y = []
    for j=1:length(Z_save)
        Y = cat(Y, vec(Z_save[j]); dims=1)
    end
    Y = cat(Y, vec(X); dims=1)
    return Float32.(Y)  # convert to Array{Float32, 1}
end

# Split 1D vector in latent space back to states Zi
function split_states(Y, Z_dims)
    L = length(Z_dims)
    Z_save = Array{Array}(undef, L)
    count = 1
    for j=1:L
        Z_save[j] = reshape(Y[count: count + prod(Z_dims[j])-1], Z_dims[j])
        count += prod(Z_dims[j])
    end
    X = reshape(Y[count: count + prod(Z_dims[end])-1], Int.(Z_dims[end].*(.5, .5, 4, 1)))
    return Z_save, X
end

#############################################################################################################
# Data and network layers

# Load MNIST
X = load("/home/pwitte3/InvertibleNetworks/data/mnist_training.jld")["X"]

# Data
nx = 32
ny = 32
n_in = 1
batchsize = 10
k = 3   # kernel size
s = 1   # stride
p = 1   # padding
hidden_factor = 2
α = 0.1f0

# Network
steps_per_scale = 10

# Scale 1
H1 = Array{HyperbolicLayer}(undef, steps_per_scale)
for j=1:steps_per_scale
    H1[j] = HyperbolicLayer(16, 16, 2, batchsize, k, s, p; action="same", α=α, hidden_factor=hidden_factor)
end
Hs1 = HyperbolicLayer(16, 16, 2, batchsize, k, s, p; action="down", α=α, hidden_factor=hidden_factor)

# Scale 2
H2 = Array{HyperbolicLayer}(undef, steps_per_scale)
for j=1:steps_per_scale
    H2[j] = HyperbolicLayer(8, 8, 4, batchsize, k, s, p; action="same", α=α, hidden_factor=hidden_factor)
end
Hs2 = HyperbolicLayer(8, 8, 4, batchsize, k, s, p; action="down", α=α, hidden_factor=hidden_factor)

# Scale 3
H3 = Array{HyperbolicLayer}(undef, steps_per_scale)
for j=1:steps_per_scale
    H3[j] = HyperbolicLayer(4, 4, 8, batchsize, k, s, p; action="same", α=α, hidden_factor=hidden_factor)
end
Hs3 = HyperbolicLayer(4, 4, 8, batchsize, k, s, p; action="down", α=α, hidden_factor=hidden_factor)

# Save and dimensions
X_save = Array{Array}(undef, 3)
X_dims = Array{Tuple}(undef, 3)

#############################################################################################################
# Network passes

# Forward pass
function forward(X)

    # Initial transform to increase no. of channels
    logdet = 0f0
    X = wavelet_squeeze(X)
    X_prev, X_curr = tensor_split(X)

    # Scale 1
    for j=1:steps_per_scale
        X_prev, X_curr = H1[j].forward(X_prev, X_curr)
    end
    X_prev, X_curr = Hs1.forward(X_prev, X_curr)
    X_save[1] = X_prev
    X_dims[1] = size(X_prev)

    # Scale 2
    X_prev, X_curr = tensor_split(X_curr)
    for j=1:steps_per_scale
        X_prev, X_curr = H2[j].forward(X_prev, X_curr)
    end
    X_prev, X_curr = Hs2.forward(X_prev, X_curr)
    X_save[2] = X_prev
    X_dims[2] = size(X_prev)

    # Scale 3
    X_prev, X_curr = tensor_split(X_curr)
    for j=1:steps_per_scale
        X_prev, X_curr = H3[j].forward(X_prev, X_curr)
    end
    X_prev, X_curr = Hs3.forward(X_prev, X_curr)
    X_save[3] = X_prev
    X_dims[3] = size(X_prev)

    # Concatenate state from all scales
    X = cat_states(X_save, X_curr)
    return X, logdet
end

function inverse(X)
    X_save, X_curr = split_states(X, X_dims)

    # Scale 3
    X_prev, X_curr = Hs3.inverse(X_save[3], X_curr)
    for j=steps_per_scale:-1:1
        X_prev, X_curr = H3[j].inverse(X_prev, X_curr)
    end
    X_curr = tensor_cat(X_prev, X_curr)

    # Scale 2
    X_prev, X_curr = Hs2.inverse(X_save[2], X_curr)
    for j=steps_per_scale:-1:1
        X_prev, X_curr = H2[j].inverse(X_prev, X_curr)
    end
    X_curr = tensor_cat(X_prev, X_curr)

    # Scale 1
    X_prev, X_curr = Hs1.inverse(X_save[1], X_curr)
    for j=steps_per_scale:-1:1
        X_prev, X_curr = H1[j].inverse(X_prev, X_curr)
    end
    X_curr = tensor_cat(X_prev, X_curr)

    return X_curr
end

# Current sample
Xi = X[:,:,:,1:10]
Y, logdet = forward(Xi)
Xi_ = inverse(Y)

#############################################################################################################
# Training

# # Loss
# function loss(H, X)
#     Y, logdet = H.forward(X)
#     f = -log_likelihood(Y) - logdet
#     ΔY = -∇log_likelihood(Y)
#     H.backward(ΔY, Y)
#     return f
# end

# # Training
# ntrain = 60000
# maxiter = 1000
# opt = Flux.ADAM(1f-2)
# fval = zeros(Float32, maxiter)
# Params = get_params(H)

# for j=1:maxiter

#     # Evaluate objective and gradients
#     idx = randperm(ntrain)[1:batchsize]
#     fval[j] = loss(H, X[:,:,:,idx])
#     mod(j, 10) == 0 && (print("Iteration: ", j, "; f = ", fval[j], "\n"))

#     # Update params
#     for p in Params
#         update!(opt, p.data, p.grad)
#     end
#     clear_grad!(H)
# end

# # Test
# testsize = 1
# idx = randperm(ntrain)[1]
# Xi = X[:,:,:,idx:idx]
# Yi = H.forward(Xi)[1]

# Yi_ = randn(Float32, nx, ny, n_in, testsize)
# Xi_ = H.inverse(Yi_)

# figure(figsize=[8,8])
# i=1
# subplot(2,2,1); imshow(Xi[:,:,1,1]); title(L"Data space: $x \sim \hat{p}_X$")
# subplot(2,2,2); imshow(Yi[:,:,1,1]); title(L"Latent space: $z = f(x)$")
# subplot(2,2,3); imshow(Xi_[:,:,1,i]); title(L"Data space: $x = f^{-1}(z)$")
# subplot(2,2,4); imshow(Yi_[:,:,1,i]); title(L"Latent space: $z \sim \hat{p}_Z$")

max_stars_repo_name: slimgroup/LoopUnrolledSeismicImaging

max_issues_repo_name: slimgroup/LoopUnrolledSeismicImaging

max_forks_repo_name: slimgroup/LoopUnrolledSeismicImaging


text:
function generate_autodiff_solver_fn(eval_obj,eval_cons,cs_eqs,cs_ineqs;hessian=true)
    function solver_fn(x)
        J = eval_obj(x)
        gJ = ForwardDiff.gradient(eval_obj, x)
        if hessian
            HJ = ForwardDiff.hessian(eval_obj, x)
        else
            HJ = zeros(length(x),length(x))
        end

        g = eval_cons(x)
        dgdx = ForwardDiff.jacobian(eval_cons, x)

        ceq = g[cs_eqs]
        c = g[cs_ineqs]
        gceq = dgdx[cs_eqs,:]
        gc = dgdx[cs_ineqs,:]

        J, ceq, c, gJ, gceq, gc, HJ
    end

    solver_fn
end

function generate_autodiff_solver_fn(eval_obj,fres,fcfg,eval_cons,gres,gcfg,cs_eqs,cs_ineqs;hessian=true)
    function solver_fn(x)
        ForwardDiff.hessian!(fres, eval_obj, x, fcfg)
        J = DiffResults.value(fres)
        gJ = DiffResults.gradient(fres)
        if hessian
            HJ = DiffResults.hessian(fres)
        else
            HJ = zeros(length(x),length(x))
        end

        ForwardDiff.jacobian!(gres, eval_cons, x, gcfg)
        g = DiffResults.value(gres)
        dgdx = DiffResults.jacobian(gres)

        ceq = g[cs_eqs]
        c = g[cs_ineqs]
        gceq = dgdx[cs_eqs,:]
        gc = dgdx[cs_ineqs,:]

        J, ceq, c, gJ, gceq, gc, HJ
    end

    solver_fn
end

function generate_autodiff_solver_fn(eval_obj,fres,eval_cons,gres,cs_eqs,cs_ineqs;hessian=true)
    function solver_fn(x)
        ForwardDiff.hessian!(fres, eval_obj, x)
        J = DiffResults.value(fres)
        gJ = DiffResults.gradient(fres)
        if hessian
            HJ = DiffResults.hessian(fres)
        else
            HJ = zeros(length(x),length(x))
        end

        ForwardDiff.jacobian!(gres, eval_cons, x)
        g = DiffResults.value(gres)
        dgdx = DiffResults.jacobian(gres)

        ceq = g[cs_eqs]
        c = g[cs_ineqs]
        gceq = dgdx[cs_eqs,:]
        gc = dgdx[cs_ineqs,:]

        J, ceq, c, gJ, gceq, gc, HJ
    end

    solver_fn
end

function generate_autodiff_solver_fn(eval_obj,eval_cons,cs_eqs,cs_ineqs,vs_num_vars;hessian=true)
    cs_num_cons = length(cs_eqs) + length(cs_ineqs)

    fres = DiffResults.HessianResult(zeros(vs_num_vars))
    fcfg = ForwardDiff.HessianConfig(eval_obj, fres, zeros(vs_num_vars))
    gres = DiffResults.JacobianResult(zeros(cs_num_cons), zeros(vs_num_vars))
    gcfg = ForwardDiff.JacobianConfig(eval_cons, zeros(vs_num_vars))

    generate_autodiff_solver_fn(eval_obj,fres,fcfg,eval_cons,gres,gcfg,cs_eqs,cs_ineqs,hessian=hessian)
end


text:
using PiecewiseIncreasingRanges, Base.Test, Grid

function test(rgs, divisor...)
    vcrgs = vcat(rgs...)
    !isempty(divisor) && (vcrgs /= divisor[1])
    rg = PiecewiseIncreasingRange(rgs, divisor...)
    @test length(rg.ranges) == 4
    @test vcrgs == rg

    yi = InterpGrid(convert(Vector{Float64}, vcrgs), BCnil, InterpLinear)
    @test_approx_eq resample(rg, 3//7) yi[1:7//3:length(yi)]
    @test_approx_eq resample(rg, 7//3) yi[1:3//7:length(yi)]
    @test_approx_eq resample(rg, 5//9) yi[1:9//5:length(yi)]
    @test_approx_eq resample(rg, 61//3) yi[1:3//61:length(yi)]

    for i = 1:length(rg)
        @test searchsortedfirst(rg, rg[i]) == i
        @test searchsortedfirst(rg, rg[i]-1//17) == i
        @test searchsortedfirst(rg, rg[i]+1//17) == i+1

        @test searchsortedlast(rg, rg[i]) == i
        @test searchsortedlast(rg, rg[i]-1//17) == i-1
        @test searchsortedlast(rg, rg[i]+1//17) == i

        for within_half_sample in (false, true)
            @test findnearest(rg, rg[i], within_half_sample) == i
            @test findnearest(rg, rg[i]+1//16, within_half_sample) == i
            @test findnearest(rg, rg[i]-1//17, within_half_sample) == i
            @test findnearest(rg, rg[i]+1//17, within_half_sample) == i
        end

        for j = i:length(rg)
            @test vcrgs[i:j] == rg[i:j]
            for k = 1:i-j+1
                @test vcrgs[i:k:j] == rg[i:k:j]
            end
        end
    end

    @test findnearest(rg, -1) == 1
    @test_throws NoNearestSampleError findnearest(rg, -1, true)
    @test findnearest(rg, 15.9) == searchsortedfirst(rg, 15)
    @test findnearest(rg, 16.1) == searchsortedfirst(rg, 17)
    @test_throws NoNearestSampleError findnearest(rg, 15.9, true)
    @test_throws NoNearestSampleError findnearest(rg, 16.1, true)
    @test findnearest(rg, 19) == length(rg)
    @test_throws NoNearestSampleError findnearest(rg, 19, true)
end

test(StepRange{Int,Int}[0:4:40, 41:1:80, 82:2:112, 114:2:120, 136:4:144], 8)
test(StepRange{Rational{Int},Rational{Int}}[0:1//2:5, 5+1//8:1//8:10, 10+1//4:1//4:14, 14+1//4:1//4:15, 17:1//2:18])
test(FloatRange{Float64}[0:1//2:5., 5+1//8:1//8:10., 10+1//4:1//4:14., 14+1//4:1//4:15., 17:1//2:18.])

# Empty test
rg = PiecewiseIncreasingRange(StepRange{Rational{Int},Rational{Int}}[])
@test searchsortedfirst(rg, 1) == 1
@test searchsortedlast(rg, 1) == 1
@test_throws NoNearestSampleError findnearest(rg, 1)
rg = PiecewiseIncreasingRange(UnitRange{Int}[0:-1])
@test isempty(rg)

# Test with non-monotonic range
@test_throws ArgumentError PiecewiseIncreasingRange(StepRange{Rational{Int},Rational{Int}}[0:1//2:1, 3//4:1//2:4])
@test_throws ArgumentError PiecewiseIncreasingRange(StepRange{Rational{Int},Rational{Int}}[0:-1//2:-1, 10:1//2:8])
@test_throws ArgumentError PiecewiseIncreasingRange(StepRange{Rational{Int},Rational{Int}}[0:1//2:1, 10:-1//2:8])
@test_throws ArgumentError PiecewiseIncreasingRange(FloatRange{Float64}[0:0.5:1, 0.75:0.5:4])
@test_throws ArgumentError PiecewiseIncreasingRange(FloatRange{Float64}[0:-0.5:-1, 10:0.5:8])
@test_throws ArgumentError PiecewiseIncreasingRange(FloatRange{Float64}[0:0.5:1, 10:-0.5:8])

# Test with empty ranges interspersed with non-empty ranges
rg = PiecewiseIncreasingRange(UnitRange{Int}[0:-1, 1:0, 0:-1, 1:3, 0:-1, 5:10])
@test rg.ranges == UnitRange{Int}[1:3, 5:10]


text:
# TODO(tim): support NA
# TODO(tim): support Nullable
# TODO(tim): support default value for missing keys

immutable CategoricalDiscretizer{N,D} <: AbstractDiscretizer{N,D}
    n2d :: Dict{N,D}    # maps natural to discrete
    d2n :: Dict{D,N}    # maps discrete to natural
end

function CategoricalDiscretizer{N,D}(natural_to_discrete::Dict{N,D})
    d2n = Dict{D,N}()
    for (k,v) in natural_to_discrete
        d2n[v] = k
    end
    CategoricalDiscretizer{N,D}(natural_to_discrete, d2n)
end
function CategoricalDiscretizer{N, D<:Integer}(data::AbstractArray{N}, ::Type{D}=Int)
    # build a label mapping N -> D <: Integer
    i = zero(D)
    n2d = @compat Dict{N, D}()
    for x in data
        if !haskey(n2d,x)
            n2d[x] = (i += 1)
        end
    end
    CategoricalDiscretizer(n2d)
end

supports_encoding{N,D}(cd::CategoricalDiscretizer{N,D}, x::N) = haskey(cd.n2d, x)
supports_decoding{N,D}(cd::CategoricalDiscretizer{N,D}, d::D) = 1 ≤ d ≤ nlabels(cd)

encode{N,D}(cd::CategoricalDiscretizer{N,D}, x::N) = cd.n2d[x]::D
encode{N,D}(cd::CategoricalDiscretizer{N,D}, x) = cd.n2d[convert(N,x)]::D
function encode{N,D}(cd::CategoricalDiscretizer{N,D}, data::AbstractArray)
    arr = Array{D}(length(data))
    for (i,x) in enumerate(data)
        arr[i] = encode(cd, x)
    end
    reshape(arr, size(data))
end

decode{N,D}(cd::CategoricalDiscretizer{N,D}, x::D) = cd.d2n[x]::N
decode{N,D}(cd::CategoricalDiscretizer{N,D}, x) = cd.d2n[convert(D,x)]::N
function decode{N,D}(cd::CategoricalDiscretizer{N,D}, data::AbstractArray{D})
    arr = Array{N}(length(data))
    for (i,d) in enumerate(data)
        arr[i] = decode(cd, d)
    end
    reshape(arr, size(data))
end

nlabels(cd::CategoricalDiscretizer) = length(cd.n2d)


text:
module LearningStrategies

import LearnBase: learn!, update!
using LinearAlgebra

export
    LearningStrategy, MetaStrategy, strategy, Offline, InfiniteNothing,
    # LearningStrategies
    Verbose, MaxIter, TimeLimit, Converged, ConvergedTo, IterFunction, Tracer, Breaker,
    # functions
    setup!, update!, hook, finished, cleanup!, learn!

#---------------------------------------------------------------------# LearningStrategy interface
"""
A LearningStrategy should implement at least one of the following methods:

- `setup!(strat, model, data)`
- `cleanup!(strat, model)`
- `hook(strat, model, i)`
- `finished(strat, model, data, i)`
- `update!(model, strat, item)`
"""
abstract type LearningStrategy end

setup!(s::LearningStrategy, model) = nothing
setup!(s::LearningStrategy, model, data) = setup!(s, model)

update!(model, s::LearningStrategy) = nothing
update!(model, s::LearningStrategy, item) = update!(model, s)
update!(model, s::LearningStrategy, i, item) = update!(model, s, item)

hook(s::LearningStrategy, model) = nothing
hook(s::LearningStrategy, model, i) = hook(s, model)
hook(s::LearningStrategy, model, data, i) = hook(s, model, i)

finished(s::LearningStrategy, model) = false
finished(s::LearningStrategy, model, i) = finished(s, model)
finished(s::LearningStrategy, model, data, i) = finished(s, model, i)

cleanup!(s::LearningStrategy, model) = nothing

#-----------------------------------------------------------------------# MetaStrategy
"""
    MetaStrategy(strats::LearningStrategy...)

A collection of learning strategies in a type-stable way.
"""
struct MetaStrategy{T <: Tuple} <: LearningStrategy
    strategies::T
end

MetaStrategy(s::LearningStrategy...) = MetaStrategy(s)

function Base.show(io::IO, s::MetaStrategy)
    print(io, "MetaStrategy")
    for s in s.strategies
        print(io, "\n  > ", s)
    end
end

setup!(ms::MetaStrategy, model, data) = foreach(s -> setup!(s, model, data), ms.strategies)
update!(model, s::MetaStrategy, i, item) = foreach(x -> update!(model, x, i, item), s.strategies)
hook(s::MetaStrategy, model, data, i) = foreach(x -> hook(x, model, data, i), s.strategies)
finished(s::MetaStrategy, model, data, i) = any(x -> finished(x, model, data, i), s.strategies)
cleanup!(s::MetaStrategy, model) = foreach(x -> cleanup!(x, model), s.strategies)

"""
    strategy(s::LearningStrategy...)
    strategy(ms::MetaStrategy, s::LearningStrategy...)

Create a MetaStrategy from LearningStrategies or add a LearningStrategy to an existing MetaStrategy.
"""
strategy(s::LearningStrategy...) = MetaStrategy(s)
strategy(ms::MetaStrategy, s::LearningStrategy...) = MetaStrategy(ms.strategies..., s...)

Base.getindex(ms::MetaStrategy, i) = ms.strategies[i]

#-----------------------------------------------------------------------# learn!
"""
    learn!(model, strategy, data) -> model

Learn a `model` from `data` using `strategy`.
New models/strategies/data types should overload at least one of the following:

- [`setup!`](@ref)
- [`update!`](@ref)
- [`hook`](@ref)
- [`finished`](@ref)
- [`cleanup!`](@ref)

# `learn!` Implementation:

    function learn!(model, s::LearningStrategy, data)
        setup!(s, model[, data])
        for (i, item) in enumerate(data)
            update!(model, s[, i], item)
            hook(s, model[, data], i)
            finished(s, model[, data], i) && break
        end
        cleanup!(s, model)
    end
"""
function learn!(model, s::LearningStrategy, data)
    setup!(s, model, data)
    for (i, item) in enumerate(data)
        update!(model, s, i, item)
        hook(s, model, data, i)
        finished(s, model, data, i) && break
    end
    cleanup!(s, model)
    model
end

#-----------------------------------------------------------------------# InfiniteNothing
# learn without input data... good for minimizing functions
struct InfiniteNothing end
Base.iterate(o::InfiniteNothing) = (nothing, 1)
Base.iterate(o::InfiniteNothing, state) = (state, 1)
learn!(model, s::LearningStrategy) = learn!(model, s, InfiniteNothing())

################################################################################## Strategies

#-----------------------------------------------------------------------# Verbose
"""
    Verbose(s::LearningStrategy)
    Verbose(s::LearningStrategy, io::IO)

Allow the LearningStrategy `s` to print output.

- Will automatically print when `finished(s, args...) == true`.
- Other methods should be overloaded to add printout.
    - For example: `update!(model, v::Verbose{MyStrategy}, item) = ...`
"""
struct Verbose{S <: LearningStrategy} <: LearningStrategy
    strategy::S
end
Base.show(io::IO, v::Verbose) = print(io, "Verbose ", v.strategy)

setup!(v::Verbose, model, data) = setup!(v.strategy, model, data)

update!(model, v::Verbose, i, item) = update!(model, v.strategy, i, item)

hook(v::Verbose, model, data, i) = hook(v.strategy, model, data, i)

function finished(v::Verbose, model, data, i)
    done = finished(v.strategy, model, data, i)
    done && @info("$(v.strategy) finished")
    done
end
function finished(v::Verbose{<:MetaStrategy}, model, data, i)
    done = finished(v.strategy, model, data, i)
    if done
        for s in v.strategy.strategies
            finished(Verbose(s), model, data, i)
        end
    end
    done
end
cleanup!(v::Verbose, model) = cleanup!(v.strategy, model)

#-----------------------------------------------------------------------# MaxIter
"""
    MaxIter(n)

Stop learning after `n` iterations.

Wrapping `MaxIter` with [`Verbose`](@ref):

```jldoctest
julia> learn!(nothing, Verbose(MaxIter(5)), 1:100)
INFO: MaxIter: 1/5
INFO: MaxIter: 2/5
INFO: MaxIter: 3/5
INFO: MaxIter: 4/5
INFO: MaxIter: 5/5
INFO: MaxIter(5) finished
```
"""
struct MaxIter <: LearningStrategy
    n::Int
end
MaxIter() = MaxIter(100)

Base.show(io::IO, s::MaxIter) = print(io, "MaxIter($(s.n))")

function update!(model, v::Verbose{MaxIter}, i, item)
  n = v.strategy.n
  l = length(string(n))
  @info("MaxIter: $(lpad(string(i), l, string(0)))/$n")
end

finished(s::MaxIter, model, data, i) = i >= s.n

#-----------------------------------------------------------------------# TimeLimit
"""
    TimeLimit(s)

Stop learning after `s` seconds.
"""
mutable struct TimeLimit <: LearningStrategy
    secs::Float64
    secs_end::Float64
    TimeLimit(secs::Number) = new(secs)
end
Base.show(io::IO, s::TimeLimit) = print(io, "TimeLimit($(s.secs))")
setup!(strat::TimeLimit, model, data) = (strat.secs_end = time() + strat.secs)
finished(strat::TimeLimit, model, data, i) = time() >= strat.secs_end

#-----------------------------------------------------------------------# Converged
"""
    Converged(f; tol = 1e-6, every = 1)

Stop learning when `norm(f(model) - lastf) ≦ tol`.  `f` must return a `Vector{Float64}`.
"""
mutable struct Converged{F <: Function} <: LearningStrategy
    f::F          # f(model)
    tol::Float64  # normdiff tolerance
    every::Int    # only check every ith iteration
    lastval::Vector{Float64}
end
Converged(f::Function; tol::Number = 1e-6, every::Int = 1) = Converged(f, tol, every, [0.0])

Base.show(io::IO, s::Converged) = print(io, "Converged($(s.f), $(s.tol), $(s.every))")

function setup!(s::Converged, model, item)
    val = s.f(model)
    s.lastval = val isa Number ? [val] : fill(0.0, length(val))
end

function finished(s::Converged, model, data, i)
    val = s.f(model)
    if i > 1 && norm(val .- s.lastval) <= s.tol
        true
    else
        copyto!(s.lastval, val)
        false
    end
end
function finished(v::Verbose{<:Converged}, model, data, i)
    done = finished(v.strategy, model, data, i)
    done && @info("Converged after $i iterations: $(v.strategy.f(model))")
    done
end

#-----------------------------------------------------------------------# ConvergedTo
"""
    ConvergedTo(f, goal; tol=1e-6, every=1)

Stop learning when `‖f(model) - goal‖ ≦ tol`.
"""
struct ConvergedTo{V} <: LearningStrategy
    f::Function   # f(model)
    tol::Float64  # normdiff tolerance
    goal::V       # goal value
    every::Int    # only check every ith iteration
end
function ConvergedTo(f::Function, goal; tol::Number = 1e-6, every::Int = 1)
    ConvergedTo(f, tol, goal, every)
end
Base.show(io::IO, s::ConvergedTo) = print(io, "ConvergedTo($(s.f), $(s.tol), $(s.goal), $(s.every))")
function finished(strat::ConvergedTo, model, data, i)
    val = strat.f(model)
    if norm(val - strat.goal) <= strat.tol
        @info("Converged after $i iterations: $val")
        true
    else
        false
    end
end

#-----------------------------------------------------------------------# IterFunction
"""
    IterFunction(f)
    IterFunction(f, b)
    IterFunction(b, f)

Call `f(model, i)` every `b` iterations at `hook` call.
Default value of `b` is 1.
"""
struct IterFunction{F<:Function} <: LearningStrategy
    f::F
    b::Int
end
IterFunction(f::Function) = IterFunction(f, 1)
IterFunction(b::Int, f::Function) = IterFunction(f, b)

Base.show(io::IO, o::IterFunction) = print(io, "IterFunction($(o.f), $(o.b))")

function hook(s::IterFunction, model, i)
    if mod1(i, s.b) == s.b
        s.f(model, i)
    end
    return
end

#-----------------------------------------------------------------------# Tracer
"""
    Tracer{T}(::Type{T}, f, b=1)

Store `f(model, i)` every `b` iterations.

To extract the data `Tracer` collected, `collect(tracer)`.
Note that this operation will copy.

```jldoctest
julia> t = Tracer(Int, (model, i) -> @show(i))
Tracer(1, #5, 0-element Array{Int64,1})

julia> learn!(nothing, strategy(MaxIter(3), t))
i = 1
i = 2
i = 3

julia> collect(t)
3-element Array{Int64,1}:
 1
 2
 3
```
"""
struct Tracer{S} <: LearningStrategy
    every::Int
    f::Function
    storage::Vector{S}
end

Tracer(::Type{S}, f::Function, every::Int = 1) where {S} = Tracer(every, f, S[])

Base.show(io::IO, s::Tracer) = print(io, "Tracer($(s.every), $(s.f), $(summary(s.storage)))")

function hook(strat::Tracer, model, i)
    if mod1(i, strat.every) == strat.every
        push!(strat.storage, strat.f(model, i))
    end
    return
end

Base.collect(t::Tracer) = collect(t.storage)

#-----------------------------------------------------------------------# Breaker
"""
    Breaker(f)

Stop learning when `f(model, i)` returns true.
"""
struct Breaker{F<:Function} <: LearningStrategy
    f::F
end
Base.show(io::IO, s::Breaker) = print(io, "Breaker($(s.f))")
finished(strat::Breaker, model, data, i)::Bool = strat.f(model, i)

end # module


text:
""" Syntax systems for generalized algebraic theories (GATs).

In general, a single theory may have many different syntaxes. The purpose of
this module to enable the simple but flexible construction of syntax systems.
"""
module Syntax
export @syntax, GATExpr, SyntaxDomainError, head, args, type_args, first, last,
  invoke_term, functor, to_json_sexpr, parse_json_sexpr, show_sexpr,
  show_unicode, show_latex

import Base: first, last
import Base.Meta: ParseError, show_sexpr
using Compat
using Match

using ..GAT: Context, Theory, TypeConstructor, TermConstructor, Typeclass
import ..GAT
import ..GAT: invoke_term
using ..Meta

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

""" Base type for expression in the syntax of a GAT.

We define Julia types for each *type constructor* in the theory, e.g., object,
morphism, and 2-morphism in the theory of 2-categories. Of course, Julia's
type system does not support dependent types, so the type parameters are
incorporated in the Julia types. (They are stored as extra data in the
expression instances.)

The concrete types are structurally similar to the core type `Expr` in Julia.
However, the *term constructor* is represented as a type parameter, rather than
as a `head` field. This makes dispatch using Julia's type system more
convenient.
"""
abstract type GATExpr{T} end

head(::GATExpr{T}) where T = T
args(expr::GATExpr) = expr.args
first(expr::GATExpr) = first(args(expr))
last(expr::GATExpr) = last(args(expr))
type_args(expr::GATExpr) = expr.type_args

function Base.:(==)(e1::GATExpr, e2::GATExpr)
  head(e1) == head(e2) && args(e1) == args(e2) && type_args(e1) == type_args(e2)
end
function Base.hash(e::GATExpr, h::UInt)
  hash(args(e), hash(head(e), h))
end

function Base.show(io::IO, expr::GATExpr)
  print(io, head(expr))
  print(io, "(")
  join(io, args(expr), ",")
  print(io, ")")
end
function Base.show(io::IO, expr::GATExpr{:generator})
  value = first(expr)
  if isnothing(value)
    show(io, value) # Value `nothing` cannot be printed
  else
    print(io, value)
  end
end

struct SyntaxDomainError <: Exception
  constructor::Symbol
  args::Vector
end

function Base.showerror(io::IO, exc::SyntaxDomainError)
  print(io, "Domain error in term constructor $(exc.constructor)(")
  join(io, exc.args, ",")
  print(io, ")")
end

# Syntax
########

""" Define a *syntax* system for a generalized algebraic theory (GAT).

A syntax system consists of Julia types (with top type `GATExpr`) for each type
constructor in the signature, plus Julia functions for

1. *Generators*: creating new generator terms, e.g., objects or morphisms
2. *Accessors*: accessing type parameters, e.g., domains and codomains
3. *Term constructors*: applying term constructors, e.g., composition and
   monoidal products

Julia code for all this is generated by the macro. Any of the methods can be
overriden with custom simplification logic.
"""
macro syntax(syntax_head, mod_name, body=nothing)
  if isnothing(body); body = Expr(:block) end
  @assert body.head == :block
  syntax_name, base_types = @match syntax_head begin
    Expr(:call, [name::Symbol, args...]) => (name, args)
    name::Symbol => (name, [])
    _ => throw(ParseError("Ill-formed syntax signature $syntax_head"))
  end
  functions = map(parse_function, strip_lines(body).args)

  expr = Expr(:call, :syntax_code, Expr(:quote, syntax_name),
              esc(Expr(:ref, :Type, base_types...)),
              esc(mod_name), esc(nameof(__module__)), functions)
  Expr(:block,
    Expr(:call, esc(:eval), expr),
    :(Core.@__doc__ $(esc(syntax_name))))
end
function syntax_code(name::Symbol, base_types::Vector{Type},
                     theory_module::Module, outer_module::Module,
                     functions::Vector)
  class = theory_module.class()
  theory = class.theory
  theory_ref = GlobalRef(parentmodule(theory_module),
                            nameof(theory_module))

  # Generate module with syntax types and type/term generators.
  mod = Expr(:module, true, name,
    Expr(:block, [
      # Prevents error about export not being at toplevel.
      # https://github.com/JuliaLang/julia/issues/28991
      LineNumberNode(0);
      Expr(:export, [cons.name for cons in theory.types]...);
      Expr(:using, Expr(:., :., :., nameof(outer_module)));
      :(theory() = $theory_ref);
      gen_types(theory, base_types);
      gen_type_accessors(theory);
      gen_term_generators(theory, outer_module);
      gen_term_constructors(theory, outer_module);
    ]...))

  # Generate toplevel functions.
  toplevel = []
  bindings = Dict{Symbol,Any}(
    c.name => Expr(:(.), name, QuoteNode(c.name)) for c in theory.types)
  syntax_fns = Dict(parse_function_sig(f) => f for f in functions)
  for f in interface(class)
    sig = parse_function_sig(f)
    bindings[:new] = Expr(:(.), name, QuoteNode(sig.name))
    if haskey(syntax_fns, sig)
      # Case 1: The method is overriden in the syntax body.
      expr = generate_function(replace_symbols(bindings, syntax_fns[sig]))
    elseif !isnothing(f.impl)
      # Case 2: The method has a default implementation in the theory.
      expr = generate_function(replace_symbols(bindings, f))
    else
      # Case 3: Call the default syntax method.
      params = [ gensym("x$i") for i in eachindex(sig.types) ]
      call_expr = Expr(:call, sig.name,
        [ Expr(:(::), pair...) for pair in zip(params, sig.types) ]...)
      body = Expr(:call, :new, params...)
      f_impl = JuliaFunction(call_expr, f.return_type, body)
      expr = generate_function(replace_symbols(bindings, f_impl))
    end
    push!(toplevel, expr)
  end
  Expr(:toplevel, mod, toplevel...)
end

""" Complete set of Julia functions for a syntax system.
"""
function interface(class::Typeclass)::Vector{JuliaFunction}
  theory = class.theory
  [ GAT.interface(class);
    [ GAT.constructor(constructor_for_generator(cons), theory)
      for cons in theory.types ]; ]
end

""" Generate syntax type definitions.
"""
function gen_type(cons::TypeConstructor, base_type::Type=Any)::Expr
  base_expr = GlobalRef(Syntax, :GATExpr)
  base_name = if base_type == Any
    base_expr
  else
    GlobalRef(parentmodule(base_type), nameof(base_type))
  end
  expr = :(struct $(cons.name){T} <: $base_name{T}
    args::Vector
    type_args::Vector{$base_expr}
  end)
  generate_docstring(strip_lines(expr, recurse=true), cons.doc)
end
function gen_types(theory::Theory, base_types::Vector{Type})::Vector{Expr}
  if isempty(base_types)
    map(gen_type, theory.types)
  else
    map(gen_type, theory.types, base_types)
  end
end

""" Generate accessor methods for type parameters.
"""
function gen_type_accessors(cons::TypeConstructor)::Vector{Expr}
  fns = []
  sym = gensym(:x)
  for (i, param) in enumerate(cons.params)
    call_expr = Expr(:call, param, Expr(:(::), sym, cons.name))
    return_type = GAT.strip_type(cons.context[param])
    body = Expr(:ref, Expr(:(.), sym, QuoteNode(:type_args)), i)
    push!(fns, generate_function(JuliaFunction(call_expr, return_type, body)))
  end
  fns
end
function gen_type_accessors(theory::Theory)::Vector{Expr}
  vcat(map(gen_type_accessors, theory.types)...)
end

""" Generate methods for syntax term constructors.
"""
function gen_term_constructor(cons::TermConstructor, theory::Theory,
                              mod::Module; dispatch_type::Symbol=Symbol())::Expr
  head = GAT.constructor(cons, theory)
  call_expr, return_type = head.call_expr, head.return_type
  if dispatch_type == Symbol()
    dispatch_type = cons.name
  end
  body = Expr(:block)

  # Create expression to check constructor domain.
  eqs = GAT.equations(cons, theory)
  if !isempty(eqs)
    clauses = [ Expr(:call,:(==),lhs,rhs) for (lhs,rhs) in eqs ]
    conj = foldr((x,y) -> Expr(:(&&),x,y), clauses)
    insert!(call_expr.args, 2,
      Expr(:parameters, Expr(:kw, :strict, false)))
    push!(body.args,
      Expr(:if,
        Expr(:(&&), :strict, Expr(:call, :(!), conj)),
        Expr(:call, :throw,
          Expr(:call, GlobalRef(Syntax, :SyntaxDomainError),
            Expr(:quote, cons.name),
            Expr(:vect, cons.params...)))))
  end

  # Create call to expression constructor.
  type_params = gen_term_constructor_params(cons, theory, mod)
  push!(body.args,
    Expr(:call,
      Expr(:curly, return_type, Expr(:quote, dispatch_type)),
      Expr(:vect, cons.params...),
      Expr(:vect, type_params...)))

  generate_function(JuliaFunction(call_expr, return_type, body))
end
function gen_term_constructors(theory::Theory, mod::Module)::Vector{Expr}
  [ gen_term_constructor(cons, theory, mod) for cons in theory.terms ]
end

""" Generate expressions for type parameters of term constructor.

Besides expanding the implicit variables, we must handle two annoying issues:

1. Add types for method dispatch where necessary (see `GAT.add_type_dispatch`)
   FIXME: We are currently only handling the nullary case (e.g., `munit()`).
   To handle the general case, we need to do basic type inference.

2. Rebind the term constructors to ensure that user overrides are preferred over
   the default term constructors.
"""
function gen_term_constructor_params(cons, theory, mod)::Vector
  expr = GAT.expand_term_type(cons, theory)
  raw_params = @match expr begin
    Expr(:call, [name::Symbol, args...]) => args
    _::Symbol => []
  end

  bindings = Dict(c.name => GlobalRef(mod, c.name) for c in theory.terms)
  params = []
  for expr in raw_params
    expr = replace_nullary_constructors(expr, theory)
    expr = replace_symbols(bindings, expr)
    push!(params, expr)
  end
  params
end
function replace_nullary_constructors(expr, theory)
  @match expr begin
    Expr(:call, [name::Symbol]) => begin
      terms = theory.terms[findall(cons -> cons.name == name, theory.terms)]
      @assert length(terms) == 1
      Expr(:call, name, terms[1].typ)
    end
    Expr(:call, [name::Symbol, args...]) =>
      Expr(:call, name, [replace_nullary_constructors(a,theory) for a in args]...)
    _ => expr
  end
end

""" Generate methods for term generators.

Generators are extra term constructors created automatically for the syntax.
"""
function gen_term_generator(cons::TypeConstructor, theory::Theory, mod::Module)::Expr
  gen_term_constructor(constructor_for_generator(cons), theory, mod;
                       dispatch_type = :generator)
end
function gen_term_generators(theory::Theory, mod::Module)::Vector{Expr}
  [ gen_term_generator(cons, theory, mod) for cons in theory.types ]
end
function constructor_for_generator(cons::TypeConstructor)::TermConstructor
  value_param = :__value__
  params = [ value_param; cons.params ]
  typ = Expr(:call, cons.name, cons.params...)
  context = merge(Context(value_param => :Any), cons.context)
  TermConstructor(cons.name, params, typ, context)
end

# Reflection
############

""" Invoke a term constructor by name in a syntax system.

This method provides reflection for syntax systems. In everyday use the generic
method for the constructor should be called directly, not through this function.
"""
function invoke_term(syntax_module::Module, constructor_name::Symbol, args...)
  theory_module = syntax_module.theory()
  theory = theory_module.class().theory
  syntax_types = Tuple(getfield(syntax_module, cons.name) for cons in theory.types)
  invoke_term(theory_module, syntax_types, constructor_name, args...)
end

""" Name of constructor that created expression.
"""
function constructor_name(expr::GATExpr)::Symbol
  if head(expr) == :generator
    nameof(typeof(expr))
  else
    head(expr)
  end
end

""" Create generator of the same type as the given expression.
"""
function generator_like(expr::GATExpr, value)::GATExpr
  invoke_term(
    syntax_module(expr), nameof(typeof(expr)), value, type_args(expr)...)
end

""" Get syntax module of given expression.
"""
syntax_module(expr::GATExpr) = parentmodule(typeof(expr))

# Functors
##########

""" Functor from GAT expression to GAT instance.

Strictly speaking, we should call these "structure-preserving functors" or,
better, "model homomorphisms of GATs". But this is a category theory library,
so we'll go with the simpler "functor".

A functor is completely determined by its action on the generators. There are
several ways to specify this mapping:

  1. Simply specify a Julia instance type for each doctrine type, using the
     required `types` tuple. For this to work, the generator constructors
     must be defined for the instance types.

  2. Explicitly map each generator term to an instance value, using the
     `generators` dictionary.

  3. For each doctrine type (e.g., object and morphism), specify a function
     mapping generator terms of that type to an instance value, using the
     `terms` dictionary.

The `terms` dictionary can also be used for special handling of non-generator
expressions. One use case for this capability is defining forgetful functors,
which map non-generators to generators.
"""
function functor(types::Tuple, expr::GATExpr;
                 generators::AbstractDict=Dict(), terms::AbstractDict=Dict())
  # Special case: look up a specific generator.
  if head(expr) == :generator && haskey(generators, expr)
    return generators[expr]
  end

  # Special case: look up by type of term (usually a generator).
  name = constructor_name(expr)
  if haskey(terms, name)
    return terms[name](expr)
  end

  # Otherwise, we need to call a term constructor (possibly for a generator).
  # Recursively evalute the arguments.
  term_args = []
  for arg in args(expr)
    if isa(arg, GATExpr)
      arg = functor(types, arg; generators=generators, terms=terms)
    end
    push!(term_args, arg)
  end

  # Invoke the constructor in the codomain category!
  theory_module = syntax_module(expr).theory()
  invoke_term(theory_module, types, name, term_args...)
end

# Serialization
###############

""" Serialize expression as JSON-able S-expression.

The format is an S-expression encoded as JSON, e.g., "compose(f,g)" is
represented as ["compose", f, g].
"""
function to_json_sexpr(expr::GATExpr; by_reference::Function = x->false)
  if head(expr) == :generator && by_reference(first(expr))
    to_json_sexpr(first(expr))
  else
    [ string(constructor_name(expr));
      [ to_json_sexpr(arg; by_reference=by_reference) for arg in args(expr) ] ]
  end
end
to_json_sexpr(x::Union{Bool,Real,String,Nothing}; kw...) = x
to_json_sexpr(x; kw...) = string(x)

""" Deserialize expression from JSON-able S-expression.

If `symbols` is true (the default), strings are converted to symbols.
"""
function parse_json_sexpr(syntax_module::Module, sexpr;
    parse_head::Function = identity,
    parse_reference::Function = x->error("Loading terms by name is disabled"),
    parse_value::Function = identity,
    symbols::Bool = true,
  )
  theory_module = syntax_module.theory()
  theory = theory_module.class().theory
  type_lens = Dict(cons.name => length(cons.params) for cons in theory.types)

  function parse_impl(sexpr::Vector, ::Val{:expr})
    name = Symbol(parse_head(symbols ? Symbol(sexpr[1]) : sexpr[1]))
    nargs = length(sexpr) - 1
    args = map(enumerate(sexpr[2:end])) do (i, arg)
      arg_kind = ((i == 1 && get(type_lens, name, nothing) == nargs-1) ||
                  arg isa Union{Bool,Number,Nothing}) ? :value : :expr
      parse_impl(arg, Val(arg_kind))
    end
    invoke_term(syntax_module, name, args...)
  end
  parse_impl(x, ::Val{:value}) = parse_value(x)
  parse_impl(x::String, ::Val{:expr}) = parse_reference(symbols ? Symbol(x) : x)
  parse_impl(x::String, ::Val{:value}) = parse_value(symbols ? Symbol(x) : x)

  parse_impl(sexpr, Val(:expr))
end

# Pretty-print
##############

""" Show the syntax expression as an S-expression.

Cf. the standard library function `Meta.show_sexpr`.
"""
show_sexpr(expr::GATExpr) = show_sexpr(stdout, expr)

function show_sexpr(io::IO, expr::GATExpr)
  if head(expr) == :generator
    print(io, repr(first(expr)))
  else
    print(io, "(")
    join(io, [string(head(expr));
              [sprint(show_sexpr, arg) for arg in args(expr)]], " ")
    print(io, ")")
  end
end

""" Show the expression in infix notation using Unicode symbols.
"""
show_unicode(expr::GATExpr) = show_unicode(stdout, expr)
show_unicode(io::IO, x::Any; kw...) = show(io, x)

function show_unicode(io::IO, expr::GATExpr; kw...)
  # By default, show in prefix notation.
  print(io, head(expr))
  print(io, "{")
  join(io, [sprint(show_unicode, arg) for arg in args(expr)], ",")
  print(io, "}")
end

function show_unicode(io::IO, expr::GATExpr{:generator}; kw...)
  print(io, first(expr))
end

function show_unicode_infix(io::IO, expr::GATExpr, op::String;
                            paren::Bool=false)
  show_unicode_paren(io, expr) = show_unicode(io, expr; paren=true)
  if (paren) print(io, "(") end
  join(io, [sprint(show_unicode_paren, arg) for arg in args(expr)], op)
  if (paren) print(io, ")") end
end

""" Show the expression in infix notation using LaTeX math.

Does *not* include `\$` or `\\[begin|end]{equation}` delimiters.
"""
show_latex(expr::GATExpr) = show_latex(stdout, expr)
show_latex(io::IO, sym::Symbol; kw...) = print(io, sym)
show_latex(io::IO, x::Any; kw...) = show(io, x)

function show_latex(io::IO, expr::GATExpr; kw...)
  # By default, show in prefix notation.
  print(io, "\\mathop{\\mathrm{$(head(expr))}}")
  print(io, "\\left[")
  join(io, [sprint(show_latex, arg) for arg in args(expr)], ",")
  print(io, "\\right]")
end

function show_latex(io::IO, expr::GATExpr{:generator}; kw...)
  # Try to be smart about using text or math mode.
  content = string(first(expr))
  if all(isletter, content) && length(content) > 1
    print(io, "\\mathrm{$content}")
  else
    print(io, content)
  end
end

function show_latex_infix(io::IO, expr::GATExpr, op::String;
                          paren::Bool=false, kw...)
  show_latex_paren(io, expr) = show_latex(io, expr, paren=true)
  sep = op == " " ? op : " $op "
  if (paren) print(io, "\\left(") end
  join(io, [sprint(show_latex_paren, arg) for arg in args(expr)], sep)
  if (paren) print(io, "\\right)") end
end

function show_latex_postfix(io::IO, expr::GATExpr, op::String; kw...)
  @assert length(args(expr)) == 1
  print(io, "{")
  show_latex(io, first(expr), paren=true)
  print(io, "}")
  print(io, op)
end

function show_latex_script(io::IO, expr::GATExpr, head::String;
                           super::Bool=false, kw...)
  print(io, head, super ? "^" : "_", "{")
  join(io, [sprint(show_latex, arg) for arg in args(expr)], ",")
  print(io, "}")
end

end


text:
module Config

using AbInitioSoftwareBase.Commands: CommandConfig
using Configurations: from_dict, option_m, @option
using Formatting: sprintf1
using Unitful: PressureUnits, VolumeUnits, Units, uparse, @u_str
using Unitful
using UnitfulAtomic

import Configurations: convert_to_option

# export Pressures, Volumes
export @uvec

abstract type UnitfulVector <: AbstractVector{Float64} end

convert_to_option(::Type{<:UnitfulVector}, ::Type{Vector}, str::AbstractString) =
    eval(Meta.parse(str))
convert_to_option(::Type{<:UnitfulVector}, ::Type{VolumeUnits}, str::AbstractString) =
    myuparse(str)

Base.size(A::UnitfulVector) = size(A.values)

Base.getindex(A::UnitfulVector, I) = getindex(A.values, I) * A.unit

Base.setindex!(A::UnitfulVector, v, I) = setindex!(A.values, v, I)

macro uvec(typename, unittype, alias, defaultunit)
    ex = quote
        struct $typename <: UnitfulVector
            values::Vector{Float64}
            unit::$unittype
        end
    end
    ex = quote
        @option $(esc(ex))
        $typename(values, unit::AbstractString = @u_str($defaultunit)) = $typename(collect(values), myuparse(unit))
    end
end

# @option "pressures" struct Pressures <: UnitfulVector
#     values::Vector{Float64}
#     unit::PressureUnits
# end

# @option "volumes" struct Volumes <: UnitfulVector
#     values::Vector{Float64}
#     unit::VolumeUnits
# end
# Volumes(values, unit::AbstractString = u"angstrom^3") = Volumes(collect(values), myuparse(unit))

myuparse(str::AbstractString) =
    uparse(filter(!isspace, str); unit_context = [Unitful, UnitfulAtomic])
myuparse(num::Number) = num  # FIXME: this might be error-prone!

end


text:
#=
This file is auto-generated. Do not edit.
=#
"""
    mutable struct Area <: AggregationTopology
        name::String
        peak_active_power::Float64
        peak_reactive_power::Float64
        load_response::Float64
        time_series_container::InfrastructureSystems.TimeSeriesContainer
        internal::InfrastructureSystemsInternal
    end

A collection of buses for control purposes.

# Arguments
- `name::String`
- `peak_active_power::Float64`
- `peak_reactive_power::Float64`
- `load_response::Float64`
- `time_series_container::InfrastructureSystems.TimeSeriesContainer`: internal time_series storage
- `internal::InfrastructureSystemsInternal`: power system internal reference, do not modify
"""
mutable struct Area <: AggregationTopology
    name::String
    peak_active_power::Float64
    peak_reactive_power::Float64
    load_response::Float64
    "internal time_series storage"
    time_series_container::InfrastructureSystems.TimeSeriesContainer
    "power system internal reference, do not modify"
    internal::InfrastructureSystemsInternal
end

function Area(name, peak_active_power=0.0, peak_reactive_power=0.0, load_response=0.0, time_series_container=InfrastructureSystems.TimeSeriesContainer(), )
    Area(name, peak_active_power, peak_reactive_power, load_response, time_series_container, InfrastructureSystemsInternal(), )
end

function Area(; name, peak_active_power=0.0, peak_reactive_power=0.0, load_response=0.0, time_series_container=InfrastructureSystems.TimeSeriesContainer(), internal=InfrastructureSystemsInternal(), )
    Area(name, peak_active_power, peak_reactive_power, load_response, time_series_container, internal, )
end

# Constructor for demo purposes; non-functional.
function Area(::Nothing)
    Area(;
        name="init",
        peak_active_power=0.0,
        peak_reactive_power=0.0,
        load_response=0.0,
        time_series_container=InfrastructureSystems.TimeSeriesContainer(),
    )
end

"""Get [`Area`](@ref) `name`."""
get_name(value::Area) = value.name
"""Get [`Area`](@ref) `peak_active_power`."""
get_peak_active_power(value::Area) = get_value(value, value.peak_active_power)
"""Get [`Area`](@ref) `peak_reactive_power`."""
get_peak_reactive_power(value::Area) = get_value(value, value.peak_reactive_power)
"""Get [`Area`](@ref) `load_response`."""
get_load_response(value::Area) = value.load_response
"""Get [`Area`](@ref) `time_series_container`."""
get_time_series_container(value::Area) = value.time_series_container
"""Get [`Area`](@ref) `internal`."""
get_internal(value::Area) = value.internal

"""Set [`Area`](@ref) `name`."""
set_name!(value::Area, val) = value.name = val
"""Set [`Area`](@ref) `peak_active_power`."""
set_peak_active_power!(value::Area, val) = value.peak_active_power = set_value(value, val)
"""Set [`Area`](@ref) `peak_reactive_power`."""
set_peak_reactive_power!(value::Area, val) = value.peak_reactive_power = set_value(value, val)
"""Set [`Area`](@ref) `load_response`."""
set_load_response!(value::Area, val) = value.load_response = val
"""Set [`Area`](@ref) `time_series_container`."""
set_time_series_container!(value::Area, val) = value.time_series_container = val


text:
"""
    LinearMixedModel

Linear mixed-effects model representation

## Fields

* `formula`: the formula for the model
* `allterms`: a vector of random-effects terms, the fixed-effects terms and the response
* `reterms`: a `Vector{AbstractReMat{T}}` of random-effects terms.
* `feterms`: a `Vector{FeMat{T}}` of the fixed-effects model matrix and the response
* `sqrtwts`: vector of square roots of the case weights.  Can be empty.
* `parmap` : Vector{NTuple{3,Int}} of (block, row, column) mapping of θ to λ
* `dims` : NamedTuple{(:n, :p, :nretrms),NTuple{3,Int}} of dimensions.  `p` is the rank of `X`, which may be smaller than `size(X, 2)`.
* `A`: an `nt × nt` symmetric `BlockMatrix` of matrices representing `hcat(Z,X,y)'hcat(Z,X,y)`
* `L`: a `nt × nt` `BlockMatrix` - the lower Cholesky factor of `Λ'AΛ+I`
* `optsum`: an [`OptSummary`](@ref) object

## Properties

* `θ` or `theta`: the covariance parameter vector used to form λ
* `β` or `beta`: the fixed-effects coefficient vector
* `λ` or `lambda`: a vector of lower triangular matrices repeated on the diagonal blocks of `Λ`
* `σ` or `sigma`: current value of the standard deviation of the per-observation noise
* `b`: random effects on the original scale, as a vector of matrices
* `u`: random effects on the orthogonal scale, as a vector of matrices
* `lowerbd`: lower bounds on the elements of θ
* `X`: the fixed-effects model matrix
* `y`: the response vector
"""
struct LinearMixedModel{T<:AbstractFloat} <: MixedModel{T}
    formula::FormulaTerm
    allterms::Vector{Union{AbstractReMat{T}, FeMat{T}}}
    reterms::Vector{AbstractReMat{T}}
    feterms::Vector{FeMat{T}}
    sqrtwts::Vector{T}
    parmap::Vector{NTuple{3,Int}}
    dims::NamedTuple{(:n, :p, :nretrms),NTuple{3,Int}}
    A::BlockMatrix{T}            # cross-product blocks
    L::BlockMatrix{T}
    optsum::OptSummary{T}
end
LinearMixedModel(f::FormulaTerm, tbl; contrasts = Dict{Symbol,Any}(), wts = []) =
    LinearMixedModel(
        f::FormulaTerm,
        Tables.columntable(tbl),
        contrasts = contrasts,
        wts = wts,
    )
function LinearMixedModel(
    f::FormulaTerm,
    tbl::Tables.ColumnTable;
    contrasts = Dict{Symbol,Any}(),
    wts = [],
)
    # TODO: perform missing_omit() after apply_schema() when improved
    # missing support is in a StatsModels release
    tbl, _ = StatsModels.missing_omit(tbl, f)
    sch = try
        schema(f, tbl, contrasts)
    catch e
        if isa(e, OutOfMemoryError)
            @warn "Random effects grouping variables with many levels can cause out-of-memory errors.  Try manually specifying `Grouping()` contrasts for those variables."
        end
        rethrow(e)
    end
    form = apply_schema(f, sch, LinearMixedModel)
    # tbl, _ = StatsModels.missing_omit(tbl, form)

    y, Xs = modelcols(form, tbl)

    return LinearMixedModel(y, Xs, form, wts)
end

"""
    LinearMixedModel(y, Xs, form)

Private constructor for a LinearMixedModel.

To construct a model, you only need the response (`y`), already assembled
model matrices (`Xs`), schematized formula (`form`) and weights (`wts`).
Everything else in the structure can be derived from these quantities.

!!! note
    This method is internal and experimental and so may change or disappear in
    a future release without being considered a breaking change.
"""
function LinearMixedModel(y::AbstractArray,
                           Xs::Tuple, # can't be more specific here without stressing the compiler
                           form::FormulaTerm, wts = [])
    y = reshape(float(y), (:, 1)) # y as a floating-point matrix
    T = promote_type(Float64, eltype(y))  # ensure that eltype of model matrices is at least Float64
    y = convert(Matrix{T}, y)

    reterms = AbstractReMat{T}[]
    feterms = FeMat{T}[]
    for (i, x) in enumerate(Xs)
        if isa(x, AbstractReMat{T})
            push!(reterms, x)
        elseif isa(x, ReMat) # this can occur in weird situation where x is a ReMat{U}
            # avoid keeping a second copy if unweighted
            z = convert(Matrix{T}, x.z)
            wtz = x.z === x.wtz ? z : convert(Matrix{T}, x.wtz)
            S = size(z, 1)
            x = ReMat{T,S}(x.trm, x.refs, x.levels, x.cnames, z, wtz,
                           convert(LowerTriangular{Float64, Matrix{Float64}}, x.λ),
                           x.inds,
                           convert(SparseMatrixCSC{T,Int32}, x.adjA),
                           convert(Matrix{T}, x.scratch)
                           )
            push!(reterms, x)
        else
            cnames = coefnames(form.rhs[i])
            push!(feterms, FeMat(x, isa(cnames, String) ? [cnames] : collect(cnames)))
        end
    end
    push!(feterms, FeMat(y, [""]))

    return LinearMixedModel(feterms, reterms, form, wts)
end

"""
    LinearMixedModel(feterms, reterms, form, wts=[])

Private constructor for a `LinearMixedModel` given already assembled fixed and random effects.

To construct a model, you only need a vector of `FeMat`s (the fixed-effects
model matrix and response), a vector of `AbstractReMat` (the random-effects
model matrices), the formula and the weights. Everything else in the structure
can be derived from these quantities.

!!! note
    This method is internal and experimental and so may change or disappear in
    a future release without being considered a breaking change.
"""
function LinearMixedModel(feterms::Vector{FeMat{T}}, reterms::Vector{AbstractReMat{T}},
                          form::FormulaTerm, wts=[]) where T

    # detect and combine RE terms with the same grouping var
    if length(reterms) > 1
        reterms = amalgamate(reterms)
    end

    sort!(reterms, by = nranef, rev = true)
    allterms = convert(Vector{Union{AbstractReMat{T},FeMat{T}}}, vcat(reterms, feterms))
    sqrtwts = sqrt.(convert(Vector{T}, wts))
    reweight!.(allterms, Ref(sqrtwts))
    A, L = createAL(allterms)
    lbd = foldl(vcat, lowerbd(c) for c in reterms)
    θ = foldl(vcat, getθ(c) for c in reterms)
    X = first(feterms)
    optsum = OptSummary(θ, lbd, :LN_BOBYQA, ftol_rel = T(1.0e-12), ftol_abs = T(1.0e-8))
    fill!(optsum.xtol_abs, 1.0e-10)
    LinearMixedModel(
        form,
        allterms,
        reterms,
        feterms,
        sqrtwts,
        mkparmap(reterms),
        (n = size(X, 1), p = X.rank, nretrms = length(reterms)),
        A,
        L,
        optsum,
        )
end

fit(
    ::Type{LinearMixedModel},
    f::FormulaTerm,
    tbl;
    wts = [],
    contrasts = Dict{Symbol,Any}(),
    verbose::Bool = false,
    REML::Bool = false,
) = fit(
    LinearMixedModel,
    f,
    Tables.columntable(tbl),
    wts = wts,
    contrasts = contrasts,
    verbose = verbose,
    REML = REML,
)

fit(
    ::Type{LinearMixedModel},
    f::FormulaTerm,
    tbl::Tables.ColumnTable;
    wts = wts,
    contrasts = contrasts,
    verbose = verbose,
    REML = REML,
) = fit!(
    LinearMixedModel(f, tbl, contrasts = contrasts, wts = wts),
    verbose = verbose,
    REML = REML,
)

fit(
    ::Type{MixedModel},
    f::FormulaTerm,
    tbl;
    wts = [],
    contrasts = Dict{Symbol,Any}(),
    verbose::Bool = false,
    REML::Bool = false,
) = fit(
    LinearMixedModel,
    f,
    tbl,
    wts = wts,
    contrasts = contrasts,
    verbose = verbose,
    REML = REML,
)

fit(
    ::Type{MixedModel},
    f::FormulaTerm,
    tbl,
    d::Normal,
    l::IdentityLink;
    wts = [],
    contrasts = Dict{Symbol,Any}(),
    verbose::Bool = false,
    REML::Bool = false,
    offset = [],
    fast::Bool = false,
    nAGQ::Integer = 1,
) = fit(
    LinearMixedModel,
    f,
    tbl,
    wts = wts,
    contrasts = contrasts,
    verbose = verbose,
    REML = REML,
)

function StatsBase.coef(m::LinearMixedModel{T}) where {T}
    piv = first(m.feterms).piv
    invpermute!(fixef!(similar(piv, T), m), piv)
end

βs(m::LinearMixedModel) = NamedTuple{(Symbol.(coefnames(m))...,)}(coef(m))

function StatsBase.coefnames(m::LinearMixedModel)
    Xtrm = first(m.feterms)
    invpermute!(copy(Xtrm.cnames), Xtrm.piv)
end

function StatsBase.coeftable(m::LinearMixedModel)
    co = coef(m)
    se = stderror!(similar(co), m)
    z = co ./ se
    pvalue = ccdf.(Chisq(1), abs2.(z))
    names = coefnames(m)

    CoefTable(
        hcat(co, se, z, pvalue),
        ["Coef.", "Std. Error", "z", "Pr(>|z|)"],
        names,
        4, # pvalcol
        3, # teststatcol
    )
end

"""
    condVar(m::LinearMixedModel)

Return the conditional variances matrices of the random effects.

The random effects are returned by `ranef` as a vector of length `k`,
where `k` is the number of random effects terms.  The `i`th element
is a matrix of size `vᵢ × ℓᵢ`  where `vᵢ` is the size of the
vector-valued random effects for each of the `ℓᵢ` levels of the grouping
factor.  Technically those values are the modes of the conditional
distribution of the random effects given the observed data.

This function returns an array of `k` three dimensional arrays,
where the `i`th array is of size `vᵢ × vᵢ × ℓᵢ`.  These are the
diagonal blocks from the conditional variance-covariance matrix,

    s² Λ(Λ'Z'ZΛ + I)⁻¹Λ'
"""
function condVar(m::LinearMixedModel{T}) where {T}
    retrms = m.reterms
    t1 = first(retrms)
    L11 = getblock(m.L, 1, 1)
    if !isone(length(retrms)) || !isa(L11, Diagonal{T,Vector{T}})
        throw(ArgumentError("code for multiple or vector-valued r.e. not yet written"))
    end
    ll = first(t1.λ)
    Ld = L11.diag
    Array{T,3}[reshape(abs2.(ll ./ Ld) .* varest(m), (1, 1, length(Ld)))]
end

function createAL(allterms::Vector{Union{AbstractReMat{T},FeMat{T}}}) where {T}
    k = length(allterms)
    sz = [isa(t, ReMat) ? size(t, 2) : rank(t) for t in allterms]
    A = BlockArray(undef_blocks, AbstractMatrix{T}, sz, sz)
    L = BlockArray(undef_blocks, AbstractMatrix{T}, sz, sz)
    for j = 1:k
        for i = j:k
            Lij = L[Block(i, j)] = densify(allterms[i]' * allterms[j])
            A[Block(i, j)] = deepcopy(isa(Lij, BlockedSparse) ? Lij.cscmat : Lij)
        end
    end
    nretrm = sum(Base.Fix2(isa, ReMat), allterms)
    for i = 2:nretrm      # check for fill-in due to non-nested grouping factors
        ci = allterms[i]
        for j = 1:(i-1)
            cj = allterms[j]
            if !isnested(cj, ci)
                for l = i:k
                    L[Block(l, i)] = Matrix(getblock(L, l, i))
                end
                break
            end
        end
    end
    A, L
end

StatsBase.deviance(m::LinearMixedModel) = objective(m)

GLM.dispersion(m::LinearMixedModel, sqr::Bool = false) = sqr ? varest(m) : sdest(m)

GLM.dispersion_parameter(m::LinearMixedModel) = true

StatsBase.dof(m::LinearMixedModel) = m.dims.p + nθ(m) + 1

function StatsBase.dof_residual(m::LinearMixedModel)::Int
    # nobs - rank(FE) - 1 (dispersion)
    # this differs from lme4 by not including nθ
    # a better estimate would be a number somewhere between the number of
    # variance components and the number of conditional modes
    # nobs, rank FE, num conditional modes, num grouping vars
    dd = m.dims
    dd.n - dd.p - 1
end

"""
    feL(m::LinearMixedModel)

Return the lower Cholesky factor for the fixed-effects parameters, as an `LowerTriangular`
`p × p` matrix.
"""
function feL(m::LinearMixedModel)
    k = m.dims.nretrms + 1
    LowerTriangular(m.L.blocks[k, k])
end

"""
    fit!(m::LinearMixedModel[; verbose::Bool=false, REML::Bool=false])

Optimize the objective of a `LinearMixedModel`.  When `verbose` is `true` the values of the
objective and the parameters are printed on stdout at each function evaluation.
"""
function fit!(m::LinearMixedModel{T}; verbose::Bool = false, REML::Bool = false) where {T}
    optsum = m.optsum
    # this doesn't matter for LMM, but it does for GLMM, so let's be consistent
    if optsum.feval > 0
        throw(ArgumentError("This model has already been fitted. Use refit!() instead."))
    end
    opt = Opt(optsum)
    optsum.REML = REML
    function obj(x, g)
        isempty(g) || throw(ArgumentError("g should be empty for this objective"))
        val = objective(updateL!(setθ!(m, x)))
        verbose && println(round(val, digits = 5), " ", x)
        val
    end
    NLopt.min_objective!(opt, obj)
    optsum.finitial = obj(optsum.initial, T[])
    fmin, xmin, ret = NLopt.optimize!(opt, copyto!(optsum.final, optsum.initial))
    ## check if small non-negative parameter values can be set to zero
    xmin_ = copy(xmin)
    lb = optsum.lowerbd
    for i in eachindex(xmin_)
        if iszero(lb[i]) && zero(T) < xmin_[i] < T(0.001)
            xmin_[i] = zero(T)
        end
    end
    if xmin_ ≠ xmin
        if (zeroobj = obj(xmin_, T[])) ≤ (fmin + 1.e-5)
            fmin = zeroobj
            copyto!(xmin, xmin_)
        end
    end
    ## ensure that the parameter values saved in m are xmin
    updateL!(setθ!(m, xmin))

    optsum.feval = opt.numevals
    optsum.final = xmin
    optsum.fmin = fmin
    optsum.returnvalue = ret
    ret == :ROUNDOFF_LIMITED && @warn("NLopt was roundoff limited")
    if ret ∈ [:FAILURE, :INVALID_ARGS, :OUT_OF_MEMORY, :FORCED_STOP, :MAXEVAL_REACHED]
        @warn("NLopt optimization failure: $ret")
    end
    m
end

function fitted!(v::AbstractArray{T}, m::LinearMixedModel{T}) where {T}
    ## FIXME: Create and use `effects(m) -> β, b` w/o calculating β twice
    Xtrm = first(m.feterms)
    vv = mul!(vec(v), Xtrm, fixef!(similar(Xtrm.piv, T), m))
    for (rt, bb) in zip(m.reterms, ranef(m))
        unscaledre!(vv, rt, bb)
    end
    v
end

StatsBase.fitted(m::LinearMixedModel{T}) where {T} = fitted!(Vector{T}(undef, nobs(m)), m)

"""
    fixef!(v::Vector{T}, m::MixedModel{T})

Overwrite `v` with the pivoted fixed-effects coefficients of model `m`

For full-rank models the length of `v` must be the rank of `X`.  For rank-deficient models
the length of `v` can be the rank of `X` or the number of columns of `X`.  In the latter
case the calculated coefficients are padded with -0.0 out to the number of columns.
"""
function fixef!(v::AbstractVector{T}, m::LinearMixedModel{T}) where {T}
    Xtrm = first(m.feterms)
    if isfullrank(Xtrm)
        ldiv!(feL(m)', copyto!(v, m.L.blocks[end, end-1]))
    else
        ldiv!(
            feL(m)',
            view(copyto!(fill!(v, -zero(T)), m.L.blocks[end, end-1]), 1:(Xtrm.rank)),
        )
    end
    v
end

"""
    fixef(m::MixedModel)

Return the fixed-effects parameter vector estimate of `m`.

In the rank-deficient case the truncated parameter vector, of length `rank(m)` is returned.
This is unlike `coef` which always returns a vector whose length matches the number of
columns in `X`.
"""
fixef(m::LinearMixedModel{T}) where {T} = fixef!(Vector{T}(undef, first(m.feterms).rank), m)

"""
    fixefnames(m::MixedModel)

Return a (permuted and truncated in the rank-deficient case) vector of coefficient names.
"""
function fixefnames(m::LinearMixedModel{T}) where {T}
    Xtrm = first(m.feterms)
    Xtrm.cnames[1:Xtrm.rank]
end

"""
    fnames(m::MixedModel)

Return the names of the grouping factors for the random-effects terms.
"""
fnames(m::MixedModel) = (fname.(m.reterms)...,)

"""
    getθ(m::LinearMixedModel)

Return the current covariance parameter vector.
"""
getθ(m::LinearMixedModel{T}) where {T} = getθ!(Vector{T}(undef, length(m.parmap)), m)

function getθ!(v::AbstractVector{T}, m::LinearMixedModel{T}) where {T}
    pmap = m.parmap
    if length(v) ≠ length(pmap)
        throw(DimensionMismatch("length(v) = $(length(v)) ≠ length(m.parmap) = $(length(pmap))"))
    end
    reind = 1
    λ = first(m.allterms).λ
    for (k, tp) in enumerate(pmap)
        tp1 = first(tp)
        if reind ≠ tp1
            reind = tp1
            λ = m.allterms[tp1].λ
        end
        v[k] = λ[tp[2], tp[3]]
    end
    v
end

function Base.getproperty(m::LinearMixedModel{T}, s::Symbol) where {T}
    if s == :θ || s == :theta
        getθ(m)
    elseif s == :β || s == :beta
        coef(m)
    elseif s == :βs || s == :betas
        βs(m)
    elseif s == :λ || s == :lambda
        getproperty.(m.reterms, :λ)
    elseif s == :σ || s == :sigma
        sdest(m)
    elseif s == :σs || s == :sigmas
        σs(m)
    elseif s == :σρs || s == :sigmarhos
        σρs(m)
    elseif s == :b
        ranef(m)
    elseif s == :objective
        objective(m)
    elseif s == :corr
        vcov(m, corr=true)
    elseif s == :vcov
        vcov(m, corr=false)
    elseif s == :PCA
        NamedTuple{fnames(m)}(PCA.(m.reterms))
    elseif s == :pvalues
        ccdf.(Chisq(1), abs2.(coef(m) ./ stderror(m)))
    elseif s == :stderror
        stderror(m)
    elseif s == :u
        ranef(m, uscale = true)
    elseif s == :lowerbd
        m.optsum.lowerbd
    elseif s == :X
        modelmatrix(m)
    elseif s == :y
        vec(last(m.feterms).x)
    elseif s == :rePCA
        rePCA(m)
    else
        getfield(m, s)
    end
end

function StatsBase.leverage(m::LinearMixedModel{T}) where {T}
    # This can be done more efficiently but reusing existing tools is easier.
    # The i'th leverage value is obtained by replacing the response with the i'th
    # basis vector, updating A and L, then taking the sum of squared values of the
    # last row of L, excluding the last position.
    yorig = copy(m.y)
    l = length(m.allterms)
    value = map(eachindex(yorig)) do i
        fill!(m.y, zero(T))
        m.y[i] = one(T)
        updateL!(reevaluateAend!(m))
        sum(j -> sum(abs2, view(m.L, Block(l, j))), 1:(l-1))
    end
    copyto!(m.y, yorig)
    updateL!(reevaluateAend!(m))
    value
end

function StatsBase.loglikelihood(m::LinearMixedModel)
    if m.optsum.REML
        throw(ArgumentError("loglikelihood not available for models fit by REML"))
    end
    -objective(m) / 2
end

lowerbd(m::LinearMixedModel) = m.optsum.lowerbd

function mkparmap(reterms::Vector{AbstractReMat{T}}) where {T}
    parmap = NTuple{3,Int}[]
    for (k, trm) in enumerate(reterms)
        n = LinearAlgebra.checksquare(trm.λ)
        for ind in trm.inds
            d,r = divrem(ind-1, n)
            push!(parmap, (k, r+1, d+1))
        end
    end
    parmap
end

function StatsBase.modelmatrix(m::LinearMixedModel)
    fe = first(m.feterms)
    if fe.rank == size(fe, 2)
        fe.x
    else
        fe.x[:, invperm(fe.piv)]
    end
end

nθ(m::LinearMixedModel) = length(m.parmap)

StatsBase.nobs(m::LinearMixedModel) = m.dims.n

"""
    objective(m::LinearMixedModel)

Return negative twice the log-likelihood of model `m`
"""
function objective(m::LinearMixedModel{T}) where {T}
    wts = m.sqrtwts
    denomdf = T(ssqdenom(m))
    val = logdet(m) + denomdf * (one(T) + log2π + log(pwrss(m) / denomdf))
    isempty(wts) ? val : val - T(2.0) * sum(log, wts)
end

StatsBase.predict(m::LinearMixedModel) = fitted(m)

Base.propertynames(m::LinearMixedModel, private::Bool = false) = (
    :formula,
    :sqrtwts,
    :A,
    :L,
    :optsum,
    :θ,
    :theta,
    :β,
    :beta,
    :λ,
    :lambda,
    :stderror,
    :σ,
    :sigma,
    :σs,
    :sigmas,
    :b,
    :u,
    :lowerbd,
    :X,
    :y,
    :corr,
    :vcov,
    :PCA,
    :rePCA,
    :reterms,
    :feterms,
    :allterms,
    :objective,
    :pvalues,
)

"""
    pwrss(m::LinearMixedModel)

The penalized, weighted residual sum-of-squares.
"""
pwrss(m::LinearMixedModel) = abs2(last(m.L))

"""
    ranef!(v::Vector{Matrix{T}}, m::MixedModel{T}, β, uscale::Bool) where {T}

Overwrite `v` with the conditional modes of the random effects for `m`.

If `uscale` is `true` the random effects are on the spherical (i.e. `u`) scale, otherwise
on the original scale
"""
function ranef!(
    v::Vector,
    m::LinearMixedModel{T},
    β::AbstractArray{T},
    uscale::Bool,
) where {T}
    (k = length(v)) == m.dims.nretrms || throw(DimensionMismatch(""))
    L = m.L
    for j = 1:k
        mul!(
            vec(copyto!(v[j], getblock(L, length(m.allterms), j))),
            getblock(L, k + 1, j)',
            β,
            -one(T),
            one(T),
        )
    end
    for i = k:-1:1
        Lii = getblock(L, i, i)
        vi = vec(v[i])
        ldiv!(adjoint(isa(Lii, Diagonal) ? Lii : LowerTriangular(Lii)), vi)
        for j = 1:(i-1)
            mul!(vec(v[j]), getblock(L, i, j)', vi, -one(T), one(T))
        end
    end
    if !uscale
        for (t, vv) in zip(m.reterms, v)
            lmul!(t.λ, vv)
        end
    end
    v
end

ranef!(v::Vector, m::LinearMixedModel, uscale::Bool) = ranef!(v, m, fixef(m), uscale)

"""
    ranef(m::LinearMixedModel; uscale=false, named=false)

Return, as a `Vector{Matrix{T}}`, the conditional modes of the random effects in model `m`.

If `uscale` is `true` the random effects are on the spherical (i.e. `u`) scale, otherwise on
the original scale.
"""
function ranef(m::LinearMixedModel{T}; uscale = false, named = false) where {T}
    reterms = m.reterms
    v = [Matrix{T}(undef, size(t.z, 1), nlevs(t)) for t in reterms]
    ranef!(v, m, uscale)
end

LinearAlgebra.rank(m::LinearMixedModel) = first(m.feterms).rank

"""
    rePCA(m::LinearMixedModel; corr::Bool=true)

Return a named tuple of the normalized cumulative variance of a principal components
analysis of the random effects covariance matrices or correlation
matrices when `corr` is `true`.

The normalized cumulative variance is the proportion of the variance for the first
principal component, the first two principal components, etc.  The last element is
always 1.0 representing the complete proportion of the variance.
"""
function rePCA(m::LinearMixedModel; corr::Bool=true)
    pca = PCA.(m.reterms, corr=corr)
    NamedTuple{fnames(m)}(getproperty.(pca,:cumvar))
end

"""
    PCA(m::LinearMixedModel; corr::Bool=true)

Return a named tuple of the principal components analysis of the random effects
covariance matrices or correlation matrices when `corr` is `true`.
"""

function PCA(m::LinearMixedModel; corr::Bool=true)
    NamedTuple{fnames(m)}(PCA.(m.reterms, corr=corr))
end

"""
    reevaluateAend!(m::LinearMixedModel)

Reevaluate the last column of `m.A` from `m.feterms`.  This function should be called
after updating the response, `m.feterms[end]`.
"""
function reevaluateAend!(m::LinearMixedModel)
    A = m.A
    trmn = reweight!(last(m.allterms), m.sqrtwts)
    nblk = length(m.allterms)
    for (j, trm) in enumerate(m.allterms)
        mul!(getblock(A, nblk, j), trmn', trm)
    end
    m
end

"""
    refit!(m::LinearMixedModel[, y::Vector]; REML=m.optsum.REML)

Refit the model `m` after installing response `y`.

If `y` is omitted the current response vector is used.
"""
function refit!(m::LinearMixedModel; REML=m.optsum.REML)
    m.optsum.feval = -1
    fit!(reevaluateAend!(m); REML=REML)
end

function refit!(m::LinearMixedModel, y; REML=m.optsum.REML)
    resp = last(m.feterms)
    length(y) == size(resp, 1) || throw(DimensionMismatch(""))
    copyto!(resp, y)
    refit!(m; REML=REML)
end

StatsBase.residuals(m::LinearMixedModel) = response(m) .- fitted(m)

StatsBase.response(m::LinearMixedModel) = vec(last(m.feterms).x)

function reweight!(m::LinearMixedModel, weights)
    sqrtwts = map!(sqrt, m.sqrtwts, weights)
    reweight!.(m.allterms, Ref(sqrtwts))
    updateA!(m)
    updateL!(m)
end

"""
    sdest(m::LinearMixedModel)

Return the estimate of σ, the standard deviation of the per-observation noise.
"""
sdest(m::LinearMixedModel) = √varest(m)

"""
    setθ!(m::LinearMixedModel, v)

Install `v` as the θ parameters in `m`.
"""
function setθ!(m::LinearMixedModel{T}, θ::Vector{T}) where {T}
    parmap, reterms = m.parmap, m.reterms
    length(θ) == length(parmap) || throw(DimensionMismatch())
    reind = 1
    λ = first(reterms).λ
    for (tv, tr) in zip(θ, parmap)
        tr1 = first(tr)
        if reind ≠ tr1
            reind = tr1
            λ = reterms[tr1].λ
        end
        λ[tr[2], tr[3]] = tv
    end
    m
end

Base.setproperty!(m::LinearMixedModel, s::Symbol, y) =
    s == :θ ? setθ!(m, y) : setfield!(m, s, y)

function Base.show(io::IO, m::LinearMixedModel)
    if m.optsum.feval < 0
        @warn("Model has not been fit")
        return nothing
    end
    n, p, q, k = size(m)
    REML = m.optsum.REML
    println(io, "Linear mixed model fit by ", REML ? "REML" : "maximum likelihood")
    println(io, " ", m.formula)
    oo = objective(m)
    if REML
        println(io, " REML criterion at convergence: ", oo)
    else
        nums = Ryu.writefixed.([-oo / 2, oo, aic(m), aicc(m), bic(m)], 4)
        fieldwd = max(maximum(textwidth.(nums)) + 1, 11)
        for label in ["  logLik", "-2 logLik", "AIC", "AICc", "BIC"]
            print(io, rpad(lpad(label, (fieldwd + textwidth(label)) >> 1), fieldwd))
        end
        println(io)
        print.(Ref(io), lpad.(nums, fieldwd))
        println(io)
    end
    println(io)

    show(io, VarCorr(m))

    print(io, " Number of obs: $n; levels of grouping factors: ")
    join(io, nlevs.(m.reterms), ", ")
    println(io)
    println(io, "\n  Fixed-effects parameters:")
    show(io, coeftable(m))
end

"""
    size(m::LinearMixedModel)

Returns the size of a mixed model as a tuple of length four:
the number of observations, the number of (non-singular) fixed-effects parameters,
the number of conditional modes (random effects), the number of grouping variables
"""
function Base.size(m::LinearMixedModel)
    dd = m.dims
    dd.n, dd.p, sum(size.(m.reterms, 2)), dd.nretrms
end

#=
"""
    sqrtpwrss(m::LinearMixedModel)

Return the square root of the penalized, weighted residual sum-of-squares (pwrss).

This is the element in the lower-right of `m.L`
"""
sqrtpwrss(m::LinearMixedModel) = last(m.L)
=#

"""
    ssqdenom(m::LinearMixedModel)

Return the denominator for penalized sums-of-squares.

For MLE, this value is the number of observations. For REML, this
value is the number of observations minus the rank of the fixed-effects matrix.
The difference is analagous to the use of n or n-1 in the denominator when
calculating the variance.
"""
function ssqdenom(m::LinearMixedModel)::Int
    n = m.dims.n
    m.optsum.REML ? n - m.dims.p : n
end

"""
    std(m::MixedModel)

Return the estimated standard deviations of the random effects as a `Vector{Vector{T}}`.

FIXME: This uses an old convention of isfinite(sdest(m)).  Probably drop in favor of m.σs
"""
function Statistics.std(m::LinearMixedModel)
    rl = rowlengths.(m.reterms)
    s = sdest(m)
    isfinite(s) ? rmul!(push!(rl, [1.0]), s) : rl
end

"""
    stderror!(v::AbstractVector, m::LinearMixedModel)

Overwrite `v` with the standard errors of the fixed-effects coefficients in `m`

The length of `v` should be the total number of coefficients (i.e. `length(coef(m))`).
When the model matrix is rank-deficient the coefficients forced to `-0.0` have an
undefined (i.e. `NaN`) standard error.
"""
function stderror!(v::AbstractVector{T}, m::LinearMixedModel{T}) where {T}
    L = feL(m)
    scr = Vector{T}(undef, size(L, 2))
    s = sdest(m)
    fill!(v, zero(T) / zero(T))  # initialize to appropriate NaN for rank-deficient case
    for i in eachindex(scr)
        fill!(scr, false)
        scr[i] = true
        v[i] = s * norm(ldiv!(L, scr))
    end
    invpermute!(v, first(m.feterms).piv)
    v
end

function StatsBase.stderror(m::LinearMixedModel{T}) where {T}
    stderror!(similar(first(m.feterms).piv, T), m)
end

"""
    updateA!(m::LinearMixedModel)

Update the cross-product array, `m.A`, from `m.allterms`

This is usually done after a reweight! operation.
"""
function updateA!(m::LinearMixedModel)
    allterms = m.allterms
    k = length(allterms)
    A = m.A
    for j = 1:k
        for i = j:k
            mul!(getblock(A, i, j), allterms[i]', allterms[j])
        end
    end
    m
end

"""
    updateL!(m::LinearMixedModel)

Update the blocked lower Cholesky factor, `m.L`, from `m.A` and `m.reterms` (used for λ only)

This is the crucial step in evaluating the objective, given a new parameter value.
"""
function updateL!(m::LinearMixedModel{T}) where {T}
    A = m.A
    L = m.L
    k = length(m.allterms)
    for j = 1:k                         # copy lower triangle of A to L
        for i = j:k
            copyto!(getblock(L, i, j),
                    getblock(A, i, j))
        end
    end
    for (j, cj) in enumerate(m.reterms)  # pre- and post-multiply by Λ, add I to diagonal
        scaleinflate!(getblock(L, j, j), cj)
        for i = (j+1):k         # postmultiply column by Λ
            rmulΛ!(getblock(L, i, j), cj)
        end
        for jj = 1:(j-1)        # premultiply row by Λ'
            lmulΛ!(cj', getblock(L, j, jj))
        end
    end
    for j = 1:k                         # blocked Cholesky
        Ljj = getblock(L, j, j)
        for jj = 1:(j-1)
            rankUpdate!(Hermitian(Ljj, :L), getblock(L, j, jj), -one(T), one(T))
        end
        cholUnblocked!(Ljj, Val{:L})
        LjjT = isa(Ljj, Diagonal) ? Ljj : LowerTriangular(Ljj)
        for i = (j+1):k
            Lij = getblock(L, i, j)
            for jj = 1:(j-1)
                mul!(Lij, getblock(L, i, jj), getblock(L, j, jj)', -one(T), one(T))
            end
            rdiv!(Lij, LjjT')
        end
    end
    m
end

"""
    varest(m::LinearMixedModel)

Returns the estimate of σ², the variance of the conditional distribution of Y given B.
"""
varest(m::LinearMixedModel) = pwrss(m) / ssqdenom(m)

"""
    _zerocorr!(m::LinearMixedModel[, trmnms::Vector{Symbol}])

Rewrite the random effects specification for the grouping factors in `trmnms` to zero correlation parameter.

The default for `trmnms` is all the names of random-effects terms.

A random effects term is in the zero correlation parameter configuration when the off-diagonal elements of
λ are all zero - hence there are no correlation parameters in that term being estimated.

Note that this is numerically equivalent to specifying a formula with `zerocorr` around each random effects
term, but the `formula`  fields in the resulting model will differ. In particular, `zerocorr!` will **not**
change the original `formula`'s terms to be of type of `ZeroCorr` because this would involve changing
immutable types.  This may have implications for software that manipulates the formula of a fitted model.

This is an internal function and may disappear in a future release without being considered a breaking change.
"""
function _zerocorr!(m::LinearMixedModel{T}, trmns) where {T}
    reterms = m.reterms
    for trm in reterms
        if fname(trm) in trmns
            zerocorr!(trm)
        end
    end
    newparmap = mkparmap(reterms)
    copyto!(m.parmap, newparmap)
    resize!(m.parmap, length(newparmap))
    optsum = m.optsum
    optsum.lowerbd = mapfoldl(lowerbd, vcat, reterms)
    optsum.initial = mapfoldl(getθ, vcat, reterms)
    optsum.final = copy(optsum.initial)
    optsum.xtol_abs = fill!(copy(optsum.initial), 1.0e-10)
    optsum.initial_step = T[]

    # the model is no longer fitted
    optsum.feval = -1

    m
end

_zerocorr!(m::LinearMixedModel) = _zerocorr!(m, fnames(m))


text:
@testset "mcs" begin

@testset "MCES" begin
    null = smilestomol("")
    @test mcesmolsize(null, null) == 0

    BuOH = smilestomol("CCCO")
    butane = smilestomol("CCCC")
    @test mcesmolsize(butane, BuOH) == 2

    hexane = smilestomol("CCCCCC")
    cyclohexane = smilestomol("C1CCCCC1")
    @test mcesmolsize(hexane, cyclohexane) == 5

    tms = smilestomol("C[Si](C)(C)C")
    tsm = smilestomol("[Si]C([Si])([Si])[Si]")
    @test mcesmolsize(tms, tsm) == 1

    sulfide = smilestomol("CSSC")
    disconn = smilestomol("CS.SC")
    @test mcesmolsize(sulfide, disconn) == 2

    # Delta-Y
    cyclopropane = smilestomol("C1CC1")
    isopropane = smilestomol("CC(C)C")
    @test mcesmolsize(cyclopropane, isopropane) == 2

    tetrahedrane = smilestomol("C12C3C1C23")
    fused = smilestomol("C1C2C1C2")
    spiro = smilestomol("C1CC12CC2")
    @test mcesmolsize(tetrahedrane, fused) == 5
    @test mcesmolsize(tetrahedrane, spiro) == 4
end

end # mcs


text:
let path = joinpath(@__DIR__, "pkgs")
    if path ∉ LOAD_PATH
        pushfirst!(LOAD_PATH, path)
    end
end

using Test
using Aqua
using AquaTesting: @testtestset, AquaTesting, with_sample_pkgs