using MAT
using Regression
using DecisionTree

type PredictionCenter
	TEST_MODE_MIXTURE_TPS #default value:0
	TEST_MODE_DATASET #default value:1
	NUM_TPS_SAMPLES  #default value:100
	testVar
	taskName
	workloadName
	lockType
	learnLock
	whichTransactionToPlot
	ioConf
	lockConf

	trainConfig
	testMode
	testConfig
	testMixture
	testMinTPS
	testMaxTPS
	testWorkloadRatio

	testSampleTPS
	testSampleTransactionCount

	throttleLatencyType
	throttleTargetLatency
	throttleTargetTransactionIndex
	throttlePenalty
	throttleIndividualTransactions #default value: false;
end
PredictionCenter()=PredictionCenter(0,0,100,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,false)

type PredictionConfig
	io_conf
	lock_conf
	transactionType
	groupingStrategy
	datasetList #default value: {};
	initialized #default value: false;
	mv # this can be grouped or ungrouped depending on group parameters.
	mvUngrouped # this is always ungrouped.
	configSummary

	transactionCount # C
	averageCpuUsage # P
	diskWrite # IO
	transactionLatency # L
	transactionLatencyPercentile
	lockWaitTime # W, NumOfWaitDueToLocks
	currentLockWait # LocksBeingWaitedFor
	TPS # TPS
	rowsChanged # RowsChanged
	pagesFlushed # PagesFlushed
	transactionMixture # Mixture
	logicalReads
	physicalReads
	physicalReadsMB
	networkSendKB
	networkRecvKB
	logWriteMB # LogIOw

	transactionCountUngrouped # C
	averageCpuUsageUngrouped # P
	diskWriteUngrouped # IO
	transactionLatencyUngrouped # L
	transactionLatencyPercentileUngrouped
	lockWaitTimeUngrouped # W, NumOfWaitDueToLocks
	currentLockWaitUngrouped # LocksBeingWaitedFor
	TPSUngrouped # TPS
	rowsChangedUngrouped # RowsChanged
	pagesFlushedUngrouped # PagesFlushed
	transactionMixtureUngrouped # Mixture
	logicalReadsUngrouped
	physicalReadsUngrouped
	physicalReadsMBUngrouped
	networkSendKBUngrouped
	networkRecvKBUngrouped
	logWriteMBUngrouped # LogIOw
end
PredictionConfig()=PredictionConfig(0,0,0,0,{},false,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0)

type typeOfPredictions
	R
	T_total
	M_total
	Vp
	V
	W
	Pcon
	totalWaits
	TimeSpentWaiting
	LocksBeingHeld
end

type typeOfFlushRate
	io_conf
	workloadName
end

type typeOfConcurrencyLevel
	lock_conf
	workloadName
end


#!/usr/bin/env python
# -*- coding: utf-8 -*-

#from __future__ import division, print_function

using Optim
include("transit_model_sampled.jl")
include("transit_orb.jl")
include("occultquad.jl")
include("kepler.jl")
#import os
#import sys
#import corner
#import emcee3
#import pickle
#import fitsio
using FITSIO
#import numpy as np
#import matplotlib.pyplot as plt
#from scipy.optimize import minimize

#using celerite
include("../src/compile_matrix_symm.jl")
include("../src/compute_likelihood.jl")
using PyPlot
#from plot_setup import setup, get_figsize
#from transit_model import RotationTerm, TransitModel

#setup()
#np.random.seed(42)

fname= "data/kplr001430163-2013011073258_llc.fits"
fits_data = FITS(fname);
t_tmp = float(read(fits_data[2],"TIME"))
sapq = read(fits_data[2],"SAP_QUALITY")
y_tmp = float(read(fits_data[2],"PDCSAP_FLUX"))
yerr_tmp = float(read(fits_data[2],"PDCSAP_FLUX_ERR"))
nt = size(t_tmp)[1]
# See how many "good" data points there are:
good = falses(nt)
good_count = 0
for i=1:nt
  if isfinite(t_tmp[i]) && isfinite(y_tmp[i]) && sapq[i] == 0
    good[i] = true
    good_count += 1
  end
end
# Now, make arrays to store the good data:
N =  1000
t = zeros(N)
y = zeros(N)
yerr = zeros(N)
count = 1
for i=1:nt
  if isfinite(t_tmp[i]) && isfinite(y_tmp[i]) && sapq[i] == 0 && count <= N
    t[count]=t_tmp[i]
    y[count]=y_tmp[i]
    yerr[count]=yerr_tmp[i]
    count +=1
  end
end
dtmax = 0.0
for i=2:N
  if t[i]-t[i-1] > dtmax
    dtmax = t[i]-t[i-1]
  end
end
# The final number of data points we have is count:
# Shift the zero-point of time to be the middle of the dataset:
t -= 0.5*(minimum(t)+maximum(t))
println("Time range: ",minimum(t)," ",maximum(t))
hdr = read_header(fits_data[2]);

#data, hdr = fitsio.read("data/kplr001430163-2013011073258_llc.fits",
#                        header=True)

texp = float(hdr["INT_TIME"] * hdr["NUM_FRM"]) / 60. / 60. / 24.

#N = 1000
#m = data["SAP_QUALITY"] == 0
#m &= np.isfinite(data["TIME"])
#m &= np.isfinite(data["PDCSAP_FLUX"])
#t = np.ascontiguousarray(data["TIME"][m], dtype=np.float64)[:N]
#y = np.ascontiguousarray(data["PDCSAP_FLUX"][m], dtype=np.float64)[:N]
#yerr = np.ascontiguousarray(data["PDCSAP_FLUX_ERR"][m], dtype=np.float64)[:N]
#t -= 0.5 * (t.min() + t.max())

true_params = [
    0.0,         # mean flux
    log(8.0),    # period
    log(0.015),  # Rp / Rs
    log(0.5),    # duration
    0.0,            # t_0
    0.5,            # impact
    0.5,            # q_1
    0.5,            # q_2
]

println("Calling transit model")
true_model = transit_model_sampled(t,true_params)
#u1 = 2.*sqrt(true_params[8])*true_params[9]
#u2 = sqrt(true_params[8])*(1.-2.*true_params[9])
#x = 1.-sqrt(1.-true_params[7]^2)
#depth = exp(2.*true_params[4])*(1.-u1*x-u2*x^2)/(1.-u1/3.-u2/6.)
#print("Depth [ppt]: ",depth*1e3)
plot(t,true_model)
println("Called transit model")
# Build the true model
#true_model = TransitModel(
#    texp,
#    0.0,
#    np.log(8.0),    # period
#    np.log(0.015),  # Rp / Rs
#    np.log(0.5),    # duration
#    0.0,            # t_0
#    0.5,            # impact
#    0.5,            # q_1
#    0.5,            # q_2
#)
#true_params = np.array(true_model.get_parameter_vector())

# Inject the transit into the data
true_transit = 1e-3*true_model + 1.0
y = y.*true_transit
println("Range of y: ",minimum(y)," ",maximum(y))
# Normalize the data
med = median(y)
y = (y./ med - 1.0).* 1e3
yerr *= 1e3 / med

read(STDIN,Char)
# Set up the GP model
#mean = TransitModel(
#    texp,
#    0.0,
#    np.log(8.0),
#    np.log(0.015),
#    np.log(0.5),
#    0.0,
#    0.5,
#    0.5,
#    0.5,
#    bounds=[
#        (-0.5, 0.5),
#        np.log([7.9, 8.1]),
#        (np.log(0.005), np.log(0.1)),
#        (np.log(0.4), np.log(0.6)),
#        (-0.1, 0.1),
#        (0, 1.0), (1e-5, 1-1e-5), (1e-5, 1-1e-5)
#    ]
#)
var = std(y)^2
pname = ["mean flux","log(period [d])","log(R_p/R_*)","log(T [d])","t_0 [d]","b","q1","q2","log(sig^2)","log(amplitude)","log(tau)","log(Pstar)","log(f)"]
gpp = [2.*log(minimum(yerr)*0.5),log(var),log(0.5*maximum(t)),log(4.5),0.0]
#minp = [-0.5,log(7.9),log(.005),log(0.4),-.1,0,  1e-5,  1e-5, log(var*0.01), log(var*0.01), log(dtmax)                ,  -8.0   ]
#maxp = [ 0.5, log(8.1), log(.100), log(0.6), .1, 1.0, 1-1e-5, 1-1e-5, log(var*100.), log(var*100.), log(maximum(t)-minimum(t)), log(5.0)]
lower = [-0.5, log(7.9), log(.005), log(0.4),-.1, 0.0,   1e-5,   1e-5, -12.0, log(var*0.01), log(dtmax)                ,               3.0*log(texp),   -8.0  ]
upper = [ 0.5, log(8.1), log(.100), log(0.6), .1, 1.0, 1-1e-5, 1-1e-5,   2.0, log(var*100.), log(maximum(t)-minimum(t)), 0.5*(maximum(t)-minimum(t)), log(5.0)]
sigma = [1e-4, 1e-4, 1e-3, 1e-3, 1e-3, 0.1, 0.1, 0.1, 0.01, 0.01, 0.1, 0.1, 0.1]

# To do:
# 1).  Add in bounds on parameters. [x]
# 2).  Write a function that calls transit calculation & calls likelihood computation. [ ]
# 3).  Compute the mean GP "prediction". [ ]
# 4).  Run MCMC. [ ]
# 5).  Make plots of lightcurve & results. [ ]

# Likelihood function

# Define width, nex.
p0 = 1
p = 2
m1 = 2(p-p0)+p0+2
width = 2m1+1
nex = (4(p-p0)+2p0+1)*(N-1)+1
aex = zeros(width,nex)
al_small= zeros(m1,nex)
indx = collect(1:nex)
bex = zeros(nex)

param = [true_params;gpp]
println("Initial parameters: ",param)
println("Lower bounds:  ",lower)
println("Upper bounds:  ",upper)

function neg_log_like(param)
# Transiting planet parameters:
ptrans = param[1:8]
# Compute the transit model:
trans_model = transit_model_sampled(t,ptrans)
# Gaussian process rotation kernel parameters:
gpp = param[9:13]
alpha_real = zeros(p)
alpha_imag = zeros(p)
beta_real = zeros(p)
beta_imag = zeros(p)
# Convert from Rotation kernel parameters to celerite kernel parameters:
f = exp(gpp[5])
# Define alpha_real, alpha_imag, beta_real, beta_imag
alpha_real[1] = exp(gpp[2])*(1.+f)/(2.+f)
alpha_real[2] = exp(gpp[2])/(2.+f)
beta_real[1] = exp(-gpp[3])
beta_real[2] = exp(-gpp[3])
beta_imag[2] = 2pi*exp(-gpp[4])
# White noise component:
#w0 = zeros(N)+exp(gpp[1])
w0 = yerr.^2+exp(gpp[1])
# Compute the banded extended matrix:
logdeta = compile_matrix_symm(alpha_real,alpha_imag,beta_real,beta_imag,w0,t,nex,aex,al_small,indx)
# Subtract the mean before computing the likelihood:
#println("y range: ",minimum(y)," ",maximum(y))
#println("tr mod : ",minimum(trans_model)," ",maximum(trans_model))
resid = y-trans_model
# Use banded LU decomposition to compute the log likelihood:
#println("aex range: ",minimum(aex)," ",maximum(aex))
#println("al range: ",minimum(al_small)," ",maximum(al_small))
#println("indx range: ",minimum(indx)," ",maximum(indx))
for i=1:nex
  bex[i]=0.0
end
log_like = compute_likelihood(p,p0,resid,aex,al_small,indx,logdeta,bex)
#println("log det a: ",logdeta)
#println("bex range: ",minimum(bex)," ",maximum(bex))
#println(param)
#println("-log(like): ",-log_like)
return -log_like
end

# Optimize the likelihood:
tic()
println("Initial log-likelihood: ",neg_log_like(param))
println("Optimizing log likelihood")
xdiff = param
result = optimize(DifferentiableFunction(neg_log_like), xdiff, lower, upper, Fminbox(), optimizer=LBFGS, optimizer_o= Optim.Options(autodiff = true))
#result = optimize(DifferentiableFunction(neg_log_like), xdiff, lower, upper, Fminbox())
#result = optimize(neg_log_like, xdiff)

toq()

#kernel = RotationTerm(
#    np.log(np.var(y)), np.log(0.5*t.max()), np.log(4.5), 0.0,
#    bounds=[
#        np.log(np.var(y) * np.array([0.01, 100])),
#        np.log([np.max(np.diff(t)), (t.max() - t.min())]),
#        np.log([3*np.median(np.diff(t)), 0.5*(t.max() - t.min())]),
#        [-8.0, np.log(5.0)],
#    ]
#)

#gp = celerite.GP(kernel, mean=mean, fit_mean=True,
#                 log_white_noise=2*np.log(0.5*yerr.min()),
#                 fit_white_noise=True)
#gp.compute(t, yerr)

param = Optim.minimizer(result)

#println("Initial log-likelihood: ",Optim.minimum(result))
println("Initial log-likelihood: ",neg_log_like(param))
#println("Initial log-likelihood: ",neg_log_like(param))
#println(Optim.minimizer(result))

read(STDIN,Char)

# Now, run a markov chain:

include("aff_inv_mcmc.jl")

# Define the model
def neg_log_like(params, y, gp):
    gp.set_parameter_vector(params)
    return -gp.log_likelihood(y)

# Optimize with random restarts
p0 = gp.get_parameter_vector()
bounds = gp.get_parameter_bounds()
r = minimize(neg_log_like, p0, method="L-BFGS-B", bounds=bounds, args=(y, gp))
gp.set_parameter_vector(r.x)
ml_params = np.array(r.x)
print("Maximum log-likelihood: {0}".format(gp.log_likelihood(y)))

# Compute the maximum likelihood predictions
x = np.linspace(t.min(), t.max(), 5000)
trend = gp.predict(y, t, return_cov=False)
trend -= gp.mean.get_value(t) - gp.mean.mean_flux
mu, var = gp.predict(y, x, return_var=True)
std = np.sqrt(var)
mean_mu = gp.mean.get_value(x)
mu -= mean_mu
wn = np.exp(gp.log_white_noise.value)
ml_yerr = np.sqrt(yerr**2 + wn)

# Plot the maximum likelihood predictions
fig, (ax1, ax2) = plt.subplots(2, 1, sharex=True, figsize=get_figsize(1, 2))
ax1.errorbar(t - t.min(), y, yerr=ml_yerr, fmt=".k", capsize=0)
ax1.plot(x - t.min(), mu)
ax1.set_ylim(-0.72, 0.72)
ax1.yaxis.set_major_locator(plt.MaxNLocator(5))
ax1.set_ylabel("raw [ppt]")

