module Dynesty

# Write your package code here.
using PyCall
using StatsBase

export NestedSampler,
       DynamicNestedSampler,
       runplot,
       traceplot,
       cornerpoints,
       cornerplot,
       boundplot,
       cornerbound,
       sample, resample_equal

const dynesty = PyNULL()

"""
    `dyplot`
Object that holds the plotting submodule of the `dynesty` package
"""
const dyplot = PyNULL()

"""
    `NestedSampler`
Julia interface to the `dynesty` NestedSampler class. Note that we do
not pass the loglikelihood or prior transform function here. Instead this
is passed to the `sample`` call.

# Example
```julia
# define a distribution
d = MvNormal(ones(10))

loglikelihood(x) = logpdf(d, x)

prior_transform(p) = -10.0 .+ 20.0.*p

smplr = NestedSampler(10)

res = sample(loglikelihood, prior_transform, smplr; dlogz=0.5)
```
"""
Base.@kwdef struct NestedSampler
    ndim::Int
    nlive::Int = 500
    bound::String = "multi"
    sample::String = "auto"
    periodic = nothing
    reflective = nothing
    update_interval::Float64 = 1.5
    first_update = nothing
    gradient = nothing
    walks::Int = 25
    facc::Float64 = 0.5
    slices::Int = 5
    fmove::Float64 = 0.9
    max_move::Int = 100
end

function NestedSampler(ndim::Int; kwargs...)
    return NestedSampler(;ndim=ndim, kwargs...)
end

"""
    `DyamicNestedSampler`
Julia interface to the `dynesty` DynamicNestedSampler class. Note that we do
not pass the loglikelihood or prior transform function here. Instead this
is passed to the `sample`` call.

# Example
```julia
# define a distribution
d = MvNormal(ones(10))

loglikelihood(x) = logpdf(d, x)

prior_transform(p) = -10.0 .+ 20.0.*p

smplr = NestedSampler(10)

# sample using dynamic nested sampling with 500 initial live points
res = sample(loglikelihood, prior_transform, smplr; dlogz=0.5, nlive_init=500)
```
"""
Base.@kwdef struct DynamicNestedSampler
    ndim::Int
    bound::String = "multi"
    sample::String = "auto"
    periodic = nothing
    reflective = nothing
    update_interval::Float64 = 1.5
    first_update = nothing
    gradient = nothing
    walks::Int = 25
    facc::Float64 = 0.5
    slices::Int = 5
    fmove::Float64 = 0.9
    max_move::Int = 100
end

function DynamicNestedSampler(ndim::Int; kwargs...)
    return DynamicNestedSampler(;ndim=ndim, kwargs...)
end

"""
    DynestyOutput
A objects that holds the dynesty output as well as the sampler.
This object can be passed to the plotting functions to produce the usual
dynesty diagnostic plots.

# Notes
We have also implemented a limited number of Julia's `Dictionary` interface so you
can access the output in the usual Dynesty manner

# Example
```julia
# define a distribution
d = MvNormal(ones(10))

loglikelihood(x) = logpdf(d, x)

# We only look at finite region of parameter space
prior_transform(p) = -5.0 .+ 10.0.*p

smplr = NestedSampler(10)

# sample using dynamic nested sampling with 500 initial live points
res = sample(loglikelihood, prior_transform, smplr; dlogz=0.5, nlive_init=500)

# fetch the samples
res[:samples]

# print the keys
keys(res)
```

"""
struct DynestyOutput{D,S}
    dict::D
    sampler::S
end

Base.getindex(d::DynestyOutput, key) = getindex(d.dict, key)
Base.keys(d::DynestyOutput) = keys(d.dict)
Base.get(f::Function, d::DynestyOutput, key) = get(f, d.dict, key)
Base.values(d::DynestyOutput) = values(d.dict)

"""
    `sample(loglikelihood, prior_transform, s::NestedSampler; kwargs...)`

Runs dynesty's NestedSampler algorithm with the specified loglikelihood and prior_transform.
The loglikelihood and prior_transform are functions. For the specific relevant kwargs see the
dynesty documentation at [https://dynesty.readthedocs.io/]
"""
function StatsBase.sample(loglikelihood, prior_transform, s::NestedSampler; kwargs...)
    dysampler = dynesty.NestedSampler(loglikelihood,
                                      prior_transform,
                                      s.ndim,;
                                      nlive = s.nlive,
                                      bound = s.bound,
                                      sample = s.sample,
                                      periodic = s.periodic,
                                      reflective = s.reflective,
                                      update_interval = s.update_interval,
                                      first_update = s.first_update,
                                      gradient = s.gradient,
                                      walks = s.walks,
                                      facc = s.facc,
                                      slices = s.slices,
                                      fmove = s.fmove,
                                      max_move = s.max_move
                                    )
    dysampler.run_nested(;kwargs...)
    return DynestyOutput(dysampler.results, dysampler)
end

"""
    `sample(loglikelihood, prior_transform, s::DynamicNestedSampler; kwargs...)`

Runs dynesty's DynamicNestedSampler algorithm with the specified loglikelihood and prior_transform.
The loglikelihood and prior_transform are functions. For the specific relevant kwargs see the
dynesty documentation at [https://dynesty.readthedocs.io/]
"""
function StatsBase.sample(loglikelihood, prior_transform, s::DynamicNestedSampler; kwargs...)
    dysampler = dynesty.DynamicNestedSampler(loglikelihood,
                                      prior_transform,
                                      s.ndim,;
                                      bound = s.bound,
                                      sample = s.sample,
                                      periodic = s.periodic,
                                      reflective = s.reflective,
                                      update_interval = s.update_interval,
                                      first_update = s.first_update,
                                      gradient = s.gradient,
                                      walks = s.walks,
                                      facc = s.facc,
                                      slices = s.slices,
                                      fmove = s.fmove,
                                      max_move = s.max_move
                                    )
    dysampler.run_nested(;kwargs...)
    return DynestyOutput(dysampler.results, dysampler)
end

"""
    resample_equal(res::DynestyOutput, nsamples::Int)
Resample the `dynesty` nested sampling run so that the samples have equal weighting.
This uses the `StatsBase` algorithm under the hood.

The results are a vector of vectors where the inner vector corresponds to the samples.
"""
function resample_equal(res::DynestyOutput, nsamples::Int)
    samples, weights = res["samples"].T, exp.(res["logwt"].T .- res["logz"][end])
    println(size(eachcol(samples)))
    sample(collect(eachcol(samples)), Weights(weights), nsamples)
end

"""
    merge(args::DynestyOutput...; print_progres=true)
Runs dynesty's merge_runs to combine multiple separate dynesty runs.
"""
function Base.merge(args::DynestyOutput...; print_progress=true)
    runs = getproperty.([args...], :dict)
    return dynesty.utils.merge_runs(runs; print_progress)
end

# internal ess estimator
function ess(res::DynestyOutput)
    weights = exp.(res["logwt"] .- res["logz"][end])
    return inv(sum(abs2, weights))
end

# Make the dynesty plots with a bit of fun metaprogramming
PLOTS = [:runplot, :traceplot, :cornerpoints, :cornerplot, :boundplot, :cornerbound]
for P in PLOTS
    @eval begin
        """
            $($P)(d::DynestyOutput; kwargs...)
        Produces the $($P) plot from the dynesty.plotting module.
        For a list of possible kwargs see the dynesty documentation at
        [dynesty.readthedocs.io]
        """
        function ($P)(d::DynestyOutput, args...; kwargs...)
            f = getproperty(dyplot, Symbol($P))
            f(d.dict, args...; kwargs...)
        end
    end
end

function __init__()
    copy!(dynesty, pyimport_conda("dynesty","dynesty"))
    copy!(dyplot, pyimport("dynesty.plotting"))
    # Define a hack to get merge to work without doing silly dict conversion
    # Not needed for new versions
    #py"""
    #import dynesty
    #def merge(x, print_progress=True):
    #    xres = [dynesty.results.Results(r) for r in x]
    #    mruns = dynesty.utils.merge_runs(xres, print_progress)
    #    return mruns
    #"""
end

end


"""
This module provides utilities to build neural networks with Flux,
along with a library of standard architectures.
"""
module FluxLib

export SimpleNet, SimpleNetHP, ResNet, ResNetHP, Gcn, GcnHP, Gin, GinHP, GraphSAGE, GraphSAGEHP

using ..AlphaZero

using CUDA
using Base: @kwdef
#using Statistics: var

import Flux
import GraphNeuralNetworks

CUDA.allowscalar(false)
array_on_gpu(::Array) = false
array_on_gpu(::CuArray) = true
array_on_gpu(arr) = error("Usupported array type: ", typeof(arr))

using Flux: relu, softmax, flatten, σ
using Flux: Chain, Dense, Conv, BatchNorm, SkipConnection, Parallel
using GraphNeuralNetworks: GCNConv, SAGEConv
import Zygote

function Flux.unbatch(g::GraphNeuralNetworks.GNNGraph, x)
  changes = g.graph_indicator[1:end-1] .!= g.graph_indicator[2:end]
  index =  [0; Array(findall(changes)); length(g.graph_indicator)]
  return [x[index[i]+1:index[i+1]] for i in 1:g.num_graphs]
end

#####
##### Flux Networks
#####

"""
    FluxNetwork <: AbstractNetwork

Abstract type for neural networks implemented using the _Flux_ framework.

The `regularized_params_` function must be overrided for all layers containing
parameters that are subject to regularization.

Provided that the above holds, `FluxNetwork` implements the full
network interface with the following exceptions:
[`Network.HyperParams`](@ref), [`Network.hyperparams`](@ref),
[`Network.forward`](@ref) and [`Network.on_gpu`](@ref).
"""
abstract type FluxNetwork <: AbstractNetwork end

function Base.copy(nn::Net) where Net <: FluxNetwork
  #new = Net(Network.hyperparams(nn))
  #Flux.loadparams!(new, Flux.params(nn))
  #return new
  return Base.deepcopy(nn)
end

Network.to_cpu(nn::FluxNetwork) = Flux.cpu(nn)

function Network.to_gpu(nn::FluxNetwork)
  CUDA.allowscalar(false)
  return Flux.gpu(nn)
end

function Network.set_test_mode!(nn::FluxNetwork, mode)
  Flux.testmode!(nn, mode)
end

Network.convert_input(nn::FluxNetwork, x) =
  Network.on_gpu(nn) ? Flux.gpu(x) : x

Network.convert_output(nn::FluxNetwork, x) = Flux.cpu(x)

Network.params(nn::FluxNetwork) = [Flux.params(nn)]

Network.set_params!(nn::FluxNetwork, weights) = Flux.loadparams!(nn, weights[1])

# This should be included in Flux
function lossgrads(f, args...)
  val, back = Zygote.pullback(f, args...)
  grad = back(Zygote.sensitivity(val))
  return val, grad
end

function Network.train!(callback, nn::FluxNetwork, opt::Adam, loss, data, n)
  optimiser = Flux.ADAM(opt.lr)
  params = Flux.params(nn)
  for (i, d) in enumerate(data)
    d = Network.convert_input_tuple(nn, d)
    l, grads = lossgrads(params) do
      loss(d...)
    end
    Flux.update!(optimiser, params, grads)
    callback(i, l)
  end
end

function Network.train!(
    callback, nn::FluxNetwork, opt::CyclicNesterov, loss, data, n)
  lr = CyclicSchedule(
    opt.lr_base,
    opt.lr_high,
    opt.lr_low, n=n)
  momentum = CyclicSchedule(
    opt.momentum_high,
    opt.momentum_low,
    opt.momentum_high, n=n)
  optimiser = Flux.Nesterov(opt.lr_low, opt.momentum_high)
  params = Flux.params(nn)
  for (i, d) in enumerate(data)
    l, grads = lossgrads(params) do
      loss(d...)
    end
    Flux.update!(optimiser, params, grads)
    optimiser.eta = lr[i]
    optimiser.rho = momentum[i]
    callback(i, l)
  end
end

regularized_params_(l) = []
regularized_params_(l::GraphNeuralNetworks.GCNConv) = [l.weight, l.bias]
regularized_params_(l::GraphNeuralNetworks.SAGEConv) = [l.weight, l.bias]
regularized_params_(l::Flux.Dense) = [l.weight]
regularized_params_(l::Flux.Conv) = [l.weight]

function Network.regularized_params(net::FluxNetwork)
  return (w for l in Flux.modules(net) for w in regularized_params_(l))
end

function Network.gc(::FluxNetwork)
  GC.gc(true)
  # CUDA.reclaim()
end

#####
##### Common functions between two-head neural networks
#####

"""
    TwoHeadNetwork <: FluxNetwork

An abstract type for two-head neural networks implemented with Flux.

Subtypes are assumed to have fields
`hyper`, `gspec`, `common`, `vhead` and `phead`. Based on those, an implementation
is provided for [`Network.hyperparams`](@ref), [`Network.game_spec`](@ref),
[`Network.forward`](@ref) and [`Network.on_gpu`](@ref), leaving only
[`Network.HyperParams`](@ref) to be implemented.
"""
abstract type TwoHeadNetwork <: FluxNetwork end

abstract type TwoHeadGraphNeuralNetwork <: TwoHeadNetwork end

Network.params(nn::TwoHeadNetwork) = [Flux.params(nn.common), Flux.params(nn.vhead), Flux.params(nn.phead)]

function Network.set_params!(nn::TwoHeadNetwork, weights)
  Flux.loadparams!(nn.common,weights[1])
  Flux.loadparams!(nn.vhead, weights[2])
  Flux.loadparams!(nn.phead, weights[3])
end

