using NCDatasets
using Formatting
using ArgParse
using Statistics
using JSON

include("constants.jl")
include("CESMReader.jl")

using .CESMReader

function mreplace(x)
    return replace(x, missing=>NaN)
end

function parse_commandline()

    s = ArgParseSettings()
    @add_arg_table s begin

        "--data-file-prefix"
            help = "Data filename prefix including folder and path until the timestamp. File extension `nc` is assumed."
            arg_type = String
            required = true

        "--data-file-timestamp-form"
            help = "Data filename timestamp form. Either `YEAR` or `YEAR_MONTH`."
            arg_type = String
            required = true

        "--output-file"
            help = "Output file."
            arg_type = String
            required = true

        "--domain-file"
            help = "Domain file."
            arg_type = String
            required = true

        "--beg-year"
            help = "Year of begin."
            arg_type = Int64
            required = true

        "--end-year"
            help = "Year of end."
            arg_type = Int64
            required = true

    end

    return parse_args(ARGS, s)
end

parsed = parse_commandline()
print(json(parsed, 4))

Dataset(parsed["domain-file"], "r") do ds
    global mask = ds["mask"][:] |> mreplace
#    global mask_idx = (mask .== 0.0)
end

let
    global swflx, nswflx, TFLUX_DIV_implied, qflx, TSAS_clim, TFLUX_bot

    if parsed["data-file-timestamp-form"] == "YEAR"
        filename_format = format("{:s}{{:04d}}.nc", joinpath(parsed["data-file-prefix"]))
        form = :YEAR
    elseif parsed["data-file-timestamp-form"] == "YEAR_MONTH"
        filename_format = format("{:s}{{:04d}}-{{:02d}}.nc", joinpath(parsed["data-file-prefix"]))
        form = :YEAR_MONTH
    end

    fh = FileHandler(filename_format=filename_format, form=form)

    beg_t = (parsed["beg-year"] - 1) * 12 + 1
    end_t = (parsed["end-year"] - 1) * 12 + 12

    swflx, nswflx, TFLUX_DIV_implied, qflx, TSAS_clim, TFLUX_bot = getData(fh, ["swflx", "nswflx", "TFLUX_DIV_implied", "qflx", "TSAS_clim", "TFLUX_bot"], (parsed["beg-year"], parsed["end-year"]), (:, :))

    swflx             = mean(swflx,             dims=(3, ))[:, :, 1]
    nswflx            = mean(nswflx,            dims=(3, ))[:, :, 1]
    TFLUX_DIV_implied = mean(TFLUX_DIV_implied, dims=(3, ))[:, :, 1]
    qflx              = mean(qflx,              dims=(3, ))[:, :, 1]
    TSAS_clim         = mean(TSAS_clim,         dims=(3, ))[:, :, 1]
    TFLUX_bot         = mean(TFLUX_bot,         dims=(3, ))[:, :, 1]

#=
    swflx[mask_idx] .= NaN
    nswflx[mask_idx] .= NaN
    TFLUX_DIV_implied[mask_idx] .= NaN
    qflx[mask_idx] .= NaN
    TSAS_clim[mask_idx] .= NaN
    TFLUX_bot[mask_idx] .= NaN
=#

    global (Nx, Ny) = size(swflx)

end

Dataset(parsed["output-file"], "c") do ds

    defDim(ds, "Nx", Nx)
    defDim(ds, "Ny", Ny)

    for (varname, vardata, vardim, attrib) in [
        ("swflx",             swflx,             ("Nx", "Ny"), Dict()),
        ("nswflx",            nswflx,            ("Nx", "Ny"), Dict()),
        ("TFLUX_DIV_implied", TFLUX_DIV_implied, ("Nx", "Ny"), Dict()),
        ("qflx",              qflx,              ("Nx", "Ny"), Dict()),
        ("TSAS_clim",         TSAS_clim,         ("Nx", "Ny"), Dict()),
        ("TFLUX_bot",         TFLUX_bot,         ("Nx", "Ny"), Dict()),
    ]

        println("Doing var: ", varname)

        var = defVar(ds, varname, Float64, vardim)
        var.attrib["_FillValue"] = 1e20

        var = ds[varname]

        for (k, v) in attrib
            var.attrib[k] = v
        end

        rng = []
        for i in 1:length(vardim)-1
            push!(rng, Colon())
        end
        push!(rng, 1:size(vardata)[end])
        var[rng...] = vardata

    end

end


function read_treebank_grammar(number_sentences::Int, Score=LogProb)
    @assert number_sentences in (1, 5, 50, 500, 972)
    grammarfile = joinpath(dirname(@__FILE__), "WSJ_data", "WSJ.$number_sentences.grammar.txt")
    Grammar(readstring(grammarfile), ["START"], Score)
end

function compute_treebank_scores(number_sentences::Int, grammar=read_treebank_grammar(number_sentences))
    @assert number_sentences in (1, 5, 50, 500, 972)
    sentences = map(split, readlines(joinpath(dirname(@__FILE__), "WSJ_data", "WSJ.$number_sentences.sentences.txt")))
    scores = zeros(GeneralizedChartparsing.score_type(grammar), length(sentences))
    for (i,s) in enumerate(sentences)
        scores[i] = score(run_chartparser(s, grammar))
    end
    scores
end

function treebank_goldstandard_scores(number_sentences::Int)
    lines = readlines(joinpath(dirname(@__FILE__), "WSJ_data", "WSJ.$number_sentences.scores.txt"))
    map(parse, lines)
end


# SPD-License-Identifier: MIT
#
# This is a thin wrapper over the Knet package to have a more high-level
# API to construct the neural network.

using CUDA
using Knet

const DEFAULT_ARRAY_TYPE = Array{Float32}

trainmode!(model) = model
testmode!(model) = model
reallocate(obj, atype) = deepcopy(obj)

gpu(obj) = reallocate(obj, CuArray{Float32})
cpu(obj) = reallocate(obj, Array{Float32})

"""
    permutedims(perm)

apply a permutation. `perm` is an array of the indices of the
dimension, e.g. `perm = [2, 1, 3]` will permute dimension 2 and 1.
"""
struct PermuteDims
    perm
end
(p::PermuteDims)(X) = permutedims(X, p.perm)

"""
    AddExtraDim(perm)

Add an extra second dimension. This is usually to use 2d-convolution on
1d signals.
"""
struct AddExtraDim end
(::AddExtraDim)(X) = reshape(X, size(X,1), 1, size(X,2), size(X,3))

"""
    RemoveExtraDim(perm)

Remove the extra second dimension. See [`AddExtraDim`](@ref).
"""
struct RemoveExtraDim end
(::RemoveExtraDim)(X) = reshape(X, size(X,1), size(X,3), size(X,4))

"""
    Dropout(pdrop)

Dropout layer. `pdrop` is the probability of a value to be set to zero.
"""
mutable struct Dropout
    pdrop
    training
end
(d::Dropout)(X) = dropout(X, d.pdrop; drop = d.training)
Dropout(pdrop) = Dropout(pdrop, true)
trainmode!(d::Dropout) = d.training = true
testmode!(d::Dropout) = d.training = false

"""
    Dense(in, out, σ = identity)

Dense layer.
"""
struct Dense
    W
    b
    σ
end
function (d::Dense)(X)
    rX = reshape(X, size(X,1), :)
    Y = d.σ.( d.W * rX .+ d.b )
    reshape(Y, :, size(X)[2:end]...)
end

Dense(in::Int, out::Int, σ = identity) =
    Dense(param(out, in, atype = DEFAULT_ARRAY_TYPE),
          param0(out, atype = DEFAULT_ARRAY_TYPE), σ)

reallocate(d::Dense, atype) =
        Dense(Param(atype(value(d.W))), Param(atype(value(d.b))), d.σ)

"""
    Conv(kernelsize, inchannels => outchannels; stride = (1,1), padding = (0,0), dilation=(1,1))

2D-convolution layer.
"""
struct Conv
    W
    b
    stride
    padding
    dilation
end

function Conv(ksize, in_out; stride = 1, padding = 0, dilation = 1)
    W = param(ksize..., in_out.first, in_out.second, atype = DEFAULT_ARRAY_TYPE)
    b = param0(1, 1, in_out.second, 1, atype = DEFAULT_ARRAY_TYPE)
    Conv(W, b, stride, padding, dilation)
end

(c::Conv)(X) = conv4(c.W, X, padding = c.padding, stride = c.stride,
                     dilation = c.dilation) .+ c.b

function reallocate(c::Conv, atype)
    Conv(Param(atype(value(c.W))), Param(atype(value(c.b))), c.stride,
         c.padding, c.dilation)
end

"""
    BatchNorm(dim, σ = identity)

Batch normalization.
"""
struct BatchNorm
    moments
    params
    σ
end
BatchNorm(dim, σ = identity) = BatchNorm(bnmoments(),
                                         DEFAULT_ARRAY_TYPE(bnparams(dim)), σ)
(bn::BatchNorm)(X) = bn.σ.(batchnorm(X, bn.moments, bn.params))

function reallocate(bn::BatchNorm, atype)
    mean, var = bn.moments.mean, bn.moments.var
    BatchNorm(
        bnmoments(;
            momentum = bn.moments.momentum,
            mean = isnothing(mean) ? nothing : atype(mean),
            var = isnothing(var) ? nothing : atype(var),
            meaninit = bn.moments.meaninit,
            varinit = bn.moments.varinit
        ),
        atype(value(bn.params)),
        bn.σ
    )
end

"""
    chain(layer1, layer2, ...)

Build a sequential model.
"""
struct Chain
    layers
    Chain(layers...) = new(layers)
end

function (c::Chain)(X)
    for layer in c.layers
        X = layer(X)
    end
    X
end

reallocate(chain::Chain, atype) = Chain((reallocate.(chain.layers, atype))...)

function Base.show(io::IO, c::Chain)
    println(io, "Chain(")
    for layer in c.layers
        println(layer)
    end
    print(io, ")")
end


"""
    $(SIGNATURES)

Convenience function to load yaml config files, commonly used to prepare the required beamforming computational objects.
"""
function loadConfigFile(file::A) where {A <: AbstractString}
  YAML.load(open(file))
end

"""
    $(SIGNATURES)

Prepare a sas2d factor to use in the factor graph where `totalPhones` is the size of the SAS array.  Pass a known `cfgd::Dict{String,} for faster load times.`
"""
function prepareSAS2DFactor(totalPhones::Int, csvWaveData::Array{<:Real};
              cfgd::Dict=loadConfigFile(joinpath(dirname(pathof(Caesar)),"config","SAS2D.yaml")), rangemodel=:Rayleigh, chirpFile=joinpath(dirname(pathof(Caesar)),"config","chirp250.txt"))
  # # @assert size(csvWaveData,2) == totalPhones
  #Constant Acoustic Params
  fFloor = cfgd["freq_lower_cutoff"];
  fCeil = cfgd["freq_upper_cutoff"];
  fSampling = cfgd["freq_sampling"];
  soundSpeed = cfgd["sound_speed"];
  nFFT_full = cfgd["nfft_full"];
  nFFT_czt = cfgd["nfft_czt"];
  azimuthDivs = cfgd["azimuth_divs"];
  rngSNRfloor = cfgd["range_snr_floor"];
  aziSNRfloor = cfgd["azi_snr_floor"];

  #Params subject to change/ Heuristics
  nPhones = totalPhones - 1;
  azimuths = LinRange(0,360,azimuthDivs)*pi/180;
  #Data preprocessing MF + CZT
  w = exp(-2im*pi*(fCeil-fFloor)/(nFFT_czt*fSampling))
  a = exp(2im*pi*fFloor/fSampling)

  chirpIn = readdlm(chirpFile,',',Float64,'\n')

  #Matched Filter on Data In
  nFFT_full = nextpow(2,size(csvWaveData,1))  # MF
  mfData = zeros(Complex{Float64}, nFFT_full, totalPhones)
  mf = prepMF(chirpIn,nFFT_full,totalPhones) # MF
  mf(csvWaveData,mfData) # MF

  cztData = zeros(Complex{Float64}, nFFT_czt,totalPhones)
  filterCZT = prepCZTFilter(nFFT_full,totalPhones,w,nFFT_czt,a)
  filterCZT(mfData,cztData)

  #CBF Filter shared by

  FFTfreqs = collect(LinRange(fFloor,fCeil,nFFT_czt))
  cfgCBF_init = CBFFilterConfig(fFloor,fCeil,nFFT_czt,totalPhones,azimuths,soundSpeed, FFTfreqs)
  cfgCBF_init_LIE = CBFFilterConfig(fFloor,fCeil,nFFT_czt,nPhones,Float64[0.0;],soundSpeed, FFTfreqs)

  sas2d = SASBearing2D()
  if rangemodel == :Rayleigh
    sas2d.rangemodel = Rayleigh(100.0)
  elseif rangemodel == :Correlator
      range_mfData = zeros(Complex{Float64}, nFFT_full, totalPhones)
      range_mf = prepMF(chirpIn,nFFT_full,totalPhones) # MF
      range_mf(csvWaveData,range_mfData) # matched filter
      ranget = [0:8000-500;]*soundSpeed/fSampling
      sas2d.rangemodel = Vector{IIF.AliasingScalarSampler}(undef, totalPhones)
      for i in 1:totalPhones
        sas2d.rangemodel[i] = IIF.AliasingScalarSampler(ranget, exp.(norm.(range_mfData[1:7501,i])), SNRfloor=rngSNRfloor)
      end
  else
      error("No can do: what is a $rangemodel rangemodel?")
  end

  #thread safe
  sas2d.cfgTotal = cfgCBF_init
  sas2d.cfgLIE = cfgCBF_init_LIE
  sas2d.cfg = deepcopy(cfgd)
  sas2d.waveformsRaw = deepcopy(csvWaveData)
  sas2d.waveformsIn = cztData
  sas2d.debugging = false

  #create memory
  prepareSASMemory(sas2d)

  return sas2d