ax2.errorbar(t - t.min(), y-trend, yerr=ml_yerr, fmt=".k", capsize=0)
ax2.plot(x - t.min(), mean_mu - gp.mean.mean_flux)
ax2.set_xlim(0, t.max()-t.min())
ax2.set_ylim(-0.41, 0.1)
ax2.yaxis.set_major_locator(plt.MaxNLocator(5))
ax2.set_ylabel("de-trended [ppt]")
ax2.set_xlabel("time [days]")
fig.savefig("transit-ml.pdf")

# Save the current state of the GP and data
with open("transit.pkl", "wb") as f:
    pickle.dump((gp, y), f, -1)

if os.path.exists("transit.h5"):
    result = input("MCMC save file exists. Overwrite? (type 'yes'): ")
    if result.lower() != "yes":
        sys.exit(0)

# Do the MCMC
def log_prob(params):
    gp.set_parameter_vector(params)
    lp = gp.log_prior()
    if not np.isfinite(lp):
        return -np.inf
    return gp.log_likelihood(y) + lp

# Initialize
print("Running MCMC sampling...")
ndim = len(ml_params)
nwalkers = 32
pos = ml_params + 1e-5 * np.random.randn(nwalkers, ndim)
lp = np.array(list(map(log_prob, pos)))
m = ~np.isfinite(lp)
while np.any(m):
    pos[m] = ml_params + 1e-5 * np.random.randn(m.sum(), ndim)
    lp[m] = np.array(list(map(log_prob, pos[m])))
    m = ~np.isfinite(lp)

# Sample
sampler = emcee3.Sampler(backend=emcee3.backends.HDFBackend("transit.h5"))
with emcee3.pools.InterruptiblePool() as pool:
    ensemble = emcee3.Ensemble(emcee3.SimpleModel(log_prob), pos, pool=pool)
    sampler.run(ensemble, 15000, progress=True)

# Plot the parameter constraints
samples = np.array(sampler.get_coords(discard=5000, flat=True, thin=13))
samples = samples[:, 1:5]
samples[:, :3] = np.exp(samples[:, :3])
truths = np.array(true_params[1:5])
truths[:3] = np.exp(truths[:3])
fig = corner.corner(samples, truths=truths,
                    labels=[r"period", r"$R_\mathrm{P}/R_\star$", r"duration",
                            r"$t_0$"])
fig.savefig("transit-corner.pdf")


module LighthouseFlux

using Zygote: Zygote
using Flux: Flux
using Lighthouse: Lighthouse, classes, log_resource_info!, log_value!

export FluxClassifier

#####
##### `FluxClassifier`
#####

struct FluxClassifier{M,O,C,P,OH,OC} <: Lighthouse.AbstractClassifier
    model::M
    optimiser::O
    classes::C
    params::P
    onehot::OH
    onecold::OC
end

"""
    FluxClassifier(model, optimiser, classes; params=Flux.params(model),
                   onehot=(label -> Flux.onehot(label, 1:length(classes))),
                   onecold=(label -> Flux.onecold(label, 1:length(classes))))

Return a `FluxClassifier <: Lighthouse.AbstractClassifier` with the given arguments:

- `model`: a Flux model. The model must additionally support LighthouseFlux's [`loss`](@ref)
  and [`loss_and_prediction`](@ref) functions.

- `optimiser`: a [Flux optimiser](https://fluxml.ai/Flux.jl/stable/training/optimisers/)

- `classes`: a `Vector` or `Tuple` of possible class values; this is the return
value of `Lighthouse.classes(::FluxClassifier)`.

- `params`: The parameters to optimise during training; generally, a `Zygote.Params`
value or a value that can be passed to `Zygote.Params`.

- `onehot`: the function used to convert hard labels to soft labels when
`Lighthouse.onehot` is called with this classifier.

- `onecold`: the function used to convert soft labels to hard labels when
`Lighthouse.onecold` is called with this classifier.
"""
function FluxClassifier(model, optimiser, classes; params=Flux.params(model),
                        onehot=(label -> Flux.onehot(label, 1:length(classes))),
                        onecold=(label -> Flux.onecold(label, 1:length(classes))))
    return FluxClassifier(Flux.testmode!(model), optimiser, classes, params, onehot,
                          onecold)
end

"""
    loss(model, batch_arguments...)

Return the scalar loss of `model` given `batch_arguments`.

This method must be implemented for all `model`s passed to [`FluxClassifier`](@ref).
"""
function loss end

"""
    loss_and_prediction(model, input_batch, other_batch_arguments...)

Return `(model_loss, model_prediction)` where:

- `model_loss` is equivalent to (and defaults to) `loss(model, input_batch, other_batch_arguments...)`.

- `model_prediction` is a matrix where the `i`th column is the soft label prediction for the `i`th
sample in `input_batch`. Thus, the numnber of columns should be `size(input_batch)[end]`, while the
number of rows is equal to the number of possible classes predicted by model. `model_prediction`
defaults to `model(input_batch)`.

This method must be implemented for all `model`s passed to [`FluxClassifier`](@ref), but
has the default return values described above, so it only needs to be overloaded if the
default definitions do not yield the expected values for a given `model` type. It
additionally may be overloaded to avoid redundant computation if `model`'s loss
function computes soft labels as an intermediate result.
"""
function loss_and_prediction(model, input_batch, other_batch_arguments...)
    return (loss(model, input_batch, other_batch_arguments...), model(input_batch))
end

#####
##### Lighthouse `AbstractClassifier` Interface
#####

Lighthouse.classes(classifier::FluxClassifier) = classifier.classes

function Lighthouse.is_early_stopping_exception(::FluxClassifier, exception)
    return exception isa Flux.Optimise.StopException
end

Lighthouse.onehot(classifier::FluxClassifier, label) = classifier.onehot(label)

Lighthouse.onecold(classifier::FluxClassifier, label) = classifier.onecold(label)

function Lighthouse.train!(classifier::FluxClassifier, batches, logger)
    Flux.trainmode!(classifier.model)
    weights = Zygote.Params(classifier.params)
    for batch in batches
        train_loss, back = log_resource_info!(logger, "train/forward_pass";
                                              suffix="_per_batch") do
            f = () -> loss(classifier.model, batch...)
            return Zygote.pullback(f, weights)
        end
        log_value!(logger, "train/loss_per_batch", train_loss)
        gradients = log_resource_info!(logger, "train/reverse_pass"; suffix="_per_batch") do
            return back(Zygote.sensitivity(train_loss))
        end
        log_resource_info!(logger, "train/update"; suffix="_per_batch") do
            Flux.Optimise.update!(classifier.optimiser, weights, gradients)
            return nothing
        end
    end
    Flux.testmode!(classifier.model)
    return nothing
end

function Lighthouse.loss_and_prediction(classifier::FluxClassifier, batch...)
    return Flux.cpu(loss_and_prediction(classifier.model, batch...))
end

#####
##### `CuIterator`
#####

Base.@deprecate_moved CuIterator CuArrays false

#####
##### miscellaneous utilities
#####

"""
    evaluate_chain_in_debug_mode(chain::Flux.Chain, input)

Evaluate `chain(input)`, printing additional debug information at each layer.
"""
function evaluate_chain_in_debug_mode(chain::Flux.Chain, input)
    for (i, layer) in enumerate(chain)
        @info "Executing layer $i / $(length(chain))..." layer size(input)
        input = layer(input)
        @info output_size = size(input)
    end
    return input
end

end # module


export AbstractMoistureModel, DryModel, EquilMoist, NonEquilMoist

#### Moisture component in atmosphere model
abstract type AbstractMoistureModel end

vars_state(::AbstractMoistureModel, ::AbstractStateType, FT) = @vars()

function atmos_nodal_update_auxiliary_state!(
    ::AbstractMoistureModel,
    m::AtmosModel,
    state::Vars,
    aux::Vars,
    t::Real,
) end
function compute_gradient_flux!(
    ::AbstractMoistureModel,
    diffusive,
    ∇transform,
    state,
    aux,
    t,
) end

function compute_gradient_argument!(
    ::AbstractMoistureModel,
    transform::Vars,
    state::Vars,
    aux::Vars,
    t::Real,
) end

internal_energy(atmos::AtmosModel, state::Vars, aux::Vars) =
    internal_energy(atmos, atmos.orientation, state, aux)

@inline function internal_energy(
    atmos::AtmosModel,
    orientation::Orientation,
    state::Vars,
    aux::Vars,
)
    Thermodynamics.internal_energy(
        density(atmos, state, aux),
        state.energy.ρe,
        state.ρu,
        gravitational_potential(orientation, aux),
    )
end

"""
    DryModel

Assumes the moisture components is in the dry limit.
"""
struct DryModel <: AbstractMoistureModel end

vars_state(::DryModel, ::Auxiliary, FT) = @vars(θ_v::FT, air_T::FT)
@inline function atmos_nodal_update_auxiliary_state!(
    moist::DryModel,
    atmos::AtmosModel,
    state::Vars,
    aux::Vars,
    t::Real,
)
    ts = new_thermo_state(atmos, state, aux)
    aux.moisture.θ_v = virtual_pottemp(ts)
    aux.moisture.air_T = air_temperature(ts)
    nothing
end

"""
    EquilMoist

Assumes the moisture components are computed via thermodynamic equilibrium.
"""
Base.@kwdef struct EquilMoist{FT, IT} <: AbstractMoistureModel
    maxiter::IT = nothing
    tolerance::FT = nothing
end

vars_state(::EquilMoist, ::Prognostic, FT) = @vars(ρq_tot::FT)
vars_state(::EquilMoist, ::Primitive, FT) = @vars(q_tot::FT)
vars_state(::EquilMoist, ::Gradient, FT) = @vars(q_tot::FT)
vars_state(::EquilMoist, ::GradientFlux, FT) = @vars(∇q_tot::SVector{3, FT})
vars_state(::EquilMoist, ::Auxiliary, FT) =
    @vars(temperature::FT, θ_v::FT, q_liq::FT, q_ice::FT)

@inline function atmos_nodal_update_auxiliary_state!(
    moist::EquilMoist,
    atmos::AtmosModel,
    state::Vars,
    aux::Vars,
    t::Real,
)
    ts = new_thermo_state(atmos, state, aux)
    aux.moisture.temperature = air_temperature(ts)
    aux.moisture.θ_v = virtual_pottemp(ts)
    aux.moisture.q_liq = PhasePartition(ts).liq
    aux.moisture.q_ice = PhasePartition(ts).ice
    nothing
end

function compute_gradient_argument!(
    moist::EquilMoist,
    transform::Vars,
    state::Vars,
    aux::Vars,
    t::Real,
)
    ρinv = 1 / state.ρ
    transform.moisture.q_tot = state.moisture.ρq_tot * ρinv
end

function compute_gradient_flux!(
    moist::EquilMoist,
    diffusive::Vars,
    ∇transform::Grad,
    state::Vars,
    aux::Vars,
    t::Real,
)
    # diffusive flux of q_tot
    diffusive.moisture.∇q_tot = ∇transform.moisture.q_tot
end

"""
    NonEquilMoist

Does not assume that the moisture components are in equilibrium.
"""
struct NonEquilMoist <: AbstractMoistureModel end

vars_state(::NonEquilMoist, ::Prognostic, FT) =
    @vars(ρq_tot::FT, ρq_liq::FT, ρq_ice::FT)
vars_state(::NonEquilMoist, ::Primitive, FT) =
    @vars(q_tot::FT, q_liq::FT, q_ice::FT)
vars_state(::NonEquilMoist, ::Gradient, FT) =
    @vars(q_tot::FT, q_liq::FT, q_ice::FT)
vars_state(::NonEquilMoist, ::GradientFlux, FT) = @vars(
    ∇q_tot::SVector{3, FT},
    ∇q_liq::SVector{3, FT},
    ∇q_ice::SVector{3, FT}
)
vars_state(::NonEquilMoist, ::Auxiliary, FT) = @vars(temperature::FT, θ_v::FT)

@inline function atmos_nodal_update_auxiliary_state!(
    moist::NonEquilMoist,
    atmos::AtmosModel,
    state::Vars,
    aux::Vars,
    t::Real,
)
    ts = new_thermo_state(atmos, state, aux)
    aux.moisture.temperature = air_temperature(ts)
    aux.moisture.θ_v = virtual_pottemp(ts)
    nothing
end

function compute_gradient_argument!(
    moist::NonEquilMoist,
    transform::Vars,
    state::Vars,
    aux::Vars,
    t::Real,
)
    ρinv = 1 / state.ρ
    transform.moisture.q_tot = state.moisture.ρq_tot * ρinv
    transform.moisture.q_liq = state.moisture.ρq_liq * ρinv
    transform.moisture.q_ice = state.moisture.ρq_ice * ρinv
end

function compute_gradient_flux!(
    moist::NonEquilMoist,
    diffusive::Vars,
    ∇transform::Grad,
    state::Vars,
    aux::Vars,
    t::Real,
)
    # diffusive fluxes of moisture variables
    diffusive.moisture.∇q_tot = ∇transform.moisture.q_tot
    diffusive.moisture.∇q_liq = ∇transform.moisture.q_liq
    diffusive.moisture.∇q_ice = ∇transform.moisture.q_ice
end


module Materials

using Images
using FileIO

using ..GfxBase

export Material, Texture
export ShadingModel, Flat, Normal
export PhysicalShadingModel, Lambertian, BlinnPhong

export get_diffuse

## Shading model type hierarchy ##
# Types are as follows:
#
# ShadingModel
#   - Flat - simply the diffuse color of the object
#   - Normal - color-code a pixel according to its surface normal
#   - PhysicalShadingModel
#       - Lambertian - diffuse shading
#       - BlinnPhong - diffuse+specular shading

abstract type ShadingModel end
struct Flat<:ShadingModel end
struct Normal <: ShadingModel end

abstract type PhysicalShadingModel <: ShadingModel end

struct Lambertian <: PhysicalShadingModel end

mutable struct BlinnPhong <: PhysicalShadingModel
    specular_color::RGB{Float32} # color of the highlight
    specular_exp::Float64 # "sharpness" exponent
end

## Texture struct definition ##
mutable struct Texture
    image_data::Array{RGB{Float32},2}
    repeat::Bool
end

""" Texture constructor - loads image data from a file"""
function Texture(image_fn::String, repeat::Bool)
    image_data = load(image_fn)
    Texture(image_data, repeat)
end

## Material struct definition ##
mutable struct Material
    shading_model::ShadingModel
    mirror_coeff::Float64
    texture::Union{Texture,Nothing}
    diffuse_color::Union{RGB{Float32},Nothing}
end

""" Get the diffuse color of a material; if the material is textured,
provide the uv coordinate on the object of that material. """
function get_diffuse(material::Material, uv::Union{Vec2,Nothing})
    if isnothing(material.texture) || isnothing(uv)
        return material.diffuse_color
    else
        return get_texture_value(material.texture, uv)
    end