function Network.forward(nn::TwoHeadGraphNeuralNetwork, state)
  c = nn.common(state, state.ndata.x)
  v = nn.vhead(state, c)
  p = nn.phead(state, c)
  p = softmax_nodes(state, p)
  p = Flux.unbatch(state, p)
  return (p, v)
end

function Network.forward(nn::TwoHeadNetwork, state)
  c = nn.common(state)
  v = nn.vhead(c)
  p = nn.phead(c)
  return (p, v)
end

# Flux.@functor does not work with abstract types
function Flux.functor(nn::Net) where Net <: TwoHeadNetwork
  children = (nn.common, nn.vhead, nn.phead)
  constructor = cs -> Net(nn.gspec, nn.hyper, cs...)
  return (children, constructor)
end

Network.hyperparams(nn::TwoHeadNetwork) = nn.hyper

Network.game_spec(nn::TwoHeadNetwork) = nn.gspec

Network.on_gpu(nn::TwoHeadNetwork) = array_on_gpu(nn.vhead[end].bias)

#####
##### Include networks library
#####

include("architectures/simplenet.jl")
include("architectures/resnet.jl")
include("architectures/gcn.jl")
include("architectures/gin.jl")
include("architectures/graphSAGE.jl")
end


import Base: convert, ==, ≈

abstract type GroundType end

const FuzzyInt = Union{Missing, UInt}
const VecOrGround = Union{GroundType, Vector}

struct Clock <: GroundType
end

ispassive(::Clock) = true

struct UInteger <: GroundType
    width::FuzzyInt
end
UInteger() = UInteger(missing)

ispassive(::UInteger) = true

struct SInteger <: GroundType
    width::FuzzyInt
end
SInteger() = SInteger(missing)

ispassive(::SInteger) = true

struct Analog <: GroundType
    width::FuzzyInt
end
Analog() = Analog(missing)

ispassive(::Analog) = true

for (x, y) in Base.product((:Clock, :UInteger, :SInteger, :Analog), (:Clock, :UInteger, :SInteger, :Analog))
    if x == y
        @eval ==(::$x, ::$y) = true
    else
        @eval ==(::$x, ::$y) = false
    end
end
≈(x::GroundType, y::GroundType) = (x == y)

ispassive(x::Vector) = all(ispassive, x)
≈(x::Vector, y::Vector) = all(x .≈ y)

struct BundleUnit
    type::F where F
    flip::Bool

    BundleUnit(type, flip) = _checkbundletype(typeof(type)) ? new(type, flip) :
        error("Cannot create BundleUnit with type == $(typeof(type)) (must be recursively terminate in a GroundType).")
end

ispassive(x::BundleUnit) = ispassive(x.type) && !x.flip
==(x::BundleUnit, y::BundleUnit) = (x.type == y.type) && (x.flip == y.flip)
==(x, ::BundleUnit) = false
==(::BundleUnit, x) = false
≈(x::BundleUnit, y::BundleUnit) = (x == y)

const Bundle = Dict{Symbol, BundleUnit}

_checkbundletype(::Type{T}) where T <: GroundType = true
_checkbundletype(::Type{BundleUnit}) = true
_checkbundletype(::Type{Bundle}) = true
_checkbundletype(::Type{Vector{T}}) where T = _checkbundletype(T)
_checkbundletype(x) = false

ispassive(x::Bundle) = all(ispassive, values(x))
==(x::Bundle, y::Bundle) = all(z -> (z[1].first == z[2].first) && (z[1].second == z[2].second), zip(x, y))
==(x, ::Bundle) = false
==(::Bundle, x) = false
≈(x::Bundle, y::Bundle) = !any(k -> !(x[k] ≈ y[k]), intersect(keys(x), keys(y)))

const FIRType = Union{VecOrGround, Bundle}
convert(::Type{BundleUnit}, x::FIRType) = BundleUnit(x, false)
convert(::Type{BundleUnit}, x::Tuple{T, Bool}) where T <: FIRType = BundleUnit(x...)


"""
    slope_array(w; reduction = true)

Generate array to store spatial slopes of solutions
"""
slope_array(w::Number; kwargs...) = deepcopy(w)

function slope_array(w::AbstractArray; reduction = true)
    nd = ndims(w)
    ids = []
    for i = 1:nd
        push!(ids, [axes(w, i) |> first, axes(w, i) |> last])
    end

    sw = ifelse(
        reduction == true,
        cat(zero(w), zero(w), dims = ndims(w) + 1),
        cat(zero(w), zero(w), zero(w), dims = ndims(w) + 1),
    )

    if w isa MArray
        sw = static_array(sw)
    end

    if w isa OffsetArray && typeof(w).types[1] <: MArray
        sw = static_array(sw)
        if ndims(sw) == 2
            sw = OffsetArray(sw, ids[1][1]:ids[1][2], axes(sw)[end])
        elseif ndims(sw) == 3
            sw = OffsetArray(sw, ids[1][1]:ids[1][2], ids[2][1]:ids[2][2], axes(sw)[end])
        elseif ndims(sw) == 4
            sw = OffsetArray(
                sw,
                ids[1][1]:ids[1][2],
                ids[2][1]:ids[2][2],
                ids[3][1]:ids[3][2],
                axes(sw)[end],
            )
        end
    end

    return sw
end


"""
    setinvalid_qualityflag!(df;
      vars=["LE","H","NEE","Tair","VPD","wind"],
      qc_suffix="_qc",
      good_quality_threshold = 1.0,
      missing_qc_as_bad = true,
      setvalmissing = true,
    )

Set records with quality flags indicating problems to false in :valid column.

# Arguments
- df: DataFrame with column :GPP
optional
- `vars=["LE","H","NEE","Tair","VPD","wind"]`: columns to theck for quality
- `qc_suffix="_qc"`: naming of the corresponding quality-flag column
- `good_quality_threshold = 1.0`: threshold in quality flag up to which
  data is considered good quality
- `missing_qc_as_bad = true`: set to false to not mark records with missing
  quality flag as invalid
- `setvalmissing = true`: set to false to prevent replacing values in value column
  corresponding to problematic quality flag to missing.

# Value
df with modified :valid and value columns.

# Example
```jldoctest; output = false
using DataFrames
df = DataFrame(
  NEE = 1:3, NEE_qc = [1,1,2],
  GPP = 10:10:30, GPP_qc = [1,missing,1])
setinvalid_qualityflag!(df; vars = ["NEE", "GPP"])
df.valid == [true, false, false]
ismissing(df.GPP[2]) && ismissing(df.NEE[3])
# output
true
```
"""
function setinvalid_qualityflag!(df; setvalmissing = true, kwargs...)
  if setvalmissing
    set_badquality_missing!(df; kwargs...)
  end
  _setinvalid_qualityflag!(df; kwargs...)
end

function set_badquality_missing!(df;
  vars=["LE","H","NEE","Tair","VPD","wind"],
  qc_suffix="_qc",
  good_quality_threshold = 1.0,
  missing_qc_as_bad = true
  )
  fqc(x,xqc) = if missing_qc_as_bad
     ifelse(!ismissing(xqc) && xqc <= good_quality_threshold, x, missing)
  else
    ifelse(ismissing(xqc) || xqc <= good_quality_threshold, x, missing)
  end
  tmp = map(vars) do var
      qcvar = var * qc_suffix
      [var, qcvar] => ByRow(fqc) => var
  end
  transform!(df, tmp...)
end

function _setinvalid_qualityflag!(df;
  vars=["LE","H","NEE","Tair","VPD","wind"],
  qc_suffix="_qc",
  good_quality_threshold = 1.0,
  missing_qc_as_bad=true
  )
  vars_qc = vars .* qc_suffix
  fvalid_var = missing_qc_as_bad ? valid_nonmissingvar_ : valid_missingvar_
  function fsel(valid, x...)
     nvar = length(x) ÷ 2
     vs = x[1:nvar]
     vs_qc = x[(nvar+1):end]
     #valid = repeat(@MVector([true]), length(vs[1]))
     for i = 1:nvar
         # TODO change when 1.7 is out which support .&&: @. valid = valid && fvalid_var(vs[i], vs_qc[i], good_quality_threshold)
         @. valid = valid & fvalid_var(vs[i], vs_qc[i], good_quality_threshold)
     end
     valid
  end
  if !hasproperty(df, :valid) df[!,:valid] .= true; end
  select!(df, :, Cols(vcat("valid", vars, vars_qc)) => fsel => :valid)
end

function valid_nonmissingvar_(x, xqc, good_quality_threshold)
  !ismissing(x) && !ismissing(xqc) && xqc <= good_quality_threshold
end
function valid_missingvar_(x, xqc, good_quality_threshold)
 !ismissing(x) && (ismissing(xqc) || xqc <= good_quality_threshold)
end

"""
    setinvalid_range!(df, var_ranges...; setvalmissing = true, ...)

Set records with values outside specified ranges to false in :valid column.

If their is no limit towards small or
large values, supply `-Inf` or `Inf` as the minimum or maximum respectively.
If there were false values in the :value column before, they are kept false.
In addition, set values outside ranges to missing.

# Arguments
- `df`: DataFrame with column :GPP
- `var_ranges`: Pair `Varname_symbol => (min,max)`: closed valid interval for
  respective column
optional
- setvalmissing: set to false to prevent replacing values in value column outside ranges to missing.

# Value
df with modified :valid and value columns.
```jldoctest; output = false
using DataFrames
df = DataFrame(NEE = [missing, 2,3], GPP = 10:10:30)
setinvalid_range!(df, :NEE => (-2.0,4.0), :GPP => (8.0,28.0))
df.valid == [false, true, false]
ismissing(df.GPP[3])
# output
true
```
"""
function setinvalid_range!(df, var_ranges::Vararg{Pair,N}; setvalmissing = true, kwargs...) where N
  if setvalmissing
    set_badrange_missing!(df, var_ranges...; kwargs...)
  end
  _setinvalid_range!(df, var_ranges...; kwargs...)
end

function set_badrange_missing!(df, var_ranges::Vararg{Pair,N}) where N
  tmp = map(var_ranges) do p
      var, (min, max) = p
      var => ByRow(x -> (ifelse(!ismissing(x) && min <= x <= max, x, missing))) => var
  end
  transform!(df, tmp...)
end

function _setinvalid_range!(df, var_ranges::Vararg{Pair,N}) where N
  function fval(valid, x...)
    for (p,xj) in zip(var_ranges, x)
      var, (min, max) = p
      #TODO simplify once 1.7 with .&& is out: @. valid = valid && !ismissing(xi) && (min <= xi <= max)
      for i in eachindex(xj)
        valid[i] = valid[i] && !ismissing(xj[i]) && (min <= xj[i] <= max)
      end
    end
    valid
  end
  vars = map(x -> x.first, var_ranges)
  if !hasproperty(df, :valid) df[!,:valid] .= true; end
  select!(df, :, SA[:valid, vars...] => fval => :valid)
end

"""
    setinvalid_nongrowingseason!(df, tGPP; kwargs...)

Set non-growseason to false in :valid column.

# Arguments
- `df`: DataFrame with columns `:GPP` and `:datetime`
- `tGPP`: scalar threshold of daily GPP (see [`get_growingseason`](@ref))
optional:
- `update_GPPd`: set to true additionally update `:GPPd_smoothed` column to
  results from [`get_growingseason`](@ref)
- and others passed to [`get_growingseason`](@ref)

# Value
df with modified columns :valid and if  `:GPPd_smoothed`,
where all non-growing season records are set to false.
"""
function setinvalid_nongrowingseason!(df, tGPP; update_GPPd_smoothed = false, kwargs...)
  if !hasproperty(df, :valid) df[!,:valid] .= true; end
  # non-copying dataframe where we can add grouping column __day
  dft = transform(df, :datetime => ByRow(Date) => :__day, copycols = false)
  function mean_nonmissing(x)
      mx = mean(skipmissing(x))
      # if there is no non-missing record in a group, mean returns nan
      ifelse(isnan(mx), missing, mx)
  end
  dfday = combine(groupby(dft, :__day), :GPP => mean_nonmissing => :GPPd, nrow)
  transform!(dfday,
    :GPPd => (x -> get_growingseason(x, tGPP; kwargs...)) => SA[:valid, :GPPd_smoothed])
  # rep(dfd.valid, each = df.nrow)
  #TODO simplify once 1.7 is out: df[!,:valid] .= df.valid .&& vcat(fill.(dfday.valid, dfday.nrow)...)
  df[!,:valid] .= df.valid .& vcat(fill.(dfday.valid, dfday.nrow)...)
  if update_GPPd_smoothed
      df[!,:GPPd_smoothed] .= vcat(fill.(dfday.GPPd_smoothed, dfday.nrow)...)
  end
  df
end