end

function prepareSASMemory(sas2d::SASBearing2D)::Nothing
  #Constant Acoustic Params
  fFloor = sas2d.cfg["freq_lower_cutoff"];
  fCeil = sas2d.cfg["freq_upper_cutoff"];
  fSampling = sas2d.cfg["freq_sampling"];
  soundSpeed = sas2d.cfg["sound_speed"];
  nFFT_full = sas2d.cfg["nfft_full"];
  nFFT_czt = sas2d.cfg["nfft_czt"];
  azimuthDivs = sas2d.cfg["azimuth_divs"];
  rngSNRfloor = sas2d.cfg["range_snr_floor"];
  aziSNRfloor = sas2d.cfg["azi_snr_floor"];
  totalPhones = size(sas2d.waveformsRaw,2)
  azimuths = LinRange(0,360,azimuthDivs)*pi/180;

  #create memory
  sas2d.threadreuse = Vector{SASREUSE}(undef, Threads.nthreads())
  for thritr in 1:Threads.nthreads()
      sas2d.threadreuse[thritr] = SASREUSE()
      sas2d.threadreuse[thritr].CBFFull = zeros(Complex{Float64}, getCBFFilter2Dsize(sas2d.cfgTotal));
      sas2d.threadreuse[thritr].CBFLIE = zeros(Complex{Float64}, getCBFFilter2Dsize(sas2d.cfgLIE));
      sas2d.threadreuse[thritr].BFOutFull = zeros(length(azimuths))
      sas2d.threadreuse[thritr].BFOutLIE = zeros(Complex{Float64},nFFT_czt)
      sas2d.threadreuse[thritr].BFOutLIETemp = zeros(Complex{Float64},nFFT_czt)
      sas2d.threadreuse[thritr].phaseshiftLOO = zeros(Complex{Float64},nFFT_czt)
      sas2d.threadreuse[thritr].arrayPos = zeros(Float64,totalPhones,2)
      sas2d.threadreuse[thritr].arrayPosLIE = zeros(Float64,totalPhones-1,2)
      sas2d.threadreuse[thritr].waveformsLIE = zeros(Complex{Float64},nFFT_czt,totalPhones-1)
      sas2d.threadreuse[thritr].waveformsLOOc = zeros(Complex{Float64},nFFT_czt)

      sas2d.threadreuse[thritr].BFtemp = zeros(Complex{Float64}, totalPhones,nFFT_czt)
      sas2d.threadreuse[thritr].LIEtemp = zeros(Complex{Float64}, totalPhones-1,nFFT_czt)

      # caching hack
      sas2d.threadreuse[thritr].hackazi = (Int[0;], Float64[0.0;])
      sas2d.threadreuse[thritr].oncebackidx = Int[0;]

      # a few more memory reuse positions
      sas2d.threadreuse[thritr].sourceXY = zeros(2)

      # beamformer update
      sas2d.threadreuse[thritr].temp1 = zeros(Complex{Float64},nFFT_czt);
      sas2d.threadreuse[thritr].temp2 = zeros(Complex{Float64},size(sas2d.waveformsIn));
      sas2d.threadreuse[thritr].temp = zeros(Complex{Float64},nFFT_czt,totalPhones-1);

      sas2d.threadreuse[thritr].dbg = SASDebug()
      reset!(sas2d.threadreuse[thritr].dbg)
  end
end

# plot(x=azimuths, y=acceptRatio, Geom.path())
# careful -- this great function depends on how you discretize the azimuths in CBF
function sampleDiscrete!(smpls::Vector{T}, BFresult::Vector{T}) where {T <: Union{Float64, Float32}}
  N = length(smpls)
  y = 1/2pi
  acceptRatio = 2pi*BFresult

  sample = zeros(1)
  sampleInt = zeros(Int, 1)

  d2r = 2*pi/180.0
  i = 0
  while true
     sample[1] = 179*rand()
     sampleInt[1] = ceil.(Int, sample[1])
     # acceptRatio[sampleInt[1]]
      if acceptRatio[sampleInt[1]] > rand()
         i += 1
         smpls[i] = sample[1] * d2r
        if i >= N
          break;
        end
      end
  end
  nothing
end

# this is a direct copy from TransformUtils
function wrapRad(th::Float64)
  if th >= pi
    th -= 2.0*pi
  end
  if th < -pi
    th += 2.0*pi
  end
  return th
end


radiation_type = typeof(0.1W*m^-2)
snowdepth_type = typeof(0.1m)
airtemperature_type = typeof(0.1K)
relhumidity_type = typeof(0.001)
windspeed_type = typeof(0.1m*s^-1)
soiltemperature_type = typeof(0.1K)
soilwaterpotential_type = typeof(0.0kPa)
soilwatercontent_type = typeof(0.001)

abstract type AbstractMicroclimate end
abstract type AbstractLayeredMicroclimate{S,I,A,R} <: AbstractMicroclimate end
abstract type AbstractMicroclimControl <: AbstractMicroclimate end
abstract type AbstractMicroclimGrid{S,I,A,R} <: AbstractLayeredMicroclimate{S,I,A,R} end
abstract type AbstractMicroclimPoint{S,I,A,R} <: AbstractLayeredMicroclimate{S,I,A,R} end

@udefault_kw @units @limits struct MicroclimControl{RA,SN,AT,RH,WS,ST,WP,WC} <: AbstractMicroclimControl
    radiation::RA          | 200.0  | W*m^-2 | (500.0, 1200.0)
    snowdepth::SN          | 0.0    | m      | (0.0, 10.0)
    airtemperature::AT     | 290.0  | K      | (260.0, 370.0)
    relhumidity::RH        | 0.7    | _      | (0.0, 1.0)
    windspeed::WS          | 0.1    | m*s^-1 | (0.0, 20.0)
    soiltemperature::ST    | 290.0  | K      | (260.0, 370.0)
    soilwaterpotential::WP | -100.0 | kPa    | (-0.0, -10000.0)
    soilwatercontent::WC   | 0.3    | _      | (0.0, 1.0)
end

struct MicroclimGrid{S,I,A,R,RA,SN,AT,RH,WS,ST,WP,WC} <: AbstractMicroclimGrid{S,I,A,R}
    radiation::RA
    snowdepth::SN
    airtemperature::AT
    relhumidity::RH
    windspeed::WS
    soiltemperature::ST
    soilwaterpotential::WP
    soilwatercontent::WC
end

MicroclimGrid{S,I,A,R}(radiation::RA, snowdepth::SN, airtemperature::AT, relhumidity::RH,
                        windspeed::WS, soiltemperature::ST, soilwaterpotential::WP,
                        soilwatercontent::WC) where {S,I,A,R,RA,SN,AT,RH,WS,ST,WP,WC} =
    MicroclimGrid{S,I,A,R,RA,SN,AT,RH,WS,ST,WP,WC}(radiation, snowdepth, airtemperature,
        relhumidity, windspeed, soiltemperature, soilwaterpotential, soilwatercontent)

struct MicroclimPoint{S,I,A,R,RA,SN,AT,RH,WS,ST,WP,WC} <: AbstractMicroclimPoint{S,I,A,R}
    radiation::RA
    snowdepth::SN
    airtemperature::AT
    relhumidity::RH
    windspeed::WS
    soiltemperature::ST
    soilwaterpotential::WP
    soilwatercontent::WC
end

MicroclimPoint{S,I,A,R}(radiation::RA, snowdepth::SN, airtemperature::AT, relhumidity::RH,
                        windspeed::WS, soiltemperature::ST, soilwaterpotential::WP,
                        soilwatercontent::WC) where {S,I,A,R,RA,SN,AT,RH,WS,ST,WP,WC} =
    MicroclimPoint{S,I,A,R,RA,SN,AT,RH,WS,ST,WP,WC}(radiation, snowdepth, airtemperature,
        relhumidity, windspeed, soiltemperature, soilwaterpotential, soilwatercontent)

# From external data without units
MicroclimPoint{S,I,A,R}(radiation::Vector{<:AbstractFloat}, snowdepth::Vector{<:AbstractFloat},
                      airtemperature::Matrix{<:AbstractFloat}, relhumidity::Matrix{<:AbstractFloat},
                      windspeed::Matrix{<:AbstractFloat}, soiltemperature::Matrix{<:AbstractFloat},
                      soilwaterpotential::Matrix{<:AbstractFloat}, soilwatercontent::Matrix{<:AbstractFloat}
                     ) where {S,I,A,R} = begin
    rad = missing_otherwise(to_radiation, radiation)
    snow = missing_otherwise(to_snowdepth, snowdepth)
    airt = missing_otherwise(to_airtemperature, airtemperature)
    relh = missing_otherwise(to_relhumidity, relhumidity)
    wind = missing_otherwise(to_windspeed, windspeed)
    soilt = missing_otherwise(to_soiltemperature, soiltemperature)
    soilwp = missing_otherwise(to_soilwaterpotential, soilwaterpotential)
    soilwc = missing_otherwise(to_soilwatercontent, soilwatercontent)

    args = (rad, snow, airt, relh, wind, soilt, soilwp, soilwc)
    if any(ismissing.(args))
        missing
    else
        MicroclimPoint{S,I,A,R,typeof.(args)...}(args...)
    end
end

# From microclim dataset in tuples and a cartesian index
# Requires conversion to float from Int
MicroclimPoint{S,I,A,R}(radiation::Vector, snowdepth::Vector, airtemperature::Tuple,
                        relhumidity::Tuple, windspeed::Tuple, soiltemperature::Tuple,
                        soilwaterpotential::Tuple, soilwatercontent::Tuple, i::CartesianIndex) where {S,I,A,R} = begin
    rad = missing_otherwise_cat(x -> to_radiation(0.1x), radiation, i)
    snow = missing_otherwise_cat(x -> to_snowdepth(0.001x), snowdepth, i)
    airt = missing_otherwise_cat(x -> to_airtemperature(0.1x), airtemperature, i)
    relh = missing_otherwise_cat(x -> to_relhumidity(0.001x), relhumidity, i)
    wind = missing_otherwise_cat(x -> to_windspeed(0.1x), windspeed, i)
    soilt = missing_otherwise_cat(x -> to_soiltemperature(0.1x), soiltemperature, i)
    soilwp = missing_otherwise_cat(x -> to_soilwaterpotential(0.1x), soilwaterpotential, i)
    soilwc = missing_otherwise_cat(x -> to_soilwatercontent(0.001x), soilwatercontent, i)

    args = (rad, snow, airt, relh, wind, soilt, soilwp, soilwc)
    if any(ismissing.(args))
        missing
    else
        MicroclimPoint{S,I,A,R,typeof.(args)...}(args...)
    end
end

# From MiriclimGrid and a cartesian index
MicroclimPoint(env::AbstractMicroclimGrid{S,I,A,R}, i::CartesianIndex) where {S,I,A,R} =
    MicroclimPoint{S,I,A,R}(radiation(env), snowdepth(env), airtemperature(env),
                   relhumidity(env), windspeed(env), soiltemperature(env),
                   soilwaterpotential(env), soilwatercontent(env), i)

missing_otherwise(f, d) =
    (length(d) == 0 || any(ismissing.(d))) ? missing : f.(d)

missing_otherwise_cat(f, ds::Tuple, i) = begin
    maybedata = map(d -> missing_otherwise_cat(f, d, i), ds)
    if any(map(ismissing, maybedata))
        missing
    else
        hcat(maybedata...)
    end
end
missing_otherwise_cat(f, d, i) =
    if length(d) == 0
        nothing
    else
        p = catpoint(d, i)
        any(ismissing.(p)) ? missing : f.(p)
    end

@inline layer_sizes(env::MicroclimPoint{S,I,NL,R}) where {S,I,NL,R} = layer_sizes(zero(NL), I, NL)
@inline layer_bounds(env::MicroclimPoint{S,I,NL,R}) where {S,I,NL,R} = layer_bounds(I, NL)

abstract type AbstractMicroclimInstant end

struct MicroclimInstant{M,T,I} <: AbstractMicroclimInstant
    microclimate::M
    t::T
    interp::I
end
MicroclimInstant(microclim, t, height::Number) =
    MicroclimInstant(microclim, t, LinearLayerInterpolators(microclim, height))

@inline layermax(layers, x::MicroclimInstant) = layermax(layers, x.t, x.interp.increment)

struct MicroclimateLayers{L<:Tuple}
    layers::L
end

# unit conversions
to_radiation(x) = x * W*m^-2
to_snowdepth(x) = x * m
to_airtemperature(x) = x * °C |> K
to_relhumidity(x) = x
to_windspeed(x) = x * m*s^-1
to_soiltemperature(x) = x * °C |> K
to_soilwaterpotential(x) = x * kPa
to_soilwatercontent(x) = x