end

""" Look up a texture value given the uv coordinates """
function get_texture_value(texture::Texture, uv::Vec2)
    # this equation is from the book
    i = convert(Int, round(uv[1]*size(texture.image_data, 1)-.5))
    j = convert(Int, round(size(texture.image_data,2)-uv[2]*size(texture.image_data,2)-.5))
    i =  i%size(texture.image_data,1)
    j =  j%size(texture.image_data,2)

    # make sure we dont index a value that isnt valid
    if i < 1
        i += size(texture.image_data,1)
    end
    if j < 1
        j += size(texture.image_data, 2)
    end
    # get rgb value from uv coord
    rgb = texture.image_data[j, i]
    return rgb
end

end # module Materials


"""
ULMFiT - Binary Sentiment Analyzer

    BinSentimentClassifier()

This is a binary sentiment classifier developed after
fine-tuning the ULMFiT language model on IMDB movie reviews dataset.

# Usage:

julia> sc = BinSentimentClassifier()

julia> doc = StringDocument("this classifier is great")

julia> sc(doc)
"positive"

"""

struct BinSentimentClassifier
    vocab::Vector
    rnn_layers::Flux.Chain
    linear_layers::Flux.Chain
end

function BinSentimentClassifier()
    BSON.@load datadep"ULMFiT Sentiment Classifier/umlfit_sentiment_en.bson" weights
    vocab_sz, em_sz = size(weights[1])
    hid_lstm_sz = 1150
    out_lstm_sz = em_sz
    clsfr_hid_sz = 50
    clsfr_out_sz = 2
    vocab = (string.(readdlm("vocabs/sc_vocab.csv", ',')))[:, 1]
    sc = SentimentClassifier(
        vocab,
        Chain(
            DroppedEmbeddings(vocab_sz, em_sz),
            LSTM(em_sz, hid_lstm_sz),
            LSTM(hid_lstm_sz, hid_lstm_sz),
            LSTM(hid_lstm_sz, out_lstm_sz)
        ),
        Chain(
            PooledDense(out_lstm_sz, clsfr_hid_sz),
            BatchNorm(clsfr_hid_sz, relu),
            Dense(clsfr_hid_sz, clsfr_out_sz, sigmoid),
            BatchNorm(clsfr_out_sz),
            softmax
        )
    )
    Flux.loadparams!(sc, weights)
    sc = mapleaves(Tracker.data, sc)
    Flux.testmode!(sc)
    return sc
end

Flux.@treelike BinSentimentClassifier

function (sc::BinSentimentClassifier)(x::TokenDocument)
    remove_case!(x)
    idxs = map(w -> indices([w], sc.vocab, "_unk_"), tokens(x))
    h = sc.rnn_layers.(idxs)
    h = sc.linear_layers(h)
    Flux.reset!(sc.rnn_layers)
    return argmax(h)[1] == 1 ? "positive" : "negative"
end


using RefractionCorrection
using Test
# using LazIO

@testset "RefractionCorrection.jl" begin

    @testset "Correct file" begin
        #MAIN#
        # ds = LazIO.open(filename)
        #
        # #subset
        # pp = points[1500000:1520000]
        #
        # #Array with coordinates
        # coords = coordinates(pp,h)
        #
        # #Separate laser pulses
        # laser_pulses = laserpulses(coords)
        #
        # #Return corrected points
        # corrected_laser_pulses, indexes_corr, indexes_un_corr  = corrected_points(laser_pulses,coords)
        #
        # #write pointcloud
        # write_pointcloud(ds,h,pp,indexes_corr,indexes_un_corr)
        @test 1 == 1
    end

    @testset "Correct" begin
        minz, maxz = -6, -5
        p1 = 0, 0, maxz
        p2 = 0, 0, minz
        _, _, z = correct(p1, p2)
        @test minz < z < maxz
    end

    @testset "Polar" begin
        p = 1., 1., 1.
        x = cart2pol(p...)
        pp = pol2cart(x...)
        @test all(p .≈ pp)
    end
end


using Oceananigans.Diagnostics: AdvectiveCFL, DiffusiveCFL
using Oceananigans.Simulations: TimeStepWizard
using Oceananigans.Utils: prettytime
using Printf

export SimpleProgressMessenger, SingleLineProgressMessenger, TimedProgressMessenger

get_Δt(Δt) = Δt
get_Δt(wizard::TimeStepWizard) = wizard.Δt

function SimpleProgressMessenger_function(simulation; LES=false, SI_units=true,
                                           initial_wall_time_seconds=1e-9*time_ns())
    model = simulation.model
    Δt = simulation.Δt

    i, t = model.clock.iteration, model.clock.time

    progress = 100 * (t / simulation.stop_time)

    current_wall_time = 1e-9 * time_ns() - initial_wall_time_seconds

    u_max = maximum(abs, model.velocities.u)
    v_max = maximum(abs, model.velocities.v)
    w_max = maximum(abs, model.velocities.w)

    if SI_units
        @info @sprintf("[%06.2f%%] i: % 6d,     time: % 10s,     Δt: % 10s,     wall time: % 8s",
                        progress, i, prettytime(t), prettytime(get_Δt(Δt)), prettytime(current_wall_time),)
        if LES
            ν_max = maximum(abs, model.diffusivity_fields.νₑ)
            @info @sprintf("          └── max(|u⃗|): [%.2e, %.2e, %.2e] m/s,     adv CFL: %.2e,     diff CFL: %.2e,     νₘₐₓ: %.2e m²/s",
                           u_max, v_max, w_max, AdvectiveCFL(Δt)(model), DiffusiveCFL(Δt)(model), ν_max)
        else
            @info @sprintf("          └── max(|u⃗|): [%.2e, %.2e, %.2e] m/s,     adv CFL: %.2e,     diff CFL: %.2e",
                           u_max, v_max, w_max, AdvectiveCFL(Δt)(model), DiffusiveCFL(Δt)(model))
        end

    else
        @info @sprintf("[%06.2f%%] i: % 6d,     time: %13.2f,     Δt: %13.2f,     wall time: % 8s",
                       progress, i, t, get_Δt(Δt), prettytime(current_wall_time))
        if LES
            ν_max = maximum(abs, model.diffusivity_fields.νₑ)
            @info @sprintf("          └── max(|u⃗|): [%.2e, %.2e, %.2e],     adv CFL: %.2e,     diff CFL: %.2e,     νₘₐₓ: %.2e",
                           u_max, v_max, w_max, AdvectiveCFL(Δt)(model), DiffusiveCFL(Δt)(model), ν_max)
        else
            @info @sprintf("          └── max(|u⃗|): [%.2e, %.2e, %.2e],     adv CFL: %.2e,     diff CFL: %.2e",
                           u_max, v_max, w_max, AdvectiveCFL(Δt)(model), DiffusiveCFL(Δt)(model))
        end
    end

    @info ""

    return nothing
end

function SimpleProgressMessenger(; kwargs...)
    return simulation -> SimpleProgressMessenger_function(simulation; kwargs...)
end

function SingleLineProgressMessenger_func(simulation; LES=false, SI_units=true,
                                          initial_wall_time_seconds=1e-9*time_ns())
    model = simulation.model
    Δt = simulation.Δt

    i, t = model.clock.iteration, model.clock.time

    progress = 100 * (t / simulation.stop_time)

    current_wall_time = 1e-9 * time_ns() - initial_wall_time_seconds

    if SI_units
        if LES
            ν_max = maximum(abs, model.diffusivity_fields.νₑ)
            @info @sprintf("[%06.2f%%] i: % 6d,     time: %10s,     Δt: %10s,     wall time: %8s,     adv CFL: %.2e,     diff CFL: %.2e,     νₘₐₓ: %.2e m²/s",
                           progress, i, prettytime(t), prettytime(get_Δt(Δt)), prettytime(current_wall_time),
                           AdvectiveCFL(Δt)(model), DiffusiveCFL(Δt)(model), ν_max)
        else
            @info @sprintf("[%06.2f%%] i: % 6d,     time: %10s,     Δt: %10s,     wall time: %8s,     adv CFL: %.2e,     diff CFL: %.2e",
                           progress, i, prettytime(t), prettytime(get_Δt(Δt)), prettytime(current_wall_time),
                           AdvectiveCFL(Δt)(model), DiffusiveCFL(Δt)(model))
        end

    else
        if LES
            ν_max = maximum(abs, model.diffusivity_fields.νₑ)
            @info @sprintf("[%06.2f%%] i: % 6d,     time: %13.2f,     Δt: %13.2f,     wall time: % 8s,     adv CFL: %.2e,     diff CFL: %.2e,     νₘₐₓ: %.2e",
                           progress, i, t, get_Δt(Δt), prettytime(current_wall_time), AdvectiveCFL(Δt)(model), DiffusiveCFL(Δt)(model), ν_max)
        else
            @info @sprintf("[%06.2f%%] i: % 6d,     time: %13.2f,     Δt: %13.2f,     wall time: % 8s,     adv CFL: %.2e,     diff CFL: %.2e",
                           progress, i, t, get_Δt(Δt), prettytime(current_wall_time), AdvectiveCFL(Δt)(model), DiffusiveCFL(Δt)(model))
        end
    end

    return nothing
end

function SingleLineProgressMessenger(; kwargs...)
    return simulation -> SingleLineProgressMessenger_func(simulation; kwargs...)
end

mutable struct TimedProgressMessenger{T, A, D, L, Δ} <: Function
    wall_time₀ :: T  # Wall time at simulation start
    wall_time⁻ :: T  # Wall time at previous calback
       adv_cfl :: A
       dif_cfl :: D
           LES :: L
            Δt :: Δ
end

TimedProgressMessenger(Δt) =
    TimedProgressMessenger(
                      1e-9 * time_ns(),
                      1e-9 * time_ns(),
                      AdvectiveCFL(Δt),
                      DiffusiveCFL(Δt),
                      false,
                      Δt)

function TimedProgressMessenger(; Δt=0, LES=false,
                       wall_time₀=1e-9*time_ns(), wall_time⁻=1e-9*time_ns(),
                       adv_cfl=nothing, dif_cfl=nothing)

    adv_cfl = adv_cfl==nothing ? AdvectiveCFL(Δt) : adv_cfl
    dif_cfl = dif_cfl==nothing ? DiffusiveCFL(Δt) : dif_cfl
    return TimedProgressMessenger(wall_time₀, wall_time⁻, adv_cfl, dif_cfl, LES, Δt)
end

function (pm::TimedProgressMessenger)(simulation)
    model = simulation.model
    Δt = simulation.Δt

    i, t = model.clock.iteration, model.clock.time

    progress = 100 * (t / simulation.stop_time)

    current_wall_time = 1e-9 * time_ns() - pm.wall_time₀
    time_since_last_callback = 1e-9 * time_ns() - pm.wall_time⁻
    wall_time_per_step = time_since_last_callback / simulation.iteration_interval
    pm.wall_time⁻ = 1e-9 * time_ns()

    u_max = maximum(abs, model.velocities.u)
    v_max = maximum(abs, model.velocities.v)
    w_max = maximum(abs, model.velocities.w)

    @info @sprintf("[%06.2f%%] iteration: % 6d, time: % 10s, Δt: % 10s, wall time: % 8s (% 8s / time step)",
                    progress, i, prettytime(t), prettytime(get_Δt(Δt)), prettytime(current_wall_time), prettytime(wall_time_per_step))

    if pm.LES
        ν_max = maximum(abs, model.diffusivity_fields.νₑ)
        @info @sprintf("          └── max(|u⃗|): [%.2e, %.2e, %.2e] m/s, CFL: %.2e, νCFL: %.2e, ν_max: %.2e m²/s",
                        u_max, v_max, w_max, pm.adv_cfl(model), pm.dif_cfl(model), ν_max)
    else
        @info @sprintf("          └── max(|u⃗|): [%.2e, %.2e, %.2e] m/s, CFL: %.2e, νCFL: %.2e",
                        u_max, v_max, w_max, pm.adv_cfl(model), pm.dif_cfl(model))
    end

    @info ""

    return nothing
end


@testset "McDonald Kreitman" begin

    @testset "mkt" begin
        testseqs = [[dna"AAAGGGCCC", dna"AGAGGGCGC"],
                    [dna"AACGGGCCG", dna"ACAGGGCCC"]]

        @test mkt(NaturalSelection.aligned_codons(testseqs[1]),
                  NaturalSelection.aligned_codons(testseqs[2])) == (2, 3, 0, 1, 1.0)
    end

end


using Compose
using Colors

tomato_bisque =
           compose(context(),
                   (context(), circle(), fill(colorant"bisque")),
                   (context(), rectangle(), fill(colorant"tomato")))
img = SVG("introspect.svg", 4inch, 4(sqrt(3)/2)inch)

draw(img, introspect(tomato_bisque))


using CUDA
using DLPack
using PyCall
using Test

@info "Installing JAX"
#
packages = "jax[cpu]"
run(PyCall.python_cmd(`-m pip install $packages`))

# Load JAX
jax = pyimport("jax")
np = pyimport("jax.numpy")
dlpack = pyimport("jax.dlpack")

@testset "DLPack.jl" begin
    v = np.asarray([1.0, 2.0, 3.0])

    dlv = DLVector{Float32}(@pycall dlpack.to_dlpack(v)::PyObject)
    tensor = dlv.manager.dl_tensor

    @test tensor.ndim == 1
    @test tensor.dtype == DLPack.jltypes_to_dtypes()[Float32]

    if DLPack.device_type(dlv) == DLPack.kDLCPU
        jv = unsafe_wrap(Array, dlv)
        jv[1] = 0  # mutate a jax's tensor
    elseif DLPack.device_type(dlv) == DLPack.kDLGPU
        jv = unsafe_wrap(CuArray, dlv)
        jv[1:1] .= 0  # mutate a jax's tensor
    end

    @test py"$np.all($v[:] == $np.asarray([0.0, 2.0, 3.0])).copy().item()"
end


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

"""
    Ensemble

An ensemble of geostatistical realizations.

The result of a geostatistical simulation solver is a
set of geospatial data (a.k.a. realizations). This set
can be interpreted as an empirical approximation of the
random field of interest.

### Notes

`Ensemble` objects are typically created by a call to the
`solve` function on a [`SimulationProblem`](@ref).
"""
struct Ensemble{𝒟,ℛ}
  domain::𝒟
  reals::ℛ
  nreals::Int

  function Ensemble{𝒟,ℛ}(domain, reals) where {𝒟,ℛ}
    n = [length(r) for r in reals]
    @assert length(unique(n)) == 1 "number of realizations must be unique"
    new(domain, reals, first(n))
  end
