module AdvancedVI

using Random: AbstractRNG

using Distributions, DistributionsAD, Bijectors
using DocStringExtensions

using ProgressMeter, LinearAlgebra

using ForwardDiff
using Tracker

const PROGRESS = Ref(true)
function turnprogress(switch::Bool)
    @info("[AdvancedVI]: global PROGRESS is set as $switch")
    PROGRESS[] = switch
end

const DEBUG = Bool(parse(Int, get(ENV, "DEBUG_ADVANCEDVI", "0")))

include("ad.jl")

using Requires
function __init__()
    @require Flux="587475ba-b771-5e3f-ad9e-33799f191a9c" begin
        apply!(o, x, Δ) = Flux.Optimise.apply!(o, x, Δ)
        Flux.Optimise.apply!(o::TruncatedADAGrad, x, Δ) = apply!(o, x, Δ)
        Flux.Optimise.apply!(o::DecayedADAGrad, x, Δ) = apply!(o, x, Δ)
    end
    @require Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" begin
        include("compat/zygote.jl")
        export ZygoteAD

        function AdvancedVI.grad!(
            vo,
            alg::VariationalInference{<:AdvancedVI.ZygoteAD},
            q,
            model,
            θ::AbstractVector{<:Real},
            out::DiffResults.MutableDiffResult,
            args...
        )
            f(θ) = if (q isa Distribution)
                - vo(alg, update(q, θ), model, args...)
            else
                - vo(alg, q(θ), model, args...)
            end
            y, back = Zygote.pullback(f, θ)
            dy = first(back(1.0))
            DiffResults.value!(out, y)
            DiffResults.gradient!(out, dy)
            return out
        end
    end
    @require ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267" begin
        include("compat/reversediff.jl")
        export ReverseDiffAD

        function AdvancedVI.grad!(
            vo,
            alg::VariationalInference{<:AdvancedVI.ReverseDiffAD{false}},
            q,
            model,
            θ::AbstractVector{<:Real},
            out::DiffResults.MutableDiffResult,
            args...
        )
            f(θ) = if (q isa Distribution)
                - vo(alg, update(q, θ), model, args...)
            else
                - vo(alg, q(θ), model, args...)
            end
            tp = AdvancedVI.tape(f, θ)
            ReverseDiff.gradient!(out, tp, θ)
            return out
        end
    end
end

export
    vi,
    ADVI,
    ELBO,
    elbo,
    TruncatedADAGrad,
    DecayedADAGrad,
    VariationalInference

abstract type VariationalInference{AD} end

getchunksize(::Type{<:VariationalInference{AD}}) where AD = getchunksize(AD)
getADtype(::VariationalInference{AD}) where AD = AD

abstract type VariationalObjective end

const VariationalPosterior = Distribution{Multivariate, Continuous}

"""
    grad!(vo, alg::VariationalInference, q, model::Model, θ, out, args...)

Computes the gradients used in `optimize!`. Default implementation is provided for
`VariationalInference{AD}` where `AD` is either `ForwardDiffAD` or `TrackerAD`.
This implicitly also gives a default implementation of `optimize!`.

Variance reduction techniques, e.g. control variates, should be implemented in this function.
"""
function grad! end

"""
    vi(model, alg::VariationalInference)
    vi(model, alg::VariationalInference, q::VariationalPosterior)
    vi(model, alg::VariationalInference, getq::Function, θ::AbstractArray)

Constructs the variational posterior from the `model` and performs the optimization
following the configuration of the given `VariationalInference` instance.

# Arguments
- `model`: `Turing.Model` or `Function` z ↦ log p(x, z) where `x` denotes the observations
- `alg`: the VI algorithm used
- `q`: a `VariationalPosterior` for which it is assumed a specialized implementation of the variational objective used exists.
- `getq`: function taking parameters `θ` as input and returns a `VariationalPosterior`
- `θ`: only required if `getq` is used, in which case it is the initial parameters for the variational posterior
"""
function vi end

function update end

# default implementations
function grad!(
    vo,
    alg::VariationalInference{<:ForwardDiffAD},
    q,
    model,
    θ::AbstractVector{<:Real},
    out::DiffResults.MutableDiffResult,
    args...
)
    f(θ_) = if (q isa Distribution)
        - vo(alg, update(q, θ_), model, args...)
    else
        - vo(alg, q(θ_), model, args...)
    end

    chunk_size = getchunksize(typeof(alg))
    # Set chunk size and do ForwardMode.
    chunk = ForwardDiff.Chunk(min(length(θ), chunk_size))
    config = ForwardDiff.GradientConfig(f, θ, chunk)
    ForwardDiff.gradient!(out, f, θ, config)
end

function grad!(
    vo,
    alg::VariationalInference{<:TrackerAD},
    q,
    model,
    θ::AbstractVector{<:Real},
    out::DiffResults.MutableDiffResult,
    args...
)
    θ_tracked = Tracker.param(θ)
    y = if (q isa Distribution)
        - vo(alg, update(q, θ_tracked), model, args...)
    else
        - vo(alg, q(θ_tracked), model, args...)
    end
    Tracker.back!(y, 1.0)

    DiffResults.value!(out, Tracker.data(y))
    DiffResults.gradient!(out, Tracker.grad(θ_tracked))
end

"""
    optimize!(vo, alg::VariationalInference{AD}, q::VariationalPosterior, model::Model, θ; optimizer = TruncatedADAGrad())

Iteratively updates parameters by calling `grad!` and using the given `optimizer` to compute
the steps.
"""
function optimize!(
    vo,
    alg::VariationalInference,
    q,
    model,
    θ::AbstractVector{<:Real};
    optimizer = TruncatedADAGrad()
)
    # TODO: should we always assume `samples_per_step` and `max_iters` for all algos?
    alg_name = alg_str(alg)
    samples_per_step = alg.samples_per_step
    max_iters = alg.max_iters

    num_params = length(θ)

    # TODO: really need a better way to warn the user about potentially
    # not using the correct accumulator
    if (optimizer isa TruncatedADAGrad) && (θ ∉ keys(optimizer.acc))
        # this message should only occurr once in the optimization process
        @info "[$alg_name] Should only be seen once: optimizer created for θ" objectid(θ)
    end

    diff_result = DiffResults.GradientResult(θ)

    i = 0
    prog = if PROGRESS[]
        ProgressMeter.Progress(max_iters, 1, "[$alg_name] Optimizing...", 0)
    else
        0
    end

    # add criterion? A running mean maybe?
    time_elapsed = @elapsed while (i < max_iters) # & converged
        grad!(vo, alg, q, model, θ, diff_result, samples_per_step)

        # apply update rule
        Δ = DiffResults.gradient(diff_result)
        Δ = apply!(optimizer, θ, Δ)
        @. θ = θ - Δ

        AdvancedVI.DEBUG && @debug "Step $i" Δ DiffResults.value(diff_result)
        PROGRESS[] && (ProgressMeter.next!(prog))

        i += 1
    end

    return θ
end

# objectives
include("objectives.jl")

# optimisers
include("optimisers.jl")

# VI algorithms
include("advi.jl")

end # module


module Preconditioners

using JuLIP: AbstractAtoms, Preconditioner, JVecs, JVecsF, Dofs, maxdist,
               constraint, bonds, cutoff, positions
using JuLIP.Potentials: PairPotential, AnalyticPotential
using JuLIP.Constraints: project!, FixedCell
using JuLIP.ASE: chemical_symbols, rnn
using PyAMG: RugeStubenSolver

import JuLIP: update!
import Base: A_ldiv_B!, A_mul_B!

# ================ AMGPrecon =====================
# this is wrapping some machinery around PyAMG

# TODO: allow direct solver as an alternative;
#       in this case we should to a cholesky factorisation
#       on `force_update!`

"""
`AMGPrecon{T}`: a preconditioner using AMG as the main solver

`AMGPrecon` stores a field `p` which is used to determine the preconditioner
matrix via
```julia
matrix(P, at::AbstractAtoms)
```
If `p` should be updated in response to an update of `at` then
the type can overload `JuLIP.Preconditioners.update_inner!`.

## Constructor:
```
AMGPrecon(p::Any, at::AbstractAtoms; updatedist=0.3, tol=1e-7)
```
## TODO:
* provide mechanism for automatically determining updatedist
from the nearest-neighbour distance in `at`
* should also have automatic updates every 10 or so iterations
"""
type AMGPrecon{T} <: Preconditioner
   p::T
   amg::RugeStubenSolver
   oldX::JVecsF
   updatedist::Float64
   tol::Float64
   updatefreq::Int
   skippedupdates::Int
end

function AMGPrecon(p, at::AbstractAtoms; updatedist=0.3, tol=1e-7, updatefreq=10)
   # make sure we don't use this in a context it is not intended for!
   @assert isa(constraint(at), FixedCell)
   P = AMGPrecon(p, RugeStubenSolver(speye(2)), copy(positions(at)),
                     updatedist, tol, updatefreq, 0)
   return force_update!(P, at)
end

A_ldiv_B!(out::Dofs, P::AMGPrecon, x::Dofs) = A_ldiv_B!(out, P.amg, x)
A_mul_B!(out::Dofs, P::AMGPrecon, f::Dofs) = A_mul_B!(out, P.amg, f)

need_update(P::AMGPrecon, at::AbstractAtoms) =
   (P.skippedupdates > P.updatefreq) ||
   (maxdist(positions(at), P.oldX) >= P.updatedist)

update!(P::AMGPrecon, at::AbstractAtoms) =
   need_update(P, at) ? force_update!(P, at) : (P.skippedupdates += 1; P)

function force_update!(P::AMGPrecon, at::AbstractAtoms)
   # perform updates of the potential p (if needed; usually not)
   P.p = update_inner!(P.p, at)
   # construct the preconditioner matrix ...
   A = project!( constraint(at), matrix(P.p, at) )
   # and the AMG solver
   P.amg = RugeStubenSolver(A, tol=P.tol)
   # remember the atom positions
   copy!(P.oldX, positions(at))
   # and remember that we just did a full update
   P.skippedupdates = 0
   return P
end

# ============== some tools to construct preconditioners ======================

function estimate_rnn(at::AbstractAtoms)
   sym = unique(chemical_symbols(at))
   return minimum([rnn(s) for s in sym])
end

# ============== Implementation of some concrete types ======================

# default implementation of update_inner!;
# most of the time no inner update is required

"update `AMGPrecon.p`, must return updated `p`"
update_inner!(p::Any, at::AbstractAtoms) = p

"return the three linear indices associated with an atom index"
atind2lininds(i::Integer) = (i-1) * 3 + [1;2;3]

"""
build the preconditioner matrix associated with a pair potential;
this is related to but not even close to the same as the hessian matrix!
"""
function matrix(p::PairPotential, at::AbstractAtoms)
   I = Int[]; J = Int[]; Z = Float64[]
   for (i, j, r, _, _) in bonds(at, cutoff(p))
      # the next 2 lines add an identity block for each atom
      # TODO: should experiment with other matrices, e.g., R ⊗ R
      ii = atind2lininds(i); jj = atind2lininds(j)
      z = p(r)
      for (a, b) in zip(ii, jj)
         append!(I, [a; a; b; b])
         append!(J, [a; b; a; b])
         append!(Z, [z; -z; -z; z])
      end
   end
   N = 3*length(at)
   return sparse(I, J, Z, N, N) + 0.001 * speye(N)
end

"""
A variant of the `Exp` preconditioner; see

### Constructor: `Exp(at::AbstractAtoms)`

Keyword arguments:

* `A=3.0`: stiffness of potential
* `r0=nothing`: if `nothing`, then it is an estimate nn distance
* `cutoff_mult`: cut-off multiplier
* `tol, updatefrew`: AMG parameters

### Reference

      D. Packwood, J. Kermode, L. Mones, N. Bernstein, J. Woolley, N. I. M. Gould,
      C. Ortner, and G. Csanyi. A universal preconditioner for simulating condensed
      phase materials. J. Chem. Phys., 144, 2016.
"""
function Exp(at::AbstractAtoms;
             A=3.0, r0=nothing, cutoff_mult=2.2, tol=1e-7, updatefreq=10)
   if r0 == nothing
      r0 = estimate_rnn(at)
   end
   rcut = r0 * cutoff_mult
   exp_shit = exp( - A*(rcut/r0 - 1.0) )
   pot = AnalyticPotential( :(exp( - $A * (r/$r0 - 1.0)) - $exp_shit),
                            cutoff = rcut )
   return AMGPrecon(pot, at, updatedist=0.2 * r0, tol=tol, updatefreq=updatefreq)
end

end # end module Preconditioners


import NNlibCUDA: batchnorm, ∇batchnorm

function (BN::Flux.BatchNorm)(x::Union{CuArray{T,2},CuArray{T,4},CuArray{T,5}},
                              cache=nothing) where T<:Union{Float32, Float64}

  @assert BN.affine "BatchNorm: only affine=true supported on gpu"
  @assert BN.track_stats "BatchNorm: only track_stats=true supported on gpu"
  @assert length(BN.β) == size(x, ndims(x)-1) "BatchNorm: input has wronng number of channels"
  return BN.λ.(batchnorm(BN.γ, BN.β, x, BN.μ, BN.σ², BN.momentum;
                  cache=cache, alpha=1, beta=0, eps=BN.ϵ,
                  training=Flux._isactive(BN)))