# Get layer specifications from type system
get_shade(layers::AbstractLayeredMicroclimate{S,I,A,R}) where {S,I,A,R} = S
get_increments(layers::AbstractLayeredMicroclimate{S,I,A,R}) where {S,I,A,R} = I
get_nextlayer(layers::AbstractLayeredMicroclimate{S,I,A,R}) where {S,I,A,R} = A
get_range(layers::AbstractLayeredMicroclimate{S,I,A,R}) where {S,I,A,R} = R

@inline catpoint(layers::Tuple{}, i::CartesianIndex) = Float64[]
@inline catpoint(layers::Tuple, i::CartesianIndex) =
    cat([[cat((y[i, :] for y in l)..., dims=1)...] for l in layers]..., dims=2)
@inline catpoint(layers::Vector, i::CartesianIndex) =
    cat((getindex(y, i, :) for y in layers)..., dims=1)

# Accessors for single layer data
for name in (:radiation, :snowdepth)
    eval(quote
        $name(env::AbstractMicroclimate) = env.$name
        $name(env::AbstractMicroclimPoint, t, interp) = lin_interp($name(env), t)
        $name(env::AbstractMicroclimInstant) = lin_interp($name(env.microclimate), env.t)
        $name(env::AbstractMicroclimControl, x, args...) = $name(env)
    end)
end

# Accessor functions for range and increment data
for (name, mode) in ((:airtemperature, :range), (:relhumidity, :range), (:windspeed, :range),
                     (:soiltemperature, :increment), (:soilwaterpotential, :increment),
                     (:soilwatercontent, :increment))
    eval(quote
        $name(env::AbstractMicroclimate) = env.$name
        $name(env::AbstractMicroclimPoint, t, interp::LinearLayerInterpolators) =
            $name(env::AbstractMicroclimPoint, t, interp.$mode)
        $name(env::AbstractMicroclimPoint, t, height::Number) =
            $name(env, t, $(Symbol(mode, :_interpolator))(height))
        $name(env::AbstractMicroclimInstant) =
            $name(env.microclimate, env.t, env.interp)
        $name(env::AbstractMicroclimInstant, interp::AbstractLayerInterpolator) =
            $name(env.microclimate, env.t, interp)
        $name(env::AbstractMicroclimControl, x, args...) = $name(env)
        $name(env::AbstractMicroclimPoint, t, interp::LinearLayerInterpolator) =
            interp_layer($name(env), t, interp)
    end)
end

# Accessors for weighted mean of 8 layer data
for name in (:soiltemperature, :soilwatercontent, :soilwaterpotential)
    eval(quote
        $(Symbol(:mean_, name))(env::AbstractMicroclimPoint, t, x) =
            weightedmean(env, $name(env), t, x)
        $(Symbol(:mean_, name))(env::AbstractMicroclimInstant) =
            weightedmean(env, $name(env.microclimate), env.t, env.interp)
        $(Symbol(:mean_, name))(env::AbstractMicroclimInstant, t::Number) =
            weightedmean(env, $name(env.microclimate), t, env.interp)
        $(Symbol(:mean_, name))(env::AbstractMicroclimInstant, interp::AbstractLayerInterpolator) =
            weightedmean(env, $name(env.microclimate), env.t, interp)
        # No mean for control, just return the control value
        $(Symbol(:mean_, name))(env::AbstractMicroclimControl, args...) = $name(env)
    end)
end


"""
	A global agents control method
"""
module Airspace_Control_module

push!(LOAD_PATH, "../src")

using Constants_module, Stats_module
using Distributions, CAS_module, Aircraft_module

include("check_status.jl")

export Airspace_Controller, Constant_Number_Controller, Constant_Spawn_Rate_Controller
export CAS_Property_Controller, Uniform_CAS, Random_CAS
export Sensor_Property_Controller, Uniform_Sensor, Random_Sensor
export airspace_initialization, airspace_control

abstract type Airspace_Controller end
mutable struct Constant_Number_Controller <: Airspace_Controller
	ac_number::Int64
end
mutable struct Constant_Spawn_Rate_Controller <: Airspace_Controller
	ac_number::Int64 # initial value
	spawn_rate::Float64 # num new ac per time step
end

abstract type CAS_Property_Controller end
mutable struct Uniform_CAS <: CAS_Property_Controller
	the_cas::CAS
end
struct Random_CAS <: CAS_Property_Controller
	the_cas::Void
	function Random_CAS()
		return new(nothing)
	end
end

abstract type Sensor_Property_Controller end
mutable struct Uniform_Sensor <: Sensor_Property_Controller
	the_sensor::Sensor
end
struct Random_Sensor <: Sensor_Property_Controller end

"""
	Initialize airspace
"""
function airspace_initialization(stats::Stats, init_num_ac::Int64, cas_ctrl::CAS_Property_Controller,
	sensor_ctrl::Sensor_Property_Controller, t::Int64)
	All_AC = Array{Aircraft, 1}(0);
	for i = 1 : init_num_ac
		ac = create_ac(cas_ctrl, sensor_ctrl, t)
		push!(All_AC, ac)
	end
	return All_AC
end

"""
	Control the airspace traffic density using constant ac number
"""
# Use airspace_control() after dynamics updated for All_AC
function airspace_control(stats::Stats, All_AC::Array{Aircraft, 1}, as_ctrl::Constant_Number_Controller,
	cas_ctrl::CAS_Property_Controller, sensor_ctrl::Sensor_Property_Controller, t::Int64)
	while length(All_AC) < as_ctrl.ac_number
		ac = create_ac(cas_ctrl, sensor_ctrl, t)
		push!(All_AC, ac)
	end
	for i = 1:length(All_AC)
		All_AC[i].desti = check_desti(All_AC[i])
		if All_AC[i].desti == true
			stats.total_nominal_route_len += All_AC[i].nominal_route_len
			stats.total_route_len += (All_AC[i].actual_route_len + dest_range)
			deleteat!(All_AC, i)
			ac = create_ac(cas_ctrl, sensor_ctrl, t)
			push!(All_AC, ac)
		end
	end
	return All_AC
end

"""
	Control the airspace traffic density using constant spawn rate
	Analogying Poisson process
"""
function airspace_control(stats::Stats, All_AC::Array{Aircraft, 1}, as_ctrl::Constant_Spawn_Rate_Controller,
	cas_ctrl::CAS_Property_Controller, sensor_ctrl::Sensor_Property_Controller, t::Int64)
	# Correction weight control
	# if correction_type == :closest && as_ctrl.spawn_rate > 5900./3600
	# 	policy.env.correction_weight = correction_weight
	# elseif correction_type == :sector && as_ctrl.spawn_rate > 7900./3600
	# 	policy.env.correction_weight = correction_weight
	# end

	# deactivate unqualified ac
	i = 1
	while i <= length(All_AC)
		All_AC[i].desti = check_desti(All_AC[i])
		if All_AC[i].desti == true
			stats.total_nominal_route_len += All_AC[i].nominal_route_len
			stats.total_route_len += (All_AC[i].actual_route_len + dest_range)
			deleteat!(All_AC, i)
		else
			i += 1
		end
	end
	# add new ac following Poisson process
	num_new_ac = rand(Poisson(as_ctrl.spawn_rate))
	for i = 1 : num_new_ac
		ac = create_ac(cas_ctrl, sensor_ctrl, t)
		# while check_colli(ac, All_AC) == true
		# 	ac = create_ac(cas_ctrl, sensor_ctrl, t)
		# end
		push!(All_AC, ac)
	end
	return All_AC
end

##################################################################################
"""
	Helpers:
"""
function init_ac_helper()
	id = Base.Random.uuid1()
	x_0 = rand() * airspace_dim
	y_0 = rand() * airspace_dim
	heading_0 = rand() * 2 * pi - pi
	turn_rate_0 = 0.
	x_dest_0 = rand() * airspace_dim
	y_dest_0 = rand() * airspace_dim
	colli_0 = false
	desti_0 = false
	advisory_0 = :COC # COC
	nominal_route_len = sqrt((x_dest_0 - x_0)^2 + (y_dest_0 - y_0)^2)
	actual_route_len = 0.
	t_to_nominal_t_0 = 0.
	d_to_nominal_d_0 = 1.
	cumulative_num_ac_encountered = 0
	num_ac_encountered_arr = Int64[]
	dist_to_dest = nominal_route_len
	return (id, x_0, y_0, heading_0, turn_rate_0, x_dest_0, y_dest_0,
		colli_0, desti_0, advisory_0, nominal_route_len, actual_route_len,
		t_to_nominal_t_0, d_to_nominal_d_0, cumulative_num_ac_encountered, num_ac_encountered_arr, dist_to_dest)
end

function init_random_cas()
	num_cas = length(subtypes(CAS))
	return random_cas = CAS_DICT[subtypes(CAS)[rand(1:num_cas)]]
end

function init_random_sensor()
	num_sensor = length(subtypes(Sensor))
	return random_sensor = subtypes(Sensor)[rand(1:num_sensor)](sensing_range)
end

function create_ac(cas_ctrl::Uniform_CAS, sensor_ctrl::Uniform_Sensor, t::Int64)
	(id, x_0, y_0, heading_0, turn_rate_0, x_dest_0, y_dest_0,
		colli_0, desti_0, advisory_0, nominal_route_len, actual_route_len,
		t_to_nominal_t_0, d_to_nominal_d_0, cumulative_num_ac_encountered,
		num_ac_encountered_arr, dist_to_dest) = init_ac_helper()
	v_0 = rand() * (max_speed - min_speed) + min_speed
	dyn_0 = Dynamics(x_0, y_0, dt, v_0, heading_0, turn_rate_0)
	return Aircraft(id, dyn_0, x_0, y_0, x_dest_0, y_dest_0, colli_0, desti_0,
		advisory_0, cas_ctrl.the_cas, sensor_ctrl.the_sensor, t, nominal_route_len, actual_route_len, dist_to_dest,
		t_to_nominal_t_0, d_to_nominal_d_0, cumulative_num_ac_encountered, num_ac_encountered_arr)
end

function create_ac(::Random_CAS, sensor_ctrl::Uniform_Sensor, t::Int64)
	(id, x_0, y_0, heading_0, turn_rate_0, x_dest_0, y_dest_0,
		colli_0, desti_0, advisory_0, nominal_route_len, actual_route_len,
		t_to_nominal_t_0, d_to_nominal_d_0, cumulative_num_ac_encountered,
		num_ac_encountered_arr, dist_to_dest) = init_ac_helper()
	rand_cas = init_random_cas()
	v_0 = rand() * (max_speed - min_speed) + min_speed
	dyn_0 = Dynamics(x_0, y_0, dt, v_0, heading_0, turn_rate_0)
	return Aircraft(id, dyn_0, x_0, y_0, x_dest_0, y_dest_0, colli_0, desti_0,
		advisory_0, rand_cas, sensor_ctrl.the_sensor, t, nominal_route_len, actual_route_len, dist_to_dest,
		t_to_nominal_t_0, d_to_nominal_d_0, cumulative_num_ac_encountered, num_ac_encountered_arr)
end

function create_ac(cas_ctrl::Uniform_CAS, ::Random_Sensor, t::Int64)
	(id, x_0, y_0, heading_0, turn_rate_0, x_dest_0, y_dest_0,
		colli_0, desti_0, advisory_0, nominal_route_len, actual_route_len,
		t_to_nominal_t_0, d_to_nominal_d_0, cumulative_num_ac_encountered,
		num_ac_encountered_arr, dist_to_dest) = init_ac_helper()
	rand_sensor = init_random_sensor()
	v_0 = rand() * (max_speed - min_speed) + min_speed
	dyn_0 = Dynamics(x_0, y_0, dt, v_0, heading_0, turn_rate_0)
	return Aircraft(id, dyn_0, x_0, y_0, x_dest_0, y_dest_0, colli_0, desti_0,
		in_airspace_0, advisory_0, rand_cas, rand_sensor, t, nominal_route_len, actual_route_len, dist_to_dest,
		t_to_nominal_t_0, d_to_nominal_d_0, cumulative_num_ac_encountered, num_ac_encountered_arr)
end

function create_ac(::Random_CAS, ::Random_Sensor, t::Int64)
	(id, x_0, y_0, heading_0, turn_rate_0, x_dest_0, y_dest_0,
		colli_0, desti_0, advisory_0, nominal_route_len, actual_route_len,
		t_to_nominal_t_0, d_to_nominal_d_0, cumulative_num_ac_encountered,
		num_ac_encountered_arr, dist_to_dest) = init_ac_helper()
	rand_cas = init_random_cas()
	rand_sensor = init_random_sensor()
	v_0 = rand() * (max_speed - min_speed) + min_speed
	dyn_0 = Dynamics(x_0, y_0, dt, v_0, heading_0, turn_rate_0)
	return Aircraft(id, dyn_0, x_0, y_0, x_dest_0, y_dest_0, colli_0, desti_0,
		in_airspace_0, advisory_0, rand_cas, rand_sensor, t, nominal_route_len, actual_route_len, dist_to_dest,
		t_to_nominal_t_0, d_to_nominal_d_0, cumulative_num_ac_encountered, num_ac_encountered_arr)