end

Ensemble(domain::𝒟, reals::ℛ) where {𝒟,ℛ} =
  Ensemble{𝒟,ℛ}(domain, reals)

==(e₁::Ensemble, e₂::Ensemble) =
  e₁.domain == e₂.domain && e₁.reals == e₂.reals

Meshes.domain(ensemble::Ensemble) = ensemble.domain

# -------------
# VARIABLE API
# -------------

Base.getindex(ensemble::Ensemble, var::Symbol) =
  ensemble.reals[var]

# -------------
# ITERATOR API
# -------------

Base.iterate(ensemble::Ensemble, state=1) =
  state > ensemble.nreals ? nothing : (ensemble[state], state + 1)
Base.length(ensemble::Ensemble) = ensemble.nreals

# --------------
# INDEXABLE API
# --------------

function Base.getindex(ensemble::Ensemble, ind::Int)
  sdomain = ensemble.domain
  sreals  = pairs(ensemble.reals)
  idata   = (; (var => reals[ind] for (var, reals) in sreals)...)
  georef(idata, sdomain)
end
Base.getindex(ensemble::Ensemble, inds::AbstractVector{Int}) =
  [getindex(ensemble, ind) for ind in inds]
Base.firstindex(ensemble::Ensemble) = 1
Base.lastindex(ensemble::Ensemble) = length(ensemble)

# -----------
# STATISTICS
# -----------

function mean(ensemble::Ensemble)
  varreals = pairs(ensemble.reals)
  μs = (; (v => mean(rs) for (v, rs) in varreals)...)
  georef(μs, ensemble.domain)
end

function var(ensemble::Ensemble)
  varreals = pairs(ensemble.reals)
  σs = (; (v => var(rs) for (v, rs) in varreals)...)
  georef(σs, ensemble.domain)
end

function quantile(ensemble::Ensemble, p::Number)
  cols = []
  for (variable, reals) in pairs(ensemble.reals)
    quantiles = map(1:nelements(ensemble.domain)) do location
      slice = getindex.(reals, location)
      quantile(slice, p)
    end
    push!(cols, variable => quantiles)
  end
  georef((; cols...), ensemble.domain)
end

quantile(ensemble::Ensemble, ps::AbstractVector) =
  [quantile(ensemble, p) for p in ps]

# -----------
# IO METHODS
# -----------

function Base.show(io::IO, ensemble::Ensemble)
  N = embeddim(ensemble.domain)
  print(io, "$(N)D Ensemble")
end

function Base.show(io::IO, ::MIME"text/plain", ensemble::Ensemble)
  println(io, ensemble)
  println(io, "  domain: ", ensemble.domain)
  println(io, "  variables: ", join(keys(ensemble.reals), ", ", " and "))
  print(  io, "  N° reals:  ", ensemble.nreals)
end


using Images, AxisArrays
using Base: OneTo
using Unitful: Time, Length, s

struct Acquisition <: AbstractAcquisition # {{{
    "Acquired images"
    images::AxisArray{<:AbstractGray,3}
    "Numerical aperature"
    NA::Real
    "Magnification"
    M::Integer
    "Fluorophore marker"
    marker::FluorophoreMarker
    "Pixel size"
    pixelsize::Length
end # }}}

function Acquisition( # {{{
    images::AbstractArray{<:Color,3},
    acquisitiontimes::AbstractVector{<:Time},
    NA,
    M,
    marker::FluorophoreMarker,
    pixelsize::Length,
)
    if !(eltype(images) <: AbstractGray)
        images = Gray.(images)
    end
    images_axis = AxisArray(
        images,
        Axis{:x}(OneTo(size(images, 1))),
        Axis{:y}(OneTo(size(images, 2))),
        Axis{:time}(acquisitiontimes),
    )
    return Acquisition(images_axis, NA, M, marker, pixelsize)
end # }}}

function Acquisition( # {{{
    images::AbstractArray{<:Color,3},
    acquisitiontimestep::Time,
    NA,
    M,
    marker::FluorophoreMarker,
    pixelsize::Length,
)
    return Acquisition(
        images,
        range(0s; length=size(images, 3), step=acquisitiontimestep),
        NA,
        M,
        marker,
        pixelsize,
    )
end # }}}

times(acquisition::AbstractAcquisition) = acquisition.images[Axis{:time}].val


JC69() = JC69rel()
JC69(λ::F; safe::Bool=true) where F <: Float64 = JC69abs(λ, safe=safe)

function JC69(θ_vec::A;
              safe::Bool=true) where A <: AbstractArray
  if length(θ_vec) == 0
    return JC69rel()
  elseif length(θ_vec) == 1
    return JC69abs(θ_vec[1], safe=safe)
  else
    error("Parameter vector length incompatiable with absolute or relative rate form of substitution model")
  end
end

function JC69rel(θ_vec::A;
                 safe::Bool=true) where A <: AbstractArray
  if safe
    if length(θ_vec) != 0
      error("Incorrect parameter vector length")
    end
  end
  return JC69rel()
end

function JC69abs(θ_vec::A;
                 safe::Bool=true) where A <: AbstractArray
  if safe
    if length(θ_vec) != 1
      error("Incorrect parameter vector length")
    end
  end
  return  JC69abs(θ_vec[1], safe=safe)
end


import Base.append!

struct ElasticCovStrat <: CovarianceStrategy
    capacity::Int
    stepsize::Int
end
function init_precompute(covstrat::ElasticCovStrat, X, y, k)
    nobs = size(X, 2)
    FullCovariancePrecompute(nobs)
end

append!(gp, x::AbstractVector, y::Float64) = append!(gp, reshape(x, :, 1), [y])
function append!(gp::GPE{X,Y,M,K,CS,D,P}, x::AbstractMatrix, y::AbstractVector) where {X,Y,M,K,CS,D,P <: ElasticPDMat}
    size(x, 2) == length(y) || error("$(size(x, 2)) observations, but $(length(y)) targets.")
    newcov = [cov(gp.kernel, gp.x, x); cov(gp.kernel, x, x) + (noise_variance(gp) + eps())*I]
    append!(gp.data, gp.kernel, gp.x, x)
    append!(gp.x, x)
    append!(gp.cK, newcov)
    gp.nobs += length(y)
    append!(gp.y, y)
    update_target!(gp, kern = false, noise = false)
end

LinearAlgebra.ldiv!(cK::ElasticPDMat, x) = ldiv!(cK.chol, x)
wrap_cK(cK::ElasticPDMat, Σbuffer, chol::Cholesky) = cK
mat(cK::ElasticPDMat) = view(cK.mat)
cholfactors(cK::ElasticPDMat) = view(cK.chol).factors

function ElasticGPE(dim; mean::Mean = MeanZero(), kernel = SE(0.0, 0.0),
                    logNoise = -2.0, kwargs...)
    x = ElasticArray(Array{Float64}(undef, dim, 0))
    y = ElasticArray(Array{Float64}(undef, 0))
    ElasticGPE(x, y, mean, kernel, logNoise; kwargs...)
end
function alloc_cK(covstrat::ElasticCovStrat, nobs)
    # create placeholder PDMat
    m    = Matrix{Float64}(undef, nobs, nobs)
    chol = Matrix{Float64}(undef, nobs, nobs)
    cK = ElasticPDMat(m, Cholesky(chol, :U, 0), capacity=covstrat.capacity, stepsize=covstrat.stepsize)
    return cK
end
function ElasticGPE(x::AbstractMatrix, y::AbstractVector, mean::Mean, kernel::Kernel,
                    logNoise = -2.0;
                    capacity = 10^3, stepsize = 10^3)
    data = ElasticKernelData(kernel, x, x, capacity=capacity, stepsize=stepsize)
    nobs = length(y)
    # create placeholder PDMat
    covstrat = ElasticCovStrat(capacity, stepsize)
    cK = alloc_cK(covstrat, nobs)
    gp = GPE(ElasticArray(x), ElasticArray(y), mean, kernel, logNoise, covstrat, data, cK)
    initialise_target!(gp)
end

function prepareappend!(kd, Xnew)
    dim, nobs_new = size(Xnew)
    nobs = kd.dims[1]
    if nobs + nobs_new > kd.capacity[1]
        kd.capacity = kd.capacity .+ kd.stepsize
        ElasticPDMats.resize!(kd)
    end
    kd, dim, nobs, nobs_new
end

#===========================
    Elastic Kernel Data
===========================#

function ElasticKernelData(k::Isotropic, X1::AbstractMatrix, X2::AbstractMatrix; capacity = 10^3, stepsize = 10^3)
    @assert X1==X2 # TODO: extend to X1 != X2
    X = X1
    kerneldata = IsotropicData(AllElasticArray(2; capacity = (capacity, capacity), stepsize = (stepsize, stepsize)))
    nobs = size(X, 2)
    distance!(view(kerneldata.R, 1:nobs, 1:nobs), k, X)
    setdimension!(kerneldata.R, nobs, 1:2)
    kerneldata
end
function append!(kerneldata::IsotropicData{<:AllElasticArray}, k::Isotropic, X::AbstractMatrix, Xnew::AbstractMatrix)
    kd, dim, nobs, nobs_new = prepareappend!(kerneldata.R, Xnew)
    distance!(view(kd, 1:nobs, nobs + 1:nobs + nobs_new), k, X, Xnew)
    copyto!(view(kd, nobs + 1:nobs + nobs_new, 1:nobs), transpose(view(kd, 1:nobs, nobs + 1:nobs + nobs_new)))
    distance!(view(kd, nobs + 1:nobs + nobs_new, nobs + 1:nobs + nobs_new), k, Xnew)
    setdimension!(kd, nobs + nobs_new, 1:2)
    kerneldata
end
append!(kerneldata, k, X, Xnew::AbstractVector) = append!(kerneldata, k, X, reshape(Xnew, :, 1))

function ElasticKernelData(k::StationaryARD, X1::AbstractMatrix, X2::AbstractMatrix; capacity = 10^3, stepsize = 10^3)
    @assert X1==X2 # TODO: extend to X1 != X2
    X = X1
    dim, nobs = size(X)
    dist_stack = AllElasticArray(3; capacity = (capacity, capacity, size(X, 1)),
                                    stepsize = (stepsize, stepsize, 0))
    for d in 1:dim
        grad_ls = view(dist_stack, 1:nobs, 1:nobs, d)
        distance!(grad_ls, SqEuclidean(), view(X, d:d, :))
    end
    setdimension!(dist_stack, nobs, 1:2)
    setdimension!(dist_stack, dim, 3)
    StationaryARDData(dist_stack)
end
function append!(kerneldata::StationaryARDData, kernel::StationaryARD, X::AbstractMatrix, Xnew::AbstractMatrix)
    kd, dim, nobs, nobs_new = prepareappend!(kerneldata.dist_stack, Xnew)
    for d in 1:dim
        grad_ls = view(kd, 1:nobs, nobs + 1:nobs + nobs_new, d)
        distance!(grad_ls, SqEuclidean(), view(X, d:d, :), view(Xnew, d:d, :))
        copyto!(view(kd, nobs + 1:nobs + nobs_new, 1:nobs, d),
                transpose(grad_ls))
        distance!(view(kd, nobs + 1:nobs + nobs_new,
                       nobs + 1:nobs + nobs_new, d), SqEuclidean(), view(Xnew, d:d, :))
    end
    setdimension!(kd, nobs + nobs_new, 1:2)
    kerneldata
end

function ElasticKernelData(k::LinArd, X1::AbstractMatrix, X2::AbstractMatrix; capacity = 10^3, stepsize = 10^3)
    @assert X1==X2 # TODO: extend to X1 != X2
    X = X1
    dim, nobs = size(X)
    XtX_d = AllElasticArray(3; capacity = (capacity, capacity, size(X, 1)),
                                    stepsize = (stepsize, stepsize, 0))
    @inbounds @simd for d in 1:dim
        for i in 1:nobs
            for j in 1:i
                XtX_d[i, j, d] = XtX_d[j, i, d] = X[d, i] * X[d, j]
            end
        end
    end
    setdimension!(XtX_d, nobs, 1:2)
    setdimension!(XtX_d, dim, 3)
    LinArdData(XtX_d)
end
function append!(kerneldata::LinArdData, kernel::LinArd, X::AbstractMatrix, Xnew::AbstractMatrix)
    kd, dim, nobs, nobs_new = prepareappend!(kerneldata.XtX_d, Xnew)
    @inbounds @simd for d in 1:dim
        for i in 1:nobs
            for j in 1:nobs_new
                kd[i, nobs + j, d] = kd[nobs + j, i, d] = X[d, i] * Xnew[d, j]
            end
        end
        for i in 1:nobs_new
            for j in 1:i
                kd[nobs + i, nobs + j, d] = kd[nobs + j, nobs + i, d] = Xnew[d, i] * Xnew[d, j]
            end
        end
    end
    setdimension!(kd, nobs + nobs_new, 1:2)
    kerneldata
end

function ElasticKernelData(k::LinIso, X1::AbstractMatrix, X2::AbstractMatrix; capacity = 10^3, stepsize = 10^3)
    @assert X1==X2 # TODO: extend to X1 != X2
    X = X1
    dim, nobs = size(X)
    XtX = AllElasticArray(2, capacity = (capacity, capacity), stepsize = (stepsize, stepsize))
    @inbounds @simd for d in 1:dim
        for i in 1:nobs
            for j in 1:i
                XtX[i, j] = XtX[j, i] += X[d, i] * X[d, j]
            end
        end
    end
    setdimension!(XtX, nobs, 1:2)
    LinIsoData(XtX)
end
function append!(kerneldata::LinIsoData, kernel::LinIso, X::AbstractMatrix, Xnew::AbstractMatrix)
    kd, dim, nobs, nobs_new = prepareappend!(kerneldata.XtX, Xnew)
    @inbounds @simd for d in 1:dim
        for i in 1:nobs
            for j in 1:nobs_new
                kd[i, j + nobs] = kd[j + nobs, i] += X[d, i] * Xnew[d, j]
            end
        end
        for i in 1:nobs_new
            for j in 1:i
                kd[i + nobs, j + nobs] = kd[j + nobs, i + nobs] += Xnew[d, i] * Xnew[d, j]
            end
        end
    end
    setdimension!(kd, nobs + nobs_new, 1:2)
    kerneldata
end


# http://web.mit.edu/puzzle/www/1999/puzzles/1Gumshoe/Warrants1/w1.2/w1.2.html
@testset "warrant 1.2" begin
    for (names, common_letter) in [
        (["racerx", "americanmaid", "kodachi", "ladyjane"], 'a'),
        (["brain", "judyjetson", "jonnyquest", "jeannette"], 'n'),
        (["kenshin", "lisasimpson", "michiganjfrog", "sheila"], 'i'),
        (["bedtimebear", "sherman", "stimpy", "mrmagoo"], 'm'),
        (["bettyboop", "sweetpollypurebred", "skeletor", "firefly"], 'e')]
        f = best_feature(names, 0)
        @test f.description == "contains at least 1 '$common_letter's"
    end