"""
    get_growingseason(GPPd, tGPP; ws=15, min_int=5, warngap=true)

Filters annual time series for growing season based on smoothed daily GPP data.

# Arguments
- `GPPd`:    daily GPP (any unit)
- `tGPP`:    GPP threshold (fraction of 95th percentile of the GPP time series).
               Takes values between 0 and 1.
optional
- `ws`:      window size used for GPP time series smoothing
- `min_int`: minimum time interval in days for a given state of growing season
- `warngap`: set to false to suppress warning on too few non-missing data

# Details
The basic idea behind the growing season filter is that vegetation is
considered to be active when its carbon uptake (GPP) is above a specified
threshold, which is defined relative to the peak GPP (95th percentile)
observed in the year.
The GPP-threshold is calculated as:

``GPP_{threshold} = quantile(GPPd,0.95)*tGPP``

GPPd time series are smoothed with a moving average to avoid fluctuations
in the delineation of the growing season. The window size defaults to 15
days, but depending on the ecosystem, other values can be appropriate.

The argument `min_int` serves to avoid short fluctuations in the
status growing season vs. no growing season by defining a minimum length
of the status. If a time interval shorter than `min_int` is labeled
as growing season or non-growing season, it is changed to the status of
the neighboring values, i.e its opposite.

# Value
A NamedTuple with entries
- `is_growingseason`: a `BitVector` of the same length as the input GPPd in which `false`
  indicate no growing season (dormant season) and `true` indicate growing season.
- `GPPd_smoothed`: smoothed GPPd
"""
function get_growingseason(GPPd,tGPP;ws=15,min_int=5,warngap=true)
  nday = length(GPPd)
  if sum(.!ismissing.(GPPd)) < 0.5*nday
    error("Expected number of available GPPd data " *
    "to be at least half the total number of days ($nday). " *
    "But was only ($(sum(.!ismissing.(GPPd)))).")
  end
  GPP_threshold = quantile(skipmissing(GPPd), 0.95)*tGPP
  #@show GPP_threshold
  # smooth GPP
  #fromR: GPPd_smoothed = filter(GPPd,method="convolution",filter=rep(1/ws,ws))
  GPPd_smoothed = moving_average(GPPd, ws)
  # set values at the beginning and end of the time series to the mean of original values
  wsd = floor(Int, ws/2)
  GPPd_smoothed[1:wsd] .= mean(skipmissing(GPPd[1:(2*wsd)]))
  GPPd_smoothed[(length(GPPd)-(wsd-1)):length(GPPd)] .=
    mean(skipmissing(GPPd[(length(GPPd)-(2*wsd-1)):length(GPPd)]))
  # check for occurence of missing values and set to mean of the values surrounding them
  imissing = findall(ismissing, GPPd_smoothed)
  if length(imissing) > 0
    warngap && length(imissing) > 10 &&
      @warn "Attention, there is a gap in 'GPPd' of length n = $(length(imissing))"
    #TODO check and correct for several gaps, see Impute package
    replace_val = mean(skipmissing(GPPd_smoothed[max(1,imissing[1] - 4):min(
      (imissing[length(imissing)] + 4),length(GPPd_smoothed))]))
    GPPd_smoothed = coalesce.(GPPd_smoothed, replace_val)
  end
  # filter daily GPP
  growseas = GPPd_smoothed .>= GPP_threshold
  # change short intervals to the surrounding values to avoid 'wrong' fluctuations
  # switch shortest interval successively and recompute interval lengths
  intervals = rle(growseas)[2]
  imin = argmin(intervals)
  while intervals[imin] < min_int
      end_int = cumsum(intervals)
      start_int = vcat(0, end_int[1:end-1]) .+ 1
      growseas[start_int[imin]:end_int[imin]] .= !growseas[start_int[imin]]
      intervals = rle(growseas)[2]
      imin = argmin(intervals)
  end
  (is_growingseason = growseas, GPPd_smoothed = GPPd_smoothed)
end

"""
    setinvalid_afterprecip!(df; min_precip = 0.01, hours_after = 24.0)

Set records after precipitation to false in `:valid` column.

# Arguments
- df: DataFrame with columns `:datetime` and `:precip` sorted by `:datetime`
  in increasing order.
optional:
- `min_precip` (in mm per timestep): minimum precip to be considered effective
  precipitation.
- `hours_after`: time after the precipitation event to be considered invalid

# Value
`df` with modified column `:valid`.
"""
function setinvalid_afterprecip!(df; min_precip = 0.01, hours_after = 24.0)
  dfo = @pipe copy(df) |>
    select!(_, :datetime, :precip) |>
    subset!(_, :precip => ByRow(>=(min_precip)); skipmissing = true) |>
    select!(_,
      :datetime => :p_start,
      :datetime => (x -> x .+ frac_hour(Second, hours_after)) => :p_end)
  dfno = get_nonoverlapping_periods(dfo)
  #
  if !hasproperty(df, :valid) df[!,:valid] .= true; end
  #ip = 1
  # since date.time is ordered we need to search for the next index only from last position
  i_start = i_end = 1
  for ip in 1:nrow(dfno)
      p_start, p_end = dfno.p_start[ip], dfno.p_end[ip]
      i_start = i_end - 1 + findfirst(>=(p_start), df.datetime[i_end:end])
      i_end = i_start - 1 + findlast(<=(p_end), df.datetime[i_start:end])
      df.valid[i_start:i_end] .= false
  end
  df
end


module Problox

using PyCall
using Conda
using MacroTools
using MacroTools: @capture

# Install ProbLog.
if PyCall.conda
	Conda.add("pip")
	pip = joinpath(Conda.BINDIR, "pip")
	run(`$pip install problog`)
else
	try
        pyimport("problog")
        pyimport("problog.program")
        pyimport("problog.logic")
        pyimport("problog.tasks")
        pyimport("problog.tasks.sample")
	catch ee
		typeof(ee) <: PyCall.PyError || rethrow(ee)
		warn("""
Python Dependencies not installed
Please either:
 - Rebuild PyCall to use Conda, by running in the julia REPL:
    - `ENV[PYTHON]=""; Pkg.build("PyCall"); Pkg.build("Problox")`
 - Or install the depencences, eg by running pip
	- `pip install problog`
	"""
		)
	end
end

problog = pyimport("problog")
pp = pyimport("problog.program")
pl = pyimport("problog.logic")
pt = pyimport("problog.tasks")
sample = pyimport("problog.tasks.sample")

# User-defined terms.
function term(str::String)
    return pl.Term(str)
end

SimpleProgram() = pp.SimpleProgram()
Var(str::String) =  pl.Var(str)
variable(s::String) = Var(String(s))
variable(s::Symbol) = Var(String(s))
Constant(str::String) =  pl.Constant(str)
AnnotatedDisjunction(p::Array{PyObject}, q::PyObject) = pl.AnnotatedDisjunction(p, q)
AnnotatedDisjunction(p::NTuple{N, PyObject}, q::PyObject) where N = pl.AnnotatedDisjunction([p...], q)

evaluate(p::PyObject) = problog.get_evaluatable().create_from(p).evaluate()

include("dsl.jl")

# DSL.
export @problox

# Lower-level API.
export SimpleProgram, Var, Constant, AnnotatedDisjunction, term

# Convenience.
query = term("query")
export query, evaluate, variable

end # module


import Distributions: rand, logpdf, pdf, estimate

function sample_chain(chain)
    parms = (Symbol.(chain.name_map.parameters)...,)
    n = size(chain, 1)*size(chain, 3)
    idx = rand(1:n)
    vals = map(x -> chain[x][idx], parms)
    return NamedTuple{parms}(vals)
end

function _posterior_predictive(m, chain, f=x -> x)
    parms = sample_chain(chain)
    return f(m(parms))
end

"""
**posterior_predictive**

Returns posterior predictive distribution and optionally applies function to samples
    on each replication
- `model`: the data generating function of a model
- `chain`: an MCMCChains chain object
- `n_samples`: the number of samples
- `f`: a function that is applied to each sample from posterior predictive distribution
Function Signature
````julia
posterior_predictive(model, chain, n_samples::Int, f=x -> x)
````
"""
function posterior_predictive(model, chain, n_samples::Int, f=x -> x)
    return map(x -> _posterior_predictive(model, chain, f), 1:n_samples)
end

find_index(actr::AbstractACTR; check_value=true, criteria...) = find_index(actr.declarative.memory; check_value, criteria...)

"""
**find_index**

Returns the index of first chunk that matches a set of criteria
- `chunks`: an array of chunks
- `criteria`: a set of keyword arguments for slot-value pairs

Function Signature
````julia
find_index(chunks::Array{<:Chunk,1}; criteria...)
````
## Example
````julia
chunks = [Chunk(animal=:dog), Chunk(animal=cat)]
find_index(chunks; animal=:dog)
````
"""
function find_index(chunks::Array{<:Chunk,1}; check_value=true, criteria...)
    for (i,c) in enumerate(chunks)
        match(c; check_value, criteria...) ? (return i) : nothing
    end
    return -100
end

"""
**find_index**

Returns the index of first chunk that matches a set of criteria
- `actr`: ACTR object
- `funs`: a set of functions
- `criteria`: a set of keyword arguments for slot-value pairs

Function Signature
````julia
find_index(actr::ACTR, funs...; criteria...)
````
## Example
````julia
chunks = [Chunk(animal=:dog), Chunk(animal=cat)]
find_index(chunks; animal=:dog)
````
"""
find_index(actr::ACTR, funs...; check_value=true, criteria...) = find_index(actr.declarative.memory, funs...; check_value, criteria...)

"""
**find_index**

Returns the index of first chunk that matches a set of criteria
- `chunks`: an array of chunks
- `funs`: a set of functions
- `criteria`: a set of keyword arguments for slot-value pairs

Function Signature
````julia
find_index(chunks::Array{<:Chunk,1}; criteria...)
````
## Example
````julia
chunks = [Chunk(animal=:dog), Chunk(animal=cat)]
find_index(chunks; animal=:dog)
````
"""
function find_index(chunks::Array{<:Chunk,1}, funs...; check_value=true, criteria...)
    for (i,c) in enumerate(chunks)
        match(c, funs...; check_value, criteria...) ? (return i) : nothing
    end
    return -100
end

"""
**find_indices**

Returns the index of first chunk that matches a set of criteria
- `actr`: an ACTR object
- `criteria`: a set of keyword arguments for slot-value pairs

Function Signature
````julia
find_indices(actr::ACTR; criteria...)
````
## Example
````julia
chunks = [Chunk(animal=:dog), Chunk(animal=:dog), Chunk(animal=cat)]
find_indices(chunks; animal=:dog)
````
"""
find_indices(actr::ACTR; check_value = true, criteria...) = find_indices(actr.declarative.memory; check_value, criteria...)

"""
**find_indices**

Returns the index of first chunk that matches a set of criteria
- `chunks`: an array of chunks
- `criteria`: a set of keyword arguments for slot-value pairs

**Function Signature**
````julia
find_indices(actr::ACTR; criteria...)
````
**Example**
````julia
chunks = [Chunk(animal=:dog), Chunk(animal=:dog), Chunk(animal=cat)]
find_indices(chunks; animal=:dog)
````
"""
function find_indices(chunks::Array{<:Chunk,1}; check_value=true, criteria...)
    idx = Int64[]
    for (i,c) in enumerate(chunks)
        match(c; check_value, criteria...) ? push!(idx, i) : nothing
    end
    return idx
end

"""
**find_indices**

Returns the index of first chunk that matches a set of criteria
- `actr`: an ACTR object
- `funs`: a set of functions
- `criteria`: a set of keyword arguments for slot-value pairs

**Function Signature**
````julia
find_indices(actr::ACTR; criteria...)
````
**Example**
````julia
chunks = [Chunk(animal=:dog), Chunk(animal=:dog), Chunk(animal=cat)]
find_indices(chunks; animal=:dog)
````
"""
find_indices(actr::ACTR, funs...; check_value=true, criteria...) = find_indices(actr.declarative.memory, funs...; check_value, criteria...)

"""
**find_indices**

Returns the index of first chunk that matches a set of criteria
- `chunks`: an array of chunks
- `funs`: a set of functions
- `criteria`: a set of keyword arguments for slot-value pairs

**Function Signature**
````julia
find_indices(actr::ACTR; criteria...)
````
** Example**
````julia
chunks = [Chunk(animal=:dog), Chunk(animal=:dog), Chunk(animal=cat)]
find_indices(chunks; animal=:dog)
````
"""
function find_indices(chunks::Array{<:Chunk,1}, funs...; check_value=true, criteria...)
    idx = Int64[]
    for (i,c) in enumerate(chunks)
        match(c, funs...; check_value, criteria...) ? push!(idx, i) : nothing
    end
    return idx
end

"""
**get_iconic_memory**

Returns array of chunks or visual objects representing iconic memory
- `actr`: an ACTR object

**Function Signature**
````julia
get_iconic_memory(actr)
````
"""
get_iconic_memory(actr) = actr.visual_location.iconic_memory

"""
**get_visicon**

Returns array of chunks or visual objects representing all visual objects
within the simulation

**Function Signature**
````julia
get_visicon(actr)
````
"""
get_visicon(actr) = actr.visicon

"""
**get_buffer**

Returns posterior predictive distribution and optionally applies function to samples
    on each replication
- `actr`: an ACTR object
- `m`: name of module as a symbol

**Function Signature**
````julia
get_buffer(actr, m)
````

**Example**
````julia
get_buffer(actr, :imaginal)
````
"""
get_buffer(actr, m) = getfield(actr, m).buffer
"""
**get_buffer**

Returns posterior predictive distribution and optionally applies function to samples
    on each replication
- `actr`: an ACTR object
- `m`: name of module as a symbol
- `v`: new value

**Function Signature**
````julia
set_buffer!(actr, m, v)
````

Example
````julia
set_buffer!(actr, :imaginal, [chunk])
````
"""
set_buffer!(actr, m, v) = getfield(actr, m).buffer = v

"""
**get_mean_activations**

Returns a vector of mean activations i.e. activation without noise
- `actr`: an ACTR object

**Function Signature**
````julia
get_mean_activations(actr)
````
"""
get_mean_activations(actr::AbstractACTR) = get_mean_activations(actr.declarative)

"""
**get_mean_activations**

Returns a vector of mean activations i.e. activation without noise
- `declarative`: an Declarative memory object

**Function Signature**
````julia
get_mean_activations(memory)
````
"""
get_mean_activations(declarative::Declarative) = get_mean_activations(declarative.memory)