end

function ChainRulesCore.rrule(::typeof(batchnorm), g, b, x, running_mean, running_var, momentum; kw...)
  y = batchnorm(g, b, x, running_mean, running_var, momentum; kw...)
  function batchnorm_pullback(Δ)
    grad = ∇batchnorm(g, b, x, Δ, running_mean, running_var, momentum; kw...)
    (NoTangent(), grad..., NoTangent(), NoTangent(), NoTangent())
  end
  y, batchnorm_pullback
end


###################################################################################
# Antony Schutz 2015, ANR - POLCA - 2016
###################################################################################
# OIFITS RELATED - Read files, create quantities of interest
###################################################################################

# path: read OIFITS data from files included in folder OIDATA.Folder
# OIDATA: structure containing all OIDATA and users informations
# indfile: restrict the number of files (as listed by dir() or ls() )
# ex: indfile = [1 3 5] will only read information from the 1st, the 3rd and the 5th file
# if empty read all the files
# indwvl: restrict the analyzed wavelengths restricted by index, ex:
# indwvl=1:20 --- indwvl= [1 5 7] --- indwvl = 1:5:20
#
# - - - - - - - - - - - - - - - - - - - restreindre les longueurs d'ondes par des valeurs en metre
#
function readoifits(OIDATA::PAINTER_Input,indfile=[],indwvl=[])

# Check for DP and T3 must be in data at least
    OIDATA.isDP = 1

# Read Directory
    Folder    = OIDATA.Folder
    CDir      = readdir(Folder)

# check files fits and oifits
    fits      = filter(r"\.fits$", CDir)
    oifits    = filter(r"\.oifits$", CDir)
    CDir      = vcat(fits,oifits)
    NbFiles   = length(CDir)

    if NbFiles == 0
      error("No fits or oifits files in folder: $Folder")
    end

# files reduction, only files indexed by "indfile" are read
    if isempty(indfile)
        indfile = 1:NbFiles
    end
    OIDATA.indwvl = indwvl

# temporary variable for closure index
    LT3           = 0
    LFil          = 0
# loop on files
# Extract Data from OIFITS or FITS Files

    for n in indfile
        name = string(Folder, "/", CDir[n])
        println(name)
        master = OIFITS.load(name)
# # # OI wavelength
        nw = 0
        for db in OIFITS.select(master, "OI_WAVELENGTH")
            nw = length(OIFITS.get_eff_wave(db))
        end
# Variables initialization
        if n == minimum(indfile)
            if isempty(OIDATA.indwvl)
                OIDATA.indwvl = 1:nw
            end
# Initialisation of "empty" array
            OIDATA.Closure_index = Array(Int, 0, 3)
            OIDATA.U = Array(Float64, 0, nw)
            OIDATA.V = Array(Float64, 0, nw)
            OIDATA.P = Array(Float64, 0, nw)
            OIDATA.W = Array(Float64, 0, nw)
            OIDATA.T3 = Array(Float64, 0, nw)
            OIDATA.T3err = Array(Float64, 0, nw)
            OIDATA.DP = Array(Float64, 0, nw)
            OIDATA.DPerr = Array(Float64, 0, nw)
            OIDATA.wvl = Array(Float64, 0, nw)
        end