end


mat(ex::SubHam{<:Tuple, <:LocalOperator}) = ex.val

# convert config to indices
Base.to_indices(::Type{L}, inds...) where L <: SiteLabel = map(x->to_indices(L, x), inds)

Base.to_indices(::Type{Bit}, val::Number) = Int(val + 1)
Base.to_indices(::Type{Spin}, val::Number) = Int(div(val + 1, 2))
Base.to_indices(::Type{Half}, val::Number) = Int(val + 0.5)

@generated function Base.to_indices(::Type{Bit}, val::NTuple{N, T}) where {N, T}
    ex = :(val[1])
    for i = 2:N
        factor = 2^(i-1)
        ex = :(val[$i] * $factor + $ex)
    end
    :(ex)
end

@generated function Base.to_indices(::Type{Spin}, val::NTuple{N, T}) where {N, T}
    ex = :(div(val[1] + 1, 2))
    for i = 2:N
        factor = 2^(i-1)
        ex = :(div(val[$i] + 1, 2) * $factor + $ex)
    end
    :(Int($ex))
end

@generated function Base.to_indices(::Type{Half}, val::NTuple{N, T}) where {N, T}
    ex = :(val[1] + 0.5)
    for i = 2:L
        factor = 2^(i-1)
        ex = :((val[$i] + 0.5) * $factor + $ex)
    end
    :(Int($ex))
end

struct NumberOfBody{K} end
NumberOfBody(::NumberOfBody{K}) where K = NumberOfBody{K}()
NumberOfBody(::LocalOperator{T, N}) where {T, N} = NumberOfBody{N}()
NumberOfBody(x::SubHam) = NumberOfBody{length(x.placeholders)}()
NumberOfBody(x::Operator) = NumberOfBody(x.expr)

struct PlaceHolderType{T} end
PlaceHolderType(::SubHam{PT}) where PT = PlaceHolderType{PT}()

# NOTE: only operator acts on a single SubHam
PlaceHolderType(x::Operator) = PlaceHolderType(x.expr)
# PlaceHolderType(x::AddOrSub) =

"""
    LOpEigen{K}

Number of eigen states in local operator.
"""
struct LOpEigen{K} end

LOpEigen(::LocalOperator{T, N, O}) where {T, N, O} = LOpEigen{O}()
LOpEigen(x::Operator) = LOpEigen(x.expr)
LOpEigen(x::AddOrSub) = (LOpEigen(x.lhs); LOpEigen(x.rhs))
LOpEigen(x::SubHam{<:Tuple, <:LocalOperator}) = LOpEigen(x.val)

function LOpEigen(x::SubHam{<:Tuple, <:Tuple})
    n = LOpEigen(first(x.val))

    # check if all the local operator has the same number of eigens
    for each in x.val
        n === LOpEigen(each) || error("number of eigen states in given local operators are not the same")
    end
    n
end

struct LocalOperatorValType{T, OpType} end
LocalOperatorValType(::MT) where {T, MT <: AbstractMatrix{T}} = LocalOperatorValType{T, MT}()
LocalOperatorValType(x::LocalOperator) = LocalOperatorValType(x.val)
LocalOperatorValType(x::SubHam{<:Tuple, <:LocalOperator}) = LocalOperatorValType(x.val)
LocalOperatorValType(x::Operator) = LocalOperatorValType(x.expr)

# multiple operators
LocalOperatorValType(x::SubHam{<:Tuple, <:Tuple}) = LocalOperatorValType(x.val...)
LocalOperatorValType(a::SubHam, x::SubHam...) = LocalOperatorValType(LocalOperatorValType(a), LocalOperatorValType(x...))

LocalOperatorValType(x::LocalOperatorValType{T}...) where T = LocalOperatorValType{T, Tuple{_op_type(x...)...}}()
_op_type(::Tuple{}) = ()
_op_type(::LocalOperatorValType{T, Op}) where {T, Op} = Op
_op_type(::LocalOperatorValType{T, Op1}, ::LocalOperatorValType{T, Op2}, ops...) where {T, Op1, Op2} = (Op1, Op2, _op_type(ops)...)

function is_other_sites_the_same(::NumberOfBody{1}, site, region, lhs::Sites, rhs::Sites)
    @inbounds for other_site in Lattices.sites(region)
        if other_site != site
            s = region[other_site]
            lhs[s] == rhs[s] || return false
        end
    end
    true
end

function is_other_sites_the_same(::NumberOfBody, sites, region, lhs::Sites, rhs::Sites)
    for other_site in Lattices.sites(region)
        # NOTE: we seperate single site because of the following:
        #
        # # chain lattice
        # julia> 1 in 1
        #        true
        #
        # # square lattice
        # julia> (1, 1) in (1, 1)
        #        false
        if !(other_site in sites)
            s = region[other_site]
            lhs[s] == rhs[s] || return false
        end
    end
    true
end

Base.getindex(ex::OnRegion, lhs::Sites, rhs::Sites) = eval_expr(ex, lhs, rhs)

## eval Operator expressions
# unzip OnRegion
eval_expr(ex::OnRegion, lhs::Sites, rhs::Sites) =
    eval_expr(ex.expr, ex.region, lhs, rhs)

# 1. dispatch to specific local operator
eval_expr(ex::Operator, region, lhs::Sites, rhs::Sites) =
    eval_expr(ex, region, LocalOperatorValType(ex), PlaceHolderType(ex), lhs, rhs)

# 2. if there is no specific method for the operator, then fall back to general methods
eval_expr(ex::Operator, region, ::LocalOperatorValType, ::PlaceHolderType, lhs, rhs) =
    eval_expr(ex, region, LOpEigen(ex), NumberOfBody(ex), PlaceHolderType(ex), lhs, rhs)

# H1 + H2
eval_expr(ex::Add, region, lhs::Sites, rhs::Sites) =
    eval_expr(ex.lhs, region, lhs, rhs) + eval_expr(ex.rhs, region, lhs, rhs)
# H1 - H2
eval_expr(ex::Sub, region, lhs::Sites, rhs::Sites) =
    eval_expr(ex.lhs, region, lhs, rhs) - eval_expr(ex.rhs, region, lhs, rhs)

# α * H
eval_expr(ex::ScalarProd, region, lhs::Sites, rhs::Sites) =
    ex.val * eval_expr(ex.expr, region, lhs, rhs)

# sum(A[i]) on vertex
# NOTE:
# The summation looks like
#
# U x I x I + I x U x I + I x I x U
#
# which means the 2x2 local operator U
# defines which config should be the same/or not
#
# So we will check if position (excludes U) are the same in each loop
# this determine if this block is zero.
# TODO:
# 1. for diagnol local operators, we can directly check if all configurations
# are the same.
# 2. for local operators like X, we can directly check if all configurations
# are opposite (regarding to LOpEigen{2} operators)
function eval_expr(
    expr::Sum{<:SubHam}, region,
    ::LOpEigen{2}, n::NumberOfBody{1}, # 2x2 local operator
    ::PlaceHolderType{Tuple{Vertex}},
    lhs::Sites{L}, rhs::Sites{L}) where L

    Aij = zero(eltype(expr))

    @inbounds for site in Lattices.sites(region)
        s = region[site]
        lhs_ind = to_indices(L, lhs[s])
        rhs_ind = to_indices(L, rhs[s])

        # check identities
        Aij += is_other_sites_the_same(n, site, region, lhs, rhs) ? expr.expr[lhs_ind, rhs_ind] : 0
    end
    Aij
end

######################### specialized methods ##################################
# # sum(X[i]) over all vertex
# function eval_expr(
#     ex::Sum{<:SubHam}, region,
#     ::LocalOperatorValType{T, <:Pauli.X},
#     ::PlaceHolderType{Tuple{Vertex}},
#     lhs::Sites{L}, rhs::Sites{L}) where {T, L}
#
#     # check if all the configurations are opposite
#     @inbounds for site in Lattices.sites(region)
#         s = region[site]
#         # NOTE: since this is for Pauli.X
#         # we can assume there is only two
#         # states, so just use !=
#         lhs[s] != rhs[s] || return zero(T)
#     end
#
#     # T(length(lhs))
#     # TODO:
#     # NOT DONE YET
# end

# sum(Z[i]) over all vertex
function eval_expr(
    ex::Sum{<:SubHam}, region,
    ::LocalOperatorValType{T, <:Pauli.Z},
    ::PlaceHolderType{Tuple{Vertex}},
    lhs::Sites{L}, rhs::Sites{L}) where {T, L}

    # check if all the configurations are the same
    @inbounds for site in Lattices.sites(region)
        s = region[site]
        lhs[s] == rhs[s] || return zero(T)
    end

    # down down -> -1
    # up   up   -> +1
    downs, ups = n_downs_ups(lhs)
    T(ups - downs)
end

# sum(Z[i]Z[j]) on neighbors
function eval_expr(
    expr::Sum{<:SubHam}, region,
    ::LOpEigen{2}, n::NumberOfBody{2}, # 2x2 local operator
    ::PlaceHolderType{Tuple{Neighbor{K}, Neighbor{K}}},
    lhs::Sites{L}, rhs::Sites{L}) where {K, L}

    Aij = zero(eltype(expr))
    @inbounds for (site_a, site_b) in Lattices.edges(region, length=K)
        s_a = region[site_a]
        s_b = region[site_b]

        lhs_ind = to_indices(L, lhs[s_a], lhs[s_b])
        rhs_ind = to_indices(L, rhs[s_a], rhs[s_b])

        Aij += is_other_sites_the_same(n, (site_a, site_b), region, lhs, rhs) ? expr.expr[lhs_ind, rhs_ind] : 0
    end
    Aij
end

# partial indexing
# TODO: this should use partial eval
# TODO: add assertion
Base.getindex(ex::OnRegion, lhs::Sites, colon::Colon) = eval_expr(ex.expr, ex.region, lhs, colon)
Base.getindex(ex::OnRegion, colon::Colon, rhs::Sites) = eval_expr(ex.expr, ex.region, colon, rhs)

function eval_expr(
    ex, region,
    lhs::Sites{L, T, N}, ::Colon) where {L, T, N}

    # TODO:
    # wrap this as an iterator?
    # for rhs in IterName
    # blabla
    # end

    nelem = 1<<length(region)
    Ai = Vector{eltype(ex)}(undef, nelem)
    rhs = Sites{L, T, N}(fill(down(L), size(lhs)))
    for i = 1:nelem
        Ai[i] = eval_expr(ex, region, lhs, rhs)
        rhs << 1
    end

    Ai
end

function eval_expr(
    ex, region,
    ::Colon, rhs::Sites{L, T, N}) where {L, T, N}

    # TODO:
    # wrap this as an iterator?
    # for rhs in IterName
    # blabla
    # end

    nelem = 1<<length(region)
    Ai = Vector{eltype(ex)}(undef, nelem)
    lhs = Sites{L, T, N}(fill(down(L), size(rhs)))
    for i = 1:nelem
        Ai[i] = eval_expr(ex, region, lhs, rhs)
        lhs << 1
    end

    Ai
end

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

# TODO:
# sum(A1[i]A2[j]...) on vertex
LocalOperator(x::SubHam{<:Tuple, <:LocalOperator}) = x.val
LocalOperator{T, 1}(x::SubHam{<:Tuple, <:LocalOperator{T, 1}}) where T = x.val
LocalOperator{T, 1}(x::SubHam{<:Tuple, <:Tuple}) where T = prod(LocalOperator{T, 1}, x.val)

# TODO: use a repeated product instead
# 1. A1[i]A2[j] does not go through (1, 1), (2, 2), ..., this should use @subham A1[i]A2[i]
# 2. parse A1[i]A2[i] => (A1 * A2)[i] in make_subham, not in SubHamExpr, we keep A1[i]A2[i] for printing

function eval_expr(ex::Sum{<:SubHam}, region,
    ::LOpEigen{2}, n::NumberOfBody{N}, # 2^Nx2^N local operator
    ::PlaceHolderType{NTuple{N, Vertex}},
    lhs::Sites{L}, rhs::Sites{L}) where {N, L}

    error("Summation over multiple vertex is not supported yet. See issue #1")
    Aij = zero(eltype(ex))

    it = Lattices.sites(region)

    it = Iterators.product((it for i=1:N)...)

    local_h = LocalOperator{eltype(ex), 1}(ex.expr)

    @inbounds for _sites in it
        lhs_ind = Tuple(to_indices(L, lhs[region[site]]) for site in _sites)
        rhs_ind = Tuple(to_indices(L, rhs[region[site]]) for site in _sites)

        Aij += is_other_sites_the_same(n, _sites, region, lhs, rhs) ? ex.expr[lhs_ind, rhs_ind] : 0
    end
    Aij
end

# struct NonZeroIterator
# end
#
# function eval_expr(
#     ex::Sum{<:SubHam}, region,
#     ::LOpEigen{2}, ::NumberOfBody{2}, ::PlaceHolderType{Tuple{Vertex}},
#     lhs::Sites{L}) where L
#
#     Ai = []
#     @inbounds for site in Lattices.sites(region)
#         s = region[site]
#         lhs_ind = to_indices(L, lhs[s])
#         rhs_ind = to_indices(L, rhs[s])
#
#         Aij += expr.expr[lhs_ind, rhs_ind]
#     end
# end


#md# This is a more symmetrical and accurate way to generate comparison wordclouds, but it may be more time consuming.
#md# ### Prepare two wordcloud objects
using WordCloud

stwords = ["us"];
cs = WordCloud.randomscheme() #:Set1_8#
as = WordCloud.randomangles() #(0,90,45,-45)#
dens = 0.5 #not too high
wca = wordcloud(
    processtext(open(pkgdir(WordCloud)*"/res/Barack Obama's First Inaugural Address.txt"), stopwords=WordCloud.stopwords_en ∪ stwords),
    colors = cs,
    angles = as,
    density = dens,
    backgroundcolor = :maskcolor,
    run = x->nothing, #turn off the initword! and placewords! in advance
)
wcb = wordcloud(
    processtext(open(pkgdir(WordCloud)*"/res/Donald Trump's Inaugural Address.txt"), stopwords=WordCloud.stopwords_en ∪ stwords),
    mask = getsvgmask(wca),
    colors = cs,
    angles = as,
    density = dens,
    backgroundcolor = :maskcolor,
    maskcolor = getmaskcolor(wca),
    font = getparameter(wca, :font),
    run = x->nothing,
)
#md# ### Make the same words the same style
samewords = getwords(wca) ∩ getwords(wcb)
println(length(samewords), " same words")
@assert !hasparameter(wca, :uniquewords)
@assert !hasparameter(wcb, :uniquewords)
setparameter!(wca, setdiff(getwords(wca), samewords), :uniquewords)
setparameter!(wcb, setdiff(getwords(wcb), samewords), :uniquewords)
for w in samewords
    setcolors!(wcb, w, getcolors(wca, w))
    setangles!(wcb, w, getangles(wca, w))