"""
    get_mean_activations(chunks::Vector{<:Chunk})

Returns a vector of mean activations i.e. activation without noise

# Arguments

- `chunks`: an array of chunks
"""
function get_mean_activations(chunks::Vector{<:Chunk})
    return map(x->x.act - x.act_noise, chunks)
end

"""
    get_name(actr::AbstractACTR)

Returns the name of an ACT-R model object.

- `actr`: an ACTR model object
"""
function get_name(actr::AbstractACTR)
    return actr.name
end

get_time(actr::AbstractACTR) = get_time(actr.scheduler)
get_time(scheduler) = scheduler.time

add_time!(actr::AbstractACTR, t) = add_time!(actr.scheduler, t)
add_time!(scheduler, t) = scheduler.time += t

reset_time!(actr::AbstractACTR) = reset_time!(actr.scheduler)
reset_time!(scheduler) = scheduler.time = 0.0

set_time!(actr::AbstractACTR, t) = set_time!(actr.scheduler, t)
set_time!(scheduler, t) = scheduler.time = t

rnd_time(μ) = rand(Uniform(μ * (2/3), μ * (4/3)))


#builds mcmc xsections, currently only used in testing
function xsectionmcmc(Θ::VolumePartsMCMC,::Type{TM}, ::Type{TV}, ::Type{T}) where {TM,TV,T}
  @unpack Xv,ts = Θ

  #=D = (ws=Xv.ws', RLws=Xv.RLws', Ã=Θ.expand[:Ã])
  xsection = dataxsection(D,Θ.ts,AbstractMatrix{T}, AbstractVector{T}, T)
  LD = G=Θ.expand[:G],σ²=Θ.expand[:σ²], ω=Θ.expand[:ω]=#

  dims = (T=size(Θ.Ã,2), K=size(Θ.Ã,1))
  xsections = Vector{NamedTuple}(undef, dims.T-1)
  for t ∈ 2:dims.T
    #count the number of stocks in each window (t and t-1)
    Nt = sum(ts.==t)

    #repeat t for each stock
    #Ainds = repeat([t],Nt)
    #LAinds = repeat([t-1], Nt)
    xinds = ts .== t
    xsections[t-1] = (
      ws = view(Θ.Xv.ws, xinds, :),
      RLws = view(Θ.Xv.RLws, xinds, :),
      v = view(Θ.Xv.v, xinds),
      #X̃ = view(Θ.X̃, xinds, :),
      #ṽ  = view(Θ.ṽ, xinds),
      Ã = view(Θ.expand[:Ã], :, xinds),
      LÃ = view(Θ.expand[:LÃ], :, xinds),
      G = view(Θ.expand[:G], :, xinds),
      σ²   = view(Θ.expand[:σ²], xinds),
      ω = view(Θ.expand[:ω], xinds),
      )
  end
  xsection::NamedTuple = (
    ws = (x->x.ws).(xsections),
    RLws = (x->x.RLws).(xsections),
    v = (x->x.v).(xsections),
    #X̃ = (x->x.X̃).(xsections),
    #ṽ = (x->x.ṽ).(xsections),

    Ã = (x->x.Ã).(xsections),
    LÃ = (x->x.LÃ).(xsections),
    G = (x->x.G).(xsections),
    σ²   = (x->x.σ²).(xsections),
    ω = (x->x.ω).(xsections),
  )

  return xsection
end

function testcapacitychains()
  hyper=(A = (a=1.1,), V=(v1=2.1, v2=2.2), M=(m1=rand(2,2), m2=3.2))
  testparams = (A = 2.3, V=rand(5), M=rand(2,5))
  cc = CapacityChains(Float64, Nburn=25, Ndraws=10, paramstemplate=testparams, hyper=hyper)
  Nparams = 16
  @assert size(cc) === (10,16)

  ###test the indices
  #prior to looping
  for i ∈ 1:10
    @assert i == cc.idx
    @assert i == cc.i
    increment!(cc)
  end

  #after the first loop of the buffer
  for i ∈ 1:15
    @assert (i-1)%10 + 1 == cc.idx
    @assert i + 10 == cc.i
    increment!(cc)
  end

  #after the first reset
  @assert cc.idx == 1
  @assert cc.i == 26
  @assert cc.previdx ==10

  ###test the state vector
  Ψ::CCState = CCState(cc)
  v = rand(Nparams)
  Ψ .= v
  @assert all(v .== Ψ)
  updateandincrement!(cc, Ψ)
  @assert all(cc.value[cc.previdx,:,1] .== Ψ)
  @assert cc.i == 27
  @assert cc.idx == 2
  @assert cc.previdx == 1
  for (n,p) ∈ pairs(Ψ.params) #test each params
    @assert all(reshape(cc.params[n][cc.previdx,:,1],size(p)) .== p)

    #=(all(reshape(cc.params[n][cc.previdx, :, 1],size(p)) .== p)) || error(
      "mismatch!!! cc.idx: $(cc.idx), previdx: $(cc.previdx)
      reshape(cc.params[$n][$(cc.previdx), :, 1],size(p)): $(reshape(cc.params[n][:,cc.previdx,1],size(p)))
      p: $(p)")=#
  end

  #test the params individually
  Ψ.A[] = testparams.A
  Ψ.V .= testparams.V
  Ψ.M .= testparams.M
  @assert all([(v->[v...;]).(values(testparams))...;] .== Ψ)
  updateandincrement!(cc,Ψ)
  @assert (cc.i, cc.idx, cc.previdx) === (28,3,2)
  for (n,p) ∈ pairs(Ψ.params) #test each params
    if n ≡ :A
      @assert reshape(cc.params[:A][cc.previdx,:,1],size(p))[] === testparams.A
    elseif n ≡ :V
      @assert all(reshape(cc.params[:V][cc.previdx,:,1],size(p)) .=== testparams.V)
    elseif n ≡ :M
      @assert all(reshape(cc.params[:M][cc.previdx,:,1],size(p)) .=== testparams.M)
    else
      @assert false
    end
  end

end