end

end # module


using DataFrames
using CSV

#Function that reads in counts matrix
function ReadCountsMatrix(filename)
    matrix = CSV.read(filename, header = true, delim = '\t')
    return matrix
end

#Function that reads in gene-isoform table
function ReadGeneIsoformRelationships(filename)
    relationships = CSV.read(filename, header = false, delim = ' ')
    return relationships
end

#Function that checks gene-isoform table matches counts matrix
function CheckGenesMatchIsoforms(countsMatrix, GeneIsoformRelationship)
    arr1 = sort(countsMatrix[:Column1])
    arr2 = sort(GeneIsoformRelationship[:Column2])
    if arr1 == arr2
        return true
    end
    return false
end

#Function that adds gene information to isoform table
function AddGeneInfo(countsMatrix, GeneIsoformRelationship)

    #Call the columns something sensible
    rename!(countsMatrix, :Column1 => :Isoforms)
    rename!(GeneIsoformRelationship, :Column1 => :Genes)
    rename!(GeneIsoformRelationship, :Column2 => :Isoforms)

    return join(GeneIsoformRelationship, countsMatrix, on = :Isoforms)
end

#Function that converts dataframe to dict, removes unexpressed isoforms and genes,
#and returns a dict of how many isoforms are expressed per gene.
function KeepGenesWithNIsoforms(countsMatrix, N)

    #Check N is set to a sensible value
    if N<1
        throw(ArgumentError("You can't simulate the expression of less than one
        isoform!\n"))
    end

    if N>4
        println("You have asked to simulate more than 4 isoforms. Not many genes
        have more than 4 isoforms, so this might violate the assumptions of some
         of iPredict's statistical models.")
    end

    #sort countsMatrix. The rationale behind this is that in theory we can make
    #this closer to O(n) instead of O(n^2)
    sort!(countsMatrix, (:Genes))

    #Set first gene
    gene = ""

    #Make empty dicts
    local dict_to_add = Dict{String, Array{Float64, 1}}()
    local output_dict = Dict{String, Any}()
    local unexpr_dict = Dict{String, Int64}()
    local NumIsos = 0
    local MaxNumIsos = 0
    local UnexprIsos = 0

    #Iterate once over countsMatrix
    for i in 1:nrow(countsMatrix)

        row = countsMatrix[i,:]

        #Check gene is still the same
        if row.Genes == gene

            #if N > numisos or N < numalreadyexpressed, continue
            if N > NumIsos
                continue
            end
            if N < length(dict_to_add)
                continue
            end

            #See whether row passes cutoff
            cutoff = sum(row[3:ncol(countsMatrix)].>=5)

            if cutoff >= 2
                #Add isoform to dictionary
                dict_to_add[row.Isoforms] = Vector(row[3:ncol(countsMatrix)])
            end

        else #new gene

            if length(dict_to_add) == N #See whether the previous gene passes
                output_dict[gene] = dict_to_add
                unexpr_dict[gene] = NumIsos - N

                if NumIsos > MaxNumIsos
                    MaxNumIsos = NumIsos
                end
            end

            #Reset gene, NumExpressed and dict_to_add
            gene = countsMatrix.Genes[i]
            dict_to_add = Dict{String, Array{Float64, 1}}()

            #Find NumIsos
            NumIsos = 0
            j = deepcopy(i)
            while countsMatrix.Genes[j] == gene
                j += 1
                NumIsos += 1
                #Make sure we don't go beyond the end of the matrix
                if j > nrow(countsMatrix)
                    break
                end
            end

            if NumIsos < N
                continue
            end

            #See whether row passes cutoff
            cutoff = sum(row[3:ncol(countsMatrix)].>=5)

            if cutoff >= 2
                #Add isoform to dictionary
                dict_to_add[row.Isoforms] = Vector(row[3:ncol(countsMatrix)])
            end
        end

        #Catch last gene in table
        if i == nrow(countsMatrix)
            if length(dict_to_add) == N
                output_dict[gene] = dict_to_add
            end
        end
    end

    #Return dictionary
    return (output_dict, unexpr_dict, MaxNumIsos)

end

#Function that converts expressed isoforms for a single gene to a dict. Used by
#oneGenePrediction()
function ConvertExpressedIsosToDict(countsMatrix)

    #Make empty dicts
    local dict_to_add = Dict{String, Array{Float64, 1}}()
    local output_dict = Dict{String, Any}()

    for row in eachrow(countsMatrix)

        #See whether row passes cutoff
        cutoff = sum(row[3:ncol(countsMatrix)].>=5)

        if cutoff >= 2
            #Add isoform to dictionary
            dict_to_add[row.Isoforms] = Vector(row[3:ncol(countsMatrix)])
        end
    end

    #Add gene info and return dictionary
    output_dict[countsMatrix.Genes[1]] = dict_to_add
    return output_dict
end


## extend RegionTrees

"""
    getnthcell(c::Cell, n)

Get the cell at the n-th level (at the depth `n`) from the parents of a cell.
If `n` is larger than the depth of `c`, return nothing
(it means, that the asked depth is below the current level).
Also return nothing, if the asked level is below the minimum level, currently set to 1.
"""
function getnthcell(c::Cell, n)
    #n < 1 && throw(BoundsError(c, n))
    n < 1 && return nothing
    maxd = c.data.depth
    n == maxd && return c
    while true
        c = parent(c)
        c === nothing && return nothing
        c.data.depth == n && return c
    end
    return nothing
end

"""
    struct RANSACCloud{A<:AbstractArray, B<:AbstractArray, C<:AbstractArray}

A struct to wrap a point cloud. Stores the vertices, the normals,
    the octree of the vertices,
    the subsets (as an array of arrays of vertice indexes),
    an array to indicate if a point is part of an already extracted primitive,
    the size of the point cloud (number of vertices),
    the weight of each octree level (populated during the construction of the octree,
    by copying the given element),
    an array to store the sum of the score of already
    extracted primitives for each octree level.
"""
struct RANSACCloud{A<:AbstractArray, B<:AbstractArray, C<:AbstractArray}
    vertices::A
    normals::B
    octree::Cell
    subsets::Vector{Vector{Int}}
    isenabled::BitArray{1}
    size::Int
    levelweight::C
    levelscore::C
    #=
    # Arguments
    - `vertices::AbstractArray`: an array of vertices.
    - `normals::AbstractArray`: an array of surface normals.
    - `subsets::Vector{Vector{Int}}`: an array consisting of index-arrays,
        that describe which point is in which subset.
    - `isenabled::BitArray{1}`: an array indicating if the given point is enabled
        (`true`) or already has been extracted (`false`).
    - `size::Int`: number of points in the cloud.
    - `levelweight::Vector{Float}`: the weight of every octree level.
    - `levelscore::Vector{Float}`: for every octree level,
        store the sum of the scores of primitives that have been extracted.
    =#
end

Base.show(io::IO, x::RANSACCloud{A,B,C}) where {A,B,C} =
    print(io, "RANSACCloud of size $(x.size)")

Base.show(io::IO, ::MIME"text/plain", x::RANSACCloud{A,B,C}) where {A,B,C} =
    print(io, "RANSACCloud of size $(x.size) & $(length(x.subsets)) subsets")

"""
    nomodRANSACCloud(vertices, normals, subsets)

Construct a `RANSACCloud` without touching the vertices, normals and subsets.
Other fields are computed.

# Arguments
- `vertices`: an array of vertices.
- `normals`: an array of surface normals.
- `subsets::Vector{Vector{Int}}`: a list of indexes for each subset.
"""
function nomodRANSACCloud(vertices, normals, subsets)
    octree = buildoctree(vertices)
    octree_d = octreedepth(octree)
    s = size(vertices, 1)
    levelscore = zeros(eltype(eltype(vertices)), (octree_d,))
    levelweight = fill!(similar(levelscore), 1/octree_d)
    RANSACCloud(vertices, normals, octree, subsets, trues(s), s, levelscore, levelweight)
end

"""
    RANSACCloud(vertices, normals, subsets; force_eltype::Union{Nothing,DataType}=nothing)

Construct a `RANSACCloud` with the given `subsets`.
Vertices and normals are converted to array of `SVector`s.

# Arguments
- `vertices`: an array of vertices.
- `normals`: an array of surface normals.
- `subsets::Vector{Vector{Int}}`: a list of indexes for each subset.
- `force_eltype::Union{Nothing,DataType}=nothing`:
    an element type can be forced for the normals and vertices
    (in practice `Float64` or `Float32`). If not specified,
    the element type of the passed `vertices` and `normals` will be used.
"""
function RANSACCloud(vertices, normals, subsets; force_eltype::Union{Nothing,DataType}=nothing)
    @assert size(vertices) == size(normals) "Every point must have a normal."
    v_type = force_eltype === nothing ? eltype(eltype(vertices)) : force_eltype
    n_type = force_eltype === nothing ? eltype(eltype(normals)) : force_eltype
    vc = [SVector{3,v_type}(v) for v in vertices]
    nc = [SVector{3,n_type}(v) for v in normals]
    nomodRANSACCloud(vc, nc, subsets)
end

"""
    RANSACCloud(vertices, normals, numofsubsets::Int)

Construct a `RANSACCloud` with `numofsubsets` number of random subsets.
Vertices and normals are converted to array of `SVector`s.

# Arguments
- `vertices`: an array of vertices.
- `normals`: an array of surface normals.
- `numofsubsets::Int`: number of subsets.
- `force_eltype::Union{Nothing,DataType}=nothing`:
    an element type can be forced for the normals and vertices
    (in practice `Float64` or `Float32`). If not specified,
    the element type of the passed `vertices` and `normals` will be used.
"""
function RANSACCloud(vertices, normals, numofsubsets::Int; force_eltype::Union{Nothing,DataType}=nothing)
    @assert numofsubsets > 0 "At least 1 subset please!"
    # subset length
    function makesubset(l, n)
        ssl = fld(l, n)
        alls = randperm(l)
        subsets = [ alls[(i-1)*ssl+1:i*ssl] for i in 1:n-1]
        push!(subsets, alls[(n-1)*ssl+1:end])
        subsets
    end
    subs = makesubset(length(vertices), numofsubsets)
    return RANSACCloud(vertices, normals, subs; force_eltype=force_eltype)
end

"""
    struct OctreeNode{B<:AbstractArray}

Store the data attached to one cell of the octree.
Stores the indices of the points that are inside of the cell,
and its own depth (depth of the root is 1).
"""
struct OctreeNode{B<:AbstractArray, I<:Integer}
    incellpoints::B
    depth::I
end

Base.show(io::IO, x::OctreeNode{A}) where {A} =
    print(io, "OctreeNode: $(length(x.incellpoints)) ps, $(x.depth) d")

Base.show(io::IO, ::MIME"text/plain", x::OctreeNode{A}) where {A} =
    print(io, """OctreeNode{$A}\n with $(length(x.incellpoints)) points, at depth $(x.depth)""")

struct OctreeRefinery{A<:AbstractArray} <: AbstractRefinery
    count::Int
    vertices::A
end

function needs_refinement(r::OctreeRefinery, cell)
    length(cell.data.incellpoints) > r.count
end

function refine_data(r::OctreeRefinery, cell::Cell, indices)
    # new boundary
    boundary = child_boundary(cell, indices)
    # points that are in the current cell
    points = @view r.vertices[cell.data.incellpoints]
    # points that are in the new cell
    bolarr = map(x -> iswithinrectangle(boundary, x), points)
    # new cell inherits ++ of the depth
    d = cell.data.depth + 1
    OctreeNode(cell.data.incellpoints[bolarr], d)
end

"""
    iswithinrectangle(rect::HyperRectangle, p)

Decide if p (3 dimensional point) is within rect.

Works only with those prisms which axes are aligned with the coordinate frame.
"Bottom/left" gives `false` and "top/right" gives `true`.
"""
function iswithinrectangle(rect::HyperRectangle, p)
    vs = vertices(rect)
    vmin = vs[1,1,1]
    vmax = vs[2,2,2]
    for i in 1:3
        vmin[i] < p[i] || return false
        vmax[i] >= p[i] || return false
    end
    return true
end

function updatelevelweight(pc, x = 9//10)
    P = pc.levelweight
    σ = pc.levelscore
    w = sum( σ[i]/P[i] for i in eachindex(P) )
    for i in eachindex(P)
        pc.levelweight[i] = x*σ[i]/(w*P[i]) + (1-x)/length(P)
    end
end

"""
    octreedepth(pc::RANSACCloud)

Compute the depth of the octree stored in the cloud.
"""
function octreedepth(pc::RANSACCloud)
    return octreedepth(pc.octree)
end

"""
    octreedepth(cell::Cell)

Iterate through the leafes and search the maximum depth (`cell.data.depth`).
"""
function octreedepth(cell::Cell)
    maxdepth = cell.data.depth
    alll = allleaves(cell)
    for a in alll
        if a.data.depth > maxdepth
            maxdepth = a.data.depth
        end
    end
    return maxdepth
end

"""
    buildoctree(vertices)

Build octree of the vertices.
"""
function buildoctree(vertices)
    l = size(vertices, 1)
    minV, maxV = findAABB(vertices)
    root=Cell(SVector{3}(minV), SVector{3}(maxV), OctreeNode(collect(1:l), 1))
    r = OctreeRefinery(8, vertices)
    adaptivesampling!(root, r)
    return root
end


"""
    add_obj_fun!(ata_model::Union{MaximinModel,MinimaxModel}; kwargs...)

# Description

It adds the objective function as specified in the `settings_file`. It requires the [`start_ata`](#ATA.start_ata) build step.

Computes the IIFs at predefined ability points using the item parameter estimates in the item bank.
For each pair of item, and ability point \$(i, θ_k)\$, an IIF\$(θ_k)_{i}\$ is computed.

# Arguments

- **`ata_model::Union{MaximinModel, MinimaxModel}`** : Required.
"""
function add_obj_fun!(ata_model::Union{MaximinModel,MinimaxModel}; kwargs...)

    try
        compute_estimated_iif!(ata_model)
        success!(ata_model, "Estimated IIFs computed.")

    catch e
        error!(ata_model, string(sprint(showerror, e)))

    end
    return nothing
end


export IntensityMap, fov, imagepixels, pixelsizes, stokes

"""
    $(TYPEDEF)
Image array type. This is an Matrix with a number of internal fields
to describe the field of view, pixel size, and the `pulse` function that
makes the image a continuous quantity.

To use it you just specify the array and the field of view/pulse
``julia
img = IntensityMap(zeros(512,512), 100.0, 100.0)
```
"""
struct IntensityMap{T,S<:AbstractMatrix, F, K<:Pulse} <: AbstractIntensityMap{T,S}
    im::S
    fovx::F
    fovy::F
    psizex::F
    psizey::F
    pulse::K
end

function IntensityMap(im, fovx, fovy, pulse=DeltaPulse())
    ny,nx = size(im)
    psizex=fovx/max(nx-1,1)
    psizey=fovy/max(ny-1,1)
    F = promote_type(typeof(fovx), typeof(fovy))
    return IntensityMap{eltype(im), typeof(im), F, typeof(pulse)}(im,
                       convert(typeof(psizex),fovx),
                       convert(typeof(psizey),fovy),
                       psizex,
                       psizey,
                       pulse)
end

# IntensityMap will obey the Comrade interface. This is so I can make easy models
visanalytic(::Type{<:IntensityMap}) = NotAnalytic()
imanalytic(::Type{<:IntensityMap}) = IsAnalytic()
isprimitive(::Type{<:IntensityMap}) = IsPrimitive()

function intensity_point(m::IntensityMap, x, y)
    dx, dy = pixelsizes(m)
    xitr,yitr = imagepixels(m)
    sum = zero(eltype(m))
    #The kernel is written in terms of pixel number so we convert x to it
    @inbounds for (i,xx) in pairs(xitr), (j,yy) in pairs(yitr)
        Δx = (x-xx)/dx
        Δy = (y-yy)/dy
        k = intensity_point(m.pulse, Δx, Δy)
        #println(Δx," ", Δy, " ",   k)
        sum += m[j, i]*k/(dx*dy)
    end
    return sum
end

function IntensityMap(im::Matrix{<:StokesVector}, fovx, fovy, pulse)
    return IntensityMap(StructArray(im), fovx, fovy, pulse)
end

function IntensityMap(im::IntensityMap{<:StokesVector}, fovx, fovy, pulse)
    return IntensityMap(im.im, fovx, fovy, pulse)
end

function stokes(im::IntensityMap{<:StokesVector}, p::Symbol)
    @assert p ∈ propertynames(im.im) "$p is not a valid stokes parameter"
    return IntensityMap(getproperty(im.im, p), im.fovx, im.fovy, im.pulse)
end

function IntensityMap(I::AbstractIntensityMap,
                      Q::AbstractIntensityMap,
                      U::AbstractIntensityMap,
                      V::AbstractIntensityMap
                      )
    @assert I.fovx == Q.fovx == U.fovx == V.fovx "Must have matching fov in RA"
    @assert I.fovy == Q.fovy == U.fovy == V.fovy "Must have matching fov in DEC"
    simg = StructArray{StokesVector{eltype(I)}}((I,Q,U,V))
    return IntensityMap(simg, I.fovx, I.fovy, I.pulse)
end

# function ChainRulesCore.rrule(::Type{<:IntensityMap}, im, fovx, fovy, pulse)
#     y = IntensityMap(im, fovx, fovy, pulse)
#     intensity_pullback(Δy) = (NoTangent(), IntensityMap(Δy.im, fovx, fovy, pulse), Δy.fovx, Δy.fovy, Δy.pulse)
#     return y, intensity_pullback
# end

fov(m::AbstractIntensityMap) = (m.fovx, m.fovy)

"""
    $(SIGNATURES)
Computes the flux of a intensity map
"""
function flux(im::AbstractIntensityMap{T,S}) where {T,S}
    f = sum(im.im)#*(flux(im.pulse))^2
    return f#*prod(pixelsizes(im))
end

"""
    $(SIGNATURES)
Computes the flux of a intensity map
"""
function flux(im::AbstractIntensityMap{T,S}) where {F,T<:StokesVector{F},S}
    I = stokes(im, :I)
    flux(I)
end

# Define the array interface
Base.IndexStyle(::Type{<: IntensityMap{T,S,K}}) where {T,S,K} = Base.IndexStyle(S)
Base.size(im::AbstractIntensityMap) = size(im.im)
Base.@propagate_inbounds Base.getindex(im::AbstractIntensityMap, i::Int) = getindex(im.im, i)
Base.@propagate_inbounds Base.getindex(im::AbstractIntensityMap, I...) = getindex(im.im, I...)
Base.@propagate_inbounds Base.setindex!(im::AbstractIntensityMap, x, i::Int) = setindex!(im.im, x, i)
Base.@propagate_inbounds Base.setindex!(im::IntensityMap, x, i::Int) = setindex!(im.im, x, i)

function Base.similar(im::IntensityMap, ::Type{T}) where{T}
    sim = similar(im.im, T)
    return IntensityMap(sim, im.fovx, im.fovy, im.pulse)
end

#function Base.similar(im::AbstractIntensityMap, ::Type{T}, dims::Dims) where {T}
#    fovx = im.psizex*last(dims)
#    fovy = im.psizey*first(dims)
#    sim = similar(im.im, T, dims)
#    return IntensityMap(sim, fovx, fovy, im.psizex, im.psizey)
#end

#Define the broadcast interface
struct IntensityMapStyle <: Broadcast.AbstractArrayStyle{2} end
IntensityMapStyle(::Val{2}) = IntensityMapStyle()

Base.BroadcastStyle(::Type{<:AbstractIntensityMap}) = IntensityMapStyle()
function Base.similar(bc::Broadcast.Broadcasted{IntensityMapStyle}, ::Type{ElType}) where ElType
    #Scan inputs for IntensityMap
    #print(bc.args)
    Im = _find_sim(bc)
    #fovxs = getproperty.(Im, Ref(:fovx))
    #fovys = getproperty.(Im, Ref(:fovy))
    #@assert all(i->i==first(fovxs), fovxs) "IntensityMap fov must be equal to add"
    #@assert all(i->i==first(fovys), fovys) "IntensityMap fov must be equal to add"
    return IntensityMap(similar(Array{ElType}, axes(bc)), Im.fovx, Im.fovy, Im.pulse)
end

#Finds the first IntensityMap and uses that as the base
#TODO: If multiply IntensityMaps maybe I should make sure they are consistent?
_find_sim(bc::Base.Broadcast.Broadcasted) = _find_sim(bc.args)
_find_sim(args::Tuple) = _find_sim(_find_sim(args[1]), Base.tail(args))
_find_sim(x) = x
_find_sim(::Tuple{}) = nothing
_find_sim(a::AbstractIntensityMap, rest) = a
_find_sim(::Any, rest) = _find_sim(rest)

#Guards to prevent someone from adding two Images with different FOV's
function Base.:+(x::AbstractIntensityMap, y::AbstractIntensityMap)
    @assert fov(x) == fov(y) "IntensityMaps must share same field of view"
    return x .+ y
end

@inline function pixelsizes(im::AbstractIntensityMap)
    ny,nx = size(im)
    return im.fovx/nx, im.fovy/ny
end

@inline function imagepixels(fovx, fovy, nx::Int, ny::Int)
    px = fovx/nx; py = fovy/ny
    return range(-fovx/2+px/2, step=px, length=nx),
           range(-fovy/2+py/2, step=py, length=ny)
end

@inline function imagepixels(im::AbstractIntensityMap)
    ny,nx = size(im)
    return imagepixels(im.fovx, im.fovy, nx, ny)
end


const DEFAULT_JOB_NAME = ""

const DEFAULT_DT_NULL = DateTime(0)
const DEFAULT_DT_NEXT_FIRE = DEFAULT_DT_NULL

"""
    JobConfig([misfire_grace_period])

Job default configuration

# Optional arguments
- misfire_grace_period (defaults to Dates.Second(1)): grace period for which a task can still be fired.
"""
mutable struct JobConfig
    misfire_grace_period

    JobConfig(; misfire_grace_period=Dates.Second(1)) = new(misfire_grace_period)
end

"""
    Job(id, action, trigger, name, priority, dt_created, dt_updated, dt_next_fire, n_triggered, config)

A job is an internal structure which store what action
should be executed when triggered.

It also store several properties such as priority level, number of time a job is triggered, when will next trigger should occur...
"""
mutable struct Job
    id::String
    action::Action
    trigger::AbstractTrigger
    name::String
    priority::Int
    dt_created::DateTime
    dt_updated::DateTime
    #dt_last_fire::DateTime
    dt_next_fire::DateTime
    n_triggered::BigInt
    config::JobConfig
end

"""
    hasnextfire(job)

Returns `true` if a task should be fired in the future.
"""
function hasnextfire(job::Job)
    job.dt_next_fire != DEFAULT_DT_NEXT_FIRE
end


# This file is a part of BAT.jl, licensed under the MIT License (MIT).

vjp_algorithm(d::AbstractDensity) = ForwardDiffAD()
vjp_algorithm(d::WrappedNonBATDensity) = jvp_algorithm(parent(d))
vjp_algorithm(d::AbstractPosteriorDensity) = vjp_algorithm(getlikelihood(d))
vjp_algorithm(d::DensityWithShape) = vjp_algorithm(parent(d))
vjp_algorithm(d::TransformedDensity) = vjp_algorithm(parent(d))
vjp_algorithm(d::RenormalizedDensity) = vjp_algorithm(parent(d))

vjp_algorithm(f::Base.Fix1{F,<:AbstractDensity}) where F = vjp_algorithm(f.x)
vjp_algorithm(f::Negative) = vjp_algorithm(negative(f))

jvp_algorithm(d::AbstractDensity) = ForwardDiffAD()
jvp_algorithm(d::WrappedNonBATDensity) = jvp_algorithm(parent(d))
jvp_algorithm(d::AbstractPosteriorDensity) = jvp_algorithm(getprior(d))
jvp_algorithm(d::DensityWithShape) = jvp_algorithm(parent(d))
jvp_algorithm(d::TransformedDensity) = jvp_algorithm(parent(d))
jvp_algorithm(d::RenormalizedDensity) = jvp_algorithm(parent(d))

jvp_algorithm(f::Base.Fix1{F,<:AbstractDensity}) where F = jvp_algorithm(f.x)
jvp_algorithm(f::Negative) = jvp_algorithm(f.density)

struct DensityWithDiff{JVP<:DifferentiationAlgorithm,D<:AbstractDensity,VJP<:DifferentiationAlgorithm} <: AbstractDensity
    vjpalg::VJP
    density::D
    jvpalg::JVP
end

@inline Base.parent(density::DensityWithDiff) = density.density

vjp_algorithm(density::DensityWithDiff) = density.vjpalg
jvp_algorithm(density::DensityWithDiff) = density.jvpalg

Base.:(|)(vjpalg::DifferentiationAlgorithm, density::AbstractDensity) = DensityWithDiff(vjpalg, density, jvp_algorithm(density))
Base.:(|)(vjpalg::DifferentiationAlgorithm, density::DensityWithDiff) = DensityWithDiff(vjpalg, parent(density), jvp_algorithm(density))

Base.:(|)(density::AbstractDensity, jvpalg::DifferentiationAlgorithm) = DensityWithDiff(vjp_algorithm(density), density, jvpalg)
Base.:(|)(density::DensityWithDiff, jvpalg::DifferentiationAlgorithm) = DensityWithDiff(vjp_algorithm(density), parent(density), jvpalg)

function Base.show(io::IO, density::DensityWithDiff)
    vjpalg = vjp_algorithm(density)
    jvpalg = jvp_algorithm(density)
    parent_vjpalg = vjp_algorithm(parent(density))
    parent_jvpalg = jvp_algorithm(parent(density))
    if vjpalg != parent_vjpalg || jvpalg == parent_jvpalg
        Base.show(io, vjpalg)
        print(io, "|")
    end
    Base.show(io, parent(density))
    if jvpalg != parent_jvpalg
        print(io, "|")
        Base.show(io, jvpalg)
    end
end

var_bounds(density::DensityWithDiff) = var_bounds(density.density)

ValueShapes.varshape(density::DensityWithDiff) = varshape(density.density)

ValueShapes.unshaped(density::DensityWithDiff) = DensityWithDiff(density.vjpalg, unshaped(density.density), density.jvpalg)

@inline DensityInterface.logdensityof(density::DensityWithDiff, v::Any) = logdensityof(density.density, v)
@inline DensityInterface.logdensityof(density::DensityWithDiff) = logdensityof(density.density)

@inline checked_logdensityof(density::DensityWithDiff, v::Any) = checked_logdensityof(density.density, v)

@inline get_initsrc_from_target(target::DensityWithDiff) = get_initsrc_from_target(target.density)

function bat_transform_impl(target::NoDensityTransform, density::DensityWithDiff, algorithm::DensityIdentityTransform)
    (result = density, trafo = IdentityVT(varshape(density)))
end

function bat_transform_impl(target::AbstractDensityTransformTarget, density::DensityWithDiff, algorithm::TransformAlgorithm)
    transformed_density, trafo = bat_transform_impl(target, parent(density), algorithm)
    (result = DensityWithDiff(vjp_algorithm(density), transformed_density, jvp_algorithm(density)), trafo = trafo)
end


#############################################################################
# ECOS.jl
# Wrapper around the ECOS solver https://github.com/ifa-ethz/ecos
# See http://github.com/JuliaOpt/ECOS.jl
#############################################################################
# ECOSSolverInterface.jl
# MathProgBase.jl interface for the ECOS.jl solver wrapper
#############################################################################

import MathProgBase
const MPB = MathProgBase

#############################################################################
# Define the MPB Solver and Model objects
export ECOSSolver
struct ECOSSolver <: MPB.AbstractMathProgSolver
    options
end
ECOSSolver(;kwargs...) = ECOSSolver(kwargs)

mutable struct ECOSMathProgModel <: MPB.AbstractConicModel
    nvar::Int                           # Number of variables
    nconstr::Int                        # Number of rows in the MPB-form A matrix
    nineq::Int                          # Number of inequalities Gx <=_K h
    neq::Int                            # Number of equalities Ax = b
    npos::Int                           # Number of positive orthant cones
    ncones::Int                         # Number of SO cones
    conedims::Vector{Int}               # Dimension of each SO cone
    nexp_cones::Int                     # Number of exponential cones
    G::ECOSMatrix                       # The G matrix (inequalties)
    A::ECOSMatrix                       # The A matrix (equalities)
    c::Vector{Float64}                  # The objective coeffs (always min)
    orig_sense::Symbol                  # Original objective sense
    h::Vector{Float64}                  # RHS for inequality
    b::Vector{Float64}                  # RHS for equality
    # Post-solve
    solve_stat::Symbol
    solve_time::Float64
    obj_val::Float64
    primal_sol::Vector{Float64}
    dual_sol_eq::Vector{Float64}
    dual_sol_ineq::Vector{Float64}
    # Maps x∈K to ECOS duals
    # .._ind maps a column to a row index
    # .._type  is the original cone type of each column
    col_map_ind::Vector{Int}
    col_map_type::Vector{Symbol}
    # Maps b-Ax∈K to ECOS duals
    # .._ind maps a row to an index
    # .._type  is the original cone type of each row
    row_map_ind::Vector{Int}
    row_map_type::Vector{Symbol}
    # To reorder solution to MPB form
    fwd_map::Vector{Int}
    options
end
ECOSMathProgModel(;kwargs...) = ECOSMathProgModel(0,0,0,0,0,0,
                                        Int[],0,
                                        ECOSMatrix(spzeros(0,0)),
                                        ECOSMatrix(spzeros(0,0)),
                                        Float64[], :Min,
                                        Float64[], Float64[],
                                        :NotSolved, 0.0, 0.0,
                                        Float64[],
                                        Float64[], Float64[],
                                        Int[], Symbol[],
                                        Int[], Symbol[], Int[],
                                        kwargs)

#############################################################################
# Begin implementation of the MPB low-level interface
# Implements
# - model
# - optimize!
# - status
# - numvar
# - numconstr
# http://mathprogbasejl.readthedocs.org/en/latest/lowlevel.html

MPB.ConicModel(s::ECOSSolver) = ECOSMathProgModel(;s.options...)
MPB.LinearQuadraticModel(s::ECOSSolver) = MPB.ConicToLPQPBridge(MPB.ConicModel(s))

function MPB.optimize!(m::ECOSMathProgModel)
    ecos_prob_ptr = setup(
        m.nvar, m.nineq, m.neq,
        m.npos, m.ncones, m.conedims, m.nexp_cones,
        m.G, m.A,
        m.c,
        m.h, m.b; m.options...)
    # Note: ECOS modifies problem data in setup() and restores it on cleanup()

    t = time()
    flag = solve(ecos_prob_ptr)
    m.solve_time = time() - t
    if flag == ECOS_OPTIMAL
        m.solve_stat = :Optimal
    elseif flag == ECOS_PINF
        m.solve_stat = :Infeasible
    elseif flag == ECOS_DINF  # Dual infeasible = primal unbounded, probably
        m.solve_stat = :Unbounded
    elseif flag == ECOS_MAXIT
        m.solve_stat = :UserLimit
    elseif flag == ECOS_OPTIMAL + ECOS_INACC_OFFSET
        m.solve_stat = :Suboptimal
    else
        m.solve_stat = :Error
    end
    # Extract solution
    ecos_prob = unsafe_wrap(Array, ecos_prob_ptr, 1)[1]
    m.primal_sol = unsafe_wrap(Array,ecos_prob.x, m.nvar)[:]
    m.dual_sol_eq   = unsafe_wrap(Array,ecos_prob.y, m.neq)[:]
    m.dual_sol_ineq = unsafe_wrap(Array,ecos_prob.z, m.nineq)[:]
    cleanup(ecos_prob_ptr, 0)
    # Compute this value after cleanup with restored c.
    m.obj_val = dot(m.c, m.primal_sol) * (m.orig_sense == :Max ? -1 : +1)
end

MPB.status(m::ECOSMathProgModel) = m.solve_stat
MPB.getobjval(m::ECOSMathProgModel) = m.obj_val
MPB.getsolution(m::ECOSMathProgModel) = m.primal_sol[m.fwd_map]

MPB.numvar(m::ECOSMathProgModel) = m.nvar
MPB.numconstr(m::ECOSMathProgModel) = m.nineq + m.neq

#############################################################################
# Begin implementation of the MPB conic interface
# Implements
# - supportedcones
# - loadconicproblem!
# - getconicdual
# http://mathprogbasejl.readthedocs.org/en/latest/conic.html

MPB.supportedcones(m::ECOSSolver) = [:Free,:Zero,:NonNeg,:NonPos,:SOC,:ExpPrimal]

function MPB.loadproblem!(m::ECOSMathProgModel, c, A, b, constr_cones, var_cones)
    if size(A,2) == 0
        Base.warn("Input matrix has no columns. ECOS is known to crash in this corner case.")
    end
    m.nconstr = size(A,1)
    # If A is sparse, we should use an appropriate "zeros"
    zeromat = isa(A,SparseMatrixCSC) ? spzeros : zeros

    # We don't support SOCRotated, SDP, or ExpDual
    bad_cones = [:SOCRotated, :SDP, :ExpDual]
    for cone_vars in constr_cones
        cone_vars[1] in bad_cones && error("Cone type $(cone_vars[1]) not supported")
    end
    for cone_vars in var_cones
        cone_vars[1] in bad_cones && error("Cone type $(cone_vars[1]) not supported")
    end

    # MathProgBase form             ECOS form
    # min  c'x                      min  c'x
    #  st b-Ax ∈ K_1                 st   Ax = b
    #        x ∈ K_2                    h-Gx ∈ K
    #
    # Mapping:
    # * For the constaints (K_1)
    #   * If :Zero cone, then treat as equality constraint in ECOS form
    #   * Otherewise trivially maps to h-Gx in ECOS form
    # * For the variables (K_2)
    #   * If :Free, do nothing
    #   * If :Zero, put in as equality constraint
    #   * If rest, stick in h-Gx
    #
    # Approach:
    # 0. Figure out a mapping between MPB and ECOS
    # 1. Map non-SOC/Exp variables to ECOS form
    # 2. Map non-SOC/Exp constraints to ECOS form
    # 3. Map SOC variables and constraints to ECOS form
    # 4. Map Exp variables and constraints to ECOS form

    # Allocate space for the ECOS variables
    num_vars = length(c)
    fwd_map = Array{Int}(undef, num_vars)  # Will be used for SOCs
    rev_map = Array{Int}(undef, num_vars)  # Need to restore sol. vec.
    idxcone = Array{Symbol}(undef, num_vars)  # We'll use this for non-SOC/Exp
    m.col_map_type = Array{Symbol}(undef, num_vars) # In MPB indices

    # Now build the mapping between MPB variables and ECOS variables
    pos = 1
    for (cone, idxs) in var_cones
        for i in idxs
            fwd_map[i]   = pos   # fwd_map = MPB idx -> ECOS idx
            rev_map[pos] = i     # rev_map = ECOS idx -> MPB idx
            idxcone[pos] = cone
            m.col_map_type[i] = cone
            pos += 1
        end
    end

    # Rearrange data into the internal ordering, make copy
    ecos_c = c[rev_map]
    ecos_A = zeromat(0,num_vars)
    ecos_b = Float64[]

    # Mapping for duals
    m.col_map_ind = zeros(Int,num_vars)
    m.row_map_ind = zeros(Int, length(b))
    m.row_map_type  = Array{Symbol}(undef, length(b))
    function update_row_map(cone_type, cur_ind)
        for (cone,idxs) in constr_cones
            if cone == cone_type
                for idx in idxs
                    m.row_map_ind[idx] = cur_ind
                    m.row_map_type[idx] = cone_type
                    cur_ind += 1
                end
            end
        end
        cur_ind
    end

    ###################################################################
    # PHASE ONE  -  MAP x ∈ K_2 to ECOS form, except SOC/Exp

    # If a variable is in :Zero cone, fix at 0 with equality constraint.
    for j = 1:num_vars
        idxcone[j] != :Zero && continue

        new_row    = zeromat(1,num_vars)
        new_row[j] = 1.0
        ecos_A     = vcat(ecos_A, new_row)
        ecos_b     = vcat(ecos_b, 0.0)
        m.col_map_ind[rev_map[j]] = length(ecos_b)
    end

    # G matrix:
    # * 1 row ∀ :NonNeg & :NonPos cones
    # * 1 row ∀ variable in SOC/Exp cones
    # We will only handle the first case here, the rest in phase 3.
    num_G_row_negpos = 0
    for j = 1:num_vars
        !(idxcone[j] == :NonNeg || idxcone[j] == :NonPos) && continue
        num_G_row_negpos += 1
        m.col_map_ind[rev_map[j]] = num_G_row_negpos
    end
    ecos_G = zeromat(num_G_row_negpos,num_vars)
    ecos_h =   zeros(num_G_row_negpos)

    # Handle the :NonNeg, :NonPos cases
    num_pos_orth = 0
    G_row = 1
    for j = 1:num_vars
        if idxcone[j] == :NonNeg
            ecos_G[G_row,j] = -1.0
            G_row += 1
            num_pos_orth += 1
        elseif idxcone[j] == :NonPos
            ecos_G[G_row,j] = +1.0
            G_row += 1
            num_pos_orth += 1
        end
    end
    @assert G_row == num_pos_orth + 1

    ###################################################################
    # PHASE TWO  -  MAP b-Ax ∈ K_1 to ECOS form, except SOC/Exp

    # Zero rows for Ax=b, NonNegPos rows to append to G,h
     eq_rows = Int[]
    pos_rows = Int[]
    neg_rows = Int[]
    for (cone,idxs) in constr_cones
        cone == :Free && error("Free cone constraints not handled")
        (cone == :SOC || cone == :ExpPrimal) && continue  # Handle later
        idxset = vec(collect(idxs))
        if cone == :Zero
            append!(eq_rows, idxset)
            continue
        end
        cone == :NonNeg && append!(pos_rows, idxset)
        cone == :NonPos && append!(neg_rows, idxset)
    end
    # Update mappings - eq, nonneg, nonpos
     eq_cur_ind = length(ecos_b) + 1
     eq_cur_ind = update_row_map(:Zero, eq_cur_ind)
    neq_cur_ind = length(ecos_h) + 1
    neq_cur_ind = update_row_map(:NonNeg, neq_cur_ind)
    neq_cur_ind = update_row_map(:NonPos, neq_cur_ind)
    # Equality constraints / Zero cones
    ecos_A = vcat(ecos_A,  A[eq_rows,rev_map])
    ecos_b = vcat(ecos_b,  b[eq_rows])
    ecos_G = vcat(ecos_G,  A[pos_rows,rev_map])
    ecos_h = vcat(ecos_h,  b[pos_rows])
    ecos_G = vcat(ecos_G, -A[neg_rows,rev_map])  # b-a'x <= 0 - flip sign,
    ecos_h = vcat(ecos_h, -b[neg_rows])          # then maps to a row in h-Gx
    G_row        += length(pos_rows) + length(neg_rows)
    num_pos_orth += length(pos_rows) + length(neg_rows)

    ###################################################################
    # PHASE THREE  -  MAP SOC variables and constraints to ECOS form

    # Handle the SOC variable cones
    # MPB  form: vector of var (y,x) is in the SOC ||x|| <= y
    # ECOS form: h - Gx ∈ Q  -->  0 - Ix ∈ Q
    num_G_row_soc = 0
    num_G_row_exp = 0
    for j = 1:num_vars
        if idxcone[j] == :SOC
            num_G_row_soc += 1
        elseif idxcone[j] == :ExpPrimal
            num_G_row_exp += 1
        end
    end
    ecos_G = vcat(ecos_G, zeromat(num_G_row_soc,num_vars))
    ecos_h = vcat(ecos_h,   zeros(num_G_row_soc))

    num_SOC_cones = 0
    SOC_conedims  = Int[]
    for (cone, idxs) in var_cones
        cone != :SOC && continue
        # Found a new SOC
        num_SOC_cones += 1
        push!(SOC_conedims, length(idxs))
        # Add the entries (carrying on from pos. orthant)
        for j in idxs
            ecos_G[G_row,fwd_map[j]] = -1.0
            m.col_map_ind[j] = G_row
            G_row += 1
        end
    end
    @assert G_row == num_pos_orth + num_G_row_soc + 1

    # Handle the SOC constraint cones
    # Collect all the rows we'll be appending to G,h
    all_rows = Int[]
    for (cone,idxs) in constr_cones
        if cone == :SOC
            num_SOC_cones += 1
            push!(SOC_conedims, length(idxs))
            idx_list   = collect(idxs)
            all_rows   = vcat(all_rows,   idx_list)
        end
    end
    neq_cur_ind = update_row_map(:SOC, neq_cur_ind)
    ecos_G = vcat(ecos_G, A[all_rows,rev_map])
    ecos_h = vcat(ecos_h, b[all_rows])

    ###################################################################
    # PHASE FOUR  -  MAP Exp variables and constraints to ECOS form

    # Note: The MPB definition of the ExpPrimal cone and the ECOS
    # defintions swap the 2nd and 3rd components, so we have to
    # manually permute the indices.

    num_exp_cones = 0
    ecos_G = vcat(ecos_G, zeromat(num_G_row_exp,num_vars))
    ecos_h = vcat(ecos_h,   zeros(num_G_row_exp))

    G_row += sum(all_rows)
    for (cone, idxs) in var_cones
        cone != :ExpPrimal && continue
        # Found a new Exp
        num_exp_cones += 1
        @assert length(idxs) == 3
        # Add the entries in permuted order
        ecos_G[G_row,fwd_map[idxs[1]]] = -1.0
        ecos_G[G_row+2,fwd_map[idxs[2]]] = -1.0
        ecos_G[G_row+1,fwd_map[idxs[3]]] = -1.0
        m.col_map_ind[idxs[1]] = G_row
        m.col_map_ind[idxs[2]] = G_row+2
        m.col_map_ind[idxs[3]] = G_row+1
        G_row += 3
    end
    @assert G_row == num_pos_orth + num_G_row_soc + sum(all_rows) + num_G_row_exp + 1

    exp_rows = Int[]
    for (cone,idxs) in constr_cones
        if cone == :ExpPrimal
            num_exp_cones += 1
            @assert length(idxs) == 3
            append!(exp_rows, [idxs[1],idxs[3],idxs[2]])

            # override update_row_map to handle the permutation
            m.row_map_type[idxs] .= :ExpPrimal
            m.row_map_ind[idxs[1]] = neq_cur_ind
            m.row_map_ind[idxs[3]] = neq_cur_ind+1
            m.row_map_ind[idxs[2]] = neq_cur_ind+2
            neq_cur_ind += 3
        end
    end
    ecos_G = vcat(ecos_G, A[exp_rows,rev_map])
    ecos_h = vcat(ecos_h, b[exp_rows])

    ###################################################################
    # Store in the ECOS structure
    m.nvar          = num_vars          # Num variable
    m.nineq         = size(ecos_G,1)    # Num inequality constraints
    m.neq           = length(ecos_b)    # Num equality constraints
    m.npos          = num_pos_orth      # Num ineq. constr. in +ve orthant
    m.ncones        = num_SOC_cones     # Num second-order cones
    m.conedims      = SOC_conedims      # Num contr. in each SOC
    m.nexp_cones    = num_exp_cones     # Num exponential cones
    m.G             = ECOSMatrix(ecos_G)
    m.A             = ECOSMatrix(ecos_A)
    m.c             = ecos_c
    m.orig_sense    = :Min
    m.h             = ecos_h
    m.b             = ecos_b
    m.fwd_map       = fwd_map           # Used to return solution
end

function MPB.getdual(m::ECOSMathProgModel)
    duals = zeros(length(m.row_map_ind))
    for (mpb_row,ecos_row) in enumerate(m.row_map_ind)
        cone = m.row_map_type[mpb_row]
        if cone == :Zero
            # This MPB constraint ended up in ECOS equality block
            duals[mpb_row] = m.dual_sol_eq[ecos_row]
        else
            # Ended up in ECOS inequality block
            if cone == :NonPos
                duals[mpb_row] = -m.dual_sol_ineq[ecos_row]
            else
                duals[mpb_row] = m.dual_sol_ineq[ecos_row]
            end
        end
    end
    return duals
end

function MPB.getvardual(m::ECOSMathProgModel)
    duals = zeros(length(m.col_map_ind))
    for (mpb_col,ecos_row) in enumerate(m.col_map_ind)
        cone = m.col_map_type[mpb_col]
        cone == :Free && continue # dual is zero
        if cone == :Zero
            # This MPB zero var ended up in ECOS equality block
            duals[mpb_col] = m.dual_sol_eq[ecos_row]
        else
            # Ended up in ECOS inequality block
            if cone == :NonPos
                duals[mpb_col] = -m.dual_sol_ineq[ecos_row]
            else
                duals[mpb_col] = m.dual_sol_ineq[ecos_row]
            end
        end
    end
    return duals
end

MPB.getsolvetime(m::ECOSMathProgModel) = m.solve_time


# clone of https://github.com/analogdevicesinc/libad9361-iio/blob/master/ad9361_baseband_auto_rate.c

export
    ad9361_get_trx_fir_enable,
    ad9361_set_trx_fir_enable,
    ad9361_baseband_auto_rate
;

# size of the string buffer
const FIR_BUF_SIZE = 8192;

# default filters
const fir_128_4 = Int16[
    -15,-27,-23,-6,17,33,31,9,-23,-47,-45,-13,34,69,67,21,-49,-102,-99,-32,69,146,143,48,-96,-204,-200,-69,129,278,275,97,-170,
    -372,-371,-135,222,494,497,187,-288,-654,-665,-258,376,875,902,363,-500,-1201,-1265,-530,699,1748,1906,845,-1089,-2922,-3424,
    -1697,2326,7714,12821,15921,15921,12821,7714,2326,-1697,-3424,-2922,-1089,845,1906,1748,699,-530,-1265,-1201,-500,363,902,875,
    376,-258,-665,-654,-288,187,497,494,222,-135,-371,-372,-170,97,275,278,129,-69,-200,-204,-96,48,143,146,69,-32,-99,-102,-49,21,
    67,69,34,-13,-45,-47,-23,9,31,33,17,-6,-23,-27,-15
];
const fir_128_2 = Int16[
    -0,0,1,-0,-2,0,3,-0,-5,0,8,-0,-11,0,17,-0,-24,0,33,-0,-45,0,61,-0,-80,0,104,-0,-134,0,169,-0,
	-213,0,264,-0,-327,0,401,-0,-489,0,595,-0,-724,0,880,-0,-1075,0,1323,-0,-1652,0,2114,-0,-2819,0,4056,-0,-6883,0,20837,32767,
	20837,0,-6883,-0,4056,0,-2819,-0,2114,0,-1652,-0,1323,0,-1075,-0,880,0,-724,-0,595,0,-489,-0,401,0,-327,-0,264,0,-213,-0,
	169,0,-134,-0,104,0,-80,-0,61,0,-45,-0,33,0,-24,-0,17,0,-11,-0,8,0,-5,-0,3,0,-2,-0,1,0,-0, 0
];
const fir_96_2 = Int16[
    -4,0,8,-0,-14,0,23,-0,-36,0,52,-0,-75,0,104,-0,-140,0,186,-0,-243,0,314,-0,-400,0,505,-0,-634,0,793,-0,
	-993,0,1247,-0,-1585,0,2056,-0,-2773,0,4022,-0,-6862,0,20830,32767,20830,0,-6862,-0,4022,0,-2773,-0,2056,0,-1585,-0,1247,0,-993,-0,
	793,0,-634,-0,505,0,-400,-0,314,0,-243,-0,186,0,-140,-0,104,0,-75,-0,52,0,-36,-0,23,0,-14,-0,8,0,-4,0
];
const fir_64_2 = Int16[
    -58,0,83,-0,-127,0,185,-0,-262,0,361,-0,-488,0,648,-0,-853,0,1117,-0,-1466,0,1954,-0,-2689,0,3960,-0,-6825,0,20818,32767,
	20818,0,-6825,-0,3960,0,-2689,-0,1954,0,-1466,-0,1117,0,-853,-0,648,0,-488,-0,361,0,-262,-0,185,0,-127,-0,83,0,-58,0
];

# returns the return code of the C function
function ad9361_set_trx_fir_enable(device::Ptr{iio_device}, enable::Bool)
    ret, channel, attribute = C_iio_device_identify_filename(device, "in_out_voltage_filter_fir_en");
    if ret < 0
        return ret;
    end
    return C_iio_channel_attr_write_bool(channel, attribute, enable);
end

# returns the return code of the C function and the value (needs to be discarded if the code < 0)
function ad9361_get_trx_fir_enable(device::Ptr{iio_device})
    ret, channel, attribute = C_iio_device_identify_filename(device, "in_out_voltage_filter_fir_en");
    if ret < 0
        return ret, false;
    end
    return C_iio_channel_attr_read_bool(channel, attribute);
end

function ad9361_baseband_auto_rate(device::Ptr{iio_device}, rate::Int64)
    if rate <= 2e7
        decimation = 4;
        taps = 128;
        fir = fir_128_4;
    elseif rate <= 4e7
        decimation = 2;
        taps = 128;
        fir = fir_128_2;
    elseif rate <= 53_333_333
        decimation = 2;
        taps = 96;
        fir = fir_96_2;
    else
        decimation = 2;
        taps = 64;
        fir = fir_64_2;
    end

    try
        global channel = C_iio_device_find_channel(device, "voltage0", true);
    catch
        return -19 # -ENODEV (I think)
    end

    ret, current_rate = C_iio_channel_attr_read_longlong(channel, "sampling_frequency");
    if ret < 0; return ret; end;
    ret, is_fir_enabled = ad9361_get_trx_fir_enable(device);
    if ret < 0; return ret; end;

    if is_fir_enabled
        if current_rate <= 25e6 / 2
            C_iio_channel_attr_write_longlong(channel, "sampling_frequency", Int64(3e6));
        end
        if (ret = ad9361_set_trx_fir_enable(device, false)) < 0
            return ret;
        end
    end

    fir_buffer = "RX 3 GAIN -6 DEC $decimation\nTX 3 GAIN 0 INT $decimation\n";
    fir_buffer *= join(map(c -> "$c,$c\n", fir));
    fir_buffer = Cuchar.(collect(fir_buffer * "\n"));

    ret = C_iio_device_attr_write_raw(device, "filter_fir_config", Ptr{Cvoid}(pointer(fir_buffer)), UInt64(sizeof(fir_buffer)));
    if ret < 0; return ret; end;

    if rate <= 25e6 / 12
        ret, value_buffer = C_iio_device_attr_read(device, "tx_path_rates");
        dac_rate = parse(Int64, match(r"(?:DAC:)(\d+)", value_buffer).captures[]);
        tx_rate  = parse(Int64, match(r"(?:TXSAMP:)(\d+)", value_buffer).captures[]);
        if tx_rate == 0; return -22; end; # -EINVAL

        max = (dac_rate / tx_rate) * 16;
        if max < taps
            C_iio_channel_attr_write_longlong(channel, "sampling_frequency", Int64(3e6));
        end
        if (ret = ad9361_set_trx_fir_enable(device, true)) < 0
            return ret;
        end
        if (ret = C_iio_channel_attr_write_longlong(channel, "sampling_frequency", rate)) < 0
            return ret;
        end
    else
        if (ret = C_iio_channel_attr_write_longlong(channel, "sampling_frequency", rate)) < 0
            return ret;
        end
        if (ret = ad9361_set_trx_fir_enable(device, true)) < 0
            return ret;
        end
    end

    return 0;
end


function QuantumOptics.expect(Obs::Union{GraphOperator, SparseOperator},
                     net::NeuralNetwork, sampler::Sampler=FullSumSampler())

    oprob = ObservablesProblem(NeuralQuantum.input_type(net), Obs)
    ov          = state(oprob, net)
    mosampl     = multithread(sampler);
    oitercache  = SamplingCache(oprob)
    moitercache = [SamplingCache(oprob) for i=1:Threads.nthreads()]
    moc         = init_sampler!(mosampl, net, ov)

    sample_net = sampler isa FullSumSampler ? sample_network_wholespace! : sample_network!
    #println(sample_net)
    init_sampler!(mosampl, net, ov, moc)
    #Threads.@threads
    for i=1:Threads.nthreads()
        s = sampler_list(mosampl)[i]
        n = moc.nets[i]
        v = moc.σs[i]
        _ic = moitercache[i]
        zero!(_ic)
        while true
            sample_net(_ic, oprob, n, v)
            !samplenext!(v, s, n, moc.caches[i]) && break
        end
    end
    zero!(oitercache)
    for ic=moitercache
        add!(oitercache, ic)
    end
    ObsEval = evaluation_post_sampling(oitercache)
    if sampler isa FullSumSampler
        return ObsEval.ObsAve[1], ObsEval.ObsVals[1]
    else
        return ObsEval.ObsAve[1], ObsEval.ObsVals[1]
    end
end


"""
Modify an .srt (subrip) subtitle file.

    using SubtitleStretcher
    cd("/Movies/")
    st = loadsubtitles("Pirates of the Caribbean The Curse of the Black Pearl (2003).eng.srt")
    stretch!(st, ["00:01:18" => "00:01:24", "2:03:08" => 7400])
"""
module SubtitleStretcher

using IterTools

import Base: convert

export loadsubtitles, savesubtitles, stretch!, Timestamp

# hold start and finish times for each subtitle
mutable struct Timestamp
    starttime::Float64  # seconds
    finishtime::Float64 #
end

# convert a timestring ("00:18:29,040") into Float64 seconds
function convert(::Type{Float64}, timestring::S where S <: AbstractString)
    if occursin(",", timestring)
        starttime, startmilli = split(timestring, ",")
    else
        starttime, startmilli = timestring, "0"
    end
    h1, m1, s1 = map(Meta.parse, split(starttime, ":"))
    return ((h1 * 60) + m1) * 60 + s1 + Meta.parse(startmilli)/1000
end

# convert two timestrings to a Timestamp
# eg "00:18:29,040 --> 00:18:30,201" -> Timestamp(1109.04, 1110.201)
function convert(::Type{Timestamp}, tstamp::S where S <: AbstractString)
    rawstarttime, rawfinishtime = split(tstamp, " --> ")
    return Timestamp(
        convert(Float64, rawstarttime),
        convert(Float64, rawfinishtime))
end

# forgot why I needed this
#convert(::Type{Timestamp}, t::Timestamp) = (t.starttime, t.finishtime)

"""
    timestamptostring(timecode)

Take part of a Timestamp and format it to the required string.

Example:

    timestamptostring(1528.32) -> "02:13:28,320"
"""
function timestamptostring(timecode)
    _min, _sec = divrem(timecode, 60)
    hr, mn    = divrem(_min, 60)
    sc, milli   = fldmod(_sec, 1)
    h = lpad(Int(floor(hr)), 2, "0")
    m = lpad(Int(floor(mn)), 2, "0")
    s = lpad(Int(floor(sc)), 2, "0")
    mill = lpad(convert(Int, round(milli, digits=3) * 1000), 3, "0")
    return "$h:$m:$s,$mill"
end

# convert both parts of a Timestamp to the full timestring
# Eg Timestamp(1528.32, 1530.163) -> "02:13:28,320 --> 02:13:30,163"
function convert(::Type{String}, tstamp::Timestamp)
    return timestamptostring(tstamp.starttime) * " --> " * timestamptostring(tstamp.finishtime)
end

"""
    loadsubtitles(filename)

Read subtitles from a file and load them into a dictionary. Files can be `.srt`
or `.sub` format.
"""
function loadsubtitles(filename)
    fn, ext = splitext(filename)
    file = open(filename) do f
        map(chomp, readlines(f)) # get rid of \r\n
    end
    subd = Dict{Int64, Tuple}()
    if ext == ".srt"
        #redo the numbers, this will correct any errors in the original
        global subcounter = 1
        # each group of lines is a subtitle
        for s in groupby(x -> x == "", file) # look for the blank line
            try # some problems with non-Unicode chars?
                if s[1] != ""
                    ts = convert(Timestamp, s[2])
                    l = length(s)
                    # more than one text line is possible
                    subd[subcounter] = (ts, [s[n] for n in 3:l])
                    subcounter += 1
                end
            catch
               println("skipping content at subtitle #$subcounter")
            end
        end
    elseif ext == ".sub"
        # each line is a subtitle
        for n in 1:length(file)
            l = file[n]
            g1, g2, g3, g4 = ["", "", "", ""]
            try # some problems with non-Unicode chars?
                # look for 3
                rm = match(r"\{(.*?)\}\{(.*?)\}(\{.*?\})(.*$)?", l)
                if rm == nothing
                    # look for 2 instead
                    rm = match(r"\{(.*?)\}\{(.*?)\}(.*$)?", l)
                    g1, g2, g4     = rm.captures
                else
                    g1, g2, g3, g4 = rm.captures
                end
            catch
                println("skipping content at line #$n")
            end
            fromtime = Meta.parse(g1)/25
            totime   = Meta.parse(g2)/25
            textlines = split(g4, "|")
            subd[n] = (Timestamp(fromtime, totime), [textlines...])
        end
    end
    return subd
end

# change value depending on one scale, to another
# eg rescale(0.5, 0, 1, 100, 200) changes the value 0.5
# (midway between 0 and 1) to 150, midway between 100 and 200
rescale(x, from_min, from_max, to_min, to_max) =
    ((x - from_min) / (from_max - from_min)) * (to_max - to_min) + to_min
rescale(x, from::NTuple{2,Number}, to::NTuple{2, Number}) =
    ((x - from[1]) / (from[2] - from[1])) * (to[2] - to[1]) + to[1]

"""
    stretch!(subd, markers)

Stretch all timestamps in a subtitle dictionary `subd`.

The `markers` argument is an array of 2 pairs of numbers:

    [oldsubtitletime1 => actualmovietime1, oldsubtitletime2 => actualmovietime2]

Example

    stretch!(subd, [1 => 1.0,  200.0 => 201.0]))
"""
function stretch!(subd, markers::Vector{Pair{T1, T2}} where T1 <: Real where T2 <: Real)
    println("...adjusting subtitle times")
    length(markers) < 2 && error("not enough markers to stretch the timecodes!")
    currentstart, requiredstart = markers[1]
    currentfinish, requiredfinish = markers[2]
    for key in sort(collect(keys(subd)))
           ts = subd[key][1]
           ttxt = subd[key][2]
           newstart = rescale(ts.starttime, currentstart, currentfinish, requiredstart, requiredfinish)
           newfinish = rescale(ts.finishtime, currentstart, currentfinish, requiredstart, requiredfinish)
           newts = Timestamp(newstart, newfinish)
           subd[key] = (newts, ttxt)
    end
    println("...subtitle times adjusted")
    return (requiredfinish - requiredstart)/(currentfinish - currentstart)
end

"""
    stretch!(subd, markers::Vector{Pair{T, T}} where T <: String)

Stretch all timestamps in a subtitle dictionary, using timestrings:

    stretch!(subd, ["00:06:44,580" =>"00:06:48",  "01:48:35,820" => "01:48:37"]

Timestrings are H:M:S:

    "00:18:29,040"
    "00:18:29"
"""
function stretch!(subd, markers::Vector{Pair{T, T}} where T <: String)
    length(markers) < 2 && error("not enough markers to stretch the timecodes!")
    cstart,  rstart = map(s -> convert(Float64, s), markers[1])
    cfinish, rfinish = map(s -> convert(Float64, s), markers[2])
    stretch!(subd, [cstart => rstart, cfinish => rfinish])
end

"""
    stretch!(subd, markers::Vector)

Stretch using mixture of numeric timestamps and strings.

    stretch!(st, [100 => "00:01:18", "2:03:08" => 7400])
"""
function stretch!(subd, markers::Vector)
    length(markers) !== 2 && error("not enough markers to stretch the timecodes!")
    r = []
    for a in [markers[1][1], markers[1][2], markers[2][1], markers[2][1]]
        if typeof(a) <: Real
            push!(r, a)
        else
            push!(r, convert(Float64, a))
        end
    end
    cstart, rstart, cfinish, rfinish = r
    stretch!(subd, [cstart => rstart, cfinish => rfinish])
end

"""
    savesubtitles(subd, filename)

Save a dictionary of subtitles to a file.
"""
function savesubtitles(subd, filename)
    open(filename, "w") do f
        for key in sort(collect(keys(subd)))
           println(f, key)
           println(f, convert(String, subd[key][1]))
           for i in subd[key][2]
               println(f, i)
           end
           println(f)
       end
    end
end

end # module


@defstruct PoolingLayer Layer (
  name :: AbstractString = "pooling",
  (bottoms :: Vector{Symbol} = Symbol[], length(bottoms) > 0),
  (tops :: Vector{Symbol} = Symbol[], length(tops) == length(bottoms)),
  (kernel :: NTuple{2, Int} = (1,1), all(broadcast(>, [kernel...], 0))),
  (stride :: NTuple{2, Int} = (1,1), all(broadcast(>, [stride...], 0))),
  (pad :: NTuple{2, Int} = (0,0), all(broadcast(>=, [pad...], 0))),
  pooling :: PoolingFunction = Pooling.Max(),
  neuron :: ActivationFunction = Neurons.Identity(),
)
@characterize_layer(PoolingLayer,
  can_do_bp  => true,
  has_neuron => true
)

type PoolingLayerState <: LayerState
  layer      :: PoolingLayer
  blobs      :: Vector{Blob}
  blobs_diff :: Vector{Blob}

  etc        :: Any
end

function setup_etc(backend::CPUBackend, layer::PoolingLayer, inputs,
    pooled_width, pooled_height)

  if isa(layer.pooling, Pooling.Max)
    masks = Array{Array}(length(inputs))
    for i = 1:length(inputs)
      width,height,channels,num = size(inputs[i])
      masks[i] = Array{Csize_t}(pooled_width[i], pooled_height[i], channels, num)
    end
    etc = masks
  else
    etc = nothing
  end
  return etc
end

function setup(backend::Backend, layer::PoolingLayer, inputs::Vector{Blob}, diffs::Vector{Blob})
  blobs = Array{Blob}(length(inputs))
  blobs_diff = Array{Blob}(length(inputs))
  pw_all = Array{Int}(length(inputs))
  ph_all = Array{Int}(length(inputs))
  for i = 1:length(inputs)
    width,height,channels,num = size(inputs[i])

    pooled_width  = round(Int, ceil(float(width +2*layer.pad[1]-layer.kernel[1]) / layer.stride[1]))+1
    pooled_height = round(Int, ceil(float(height+2*layer.pad[2]-layer.kernel[2]) / layer.stride[2]))+1

    # make sure the last pooling is not purely pooling padded area
    if layer.pad[1] > 0 || layer.pad[2] > 0
      if ((pooled_height-1) * layer.stride[2] >= height + layer.pad[2])
        pooled_height -= 1
      end
      if ((pooled_width-1) * layer.stride[1] >= width + layer.pad[1])
        pooled_width -= 1
      end
    end

    dtype = eltype(inputs[i])

    blobs[i] = make_blob(backend,dtype, (pooled_width,pooled_height,channels,num))
    if isa(diffs[i], NullBlob)
      blobs_diff[i] = NullBlob() # don't need back propagation unless bottom layer want
    else
      blobs_diff[i] = make_blob(backend,dtype, (pooled_width,pooled_height,channels,num))
    end

    pw_all[i] = pooled_width
    ph_all[i] = pooled_height
  end

  etc = setup_etc(backend, layer, inputs, pw_all, ph_all)
  state = PoolingLayerState(layer, blobs, blobs_diff, etc)
end
function shutdown(backend::CPUBackend, state::PoolingLayerState)
  map(destroy, state.blobs)
  map(destroy, state.blobs_diff)
end

function forward(backend::CPUBackend, state::PoolingLayerState, inputs::Vector{Blob})
  forward(backend, state.layer.pooling, state, inputs)
end

function forward(backend::CPUBackend, pool::StdPoolingFunction, state::PoolingLayerState, inputs::Vector{Blob})
  for i = 1:length(inputs)
    input = inputs[i].data
    output = state.blobs[i].data

    if isa(pool, Pooling.Max)
      max_pooling_forward(input, output, state.etc[i], state.layer)
    elseif isa(pool, Pooling.Mean)
      mean_pooling_forward(input, output, state.layer)
    end
  end
end

function backward(backend::CPUBackend, state::PoolingLayerState, inputs::Vector{Blob}, diffs::Vector{Blob})
  backward(backend, state.layer.pooling, state, inputs, diffs)
end

function backward(backend::CPUBackend, pool::StdPoolingFunction, state::PoolingLayerState,
    inputs::Vector{Blob}, diffs::Vector{Blob})

  for i = 1:length(inputs)
    diff = diffs[i]
    if !isa(diff, NullBlob)
      if isa(pool, Pooling.Max)
        max_pooling_backward(diff.data, state.blobs_diff[i].data, state.etc[i], state.layer)
      elseif isa(pool, Pooling.Mean)
        mean_pooling_backward(diff.data, state.blobs_diff[i].data, state.layer)
      end
    end
  end
end


# This file was generated, do not modify it. # hide
logdensity(Normal(), Lebesgue(ℝ), 1.0)