end
#md# ### Put the same words at same position
initwords!(wca)
initwords!(wcb)
keep(wca, samewords) do
    placewords!(wca)
    fit!(wca, 1000)
end
pin(wca, samewords) do
    placewords!(wca) #place other words
end
centers = getpositions(wca, samewords, type=getcenter)
setpositions!(wcb, samewords, centers, type=setcenter!) #manually initialize the position,
pin(wcb, samewords) do
    placewords!(wcb) #place other words
end
#md# ### Fit them all
function syncposition(samewords, pos, wca, wcb)
    pos2 = getpositions(wca, samewords, type=getcenter)
    if pos != pos2
        setpositions!(wcb, samewords, pos2, type=setcenter!)
        setstate!(wcb, :placewords!)
    end
    pos2
end
function pinfit!(wc, samewords, ep1, ep2)
    pin(wc, samewords) do
        fit!(wc, ep1)
    end
    fit!(wc, ep2, teleporting=getparameter(wc, :uniquewords)) #only teleport the unique words
end
pos = getpositions(wca, samewords, type=getcenter)
while getparameter(wca, :epoch) < 2000 && getparameter(wcb, :epoch) < 2000
    global pos
    pinfit!(wca, samewords, 200, 50)
    pos = syncposition(samewords, pos, wca, wcb)
    pinfit!(wcb, samewords, 200, 50)
    pos = syncposition(samewords, pos, wcb, wca)
    if getstate(wca) == getstate(wcb) == :fit!
        break
    end
end
println("Takes $(getparameter(wca, :epoch)) and $(getparameter(wca, :epoch)) epochs")
WordCloud.printcollisions(wca)
WordCloud.printcollisions(wcb)
#md#
ma = paint(wca)
mb = paint(wcb)
h,w = size(ma)
space = fill(mb[1], (h, w÷20))
try mkdir("address_compare2") catch end
println("results are saved in address_compare2")
WordCloud.save("address_compare2/compare2.png", [ma space mb])
#eval# try rm("address_compare2", force=true, recursive=true) catch end
gif = WordCloud.GIF("address_compare2")
record(wca, "Obama", gif)
record(wcb, "Trump", gif)
WordCloud.Render.generate(gif, framerate=1)
wca, wcb
#eval# runexample(:compare2)
#md# ![](address_compare2/compare2.png)
#md# ![](address_compare2/result.gif)


@testset "Index" begin
    @testset "NoIndex" begin
        @test project(NoIndex(), ()) === NoIndex()
        @test project(NoIndex(), (:a,)) === NoIndex()
        @test project(NoIndex(), (:a, :b)) === NoIndex()
    end

    @testset "UniqueIndex" begin
        @test project(UniqueIndex{(:a, :b)}(), (:a, :b)) === UniqueIndex{(:a, :b)}()
        @test project(UniqueIndex{(:a, :b)}(), (:b, :a)) === UniqueIndex{(:a, :b)}()
        @test project(UniqueIndex{(:a, :b)}(), (:a,)) === NoIndex()
        @test project(UniqueIndex{(:a, :b)}(), (:b,)) === NoIndex()
    end

    @testset "SortIndex" begin
        v = [3,2,1]
        @test project(SortIndex{(:a, :b)}(v), (:a, :b)) === SortIndex{(:a, :b)}(v)
        @test project(SortIndex{(:a, :b)}(v), (:b, :a)) === SortIndex{(:a, :b)}(v)
        @test project(SortIndex{(:a, :b)}(v), (:a,)) === SortIndex{(:a,)}(v)
        @test project(SortIndex{(:a, :b)}(v), (:b,)) === NoIndex()
    end

    @testset "UniqueSortIndex" begin
        v = [3,2,1]
        @test project(UniqueSortIndex{(:a, :b)}(v), (:a, :b)) === UniqueSortIndex{(:a, :b)}(v)
        @test project(UniqueSortIndex{(:a, :b)}(v), (:b, :a)) === UniqueSortIndex{(:a, :b)}(v)
        @test project(UniqueSortIndex{(:a, :b)}(v), (:a,)) === SortIndex{(:a,)}(v)
        @test project(UniqueSortIndex{(:a, :b)}(v), (:b,)) === NoIndex()
    end

    @testset "HashIndex" begin
        d = Dict((3, 2.0) => [1, 3], (4, 1.0) => [2])
        @test project(HashIndex{(:a, :b)}(d), (:a, :b)) === HashIndex{(:a, :b)}(d)
        @test project(HashIndex{(:a, :b)}(d), (:b, :a)) === HashIndex{(:a, :b)}(d)
        @test project(HashIndex{(:a, :b)}(d), (:a,)) === NoIndex()
        @test project(HashIndex{(:a, :b)}(d), (:b,)) === NoIndex()
    end

    @testset "UniqueHashIndex" begin
        d = Dict((3, 2.0) => 1, (4, 1.0) => 2)
        @test project(UniqueHashIndex{(:a, :b)}(d), (:a, :b)) === UniqueHashIndex{(:a, :b)}(d)
        @test project(UniqueHashIndex{(:a, :b)}(d), (:b, :a)) === UniqueHashIndex{(:a, :b)}(d)
        @test project(UniqueHashIndex{(:a, :b)}(d), (:a,)) === NoIndex()
        @test project(UniqueHashIndex{(:a, :b)}(d), (:b,)) === NoIndex()
    end
end

@testset "Multiple indexes" begin
    @testset "Removing NoIndex" begin
        @test clean(()) === ()
        @test clean((NoIndex(),)) === ()
        @test clean((UniqueIndex{:a}(),)) === (UniqueIndex{:a}(),)
        @test clean((NoIndex(), UniqueIndex{:a}(), NoIndex(), UniqueIndex{:b}())) === (UniqueIndex{:a}(), UniqueIndex{:b}())
    end

    @testset "Projecting multiple indexes" begin
        p = Project{(:a,)}()

        @test p(()) === ()
        @test p((NoIndex(),)) === ()
        @test p((UniqueIndex{(:a,)}(),)) === (UniqueIndex{(:a,)}(),)
        @test p((UniqueIndex{(:b,)}(),)) === ()
        @test p((UniqueIndex{(:a,:b)}(),)) === ()
        @test p((UniqueIndex{(:a,)}(), UniqueIndex{(:b,)}())) === (UniqueIndex{(:a,)}(),)
    end
end


module DistributedBoosting
using HDF5

type Unibeta
	unibeta::Array{Float64,1}
	n::Array{Int,1}
	labels::Array{ASCIIString,1}
end

type Covarmat
	covarmat::Array{Float64,2}
	labels::Array{ASCIIString,1}
end

abstract Abstractdiscohort

type Discohort <: Abstractdiscohort
	x::Array{Float64,2}
	y::Array{Float64,1}
	labels::Array{ASCIIString,1}
end

type Remotecohort <: Abstractdiscohort #hier nochmal schauen
	filename::ASCIIString
	unibeta::Unibeta
	covarmat::Covarmat
end

type Boostscratch
	actualstepno::Int
	actualbeta::Array{Float64,1}
	actualscore::Array{Float64,1}
	pooledunibeta::Unibeta
	pooledcovarmat::Covarmat
	actualupdates::Array{Float64,1}
	selections::Array{ASCIIString,1}
	wantedlabels::Array{ASCIIString,1}
	nu::Float64
	stepno::Int
	actualsel::Int
	actualnom::Array{Float64,1}
	a::Int
	x::Int
end

include("boost.jl")
include("boosting.jl")
include("calc.jl")
include("fetch_pool.jl")
include("getselections.jl")
include("pooling.jl")
include("readwrite.jl")

end


py"""
import os
import torch
import time
import numpy as np
import logging
from torch.optim.lr_scheduler import ReduceLROnPlateau
import torch.nn as nn

def save_checkpoint(state, is_best, checkpoint_dir, logger=None):

    def log_info(message):
        if logger is not None:
            logger.info(message)

    if not os.path.exists(checkpoint_dir):
        log_info(
            f"Checkpoint directory does not exists. Creating {checkpoint_dir}")
        os.mkdir(checkpoint_dir)

    last_file_path = os.path.join(checkpoint_dir, 'last_checkpoint.pt')
    log_info(f"Saving last checkpoint to '{last_file_path}'")
    torch.save(state, last_file_path)
    if is_best:
        best_file_path = os.path.join(checkpoint_dir, 'best_checkpoint.pt')
        log_info(f"Saving best checkpoint to '{best_file_path}'")
        shutil.copyfile(last_file_path, best_file_path)

class Trainer:
    def __init__(self,
                 model,
                 learning_rate = 2e-3,
                 num_epochs = 20,
                 ckpt_dir = os.path.join(os.getcwd(), 'checkpoint_res'),
                 ):

        self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
        self.model = model.to(self.device)
        self.optimizer = torch.optim.Adam(
                                    filter(lambda p: p.requires_grad, self.model.parameters()),
                                    learning_rate)  # leave betas and eps by default
        self.lr_scheduler = ReduceLROnPlateau(self.optimizer, verbose=True, factor= 0.2)
        self.criterion = nn.MSELoss()
        # self.criterion = nn.L1Loss()
        self.num_epochs = num_epochs
        self.loss_history = []
        self.ckpt_dir = ckpt_dir

    def train(self, img, target):
        img.requires_grad = True
        target.requires_grad = True
        img = img.to(self.device)
        target = target.to(self.device)
        initial_loss = self.criterion(img, target)
        print('Initial loss: ', initial_loss)
        self.loss_history.append(initial_loss.item())

        for i in range(self.num_epochs):
            self.model.train()
            start_t = time.time()
            self.optimizer.zero_grad()
            out = self.model(img)
            loss = self.criterion(out, target)
            loss.backward()
            self.optimizer.step()
            end_t = time.time()
            self.loss_history.append(loss.item())

            print('(Epoch {} / {}) train loss: {:.4f} time per epoch: {:.1f}s current lr: {}'.format(
                i + 1, self.num_epochs, loss, end_t - start_t, self.optimizer.param_groups[0]['lr']))

            if (i + 1) % 5 == 0:
                print('Save the current model to checkpoint!')
                save_checkpoint(self.model.state_dict(), is_best= False, checkpoint_dir= self.ckpt_dir)
                torch.save(self.loss_history, os.path.join(self.ckpt_dir, 'train_loss.pt'))

            self.lr_scheduler.step(loss)
"""


# Left division
@inline \(mat::AbstractMatrix, tvec::TransposedVector) = error("Cannot left-divide transposed vector by matrix")

# These four are trivial definitions in Base, resorting to transposition:
#   A_ldiv_Bc
#   A_ldiv_Bt
#   At_ldiv_Bt
#   Ac_ldiv_Bc

# The others need to be overloaded...
@inline At_ldiv_B(mat::AbstractMatrix, tvec::TransposedVector) = error("Cannot left-divide transposed vector by matrix")
@inline Ac_ldiv_B(mat::AbstractMatrix, tvec::TransposedVector) = error("Cannot left-divide transposed vector by matrix")


# This file establishes the default initial conditions, boundary conditions and sources
# for the baroclinicwave_problem experiment, following [Ullrich2014](@cite)

# Override default CLIMAParameters for consistency with literature on this case
CLIMAParameters.Planet.press_triple(::EarthParameterSet) = 610.78

struct BaroclinicWaveProblem{BCS, ISP, ISA, WP, BS, MP} <: AbstractAtmosProblem
    boundaryconditions::BCS
    init_state_prognostic::ISP
    init_state_auxiliary::ISA
    perturbation::WP
    base_state::BS
    moisture_profile::MP
end
function BaroclinicWaveProblem(
    physics::AtmosPhysics;
    boundaryconditions = (AtmosBC(physics;), AtmosBC(physics;)),
    perturbation = nothing,
    base_state = nothing,
    moisture_profile = nothing,
)
    # Set up defaults
    if isnothing(perturbation)
        perturbation = DeterministicPerturbation()
    end
    if isnothing(base_state)
        base_state = BCWaveBaseState()
    end
    if isnothing(moisture_profile)
        moisture_profile = MoistLowTropicsMoistureProfile()
    end

    problem = (
        boundaryconditions,
        init_gcm_experiment!,
        (_...) -> nothing,
        perturbation,
        base_state,
        moisture_profile,
    )
    return BaroclinicWaveProblem{typeof.(problem)...}(problem...)
end

problem_name(::BaroclinicWaveProblem) = "BaroclinicWave"

setup_source(::BaroclinicWaveProblem) = (Gravity(), Coriolis())


"""
    readop4(filename::String)

Searc file and find matrix names (no matname) in Nastran binary OP4 file.
"""
function readop4(filename::String)
    matrices = readop4(filename, MAT_NAMES)

    return collect(keys(matrices))
end

"""
    readop4(filename::String, matname::String)

Read specified matrix `matname` from Nastran binary OP4 file into sparse matrices.
"""
function readop4(filename::String, matname::String)
    file = FortranFiles.FortranFile(filename)

    matrices = Dict{String, Union{SparseMatrixCSC, Nothing}}()

    while !eof(file.io)
        (ncol, nrow, mform, mtype, mname) = FortranFiles.read(file, Int32, Int32, Int32, Int32, FortranFiles.FString{8})
        mname = FortranFiles.trimstring(mname)

        colptr = zeros(Int32, ncol + 1)#Vector{Int32}(ncol + 1)
        rowind = Vector{Int32}(undef, 0)
        data = Vector{Float64}(undef, 0)

        # pre allocate and then resize!, maybe not worth the effort
        sizehint!(rowind, 2ncol)
        sizehint!(data, 2ncol)

        icol_count::Int32 = 1
        if mname == matname || matname == MAT_ALL
            icol_count, matrices[mname] = readop4!(file, data, icol_count, colptr, rowind, ncol, nrow)
        else
            # not sough matrix, keep looking
            spool!(file, ncol)

            if matname == MAT_NAMES
                matrices[mname] = nothing
            end
        end
    end

    if isempty(matrices)
        println("No matrices found, try searching the file for available matrices.")
    end

    return matrices
end