function testA₀(Θ::AbstractVolumePartsMCMC{TM,TV,T},
  Ntest::Int=PARAM[:mcmcntest]) where {TM,TV,T}

  @unpack Ψ, Xv, xsection, Ã = deepcopy(Θ)
  @unpack A₀, G, ω, σ² = Ψ
  #@info "hyper keys=$(keys(Θ.cc.hyper))"
  @unpack μA₀, VA₀inv = Θ.cc.hyper[:A₀]

  dims = (K=size(Ã, 1), T=size(Ã, 2), )

  A₀ver = Matrix(undef, Ntest, dims.K)
  A₀test = similar(A₀ver)

  function verA₀!()
    Ã .= hcat(ones(T, dims.K), cumprod(G, dims=2))

    #store the cross-sectional matrices
    #Λₜ = Vector{TM}(undef, dims.T-1)
    #Āₜ = Vector{TV}(undef, dims.T-1)
    VA₀invmat = Matrix(VA₀inv, dims.K, dims.K)
    μA₀vec = fill(μA₀, dims.K)
    Λₜ = [Matrix{T}(undef, dims.K, dims.K) for t ∈ 1:(dims.T-1)]
    ΛĀₜ = [Vector{T}(undef, dims.K) for t ∈ 1:(dims.T-1)]
    τₜ = 1 ./ σ²
    Threads.@threads for t ∈ 1:(dims.T-1)
      X = abs.(xsection[:ws,t] .* xsection[:Ã,t]' .- xsection[:RLws,t] .* xsection[:LÃ,t]')
      XtX = X' * X
      ṽ = xsection[:v,t] .- xsection[:ω,t]
      Â = cholesky(XtX)\(X'*ṽ)
      ΛĀₜ[t] .= τₜ[t] .* (XtX*Â .+ VA₀inv*μA₀vec)
      Λₜ[t] .= τₜ[t] .* (XtX .+ VA₀invmat)

      #=if !(Λₜ[t] ≈ Λₜ[t]') #fixes some roundoff error
        println("XtX: ")
        printmln(XtX)
        println("Λₜ[t]: ")
        printmln(Λₜ[t])
        error("$t: something is wrong with Λₜ[t]: !(Λₜ[t] ≈ Λₜ[t]')")
      end=#
    end

    ΛĀ = sum(ΛĀₜ)
    Λ = sum(Λₜ)

    Σ = Λ\I

    if Σ ≈ Σ' #fixes some roundoff error
      Σ .= (Σ .+ Σ') ./ 2
    else
      printmln(Σ)
      error("something is wrong with Σ: !(Σ ≈ Σ')")
    end

    Ā = Σ*ΛĀ
    #@info sum(abs.(Σ .- Σ'))
    d = MvNormal(Ā, Σ)
    return rand(d)

    #=for k ∈ shuffle(1:dims.K)
      notk = k .≠ 1:dims.K
      Σkkinv = cholesky(Σ[notk,notk])\I

      μk = Ā[k] + Σ[k:k, notk] * Σkkinv * (A₀[notk] .- Ā[notk])=#

  end

  #verA₀!()
  #@assert false

  if PARAM[:mcmcbenchtest]
    print("time updateA₀!: ")
    @btime updateA₀!($Θ)
    print("time verA₀!: ")
    @btime $verA₀!()
  end

  for i ∈ 1:Ntest
    A₀test[i,:] .= updateA₀!(Θ)
    A₀ver[i,:] .= verA₀!()
  end

  #halfNtest = Ntest ÷ 2
  #A₀vermean = mean(A₀ver[halfNtest:end, :], dims=1)
  #A₀testmean = mean(A₀test[halfNtest:end, :], dims=1)

  A₀vermean = mean(A₀ver, dims=1)
  A₀testmean = mean(A₀test, dims=1)

  Δ = abs.(A₀vermean .- A₀testmean)
  @info "
    A₀vermean: $A₀vermean
    A₀testmean: $A₀testmean
    Δ: $Δ"

end

function testposteriors(Θ::AbstractVolumePartsMCMC)
  testA₀(Θ)
end

#we use a lot of views in VolumeParts- this tests them
function testmcmcexpand(Θ::AbstractVolumePartsMCMC{TM,TV,T}, xsectionver) where {TM, TV, T}
  dims = (K=size(Θ.Ã,1), T=size(Θ.Ã,2))
  @unpack ts=Θ

  #these sort the views as though they were first divided into xsections
  function sorttvert(M::AbstractMatrix, ts)
    allts = sort!(unique(ts))
    reduce(vcat, (t->view(M, ts .== t, :)).(allts))
  end
  function sorttvert(M::AbstractVector, ts)
    allts = sort!(unique(ts))
    reduce(vcat, (t->view(M, ts .== t)).(allts))
  end

  function sortthoriz(M::AbstractMatrix, ts)
    tst = ts'
    allts = sort!(unique(ts))
    reduce(hcat, (t->view(M, :, vec(tst .== t))).(allts))
  end

  #concatenate the cross-sections in xsection to verify the Θ.expand views
  vdata = [:ws, :RLws, :v, :σ², :ω, ]
  hdata = [:Ã, :LÃ, :G, ]
  for s ∈ [vdata; hdata]
    local catfunc
    local sortfunc
    if s ∈ vdata #check if we arranged the data vertically or horizontally
      catfunc = vcat
      sortfunc = sorttvert
    elseif s ∈ hdata
      catfunc = hcat
      sortfunc = sortthoriz
    else
      @assert false
    end

    #check the complete expanded view
    vermatrix = reduce(catfunc, xsectionver[s])
    testmatrix = sortfunc(Θ.expand[s], ts)
    testxmatrix = reduce(catfunc, (t->Θ.xsection[s,t]).(1:(dims.T-1)))
    #@assert size(vermatrix) === size(testmatrix) === size(testxmatrix)
    all(vermatrix .== testmatrix .== testxmatrix)||error(
      "!all(vermatrix .== testmatrix .== testxmatrix)
      vermatrix[1:3,1:3]: $(vermatrix[1:3,1:3])
      testmatrix[1:3,1:3]: $(testmatrix[1:3,1:3])
      testxmatrix[1:3,1:3]: $(testxmatrix[1:3,1:3])
      "
    )

  #concatenate the cross-sections in xsection to verify the Θ.expand views (now for each t)
    for t ∈ 1:(dims.T-1)

      partts = @view ts[ts .≥ t + 1]
      vermatrix = reduce(catfunc, xsectionver[s][t:(dims.T-1)])
      testmatrix = sortfunc(Θ.expand[s,t], partts)
      testxmatrix = reduce(catfunc, (t->Θ.xsection[s,t]).(t:(dims.T-1)))
      @assert size(vermatrix) === size(testmatrix) === size(testxmatrix)
      @assert all(vermatrix .== testmatrix .== testxmatrix)

      @assert all(Θ.xsection[s,t] .== xsectionver[s][t])
    end
  end
return nothing
end

function testmcmc(panel::AbstractDataFrame,  zs::ZSpec,
    ::Type{T}=PARAM[:mcmctype],
    ::Type{TV}=PARAM[:mcmcgpu] ? CUDA.CuVector{T} : Vector{T},
    ::Type{TM}=PARAM[:mcmcgpu] ? CUDA.CuMatrix{T} : Matrix{T};
    prior::Symbol=PARAM[:mcmcprior],
    Nburn::Int=PARAM[:mcmcnburn],
    Ndraws::Int=PARAM[:mcmcndraws]
    ) where {
      T<:Real, TV<:AbstractVector{T}, TM<:AbstractMatrix{T}}

  testcapacitychains()

  @unpack Zms = zs
  Fw::Vector{Symbol} = Fweights(Zms, :Fw)
  FRLw::Vector{Symbol} = Fweights(Zms, :FRLw)

  #form the data into a more usable format and expand the arrays
  p = modelparts(panel, zs, T, TV, TM,
    weightdata = (Fw=Fw, FRLw=FRLw),
    Fvol=Zms.Fvol)

  dims = p.dims
  ws::TM = p.dat[:Fw]
  RLws::TM = p.dat[:FRLw]
  v::TV = p.dat[:Fvol]
  ts::Vector{Int} = p.ts
  Ascale::Vector{T} = p.Ascale
  Nexpanded::Int = length(ts)
  #this creates the factored objects- which are accurate given that initially G=1
  absΔ(wsᵢ::CuArray,RLwsᵢ)=CUDA.abs(wsᵢ-RLwsᵢ)
  absΔ(wsᵢ,RLwsᵢ)=abs(wsᵢ-RLwsᵢ)

  #now test the level versions
  Xv = XYIter(ws ,RLws, v, ts, Ascale)
  hyper=PARAM[prior]
  initial=(
    A₀ = TV(rand(T, dims.K)),
    G = TM(ones(T, dims.K, dims.T-1)),
    σ² = rand(T, dims.T-1),
    ω = rand(T, dims.T-1)
    )
  cc = CapacityChains(T, Nburn=Nburn, Ndraws=Ndraws, paramstemplate=initial, hyper=hyper)

  Θ = VolumePartsMCMC(TM, TV, T; initial, cc, Xv, ts, )

  xsection = xsectionmcmc(Θ, TM, TV, T)

  testmcmcexpand(Θ,xsection)

  testposteriors(Θ)

#=  testG(Θ, Xvlevel, _Xv, xsectionlevel)
  @assert Θ.G |> Matrix! ≈ Θ.A[:,2:end] ./ Θ.A[:,1:(end-1)] |> Matrix!

  #WARNING WARNING WARNING uncomment below!!!
  testprediction(Θ, Xvlevel, deepcopy(_Xv))
  testpredictionzygote(deepcopy(Θ),Xvlevel,deepcopy(_Xv))
  testiterloss(deepcopy(Θ),Xvlevel,deepcopy(_Xv))

  if !(TM<:CuArray)

    Xvdemean::XYIterDemean = XYIterDemean(ws ,RLws, v, ts, Ascale)
    Ddemean = (ws=Xvdemean.ws, RLws=Xvdemean.RLws, ṽ = Xvdemean.ṽ,)
    xsectiondemean = dataxsection(Ddemean,Θ.ts,TM, TV, T)
    validatexsection(Θ, Xvdemean, _Xv, p, xsectiondemean)

    testG(Θ, Xvdemean, _Xv, xsectiondemean)
    @assert Θ.G |> Matrix! ≈ Θ.A[:,2:end] ./ Θ.A[:,1:(end-1)] |> Matrix!

    testXv(Θ, Xvlevel, Xvdemean, _Xv)

    testprediction(Θ, Xvdemean, deepcopy(_Xv))
    testpredictionzygote(deepcopy(Θ),Xvdemean,deepcopy(_Xv))
    testiterloss(deepcopy(Θ),Xvdemean,deepcopy(_Xv))
  else
    #create cpu versions
    pcpu = modelparts(panel, zs,T, Vector{T}, Matrix{T},
      weightdata = (Fw=Fw, FRLw=FRLw),
      Fvol=Zms.Fvol)

    wscpu::Matrix{T} = pcpu.dat[:Fw]
    RLwscpu::Matrix{T} = pcpu.dat[:FRLw]
    vcpu::Vector{T} = pcpu.dat[:Fvol]

    Xvlevelcpu::XYIterLevel = XYIterLevel(wscpu ,RLwscpu, vcpu, ts, Ascale)
    Θcpu = VolumePartsIter(dims.T, dims.K, ts, T, Matrix{T}, Vector{T})
    Θcpu.G .= Θ.G |> Matrix
    Θcpu.A .= Θ.A |> Matrix
    _Xvcpu = XYIterAlloc(Θ, Xvlevelcpu)

    testcudazygote(deepcopy(Θ),Xvlevel,deepcopy(_Xv), Θcpu, Xvlevelcpu,  _Xvcpu)
  end=#

  return nothing
end


#-------------------------------------------------------------------------------
# File: for-test.jl
# Description: This file contains bad complexity functions relatives to moves.
# Date: November 4, 2019
# Author: Jonathan Fontaine, Killian Fretaud, Rémi Garcia,
#         Boualem Lamraoui, Benoît Le Badezet, Benoit Loger
#-------------------------------------------------------------------------------

"""
    TEST_move_exchange!(solution::Solution, car_pos_a::Int, car_pos_b::Int, instance::Instance)

Bad complexity version for function move_exchange.
"""
function TEST_move_exchange!(solution::Solution, car_pos_a::Int, car_pos_b::Int, instance::Instance)
    if car_pos_a > car_pos_b
        return TEST_move_exchange!(solution, car_pos_b, car_pos_a, instance)
    end
    if car_pos_a == car_pos_b
        return solution
    end
    solution.sequence[car_pos_a], solution.sequence[car_pos_b] = solution.sequence[car_pos_b], solution.sequence[car_pos_a]
    update_matrices!(solution, instance)

    return solution
end

"""
    TEST_cost_move_exchange(solution::Solution, car_pos_a::Int, car_pos_b::Int,
                       instance::Instance, objective::Int)

Bad complexity version for function cost_move_exchange.
"""
function TEST_cost_move_exchange(solution::Solution, car_pos_a::Int, car_pos_b::Int,
                            instance::Instance, objective::Int)
    s = deepcopy(solution)
    s = TEST_move_exchange!(s, car_pos_a, car_pos_b, instance)

    return cost(s, instance, objective) - cost(solution, instance, objective)
end

"""
    TEST_move_insertion!(solution::Solution, old_index::Int, new_index::Int, instance::Instance)

Bad complexity version for function move_insertion.
"""
function TEST_move_insertion!(solution::Solution, old_index::Int, new_index::Int, instance::Instance)

    car_inserted = solution.sequence[old_index]
    if old_index < new_index
        for car_moved_pos in old_index:(new_index-1)
            solution.sequence[car_moved_pos] = solution.sequence[car_moved_pos+1]
        end
        solution.sequence[new_index] = car_inserted
    end
    if old_index > new_index
        for car_moved_pos in old_index:-1:(new_index+1)
            solution.sequence[car_moved_pos] = solution.sequence[car_moved_pos-1]
        end
        solution.sequence[new_index] = car_inserted
    end
    update_matrices!(solution, instance)

    return solution
end

"""
    TEST_cost_move_insertion(solution::Solution, car_pos_a::Int,
                        instance::Instance, objective::Int)

Bad complexity version for function cost_move_insertion.
"""
function TEST_cost_move_insertion(solution::Solution, car_pos_a::Int,
                             instance::Instance, objective::Int)

    cost_on_objective = zeros(Int, instance.nb_cars, 3)
    for i in (instance.nb_late_prec_day+1):instance.nb_cars
        s = deepcopy(solution)
        s = TEST_move_insertion!(s, car_pos_a, i, instance)
        c = cost(s, instance, objective) - cost(solution, instance, objective)
        cost_on_objective[i, :] = c
    end

    return cost_on_objective
end


# create a new "build_series_args" which converts all inputs into xs = Any[xitems], ys = Any[yitems].
# Special handling for: no args, xmin/xmax, parametric, dataframes
# Then once inputs have been converted, build the series args, map functions, etc.
# This should cut down on boilerplate code and allow more focused dispatch on type
# note: returns meta information... mainly for use with automatic labeling from DataFrames for now

const FuncOrFuncs = Union{Function, AVec{Function}}

all3D(d::KW) = trueOrAllTrue(st -> st in (:contour, :contourf, :heatmap, :surface, :wireframe, :contour3d, :image), get(d, :seriestype, :none))

# missing
convertToAnyVector(v::Void, d::KW) = Any[nothing], nothing

# fixed number of blank series
convertToAnyVector(n::Integer, d::KW) = Any[zeros(0) for i in 1:n], nothing

# numeric vector
convertToAnyVector(v::AVec{T}, d::KW) where {T<:Number} = Any[v], nothing

# string vector
convertToAnyVector(v::AVec{T}, d::KW) where {T<:AbstractString} = Any[v], nothing

function convertToAnyVector(v::AMat, d::KW)
    if all3D(d)
        Any[Surface(v)]
    else
        Any[v[:,i] for i in 1:size(v,2)]
    end, nothing
end

# function
convertToAnyVector(f::Function, d::KW) = Any[f], nothing

# surface
convertToAnyVector(s::Surface, d::KW) = Any[s], nothing

# # vector of OHLC
# convertToAnyVector(v::AVec{OHLC}, d::KW) = Any[v], nothing

# dates
convertToAnyVector(dts::AVec{D}, d::KW) where {D<:Union{Date,DateTime}} = Any[dts], nothing

# list of things (maybe other vectors, functions, or something else)
function convertToAnyVector(v::AVec, d::KW)
    if all(x -> typeof(x) <: Number, v)
        # all real numbers wrap the whole vector as one item
        Any[convert(Vector{Float64}, v)], nothing
    else
        # something else... treat each element as an item
        vcat(Any[convertToAnyVector(vi, d)[1] for vi in v]...), nothing
        # Any[vi for vi in v], nothing
    end
end

convertToAnyVector(t::Tuple, d::KW) = Any[t], nothing

function convertToAnyVector(args...)
    error("In convertToAnyVector, could not handle the argument types: $(map(typeof, args[1:end-1]))")
end

# --------------------------------------------------------------------

# TODO: can we avoid the copy here?  one error that crops up is that mapping functions over the same array
#       result in that array being shared.  push!, etc will add too many items to that array

compute_x(x::Void, y::Void, z)      = 1:size(z,1)
compute_x(x::Void, y, z)            = 1:size(y,1)
compute_x(x::Function, y, z)        = map(x, y)
compute_x(x, y, z)                  = copy(x)

# compute_y(x::Void, y::Function, z)  = error()
compute_y(x::Void, y::Void, z)      = 1:size(z,2)
compute_y(x, y::Function, z)        = map(y, x)
compute_y(x, y, z)                  = copy(y)

compute_z(x, y, z::Function)        = map(z, x, y)
compute_z(x, y, z::AbstractMatrix)  = Surface(z)
compute_z(x, y, z::Void)            = nothing
compute_z(x, y, z)                  = copy(z)

nobigs(v::AVec{BigFloat}) = map(Float64, v)
nobigs(v::AVec{BigInt}) = map(Int64, v)
nobigs(v) = v

@noinline function compute_xyz(x, y, z)
    x = compute_x(x,y,z)
    y = compute_y(x,y,z)
    z = compute_z(x,y,z)
    nobigs(x), nobigs(y), nobigs(z)
end

# not allowed
compute_xyz(x::Void, y::FuncOrFuncs, z)       = error("If you want to plot the function `$y`, you need to define the x values!")
compute_xyz(x::Void, y::Void, z::FuncOrFuncs) = error("If you want to plot the function `$z`, you need to define x and y values!")
compute_xyz(x::Void, y::Void, z::Void)        = error("x/y/z are all nothing!")

# --------------------------------------------------------------------


###########
## asArr ##
###########

## asArr: try conversion to Array, implicitly dealing with NAs
##
## using DataFrames, columns may be of different type and hence are
## possibly not compatible with each other as columns of an array

function asArr{T}(da::DataArray{T, 1}, typ::Type=Any,
                  replaceNA=NA)
    ## convert DataArray to typ and replace NAs with replaceNA
    vals = convert(Array{typ, 1}, da.data)
    vals[da.na] = replaceNA
    return vals
end

function asArr(df::DataFrame, typ::Type=Any,
               replaceNA=NA)
    nObs, nAss = size(df)
    vals = Array(typ, nObs, nAss)
    ## copying into this array implicitly forces conversion to this
    ## type afterwards
    for ii=1:nAss
        if isa(df.columns[ii], DataArray) # deal with NAs
            if eltype(df.columns[ii]) == NAtype # column of NAs only
                vals[:, ii] = replaceNA
            else
                vals[:, ii] = df.columns[ii].data
                vals[df.columns[ii].na, ii] = replaceNA
            end
        else # no NAs in simple Array
            vals[:, ii] = df.columns[ii]
        end
    end
    return vals
end

function asArr(dfr::DataFrameRow, typ::Type=Any,
               replaceNA=NA)
    nObs = length(dfr)
    res = Array(typ, 1, nObs)
    for ii=1:nObs
        if isna(dfr[ii])
            res[1, ii] = replaceNA
        else
            res[1, ii] = dfr[ii]
        end
    end
    return res
end

function asArr(tn::AbstractTimedata, typ::Type=Any,
               replaceNA=NA)
    return asArr(tn.vals, typ, replaceNA)
end

## problem: columns with NAs only
## problem: NaN in boolean context will be true
## problem: DataArray{Any, 1} also exist: da.data then contains NAs

#############
## Timenum ##
#############

import Base.convert
## to Timedata: conversion upwards - always works
function convert(::Type{Timedata}, tn::Timenum)
    return Timedata(tn.vals, tn.idx)
end

## to Timematr: conversion downwards - fails for NAs
function convert(::Type{Timematr}, tn::Timenum, rmNA = false)
    if rmNA
        tm = narm(tn)
        tm = Timematr(tm.vals, tm.idx)
    else
        tm = Timematr(tn.vals, tn.idx)
    end
    return tm
end

##############
## Timematr ##
##############

## to Timedata: conversion upwards - always works
function convert(::Type{Timedata}, tm::Timematr)
    Timedata(tm.vals, tm.idx)
end

## to Timenum: conversion upwards - always works
function convert(::Type{Timenum}, tm::Timematr)
    Timenum(tm.vals, tm.idx)
end

##############
## Timedata ##
##############

## to Timenum: conversion downwards - fails for non-numeric values
function convert(::Type{Timenum}, td::Timedata)
    Timenum(td.vals, td.idx)
end

## to Timematr: conversion downwards - fails for NAs
function convert(::Type{Timematr}, td::Timedata)
    Timematr(td.vals, td.idx)
end

################
## DataFrames ##
################

## convert DataFrame with dates column (as String) to TimeData object

function convert(::Type{AbstractTimedata}, df::DataFrame)
    ## test if some column already is of type Date
    if any(eltypes(df) .== Date)
        ## take first occuring column
        dateCol = find(eltypes(df) .== Date)[1]
        idx = convert(Array, df[dateCol])
        delete!(df, [dateCol])

    else ## find column that contain dates as String

        # find columns that have been parsed as Strings by readtable
        col_to_test = Array(Symbol, 0)

        nCols = size(df, 2)
        for ii=1:nCols
            isa(df[1, ii], String)?
            push!(col_to_test, names(df)[ii]):
            nothing
        end

        # test each column's data to see if Datetime will parse it
        col_that_pass = Array(Symbol, 0)

        for colname in col_to_test
            d = match(r"[-|\s|\/|.]", df[1, colname])
            d !== nothing? (bar = split(df[1, colname], d.match)): (bar = [])
            if length(bar) == 3
                push!(col_that_pass, colname)
            end
        end

        # parse first column that passes the Datetime regex test
        idx = Date[Date(d) for d in df[col_that_pass[1]]] # without
        # Date it would fail chkIdx
        # in constructor

        delete!(df, [col_that_pass[1]])
    end

    ## try whether DataFrame fits subtypes
    try
        td = Timematr(df, idx)
        return td
    catch
        try
            td = Timenum(df, idx)
            return td
        catch
            td = Timedata(df, idx)
            return td
        end
    end

    return td
end

################
## TimeArrays ##
################

function convert(::Type{AbstractTimedata}, ta::TimeArray)

    namesAsSymbols = [DataFrames.identifier(nam) for nam in ta.colnames]

    df = composeDataFrame(ta.values, namesAsSymbols)
    idx = ta.timestamp

    try
        td = Timematr(df, idx)
        return td
    catch
        try
            td = Timenum(df, idx)
            return td
        catch
            td = Timedata(df, idx)
            return td
        end
    end
    td
end

function convert(::Type{TimeArray}, tn::AbstractTimenum)
    dats = idx(tn)
    nams = UTF8String[string(names(tn)[ii]) for ii = 1:size(tn, 2)]
    return TimeArray(dats, asArr(tn, Float64, NaN), nams)
end


#!/usr/bin/env julia

"""
	Consensus checking for variant positions across samples
"""
file = ARGS[1] # .vcf/.vcf.gz file
include("../src/read.jl")

vcf = read(file)
header = []
names = split(vcf[1])[5:end]  # assumes the first line as header
push!(header, "Position")
push!(header, "Consensus")
append!(header, names)

for col in header
	print(col, '\t')
end
println()

for line in vcf[2:end]
	arr =[]
	cols = split(line)
	print(cols[2], '\t')
	alts = []
	for sample in cols[5:end]
		altRead = parse(Int64, split(sample, ':')[1])
		if altRead >= 3
			push!(alts, 1)
		else
			push!(alts, 0)
		end
	end

	@printf("%6.3f", consensus)
	for i in alts
		print('\t', i)
	end
	println()
end


# gendata.jl
# generate data

isfile("data/set2.jld") || generate_data(name="set2", p=2)
isfile("data/set5.jld") || generate_data(name="set5", p=5)

say("done.")


using EchoviewEcs

using Test

filename = joinpath(dirname(@__FILE__),
                  "data/sample.ecs")

calibrations = load(filename)

@test length(calibrations) == 4
@test calibrations[1]["Frequency"] == 18.0
@test calibrations[2]["Frequency"] == 38.0
@test calibrations[3]["Frequency"] == 120.0
@test calibrations[4]["Frequency"] == 200.0
@test calibrations[4]["EK60SaCorrection"] == -1.54f0

@test calibrations[1]["SoundSpeed"] == 1520f0 # Specified in SourceCal
@test calibrations[2]["SoundSpeed"] == 1541.21f0 # Inheritted from FileSet


module TestUtilities

using MLJIteration
using MLJBase
using ..DummyModel
using Test

@testset "get and set iteration paramter" begin
    model = DummyIterativeModel()
    X, y = make_dummy(N=2)
    mach = machine(model, X, y)
    iter = MLJBase.iteration_parameter(model)
    @test MLJIteration.rget(model, iter) == 10
    @test MLJIteration.rget(mach, iter) == 10
    MLJIteration.rset!(model, iter, 42)
    @test MLJIteration.rget(model, iter) == 42
    @test MLJIteration.rget(mach, iter) == 42
    MLJIteration.rset!(mach, iter, 17)
    @test MLJIteration.rget(model, iter) == 17
    @test MLJIteration.rget(mach, iter) == 17
end

end

true


"""
    bench_core2core_latency([cpuids = 0:Sys.CPU_THREADS-1; nbench = 5, nsamples::Integer = 100, mode::Symbol = :min])
A tool for measuring core-to-core latency (i.e. inter-core latency) in nanoseconds.

The measured latencies correspond to a full roundtrip between two cores. Divide them by two to obtain an estimate for the time needed to fetch data from another core.

**Important:** At least two Julia threads are required (`julia -t2`)!

**Refs:** Largely inspired by [rigtorp/c2clat](https://github.com/rigtorp/c2clat) and [ajakubek/core-latency](https://github.com/ajakubek/core-latency).
"""
function bench_core2core_latency(cpuids=0:(Sys.CPU_THREADS - 1); nbench=5, kwargs...)
    # check validity of cpuids input
    for c in cpuids
        if c < 0 || c > Sys.CPU_THREADS
            @error("CPU IDs must all be non-negative and ≤ Sys.CPU_THREADS.")
        end
    end
    # backup current thread affinity
    pinning_before = getcpuids()
    # run benchmarks
    ncpuids = length(cpuids)
    latencies = zeros(ncpuids, ncpuids)
    for b in 1:nbench
        for (j, cpu2) in pairs(cpuids)
            for (i, cpu1) in pairs(cpuids)
                @inbounds latencies[i, j] += Core2CoreLatency.run_bench(
                    cpu1, cpu2; kwargs...
                )
            end
        end
    end
    latencies ./= nbench
    # restore previous thread affinity
    pinthreads(pinning_before)
    return latencies
end


# This file was generated, do not modify it. # hide
fit!(y1_oos, verbosity=0)

figure(figsize=(8,6))
step(xsort, ysort, label="truth", where="mid")
plot(x, y1_oos(), ls="none", marker="o", label="linear oos")
legend()

savefig(joinpath(@OUTPUT, "s2.svg")) # hide


using Mmap

struct DAF
    filename::String
    array::Vector{UInt8}
    little::Bool
    id::String
    nd::Int32
    ni::Int32
    name::String
    first::Int32
    last::Int32
    ss::Int32
    nc::Int32
end

const RECORD_LENGTH = 1024

function getrecord(array, n)
    array[1+RECORD_LENGTH*(n-1):RECORD_LENGTH*n]
end

getrecord(daf::DAF, n) = getrecord(daf.array, n)

function DAF(filename)
    if !isfile(filename)
        throw(ArgumentError("'$filename' does not exist."))
    end
    array = Mmap.mmap(filename, Vector{UInt8})
    fr = readfilerecord(getrecord(array, 1))
    DAF(filename, array, fr...)
end

@inline function reinterpret_getindex(::Type{T}, b::Vector{UInt8}, i, le::Bool) where T
    # Read in host endian
    @boundscheck if i + sizeof(T) - 1 > length(b)
        throw(BoundsError(b, i + sizeof(T) - 1))
    end
    v = unsafe_load(Ptr{T}(pointer(b, i)))
    le ? htol(v) : hton(v)
end

@inline function reinterpret_getindex(::Type{T}, b::Vector{UInt8},
                                      idxs::NTuple{N,Int}, le::Bool) where {N, T}
    @inbounds ntuple(i->reinterpret_getindex(T, b, idxs[i], le), Val(N))
end

function readint(record, address, littleendian=true)
    reinterpret_getindex(Int32, record, address + 1, littleendian)
end

function readascii(record, address, len)
    rstrip(String(record[address+1:address+len]))
end

function islittleendian(record, legacy)
    if legacy
        nd = readint(record, 8)
        if nd == 2
            return true
        end
        nd = readint(record, 8, false)
        if nd == 2
            return false
        else
            error("Endianess could not be detected.")
        end
    end

    endianess = readascii(record, 88, 8)
    endianess == "LTL-IEEE"
end

function readfilerecord(record)
    ftpstr = b"FTPSTR:\r:\n:\r\n:\r\x00:\x81:\x10\xce:ENDFTP"
    id = readascii(record,0,8)
    legacy = id == "NAIF/DAF"
    little = islittleendian(record, legacy)
    nd = readint(record, 8, little)
    ni = readint(record, 12, little)
    name = readascii(record, 16, 60)
    first = readint(record, 76, little)
    last = readint(record, 80, little)
    ss = nd + div(ni+1, 2)
    nc = 8*ss
    if ~legacy
        if ftpstr != record[699+1:699+28]
            error("This DAF file is damaged.")
        end
    end
    little, id, nd, ni, name, first, last, ss, nc
end

function summaryheader(record, little)
    next, nsum = reinterpret_getindex(Float64, record, (1, 17), little)
    round(Int32, next), round(Int32, nsum)
end

function addsummaries!(summaries, record, names, nsum, ss, nc)
    for i = 1:nsum
        push!(summaries, (
            rstrip(String(names[1+(i-1)*nc:i*nc])),
            record[25+(i-1)*ss*8:25+i*ss*8]
        ))
    end
end

function getsummaries(daf::DAF)
    summaries = Tuple{String,Vector{UInt8}}[]
    record = getrecord(daf, daf.first)
    names = getrecord(daf, daf.first+1)
    next, nsum = summaryheader(record, daf.little)
    addsummaries!(summaries, record, names, nsum, daf.ss, daf.nc)
    while next != 0
        record = getrecord(daf, next)
        names = getrecord(daf, next+1)
        next, nsum = summaryheader(record, daf.little)
        addsummaries!(summaries, record, names, nsum, daf.ss, daf.nc)
    end
    summaries
end


function AESCBC(plaintext, iv::Array{UInt8, 1}, key::AbstractAESKey, cache::AbstractAESCache)
	len = length(plaintext)
	pad = 16 - (len % 16)
	ciphertext = similar(Array{UInt8, 1}, pad + len)
	AESCBC!(ciphertext, plaintext, iv, key, cache)
end

function AESCBC!(ciphertext, plaintext, iv::Array{UInt8, 1}, key::AbstractAESKey, cache::AbstractAESCache)
	len = length(plaintext)
	pad = 16 - (len % 16)
	for i in 1:len
		ciphertext[i] = UInt8(plaintext[i])
	end
	for i in 1:pad
		ciphertext[len+i] = pad
	end
	iters = Int((len + pad) / 16)
	ciphertextblock = AESBlock(ciphertext)
	prevblock = AESBlock(ciphertext)
	# iteration 1 start
	ciphertextblock .⊻= iv
	AESEncryptBlock!(ciphertextblock, ciphertextblock, key, cache)
	increment!(ciphertextblock)
	# iteration 1 end
	for i in 2:iters
		@. ciphertextblock ⊻= prevblock
		AESEncryptBlock!(ciphertextblock, ciphertextblock, key, cache)
		increment!(ciphertextblock)
		increment!(prevblock)
	end
	ciphertext
end

function AESCBC_D(ciphertext::Array{UInt8, 1}, iv::Array{UInt8, 1}, key::AbstractAESKey, cache::AbstractAESCache; remove_pad=true)
	len = length(ciphertext)
	iters = Int(len / 16)
	plaintext = similar(Array{UInt8, 1}, len)
	AESCBC_D!(plaintext, ciphertext, iv, key, cache; remove_pad=remove_pad)
end

function AESCBC_D!(plaintext, ciphertext::Array{UInt8, 1}, iv::Array{UInt8, 1}, key::AbstractAESKey, cache::AbstractAESCache; remove_pad=true)
	len = length(ciphertext)
	iters = Int(len / 16)
	plaintextblock = AESBlock(plaintext)
	ciphertextblock = AESBlock(ciphertext)
	prevblock = AESBlock(ciphertext)
	# iteration 1 start
	AESDecryptBlock!(plaintextblock, ciphertextblock, key, cache)
	plaintextblock .⊻= iv
	increment!(plaintextblock)
	increment!(ciphertextblock)
	# iteration 1 end
	for i in 2:iters
		AESDecryptBlock!(plaintextblock, ciphertextblock, key, cache)
		plaintextblock .⊻= prevblock
		increment!(plaintextblock)
		increment!(ciphertextblock)
		increment!(prevblock)
	end
	if remove_pad
		pad = plaintext[end]
		@view(plaintext[1:len-pad])
	else
		plaintext
	end
end


#module Plotfuncs
    using .Plots
    using LinearAlgebra
    using StatsBase
    #using ..TightBinding
    export plot_lattice_2d,calc_band_plot,plot_DOS,calc_band_plot_finite
    export plot_fermisurface_2D
    export get_DOS

    function plot_lattice_2d(lattice)
        dim = lattice.dim
        lw = 0.5
        ls = :dash
        colors = ["red","blue","orange","brown","yellow"]
        println("Plot the 2D lattice structure")
        if dim != 2
            println("Error! dim should be 2 to plot the lattice. dim is $dim")
            return
        end
        pls = plot(size=(500,500), legend=false)
        r0 = [0,0]
        r1 = lattice.vectors[1]
        r2 = lattice.vectors[2]

        na = 2
        nb = 2

        for ia = 0:na
            for ib=0:nb
                rorigin = lattice.vectors[1]*ia +  lattice.vectors[2]*ib
                plot!([r0[1]+rorigin[1], r1[1]+rorigin[1]],
                [r0[2]+rorigin[2], r1[2]+rorigin[2]], color="red", lw=lw, ls=ls)
                plot!([r0[1]+rorigin[1], r2[1]+rorigin[1]],
                [r0[2]+rorigin[2], r2[2]+rorigin[2]], color="red", lw=lw, ls=ls)
                plot!([r1[1]+rorigin[1], r1[1]+r2[1]+rorigin[1]],
                [r1[2]+rorigin[2], r1[2]+r2[2]+rorigin[2]], color="red", lw=lw, ls=ls)
                plot!([r2[1]+rorigin[1], r1[1]+r2[1]+rorigin[1]],
                [r2[2]+rorigin[2], r1[2]+r2[2]+rorigin[2]], color="red", lw=lw, ls=ls)

                for i=1:lattice.numatoms
                    x,y = get_position(lattice,lattice.positions[i])
                    plot!([x+rorigin[1]],[y+rorigin[2]],marker=:circle,color=colors[i])
                end

                for i=1:lattice.numhopps
                    ijpositions = lattice.hoppings[i].ijpositions
                    x,y = get_position(lattice,ijpositions)
                    ijatoms = lattice.hoppings[i].ijatoms
                    x0,y0 = get_position(lattice,lattice.positions[ijatoms[1]])
                    plot!([x0+rorigin[1], x0+x+rorigin[1]],
                    [y0+rorigin[2], y0+y+rorigin[2]], color="black", lw=1)
                end
            end
        end
        plot!(aspect_ratio=:equal)
        return pls
    end

    """
        plot_fermisurface_2D(lattice::Lattice;Eshift = 0.0,nk = 20)

    Show the Fermi surface for 2D system. \\
    This plots the contour with E=0.0+Eshift \\
    Eshift: the energy shift from the chemical potential lattice.μ \\
    nk: the number of meshes in 2D momentum space. The total mesh is nk x nk.
    """
    function plot_fermisurface_2D(lattice;Eshift = 0.0,nk = 50)
        ham = hamiltonian_k(lattice) #Construct the Hamiltonian
        dim = lattice.dim
        n = lattice.numatoms

        if dim != 2
            println("This function only supports 2D case. lattice.dim should be 2. But dim = $dim .")
            return
        end

        k1s = range(-0.51, stop = 0.51, length = nk)
        k2s = range(-0.51, stop = 0.51, length = nk)
        kx_vec = zeros(Float64,nk)
        ky_vec = zeros(Float64,nk)
        energies = zeros(Float64,n,nk,nk)
        colors = ["red","blue","orange","brown","yellow"]
        kxmin,kymin = get_position_kspace(lattice,[-0.5,-0.5])
        kxmax,kymax = get_position_kspace(lattice,[0.5,0.5])

        for i1=1:nk
            for i2=1:nk
                kx,ky = get_position_kspace(lattice,[k1s[i1],k2s[i2]])
                kx_vec[i1] = kx
                ky_vec[i2] = ky
                energy = dispersion(dim,n,ham,[kx,ky])
                for j=1:n
                    energies[j,i1,i2] = energy[j]
                end

            end
        end
#        println("energies ",energies[1,:,:])
        #println(kx_vec)
        i = 1
        pls = contour(kx_vec,ky_vec,energies[i,:,:],levels=[Eshift],
            xlabel="kx",ylabel="ky",aspect_ratio=1,
            xlims=(kxmin,kxmax),ylims=(kymin,kymax)
            )
        for i=2:n
            pls = contour!(kx_vec,ky_vec,energies[i,:,:],levels=[Eshift],
                xlabel="kx",ylabel="ky",aspect_ratio=1,
                xlims=(kxmin,kxmax),ylims=(kymin,kymax)
                )
        end
        return pls

    end

    function calc_band_plot(klines,lattice)
        ham = hamiltonian_k(lattice)
        numlines = klines.numlines
        dim = lattice.dim
        klength = 0.0
        vec_k = []
        numatoms = lattice.numatoms
        energies = zeros(Float64,numatoms,0)
        #plot(x,xticks=([4,5],["G","F"]))
        xticks_values = []
        xticks_labels = []

        for i=1:numlines
            push!(xticks_values,klength)
            push!(xticks_labels,klines.kpoints[i].name_start)
            kmin = klines.kpoints[i].kmin
            kmax = klines.kpoints[i].kmax
            nk = klines.kpoints[i].nk

            kmin_real = kmin[:]
            kmax_real = kmax[:]

            kdistance = sqrt(sum((kmin_real .- kmax_real).^2))
            vec_k_temp,energies_i = calc_band(kmin_real,kmax_real,nk,lattice,ham)
            energies_i = sort(energies_i,dims=1)
            vec_k_i = range(klength,length=nk,stop=klength+kdistance)
            vec_k = vcat(vec_k,vec_k_i)
            energies = hcat(energies,energies_i)
            klength += kdistance
            push!(xticks_values,klength)
            push!(xticks_labels,klines.kpoints[i].name_end)
        end

        pls = plot(vec_k[:],[energies[i,:] for i=1:numatoms],
                        xticks = (xticks_values,xticks_labels))

        return pls
    end

    function calc_band_plot_finite(klines,lattice,direction;periodic = true,nsites = 20)
        ham = hamiltonian_k_1d(lattice,direction,periodic=periodic,nsites=nsites)
        numlines = klines.numlines
        dim = lattice.dim-1
        klength = 0.0
        vec_k = Float64[]
        numatoms = lattice.numatoms
        N = numatoms*nsites
        energies = zeros(Float64,0,N)
        #plot(x,xticks=([4,5],["G","F"]))
        xticks_values = []
        xticks_labels = []

        for i=1:numlines

            push!(xticks_values,klength)
            push!(xticks_labels,klines.kpoints[i].name_start)
            kmin = klines.kpoints[i].kmin
            kmax = klines.kpoints[i].kmax
            nk = klines.kpoints[i].nk

            kmin_real = kmin[:]
            kmax_real = kmax[:]

            kdistance = sqrt(sum((kmin_real .- kmax_real).^2))

            vec_k_temp = zeros(Float64,dim,nk)

            for idim=1:dim
                k = range(kmin_real[idim],length=nk,stop=kmax_real[idim])
                vec_k_temp[idim,:] = k[:]
            end
            energies_i = zeros(Float64,nk,N)
            for ik=1:nk
                energies_i[ik,:] = eigen(ham(vec_k_temp[:,ik])).values[:]
            end

            vec_k_i = range(klength,length=nk,stop=klength+kdistance)
            vec_k = vcat(vec_k,vec_k_i)
            energies = vcat(energies,energies_i)
            klength += kdistance
            push!(xticks_values,klength)
            push!(xticks_labels,klines.kpoints[i].name_end)
        end

        pls = plot(vec_k[:],energies,
                        xticks = (xticks_values,xticks_labels),legend=false)

        return pls
    end

    function plot_DOS(lattice,nk;nbins=100)
        dim = lattice.dim
        hist = get_DOS(lattice,nk;nbins=nbins)
        pls = bar(hist.edges[1][1:end],hist.weights/nk^dim)
        return pls

        #=
        dim = lattice.dim
        ham = hamiltonian_k(lattice)
        n = lattice.numatoms

        if dim == 1
            energies = zeros(Float64,n,nk)
            dk = 1/(nk-1)
            for ik = 1:nk
                k = (ik-1)*dk
            end
        elseif dim == 2
            energies = zeros(Float64,n,nk,nk)
            dk1 = 1/(nk-1)
            dk2 = 1/(nk-1)
            for ik1 = 1:nk
                k1 = (ik1-1)*dk1
                for ik2 = 1:nk
                    k2 = (ik2-1)*dk2
                    kx,ky = get_position_kspace(lattice,[k1,k2])
                    energy = dispersion(dim,n,ham,[kx,ky])
                    energies[:,ik1,ik2] = energy[:]

                end
            end
        elseif dim == 3
            energies = zeros(Float64,nk,nk,nk)
            dk1 = 1/(nk-1)
            dk2 = 1/(nk-1)
            dk3 = 1/(nk-1)

            for ik1 = 1:nk
                k1 = (ik1-1)*dk1
                for ik2 = 1:nk
                    k2 = (ik2-1)*dk2
                    for ik3=1:nk
                        k3= (ik3-1)*dk3
                        kx,ky,kz = get_position_kspace(lattice,[k1,k2,k3])
                        energy = dispersion(dim,n,ham,[kx,ky,kz])
                        energies[:,ik1,ik2,ik3] = energy[:]
                    end

                end
            end
        end
        energies = vec(energies)
#            println(energies)
        pls = histogram(energies,bins=nbins,normalize=true)

        return pls
        =#
    end

#end


# Generic functions that apply to all (or most) objects

set_pref_scale!(switches :: AbstractEntrySwitches{F1},
    prefScale :: F1) where F1 =
    switches.entryPrefScale = prefScale;

"""
	$(SIGNATURES)

Scale entry probabilities by (j,l) -> c to bound away from 0.
Bound entry probs by (j,l) away from 1.
Does not guarantee particular min values (that is hard to do).
"""
function scale_entry_probs!(entryProb_jlcM :: Array{F1, 3},
    minEntryProb :: F1, maxEntryProb :: F1) where F1 <: Real

    J, nl, nc = size(entryProb_jlcM);
    for iType = 1 : J
        for l = 1 : nl
            prob_cV =
                max.(minEntryProb, entryProb_jlcM[iType, l, :]);

            pSum = sum(prob_cV);
            if pSum > maxEntryProb
                prob_cV .*= (maxEntryProb / pSum);
            end
            entryProb_jlcM[iType, l, :] .= prob_cV;
        end
        # for ic = 1 : nc
        #     # Scale so that total entry prob j -> c is at least minEntryProb
        #     prob_lV = entryProb_jlcM[iType, :, ic];
        #     pSum = sum(prob_lV);
        #     if pSum < 0.000001
        #         # Basically all zero entries. Set to a constant
        #         prob_lV = fill(minEntryProb / nl, nl);
        #     elseif pSum < minEntryProb
        #         prob_lV .*= (minEntryProb / pSum);
        #     end
        #     entryProb_jlcM[iType, :, ic] = prob_lV;
        # end

        # # Scale so that entry prob j <= maxEntryProb
        # pSum = sum(entryProb_jlcM[iType,:,:]);
		# if pSum > maxEntryProb
		# 	entryProb_jlcM[iType,:,:] .*= (maxEntryProb / pSum);
        # end
	end
	return nothing
end

## ---------------  Constructors

"""
    `init_entry_decision(objId :: ObjectId, switches :: AbstractEntrySwitches)`

Initializes an `AbstractEntryDecision` from its switches.
"""
function init_entry_decision end

# Initialize entry preference scale parameter.
function init_entry_prefscale(switches :: AbstractEntrySwitches{F1},
    st :: SymbolTable) where F1
    entryPrefScale = switches.entryPrefScale;
    pEntryPref = Param(:entryPrefScale,
    LatexLH.description(st, :pScaleEntry), latex(st, :pScaleEntry), entryPrefScale, entryPrefScale, F1(0.1), F1(3.0),
        switches.calEntryPrefScale);
end

## --------------  Solve entry decisions

"""
    $(SIGNATURES)

Entry probability for a student who is admitted to colleges in `admitV`.
Returns: Entry prob by [type, college], expected value at decision stage by type.
This function does not handle the sequential nature of admissions. It is mainly here for unified interface.
It defaults to the one step entry protocol. If that does not apply for a protocol, need to define a new method (e.g. two step entry).
"""
function entry_probs(e :: AbstractEntryDecision{F1},
    vWork_jV :: Vector{F1}, vCollege_jcM :: Matrix{F1}, admitV;
    prefShocks :: Bool = true) where F1 <: Real

    if prefShocks
        prefScale = entry_pref_scale(e);
    else
        prefScale = zero(F1);
    end
    return one_step_entry_probs(prefScale, vWork_jV, vCollege_jcM, admitV);
end

# The same for one student
function entry_probs(e :: AbstractEntryDecision{F1},
    vWork :: F1, vCollege_cV :: Vector{F1}, admitV;
    prefShocks :: Bool = true) where F1

    if prefShocks
        prefScale = entry_pref_scale(e);
    else
        prefScale = zero(F1);
    end
    return one_step_entry_probs(prefScale, vWork, vCollege_cV, admitV)
end

# The same shutting down preference shocks
# function entry_probs_no_pref_shocks(e :: AbstractEntryDecision{F1},
#     vWork :: F1, vCollege_cV :: Vector{F1}, admitV) where F1

#     return one_step_entry_probs(zero(F1), vWork, vCollege_cV, admitV)
# end

# Generic entry decision. One step. Given pref shock scale.
# Works for any entry protocol where entry works in one step (work/study and which college at the same time).
function one_step_entry_probs(entryPrefScale :: F1,
    vWork_jV :: Vector{F1}, vCollege_jcM :: Matrix{F1}, admitV) where F1 <: Real

    J, nc = size(vCollege_jcM);
    prob_jxM = zeros(F1, J, nc);
    if isempty(admitV)
        eVal_jV = copy(vWork_jV);
    else
        if entryPrefScale > 0.00001
            # Prob of work in column 1. Then admitted colleges.
            probM, eVal_jV = EconLH.extreme_value_decision(
                hcat(vWork_jV, vCollege_jcM[:, admitV]),
                entryPrefScale; demeaned = true);
            for j = 1 : J
                probV = probM[j,:];
                make_valid_probs!(probV);
                prob_jxM[j, admitV] .= probV[2 : end];
            end
        else
            eVal_jV, icV = max_choices(vWork_jV, vCollege_jcM, admitV);
            for j = 1 : J
                if icV[j] > 0
                    prob_jxM[j, icV[j]] = one(F1);
                end
            end
        end
    end
    return prob_jxM, eVal_jV
end

# The same for one type.
# Preference shocks may be zero.
function one_step_entry_probs(entryPrefScale :: F1,
    vWork :: F1, vCollege_cV :: Vector{F1}, admitV) where F1 <: Real

    nc = length(vCollege_cV);
    prob_cV = zeros(F1, nc);
    if isempty(admitV)
        eVal = vWork;
    else
        if entryPrefScale > 0.00001
            probV, eVal = EconLH.extreme_value_decision_one(
                vcat(vWork, vCollege_cV[admitV]), entryPrefScale;
                demeaned = true);
            make_valid_probs!(probV);
            prob_cV[admitV] .= probV[2 : end];
        else
            # No preference shocks - just choose the best option
            eVal, ic = max_choice(vWork, vCollege_cV, admitV);
            if ic > 0
                prob_cV[ic] = one(F1);
            end
        end
    end
    return prob_cV, eVal
end

"""
	$(SIGNATURES)

Entry probs for one student across all admissions sets. Handles the case where some colleges are full.
Always for multiple locations (matrix inputs). But one location is allowed.

# Outputs
- `entryProb_clM`: probability of entering each college
- `eVal`: expected value
- `entryProbBest_clM`: fraction of students who enter college `c, l` and for who `c` is the best available college (in any location). For the top college, this is identical to the entry rate.
"""
function entry_decisions_one_student(entryS :: AbstractEntryDecision{F1},
    admissionS :: AbstractAdmissionsRule{I1, F1},
    admProbFct :: AF1,
    vWork :: F1, vCollege_cV :: AbstractVector{F1},
    endowPct :: F1, full_clM :: AbstractMatrix{Bool}, l :: Integer;
    prefShocks :: Bool = true)  where
    {AF1 <: AbstractAdmProbFct{<: Real}, I1, F1}

    nl = n_locations(entryS);
    nc = n_colleges(entryS);
    @check size(full_clM) == (nc, nl)

    # Value of college is the same for all locations, except local
    vCollege_clM = repeat(vCollege_cV, outer = (1, nl));
    vCollege_clM[:, l] .+= value_local(entryS);

    entryProb_clM = zeros(F1, nc, nl);
    eVal = zero(F1);
    # Fraction of students who attend the best available college (regardless of location). By quality attended. For the best college, this equals `entryProb_cV`.
    entryProbBest_clM = zeros(F1, nc, nl);
    # Admission rule gives admission to one college type in all locations
    for (iSet, admitV) in enumerate(admissionS)
        # Prob that each person draws this college set
        probSet = prob_coll_set(admissionS, admProbFct, iSet, endowPct);
        avail_clM = available_colleges(entryS, full_clM, admitV, l);
        @check size(avail_clM) == size(vCollege_clM)

        # Entry probs for this set
        prob_clV, eValSet =
            entry_probs(entryS, vWork, vec(vCollege_clM), vec(avail_clM);
                prefShocks = prefShocks);
        # Tested separately that `reshape` undoes `vec`
        prob_clM = reshape(prob_clV, nc, nl);
        entryProb_clM .+= probSet .* prob_clM;
        eVal += probSet * eValSet;

        # Fraction attending best available.
        cBest = best_available(avail_clM);
        entryProbBest_clM[cBest,:] .+= probSet .* prob_clM[cBest, :];
    end

    make_valid_probs!(entryProb_clM);
    make_valid_probs!(entryProbBest_clM);
    # Deal with rounding errors that arise during scaling.
    bracket_array!(entryProbBest_clM, zeros(F1, nc, nl), entryProb_clM);
    return entryProb_clM, eVal, entryProbBest_clM
end

# Can only attend colleges that are not full, where student is admitted.
function available_colleges(entryS :: AbstractEntryDecision{F1},
    full_clM :: AbstractMatrix{Bool},
    admitV, l :: Integer) where F1

    # Do not use falses here. It creates a BitArray
    avail_clM = fill(false, size(full_clM));
    # Admitted counts for all locations
    avail_clM[admitV, :] .= true;
    avail_clM[full_clM] .= false;
    # If there are local-only colleges, mark those as not available
    mark_local_only_colleges!(entryS, avail_clM, l);
    return avail_clM
end

function best_available(avail_clM)
    cBest = 1;
    for ic = 2 : size(avail_clM, 1)
        if any(avail_clM[ic,:])
            cBest = ic;
        end
    end
    return cBest
end

# Mark local only colleges as unavailable when not local.
function mark_local_only_colleges!(entryS :: AbstractEntryDecision{F1},
    avail_clM, l) where F1

    idxV = local_only_colleges(entryS);
    if !isempty(idxV)
        for ic ∈ idxV
            aLocal = avail_clM[ic, l];
            avail_clM[ic, :] .= false;
            avail_clM[ic, l] = aLocal;
        end
    end
    return nothing
end

# The same with one location. Simply calls the multi-location code with one location dimension.
function entry_decisions_one_student(entryS :: AbstractEntryDecision{F1},
    admissionS :: AbstractAdmissionsRule{I1, F1},
    admProbFct :: AF1,
    vWork :: F1, vCollege_cV :: AbstractVector{F1},
    endowPct :: F1, full_cV :: AbstractVector{Bool};
    prefShocks :: Bool = true)  where
    {AF1 <: AbstractAdmProbFct{<: Real}, I1, F1}

    # Just solve the multi-location version for the first location.
    entryProb_clM, eVal, entryProbBest_clM =
        entry_decisions_one_student(entryS, admissionS, admProbFct,
        vWork, vCollege_cV, endowPct,  repeat(full_cV, outer = (1,1)), 1;
        prefShocks = prefShocks);

    make_valid_probs!(entryProb_clM)
    return vec(entryProb_clM), eVal, entryProbBest_clM
end

## ------------  Implied outcomes

# """
# 	$(SIGNATURES)

# Compute college enrollment from type mass and entry probabilities.
# """
# function college_enrollment(entryProb_jcM :: Matrix{F1},
#     typeMass_jV) where F1 <: Real

#     J, nc = size(entryProb_jcM);
#     enrollV = zeros(F1, nc);
#     for ic = 1 : nc
#         enrollV[ic] = sum(entryProb_jcM[:, ic] .* typeMass_jV);
#     end
#     return enrollV
# end

# function college_enrollment(e :: AbstractEntryDecision{F1},
#     entryProb_jcM :: Matrix{F1}) where F1 <: Real
#     return college_enrollment(e.switches, entryProb_jcM);
# end

# function college_enrollment(e :: AbstractEntrySwitches{F1},
#     entryProb_jcM :: Matrix{F1}) where F1 <: Real

#     @assert n_locations(e) == 1
#     J = size(entryProb_jcM, 1);
#     return college_enrollment(entryProb_jcM, type_masses(e));
# end

# function college_enrollment(e :: AbstractEntryDecision{F1},
#     entryProb_cV :: Vector{F1},  j :: Integer) where F1

#     @assert n_locations(e) == 1
#     return entryProb_cV .* type_mass(e, j);
# end

# """
#     $(SIGNATURES)

# Return `Bool` vector that indicates which colleges are full. Only matters for entry structures with capacity constraints.
# """
# function colleges_full(e :: AbstractEntryDecision{F1},
#     entryProb_jcM :: Matrix{F1}) where F1 <: Real

#     @assert n_locations(e) == 1  "Not valid for multiple locations"
#     return college_enrollment(e, entryProb_jcM) .>= capacities(e)
# end

# --------------


# T is the bitpattern type for logging, i is the logbook number
struct Sherlog16{T<:AbstractFloat,i} <: AbstractSherlog
    val::T
end

# conversions back from Sherlog16
Base.Float64(x::Sherlog16) = Float64(x.val)
Base.Float32(x::Sherlog16) = Float32(x.val)
Base.Float16(x::Sherlog16) = Float16(x.val)
Base.Int64(x::Sherlog16) = Int64(x.val)
Base.Int32(x::Sherlog16) = Int32(x.val)
Base.Int16(x::Sherlog16) = Int16(x.val)

# otherwise conversion to itself is ambigious
Sherlog16{T,i}(x::Sherlog16) where {T<:AbstractFloat,i} = Sherlog16{T,i}(x.val)
Sherlog16(x::Sherlog16{T,i}) where {T<:AbstractFloat,i} = Sherlog16{T,i}(x.val)

# generator functions from Float, Int and Bool
Sherlog16(x::AbstractFloat) = Sherlog16{Float16,1}(x)   # use Float16 and logbook #1 as standard
Sherlog16{T}(x::AbstractFloat) where T = Sherlog16{T,1}(x)      # use logbook #1 if no logbook provided
Sherlog16(x::Integer) = Sherlog16{Float16,1}(x)           # same for integers
Sherlog16{T}(x::Integer) where T = Sherlog16{T,1}(x)
Sherlog16{T}(x::Bool) where T = if x Sherlog16{T,1}(1) else Sherlog16{T,1}(0) end

Base.promote_rule(::Type{Int64},::Type{Sherlog16{T,i}}) where {T,i} = Sherlog16{T,i}
Base.promote_rule(::Type{Int32},::Type{Sherlog16{T,i}}) where {T,i} = Sherlog16{T,i}
Base.promote_rule(::Type{Int16},::Type{Sherlog16{T,i}}) where {T,i} = Sherlog16{T,i}

Base.promote_rule(::Type{Float64},::Type{Sherlog16{T,i}}) where {T,i} = Sherlog16{T,i}
Base.promote_rule(::Type{Float32},::Type{Sherlog16{T,i}}) where {T,i} = Sherlog16{T,i}
Base.promote_rule(::Type{Float16},::Type{Sherlog16{T,i}}) where {T,i} = Sherlog16{T,i}

Base.bitstring(x::Sherlog16) = bitstring(x.val)
Base.show(io::IO,x::Sherlog16) = print(io,"Sherlog16(",string(x.val),")")

Base.eps(::Type{Sherlog16{T,i}}) where {T,i} = eps(Float64)
Base.eps(x::Sherlog16) = eps(x.val)

Base.typemin(::Sherlog16{T,i}) where {T,i} = Sherlog16{T,i}(typemin(T))
Base.typemax(::Sherlog16{T,i}) where {T,i} = Sherlog16{T,i}(typemax(T))
Base.floatmin(::Sherlog16{T,i}) where {T,i} = Sherlog16{T,i}(floatmin(T))
Base.floatmax(::Sherlog16{T,i}) where {T,i} = Sherlog16{T,i}(floatmax(T))
Base.precision(::Sherlog16{T,i}) where {T,i} = Sherlog16{T,i}(precision(T))

function +(x::Sherlog16{T,i},y::Sherlog16{T,i}) where {T,i}
    r = x.val + y.val
    log_it(T,r,i)
    return Sherlog16{T,i}(r)
end

function -(x::Sherlog16{T,i},y::Sherlog16{T,i}) where {T,i}
    r = x.val - y.val
    log_it(T,r,i)
    return Sherlog16{T,i}(r)
end

function *(x::Sherlog16{T,i},y::Sherlog16{T,i}) where {T,i}
    r = x.val * y.val
    log_it(T,r,i)
    return Sherlog16{T,i}(r)
end

function /(x::Sherlog16{T,i},y::Sherlog16{T,i}) where {T,i}
    r = x.val / y.val
    log_it(T,r,i)
    return Sherlog16{T,i}(r)
end

function ^(x::Sherlog16{T,i},y::Sherlog16{T,i}) where {T,i}
    r = x.val ^ y.val
    log_it(T,r,i)
    return Sherlog16{T,i}(r)
end

for O in ( :(-), :(+),
           :sign,
           :prevfloat, :nextfloat,
           :round, :trunc, :ceil, :floor,
           :inv, :abs, :sqrt, :cbrt,
           :exp, :expm1, :exp2, :exp10,
           :log, :log1p, :log2, :log10,
           :rad2deg, :deg2rad, :mod2pi, :rem2pi,
           :sin, :cos, :tan, :csc, :sec, :cot,
           :asin, :acos, :atan, :acsc, :asec, :acot,
           :sinh, :cosh, :tanh, :csch, :sech, :coth,
           :asinh, :acosh, :atanh, :acsch, :asech, :acoth,
           :sinc, :sinpi, :cospi,
           :sind, :cosd, :tand, :cscd, :secd, :cotd,
           :asind, :acosd, :atand, :acscd, :asecd, :acotd
          )
    @eval begin
        function Base.$O(x::Sherlog16{T,i}) where {T,i}
            r = $O(x.val)
            log_it(T,r,i)
            return Sherlog16{T,i}(r)
        end
    end
end

for O in ( :(<), :(<=))
    @eval begin
        Base.$O(x::Sherlog16, y::Sherlog16) = $O(x.val, y.val)
    end
end