# # # OI wavelength
        for db in OIFITS.select(master, "OI_WAVELENGTH")
            OIDATA.wvl = vcat(OIDATA.wvl,OIFITS.get_eff_wave(db)')
        end
# # # OI VIS
        if isempty(OIFITS.select(master, "OI_VIS"))
            println("OI_VIS field is missing, PAINTER needs differential visibilities")
            println(" --- will try to work only with closure phases ")
            OIDATA.isDP = 0
        else

            for dbvis1 in OIFITS.select(master, "OI_VIS")
                OIDATA.DP = vcat(OIDATA.DP, OIFITS.get_visphi(dbvis1)')
                OIDATA.DPerr = vcat(OIDATA.DPerr, OIFITS.get_visphierr(dbvis1)')
            end

        end

# # OI VIS 2
        if isempty(OIFITS.select(master, "OI_VIS2"))
            error("OI_VIS2 field is missing, PAINTER needs squared visibilities")

        else
            U_v21 = Array(Float64, 0)

            for dbvis2 in OIFITS.select(master,  "OI_VIS2")
                U,V  = spatialfrequencies(OIFITS.get_ucoord(dbvis2), OIFITS.get_vcoord(dbvis2), OIFITS.get_eff_wave(dbvis2))
                U_v21 = vcat(U_v21, OIFITS.get_ucoord(dbvis2)) # for closure index
                OIDATA.U = vcat(OIDATA.U, U)
                OIDATA.V = vcat(OIDATA.V, V)
                OIDATA.P = vcat(OIDATA.P, OIFITS.get_vis2data(dbvis2)')
                OIDATA.W = vcat(OIDATA.W, OIFITS.get_vis2err(dbvis2)')
            end
        end

# # OI T3
        if isempty(OIFITS.select(master, "OI_T3"))
            error("OI_T3 field is missing, PAINTER needs closure phases")

        else
            U_t31 = Array(Float64, 0)
            U_t32 = Array(Float64, 0)
            for dbvis3 in OIFITS.select(master, "OI_T3")
                OIDATA.T3 = vcat(OIDATA.T3, OIFITS.get_t3phi(dbvis3)')
                OIDATA.T3err = vcat(OIDATA.T3err, OIFITS.get_t3phierr(dbvis3)')
# Closure index
                U_t31 = OIFITS.get_u1coord(dbvis3)
                U_t32 = OIFITS.get_u2coord(dbvis3)
                (OIDATA.Closure_index, LT3, LFil) = findclosureindex(U_v21, U_t31, U_t32, OIDATA.Closure_index, LT3, LFil)
            end
        end
    end

## FINAL OUTPUT - Wavelength reduction
    unitconvert = pi / 180
    OIDATA.FilesName = CDir
    OIDATA.indfile = indfile
    OIDATA.U = OIDATA.U[:, OIDATA.indwvl]
    OIDATA.V = OIDATA.V[:, OIDATA.indwvl]
    OIDATA.P = OIDATA.P[:, OIDATA.indwvl]
    OIDATA.W = OIDATA.W[:, OIDATA.indwvl]

    OIDATA.T3 = OIDATA.T3[:, OIDATA.indwvl] * unitconvert
    OIDATA.T3err = OIDATA.T3err[:, OIDATA.indwvl] * unitconvert
    if OIDATA.isDP == 1
        OIDATA.DP = OIDATA.DP[:, OIDATA.indwvl] * unitconvert
        OIDATA.DPerr = OIDATA.DPerr[:, OIDATA.indwvl] * unitconvert
    end
    OIDATA.wvl = vec(mean(OIDATA.wvl[:, OIDATA.indwvl], 1))
# # compute Differential Phases as needed for PAINTER and adjust angle unit
    OIDATA.nb = size(OIDATA.U, 1)
    OIDATA.nw = size(OIDATA.U, 2)

# # Create cluster of independent CLosure Phases and associated
    OIDATA = independentT3data(OIDATA)

    return OIDATA
end

function independentT3data(OIDATA::PAINTER_Input)
  Closure_index = OIDATA.Closure_index
  Cluster = independentT3(Closure_index)
  baseNb = basesincluster(Cluster)
  orderedCluster = makeclusterordered(Cluster,baseNb)
  rowt3 = t3row(Closure_index,Cluster)
  Nt3indep = length(baseNb)
  Xi = Dict{}()
  K = Dict{}()
  H = Dict{}()
  for n in 1:Nt3indep
      HDP = phasetophasediff(orderedCluster[n], OIDATA.nw, length(baseNb[n]), 0, 1, OIDATA.dptype, OIDATA.dpprm)
      T3 = OIDATA.T3[rowt3[n],:]
      T3err = OIDATA.T3err[rowt3[n],:]
      if OIDATA.isDP == 1
          DP = OIDATA.DP[baseNb[n],:]
          DPerr = OIDATA.DPerr[baseNb[n],:]
          Xi[n] = vcat(vec(T3),HDP*vec(DP)  )
          K[n] = vcat(vec(T3err),abs(HDP)*vec(DPerr)  )
      else
          Xi[n] = vec(T3)
          K[n] = vec(T3err)
      end
  end
  OIDATA.Xi = Xi
  OIDATA.K = K
  OIDATA.baseNb = baseNb
  OIDATA.orderedCluster= orderedCluster
  return OIDATA
end
###################################################################################
# FUNCTIONS FOR OIFITS DATA TRANSFORMATION
###################################################################################
# Spatial Frequencies construction
# ---------------------------------------------------------------------------------
function spatialfrequencies(U_Coord::Vector,V_Coord::Vector,wvl::Vector)
# Create spatiale frequencies from bases coordinates and wavelength
# (U,V) = SpatialFrequencies(U_Coord::Vector,V_Coord::Vector,wvl::Vector)
# U_Coord:  U coordinate of V2 (bases not spatial frequencies) vector (U not divided by wavelength)
# V_Coord:  V coordinate of V2 (bases not spatial frequencies) vector (U not divided by wavelength)
# wvl    :  vector of effective used wavelength
    U_3D = repmat(U_Coord, 1, length(wvl)) ./ repmat(wvl', length(U_Coord), 1)
    V_3D = repmat(V_Coord, 1, length(wvl)) ./ repmat(wvl', length(V_Coord), 1)
    return U_3D,V_3D
end
# ---------------------------------------------------------------------------------
# PHASE CLOSURE INDEX (sample's index) creation
# ---------------------------------------------------------------------------------
function findclosureindex(V2_U_Coord::Vector,T3_U1_Coord::Vector,T3_U2_Coord::Vector,Closure_index::Matrix,LT3::Integer,LFil::Integer)
# update Closure_index array for each files
# deduce the index in sample related to phases vector, from U1 and U2 coord of phi T3 to U coord of V2
# Closure_index = CreateClosureIndex(V2_U_Coord::Vector,T3_U1_Coord::Vector,T3_U2_Coord::Vector)
# V2_U_Coord:  U coordinate of V2 (bases not spatial frequencies) vector (U not divided by wavelength)
# T3_U1_Coord: U coordinate of T3 first base (not spatial frequencies) vector (U1 not divided by wavelength)
# T3_U2_Coord: U coordinate of T3 second base (not spatial frequencies) vector (U2 not divided by wavelength)
# U3 is given by default
# LT3 and LFil are used to increment correctly information from several files
# Closure_index: index in tabs of associated visibilities
    utmp = round(1000 * V2_U_Coord )
    U1 = round(1000 * T3_U1_Coord)
    U2 = round(1000 * T3_U2_Coord)
    U3 = round(1000 * (T3_U1_Coord + T3_U2_Coord))
    for k in 1:length(U1)
        pos1 = LT3 + find(x->(x == U1[k]), utmp)
        pos2 = LT3 + find(x->(x == U2[k]), utmp)
        pos3 = LT3 + find(x->(x == U3[k]), utmp)
        Closure_index = vcat(Closure_index, hcat(pos1, pos2, pos3))
    end
    LT3 = maximum(Closure_index[:])
    LFil = LFil + length(U1)
    return Closure_index, LT3, LFil
end
# ---------------------------------------------------------------------------------
# PAINTER used Differential Visibilities, creation from Differential Visibilities
# ---------------------------------------------------------------------------------
function diffphi(DiffPhi::Array, HDP::SparseMatrixCSC)
# DiffPhi: Differential visibilities as given in OIFITS Data
# DiffPhi is matrix (Nbase*Nwvl)
# All channels reference are assumed to be equal
# defined as in equation 21 of PAINTER (1)
    DiffPhiAB = HDP * vec(DiffPhi)
end
# ---------------------------
# All channels are assumed to be independent
# variance of sum of random variables = sum of random variables variance
function diffphierr(DiffPhierr::Array, HDP::SparseMatrixCSC)
    DiffPhiABerr = abs(HDP) * vec(DiffPhierr)
end
# ---------------------------------------------------------------------------------
# INDEPENDENT Phases Closures from index -- to put in PAINTEROIFITS.JL
# ---------------------------------------------------------------------------------
function independentT3(Closure_index::Matrix)
    Cluster = Dict{}()
    Index = collect(1:size(Closure_index,1))
    m=0
    while sum(Index)>0
        m+=1
        idxinit = find(Index.>0)[1]
        Cluster[m] = Closure_index[Index[idxinit],:]
        Index[idxinit] = 0
        nn = 0
        while true
            nn+=1
            init = vec(Cluster[m][nn,:])
            Cluster[m], Index = searchfromclusterinlist(Cluster[m],init,Index,Closure_index)
            l = size(Cluster[m],1)
            if  nn==l
                break
            end

          end
      end
    return Cluster
end
function searchfromclusterinlist(Cluster::Matrix,init::Vector,Index::Vector,Closure_index::Matrix)
    for n in Index
      if n>0
        if( !isempty( find( init[1] .== Closure_index[n,:] ) )
          ||!isempty( find( init[2] .== Closure_index[n,:] ) )
          ||!isempty( find( init[3] .== Closure_index[n,:] ) ) )
            Cluster = vcat( Cluster , Closure_index[n,:] )
            Index[n] = 0
        end
      end
    end
    return Cluster, Index
end
# ---------------------------------------------------------------------------------
# Number of bases in each cluster - tools to create smaller T3 and DP Matrices
# ---------------------------------------------------------------------------------
function basesincluster(Cluster::Dict)
    Nc = length(Cluster)
    basenb = Dict{}()
    for n in 1:Nc
        vecclust = sort( vec( Cluster[n] ) )
        indpos = vcat(1, diff( vecclust ).>0)
        basenb[n] = round(Int, vecclust[ (vecclust.*indpos).>0 ])
    end
    return basenb
end
# ---------------------------------------------------------------------------------
# Sort bases in Cluster from 1 to Nb base in Cluster
# ---------------------------------------------------------------------------------
function makeclusterordered(Cluster::Dict,baseNb::Dict)
  Nc = length(Cluster)
  theorderedcluster = Dict{}()
  for n in 1:Nc
      thecluster = Cluster[n]
      baseinthecluster = baseNb[n]
      numberofindependentbaseinthecluster = length(baseinthecluster)
      nt3 = size(thecluster,1)
      orderedbaseinthecluster = collect(1:numberofindependentbaseinthecluster)
      thevectorizedcluster = vec(thecluster)
      theorderedvectorizedcluster = zeros(thevectorizedcluster)
      for m in 1:numberofindependentbaseinthecluster
          cond = find( thevectorizedcluster .==  baseinthecluster[m])
          theorderedvectorizedcluster[cond] = orderedbaseinthecluster[m]
        end
      theorderedcluster[n] = reshape(theorderedvectorizedcluster,nt3,3)
  end
  return theorderedcluster
end
# ---------------------------------------------------------------------------------
# T3 rows related to H
# ---------------------------------------------------------------------------------
function t3row(Closure_index::Matrix,Cluster::Dict)
  Nc = length(Cluster)
  rowt3 = Dict{}()
  for n in 1:Nc
      thecluster = Cluster[n]
      lc = size(thecluster,1)
      rowt3[n] = zeros(lc)
      for m in 1:lc
          rowt3[n][m] = find( prod(Closure_index .== thecluster[m,:],2))[1]
      end
      rowt3[n] = round(Int,rowt3[n])
  end
  return rowt3
end


module KaTeX

const assetsdir = joinpath(@__DIR__, "..", "assets")
const assets = ["auto-render.min.js", "katex.min.css", "katex.min.js"]
const fontsdir = joinpath(assetsdir, "fonts")

using WebIO: Scope, onimport, onjs, node
using JSExpr: @js
using AssetRegistry: register
using Observables: AbstractObservable, Observable

const katex_min_css = joinpath(KaTeX.assetsdir, "katex.min.css")
const katex_min_js = joinpath(KaTeX.assetsdir, "katex.min.js")

"""
`latex(txt)`

Render `txt` in LaTeX using KaTeX. Backslashes need to be escaped:
`latex("\\\\sum_{i=1}^{\\\\infty} e^i")`
"""
function latex(txt)
    (txt isa AbstractObservable) || (txt = Observable(txt))
    register(assetsdir)
    w = Scope(imports=[
        katex_min_js,
        katex_min_css
    ])

    w["value"] = txt

    onimport(w, @js function (k)
        this.k = k
        this.container = this.dom.childNodes[0]
        k.render($(txt[]), this.container)
    end)

    onjs(w["value"], @js (txt) -> this.k.render(txt, this.container))

    w.dom = node(:div)
    w
end

end # module


import DataStructures: OrderedDict

"""
```
RealBondMkup{T} <: AbstractModel{T}
```

### Fields

#### Parameters and Steady-States
* `parameters::Vector{AbstractParameter}`: Vector of all time-invariant model
  parameters.

* `steady_state::Vector{AbstractParameter}`: Model steady-state values, computed
  as a function of elements of `parameters`.

* `keys::OrderedDict{Symbol,Int}`: Maps human-readable names for all model
  parameters and steady-states to their indices in `parameters` and
  `steady_state`.

#### Inputs to Measurement and Equilibrium Condition Equations

The following fields are dictionaries that map human-readable names to row and
column indices in the matrix representations of of the measurement equation and
equilibrium conditions.

* `endogenous_states::OrderedDict{Symbol,Int}`: Maps each state to a column in
  the measurement and equilibrium condition matrices.

* `exogenous_shocks::OrderedDict{Symbol,Int}`: Maps each shock to a column in
  the measurement and equilibrium condition matrices.

* `expected_shocks::OrderedDict{Symbol,Int}`: Maps each expected shock to a
  column in the measurement and equilibrium condition matrices.

* `equilibrium_conditions::OrderedDict{Symbol,Int}`: Maps each equlibrium
  condition to a row in the model's equilibrium condition matrices.

* `observables::OrderedDict{Symbol,Int}`: Maps each observable to a row in the
  model's measurement equation matrices.

#### Model Specifications and Settings

* `spec::String`: The model specification identifier, \"an_schorfheide\", cached
  here for filepath computation.

* `subspec::String`: The model subspecification number, indicating that some
  parameters from the original model spec (\"ss0\") are initialized
  differently. Cached here for filepath computation.

* `settings::Dict{Symbol,Setting}`: Settings/flags that affect computation
  without changing the economic or mathematical setup of the model.

* `test_settings::Dict{Symbol,Setting}`: Settings/flags for testing mode

#### Other Fields

* `rng::MersenneTwister`: Random number generator. Can be is seeded to ensure
  reproducibility in algorithms that involve randomness (such as
  Metropolis-Hastings).

* `testing::Bool`: Indicates whether the model is in testing mode. If `true`,
  settings from `m.test_settings` are used in place of those in `m.settings`.

* `observable_mappings::OrderedDict{Symbol,Observable}`: A dictionary that
  stores data sources, series mnemonics, and transformations to/from model units.
  DSGE.jl will fetch data from the Federal Reserve Bank of
  St. Louis's FRED database; all other data must be downloaded by the
  user. See `load_data` and `Observable` for further details.

"""
mutable struct RealBondMkup{T} <: AbstractModel{T}
    parameters::ParameterVector{T}                         # vector of all time-invariant model parameters
    steady_state::ParameterVector{T}                       # model steady-state values

    # Temporary to get it to work. Need to
    # figure out a more flexible way to define
    # "grids" that are not necessarily quadrature
    # grids within the model
    grids::OrderedDict{Symbol,Union{Grid, Vector}}

    keys::OrderedDict{Symbol,Int}                          # human-readable names for all the model
                                                           # parameters and steady-states

    state_variables::Vector{Symbol}                        # Vector of symbols of the state variables
    jump_variables::Vector{Symbol}                         # Vector of symbols of the jump variables
    normalized_model_states::Vector{Symbol}                # All of the distributional model
                                                           # state variables that need to be normalized
    endogenous_states_unnormalized::OrderedDict{Symbol,UnitRange} # Vector of unnormalized
                                                           # ranges of indices
    endogenous_states::OrderedDict{Symbol,UnitRange}       # Vector of ranges corresponding
                                                           # to normalized (post Klein solution) indices
    exogenous_shocks::OrderedDict{Symbol,Int}              #
    expected_shocks::OrderedDict{Symbol,Int}               #
    equilibrium_conditions::OrderedDict{Symbol,UnitRange}  #
    model_states_augmented::OrderedDict{Symbol, Int}
    observables::OrderedDict{Symbol,Int}                   #

    spec::String                                           # Model specification number (eg "m990")
    subspec::String                                        # Model subspecification (eg "ss0")
    settings::Dict{Symbol,Setting}                         # Settings/flags for computation
    test_settings::Dict{Symbol,Setting}                    # Settings/flags for testing mode
    rng::MersenneTwister                                   # Random number generator
    testing::Bool                                          # Whether we are in testing mode or not

    observable_mappings::OrderedDict{Symbol, Observable}
end

description(m::RealBondMkup) = "RealBondMkup, $(m.subspec)"

"""
`init_model_indices!(m::RealBondMkup)`

Arguments:
`m:: RealBondMkup`: a model object

Description:
Initializes indices for all of `m`'s states, shocks, and equilibrium conditions.
"""
function init_model_indices!(m::RealBondMkup)

    # Endogenous states
    endogenous_states = collect([
    # These states corresp. to the following in the original notation
    #    MUP,   LIIP,    ZP,    MONP,    ELLP,  RRP
        :μ′_t, :l_i′_t, :z′_t, :mon′_t, :l′_t, :R′_t,
    #    IIP    WWP    PIP    TTP
        :i′_t, :w′_t, :π′_t, :t′_t])

    # Exogenous shocks
    exogenous_shocks = collect([:z_sh, :mon_sh, :mkp_sh])

    # Equilibrium conditions
    equilibrium_conditions = collect([
        :eq_euler, :eq_kolmogorov_fwd, :eq_market_clearing, :eq_TFP,
        :eq_phillips, :eq_taylor, :eq_fisher, :eq_transfers, :eq_monetary_policy,
        :eq_markup, :eq_lagged_nominal_rate])

    model_states_augmented = [:y_t1]

    # Observables
    observables = keys(m.observable_mappings)

    ########################################################################################
    # Setting indices of endogenous_states and equilibrium conditions manually for now
    nx = get_setting(m, :nx)
    ns = get_setting(m, :ns)
    nxns = nx*ns
    endo = m.endogenous_states_unnormalized
    eqconds = m.equilibrium_conditions

    # State variables
    endo[:μ′_t]   = 1:nxns
    endo[:l_i′_t] = nxns+1:nxns+1
    endo[:z′_t]   = nxns+2:nxns+2
    endo[:mon′_t] = nxns+3:nxns+3
    endo[:mkp′_t] = nxns+4:nxns+4

    # Jump variables
    endo[:l′_t]   = nxns+5:2*nxns+4
    endo[:R′_t]   = 2*nxns+5:2*nxns+5
    endo[:i′_t]   = 2*nxns+6:2*nxns+6
    endo[:w′_t]   = 2*nxns+7:2*nxns+7
    endo[:π′_t]   = 2*nxns+8:2*nxns+8
    endo[:t′_t]   = 2*nxns+9:2*nxns+9

    eqconds[:eq_euler]              = 1:nxns
    eqconds[:eq_kolmogorov_fwd]     = nxns+1:2*nxns
    eqconds[:eq_market_clearing]    = 2*nxns+1:2*nxns+1
    eqconds[:eq_TFP]                = 2*nxns+2:2*nxns+2
    eqconds[:eq_phillips]           = 2*nxns+3:2*nxns+3
    eqconds[:eq_taylor]             = 2*nxns+4:2*nxns+4
    eqconds[:eq_fisher]             = 2*nxns+5:2*nxns+5
    eqconds[:eq_transfers]          = 2*nxns+6:2*nxns+6
    eqconds[:eq_monetary_policy]    = 2*nxns+7:2*nxns+7
    eqconds[:eq_markup]             = 2*nxns+8:2*nxns+8
    eqconds[:eq_lagged_nominal_rate]= 2*nxns+9:2*nxns+9
    ########################################################################################

    m.normalized_model_states = [:μ′_t]
    m.endogenous_states = deepcopy(endo)
    m.state_variables = m.endogenous_states.keys[get_setting(m, :state_indices)]
    m.jump_variables = m.endogenous_states.keys[get_setting(m, :jump_indices)]

    for (i,k) in enumerate(exogenous_shocks);            m.exogenous_shocks[k]            = i end
    for (i,k) in enumerate(model_states_augmented); m.model_states_augmented[k]   = i+length(endogenous_states) end
    for (i,k) in enumerate(observables);                 m.observables[k]                 = i end
end

function RealBondMkup(subspec::String="ss0";
                   custom_settings::Dict{Symbol, Setting} = Dict{Symbol, Setting}(),
                   testing = false)

    # Model-specific specifications
    spec               = "real_bond_mkup"
    subspec            = subspec
    settings           = Dict{Symbol,Setting}()
    test_settings      = Dict{Symbol,Setting}()
    rng                = MersenneTwister(0)

    # initialize empty model
    m = RealBondMkup{Float64}(
            # model parameters and steady state values
            Vector{AbstractParameter{Float64}}(), Vector{Float64}(),
            # grids and keys
            OrderedDict{Symbol,Grid}(), OrderedDict{Symbol,Int}(),

            # normalized_model_states, state_inds, jump_inds
            Vector{Symbol}(), Vector{Symbol}(), Vector{Symbol}(),

            # model indices
            # endogenous states unnormalized, endogenous states normalized
            OrderedDict{Symbol,UnitRange}(), OrderedDict{Symbol,UnitRange}(),
            OrderedDict{Symbol,Int}(), OrderedDict{Symbol,Int}(),
            OrderedDict{Symbol,UnitRange}(), OrderedDict{Symbol,Int}(), OrderedDict{Symbol, Int}(),

            spec,
            subspec,
            settings,
            test_settings,
            rng,
            testing,
            OrderedDict{Symbol,Observable}())
    # Set settings
    model_settings!(m)

    # default_test_settings!(m)
    for custom_setting in values(custom_settings)
        m <= custom_setting
    end

    # Set observable transformations
    init_observable_mappings!(m)

    # Initialize parameters
    init_parameters!(m)

    # Initialize grids
    init_grids!(m)

    # Initialize model indices
    init_model_indices!(m)

    # # Initialize subspecs
    init_subspec!(m)

    # Solve for the steady state
    steadystate!(m)

    # So that the indices of m.endogenous_states reflect the normalization
    normalize_model_state_indices!(m)

    return m
end

"""
```
init_parameters!(m::RealBondMkup)
```

Initializes the model's parameters, as well as empty values for the steady-state
parameters (in preparation for `steadystate!(m)` being called to initialize
those).
"""
function init_parameters!(m::RealBondMkup)

    # Steady State Parameters
    m <= parameter(:R, 1.04, fixed = true,
                   description = "R: Steady-state gross real interest rate.", tex_label = "R")
    m <= parameter(:γ, 1.0, fixed = true,
                   description = "γ: CRRA Parameter.", tex_label = "\\gamma")
    m <= parameter(:ν, 1.0, fixed = true,
                   description = "Inverse Frisch elasticity of labor supply.", tex_label = "\\nu")
    m <= parameter(:abar, -0.5, fixed = true,
                   description = "Borrowing floor.", tex_label = "\\bar{a}")

    # AR for observables (Non-steadystate)
    m <= parameter(:ρ_z, 0.03081565350294113, (1e-3, 0.999), (1e-3, 0.999), SquareRoot(), BetaAlt(0.5, 0.2),
                   fixed=false,
                   description="ρ_z: AR(1) coefficient in the technology process.",
                   tex_label="\\rho_z")
    m <= parameter(:ρ_mon, 0.5, (1e-3, 0.999), (1e-3, 0.999), SquareRoot(), BetaAlt(0.5, 0.2),
                   fixed = false,
                   description = "ρ_mon: Persistence of monetary policy shock",
                   tex_label = "\\rho_{mon}")
    m <= parameter(:ρ_mkp, 0.5, (1e-3, 0.999), (1e-3, 0.999), SquareRoot(), BetaAlt(0.5, 0.2),
                   fixed = false,
                   tex_label = "\\rho_{mkp}", description = "ρ_mkp: Persistence of the markup shock")

    m <= parameter(:σ_z, sqrt(.007), (1e-8, 5.), (1e-8, 5.), ModelConstructors.Exponential(), RootInverseGamma(2, 0.10),
                   fixed=false,
                   description="σ_z: The standard deviation of the process describing the stationary component of productivity.",
                   tex_label="\\sigma_{z}")
    m <= parameter(:σ_mon, 0.2380, (1e-8, 5.), (1e-8, 5.), ModelConstructors.Exponential(), RootInverseGamma(2, 0.10),
                   fixed = false,
                   description="σ_mon: The standard deviation of the monetary policy shock.",
                   tex_label="\\sigma_{mon}")
    m <= parameter(:σ_mkp, 0.1314, (1e-8, 5.), (1e-8, 5.), ModelConstructors.Exponential(), RootInverseGamma(2, 0.10),
                   fixed=false,
                   description="σ_mkp: The mean of the process that generates the price elasticity of the composite good. Specifically, the elasticity is (1+λ_{f,t})/(λ_{f_t}).",
                   tex_label="\\sigma_{\\lambda_f}")

    # More non-steadystate parameters
    m <= parameter(:ρ_tay, 0.5, (1e-5, 0.999), (1e-5, 0.999), SquareRoot(), BetaAlt(0.5, 0.2),
                   fixed = false,
                   description = "ρ_tay: Persistence in the taylor rule",
                   tex_label = "rho_{tay}")
    m <= parameter(:κ, 0.5, (1e-5, 0.999), (1e-5, 0.999), SquareRoot(), Distributions.Uniform(0.0, 1.0),
                   fixed = false,
                   description = "κ: The slope of the Phillips curve",
                   tex_label = "\\kappa")
    m <= parameter(:phipi, 1.5, (1e-5, 10.), (1e-5, 10.00), ModelConstructors.Exponential(), Normal(1.5, 0.25),
                   fixed = false,
                   description = "phipi: The slope of the taylor rule",
                   tex_label = "\\phi_\\pi")

    # More steadystate parameters
    m <= parameter(:μ_s, 0., fixed = true, description = "μ_s: Mu of log normal in income")
    m <= parameter(:σ_s, 0.1, fixed = true,
                   description = "σ_s: Sigma of log normal in income")
    m <= parameter(:e_y, 0., fixed = true, description = "e_y: Measurement error on GDP",
                   tex_label = "e_y")

    # Setting steady-state parameters
    nx = get_setting(m, :nx)
    ns = get_setting(m, :ns)
    nxns = nx*ns

    m <= SteadyStateParameterGrid(:lstar, fill(NaN, nxns), description = "Steady-state expected discounted
                                  marginal utility of consumption", tex_label = "l_*")
    m <= SteadyStateParameterGrid(:cstar, fill(NaN, nxns), description = "Steady-state consumption",
                                  tex_label = "c_*")
    m <= SteadyStateParameterGrid(:ηstar, fill(NaN, nxns), description = "Steady-state
                                  level of labor supply",
                                  tex_label = "\\eta_*")
    m <= SteadyStateParameterGrid(:μstar, fill(NaN, nxns), description = "Steady-state cross-sectional
                                  density of cash on hand", tex_label = "\\mu_*")

    # Figure out a better description for this...
    m <= SteadyStateParameterGrid(:χstar, fill(NaN, nxns), description = "Steady-state
                                  solution for constrained consumption and labor supply",
                                  tex_label = "\\chi_*")

    m <= SteadyStateParameter(:βstar, NaN, description = "Steady-state discount factor",
                              tex_label = "\\beta_*")
end

"""
```
init_grids!(m::RealBondMkup)
```
"""
function init_grids!(m::RealBondMkup)
    xscale  = get_setting(m, :xscale)
    xlo     = get_setting(m, :xlo)
    xhi     = get_setting(m, :xhi)
    nx      = get_setting(m, :nx)

    ns      = get_setting(m, :ns)
    λ       = get_setting(m, :λ)

    ehi     = get_setting(m, :ehi)

    grids = OrderedDict()

    # Cash on hand grid
    grids[:xgrid] = Grid(uniform_quadrature(xscale), xlo, xhi, nx, scale = xscale)

    # Skill grid
    lsgrid, sprob, sscale = tauchen86(m[:μ_s].value, m[:σ_s].value, ns, λ)
    swts = (sscale/ns)*ones(ns)
    sgrid = exp.(lsgrid) .+ ehi
    grids[:sgrid] = Grid(sgrid, swts, sscale)

    # Density of skill across skill grid
    grids[:ggrid] = sprob./swts

    # Total grid vectorized across both dimensions
    grids[:sgrid_total] = kron(sgrid, ones(nx))
    grids[:xgrid_total] = kron(ones(ns), grids[:xgrid].points)
    grids[:weights_total] = kron(swts, grids[:xgrid].weights)

    m.grids = grids
end

function model_settings!(m::RealBondMkup)
    default_settings!(m)

    # Defaults
    # Data settings for released and conditional data. Default behavior is to set vintage
    # of data to today's date.
    vint = Dates.format(now(), DSGE_DATE_FORMAT)
    m <= Setting(:data_vintage, vint, true, "vint", "Data vintage")

    saveroot = normpath(joinpath(dirname(@__FILE__), "../../../../","save"))
    datapath = normpath(joinpath(dirname(@__FILE__), "../../../../","save","input_data"))

    m <= Setting(:saveroot, saveroot, "Root of data directory structure")
    m <= Setting(:dataroot, datapath, "Input data directory path")

    # Solution method
    m <= Setting(:solution_method, :klein)

    # Likelihood method
    m <= Setting(:use_chand_recursion, true)

    # Anticipated shocks
    m <= Setting(:n_anticipated_shocks, 0,
                 "Number of anticipated policy shocks")
    # m <= Setting(:n_anticipated_shocks_padding, 20,
                 # "Padding for anticipated policy shocks")

    # Whether one wishes to re-compute βstar
    m <= Setting(:use_last_βstar, false, "Flag to avoid recomputing steady-state.")

    # Number of states and jumps
    m <= Setting(:normalize_distr_variables, true, "Whether or not to perform the
                 normalization of the distributional states in the Klein solution step")
    m <= Setting(:n_predetermined_variables, 0, "Number of predetermined variables after
                 removing the densities. Calculated with the Jacobian normalization.
                 This set to 0 is just the default setting, since it will always be
                 overwritten once the Jacobian is calculated.")

    m <= Setting(:state_indices, 1:5, "Which indices of m.endogenous_states correspond to
                 backward looking state variables")
    m <= Setting(:jump_indices, 6:11, "Which indices of m.endogenous_states correspond to jump
                 variables")

    # Mollifier setting parameters
    m <= Setting(:In, 0.443993816237631, "Normalizing constant for the mollifier")
    m <= Setting(:elo, 0.0, "Lower bound on stochastic consumption commitments")
    m <= Setting(:ehi, 1.0, "Upper bound on stochastic consumption commitments")

    # x: Cash on Hand Grid Setup
    m <= Setting(:nx, 150, "Cash on hand distribution grid points")
    m <= Setting(:xlo, -0.5 - 1.0, "Lower bound on cash on hand")
    m <= Setting(:xhi, 4.0, "Upper Bound on cash on hand")
    m <= Setting(:xscale, get_setting(m, :xhi) - get_setting(m, :xlo), "Size of the xgrid")

    # s: Skill Distribution/ "Units of effective labor" Grid Setup
    m <= Setting(:ns, 2, "Skill distribution grid points")
    m <= Setting(:λ, 3.0, "The λ parameter in the Tauchen distribution calculation")

    # Total grid x*s
    m <= Setting(:n, get_setting(m, :nx) * get_setting(m, :ns), "Total grid size, multiplying
                 across grid dimensions.")

    # Function-valued variables include distributional variables
    m <= Setting(:n_function_valued_backward_looking_states, 1, "Number of function-valued
                 backward looking state variables")
    m <= Setting(:n_backward_looking_distributional_vars, 1, "Number of state variables that are
                 distributional variables.")
    m <= Setting(:n_function_valued_jumps, 1, "Number of function-valued jump variables")
    m <= Setting(:n_jump_distributional_vars, 0, "Number of jump variables that are distributional
                 variables.")

    # Note, these settings assume normalization.
    # The n degrees of freedom removed depends on the distributions/dimensions
    # of heterogeneity that we have discretized over, in this case,
    # cash on hand and the skill distribution. In general the rule of
    # thumb is, remove one degree of freedom for the first endogenous distribution (cash on
    # hand), then one additional degree of freedom for each exogenous distribution (skill
    # distribution). Multiple endogenous distributions only permit removing a single degree
    # of freedom since it is then non-trivial to obtain the marginal distributions.
    m <= Setting(:n_degrees_of_freedom_removed, 2, "Number of degrees of freedom from the
                 distributional variables to remove.")
    n_dof_removed = get_setting(m, :n_degrees_of_freedom_removed)

    ####################################################
    # Calculating the number of backward-looking states
    ####################################################
    n_backward_looking_distr_vars = get_setting(m, :n_backward_looking_distributional_vars)
    m <= Setting(:backward_looking_states_normalization_factor,
                 n_dof_removed*n_backward_looking_distr_vars, "The number of dimensions removed from the
                 backward looking state variables for the normalization.")

    n = get_setting(m, :n)
    n_backward_looking_vars = length(get_setting(m, :state_indices))
    n_backward_looking_function_valued_vars = get_setting(m, :n_function_valued_backward_looking_states)
    n_backward_looking_scalar_vars = n_backward_looking_vars - n_backward_looking_function_valued_vars

    m <= Setting(:n_backward_looking_states, n*n_backward_looking_distr_vars +
                 n_backward_looking_scalar_vars - get_setting(m, :backward_looking_states_normalization_factor),
                 "Number of state variables, in the true sense (fully
                  backward looking) accounting for the discretization across the grid
                  of function-valued variables and the normalization of
                  distributional variables.")

    ##################################
    # Calculating the number of jumps
    ##################################
    n_jump_distr_vars = get_setting(m, :n_jump_distributional_vars)
    m <= Setting(:jumps_normalization_factor,
                 n_dof_removed*n_jump_distr_vars, "The number of dimensions removed from the
                 jump variables for the normalization.")

    n_jump_vars = length(get_setting(m, :jump_indices))
    n_jump_function_valued_vars = get_setting(m, :n_function_valued_jumps)
    n_jump_scalar_vars = n_jump_vars - n_jump_function_valued_vars

    m <= Setting(:n_jumps, n*n_jump_function_valued_vars +
                 n_jump_scalar_vars - get_setting(m, :jumps_normalization_factor),
                 "Number of jump variables (forward looking) accounting for
                  the discretization across the grid of function-valued variables and the
                  normalization of distributional variables.")

    m <= Setting(:n_model_states, get_setting(m, :n_backward_looking_states) + get_setting(m, :n_jumps),
                 "Number of 'states' in the state space model. Because backward and forward
                 looking variables need to be explicitly tracked for the Klein solution
                 method, we have n_states and n_jumps")

    # Reduction
    m <= Setting(:reduce_state_vars, true, "Reduce state variables or not")
    m <= Setting(:reduce_v, true, "Reduce value function or not")
    m <= Setting(:krylov_dim, 5, "Krylov reduction dimension")
    m <= Setting(:F, identity, "Function applied during block arnoldi deflation")
    m <= Setting(:n_knots, 12, "Number of knot points")
    m <= Setting(:c_power, 7, "Amount of bending of knot point locations to make them nonuniform")

end


"""
    phase(loc::String, type::DataType, digit::Int)
Takes a String of numbers returns a typed locus appropriate for PopGen.jl as used in the
`genepop` and `csv` file parsers. Use `type` to specify output type (`Int8` or `Int16`),
and `digit` to specify the number of digits/characters used per allele in a locus.

## Examples
`ph_locus = phase("128114", Int16, 3)`

`map(i -> phase(i, Int16, 3), ["112131", "211112", "001003", "516500"])`

`[phase(i, Int8, 2) for i in ["0101", "0103", "0202", "0103"]]`
"""
function phase(loc::String, type::DataType, digit::Int)
    phased = [join(k) for k in Iterators.partition(loc, digit)]
    typed = parse.(type, phased) |> sort!
    tupled = Tuple(typed)
    iszero(tupled |> collect) && return missing
    return tupled
end

### Variant Call Format parsing ###

"""
    bcf(infile::String)
Load a BCF file into memory as a PopObj object. Population and [optional]
location information need to be provided separately.
- `infile` : path to BCF file
"""
function bcf(infile::String)
    bcf_file = BCF.Reader(open(infile, "r"))

    # get sample names from header
    sample_names = header(bcf_file).sampleID

    # fill in pop/lat/long with missing
    population = fill("missing", length(sample_names))
    loc_xy = Vector{Union{Missing, Float32}}()
    append!(loc_xy, fill(missing, length(sample_names)))

    # get loci names
    locinames = Vector{String}()

    ## array of genotypes
    locus_array = Vector{Vector{Union{Missing, Tuple{Vararg}}}}()

    # get genotypes
    for record in bcf_file
        # fix locus names to be syntax-safe
        chr_safe = replace(BCF.chrom(record), r"\.|\-|\=|\/" => "_")
        chr_safer = replace(chr_safe, "|" => "_")
        pos = BCF.pos(record) |> string
        push!(locinames, chr_safer*"_"*pos)

        # get the genotype information
        geno_raw = [split(i, ('/', '|')) for i in BCF.genotype(record, :, "GT")] |> sort

        # change missing data "." to "-1"
        geno_corr_miss = map(i -> replace(i, "." => "-1"), geno_raw)

        # convert everything to an integer
        geno_int = map(i -> parse.(Int8, i), geno_corr_miss)

        # add 1 to shift genos so 0 is 1 and -1 is 0 etc.
        geno_shift = map(i -> i .+ Int8(1), geno_int)
        geno_final = [replace(i, 0 => missing) for i in geno_shift]
        geno_tuple = [Tuple(i) for i in geno_final]
        push!(locus_array, geno_tuple)
    end
    close(vcf_file)

    # intelligently scan for ploidy
    ploidy = Vector{Int8}()
    @inbounds for i in 1:length(sample_names)
        @inbounds for j in 1:length(locinames)
            locus_array[j][i] === missing && continue   # if missing, go to next locus
            ploid = length(locus_array[j][i])   # if not, get the # of alleles
            push!(ploidy, ploid)    # push that to the ploidy vector
            break   # break out of the loop and begin next sample
        end
    end

    @info "\n$(abspath(infile))
$(length(sample_names)) samples detected
population info must be added <---
$(length(locinames)) loci detected"

    # create loci dataframe
    loci_df = DataFrame([j => k for (j,k) in zip(Symbol.(locinames), locus_array)])

    # create samples df
    samples_df = DataFrame(
        name = sample_names,
        population = population,
        ploidy = ploidy,
        latitude = loc_xy,
        longitude = loc_xy
    )
    PopObj(samples_df, loci_df)
end

### delimited parsing ###

"""
    delimited(infile::String; delim::Union{Char,String,Regex}, digits::Int64 = 3)
Load a delimited-type file into memory as a PopObj object. *There should be no empty cells
in your file*
### Arguments
- `infile` : path to file
- `delim` : delimiter characters. By default auto-parses comma, space, or tab (use ONLY one!)
- `digits` : number of digits denoting each allele (default = 3)
- `marker` : "snp" (default) or "msat"

### File formatting:
- First row is column names in this order:
    1. name
    2. population
    3. longitude
    4. latitude
    5. locus_1_name
    6. locus_2_name
    7. etc...

#### Formatting location data
If location data is missing for a sample (which is ok!), make sure the value is written
as *0*, otherwise there will be transcription errors!

## Example
`lizardsCA = Read.delimited("CA_lizards.csv", digits = 3);`

### Formatting example:
name,population,long,lat,Locus1,Locus2,Locus3   \n
sierra_01,mountain,11.11,-22.22,001001,002002,001001   \n
sierra_02,mountain,11.12,-22.21,001001,001001,001002   \n
snbarb_03,coast,0,0,001001,001001,001002 \n
snbarb_02,coast,11.14,-22.24,001001,001001,001001 \n
snbarb_03,coast,11.15,0,001002,001001,001001 \n
"""
function delimited(
    infile::String;
    delim::Union{Char,String,Regex} = r"\,|\t|\s",
    digits::Int = 3,
    marker = "snp",
)

    # instantiate empty vectors
    indnames = Vector{String}()
    popid = Vector{Union{String,Int32}}()
    ploidy = Vector{Int8}()
    locx = Vector{Union{Missing, Float32}}()
    locy = Vector{Union{Missing, Float32}}()
    sample_geno_array = Array{Vector{Union{Missing, Tuple}},1}()

    # establish types for genotypes
    if lowercase(marker) == "snp"
        geno_type = Int8
    else
        geno_type = Int16
    end

    file = open(infile, "r")

    # get header information
    col_names = map(join, split(readline(file), delim))
    replace!(col_names, "." => "_")
    locinames = col_names[5:end]

    @inbounds for ln in eachline(file)
        tmp = map(join, split(ln, delim))
        # just in case -9 = missing
        replace!(tmp, "-9" => "0"^digits)
        push!(indnames, popfirst!(tmp))
        push!(popid, popfirst!(tmp))
        push!(locx, parse(Float32, popfirst!(tmp)))
        push!(locy, parse(Float32, popfirst!(tmp)))

        # phase the loci
        phasedloci = map(i -> phase(i, geno_type, digits), tmp)
        push!(sample_geno_array, phasedloci)

        # get the ploidy via # of alleles in the first non-missing locus
        sample_ploidy = (skipmissing(phasedloci) |> collect)[1]  |> length
        push!(ploidy, sample_ploidy)

    end
    close(file)

    @info "\n$(abspath(infile))
$(length(indnames)) samples detected
$(length(unique(popid))) populations detected
$(length(locinames)) loci detected"

    # convert the entire thing into a matrix
    geno_matrix = hcat(sample_geno_array...)

    # convert is back into an array of arrays, transposed across loci
    # instead of samples
    loci_array = collect.(eachrow(geno_matrix))

    # replace 0.0 with missing in location
    replace!(locx, 0.0 => missing)
    replace!(locy, 0.0 => missing)

    # genotype DataFrame
    loci_df = DataFrame([j => k for (j,k) in zip(Symbol.(locinames), loci_array)])

    samples_df = DataFrame(
        name = indnames,
        population = popid,
        ploidy = ploidy,
        longitude = locx,
        latitude = locy,
    )

    return PopObj(samples_df, loci_df)
end

const csv = delimited

### GenePop parsing ###
"""
    genepop(infile::String; digits::Int = 3, popsep::String = "POP", marker::String = "snp")
Load a Genepop format file into memory as a PopObj object.
### Arguments
- `infile` : path to Genepop file
- `digits` : number of digits denoting each allele
- `popsep` : word that separates populations in `infile` (default: "POP")
- `marker` : "snp" (default) or "msat"
### File must follow standard Genepop formatting:
- First line is a comment (and skipped)
- Loci are listed after first line as one-per-line without commas or in single comma-separated row
- A line with a particular keyword (default "POP") must delimit populations
- Sample name is followed by a comma and space or tab (", ")
- File is tab or space delimted

## Example

 `waspsNY = genepop("wasp_hive.gen", digits = 3, popsep = "POP")`

### Genepop file example:
Wasp populations in New York \n
Locus1\n
Locus2\n
Locus3\n
POP\n
Oneida_01,  250230 564568 110100\n
Oneida_02,  252238 568558 100120\n
Oneida_03,  254230 564558 090100\n
POP\n
Newcomb_01,  254230 564558 080100\n
Newcomb_02,  000230 564558 090080\n
Newcomb_03,  254230 000000 090100\n
Newcomb_04,  254230 564000 090120\n
"""
function genepop(
    infile::String;
    digits::Int = 3,
    popsep::String = "POP",
    marker::String = "snp",
)
    if lowercase(marker) == "snp"
        geno_type = Int8
    else
        geno_type = Int16
    end

    gpop = split(open(readlines, infile)[2:end], popsep)

    @info "\n$(abspath(infile))
$(sum(length.(gpop[2:end]))) samples detected
$(length(gpop)-1) populations detected using \"$popsep\" seperator
$(length(gpop[1])) loci detected"

    if length(gpop[1]) == 1     # loci horizontally stacked.
        locinames = strip.(split(gpop[1] |> join, ",") |> Vector{String})
        replace!(locinames, "." => "_")
    else                        # loci vertically stacked
        locinames = replace(gpop[1], "." => "_")
    end

    sample_meta_names = [:name, :population, :ploidy, :latitude, :longitude]
    sample_meta_array = [
        Vector{String}(),
        Vector{Union{String,Int}}(),
        Vector{Int8}(),
        Vector{Union{Missing,Float32}}(),
        Vector{Union{Missing, Float32}}()
    ]

    sample_geno_array = Vector{Vector{Union{Missing, Tuple}}}()

    for i = 2:length(gpop)
        for j in gpop[i]
            samplerow = map(join, split(strip(j), r"\,\t|\,\s|\s|\t"))

            # just in case missing genotypes are coded as -9 (msats)
            replace!(samplerow, "-9" => "0"^digits)

            # sample name and population
            push!(sample_meta_array[1], popfirst!(samplerow))
            push!(sample_meta_array[2], i-1)

            # phase the loci into tuples
            converted_row = map(k -> phase(k, geno_type, digits), samplerow)
            push!(sample_geno_array, converted_row)

            # get the ploidy via # of alleles in the first non-missing locus
            push!(sample_meta_array[3], length((skipmissing(converted_row) |> collect)[1]))

            # long and lat
            push!(sample_meta_array[4], missing) ; push!(sample_meta_array[5], missing)
        end
    end

    # convert the entire thing into a matrix
    geno_matrix = hcat(sample_geno_array...)

    # convert back into an array of arrays of loci instead of samples
    loci_array = collect.(eachrow(geno_matrix))

    # create the two dataframes necessary for a PopObj
    samples = DataFrame([j => k for (j,k) in zip(sample_meta_names, sample_meta_array)])
    loci = DataFrame([j => k for (j,k) in zip(Symbol.(locinames), loci_array)])

    return PopObj(samples, loci)
end

#= Alternative line-by-line genepop reader

# Alternative line-by-line reader
## NOT EXPORTED

function gpop2(
    infile::String;
    digits::Int = 3,
    popsep::String = "POP",
    marker::String = "snp",
)

    if lowercase(marker) == "snp"
        geno_type = Int8
    else
        geno_type = Int16
    end

    sample_meta_names = [:name, :population, :ploidy, :latitude, :longitude]
    sample_meta_array = [
        Vector{String}(),
        Vector{Union{String,Int}}(),
        Vector{Int8}(),
        Vector{Union{Missing,Float32}}(),
        Vector{Union{Missing, Float32}}()
    ]

    locinames = Vector{String}()
    sample_geno_array = Vector{Vector{Union{Missing, Tuple}}}()
    popcount = 1

    gpop = open(infile, "r")

    # skip first line
    readline(gpop)

    # test for loci horizontal/vertical formatting
    locus_name = readline(gpop)
    if occursin(",", locus_name) == true
        append!(locinames, strip.(split(locus_name |> join, ",")))
        replace!(locinames, "." => "_")
    else
        while locus_name != popsep
            push!(locinames, locus_name)
            locus_name = readline(gpop)
        end
    end

    for ln in readlines(gpop)
        if ln == popsep
            popcount += 1
            continue
        else
            samplerow = map(join, split(strip(ln), r"\,\t|\,\s|\s|\t"))

            # just in case missing genotypes are coded as -9 (msats)
            replace!(samplerow, "-9" => "0"^digits)

            # sample name and population
            push!(sample_meta_array[1], popfirst!(samplerow))
            push!(sample_meta_array[2], popcount)

            # phase the loci into tuples
            converted_row = map(k -> phase(k, geno_type, digits), samplerow)
            push!(sample_geno_array, converted_row)

            # get the ploidy via # of alleles in the first non-missing locus
            push!(sample_meta_array[3], length((skipmissing(converted_row) |> collect)[1]))

            # long and lat
            push!(sample_meta_array[4], missing) ; push!(sample_meta_array[5], missing)
        end
    end

    @info "\n$(abspath(infile))
$(length(sample_meta_array[1])) samples detected
$(popcount) populations detected using \"$popsep\" seperator
$(length(locinames)) loci detected"
    # convert the entire thing into a matrix
    geno_matrix = hcat(sample_geno_array...)

    # convert back into an array of arrays of loci instead of samples
    loci_array = collect.(eachrow(geno_matrix))

    # create the two dataframes necessary for a PopObj
    samples = DataFrame([j => k for (j,k) in zip(sample_meta_names, sample_meta_array)])
    loci = DataFrame([j => k for (j,k) in zip(Symbol.(locinames), loci_array)])

    return PopObj(samples, loci)
end
=#

### VCF parsing ###

"""
    vcf(infile::String)
Load a VCF file into memory as a PopObj object. Population and [optional]
location information need to be provided separately.
- `infile` : path to VCF file
"""
function vcf(infile::String)
    vcf_file = VCF.Reader(open(infile, "r"))

    # get sample names from header
    sample_names = header(vcf_file).sampleID

    # fill in pop/lat/long with missing
    population = fill("missing", length(sample_names))
    loc_xy = Vector{Union{Missing, Float32}}()
    append!(loc_xy, fill(missing, length(sample_names)))

    # get loci names
    locinames = Vector{String}()

    ## array of genotypes
    locus_array = Vector{Vector{Union{Missing, Tuple{Vararg}}}}()

    # get genotypes
    for record in vcf_file
        # fix locus names to be syntax-safe
        chr_safe = replace(VCF.chrom(record), r"\.|\-|\=|\/" => "_")
        chr_safer = replace(chr_safe, "|" => "_")
        pos = VCF.pos(record) |> string
        push!(locinames, chr_safer*"_"*pos)

        # get the genotype information
        geno_raw = [split(i, ('/', '|')) for i in VCF.genotype(record, :, "GT")] |> sort

        # change missing data "." to "-1"
        geno_corr_miss = map(i -> replace(i, "." => "-1"), geno_raw)

        # convert everything to an integer
        geno_int = map(i -> parse.(Int8, i), geno_corr_miss)

        # add 1 to shift genos so 0 is 1 and -1 is 0 etc.
        geno_shift = map(i -> i .+ Int8(1), geno_int)
        geno_final = [replace(i, 0 => missing) for i in geno_shift]
        geno_tuple = [Tuple(i) for i in geno_final]
        push!(locus_array, geno_tuple)
    end
    close(vcf_file)

    # intelligently scan for ploidy
    ploidy = Vector{Int8}()
    @inbounds for i in 1:length(sample_names)
        @inbounds for j in 1:length(locinames)
            locus_array[j][i] === missing && continue   # if missing, go to next locus
            ploid = length(locus_array[j][i])   # if not, get the # of alleles
            push!(ploidy, ploid)    # push that to the ploidy vector
            break   # break out of the loop and begin next sample
        end
    end

    @info "\n$(abspath(infile))
$(length(sample_names)) samples detected
population info must be added <---
$(length(locinames)) loci detected"

    # create loci dataframe
    loci_df = DataFrame([j => k for (j,k) in zip(Symbol.(locinames), locus_array)])

    # create samples df
    samples_df = DataFrame(
        name = sample_names,
        population = population,
        ploidy = ploidy,
        latitude = loc_xy,
        longitude = loc_xy
    )
    PopObj(samples_df, loci_df)
end

"""
    PopGen.read(infile::String; kwargs...)
Wraps `csv()`, `genepop()`, `bcf()`, and `vcf()` to read a file in as a `PopObj`. File type is
inferred from the file extension:
- delimited: .csv | .tsv | .txt
- genepop: .gen | .genepop
- variant call format: .vcf
This function uses the same keyword arguments (and defaults) as the file importing
functions it wraps; please see their respective docstrings in the Julia help console.
(e.g. `?genepop`) for specific usage details.

## Examples
`read("cavernous_assfish.gen", markers = "msat", digits = 3)`

`read("bos_tauros.csv")`

`read("juglans_nigra.vcf")`
"""
function Base.read(infile::String; kwargs...)
    kwargs = Dict(kwargs)
    ext = split(infile, ".")[end]
    if ext == "gen" || ext == "genepop"
        if haskey(kwargs, :digits)
            digits = kwargs[:digits]
        else
            digits = 3
        end
        if haskey(kwargs, :popsep)
            popsep = kwargs[:popsep]
        else
            popsep = "POP"
        end
        if haskey(kwargs, :marker)
            marker = kwargs[:marker]
        else
            marker = "snp"
        end
        return genepop(infile, digits = digits, popsep = popsep, marker = marker)

    elseif ext == "csv" || ext == "txt" || ext == "tsv"
        if haskey(kwargs, :delim)
            delim = kwargs[:delim]
        else
            delim = r"\,|\s|\t"
        end
        if haskey(kwargs, :digits)
            digits = kwargs[:digits]
        else
            digits = 3
        end
        if haskey(kwargs, :marker)
            marker = kwargs[:marker]
        else
            marker = "snp"
        end
        return csv(infile, delim = delim, digits = digits, marker = marker)

    elseif ext == "vcf" || ext == "bcf"
            ext == "vcf" && return vcf(infile)
            ext == "bcf" && return bcf(infile)
    else
        @error "file type not recognized by filename extension \n delimited: .csv | .tsv | .txt \n genepop: .gen | .genepop \n variant call formant: .bcf | .vcf"
    end
end


#module ParallelMatrix

export SparseBinMatrix, ParallelLogic, ParallelSBM, balanced_parallelsbm
export sort_hilbert
export AtA_mul_B!

using Compat

type SparseBinMatrix
  m::Int64
  n::Int64
  mrange::UnitRange{Int32}
  nrange::UnitRange{Int32}
  ## data
  rows::Vector{Int32}
  cols::Vector{Int32}
end

function SparseBinMatrix(m, n, rows::Vector{Int32}, cols::Vector{Int32})
  length(rows) == length(cols) || throw(DimensionMismatch("length(rows) must equal length(cols)"))
  return SparseBinMatrix(m, n, minimum(rows):maximum(rows), minimum(cols):maximum(cols), rows, cols)
end
SparseBinMatrix(rows::Vector{Int32}, cols::Vector{Int32}) = SparseBinMatrix(maximum(rows), maximum(cols), rows, cols)
SparseBinMatrix(rows::Vector{Int64}, cols::Vector{Int64}) = SparseBinMatrix(convert(Vector{Int32}, rows), convert(Vector{Int32}, cols))

import Base.getindex
function getindex(sbm::SparseBinMatrix, rows::Vector{Bool})
  length(rows) == sbm.m || throw(DimensionMismatch("length(rows) must equal size(sbm,1)"))
  idx = rows[sbm.rows]
  colidx = sbm.cols[idx]
  rsum   = cumsum(rows)
  rowidx = rsum[ sbm.rows[idx] ]
  out = SparseBinMatrix(sum(rows), sbm.n, one(Int32):convert(Int32,sum(rows)), sbm.nrange, rowidx, colidx)
  return out
end

function getindex(sbm::SparseBinMatrix, rows::Vector{Bool}, cols::Colon)
  return getindex(sbm, rows)
end

function getindex(sbm::SparseBinMatrix, rows::Vector{Bool}, cols::UnitRange{Int64})
  return getindex(sbm, rows)
end

import Base.ndims
ndims(sbm::SparseBinMatrix) = 2

type ParallelLogic
  ## for parallel compute
  mblocks::Vector{UnitRange{Int32}}
  nblocks::Vector{UnitRange{Int32}}
  mblock_order::Vector{Int64}
  nblock_order::Vector{Int64}

  ## local vectors
  localm::Vector{Float64}
  localn::Vector{Float64}

  ## semaphores
  tmp::SharedVector{Float64}
  sems::Vector{SharedArray{UInt32,1}}

  ## constructor for non-shared elements
  ParallelLogic(mblocks, nblocks, mblock_order, nblock_order, localm, localn) = new(mblocks, nblocks, mblock_order, nblock_order, localm, localn)
  ParallelLogic(mblocks, nblocks, mblock_order, nblock_order, localm, localn, tmp, sems) = new(mblocks, nblocks, mblock_order, nblock_order, localm, localn, tmp, sems)
end

nonshared(A::ParallelLogic) = ParallelLogic(A.mblocks, A.nblocks, A.mblock_order, A.nblock_order, A.localm, A.localn)

type ParallelSBM
  m::Int64
  n::Int64
  pids::Vector{Int64}
  sbms::Vector{RemoteRef}  ## SparseBinMatrix
  logic::Vector{RemoteRef} ## ParallelLogic

  numblocks::Int                ## number of blocks, each has semaphore
  tmp::SharedVector{Float64}    ## for storing middle vector in A'A
  sh1::SharedVector{Float64}    ## length(sh1) = A.n
  sh2::SharedVector{Float64}    ## length(sh2) = A.n
  sems::Vector{SharedArray{UInt32,1}} ## semaphores

  ## constructors
  ParallelSBM(m, n, pids, sbms, logic, numblocks) = new(m, n, pids, sbms, logic, numblocks)
  ParallelSBM(m, n, pids, sbms, logic, numblocks, tmp, sh1, sh2, sems) = new(m, n, pids, sbms, logic, numblocks, tmp, sh1, sh2, sems)
end

nonshared(A::ParallelSBM) = ParallelSBM(A.m, A.n, A.pids, A.sbms, A.logic, A.numblocks)

function make_sems(numblocks::Int, pids::Vector{Int})
  sems = SharedArray{UInt32,1}[SharedArray(UInt32, 16, pids=pids) for i=1:numblocks]
  for sem in sems
    sem_init(sem)
  end
  return sems
end

function block_order(counts, blocks)
  bcounts = counts[blocks]
  return blocks[sortperm(bcounts, rev=true)]
end

function ParallelSBM(rows::Vector{Int}, cols::Vector{Int}, pids::Vector{Int}, weights=ones(length(pids)), m=convert(Int32, maximum(rows)), n=convert(Int32, maximum(cols)), numblocks=length(pids)*2)
  return ParallelSBM(convert(Vector{Int32}, rows), convert(Vector{Int32}, cols), pids, weights=weights, m=m, n=n, numblocks=numblocks)
end

function ParallelSBM(rows::Vector{Int32}, cols::Vector{Int32}, pids::Vector{Int}=Int[]; weights=ones(length(pids)), m=maximum(rows), n=maximum(cols), numblocks=length(pids)*2 )
  length(rows) == length(cols) || throw(DimensionMismatch("length(rows) must equal length(cols)"))

  shtmp = SharedArray(Float64, convert(Int, m), pids=pids)
  sh1   = SharedArray(Float64, convert(Int, n), pids=pids)
  sh2   = SharedArray(Float64, convert(Int, n), pids=pids)
  sems  = make_sems(numblocks, pids)
  ps = ParallelSBM(m, n, pids, RemoteRef[], RemoteRef[], numblocks, shtmp, sh1, sh2, sems)
  ranges  = make_lengths(length(rows), weights)
  mblocks = make_blocks(m, convert(Int32, numblocks) )
  nblocks = make_blocks(n, convert(Int32, numblocks) )
  mblock_grid = zeros(Int, numblocks, length(pids))
  nblock_grid = zeros(Int, numblocks, length(pids))
  for i in 1:length(pids)
    sbm = SparseBinMatrix(m, n, rows[ranges[i]], cols[ranges[i]])
    sbm_ref = @spawnat pids[i] fetch(sbm)
    push!(ps.sbms, sbm_ref)
    mblock_grid[:,i] = [! isempty(intersect(sbm.mrange, i)) for i in mblocks]
    nblock_grid[:,i] = [! isempty(intersect(sbm.nrange, i)) for i in nblocks]
  end
  mblock_counts = vec(sum(mblock_grid, 2))
  nblock_counts = vec(sum(nblock_grid, 2))
  for i in 1:length(pids)
    mb = block_order(mblock_counts, find(mblock_grid[:,i]))
    nb = block_order(nblock_counts, find(nblock_grid[:,i]))
    pl_ref = @spawnat pids[i] ParallelLogic(mblocks, nblocks, mb, nb, zeros(m), zeros(n), ps.tmp, ps.sems )
    push!(ps.logic, pl_ref)
  end
  return ps
end

## copies ParallelSBM to new pids
function copyto(F::ParallelSBM, pids::Vector{Int})
  length(pids) == length(F.pids) || throw(DimensionMismatch("length(pids)=$(length(pids)) must equal length(F.pids)=$(length(F.pids))"))

  shtmp = SharedArray(Float64, size(F, 1), pids=pids)
  sh1   = SharedArray(Float64, size(F, 2), pids=pids)
  sh2   = SharedArray(Float64, size(F, 2), pids=pids)
  sems  = make_sems(F.numblocks, pids)
  ps    = ParallelSBM(F.m, F.n, pids, RemoteRef[], RemoteRef[], F.numblocks, shtmp, sh1, sh2, sems)

  for i in 1:length(F.sbms)
    push!(ps.sbms, @spawnat(pids[i], fetch(F.sbms[i])) )

    l = fetch(@spawnat F.pids[i] nonshared(fetch(F.logic[i])))
    logic_ref = @spawnat pids[i] ParallelLogic(l.mblocks, l.nblocks, l.mblock_order, l.nblock_order, zeros(F.m), zeros(F.n), ps.tmp, ps.sems)
    push!(ps.logic, logic_ref)
  end
  return ps
end

sem_init(x::SharedArray)    = ccall(:sem_init, Cint, (Ptr{Void}, Cint, Cuint), x, 1, one(UInt32))
sem_wait(x::SharedArray)    = ccall(:sem_wait, Cint, (Ptr{Void},), x)
sem_trywait(x::SharedArray) = ccall(:sem_trywait, Cint, (Ptr{Void},), x)
sem_post(x::SharedArray)    = ccall(:sem_post, Cint, (Ptr{Void},), x)

gmean(x) = prod(x) ^ (1 / length(x))
pretty(x) = "[" * join([@sprintf("%.3f", i) for i in x], ", ") * "]"

function balanced_parallelsbm(rows::Vector{Int32}, cols::Vector{Int32}, pids::Vector{Int}; numblocks=length(pids)*2, niter=4, ntotal=30, keeplast=4, verbose=false)
  weights = ones(length(pids))
  y = SharedArray(Float64, convert(Int, maximum(rows)) )
  x = SharedArray(Float64, convert(Int, maximum(cols)) )
  local psbm
  for i = 1:niter
    psbm   = ParallelSBM(rows, cols, pids, numblocks=numblocks, weights=weights)
    times  = A_mul_B!_time(y, psbm, x, ntotal)
    ctimes = vec(mean(times[:,end-keeplast:end], 2))
    meantime  = gmean(ctimes)
    weights .*= (meantime ./ ctimes) .^ (1/(1+0.2*i))
    weights   = weights ./ sum(weights)
    verbose && println("$i. ctimes  = ", pretty(ctimes) )
    verbose && println("$i. weights = ", pretty(weights) )
  end
  return psbm
end

function make_blocks(n::Int32, nblocks::Int32)
  bsize   = 8 * ceil(n / nblocks / 8)
  if (bsize-1) * nblocks > n
    bsize = ceil(n / nblocks)
  end
  return  UnitRange{Int32}[convert(Int32, 1+(i-1)*bsize) : convert(Int32, min(n, i*bsize))
                           for i in 1:nblocks]
end

function make_lengths(total::Int, weights)
    wnorm  = weights ./ sum(weights)
    n      = convert(Vector{Int64}, round(wnorm * total))
    excess = sum(n) - total
    i = 0
    while excess != 0
        n[i+1] -= sign(excess)
        excess  = sum(n) - total
        i = (i + 1) % length(n)
    end
    ranges = UnitRange{Int}[]
    k = 1
    for len in n
      push!(ranges, k:(k+len-1))
      k += len
    end
    return ranges
end

## waits 100x to see that x[n] == value
function busywait(x::SharedArray{Int,1}, n::Int, value::Int, ntimes=100)
  for i=1:ntimes
    if x[n] != value
      return false
    end
  end
  return true
end

import Base.size
import Base.isempty
size(X::SparseBinMatrix) = (X.m, X.n)
size(X::SparseBinMatrix, d::Int) = d==1 ? X.m : X.n
isempty(X::SparseBinMatrix)      = X.m == 0 || X.n == 0

size(X::ParallelSBM) = (X.m, X.n)
size(X::ParallelSBM, d::Int) = d==1 ? X.m : X.n
isempty(X::ParallelSBM)      = X.m == 0 || X.n == 0

import Base.A_mul_B!
import Base.At_mul_B!
import Base.At_mul_B
import Base.*

## multiplication: y = A * x
*{Tx}(A::SparseBinMatrix, x::AbstractArray{Tx,1}) = (y = zeros(Tx, A.m); A_mul_B!(y, A, x); y)
*{Tx}(A::ParallelSBM,     x::AbstractArray{Tx,1}) = (y = zeros(Tx, A.m); A_mul_B!(y, A, x); y)

function A_mul_B!{Tx}(y::AbstractArray{Tx,1}, A::SparseBinMatrix, x::AbstractArray{Tx,1})
    A.n == length(x) || throw(DimensionMismatch("A.n=$(A.n) must equal length(x)=$(length(x))"))
    A.m == length(y) || throw(DimensionMismatch("A.m=$(A.m) must equal length(y)=$(length(y))"))
    fill!(y, zero(Tx) )
    rows = A.rows
    cols = A.cols
    @inbounds for i = 1:length(rows)
        y[rows[i]] += x[cols[i]]
    end
    return
end

## multiplication: y = A' * x
At_mul_B{Tx}(A::SparseBinMatrix, x::AbstractArray{Tx,1}) = (y = zeros(Tx, A.n); At_mul_B!(y, A, x); y)

function At_mul_B!{Tx}(y::AbstractArray{Tx,1}, A::SparseBinMatrix, x::AbstractArray{Tx,1})
    A.n == length(y) || throw(DimensionMismatch("A.n=$(A.n) must equal length(y)=$(length(y))"))
    A.m == length(x) || throw(DimensionMismatch("A.m=$(A.m) must equal length(x)=$(length(x))"))
    fill!(y, zero(Tx) )
    rows = A.rows
    cols = A.cols
    @inbounds for i = 1:length(rows)
        y[cols[i]] += x[rows[i]]
    end
    return
end

function A_mul_B!{Tx}(y::SharedArray{Tx,1}, A::ParallelSBM, x::SharedArray{Tx,1})
  A.n == length(x) || throw(DimensionMismatch("A.n=$(A.n) must equal length(x)=$(length(x))"))
  A.m == length(y) || throw(DimensionMismatch("A.m=$(A.m) must equal length(y)=$(length(y))"))
  y[1:end] = zero(Tx)
  np = length(A.pids)
  @sync begin
    for p in 1:length(A.pids)
      pid = A.pids[p]
      if pid != myid() || np == 1
        @async begin
          remotecall_wait(pid, partmul_ref, y, A.sbms[p], A.logic[p], x)
        end
      end
    end
  end
  ## done
end

function At_mul_B!{Tx}(y::SharedArray{Tx,1}, A::ParallelSBM, x::SharedArray{Tx,1})
  A.m == length(x) || throw(DimensionMismatch("A.m=$(A.m) must equal length(x)=$(length(x))"))
  A.n == length(y) || throw(DimensionMismatch("A.n=$(A.n) must equal length(y)=$(length(y))"))
  y[1:end] = zero(Tx)
  np = length(A.pids)
  @sync begin
    for p in 1:length(A.pids)
      pid = A.pids[p]
      if pid != myid() || np == 1
        @async begin
          remotecall_wait(pid, partmul_t_ref, y, A.sbms[p], A.logic[p], x)
        end
      end
    end
  end
end

import Base.*

function *{Tx}(A::ParallelSBM, x::AbstractArray{Tx,1})
  A.n == length(x) || throw(DimensionMismatch("A.n=$(A.n) must equal length(x)=$(length(x))"))
  A.sh1[1:end] = x
  A_mul_B!(A.tmp, A, A.sh1)
  y = zeros(Tx, length(A.tmp))
  y[1:end] = A.tmp
  return y
end

function At_mul_B{Tx}(A::ParallelSBM, x::AbstractArray{Tx,1})
  A.m == length(x) || throw(DimensionMismatch("A.m=$(A.m) must equal length(x)=$(length(x))"))
  A.tmp[1:end] = x
  At_mul_B!(A.sh1, A, A.tmp)
  y = zeros(Tx, A.n)
  y[1:end] = A.sh1
  return y
end

function prod_copy!{Tx}(y::SharedArray{Tx,1}, v::Tx, x::SharedArray{Tx,1})
  length(y) == length(x) || throw(DimensionMismatch("length(y)=$(length(y)) must equal length(x)=$(length(x))"))
  @inbounds @simd for i = 1:length(y)
    y[i] = v * x[i]
  end
  return nothing
end

## computes A'A*x + lambda*x
function AtA_mul_B!{Tx}(y::SharedArray{Tx,1}, A::ParallelSBM, x::SharedArray{Tx,1}, lambda::Float64)
  A.n == length(x) || throw(DimensionMismatch("A.n=$(A.n) must equal length(x)=$(length(x))"))
  A.n == length(y) || throw(DimensionMismatch("A.n=$(A.n) must equal length(y)=$(length(y))"))
  tmp = A.tmp
  tmp[1:end] = zero(Tx)
  np = length(A.pids)
  ## doing tmp = A * x
  @sync begin
    for p in 1:length(A.pids)
      pid = A.pids[p]
      if pid != myid() || np == 1
        @async begin
          remotecall_wait(pid, partmul_ref, tmp, A.sbms[p], A.logic[p], x)
        end
      end
    end
    prod_copy!(y, lambda, x)
  end
  ## doing y += A' * tmp
  @sync begin
    for p in 1:length(A.pids)
      pid = A.pids[p]
      if pid != myid() || np == 1
        @async begin
          remotecall_wait(pid, partmul_t_ref, y, A.sbms[p], A.logic[p], tmp)
        end
      end
    end
  end
  return nothing
end

function A_mul_B!_time{Tx}(y::SharedArray{Tx,1}, A::ParallelSBM, x::SharedArray{Tx,1}, ntimes::Int)
  A.n == length(x) || throw(DimensionMismatch("A.n=$(A.n) must equal length(x)=$(length(x))"))
  A.m == length(y) || throw(DimensionMismatch("A.m=$(A.m) must equal length(y)=$(length(y))"))
  ptime = zeros(length(A.pids), ntimes)
  for i = 1:ntimes
    y[1:end] = zero(Tx)
    ## clearing warmup results
    @sync begin
      for p in 1:length(A.pids)
        pid = A.pids[p]
        if pid != myid() || length(A.pids) == 1
          @async begin
            ptime[p, i] += remotecall_fetch(pid, partmul_time, y, A.sbms[p], A.logic[p], x)
          end
        end
      end
    end
  end
  return ptime
end

function copy!{Tx}(to::AbstractArray{Tx,1}, from::AbstractArray{Tx,1}, range)
  @inbounds @simd for i in range
    to[i] = from[i]
  end
  return nothing
end

function add!{Tx}(to::AbstractArray{Tx,1}, from::AbstractArray{Tx,1}, range)
  @inbounds @simd for i in range
    to[i] += from[i]
  end
  return nothing
end

function rangefill!{Tx}(x::AbstractArray{Tx,1}, v::Tx, range)
  @inbounds @simd for i in range
    x[i] = v
  end
  return nothing
end

function partmul_time{Tx}(y::SharedArray{Tx,1}, Aref::RemoteRef, logicref::RemoteRef, x::SharedArray{Tx,1})
  A     = fetch(Aref)::SparseBinMatrix
  logic = fetch(logicref)::ParallelLogic
  tic();
  partmul(y, A, logic, x)
  return toq()
end

function partmul_ref{Tx}(y::SharedArray{Tx,1}, Aref::RemoteRef, logicref::RemoteRef, x::SharedArray{Tx,1})
  A     = fetch(Aref)::SparseBinMatrix
  logic = fetch(logicref)::ParallelLogic
  partmul(y, A, logic, x)
end

function partmul_t_ref{Tx}(y::SharedArray{Tx,1}, Aref::RemoteRef, logicref::RemoteRef, x::SharedArray{Tx,1})
  A     = fetch(Aref)::SparseBinMatrix
  logic = fetch(logicref)::ParallelLogic
  partmul_t(y, A, logic, x)
end

## assumes sizes are correct
function partmul{Tx}(y::SharedArray{Tx,1}, A::SparseBinMatrix, logic::ParallelLogic, x::SharedArray{Tx,1})
  ylocal = logic.localm
  xlocal = logic.localn
  rangefill!(ylocal, zero(Tx), A.mrange)
  copy!(xlocal, x, A.nrange)
  ## standard y = A * x
  rows = A.rows
  cols = A.cols
  @inbounds for i = 1:length(rows)
      ylocal[rows[i]] += xlocal[cols[i]]
  end
  ## adding the result to shared array
  addshared!(y, ylocal, logic.sems, logic.mblocks, logic.mblock_order, A.mrange)
  return nothing
end

## assumes sizes are correct
function partmul_t{Tx}(y::SharedArray{Tx,1}, A::SparseBinMatrix, logic::ParallelLogic, x::SharedArray{Tx,1})
  ylocal = logic.localn
  xlocal = logic.localm
  rangefill!(ylocal, zero(Tx), A.nrange)
  copy!(xlocal, x, A.mrange)
  ## standard y = A' * x
  rows = A.cols
  cols = A.rows
  @inbounds for i = 1:length(rows)
      ylocal[rows[i]] += xlocal[cols[i]]
  end
  ## adding the result to shared array
  addshared!(y, ylocal, logic.sems, logic.nblocks, logic.nblock_order, A.nrange)
  return nothing
end

## does y += x
function addshared!{Tx}(y::SharedArray{Tx,1}, x::AbstractArray{Tx,1}, sems, ranges, order, yrange)
  blocks = copy(order)
  nblocks = length(blocks)
  pid = myid()
  while true
    i = findfirst(blocks)
    i <= 0 && return nothing ## done
    for j = i:nblocks
      block = blocks[j]
      block == 0 && continue
      ## try to get a lock
      if sem_trywait( sems[block] ) < 0
        ## didn't get lock
        continue
      end
      ## copying result to shared array
      add!(y, x, intersect(ranges[blocks[j]], yrange) )
      sem_post( sems[block] )
      blocks[j] = 0
      break
    end
  end
  return nothing
end

######## parallel operations on Frefs
function solve_cg2(Frefs::Vector{RemoteRef}, rhs::Matrix{Float64}, lambda::Float64; tol=1e-6, maxiter=size(rhs,1))
  beta = zeros(size(rhs,1), size(rhs,2))
  D    = size(rhs,2)
  i    = 1
  # function to produce the next work item from the queue.
  # in this case it's just an index.
  nextidx() = (idx=i; i+=1; idx)
  @sync begin
    for ref in Frefs
      @async begin
        while true
          idx = nextidx()
          idx > D && break
          beta[:,idx] = remotecall_fetch(ref.where, cg_AtA_ref, ref, rhs[:,idx], lambda, tol, maxiter)
        end
      end
    end
  end
  return beta
end

function A_mul_B_ref(Fref::RemoteRef, x::AbstractVector{Float64})
  F = fetch(Fref)
  return F * x
end

function At_mul_B_ref(Fref::RemoteRef, x::AbstractVector{Float64})
  F = fetch(Fref)
  return At_mul_B(F, x)
end

## computes y = F * x in parallel (along columns of x)
function Frefs_mul_B(Frefs::Vector{RemoteRef}, x::Matrix{Float64})
  m, n = fetch( @spawnat Frefs[1].where size(fetch(Frefs[1])) )
  n == size(x, 1) || throw(DimensionMismatch("Frefs.n=$(n) must equal length(x)=$(length(x))"))
  y = zeros(m, size(x,2))
  D = size(x, 2)
  i = 1
  nextidx() = (idx=i; i+=1; idx)
  @sync begin
    for ref in Frefs
      @async begin
        while true
          idx = nextidx()
          idx > D && break
          y[:,idx] = remotecall_fetch(ref.where, A_mul_B_ref, ref, x[:,idx])
        end
      end
    end
  end
  return y
end

## computes y = F' * x in parallel (along columns of x)
function Frefs_t_mul_B(Frefs::Vector{RemoteRef}, x::Matrix{Float64})
  m, n = fetch( @spawnat Frefs[1].where size(fetch(Frefs[1])) )
  m == size(x, 1) || throw(DimensionMismatch("Frefs.m=$(m) must equal length(x)=$(length(x))"))
  y = zeros(n, size(x,2))
  D = size(x, 2)
  i = 1
  nextidx() = (idx=i; i+=1; idx)
  @sync begin
    for ref in Frefs
      @async begin
        while true
          idx = nextidx()
          idx > D && break
          y[:,idx] = remotecall_fetch(ref.where, At_mul_B_ref, ref, x[:,idx])
        end
      end
    end
  end
  return y
end

######## Hilbert ordering ########

## convert (x,y) to d
function xy2d(n, x, y)
    local rx::Bool, ry::Bool
    d = 0::Int64
    s = div(n, 2)
    while s > 0
        rx = (x & s) > 0;
        ry = (y & s) > 0;
        d += s * s * ((3 * rx) $ ry);
        x, y = rot(s, x, y, rx, ry);
        s = div(s, 2)
    end
    return d
end

function rot(n, x, y, rx::Bool, ry::Bool)
    if ry == false
        if rx == true
            x = n - 1 - x
            y = n - 1 - y
        end
        return (y, x)
    end
    return (x, y)
end

function sort_hilbert(rows, cols)
  maxrc = max(maximum(rows), maximum(cols))
  n = 2 ^ round(Int, ceil(log2( maxrc )))
  h = zeros(Int, length(rows))
  @inbounds @simd for i = 1:length(h)
    h[i] = xy2d(n, rows[i] - 1, cols[i] - 1)
  end
  hsorted = sortperm(h)
  return rows[hsorted], cols[hsorted]
end


"""
    Layer(transformations::Tuple, data, positional::Tuple, named::NamedTuple)

Algebraic object encoding a single layer of a visualization. It is composed of a dataset,
positional and named arguments, as well as transformations to be applied to those.
`Layer` objects can be multiplied, yielding a novel `Layer` object, or added,
yielding a [`AlgebraOfGraphics.Layers`](@ref) object.
"""
struct Layer
    transformations::Tuple
    data::Any
    positional::Tuple
    named::NamedTuple
end

Layer(transformations::Tuple=()) = Layer(transformations, nothing, (), (;))

data(df) = Layer((), columns(df), (), (;))
mapping(args...; kwargs...) = Layer((), nothing, args, values(kwargs))

function Base.:*(l1::Layer, l2::Layer)
    t1, t2 = l1.transformations, l2.transformations
    d1, d2 = l1.data, l2.data
    p1, p2 = l1.positional, l2.positional
    n1, n2 = l1.named, l2.named
    transformations = (t1..., t2...)
    data = isnothing(d2) ? d1 : d2
    positional =(p1..., p2...)
    named = merge(n1, n2)
    return Layer(transformations, data, positional, named)
end


using WAV

struct Sound
  samples
  sr
end

"""
  Constructor for a `Sound` object
    Sound(fname)

  # Args
  * **fname** The filename of the wav file to load

  # Returns
  A `Sound` object containing the samples and the the sampling rate

  Examples
  =========

  s = Sound("sound.wav")
"""
function Sound(fname)
  s, sr = wavread(fname)
  s = vec(s)
  return Sound(s, sr)
end


"""
	EmptyNetwork <: AbstractNetwork

A network with no connections.

# Fields
- N              :: Integer -> Number of nodes
"""
struct EmptyNetwork <: AbstractNetwork
	N :: Integer
	_props :: NetworkProperties
end

"""
	EmptyNetwork(N::Integer)

Create an empty network with `N` nodes
"""
function EmptyNetwork(N::Integer)
	(N <= 0) && throw(ArgumentError("Number of nodes must be a positive integer!"))
	EmptyNetwork(N, NetworkProperties(N, false))
end

show(network::EmptyNetwork) = println("Empty Network\n- N = $(network.N)")


## A small driver file to run the 4
# synthetic data tests in parallel
# data calibrated to the RossmanKaggleData
#Run like this
#julia -p 4 -L syntheticDataHarness.jl test_SyntheticData.jl numRuns outPathStub d usePoisson
#passed arguments
	#ARGS[1] numRuns
	#ARGs[2] partial path for output.
	#ARGS[3] d
	#ARGS[4] bool for using the poisson amount of data assumption.

const numRuns = parse(Int, ARGS[1])
const spath = ARGS[2]
const param_path = "../RossmanKaggleData/CleanedData/"
const d = parse(Int, ARGS[3])
const usePoisson = parse(Bool, ARGS[4])

const s = .95
K_grid = vcat(round.(Int, 2 .^(3:.25:10)), 1115)
N_grid = [10]

outPath = "$(spath)_syntheticRossman_$(s)_$(usePoisson)_$(4*numRuns)"
#old naming specification.  keep for a bit.
#outPath = "$(spath)_Ross_$(maximum(K_grid))_$(d)_$(N)_$(s)_$(usePoisson)_$(4*numRuns)"

#First read in the data and parse it appropriately
#Do this on a single processor bc it should be fast.
ps_full = readdlm("../RossmanKaggleData/CleanedData/ps_full$(d).csv", ',')
supp_full = readdlm("../RossmanKaggleData/CleanedData/support$(d).csv", ',')

start_time = time_ns()
file_a = @spawn convInKtest(numRuns, K_grid, supp_full, ps_full, "$(outPath)_a_", N_grid, s, seed=8675309, usePoisson=usePoisson)
file_b = @spawn convInKtest(numRuns, K_grid, supp_full, ps_full, "$(outPath)_b_", N_grid, s, seed=5167462266, usePoisson=usePoisson)
file_c = @spawn convInKtest(numRuns, K_grid, supp_full, ps_full, "$(outPath)_c_", N_grid, s, seed=5164174290, usePoisson=usePoisson)
file_d = @spawn convInKtest(numRuns, K_grid, supp_full, ps_full, "$(outPath)_d_", N_grid, s, seed=112456059, usePoisson=usePoisson)

# ######
file_a = fetch(file_a)
file_b = fetch(file_b)
file_c = fetch(file_c)
file_d = fetch(file_d)

time_stamp = (time_ns() - start_time) * 1e-9

##read everyone in, throw away a line
data, header = readdlm(file_a, ',', header=true)

data_t = readdlm(file_b, ',', skipstart=1)
data_t[:, 1] .+= numRuns
data = vcat(data, data_t)

data_t = readdlm(file_c, ',', skipstart=1)
data_t[:, 1] .+= 2numRuns
data = vcat(data, data_t)

data_t = readdlm(file_d, ',', skipstart=1)
data_t[:, 1] .+= 3numRuns
data = vcat(data, data_t)

f = open("$(outPath).csv", "w")
writedlm(f, header, ',')
writedlm(f, data, ',')
close(f)

println("Num Paths: \t $(4*numRuns) \t Time:", time_stamp )


struct Rxn{T<:Real}
    id::String
    S::Vector{T}
    mets::Vector
    c::T
    lb::T
    ub::T

    function Rxn{T}(id::String;
            mets = [],
            S = T[],
            c = zero(T),
            lb = zero(T),
            ub = zero(T)
        ) where {T<:Real}

        length(S) != length(mets) && error("'S' and 'mets' must have the same length")
        new{T}(id, S, mets, c, lb, ub)
    end
end

Rxn(id::String; kwargs...) = Rxn{Float64}(id; kwargs...)


###############################################################################
# BUILD TMLE FROM .TOML

function buildmodels(config)
    models = Dict()
    for (modelname, hyperparams) in config
        if !(modelname in ("resampling", "measures"))
            modeltype = eval(Symbol(modelname))
            paramnames = Tuple(Symbol(x[1]) for x in hyperparams)
            counter = 1
            for paramvals in Base.Iterators.product(values(hyperparams)...)
                model = modeltype(;NamedTuple{paramnames}(paramvals)...)
                models[Symbol(modelname*"_$counter")] = model
                counter += 1
            end
        end
    end
    return models
end

function stack_from_config(config::Dict, metalearner; adaptive_cv=true)
    # Define the resampling strategy
    resampling = config["resampling"]
    resampling = eval(Symbol(resampling["type"]))(nfolds=resampling["nfolds"])
    if adaptive_cv
        resampling = AdaptiveCV(resampling)
    end

    # Define the internal cross validation measures to report
    measures = (haskey(config, "measures") && size(config["measures"], 1) > 0) ?
                    [getfield(MLJBase, Symbol(fn)) for fn in config["measures"]] :
                    nothing

    # Define the models library
    models = buildmodels(config)

    # Define the Stack
    Stack(;metalearner=metalearner, resampling=resampling, measures=measures, models...)
end

function estimator_from_toml(config::Dict, queries, target_type; adaptive_cv=true)
    # Parse estimator for the propensity score
    metalearner = LogisticClassifier(fit_intercept=false)
    if length(first(queries).case) > 1
        G = FullCategoricalJoint(stack_from_config(config["G"], metalearner, adaptive_cv=adaptive_cv))
    else
        G = stack_from_config(config["G"], metalearner, adaptive_cv=adaptive_cv)
    end

    # Parse estimator for the outcome regression
    if target_type == Real
        metalearner =  LinearRegressor(fit_intercept=false)
        Q̅ = stack_from_config(config["Qcont"], metalearner, adaptive_cv=adaptive_cv)
    elseif target_type == Bool
        metalearner = LogisticClassifier(fit_intercept=false)
        Q̅ = stack_from_config(config["Qcat"], metalearner, adaptive_cv=adaptive_cv)
    else
        throw(ArgumentError("The type of the outcomes: $target_type, should be either a Float or a Bool"))
    end

    threshold = haskey(config, "threshold") ? config["threshold"] : 0.005
    return TMLEstimator(Q̅, G, queries...; threshold=threshold)
end