"""
    readop4!(
        file::FortranFiles.FortranFile,
        data::Vector,
        icol_count::Int32,
        colptr::Vector,
        rowind::Vector,
        ncol::Int32,
        nrow::Int32,
    )

Read Nastran binary OP4 file and populate sparse matrix.
"""
function readop4!(
    file::FortranFiles.FortranFile,
    data::Vector,
    icol_count::Int32,
    colptr::Vector,
    rowind::Vector,
    ncol::Int32,
    nrow::Int32,
)
    while true
        @fread file icol::Int32 irow::Int32 nword::Int32 iline::Array{Int32}(undef, nword)

        p::Int32 = 1
        m::Int32 = 1

        if icol > ncol
            # read one extra line after end to terminate matrix
            colptr[ncol + 1] = length(rowind) + 1
            fillcolptr!(colptr, icol, icol_count)
            return icol_count, SparseMatrixCSC(nrow, ncol, convert(Vector{Int}, colptr), convert(Vector{Int}, rowind), data)
        end

        if ncol < COL #|| nrow > 0 # small matrix, comment for now version diff.
            p = readsmallop4!(data, iline, p, nword, colptr, icol, rowind)
        else # big matrix
            p, m = readbigop4!(data, iline, p, nword, colptr, icol, rowind)
        end

        # correct col_ptr
        fillcolptr!(colptr, icol, icol_count)
        icol_count = icol + 1
    end
end

"""
    spool(file::FortranFiles.FortranFile, ncol::Int32)

If nothing found, keep looking for data.
"""
function spool!(file::FortranFiles.FortranFile, ncol::Int32)
    icol2::Int32 = 1
    while icol2 < ncol + 1
        # spool forward
        icol2 = FortranFiles.read(file, Int32)
    end

    return nothing
end

"""
    readsmallop4!(data, iline, p, nword, colptr, icol, rowind)

Read small OP4 matrix.
"""
function readsmallop4!(data, iline, p, nword, colptr, icol, rowind)
    while p <= nword
        if p == 1
            colptr[icol] = length(rowind) + 1
        end
        slength = getlength(iline[p])
        wsize = Int32(slength/2)
        irow = getirow(iline[p])
        #println("p: ", p, ", irow: ", irow, ", slength: ", slength, ", wsize: ", wsize, ", nword: ", nword)

        append!(rowind, irow:(irow + wsize - 1))
        append!(data, reinterpret(Float64, iline[(p + 1):(p + slength)]))

        p = p + slength + 1
    end

    return p
end

"""
    readbigop4!(data, iline, p, nword, colptr, icol, rowind)

Read big OP4 matrix.
"""
function readbigop4!(data, iline, p, nword, colptr, icol, rowind)
    line = reinterpret(Float64, iline)
    while p <= nword
        if p == 1
            colptr[icol] = length(rowind) + 1
        end
        slength = iline[p] - 1
        irow = iline[p+1]
        wsize = Int32(slength/2)
        #println("p: ", p, ", m: " , m, ", irow: ", irow, ", slength: ", slength, ", wsize: ", wsize)

        append!(rowind, irow:(irow + wsize - 1))
        append!(data, line[(m + 1):(m + wsize)])

        p = p + slength + 2
        m = m + wsize + 1
    end

    return p, m
end

"""
    getlength(head::Int32)

Get record length, see Nastran documentation.
"""
getlength(head::Int32) = trunc(Int32, head*iCOL) - 1

"""
    getirow(head::Int32)

Get row index, see Nastran documentation.
"""
getirow(head::Int32) = head - COL*(getlength(head) + 1)

"""
    fillcolptr!(colptr, icol, icol_count)

Fill column pointer, see Nastran documentation.
"""
function fillcolptr!(colptr, icol, icol_count)
    while icol_count < icol
        colptr[icol_count] = colptr[icol]
        icol_count = icol_count + 1
    end

    return nothing
end


# Abstract representation of demand.

#=
Demands are parameterized in terms of how they are registered in time and where
they are located.  Functions will be provided to convert `Demand` into the
appropriate type `StaticLoad`, `InterruptibleLoad`, etc. that is properly
located at a `Bus`.
=#

# Import packages.

using Dates, TimeSeries

# Timeseries for allowable ranges of demand.

"""
An "envelope" of minimum and maximum allowable demands at each time point.

# Type parameters
- `T <: TimeType`: timestamp

# Example
```
example = TimeArray(
    [Time(0) , Time(8)   , Time(9) , Time(17)  , Time(18)],
    [(5., 8.), (10., 12.), (8., 8.), (11., 13.), (3., 4.)]
) :: Envelope{Time}
```
"""
const Envelope{T <: TimeType} = TimeArray{Tuple{Float64,Float64},1,T,Array{Tuple{Float64,Float64},1}}

"""
An "envelope" of minimum and maximum allowable demands at each time point, with a location for the demand.

# Type parameters
- `T <: TimeType`: timestamp
- `L`            : network location

# Example
```
example = TimeArray(
    [Time(0)             , Time(8)             , Time(9)                 , Time(17)           , Time(18)            ],
    [("Home #23", 5., 8.), ("Road #14", 0., 0.), ("Workplace #3", 8., 8.), ("Road #9", 0., 0.), ("Home #23", 3., 4.)]
) :: LocatedEnvelope{Time}
```
"""
const LocatedEnvelope{T <: TimeType, L} = TimeArray{Tuple{L,Float64,Float64},1,T,Array{Tuple{L,Float64,Float64},1}}

# Abstract interface for demand.

"""
The most abstract type of demand.

# Type parameters
- `T <: TimeType`: timestamp
- `L`            : network location
"""
abstract type Demand{T <: TimeType, L} end

"""
The "envelope" of minimum and maximum allowable demands at each time pont, with a location for the demand.

This must be implemented by subtypes of `Demand`.

# Arguments
- `demand :: Demand{T,L}`: the demand
"""
function envelope(demand :: Demand{T,L}) :: LocatedEnvelope{T, L} where L where T <: TimeType
end


using ArgParse
using CSV
using DataFrames
using DelimitedFiles
using LinearAlgebra
using Mamba
using MicrobiomeMixedModels
using Random
using YAML

function parse_commandline()
    s = ArgParseSettings("Fit a model to real-world data.")
    @add_arg_table! s begin
        "input"
            help = "Directory from which to read X.csv, Y.csv, and Z.csv."
        "output"
            help = "Directory to which results are written."
        "--iters", "-i"
            arg_type = Int
            default = 500
            help = "Number of MCMC iterations."
        "--burnin", "-b"
            arg_type = Int
            default = 0
            help = "Number of initial MCMC iterations to remove as burn-in."
        "--thin", "-t"
            arg_type = Int
            default = 1
            help = "Retain one of every 'thin' iterations."
        "--chains", "-c"
            arg_type = Int
            default = 1
            help = "Number of MCMC chains to run."
        "--seed"
            arg_type = Int
            default = 1
            help = "Reseed the random number generator."
        "--monitor"
            nargs = '*'
            help = "YAML defining the nodes to monitor through MCMC."
        "--inits"
            nargs = '*'
            help = "YAML defining initial values of stochastic nodes."
        "--hyper"
            nargs = '*'
            help = "YAML defining hyperparameters of stochastic node priors."
        "--factors"
            arg_type = Int
            default = 0
            help = "Number of factors to use in fitting the MIMIX model."
        "--loadings"
            default = "DL"
            help = "Distributions on loadings (G for Gaussian priors, DL for Dirichlet-Laplace priors)."
        "--permanova"
            help = "Run PERMANOVA with vegan::adonis() in R."
            action = :store_true
        "--post-pred-check"
            help = "Perform posterior predictive checks, requires at least one of :Y, :ϕ, or :θ are monitored."
            action = :store_true
    end
    return parse_args(s)
end

function read_data(dir; L=0)
    println("Reading X.csv")
    X = convert(Matrix{Int}, DataFrame!(CSV.File(joinpath(dir, "X.csv"), datarow=1)))
    println("Reading Y.csv")
    Y = convert(Matrix{Int}, DataFrame!(CSV.File(joinpath(dir, "Y.csv"), datarow=1)))
    println("Reading Z.csv")
    Z = convert(Matrix{Int}, DataFrame!(CSV.File(joinpath(dir, "Z.csv"), datarow=1)))
    N, K = size(Y)
    m = sum(Y, dims=2)
    q = maximum(vec(Z))
    data = Dict{Symbol, Any}(
        :X => X,
        :Y => Y,
        :Z => Z,
        :N => N,
        :K => K,
        :m => m,
        :q => q,
        :L => L
    )
    data[:p] = size(X, 2)
    data[:num_blocking_factors] = size(Z, 2)
    data[:blocking_factor] = Dict{Int, Int}()
    for level in unique(Z)
        bf = findall(vec(any(Z .== level, dims=1)))
        @assert length(bf) == 1
        data[:blocking_factor][level] = bf[1]
    end

    return data
end

args = parse_commandline()

input = abspath(args["input"])
@assert isdir(input)

output = abspath(args["output"])
Random.seed!(args["seed"])

if args["permanova"]
    data = read_data(input)

    println("Beginning PERMANOVA test with vegan::adonis in R.")
    sim = permanova(data)
    global_results = DataFrame(pval = sim[:pval])

    global_test_path = joinpath(output, "global-test.tsv")
    println("Writing global test results to $global_test_path")
    CSV.write(global_test_path, global_results, delim='\t')
else  # mimix
    @assert 0 < args["iters"]   "Iters must be positive"
    @assert 0 <= args["burnin"] "Burn-in must be non-negative"
    @assert 0 < args["thin"]    "Thin must be positive"
    @assert 0 < args["chains"]  "Chains must be positive"

    factors = args["factors"]
    loadings = args["loadings"]
    if (factors > 0) & (loadings == "DL")
        model_type = MIMIX(factors)
    elseif (factors > 0) & (loadings == "G")
        model_type = MIMIXGaussian(factors)
    elseif factors == 0
        model_type = MIMIXNoFactors()
    else
        ValueError("--factors requires a non-negative integer")
    end

    monitor_conf = load_config(AbstractString[abspath(data_path) for data_path in args["monitor"]])
    hyper_conf = load_config(AbstractString[abspath(data_path) for data_path in args["hyper"]])
    inits_conf = load_config(AbstractString[abspath(data_path) for data_path in args["inits"]])

    if args["post-pred-check"] && !monitor_conf[:θ] && (!monitor_conf[:θ_mean] || !monitor_conf[:θ_var])
        error("Posterior prediction checks require monitoring θ or both θ_mean and θ_var, please edit the monitor config.")
    end

    model = get_model(model_type, monitor_conf, hyper_conf)

    data = read_data(input; L = factors)

    inits = get_inits(model_type, inits_conf, data)
    inits = [inits for _ in 1:args["chains"]]

    println("Beginning MCMC sampling")
    mcmc_kwargs = Dict(Symbol(key) => args[key] for key in ["burnin", "thin", "chains"])
    sim = mcmc(model, data, inits, args["iters"]; mcmc_kwargs...)

    monitor_Y = pop!(monitor_conf, :Y)
    monitor_ϕ = pop!(monitor_conf, :ϕ)
    monitor_θ = pop!(monitor_conf, :θ)
    monitor_θ_mean = pop!(monitor_conf, :θ_mean)
    monitor_θ_var = monitor_conf[:θ_var]
    monitor_Λ = pop!(monitor_conf, :Λ)

    for (param, monitor) in monitor_conf
        if monitor
            post = get_post(sim, data, param)
            param = String(param)
            for (letter, name) in MicrobiomeMixedModels.latin
                if startswith(param, letter)
                    param = replace(param, letter => name)
                end
            end
            post_path = joinpath(output, "$param.tsv")
            println("Saving posterior samples of $param to $post_path")
            CSV.write(post_path, post, delim='\t')
        end
    end

    if monitor_Λ
        # Lambda is large, so rather than save every sample we save the posterior mean
        Λ = get_post(sim, data, :Λ)
        Λ_values = convert(Matrix, Λ)
        Λ_postmean = reshape(vec(mean(Λ_values, dims=1)), data[:L], data[:K])
        post_path = joinpath(output, "Lambda-postmean.tsv")
        println("Saving posterior mean of Lambda to $post_path")
        writedlm(post_path, Λ_postmean)

        # We can also store the posterior mean of Lambda' * Lambda
        # ΛtΛs = Matrix{Float64}[]
        # for i in 1:size(Λ_values, 1)
        #     Λ_mat = reshape(Λ_values[i, :], data[:L], data[:K])
        #     push!(ΛtΛs, transpose(Λ_mat) * Λ_mat)
        # end
        # ΛtΛ_postmean = sum(ΛtΛs, dims=3)
        # ΛtΛ_postmean ./= size(Λ_values, 1)
        # post_path = joinpath(output, "LambdatLambda-postmean.tsv")
        # println("Saving posterior mean of Lambda' * Lambda to $post_path")
        # writedlm(post_path, ΛtΛ_postmean)
    end

    if args["post-pred-check"]
        # Posterior predictive checks on sparsity, overdispersion, and ecological diversity (alpha and beta)

        function shannon_diversity(x::Vector{T}) where T <: Real
            p = x ./ sum(x)
            p = p[p .> 0.]
            return -sum(p .* log.(p))
        end

        function simpson_diversity(x::Vector{T}) where T <: Real
            p = x ./ sum(x)
            return sum(abs2, p)
        end

        function braycurtis_diversity(x::Vector{T}, y::Vector{T}) where T <: Real
            c = zero(eltype(x))
            s = zero(eltype(x))
            for i in eachindex(x, y)
                c += min(x[i], y[i])
                s += x[i] + y[i]
            end
            return 2c / s
        end

        function jaccard_diversity(x::Vector{T}, y::Vector{T}) where T <: Real
            p = 0
            q = 0
            r = 0
            z = zero(T)
            for i in eachindex(x, y)
                if (x[i] > z) && (y[i] > z)
                    p += 1
                elseif (x[i] > z) && (y[i] == z)
                    q += 1
                elseif (x[i] == z) && (y[i] > z)
                    r += 1
                end
            end
            return p / (p + q + r)
        end

        function calculate_mean_beta_div_by_group(Y, groups, diversity_fn)
            @assert size(Y, 1) == length(groups)
            mean_div = zeros(maximum(group_ids))
            for group_id in unique(group_ids)
                Y_group = Y[group_id .== group_ids, :]
                pairwise_div = Float64[]
                for i in 1:size(Y_group, 1)
                    for j in i+1:size(Y_group, 1)
                        push!(pairwise_div, diversity_fn(Y_group[i, :], Y_group[j, :]))
                    end
                end
                mean_div[group_id] = mean(pairwise_div)
            end
            return mean_div
        end

        matchrows(A::Matrix, B::Matrix) = mapslices(findall, hcat([all(A .== transpose(B[i, :]), dims=2) for i in 1:size(B, 1)]...), dims=2)

        full_combos = hcat(data[:X], data[:Z])
        group_ids = vec(matchrows(full_combos, unique(full_combos, dims=1)))  # each group is unique combo of treatment, site, and blocks

        obs_max_count = vec(maximum(data[:Y], dims=2) ./ sum(data[:Y], dims=2))
        obs_mean_count = vec(var(data[:Y], dims=2))
        obs_prop_eq_zero = vec(mean(data[:Y] .== 0.0, dims=2))
        obs_prop_leq_one = vec(mean(data[:Y] .<= 1.0, dims=2))
        obs_prop_leq_two = vec(mean(data[:Y] .<= 2.0, dims=2))
        obs_shannon_div = mapslices(shannon_diversity, data[:Y], dims=2)
        obs_simpson_div = mapslices(simpson_diversity, data[:Y], dims=2)
        obs_braycurtis_div = reshape(calculate_mean_beta_div_by_group(data[:Y], group_ids, braycurtis_diversity), maximum(group_ids), 1)
        obs_jaccard_div = reshape(calculate_mean_beta_div_by_group(data[:Y], group_ids, jaccard_diversity), maximum(group_ids), 1)
        writedlm(joinpath(output, "obs-max-count.tsv"), obs_max_count)
        writedlm(joinpath(output, "obs-mean-count.tsv"), obs_mean_count)
        writedlm(joinpath(output, "obs-prop-eq-zero.tsv"), obs_prop_eq_zero)
        writedlm(joinpath(output, "obs-prop-leq-one.tsv"), obs_prop_leq_one)
        writedlm(joinpath(output, "obs-prop-leq-two.tsv"), obs_prop_leq_two)
        writedlm(joinpath(output, "obs-shannon-div.tsv"), obs_shannon_div)
        writedlm(joinpath(output, "obs-simpson-div.tsv"), obs_simpson_div)
        writedlm(joinpath(output, "obs-braycurtis-div.tsv"), obs_braycurtis_div)
        writedlm(joinpath(output, "obs-jaccard-div.tsv"), obs_jaccard_div)

        post_pred_checks = Dict{String, Matrix}()

        if monitor_θ
            θ_post = convert(Matrix, get_post(sim, data, :θ))
            post_pred_checks["theta_mcmc"] = θ_post
        end

        if monitor_θ_mean && monitor_θ_var
            θ_mean_post = convert(Matrix, get_post(sim, data, :θ_mean))
            θ_var_post = convert(Matrix, get_post(sim, data, :θ_var))
            iters = size(θ_mean_post, 1)
            θ_post = zeros(iters, data[:N] * data[:K])
            for iter in 1:iters
                θ_mean = reshape(θ_mean_post[iter, :], data[:N], data[:K])
                θ_new = [rand(MvNormal(θ_mean[i, :], sqrt.(θ_var_post[iter, :]))) for i in 1:data[:N]]
                θ_post[iter, :] = vec(transpose(hcat(θ_new...)))
            end
            post_pred_checks["theta_prior"] = θ_post
        end

        for post_pred_check in post_pred_checks
            name, θ_post = post_pred_check

            n_iter = size(θ_post, 1)
            max_count = zeros(n_iter, data[:N])
            mean_count = zeros(n_iter, data[:N])
            prop_eq_zero = zeros(n_iter, data[:N])
            prop_leq_one = zeros(n_iter, data[:N])
            prop_leq_two = zeros(n_iter, data[:N])
            shannon_div = zeros(n_iter, data[:N])
            simpson_div = zeros(n_iter, data[:N])
            braycurtis_div = zeros(n_iter, maximum(group_ids))
            jaccard_div = zeros(n_iter, maximum(group_ids))
            for iter in 1:n_iter
                ϕ = clr(reshape(θ_post[iter, :], data[:N], data[:K]))
                Y_pred_iter = transpose(hcat([rand(Multinomial(data[:m][i], ϕ[i, :])) for i in 1:data[:N]]...))
                max_count[iter, :] = vec(maximum(Y_pred_iter, dims=2) ./ sum(Y_pred_iter, dims=2))
                mean_count[iter, :] = vec(var(Y_pred_iter, dims=2))
                prop_eq_zero[iter, :] = vec(mean(Y_pred_iter .== 0.0, dims=2))
                prop_leq_one[iter, :] = vec(mean(Y_pred_iter .<= 1.0, dims=2))
                prop_leq_two[iter, :] = vec(mean(Y_pred_iter .<= 2.0, dims=2))
                shannon_div[iter, :] = mapslices(shannon_diversity, Y_pred_iter, dims=2)
                simpson_div[iter, :] = mapslices(simpson_diversity, Y_pred_iter, dims=2)
                braycurtis_div[iter, :] = calculate_mean_beta_div_by_group(Y_pred_iter, group_ids, braycurtis_diversity)
                jaccard_div[iter, :] = calculate_mean_beta_div_by_group(Y_pred_iter, group_ids, jaccard_diversity)
            end

            if !isdir(joinpath(output, name))
                mkdir(joinpath(output, name))
            end

            writedlm(joinpath(output, name, "post-pred-max-count.tsv"), max_count)
            writedlm(joinpath(output, name, "post-pred-mean-count.tsv"), mean_count)
            writedlm(joinpath(output, name, "post-pred-prop-eq-zero.tsv"), prop_eq_zero)
            writedlm(joinpath(output, name, "post-pred-prop-leq-one.tsv"), prop_leq_one)
            writedlm(joinpath(output, name, "post-pred-prop-leq-two.tsv"), prop_leq_two)
            writedlm(joinpath(output, name, "post-pred-shannon-div.tsv"), shannon_div)
            writedlm(joinpath(output, name, "post-pred-simpson-div.tsv"), simpson_div)
            writedlm(joinpath(output, name, "post-pred-braycurtis-div.tsv"), braycurtis_div)
            writedlm(joinpath(output, name, "post-pred-jaccard-div.tsv"), jaccard_div)
        end
    end
end


"""
    AggregationStack{T <: Tuple{Vararg{AbstractAggregation}}} <: AbstractAggregation

A container that implements a concatenation of one or more `AbstractAggregation`s.

Construct with e.g. `mean_aggregation([t::Type, ]d)`, `max_aggregation([t::Type, ]d)` for single
operators and with e.g. `pnormlse_aggregation([t::Type, ]d)` for concatenations. With these calls
all parameters inside operators are initialized randomly as `Float32` arrays, unless type `t` is
further specified. It is also possible to call the constructor directly, see Examples.

Intended to be used as a functor:

    (a::AggregationStack)(x, bags[, w])

where `x` is either `Missing`, `AbstractMatrix` or [`ArrayNode`](@ref),
`bags` is [`AbstractBags`](@ref) structure and optionally `w` is an `AbstractVector` of weights.

TODO mention that it can be constructed by vcatting operators as well.
mention flattening

# Examples
```jldoctest
julia> a = AggregationStack(SegmentedMean(2), SegmentedMax(2))
AggregationStack:
 SegmentedMean(ψ = Float32[0.0, 0.0])
 SegmentedMax(ψ = Float32[0.0, 0.0])

julia> a(Float32[0 1 2; 3 4 5], bags([1:1, 2:3]))
4×2 Matrix{Float32}:
 0.0  1.5
 3.0  4.5
 0.0  2.0
 3.0  5.0

julia> a = AggregationStack(SegmentedMean(2), AggregationStack(SegmentedMax(2)))
AggregationStack:
 SegmentedMean(ψ = Float32[0.0, 0.0])
 SegmentedMax(ψ = Float32[0.0, 0.0])

julia> a = SegmentedMeanMax(Float32, 2)
AggregationStack:
 SegmentedMean(ψ = Float32[0.0, 0.0])
 SegmentedMax(ψ = Float32[0.0, 0.0])

julia> vcat(SegmentedMean(2), SegmentedMax(2))
AggregationStack:
 SegmentedMean(ψ = Float32[0.0, 0.0])
 SegmentedMax(ψ = Float32[0.0, 0.0])

```

See also: [`AbstractAggregation`](@ref), [`SegmentedSum`](@ref), [`SegmentedMax`](@ref),
    [`SegmentedMean`](@ref), [`SegmentedPNorm`](@ref), [`SegmentedLSE`](@ref).
"""
struct AggregationStack{T <: Tuple{Vararg{AbstractAggregation}}} <: AbstractAggregation
    fs::T
    function AggregationStack(fs::Tuple{Vararg{AbstractAggregation}})
        ffs = _flatten_agg(fs)
        new{typeof(ffs)}(ffs)
    end
end

_flatten_agg(t) = tuple(vcat(map(_flatten_agg, t)...)...)
_flatten_agg(a::AggregationStack) = vcat(map(_flatten_agg, a.fs)...)
_flatten_agg(a::AbstractAggregation) = [a]

AggregationStack(fs::AbstractAggregation...) = AggregationStack(fs)

Flux.@functor AggregationStack

function (a::AggregationStack)(x::Union{AbstractArray, Missing}, bags::AbstractBags, args...)
    reduce(vcat, (f(x, bags, args...) for f in a.fs))
end

Flux.@forward AggregationStack.fs Base.getindex, Base.firstindex, Base.lastindex, Base.first,
    Base.last, Base.iterate, Base.eltype

Base.length(a::AggregationStack) = sum(length.(a.fs))
Base.size(a::AggregationStack) = tuple(sum(only, size.(a.fs)))

function Base.show(io::IO, ::MIME"text/plain", a::AggregationStack)
    print(io, "AggregationStack:\n")
    Base.print_array(io, a.fs |> collect)
end

function Base.show(io::IO, a::AggregationStack)
    if get(io, :compact, false)
        print(io, "AggregationStack")
    else
        s = join(repr.(a.fs; context=io), "; ")
        print(io, "[", s, "]")
    end
end


function RBMs.zerosum!(rbm::DenseConvRBM)
    RBMs.zerosum_visible!(rbm)
    RBMs.zerosum_hidden!(rbm)
    return rbm
end

function RBMs.zerosum_visible!(rbm::DenseConvRBM)
    if rbm.visible_dense isa RBMs.Potts
        RBMs.zerosum!(rbm.visible_dense)
        RBMs.zerosum!(rbm.w_dense; dims = 1)
    end
    if rbm.visible_conv isa RBMs.Potts
        RBMs.zerosum!(rbm.visible_conv)
        RBMs.zerosum!(rbm.w_conv; dims = 1)
    end
    return rbm
end

function RBMs.zerosum_hidden!(rbm::DenseConvRBM)
    if hidden(rbm) isa RBMs.Potts
        RBMs.zerosum!(hidden(rbm))
        RBMs.zerosum!(rbm.w_dense; dims = ndims(rbm.w_dense) - ndims(hidden(rbm)) + 1)
        RBMs.zerosum!(rbm.w_conv; dims = ndims(rbm.w_conv) - ndims(hidden(rbm)) + 1)
    end
    return rbm
end


# see web page for full listing
# https://code.google.com/p/yahoo-finance-managed/wiki/enumQuoteProperty

AfterHoursChangeRealtime = "c8"
AnnualizedGain = "g3"
Ask = "a0"
AskRealtime = "b2"
AskSize = "a5"
AverageDailyVolume = "a2"
Bid = "b0"
BidRealtime = "b3"
BidSize = "b6"
BookValuePerShare = "b4"
Change = "c1"
Change_ChangeInPercent = "c0"
ChangeFromFiftydayMovingAverage = "m7"
ChangeFromTwoHundreddayMovingAverage = "m5"
ChangeFromYearHigh = "k4"
ChangeFromYearLow = "j5"
ChangeInPercent = "p2"
ChangeInPercentRealtime = "k2"
ChangeRealtime = "c6"
Commission = "c3"
Currency = "c4"
DaysHigh = "h0"
DaysLow = "g0"
DaysRange = "m0"
DaysRangeRealtime = "m2"
DaysValueChange = "w1"
DaysValueChangeRealtime = "w4"
DividendPayDate = "r1"
TrailingAnnualDividendYield = "d0"
TrailingAnnualDividendYieldInPercent = "y0"
DilutedEPS = "e0"
EBITDA = "j4"
EPSEstimateCurrentYear = "e7"
EPSEstimateNextQuarter = "e9"
EPSEstimateNextYear = "e8"
ExDividendDate = "q0"
FiftydayMovingAverage = "m3"
SharesFloat = "f6"
HighLimit = "l2"
HoldingsGain = "g4"
HoldingsGainPercent = "g1"
HoldingsGainPercentRealtime = "g5"
HoldingsGainRealtime = "g6"
HoldingsValue = "v1"
HoldingsValueRealtime = "v7"
LastTradeDate = "d1"
LastTradePriceOnly = "l1"
LastTradeRealtimeWithTime = "k1"
LastTradeSize = "k3"
LastTradeTime = "t1"
LastTradeWithTime = "l0"
LowLimit = "l3"
MarketCapitalization = "j1"
MarketCapRealtime = "j3"
MoreInfo = "i0"
Name = "n0"
Notes = "n4"
OneyrTargetPrice = "t8"
Open = "o0"
OrderBookRealtime = "i5"
PEGRatio = "r5"
PERatio = "r0"
PERatioRealtime = "r2"
PercentChangeFromFiftydayMovingAverage = "m8"
PercentChangeFromTwoHundreddayMovingAverage = "m6"
ChangeInPercentFromYearHigh = "k5"
PercentChangeFromYearLow = "j6"
PreviousClose = "p0"
PriceBook = "p6"
PriceEPSEstimateCurrentYear = "r6"
PriceEPSEstimateNextYear = "r7"
PricePaid = "p1"
PriceSales = "p5"
Revenue = "s6"
SharesOwned = "s1"
SharesOutstanding = "j2"
ShortRatio = "s7"
StockExchange = "x0"
Symbol = "s0"
TickerTrend = "t7"
TradeDate = "d2"
TradeLinks = "t6"
TradeLinksAdditional = "f0"
TwoHundreddayMovingAverage = "m4"
Volume = "v0"
YearHigh = "k0"
YearLow = "j0"
YearRange = "k2"


using MPIMeasurements
using HDF5

scanner = MPIScanner("HeadScanner.toml")
daq = getDAQ(scanner)

params = toDict(daq.params)

params["studyName"]="TestTobi"
params["studyDescription"]="A very cool measurement"
params["scannerOperator"]="Tobi"
params["dfStrength"]=[5e-3]  #[5e-3]
params["acqNumAverages"]=1000
#params["calibFieldToVolt"]=[1]

params["acqFFSequence"] = "HeadScanner102Triangle"
#params["acqFFValues"] = repeat( cat(1,x,reverse(x[2:end-1])),inner=5)
params["acqNumPeriodsPerFrame"]=length(params["acqFFValues"])

robot = getRobot(scanner)
S,S2 = measurementContReadAndSave(daq, robot, params, controlPhase=true, showFT=true)

h5write("systemMatrixNew.h5","/matrix",S)
h5write("phantom.h5","/matrix",S2)