text:
## Glue code for Optim.jl and csminwel.jl
immutable OptimSolver <: MathProgBase.SolverInterface.AbstractMathProgSolver
    optimizer
    options
end

function OptimSolver{T<:Optimizer}(::Type{T}; kwargs...)
  OptimSolver(T(), kwargs)
end

MathProgBase.NonlinearModel(s::OptimSolver) = OptimMathProgModel(s; s.options...)

mutable struct OptimProblem
  out
  eval_f::Function
  eval_grad_f::Function
  sense::Symbol
end

mutable struct OptimMathProgModel <: MathProgBase.SolverInterface.AbstractNonlinearModel
    options
    inner::OptimProblem
    numVar::Int
    varLB::Vector{Float64}
    varUB::Vector{Float64}
    initial_x
    sense::Int32
    d::MathProgBase.SolverInterface.AbstractNLPEvaluator
    s::Optimizer
    solve_time::Float64
    function OptimMathProgModel(s; options...)
           model = new()
           model.options = options
           model.s = s.optimizer
           model
    end
end

MathProgBase.SolverInterface.setwarmstart!(m::OptimMathProgModel, x) = copy!(m.initial_x, x)

MathProgBase.SolverInterface.loadproblem!(m::OptimMathProgModel, numVar::Int, x_l, x_u,  sense::Symbol,
                      d::MathProgBase.SolverInterface.AbstractNLPEvaluator) =
                      MathProgBase.SolverInterface.loadproblem!(m, numVar, 0, x_l, x_u, [], [], sense, d)

MathProgBase.SolverInterface.loadproblem!(m::OptimMathProgModel, numVar::Int,  sense::Symbol,
                      d::MathProgBase.SolverInterface.AbstractNLPEvaluator) =
                      MathProgBase.SolverInterface.loadproblem!(m, numVar, 0, [-Inf for j in 1:numVar], [+Inf for j in numVar], [], [], sense, d)

function MathProgBase.SolverInterface.loadproblem!(m::OptimMathProgModel,
                                                   numVar::Integer,
                                                   numConstr::Integer,
                                                   x_l, x_u, g_lb, g_ub,
                                                   sense::Symbol,
                                                   d::MathProgBase.SolverInterface.AbstractNLPEvaluator)

    @assert numConstr == 0  "OptimProblem can only solve unconstrained problems"
    @assert length(x_l) == length(x_u) "Lower and upper bounds have inconsistent dims"
    @assert all(x_l .<= x_u) "Lower and upper bounds have inconsistent dims"
    m.varLB, m.varUB = x_l, x_u
    m.numVar = length(x_l)
    @assert m.numVar == length(m.varUB)
    m.initial_x = zeros(m.numVar)
    @assert sense == :Min || sense == :Max
    s = (sense == :Min) ? 1 : -1
    m.sense = s
    # Objective callback
    eval_f_cb(x) = s*MathProgBase.eval_f(d, x)
    # Objective gradient callback
    gr_cb(grad_f, x) = s*MathProgBase.eval_grad_f(d, grad_f, x)
    ad = any(map(x -> first(x), m.options) .== :autodiff)
    ng = :Grad ∈ MathProgBase.features_available(d)
    eval_grad_f_cb = getgradient(eval_f_cb, gr_cb, m.s, Val{ng}, Val{ad})
    m.inner = OptimProblem(nothing, eval_f_cb, eval_grad_f_cb, sense)
    m.d = d
    m.solve_time = 0.0
    m
end

## BFGS
getgradient(fcn, grad_f, s::Optimizer, ::Type{Val{true}}, ::Type{Val{false}}) = grad_f

function getgradient(fcn, grad_f, s::Optimizer, ::Type{Val{false}}, ::Type{Val{true}})
    gradient(gr, x) = ForwardDiff.gradient!(gr, fcn, x)
    gradient
end

function getgradient(fcn, grad_f, s::Optimizer, ::Type{Val{false}}, ::Type{Val{false}})
    gradient(gr, x) = gr[:] = Calculus.gradient(fcn, x)
    gradient
end

function unbounded(m::OptimMathProgModel)
    (all(MathProgBase.SolverInterface.getvarLB(m) == -Inf) && all(MathProgBase.SolverInterface.getvarUB(m) == +Inf)) ? true : false
end

function getoptimizer(s::Optim.Optimizer)
    notworking = (Newton, NewtonTrustRegion)
    for j in notworking
        isa(s, j) && error("Optimizer $j does not work with Fminbox")
    end
    m = (AcceleratedGradientDescent,
         BFGS,
         ConjugateGradient,
         GradientDescent,
         LBFGS,
         MomentumGradientDescent,
         NelderMead)
    out = BFGS
    for j in m
        isa(s, j) && (out = j; break)
    end
    out
end

function MathProgBase.SolverInterface.optimize!(m::OptimMathProgModel)
    if unbounded(m)
        out = Optim.optimize(m.inner.eval_f,
                             m.inner.eval_grad_f,
                             m.initial_x, m.s,
                             Optim.Options(;m.options...))
    else
        method = getoptimizer(m.s)
        out = optimize(OnceDifferentiable(m.inner.eval_f, m.inner.eval_grad_f),
                 m.initial_x, MathProgBase.SolverInterface.getvarLB(m),
                 MathProgBase.SolverInterface.getvarUB(m), Fminbox(), optimizer = getoptimizer(m.s))
    end
    m.inner.out = out
end

MathProgBase.SolverInterface.getsense(m::OptimMathProgModel) = m.sense
MathProgBase.SolverInterface.numvar(m::OptimMathProgModel) = m.numVar
MathProgBase.SolverInterface.numconstr(m::OptimMathProgModel) = 0
MathProgBase.SolverInterface.status(m::OptimMathProgModel) = Optim.converged(m.inner.out)
MathProgBase.SolverInterface.getobjval(m::OptimMathProgModel) = minimum(m.inner.out)
MathProgBase.SolverInterface.getsolution(m::OptimMathProgModel) = minimizer(m.inner.out)
MathProgBase.SolverInterface.getvarLB(m::OptimMathProgModel) = m.varLB
MathProgBase.SolverInterface.getvarUB(m::OptimMathProgModel) = m.varUB
MathProgBase.SolverInterface.setvarLB!(m::OptimMathProgModel, x) = copy!(m.varLB, x)
MathProgBase.SolverInterface.setvarUB!(m::OptimMathProgModel, x) = copy!(m.varUB, x)


text:
_range(v0, v1, N::Int) = range(v0, v1; length = N)
_range(v0, v1, s::Float64) = range(v0, v1; step = s)
_range(b::Tuple, δ) = _range(b..., δ)

## ------------------------------------------------------
struct BoxGrid
    dims_ids::Vector{Symbol}
    dims_bounds::Vector{Tuple}
    dims_δs::Vector{Union{Int, Float64}}
    dims_ranges::Vector{StepRangeLen}
    dims_steps::Vector{Float64}
    iter::Base.Iterators.ProductIterator
    D::Int

    function BoxGrid(dat::Vector)
        isempty(dat) && error("No dimension info")
        dims_ids = Symbol[]
        dims_bounds = Tuple[]
        dims_δs = Union{Int, Float64}[]
        for (id, lb, ub, δ) in dat
            (lb > ub) &&
                error("Bad bounds, lb >= ub. At $id got ", (lb, ub))
            push!(dims_ids, id)
            push!(dims_bounds, (lb, ub))
            push!(dims_δs, δ)
        end
        dims_ranges = _range.(dims_bounds, dims_δs)
        dims_steps = step.(dims_ranges)
        iter = Iterators.product(dims_ranges...)
        D = length(dims_ids)
        new(
            dims_ids, dims_bounds, dims_δs,
            dims_ranges, dims_steps, iter, D
        )
    end
end
function Base.show(io::IO, b::BoxGrid)
    println(io, "BoxGrid(;")
    tab = "    "
    println(io, tab, "# id = (lb, ub, δ)")
    params = String[]
    for (id, (lb, ub), δ) in zip(b.dims_ids, b.dims_bounds, b.dims_δs)
        push!(params, string(tab, id, " = ", (lb, ub, δ)))
    end
    println(io, join(params, ",\n"))
    println(io, ")")
end
Base.show(b::BoxGrid) = show(stdout, b)

## ------------------------------------------------------
# Iterator Interface
Base.IteratorSize(::BoxGrid) = Base.IteratorSize(Base.Iterators.ProductIterator)
Base.IteratorEltype(::BoxGrid) = Base.IteratorEltype(Base.Iterators.ProductIterator)
Base.eltype(b::BoxGrid) = eltype(b.iter)
Base.length(b::BoxGrid) = length(b.iter)
Base.iterate(b::BoxGrid) = iterate(b.iter)
Base.iterate(b::BoxGrid, state) = iterate(b.iter, state)
Base.size(b::BoxGrid) = tuple(length.(b.dims_ranges)...)
Base.size(b::BoxGrid, dim) = size(b)[dim]

## ------------------------------------------------------
dim_ids(b::BoxGrid) = b.dims_ids
dim_ids(b::BoxGrid, dim) = b.dims_ids[dim_ci(b, dim)]
dim_ci(b::BoxGrid, dimid::Symbol) = findfirst(isequal(dimid), dim_ids(b))
dim_ci(b::BoxGrid, dim::Int) = findfirst(isequal(dim), eachindex(dim_ids(b)))
dim_ci(b::BoxGrid, dims) = dim_ci.([b], dims)

dim_range(b::BoxGrid, dim) = b.dims_ranges[dim_ci(b, dim)]
dim_step(b::BoxGrid, dim) = b.dims_steps[dim_ci(b, dim)]

ext: jl


text:
# __precompile__()

module LambertW

import Base: convert
export lambertw, lambertwbp

const omega_const_bf_ = Ref{BigFloat}()

function __init__()
    # allocate storage for this BigFloat constant each time this module is loaded
    omega_const_bf_[] =
        parse(BigFloat, "0.5671432904097838729999686622103555497538157871865125081351310792230457930866845666932194")
end

### Lambert W function

"""
    lambertw(z::Complex{T}, k::V=0) where {T<:Real, V<:Integer}
    lambertw(z::T, k::V=0) where {T<:Real, V<:Integer}

Compute the `k`th branch of the Lambert W function of `z`. If `z` is real, `k` must be
either `0` or `-1`. For `Real` `z`, the domain of the branch `k = -1` is `[-1/e, 0]` and the
domain of the branch `k = 0` is `[-1/e, Inf]`. For `Complex` `z`, and all `k`, the domain is
the complex plane.

```jldoctest
julia> lambertw(-1/e, -1)
-1.0

julia> lambertw(-1/e, 0)
-1.0

julia> lambertw(0, 0)
0.0

julia> lambertw(0, -1)
-Inf

julia> lambertw(Complex(-10.0, 3.0), 4)
-0.9274337508660128 + 26.37693445371142im
```
"""
lambertw(z, k::Integer=0, maxits::Integer=1000) = _lambertw(z, k, maxits)

# lambertw(e + 0im, k) is ok for all k
# Maybe this should return a float. But, this should cause no type instability in any case
function _lambertw(::typeof(MathConstants.e), k, maxits)
    k == 0 && return 1
    throw(DomainError(k))
end
_lambertw(x::Irrational, k, maxits) = _lambertw(float(x), k, maxits)
function _lambertw(x::Union{Integer, Rational}, k, maxits)
    if k == 0
        x == 0 && return float(zero(x))
        x == 1 && return convert(typeof(float(x)), LambertW.omega) # must be a more efficient way
    end
    return _lambertw(float(x), k, maxits)
end

### Real z

function _lambertw(x::Real, k, maxits)
    k == 0 && return lambertw_branch_zero(x, maxits)
    k == -1 && return lambertw_branch_one(x, maxits)
    throw(DomainError(k, "lambertw: real x must have branch k == 0 or k == -1"))
end

# Real x, k = 0
# This appears to be inferrable with T=Float64 and T=BigFloat, including if x=Inf.
# There is a magic number here. It could be noted, or possibly removed.
# In particular, the fancy initial condition selection does not seem to help speed.
function lambertw_branch_zero(x::T, maxits)::T where T<:AbstractFloat
    isnan(x) && return(NaN)
    x == Inf && return Inf # appears to return convert(BigFloat, Inf) for x == BigFloat(Inf)
    one_t = one(T)
    oneoe = -one_t / convert(T, MathConstants.e)  # The branch point
    x == oneoe && return -one_t
    oneoe <= x || throw(DomainError(x))
    itwo_t = 1 / convert(T, 2)
    if x > one_t
        lx = log(x)
        llx = log(lx)
        x0 = lx - llx - log(one_t - llx / lx) * itwo_t
    else
        x0 = (567//1000) * x
    end
    return lambertw_root_finding(x, x0, maxits)
end

# Real x, k = -1
function lambertw_branch_one(x::T, maxits) where T<:Real
    oneoe = -one(T) / convert(T, MathConstants.e)
    x == oneoe && return -one(T) # W approaches -1 as x -> -1/e from above
    oneoe <= x || throw(DomainError(x))  # branch domain exludes x < -1/e
    x == zero(T) && return -convert(T, Inf) # W decreases w/o bound as x -> 0 from below
    x < zero(T) || throw(DomainError(x))
    return lambertw_root_finding(x, log(-x), maxits)
end

### Complex z

_lambertw(z::Complex{<:Integer}, k, maxits) = _lambertw(float(z), k, maxits)
# choose initial value inside correct branch for root finding
function _lambertw(z::Complex{T}, k, maxits) where T<:Real
    one_t = one(T)
    local w::Complex{T}
    pointseven = 7//10
    if abs(z) <= one_t/convert(T, MathConstants.e)
        if z == 0
            k == 0 && return z
            return complex(-convert(T, Inf), zero(T))
        end
        if k == 0
            w = z
        elseif k == -1 && imag(z) == 0 && real(z) < 0
            w = complex(log(-real(z)), 1//10^7) # need offset for z ≈ -1/e.
        else
            w = log(z)
            k != 0 ? w += complex(0, k * 2 * pi) : nothing
        end
    elseif k == 0 && imag(z) <= pointseven && abs(z) <= pointseven
        w = abs(z+ 1//2) < 1//10 ? imag(z) > 0 ? complex(pointseven, pointseven) : complex(pointseven, -pointseven) : z
    else
        if real(z) == convert(T, Inf)
            k == 0 && return z
            return z + complex(0, 2*k*pi)
        end
        real(z) == -convert(T, Inf) && return -z + complex(0, (2*k+1)*pi)
        w = log(z)
        k != 0 ? w += complex(0, 2*k*pi) : nothing
    end
    return lambertw_root_finding(z, w, maxits)
end

### root finding, iterative solution

# Use Halley's root-finding method to find
# x = lambertw(z) with initial point x0.
function lambertw_root_finding(z::T, x0::T, maxits) where T <: Number
    two_t = convert(T, 2)
    x = x0
    lastx = x
    lastdiff = zero(T)
    converged::Bool = false
    for i in 1:maxits
        ex = exp(x)
        xexz = x * ex - z
        x1 = x + 1
        x -= xexz / (ex * x1 - (x + two_t) * xexz / (two_t * x1 ))
        xdiff = abs(lastx - x)
        if xdiff <= 3 * eps(abs(lastx)) || lastdiff == xdiff  # second condition catches two-value cycle
            converged = true
            break
        end
        lastx = x
        lastdiff = xdiff
    end
    converged || @warn("lambertw with z=", z, " did not converge in ", maxits, " iterations.")
    return x
end

### Inverse of Lambert W function

"""
    finv(::typeof(lambertw)) -> Function

return the functional inverse of the Lambert W function.
"""
finv(::typeof(lambertw)) = z -> z * exp(z)

### omega constant

const omega_const_ = 0.567143290409783872999968662210355
# The BigFloat `omega_const_bf_` is set via a literal in the function __init__ to prevent a segfault

# compute omega constant via root finding
# We could compute higher precision. This converges very quickly.
function omega_const(::Type{BigFloat})
    precision(BigFloat) <= 256 && return omega_const_bf_[]
    myeps = eps(BigFloat)
    oc = omega_const_bf_[]
    for i in 1:100
        nextoc = (1 + oc) / (1 + exp(oc))
        abs(oc - nextoc) <= myeps && break
        oc = nextoc
    end
    return oc
end

"""
    omega
    ω

The constant defined by `ω exp(ω) = 1`.

```jldoctest
julia> ω
ω = 0.5671432904097...

julia> omega
ω = 0.5671432904097...

julia> ω * exp(ω)
1.0

julia> big(omega)
5.67143290409783872999968662210355549753815787186512508135131079223045793086683e-01
```
"""
const ω = Irrational{:ω}()
@doc (@doc ω) omega = ω

Base.Float64(::Irrational{:ω}) = omega_const_  # FIXME: This is very slow. Why ?
Base.Float32(::Irrational{:ω}) = Float32(omega_const_)
Base.Float16(::Irrational{:ω}) = Float16(omega_const_)
Base.BigFloat(o::Irrational{:ω}) = omega_const(BigFloat)

### Expansion about branch point x = -1/e

# Refer to the paper "On the Lambert W function". In (4.22)
# coefficients μ₀ through μ₃ are given explicitly. Recursion relations
# (4.23) and (4.24) for all μ are also given. This code implements the
# recursion relations.

# (4.23) and (4.24) give zero based coefficients.
cset(a, i, v) = a[i+1] = v
cget(a, i) = a[i+1]

# (4.24)
function compute_a_coeffs(k, m, a)
    sum0 = zero(eltype(m))
    for j in 2:(k - 1)
        sum0 += cget(m, j) * cget(m, k + 1 - j)
    end
    cset(a, k, sum0)
    return sum0
end

# (4.23)
function compute_m_coefficients(k, m, a)
    kt = convert(eltype(m), k)
    mk = (kt - 1) / (kt + 1) *(cget(m, k - 2) / 2 + cget(a, k - 2) / 4) -
        cget(a, k) / 2 - cget(m, k - 1) / (kt + 1)
    cset(m, k, mk)
    return mk
end

# We plug the known value μ₂ == -1//3 for (4.22) into (4.23) and
# solve for α₂. We get α₂ = 0.
# Compute array of coefficients μ in (4.22).
# m[1] is μ₀
function compute_branch_point_coeffs(T::DataType, n::Int)
    a =  Array{T}(undef, n)
    m =  Array{T}(undef, n)
    cset(a, 0, 2)  # α₀ literal in paper
    cset(a, 1, -1) # α₁ literal in paper
    cset(a, 2, 0)  # α₂ get this by solving (4.23) for alpha_2 with values printed in paper
    cset(m, 0, -1) # μ₀ literal in paper
    cset(m, 1, 1)  # μ₁ literal in paper
    cset(m, 2, -1//3) # μ₂ literal in paper, but only in (4.22)
    for i in 3:(n - 1)  # coeffs are zero indexed
        compute_a_coeffs(i, m, a)
        compute_m_coefficients(i, m, a)
    end
    return m
end

const BRANCH_POINT_COEFFS_FLOAT64 = compute_branch_point_coeffs(Float64, 500)

# Base.Math.@horner requires literal coefficients
# It cannot be used here because we have an array of computed coefficients
function horner(x, coeffs::AbstractArray, n)
    n += 1
    ex = coeffs[n]
    for i = (n - 1):-1:2
        ex = :(muladd(t, $ex, $(coeffs[i])))
    end
    ex = :( t * $ex)
    return Expr(:block, :(t = $x), ex)
end

# write functions that evaluate the branch point series
# with `num_terms` number of terms.
for (func_name, num_terms) in (
    (:wser3, 3), (:wser5, 5), (:wser7, 7), (:wser12, 12),
    (:wser19, 19), (:wser26, 26), (:wser32, 32),
    (:wser50, 50), (:wser100, 100), (:wser290, 290))
    iex = horner(:x, BRANCH_POINT_COEFFS_FLOAT64, num_terms)
    @eval function ($func_name)(x) $iex end
end

# Converges to Float64 precision
# We could get finer tuning by separating k=0, -1 branches.
# Why is wser5 omitted ?
# p is the argument to the series which is computed
# from x before calling `branch_point_series`.
function branch_point_series(p, x)
    x < 4e-11 && return wser3(p)
    x < 1e-5 && return wser7(p)
    x < 1e-3 && return wser12(p)
    x < 1e-2 && return wser19(p)
    x < 3e-2 && return wser26(p)
    x < 5e-2 && return wser32(p)
    x < 1e-1 && return wser50(p)
    x < 1.9e-1 && return wser100(p)
    x > 1 / MathConstants.e && throw(DomainError(x))  # radius of convergence
    return wser290(p)  # good for x approx .32
end

# These may need tuning.
function branch_point_series(p::Complex{T}, z) where T<:Real
    x = abs(z)
    x < 4e-11 && return wser3(p)
    x < 1e-5 && return wser7(p)
    x < 1e-3 && return wser12(p)
    x < 1e-2 && return wser19(p)
    x < 3e-2 && return wser26(p)
    x < 5e-2 && return wser32(p)
    x < 1e-1 && return wser50(p)
    x < 1.9e-1 && return wser100(p)
    x > 1 / MathConstants.e && throw(DomainError(x))  # radius of convergence
    return wser290(p)
end

function _lambertw0(x) # 1 + W(-1/e + x)  , k = 0
    ps = 2 * MathConstants.e * x
    series_arg = sqrt(ps)
    branch_point_series(series_arg, x)
end

function _lambertwm1(x) # 1 + W(-1/e + x)  , k = -1
    ps = 2 * MathConstants.e * x
    series_arg = -sqrt(ps)
    branch_point_series(series_arg, x)
end

"""
    lambertwbp(z, k=0)

Compute accurate value of `1 + W(-1/e + z)`, for `abs(z)` in `[0, 1/e]` for `k` either `0` or `-1`.
This function is faster and more accurate near the branch point `-1/e` between `k=0` and `k=1`.
The result is accurate to Float64 precision for abs(z) < 0.32.
If `k=-1` and `imag(z) < 0`, the value on the branch `k=1` is returned.

```jldoctest
julia> lambertw(-1/e + 1e-18, -1)
-1.0

julia> lambertwbp(1e-18, -1)
-2.331643983409312e-9

# Same result, but 1000 times slower
julia> convert(Float64, (lambertw(-BigFloat(1)/e + BigFloat(10)^(-18), -1) + 1))
-2.331643983409312e-9
```

!!! note
    `lambertwbp` uses a series expansion about the branch point `z=-1/e` to avoid loss of precision.
    The loss of precision in `lambertw` is analogous to the loss of precision
    in computing the `sqrt(1-x)` for `x` close to `1`.
"""
function lambertwbp(x::Number, k::Integer)
    k == 0 && return _lambertw0(x)
    k == -1 && return _lambertwm1(x)
    throw(ArgumentError("expansion about branch point only implemented for k = 0 and -1."))
end
lambertwbp(x::Number) = _lambertw0(x)

end #module


text:
using PyPlot, JUDI, LinearAlgebra, FFTW, JOLI, DSP
using IterativeSolvers

using Shampoo

n = (201, 201)   # (x,y,z) or (x,z)
d = (5f0, 5f0)
o = (0f0, 0f0)

extentx = (n[1]-1)*d[1]
extentz = (n[2]-1)*d[2]

v0 = 2f0*ones(Float32,n)
v = deepcopy(v0)
v[101,101] = 2.1f0  # dirac perturbation

# Slowness squared [s^2/km^2]
m = convert(Array{Float32,2},(1f0 ./ v).^2)	# true model
m0 = convert(Array{Float32,2},(1f0 ./ v0).^2)
dm = vec(m-m0)

model = Model(n, d, o, m)
model0 = Model(n, d, o, m0)

# Model structure

dtS = 2f0
timeS = 2000f0
nt = Int64(timeS/dtS)+1

fmin = 10f0
fmax = 60f0
cfreqs = (fmin, fmin+5f0, fmax-5f0, fmax) # corner frequencies
wavelet = reshape(cfreq_wavelet(500, nt, dtS/1f3, cfreqs; edge=hanning),:,1)
#wavelet = ricker_wavelet(timeS, dtS, 0.03f0)

################################### Source/receiver geometry ######################################

# Set up receiver geometry

nrec = 512

xrec, zrec = circleShape(extentx / 2,extentz/2, extentz/2-5*d[2],nrec)
yrec = 0f0 #2d so always 0

# receiver sampling and recording time
timeR = timeS   # receiver recording time [ms]
dtR   = dtS
# Set up receiver structure
nsrc = 32

xsrc, zsrc = circleShape(extentx / 2,extentz/2, extentz/2-5*d[2],nsrc)
ysrc = range(0f0,stop=0f0,length=nsrc) #2d so always 0

recGeometry = Geometry(xrec, yrec, zrec; dt=dtR, t=timeR, nsrc=nsrc)

srcGeometry = Geometry(convertToCell(xsrc), convertToCell(ysrc), convertToCell(zsrc); dt=dtS, t=timeS)

q = judiVector(srcGeometry, wavelet)

mute_matrix = zeros(Float32,n)

for i = 1:n[1]
	for j = 1:n[2]
		mute_matrix[i,j] = ((i-201)^2+(j-201)^2<190^2)
	end
end

M = MaskOp(n,mute_matrix)

# Set up info structure for linear operators
ntComp = get_computational_nt(recGeometry, model0)
info = Info(prod(n), nsrc, ntComp)

opt = Options(free_surface=false,isic=false)

Pr  = judiProjection(info, recGeometry)
F0  = judiModeling(info, model0; options=opt) # modeling operator background model
Ps = judiProjection(info, srcGeometry)
# Combined operators
F0 = Pr*F0*adjoint(Ps)

J = judiJacobian(F0, q)
P = DiffOp(nt,nsrc,nrec)

PyPlot.rc("font", family="serif"); PyPlot.rc("xtick", labelsize=8); PyPlot.rc("ytick", labelsize=8)
figure(); imshow(reshape(sqrt.(1f0./m), n[1], n[2])',extent=(0,extentx,extentz,0));
PyPlot.scatter(xrec,zrec,marker=".",color="red",label="receivers",s=1);
PyPlot.scatter(xsrc,zsrc,marker=".",color="white",label="sources",s=1);
legend(loc=3,fontsize=8);
xlabel("Lateral position [m]", fontsize=8);ylabel("Horizontal position [m]", fontsize=8);
fig = PyPlot.gcf();
title("Experimental setup")

######################################### Set up JUDI operators ###################################

d_lin = J*dm

dm1 = J'*d_lin

dm2 = J'*P*J*dm

figure();
subplot(2,1,1);
PyPlot.plot(wavelet);title("zoomed-in wavelet in time domain");
subplot(2,1,2);
PyPlot.plot(abs.(fftshift(fft(wavelet))));title("wavelet in frequency domain");

figure();
subplot(2,2,1)
imshow(dm1.data');title("J'*J*dm")
subplot(2,2,2)
imshow(dm2.data');title("J'*cumsum*J*dm")
subplot(2,2,3)
imshow(abs.(fftshift(fft(dm1.data')))/norm(abs.(fftshift(fft(dm1.data'))),Inf),cmap="jet",vmin=0,vmax=1)
subplot(2,2,4)
imshow(abs.(fftshift(fft(dm2.data')))/norm(abs.(fftshift(fft(dm2.data'))),Inf),cmap="jet",vmin=0,vmax=1)

maxit = 10

x1 = 0f0 .* dm
start1 = J'*P*d_lin
x1,his1 = cg!(x1,M'*J'*P*J*M,start1,maxiter=maxit,log=true,verbose=true)
x2 = 0f0 .* dm
start2 = J'*d_lin
x2,his2 = cg!(x2,M'*J'*J*M,start2,maxiter=maxit,log=true,verbose=true)
res1 = his1[:resnorm]
res2 = his2[:resnorm]

u1 = 20*log.(res1./norm(start1))
u2 = 20*log.(res2./norm(start2))

figure();
xlabel("iterations")
ylabel("ratio")
PyPlot.plot(res1/norm(start1),label="CG")
PyPlot.plot(res2/norm(start2),label="P-CG")
title("normalized least squares residual")
legend()

figure();
xlabel("iterations")
ylabel("dB")
PyPlot.plot(u1,label="CG")
PyPlot.plot(u2,label="P-CG")
title("normalized log-based least squares residual")
legend()

figure();
subplot(2,2,1)
imshow(reshape(M*x1,n)');title("CG")
subplot(2,2,2)
imshow(reshape(M*x2,n)');title("P-CG")
subplot(2,2,3)
imshow(abs.(fftshift(fft(reshape(M*x2,n)')))/norm(abs.(fftshift(fft(reshape(M*x2,n)'))),Inf),cmap="jet",vmin=0,vmax=1)
subplot(2,2,4)
imshow(abs.(fftshift(fft(reshape(M*x2,n)')))/norm(abs.(fftshift(fft(reshape(M*x2,n)'))),Inf),cmap="jet",vmin=0,vmax=1)


text:
"""
vectorize single input
"""

function vectorize!(v, dims, x::Union{Array, Real, <:Tuple{Vararg{Real}}})
    x = copy(x)
    if x isa Tuple
        n = length(x)
        for i = 1:n
            v[i] = convert(Float64, x[i])
        end
    elseif ndims(x) == 0
        v[1] = x
    else
        v = reshape(convert.(Float64, x), dims, 1)
    end
    return v
end

function vectorize!(v, dims, x::MvNormal)
    μ = mean(x)
    σ = diag(cov(x))
    dims_half = Int64(dims/2)
    v[1:dims_half] = μ
    v[dims_half+1:dims] = σ
    return v
end

function build_action_set(action_array::Array)
    n_actions = size(action_array)[end]
    action_set = Set()
    for i=1:n_actions
        array_slice = selectdim(action_array, ndims(action_array), i)
        union!(action_set, Set(array_slice))
    end
    return action_set
end

function acquire_max_action(opt, lowerbounds, upperbounds, restarts)
    maxf = -Inf
    maxx = lowerbounds
    seq = BayesianOptimization.ScaledLHSIterator(lowerbounds, upperbounds, restarts)
    for x0 in seq
        f, x, ret = NLopt.optimize(opt, x0)
        ret == NLopt.FORCED_STOP && throw(InterruptException())
        if f > maxf
            maxf = f
            maxx = x
        end
    end
    maxf, maxx
end

function action_function(acq, b, model)
    x -> begin
        x_test = [reshape(x, :, 1); b]
        μ, σ² = GaussianProcesses.predict_f(model, x_test)
        acq.(μ[1], max(σ²[1], 0.0))
    end
end

function next_action(o::BOActionSelector, p::Union{POMDP, MDP}, b, bnode)
    if haskey(o.b_a_dict, b)
        acts = o.b_a_dict[b]
    else
        if o.discrete_actions
            acts = build_action_set(POMDPs.actions(p, b))
        else
            acts = POMDPs.actions(p, b)
        end
        o.b_a_dict[b] = acts
    end
    b_node_id = bnode.index
    tree = bnode.tree
    if o.gp isa Dict
        if haskey(o.gp, b_node_id)
            gp = o.gp[b_node_id]
        else
            mean = isnothing(o.mean_params) ? o.mean_type() : o.mean_type(o.mean_params...)
            kernel = isnothing(o.kernel_params) ? o.kernel_type() : o.kernel_type(o.kernel_params...)
            x = zeros(Float64, o.action_dims, 0)
            y = zeros(Float64, 0)
            gp = GPE(x, y,
                    mean,
                    kernel,
                    o.log_noise)
            o.gp[b_node_id] = gp
        end
        belief_x = false
    else
        gp = o.gp
        belief_x = true
    end
    n_obs = size(gp.x, 2)
    if n_obs == 0
        if typeof(o.initial_action) <: Function
            action = o.initial_action(gp, acts)
        elseif typeof(o.initial_action) <: Policy
            action = POMDPs.action(o.initial_action, b)
        else
            action = o.initial_action in acts ? o.initial_action : throw(AssertionError("The provided initial action is not in the set returned by actions(pomdp, b)"))
        end
    else
        act_idxs = tree.b_children[b_node_id]
        obs_acts = tree.a_labels[act_idxs]
        obs_q = tree.q[act_idxs]
        if length(obs_q) >= 1
            o.acquisition_function.τ = maximum(obs_q)
        else
            o.acquisition_function.τ = -1e4 #This is the Lower Bound
        end
        if o.discrete_actions
            new_acts = collect(setdiff(acts, obs_acts))
            if isempty(new_acts)
                return rand(acts)
            end
            n_acts = length(new_acts)
            if belief_x
                vector_belief = zeros(Float64, o.belief_dims)
                vector_belief = vectorize!(vector_belief, o.belief_dims, b)
                vector_belief .*= o.belief_λ
                x = zeros(Float64, o.action_dims + o.belief_dims, n_acts)
                for i = 1:n_acts
                    x[1:o.action_dims,i] = vectorize!(x[1:o.action_dims,i], o.action_dims, new_acts[i])
                    x[o.action_dims+1:o.action_dims+o.belief_dims, i] = vector_belief
                end
            else
                x = zeros(Float64, o.action_dims, n_acts)
                for i = 1:n_acts
                    x[:,i] = vectorize!(x[:,i], o.action_dims, new_acts[i])
                end
            end
            mu, sig = GaussianProcesses.predict_f(gp, x)
            scores = o.acquisition_function.(mu, sig)
            act_idx = argmax(scores)
            action = new_acts[act_idx]
        else
            if o.gp isa Dict
                f = BayesianOptimization.wrap_gradient(BayesianOptimization.acquisitionfunction(o.acquisition_function, gp))
                NLopt.max_objective!(o.optim, f)
                _, action = acquire_max_action(o.optim, o.lower_bounds, o.upper_bounds, 100) # restarts)
            else
                vector_belief = zeros(Float64, o.belief_dims)
                vector_belief = vectorize!(vector_belief, o.belief_dims, b)
                vector_belief .*= o.belief_λ
                f = BayesianOptimization.wrap_gradient(action_function(o.acquisition_function, vector_belief, gp))
                NLopt.max_objective!(o.optim, f)
                _, action = acquire_max_action(o.optim, o.lower_bounds, o.upper_bounds, 100) # restarts)
            end
        end
    end
    return action
end

function update_gp!(gp::GPLA, tree, bnode, new_node, λ, buffer) #, dim::Val{M}, nobs::Val{N}) where {M, N}
    n_obs = length(tree.n_a)
    n_buffer = size(buffer[1], 2)
    tot_dim = beliefdim(gp) + actiondim(gp)
    if new_node
        gp.x = zeros(Float64, tot_dim, n_obs + n_buffer)
        gp.y = zeros(Float64, n_obs + n_buffer)
        for i = 1:n_obs
            gp.x[1:actiondim(gp), i] = vectorize!(gp.x[1:actiondim(gp), i], actiondim(gp), tree.a_labels[i])
            gp.x[actiondim(gp)+1:tot_dim, i] = vectorize!(gp.x[actiondim(gp)+1:tot_dim, i], beliefdim(gp), tree.p_labels[i])
            gp.x[actiondim(gp)+1:tot_dim, i] .*= λ
            gp.y[i] = tree.q[i] - tree.v[tree.p_nodes[i]]
        end
        gp.x[:,n_obs+1:n_buffer+n_obs] = buffer[1]
        gp.y[n_obs+1:n_buffer+n_obs] = buffer[2]
        gp.kdtree = KDTree(NearestNeighbors.copy_svec(Float64, gp.x, Val(tot_dim)))
    else
        for i = 1:n_obs
            gp.y[i] = tree.q[i] - tree.v[tree.p_nodes[i]]
        end
    end
    return gp
end

function update_gp!(gp::Dict, tree, bnode, new_node, λ)
    gpn = gp[bnode]
    act_idxs = tree.b_children[bnode]
    obs_acts = tree.a_labels[act_idxs]
    q_vals = tree.q[act_idxs]
    n_obs = length(obs_acts)
    x = zeros(Float64, gpn.dim, n_obs)
    for i = 1:n_obs
        x[:,i] = vectorize!(x[:,i], gpn.dim, obs_acts[i])
    end
    y = q_vals
    if new_node
        GaussianProcesses.fit!(gpn, x, y)
    else
        gpn.x = x
        gpn.y = y
        GaussianProcesses.update_mll!(gpn)
    end
    gp[bnode] = gpn
end

function update_gp_buffer!(action_selector::BOActionSelector)
    gp = action_selector.gp
    n = action_selector.buffer_size
    n_obs = size(gp.x, 2)
    if n_obs > n
        idxs = sample([1:n_obs;], n, replace=false)
        x_buffer = gp.x[:,idxs]
        y_buffer = gp.y[idxs]
    else
        x_buffer = gp.x
        y_buffer = gp.y
    end
    action_selector.obs_buffer = (x_buffer, y_buffer)
end

function reset_gp!(gp::GPLA)
    tot_dim = beliefdim(gp) + actiondim(gp)
    gp.x = zeros(Float64, tot_dim, 0)
    gp.y = zeros(Float64, 0)
    sx = NearestNeighbors.copy_svec(Float64, zeros(Float64, tot_dim, 0), Val(tot_dim))
    gp.kdtree = KDTree(sx, Euclidean())
end

function optimize_gp!(gp::GPLA, p)
    mean_results = nothing
    kernel_results = GaussianProcesses.optimize!(gp, method=Optim.ParticleSwarm(), domean = false, time_limit = 30.0,  allow_f_increases=true) #, g_tol=1e-7)
    return mean_results, kernel_results
end

function optimize_gp!(gp::Dict, p)
    x = zeros(Float64, size(gp[1].x, 1), 0)
    y = zeros(Float64, 0)
    for (key, value) in gp
        x = [x value.x]
        y = [y; value.y]
    end
    gpla = GPLA(x, y, 5, gp[1].dim, 0, gp[1].mean, gp[1].kernel, gp[1].logNoise.value)
    results = GaussianProcesses.optimize!(gpla, method=Optim.LBFGS(), time_limit = 60.0,  allow_f_increases=true)#,  g_tol=1e-4)

    mean_params = GaussianProcesses.get_params(gpla.mean)
    kernel_params = GaussianProcesses.get_params(gpla.kernel)
    noise_params = GaussianProcesses.get_params(gpla.logNoise)

    p.solver.action_selector.mean_params = mean_params
    p.solver.action_selector.kernel_params = kernel_params
    p.solver.action_selector.log_noise = noise_params[1]
    return results
end


text:
module TEMPEST1DInversion
import ..AbstractOperator.get_misfit, ..main
import ..AbstractOperator.Sounding
import ..AbstractOperator.makeoperator
import ..AbstractOperator.make_tdgp_opt

using ..AbstractOperator, ..AEM_VMD_HMD, Statistics
using PyPlot, LinearAlgebra, ..CommonToAll, Random, DelimitedFiles, Distributed, Dates

import ..Model, ..Options, ..OptionsStat, ..OptionsNonstat
import ..ModelNuisance, ..OptionsNuisance, ..findidxnotzero

const μ₀ = 4*pi*1e-7

mutable struct Bfield<:Operator1D
    F          :: AEM_VMD_HMD.HField
    dataHx     :: Array{Float64, 1}
    dataHz     :: Array{Float64, 1}
    useML      :: Bool
    σx         :: Array{Float64, 1}
    σz         :: Array{Float64, 1}
    z          :: Array{Float64, 1}
    nfixed     :: Int
    ρ          :: Array{Float64, 1}
    selectx    :: Array{Bool, 1}
    selectz    :: Array{Bool, 1}
    ndatax     :: Int
    ndataz     :: Int
    rx_roll    :: Float64
    rx_pitch   :: Float64
    rx_yaw     :: Float64
    tx_roll    :: Float64
    tx_pitch   :: Float64
    tx_yaw     :: Float64
    Rot_rx     :: Array{Float64,2}
	x_rx       :: Float64
	y_rx       :: Float64
	mhat       :: Array{Float64, 1}
	Hx         :: Array{Float64, 1}
	Hy         :: Array{Float64, 1}
	Hz         :: Array{Float64, 1}
	addprimary :: Bool
	peakcurrent:: Float64
    vectorsum  :: Bool
end

# If needed to make z axis flip to align with GA-AEM
const Roll180 = [1. 0. 0.
				 0 -1  0.
				 0  0 -1]

const fTinv = 1e15

function Bfield(;
				dataHx = zeros(0),
				dataHz = zeros(0),
				σx = zeros(0),
				σz = zeros(0),
				selectx = zeros(Bool, 0),
				selectz = zeros(Bool, 0),
				useML = false,
				nfixed = 1,
				times = [],
				ramp  = [],
				z = zeros(0),
				ρ = zeros(0),
				ndatax = 0,
				ndataz = 0,
				ntimesperdecade = 10,
				nfreqsperdecade = 5,
				nkᵣeval = 60,
				zTx = -120,
				zRx = -80,
				x_rx = -115.0,
				y_rx = 0.,
				rx_roll = 0.,
			    rx_pitch = 0.,
			    rx_yaw = 0.,
			    tx_roll = 0.,
			    tx_pitch = 0.,
			    tx_yaw = 0.,
				order_tx = "ypr",
				order_rx = "ypr",
				strictgeometry = true,
				addprimary = false,
				peakcurrent = 0.5,
                vectorsum = false)

	@assert !isempty(times)
	@assert(!isempty(ramp))
	@assert zTx<0
	if strictgeometry
		@assert zRx>zTx # receiver below transmitter
		@assert x_rx<0 # receiver behind transmitter
	end
	Rot_rx = makerotationmatrix(order=order_rx,yaw=rx_yaw, pitch=rx_pitch, roll=rx_roll,doinv=true)
	Rot_tx = makerotationmatrix(order=order_tx,yaw=tx_yaw, pitch=tx_pitch, roll=tx_roll)
	F = AEM_VMD_HMD.HFieldDHT(;
	                      ntimesperdecade = ntimesperdecade,
	                      nfreqsperdecade = nfreqsperdecade,
						  freqlow=1e-5,
						  nkᵣeval = nkᵣeval,
	                      times  = times,
	                      ramp   = ramp,
	                      zTx    = zTx,
	                      rRx    = sqrt(x_rx^2 + y_rx^2),
	                      zRx    = zRx,
						  modelprimary = false,
						  getradialH = true,
						  getazimH = true,
						  provideddt = false)
	mhat = Rot_tx*[0,0,1] # dirn cosines in inertial frame for VMDz
	Hx, Hy, Hz = map(x->zeros(size(times)), 1:3)
	Bfield(F, dataHx, dataHz, useML,σx, σz, z, nfixed, copy(ρ), selectx, selectz,
			ndatax, ndataz, rx_roll, rx_pitch, rx_yaw, tx_roll, tx_pitch, tx_yaw,
			Rot_rx, x_rx, y_rx, mhat, Hx, Hy, Hz, addprimary, peakcurrent, vectorsum)
end

#TODO for nuisance moves in an MCMC chain
function update_geometry(tempest::Bfield, geovec::Array{Float64,1},
	order_rx = "ypr", order_tx = "ypr")
	length(geovec) == 10 ||
		throw(DimensionMismatch("TEMPEST geometry set with vector of wrong length."))
	zTx = geovec[1]
	zRx = geovec[2]
	x_rx = geovec[3]
	y_rx = geovec[4]

	rx_roll = geovec[5]
	rx_pitch = geovec[6]
	rx_yaw = geovec[7]

	tx_roll = geovec[8]
	tx_pitch = geovec[9]
	tx_yaw = geovec[10]

	#make new rotation matrices
	Rot_rx = makerotationmatrix(order = order_rx,
		yaw = rx_yaw, pitch = rx_pitch, roll = rx_roll, doinv = true)
	Rot_tx = makerotationmatrix(order = order_tx,
		yaw = tx_yaw, pitch = tx_pitch, roll = tx_roll)

	#do update on internal VMD model
	AEM_VMD_HMD.update_ZR!(tempest.F, zTx, zRx, nothing, sqrt(x_rx^2 + y_rx^2))
	tempest.x_rx = x_rx
	tempest.y_rx = y_rx
	tempest.Rot_rx = Rot_rx
	tempest.mhat = Rot_tx*[0,0,1]

	tempest.rx_roll = rx_roll
	tempest.rx_pitch = rx_pitch
	tempest.rx_yaw = rx_yaw

	tempest.tx_roll = tx_roll
	tempest.tx_pitch = tx_pitch
	tempest.tx_yaw = tx_yaw

	nothing
end

function getfieldTD!(tempest::Bfield, z::Array{Float64, 1}, ρ::Array{Float64, 1})
	AEM_VMD_HMD.getfieldTD!(tempest.F,  z, ρ)
	reducegreenstensor!(tempest)
	nothing
end

function returnprimary!(tempestin;fillrho=1e12)
# only useful for synthetics I guess
    tempest = deepcopy(tempestin)
    tempest.ρ .= 10 # no contrasts at all
    tempest.addprimary = true
    getfieldTD!(tempest, tempest.z, tempest.ρ)
	tempest.Hx, tempest.Hy, tempest.Hz
end

#match API for SkyTEM inversion getfield
function getfield!(m::Model, tempest::Bfield)
	getfield!(m.fstar, tempest)
	nothing
end
function getfield!(m::Array{Float64}, tempest::Bfield)
    copyto!(tempest.ρ, tempest.nfixed+1:length(tempest.ρ), 10 .^m, 1:length(m))
	getfieldTD!(tempest, tempest.z, tempest.ρ)
    nothing
end

# set the field given a conductivity model (GP parametrisation)
# and nuisance model (vector)
function getfield!(m::Model, mn::ModelNuisance, tempest::Bfield)
	update_geometry(tempest, mn.nuisance)
	getfield!(m, tempest)
end

function getfield!(m::Array{Float64}, mn::Array{Float64}, tempest::Bfield)
    copyto!(tempest.ρ, tempest.nfixed+1:length(tempest.ρ), 10 .^m, 1:length(m))
	update_geometry(tempest, vec(mn))
	getfieldTD!(tempest, tempest.z, tempest.ρ)
    nothing
end

function reducegreenstensor!(tempest)
	x, y   = tempest.x_rx, tempest.y_rx
	r      = tempest.F.rRx
	J1h    = tempest.F.dBazdt
	J1v    = tempest.F.dBrdt
	J0v    = tempest.F.dBzdt
	mhat   = tempest.mhat
	Rot_rx = tempest.Rot_rx
	Hx, Hy, Hz = tempest.Hx, tempest.Hy, tempest.Hz
	currentfac = tempest.peakcurrent

	y2mx2 = y^2-x^2
	r3    = r^3
	x2    = x^2
	r2    = r^2
	xy    = x*y
	y2    = y^2
	z     = tempest.F.zRx - tempest.F.zTx
	R2    = r2 + z^2
	fpiR5 = 4pi*sqrt(r2 + z^2)^5
	xz    = x*z
	yz    = y*z

	HMDx = [y2mx2/r3*J1h + x2/r2*J0v,
		    -2xy/r3*J1h  + xy/r2*J0v,
		    -x/r*J1v                       ]

	HMDy = [HMDx[2],
	 		-y2mx2/r3*J1h + y2/r2*J0v,
	 		-y/r*J1v		    		   ]

	VMDz = [x/r*J1v,
			y/r*J1v,
			J0v                            ]

	if tempest.addprimary
		HMDxp = [3x2 - R2, 3xy     , 3xz      ]/fpiR5
		HMDyp = [3xy     , 3y2 - R2, 3yz      ]/fpiR5
		VMDzp = [3xz     , 3yz     , 3z^2 - R2]/fpiR5
		for idim in 1:3
			HMDx[idim] .+= currentfac*HMDxp[idim]
			HMDy[idim] .+= currentfac*HMDyp[idim]
			VMDz[idim] .+= currentfac*VMDzp[idim]
		end
	end

	Hx[:], Hy[:], Hz[:] = Rot_rx*Roll180*[HMDx HMDy VMDz]*Roll180*mhat
	nothing
end

function makerotationmatrix(;yaw=0.0,roll=0.0,pitch=0.0, order="lala", doinv = false)
    ordervector = ["ypr","yrp","rpy","ryp","pyr","pry"]
    @assert any(ordervector .== order) """not a valid order such as "pry" """
	# y --> z or [0,1,0] --> [0,0,1]
    Rr = [ 1.           0            0
           0            cosd(roll)  -sind(roll)
           0            sind(roll)   cosd(roll)  ]
	# z --> x or [0,0,1] --> [1,0,0]
    Rp = [ cosd(pitch)  0            sind(pitch)
           0            1            0.
          -sind(pitch)  0            cosd(pitch) ]
	# x --> y or [1,0,0] --> [0,1,0]
    Ry = [ cosd(yaw)   -sind(yaw)    0
           sind(yaw)    cosd(yaw)    0
           0            0            1.          ]
    # Applying them to a column vector from right to left
    # e.g., ypr is r(p(y(v)))
    if order     == "ypr"
        Rot = Rr*Rp*Ry
    elseif order == "yrp"
        Rot = Rp*Rr*Ry
    elseif order == "rpy"
        Rot = Ry*Rp*Rr
    elseif order == "ryp"
        Rot = Rp*Ry*Rr
    elseif order == "pyr"
        Rot = Rr*Ry*Rp
    else
        Rot = Ry*Rr*Rp
    end

    if doinv
        Rot'
    else
        Rot
    end
end

function get_misfit(m::Model, opt::Options, tempest::Bfield)
	chi2by2 = 0.0;
	if !opt.debug
		getfield!(m, tempest)
		idxx, idxz = tempest.selectx, tempest.selectz
        chi2by2 = getchi2by2([tempest.Hx[idxx]; tempest.Hz[idxz]],
                        [tempest.dataHx[idxx]; tempest.dataHz[idxz]],
                        [tempest.σx[idxx];tempest.σz[idxz]],
                        tempest.useML, tempest.ndatax + tempest.ndataz)
	end
	return chi2by2
end
#new function for vector sum elements to be called for plotting
function get_fm(Bx,Bz)
    fm = sqrt.(Bx.^2 + Bz.^2)
    fm
end
function get_dSigma(DBx,DBz,σx,σz)
    d = sqrt.(DBx.^2 + DBz.^2)
    σ = sqrt.((σx.^2).*DBx.^2 + (σz.^2).*DBz.^2)./d
    d,σ
end

function get_misfit(m::AbstractArray, mn::AbstractVector, opt::Union{Options,OptionsNuisance}, tempest::Bfield)
    # this is useful when reading history files and debugging
	chi2by2 = 0.0;
	if !opt.debug
		getfield!(m, mn, tempest)
		idxx, idxz = tempest.selectx, tempest.selectz
        if tempest.vectorsum
            fm = get_fm(tempest.Hx,tempest.Hz)
            d,σ = get_dSigma(tempest.dataHx,tempest.dataHz,tempest.σx,tempest.σz)
            chi2by2 = getchi2by2(fm, d, σ, tempest.useML, tempest.ndatax)
        else
            chi2by2 = getchi2by2([tempest.Hx[idxx]; tempest.Hz[idxz]],
                        [tempest.dataHx[idxx]; tempest.dataHz[idxz]],
                        [tempest.σx[idxx];tempest.σz[idxz]],
                        tempest.useML, tempest.ndatax + tempest.ndataz)
        end
	end
	return chi2by2
end

function get_misfit(m::Model, mn::ModelNuisance, opt::Union{Options,OptionsNuisance}, tempest::Bfield)
    # actually used in McMC
    get_misfit(m.fstar, mn.nuisance, opt, tempest)
end

function getchi2by2(fm, d, σ, useML, ndata)
    r = (fm - d)./σ
    if useML
        chi2by2 = 0.5*ndata*log(norm(r)^2)
    else
        chi2by2 = 0.5*norm(r)^2
    end
end

function plotmodelfield!(tempest::Bfield, z::Array{Float64,1}, ρ::Array{Float64,1})
	getfieldTD!(tempest, z, ρ)
	figure(figsize=(8,6))
	s = subplot(121)
	s.step(vcat(ρ[2:end],ρ[end]), vcat(z[2:end], z[end]+50))
	s.set_xscale("log")
	s.invert_xaxis()
	s.invert_yaxis()
	xlabel("ρ")
	ylabel("depth m")
	grid(true, which="both")
	s1 = subplot(122)
	semilogx(tempest.F.times, abs.(μ₀*tempest.Hz)*fTinv, label="Bz")
	semilogx(tempest.F.times, abs.(μ₀*tempest.Hx)*fTinv, label="Bx")
	if !isempty(tempest.σx)
		errorbar(tempest.F.times, μ₀*abs.(vec(tempest.dataHz))*fTinv, yerr = μ₀*2vec(tempest.σz)*fTinv,
                         linestyle="none", marker=".", elinewidth=0, capsize=3)
		errorbar(tempest.F.times, μ₀*abs.(vec(tempest.dataHx))*fTinv, yerr = μ₀*2vec(tempest.σx)*fTinv,
						 linestyle="none", marker=".", elinewidth=0, capsize=3)
	end
	xlabel("time s")
	ylabel("B field 10⁻¹⁵ T")
	legend()
	grid(true, which="both")
	!tempest.addprimary && s1.set_yscale("log")
	nicenup(gcf())
	nothing
end
function plotmodelfield!(tempest::Bfield, Ρ::Vector{Array{Float64}};
                        figsize=(8,5), dz=-1., onesigma=true,
                        extendfrac=-1., fsize=10, alpha=0.1)
	times, f, ax, nfixed, z  = setupaxis(tempest, Ρ, figsize, dz, onesigma,
                          extendfrac, fsize, alpha)
    for ρ in Ρ
        getfield!(ρ, tempest)
        tempest.Hz[.!tempest.selectz] .= NaN
        tempest.Hx[.!tempest.selectx] .= NaN
        ax[1].step(log10.(tempest.ρ[2:end]), tempest.z[2:end], "-k", alpha=alpha)
        ax[2].semilogx(times,μ₀*abs.(tempest.Hz)*fTinv, "k", alpha=alpha, markersize=2)
        ax[2].semilogx(times,μ₀*abs.(tempest.Hx)*fTinv, "k", alpha=alpha, markersize=2)
    end
	finishaxis(ax, f, z, dz, extendfrac, nfixed, fsize)
end

function plotmodelfield!(tempest::Bfield, Ρ::Vector{Array{Float64}},
						mn::Array{Float64,2};
                        figsize=(8,5), dz=-1., onesigma=true,
                        extendfrac=-1., fsize=10, alpha=0.1)
	@assert length(Ρ) == size(mn, 1)
	times, f, ax, nfixed, z  = setupaxis(tempest, Ρ, figsize, dz, onesigma,
                          extendfrac, fsize, alpha)
    for (i, ρ) in enumerate(Ρ)
        getfield!(ρ, mn[i,:], tempest)
        tempest.Hz[.!tempest.selectz] .= NaN
        tempest.Hx[.!tempest.selectx] .= NaN
        if tempest.vectorsum
            fm = get_fm(tempest.Hx,tempest.Hz)
            ax[1].step(log10.(tempest.ρ[2:end]), tempest.z[2:end], "-k", alpha=alpha)
            ax[2].semilogx(times,μ₀*abs.(fm)*fTinv, "k", alpha=alpha, markersize=2) #check with Anand
        else
            ax[1].step(log10.(tempest.ρ[2:end]), tempest.z[2:end], "-k", alpha=alpha)
            ax[2].semilogx(times,μ₀*abs.(tempest.Hz)*fTinv, "k", alpha=alpha, markersize=2)
            ax[2].semilogx(times,μ₀*abs.(tempest.Hx)*fTinv, "k", alpha=alpha, markersize=2)
        end
    end
	finishaxis(ax, f, z, dz, extendfrac, nfixed, fsize)
end

function setupaxis(tempest::Bfield, Ρ,
                   figsize, dz, onesigma,
                   extendfrac, fsize, alpha)
    @assert all((dz, extendfrac) .> 0)
    sigma = onesigma ? 1.0 : 2.0
    f = figure(figsize=figsize)
    ax = Vector{PyPlot.PyObject}(undef, 3)
    ax[1] = subplot(121)
    ρmin, ρmax = extrema(vcat(Ρ...))
    delρ = ρmax - ρmin
    ax[1].set_xlim(ρmin-0.1delρ,ρmax+0.1delρ)
    nfixed, z = tempest.nfixed, tempest.z
    ax[1].plot([ρmin-0.1delρ,ρmax+0.1delρ], z[nfixed+1]*[1., 1], color="b")
    ax[2] = subplot(122)
	times = tempest.F.times
    Hz, Hx = tempest.Hz, tempest.Hx
    dataHx, σx = tempest.dataHx, tempest.σx
	dataHz, σz = tempest.dataHz, tempest.σz
    if tempest.vectorsum
        d,σ = get_dSigma(dataHx,dataHz,σx,σz)
        ax[2].errorbar(times,μ₀*abs.(d)*fTinv, yerr = μ₀*sigma*σ*fTinv,
                        linestyle="none", marker=".", elinewidth=0, capsize=5, label="B")
    else
        ax[2].errorbar(times, μ₀*abs.(dataHz)*fTinv, yerr = μ₀*sigma*σz*fTinv,
                        linestyle="none", marker=".", elinewidth=0, capsize=3, label="Bz")
        ax[2].errorbar(times, μ₀*abs.(dataHx)*fTinv, yerr = μ₀*sigma*σx*fTinv,
                        linestyle="none", marker=".", elinewidth=0, capsize=3, label="Bx")
    end
	times, f, ax, nfixed, z
end

function finishaxis(ax, f, z, dz, extendfrac, nfixed, fsize)
	ax[1].grid()
    ax[1].set_ylabel("Depth m")
    ax[1].plot(xlim(), z[nfixed+1]*[1, 1], "--k")
    if dz > 0
        axn = ax[1].twinx()
        ax[1].get_shared_y_axes().join(ax[1],axn)
        yt = ax[1].get_yticks()[ax[1].get_yticks().>=z[nfixed+1]]
        axn.set_yticks(yt)
        axn.set_ylim(ax[1].get_ylim()[end:-1:1])
        axn.set_yticklabels(string.(Int.(round.(getn.(yt .- z[nfixed+1], dz, extendfrac)))))
    end
    axn.set_ylabel("Depth index", rotation=-90)
    ax[1].set_xlabel("Log₁₀ρ")
    ax[1].set_title("Model")
    ax[2].set_ylabel("B field 10¹⁵")
    ax[2].set_xlabel("Time (s)")
    ax[2].set_title("Transient response")
    ax[2].legend()
    ax[2].grid()
    ax[1].invert_xaxis()
    nicenup(f, fsize=fsize)
end

#for synthetics
function set_noisy_data!(tempest::Bfield, z::Array{Float64,1}, ρ::Array{Float64,1};
	noisefracx = 0.02, noisefracz = 0.02, rseed=123,
	halt_X = nothing, halt_Z = nothing)
	if halt_X != nothing
        @assert length(halt_X) == length(tempest.F.times)
    else
        halt_X = zeros(length(tempest.F.times))
    end
    if halt_Z != nothing
        @assert length(halt_Z) == length(tempest.F.times)
    else
        halt_Z = zeros(length(tempest.F.times))
    end
	primaryflag = tempest.addprimary
	if tempest.addprimary
		# adds noise only proportional to secondary field
		# when we compute full field next
		tempest.addprimary = false
	end
	getfieldTD!(tempest, z, ρ)
	σx = sqrt.((noisefracx*abs.(tempest.Hx)).^2 + (halt_X/μ₀).^2)
	σz = sqrt.((noisefracz*abs.(tempest.Hz)).^2 + (halt_Z/μ₀).^2)
	# reset the tempest primary field modeling flag to original
	tempest.addprimary = primaryflag
	set_noisy_data!(tempest, z, ρ, σx, σz, rseed=rseed)
	plotmodelfield!(tempest, z, ρ)
	nothing
end

function set_noisy_data!(tempest::Bfield, z::Array{Float64,1}, ρ::Array{Float64,1},
	σx, σz;rseed = 123)
	Random.seed!(rseed)
	getfieldTD!(tempest, z, ρ)
	set_noisy_data!(tempest,
		dataHx = tempest.Hx + σx.*randn(size(σx)),
		dataHz = tempest.Hz + σz.*randn(size(σz)),
		σx = σx,
		σz = σz)
	nothing
end

function set_noisy_data!(tempest::Bfield;
						dataHz = zeros(0), dataHx = zeros(0),
						σz = zeros(0), σx = zeros(0))

    @assert size(σx) == size(dataHx)
    @assert size(σz) == size(dataHz)
    ndatax  = sum(.!isnan.(dataHx))
    ndataz  = sum(.!isnan.(dataHz))
    selectx = .!isnan.(dataHx)
    selectz = .!isnan.(dataHz)

	tempest.σx = σx
	tempest.σz = σz
	tempest.dataHx = dataHx
	tempest.dataHz = dataHz
	tempest.selectx = selectx
	tempest.selectz = selectz
	tempest.ndatax = ndatax
	tempest.ndataz = ndataz
	nothing
end

mutable struct TempestSoundingData <: Sounding
    sounding_string :: String
    X       :: Float64
    Y       :: Float64
	Z       :: Float64
    fid     :: Float64
    linenum :: Int
    x_rx    :: Float64
	y_rx    :: Float64
	z_rx    :: Float64
	roll_rx :: Float64
	pitch_rx:: Float64
	yaw_rx  :: Float64
	roll_tx :: Float64
	pitch_tx:: Float64
	yaw_tx  :: Float64
    z_tx    :: Float64
    times   :: Array{Float64, 1}
    ramp    :: Array{Float64, 2}
    σ_x     :: Array{Float64, 1}
    σ_z     :: Array{Float64, 1}
    Hx_data :: Array{Float64, 1}
    Hz_data :: Array{Float64, 1}
end

function getnufromsounding(t::TempestSoundingData)
    # 1 zTx
    # 2 zRx
    # 3 x_rx
    # 4 y_rx
    # 5 rx_roll
    # 6 rx_pitch
    # 7 rx_yaw
    # 8 tx_roll
    # 9 tx_pitch
    # 10 tx_yaw
    return [t.z_tx, t.z_rx, t.x_rx, t.y_rx, t.roll_rx, t.pitch_rx, t.yaw_rx,
            t.roll_tx, t.pitch_tx, t.yaw_tx]
end

function read_survey_files(;
    fname_dat="",
    fname_specs_halt="",
    frame_height = -2,
    frame_dz = -2,
    frame_dx = -2,
    frame_dy = -2,
	Hxp = -1,
	Hzp = -1,
    Hxs = [-2, 2],
    Hzs = [-2, 2],
    units = 1e-15,
	yaw_rx = -1,
	pitch_rx = -1,
	roll_rx = -1,
	yaw_tx = -1,
	pitch_tx = -1,
	roll_tx = -1,
    figsize = (16,11),
    makesounding = false,
    dotillsounding = nothing,
    startfrom = 1,
    skipevery = 1,
    multnoise = 0.02,
    X = -1,
    Y = -1,
	Z = -1,
    fid = -1,
    linenum = -1,
	fsize = 13)

    @assert frame_height > 0
    @assert frame_dz > 0
    @assert frame_dx > 0
    @assert frame_dy > 0
    @assert all(Hxs .> 0)
    @assert all(Hzs .> 0)
	@assert Hxp > 0
	@assert Hzp > 0
    @assert X > 0
    @assert Y > 0
	@assert Z > 0
    @assert linenum > 0
    @assert fid > 0
	@assert pitch_rx > 0
	@assert roll_rx > 0
	@assert yaw_rx > 0
	@assert pitch_tx > 0
	@assert roll_tx > 0
	@assert yaw_tx > 0
	@assert 0 < multnoise < 1.0

    @info "reading $fname_dat"
    if dotillsounding!= nothing
        soundings = readdlm(fname_dat)[startfrom:skipevery:dotillsounding,:]
    else
        soundings = readdlm(fname_dat)[startfrom:skipevery:end,:]
    end
    easting = soundings[:,X]
    northing = soundings[:,Y]
	topo = soundings[:,Z]
    fiducial = soundings[:,fid]
    whichline = soundings[:,linenum]

	d_Hx = soundings[:,Hxs[1]:Hxs[2]] # secondary field
    d_Hz = soundings[:,Hzs[1]:Hzs[2]] # secondary field
    σ_Hx = multnoise*d_Hx # noise proportional to 2ndary
    σ_Hz = multnoise*d_Hz # noise proportional to 2ndary
	d_Hx .+= soundings[:,Hxp] # add primary field
	d_Hz .+= soundings[:,Hzp] # add primary field

    z_tx = soundings[:,frame_height]
    z_rx = -(z_tx + soundings[:,frame_dz]) # Flipping to my earth geometry
    z_tx = -z_tx # Flipping to my earth geometry
    x_rx = soundings[:,frame_dx]
	y_rx = -soundings[:,frame_dy] # Flipping to my earth geometry

	d_pitch_rx = -soundings[:,pitch_rx] # Flipping to GA-AEM geometry
	d_yaw_rx = -soundings[:,yaw_rx] # Flipping to GA-AEM geometry
	d_roll_rx = soundings[:,roll_rx]

	d_pitch_tx = -soundings[:,pitch_tx] # Flipping to GA-AEM geometry
	d_yaw_tx = -soundings[:,yaw_tx] # Flipping to GA-AEM geometry
	d_roll_tx = soundings[:,roll_tx]

    @info "reading $fname_specs_halt"
    include(fname_specs_halt)
    @assert size(d_Hx, 2) == length(times)
    @assert size(d_Hz, 2) == length(times)
    @assert size(d_Hx, 2) == length(Hx_add_noise)
    @assert size(d_Hz, 2) == length(Hz_add_noise)
	σ_Hx .*= units
	σ_Hz .*= units
    Hx_add_noise[:] .*= units
    Hz_add_noise[:] .*= units
    d_Hx[:]     .*= units
    d_Hz[:]     .*= -units # Flipping the Z component to align with GA_AEM rx
    f = figure(figsize=figsize)
    ax = Array{Any, 1}(undef, 4)
    ax[1] = subplot(2,2,1)
    nsoundings = size(soundings, 1)
    plot_dHx = permutedims(d_Hx)
    # plot_dHx[plot_dHx .<0] .= NaN
    pcolormesh(1:nsoundings, times, plot_dHx, shading="nearest")
    xlabel("sounding #")
    cbHx = colorbar()
    cbHx.ax.set_xlabel("Bx", fontsize=fsize)
    ylabel("time s")
    axHx = ax[1].twiny()
    axHx.semilogy(Hx_add_noise, times, "-k")
	axHx.semilogy(Hx_add_noise, times, "--w")
    axHx.set_xlabel("high alt noise")
    ax[2] = subplot(2,2,3,sharex=ax[1], sharey=ax[1])
    plot_dHz = permutedims(d_Hz)
    # plot_dHz[plot_dHM .<0] .= NaN
    pcolormesh(1:nsoundings, times, plot_dHz, shading="nearest")
    xlabel("sounding #")
    cbHz = colorbar()
    cbHz.ax.set_xlabel("Bz", fontsize=fsize)
    ylabel("time s")
    ax[2].invert_yaxis()
    axHz = ax[2].twiny()
    axHz.semilogy(Hz_add_noise, times, "-k")
	axHz.semilogy(Hz_add_noise, times, "--w")
    axHz.set_xlabel("high alt noise")
    ax[3] = subplot(2,2,2, sharex=ax[1])
	plot(1:nsoundings, z_tx, label="z_tx", "--")
	plot(1:nsoundings, z_rx, label="z_rx")
	plot(1:nsoundings, x_rx, label="x_rx")
	plot(1:nsoundings, y_rx, label="y_rx")
    legend()
    xlabel("sounding #")
    ylabel("earth frame geometry m")
	ax[3].grid()
    ax[3].invert_yaxis()
    ax[4] = subplot(2,2,4, sharex=ax[1])
    ax[4].plot(1:nsoundings, d_yaw_tx, label="yaw_tx", "--")
	ax[4].plot(1:nsoundings, d_pitch_tx, label="pitch_tx", "--")
	ax[4].plot(1:nsoundings, d_roll_tx, label="roll_tx", "--")
	ax[4].plot(1:nsoundings, d_yaw_rx, label="yaw_rx")
	ax[4].plot(1:nsoundings, d_pitch_rx, label="pitch_rx")
	ax[4].plot(1:nsoundings, d_roll_rx, label="roll_rx")
	ax[4].grid()
    xlabel("sounding #")
    ylabel("rotations degrees")
	legend()
	nicenup(f, fsize=fsize)
    if makesounding
        s_array = Array{TempestSoundingData, 1}(undef, nsoundings)
        for is in 1:nsoundings
            l, fi = Int(whichline[is]), fiducial[is]
            @info "read $is out of $nsoundings"
            dHx, dHz = vec(d_Hx[is,:]), vec(d_Hz[is,:])
            σHx = sqrt.(vec(σ_Hx[is,:].^2) + Hx_add_noise.^2)
            σHz = sqrt.(vec(σ_Hz[is,:].^2) + Hz_add_noise.^2)
            s_array[is] = TempestSoundingData(
				"sounding_$(l)_$fi", easting[is], northing[is],
				topo[is], fi, l,
				x_rx[is], y_rx[is], z_rx[is],
                d_roll_rx[is], d_pitch_rx[is], d_yaw_rx[is],
				d_roll_tx[is], d_pitch_tx[is], d_yaw_tx[is],
				z_tx[is],
                times, ramp,
                σHx, σHz, dHx, dHz
                )
        end
        return s_array
    end
end

function makeoperator( sounding::TempestSoundingData;
                       zfixed   = [-1e5],
                       ρfixed   = [1e12],
                       zstart = 0.0,
                       extendfrac = 1.06,
                       useML = false,
                       dz = 2.,
                       ρbg = 10,
                       nlayers = 40,
                       ntimesperdecade = 10,
                       nfreqsperdecade = 5,
                       showgeomplot = false,
                       plotfield = false,
					   addprimary = true,
                       vectorsum = false)
    @assert extendfrac > 1.0
    @assert dz > 0.0
    @assert ρbg > 0.0
    @assert nlayers > 1
    nmax = nlayers+1

    zall, znall, zboundaries = setupz(zstart, extendfrac, dz=dz, n=nlayers, showplot=showgeomplot)
    z, ρ, nfixed = makezρ(zboundaries; zfixed=zfixed, ρfixed=ρfixed)
    ρ[z.>=zstart] .= ρbg
    ## Tempest operator creation from sounding data
	aem = Bfield(
    zTx = sounding.z_tx, zRx = sounding.z_rx,
	x_rx = sounding.x_rx, y_rx = sounding.y_rx,
    rx_roll = sounding.roll_rx, rx_pitch = sounding.pitch_rx, rx_yaw = sounding.yaw_rx,
    tx_roll = sounding.roll_tx, tx_pitch = sounding.pitch_tx, tx_yaw = sounding.yaw_tx,
	ramp = sounding.ramp, times = sounding.times, useML = useML,
	z=z, ρ=ρ,
	addprimary = addprimary, #this ensures that the geometry update actually changes everything that needs to be
    vectorsum = vectorsum
	)

	set_noisy_data!(aem,
		dataHx = sounding.Hx_data/μ₀,
		dataHz = sounding.Hz_data/μ₀,
		σx = sounding.σ_x/μ₀,
		σz = sounding.σ_z/μ₀)

	plotfield && plotmodelfield!(aem, z, ρ)

    aem, znall
end
function make_tdgp_opt(sounding::TempestSoundingData;
                    rseed = nothing,
                    znall = [1],
                    fileprefix = "sounding",
                    nmin = 2,
                    nmax = 40,
                    K = GP.Mat32(),
                    demean = false,
                    sampledc = true,
                    sddc = 0.01,
                    sdpos = 0.05,
                    sdprop = 0.05,
                    fbounds = [-0.5 2.5],
                    λ = [2],
                    δ = 0.1,
                    pnorm = 2,
                    save_freq = 50,
					nuisance_sdev   = [0.],
					nuisance_bounds = [0. 0.],
					updatenuisances = true,
                    C = nothing,
					dispstatstoscreen = false,
					restart = false
                    )
    sdev_dc = sddc*diff(fbounds, dims=2)[:]
    sdev_pos = [sdpos*abs(diff([extrema(znall)...])[1])]
    sdev_prop = sdprop*diff(fbounds, dims=2)[:]
    xall = permutedims(collect(znall))
    xbounds = permutedims([extrema(znall)...])

	history_mode = "w"
	restart && (history_mode = "a")

    updatenonstat = false
    needλ²fromlog = false
    if rseed != nothing
        Random.seed!(rseed)
    end

    opt = OptionsStat(fdataname = fileprefix*"_",
                            nmin = nmin,
                            nmax = nmax,
                            xbounds = xbounds,
                            fbounds = fbounds,
                            xall = xall,
                            λ = λ,
                            δ = δ,
                            demean = demean,
                            sampledc = sampledc,
                            sdev_dc = sdev_dc,
                            sdev_prop = sdev_prop,
                            sdev_pos = sdev_pos,
                            pnorm = pnorm,
                            quasimultid = false,
                            K = K,
                            save_freq = save_freq,
                            needλ²fromlog = needλ²fromlog,
                            updatenonstat = updatenonstat,
							updatenuisances = updatenuisances,
                            dispstatstoscreen = dispstatstoscreen,
							history_mode = history_mode
                            )

	bounds = nuisance_bounds .+ [
      sounding.z_tx
      sounding.z_rx
      sounding.x_rx
      sounding.y_rx
      sounding.roll_rx
      sounding.pitch_rx
      sounding.yaw_rx
      sounding.roll_tx
      sounding.pitch_tx
      sounding.yaw_tx]

	optn = OptionsNuisance(opt;
                        sdev = copy(nuisance_sdev),
						bounds = bounds, C = C,
						updatenuisances = updatenuisances)

    opt, optn
end

function makeoperatorandoptions(soundings::Array{TempestSoundingData, 1};
                                                    nplot = 2,
                                                    zfixed   = [-1e5],
                                                    ρfixed   = [1e12],
                                                    zstart = 0.0,
                                                    extendfrac = 1.06,
                                                    useML = false,
                                                    dz = 2.,
                                                    ρbg = 10,
                                                    nlayers = 40,
                                                    ntimesperdecade = 10,
                                                    nfreqsperdecade = 5,
                                                    showgeomplot = false,
                                                    plotfield = true,
                                                    nmin = 2,
                                                    nmax = 40,
                                                    K = GP.Mat32(),
                                                    demean = false,
                                                    sampledc = true,
                                                    sddc = 0.01,
                                                    sdpos = 0.05,
                                                    sdprop = 0.05,
                                                    fbounds = [-0.5 2.5],
                                                    λ = [2],
                                                    δ = 0.1,
                                                    save_freq = 50,
                                                    nuisance_sdev   = [0.],
                                                    nuisance_bounds = [0. 0.],
                                                    C = nothing,
                                                    updatenuisances = true,
                                                    dispstatstoscreen = false,
                                                    vectorsum = false
                                                    )
    for idx in randperm(length(soundings))[1:nplot]
            aem, znall = makeoperator(soundings[idx],
                                   zfixed = zfixed,
                                   ρfixed = ρfixed,
                                   zstart = zstart,
                                   extendfrac = extendfrac,
                                   dz = dz,
                                   ρbg = ρbg,
                                   nlayers = nlayers,
                                   ntimesperdecade = ntimesperdecade,
                                   nfreqsperdecade = nfreqsperdecade,
                                   showgeomplot = showgeomplot,
                                   plotfield = plotfield,
                                   useML = useML,
                                   vectorsum = vectorsum)

            opt, optn = make_tdgp_opt(soundings[idx],
                                    znall = znall,
                                    fileprefix = soundings[idx].sounding_string,
                                    nmin = nmin,
                                    nmax = nmax,
                                    K = K,
                                    demean = demean,
                                    sampledc = sampledc,
                                    sddc = sddc,
                                    sdpos = sdpos,
                                    sdprop = sdprop,
                                    fbounds = fbounds,
                                    save_freq = save_freq,
                                    λ = λ,
                                    δ = δ,
                                    nuisance_bounds = nuisance_bounds,
                                    nuisance_sdev = nuisance_sdev,
                                    updatenuisances = updatenuisances,
                                    C = C,
                                    dispstatstoscreen = dispstatstoscreen
                                    )
    end

end

function summarypost(soundings::Array{TempestSoundingData, 1};
        qp1=0.05,
        qp2=0.95,
        burninfrac=0.5,
        zstart = 0.0,
        extendfrac = 1.06,
        useML = false,
        dz = 2.,
        nlayers = 40,
        nmin = 2,
        nmax = 40,
        K = GP.Mat32(),
        demean = false,
        sampledc = true,
        sddc = 0.01,
        sdpos = 0.05,
        sdprop = 0.05,
        fbounds = [-0.5 2.5],
        λ = [2],
        δ = 0.1,
        save_freq = 50,
        nuisance_sdev   = [0.],
        nuisance_bounds = [0. 0.],
        updatenuisances = true)

    @assert extendfrac > 1.0
    @assert dz > 0.0
    @assert nlayers > 1
    zall, znall, = setupz(zstart, extendfrac, dz=dz, n=nlayers)

    linename = "_line_$(soundings[1].linenum)_summary.txt"
    fnames = ["rho_low", "rho_mid", "rho_hi", "rho_avg",
        "ddz_mean", "ddz_sdev", "phid_mean", "phid_sdev",
        "nu_low", "nu_mid", "nu_high"].*linename

    idxnotzero = findidxnotzero(length(nuisance_sdev), nuisance_bounds)
    nunominal = zeros(length(idxnotzero), length(soundings))
    if isfile(fnames[1])
        @warn fnames[1]*" exists, reading stored values"
        pl, pm, ph, ρmean,
        vdmean, vddev, χ²mean, χ²sd,
        nulow, numid, nuhigh = map(x->readdlm(x), fnames)
        for idx = 1:length(soundings)
            nunominal[:,idx] = getnufromsounding(soundings[idx])[idxnotzero]
        end
    else
        # this is a dummy operator for plotting
        aem, = makeoperator(soundings[1])
        pl, pm, ph, ρmean, vdmean, vddev = map(x->zeros(length(zall), length(soundings)), 1:6)
        χ²mean, χ²sd = zeros(length(soundings)), zeros(length(soundings))
        nulow, numid, nuhigh  = map(x->zeros(length(idxnotzero), length(soundings)), 1:3)
        for idx = 1:length(soundings)
            opt, optn = make_tdgp_opt(soundings[idx],
                                        znall = znall,
                                        fileprefix = soundings[idx].sounding_string,
                                        nmin = nmin,
                                        nmax = nmax,
                                        K = K,
                                        demean = demean,
                                        sampledc = sampledc,
                                        sddc = sddc,
                                        sdpos = sdpos,
                                        sdprop = sdprop,
                                        fbounds = fbounds,
                                        save_freq = save_freq,
                                        λ = λ,
                                        δ = δ,
                                        nuisance_bounds = nuisance_bounds,
                                        nuisance_sdev = nuisance_sdev,
                                        updatenuisances = updatenuisances,
                                        dispstatstoscreen = false)
            opt.xall[:] .= zall
            @info "$idx out of $(length(soundings))\n"
            opt.xall[:] .= zall
            nunominal[:,idx] = mean(optn.bounds[idxnotzero,:], dims=2)
            pl[:,idx], pm[:,idx], ph[:,idx], ρmean[:,idx],
            vdmean[:,idx], vddev[:,idx] = CommonToAll.plot_posterior(aem, opt, burninfrac=burninfrac,
                                                    qp1=qp1, qp2=qp2,
                                                    doplot=false)
            h, nuquants = plot_posterior(aem, optn, burninfrac=burninfrac, doplot=false)
            nulow[:, idx]  .= nuquants[:, 1]
            numid[:, idx]  .= nuquants[:, 2]
            nuhigh[:, idx] .= nuquants[:, 3]
            χ² = 2*CommonToAll.assembleTat1(opt, :U, temperaturenum=1, burninfrac=burninfrac)
            ndata = sum(.!isnan.(soundings[idx].Hx_data)) +
                    sum(.!isnan.(soundings[idx].Hz_data))
            χ²mean[idx] = mean(χ²)/ndata
            χ²sd[idx]   = std(χ²)/ndata
            if useML
                χ²mean[idx] -= log(ndata)
                χ²sd[idx]   -= log(ndata) # I think, need to check
            end
        end
        # write in grid format
        for (fname, vals) in Dict(zip(fnames, [pl, pm, ph, ρmean, vdmean, vddev, χ²mean, χ²sd, nulow, numid, nuhigh]))
            writedlm(fname, vals)
        end
        # write in x, y, z, rho format
        for (i, d) in enumerate([pl, pm, ph, ρmean])
            xyzrho = makearray(soundings, d, zall)
            writedlm(fnames[i][1:end-4]*"_xyzrho.txt", xyzrho)
        end
    end
    pl, pm, ph, ρmean, vdmean, vddev, χ²mean, χ²sd, zall, nulow, numid, nuhigh, nunominal
end

function plotsummarygrids1(meangrid, phgrid, plgrid, pmgrid, gridx, gridz, topofine, R, Z, χ²mean, χ²sd, lname; qp1=0.05, qp2=0.95,
    figsize=(10,10), fontsize=12, cmap="viridis", vmin=-2, vmax=0.5, Eislast=true, Nislast=true,
    topowidth=2, idx=nothing, omitconvergence=false, useML=false, preferEright=false, preferNright=false,
    saveplot=false)

    f = figure(figsize=figsize)
    dr = diff(gridx)[1]
    f.suptitle(lname*" Δx=$dr m, Fids: $(length(R))", fontsize=fontsize)
    nrows = omitconvergence ? 4 : 5
    icol = 1
    s = Array{Any, 1}(undef, nrows)
    if !omitconvergence
        s[icol] = subplot(nrows, 1, icol)
        if useML
            plot(R, exp.(χ²mean))
            semilogy(R, ones(length(R)), "--k")
            ylabel("variance factor")
            title("Max likelihood variance adjustment")
        else
            plot(R, χ²mean)
            plot(R, ones(length(R)), "--k")
            fill_between(R, vec(χ²mean-χ²sd), vec(χ²mean+χ²sd), alpha=0.5)
            ylabel(L"ϕ_d")
            title("Data misfit")
        end
        icol += 1
    end
    s[icol] = omitconvergence ? subplot(nrows, 1, icol) : subplot(nrows, 1, icol, sharex=s[icol-1])
    imshow(plgrid, cmap=cmap, aspect="auto", vmax=vmax, vmin = vmin,
    extent=[gridx[1], gridx[end], gridz[end], gridz[1]])
    plot(gridx, topofine, linewidth=topowidth, "-k")
    idx == nothing || plotprofile(s[icol], idx, Z, R)
    title("Percentile $(round(Int, 100*qp1)) conductivity")
    ylabel("Height m")
    icol += 1
    s[icol] = subplot(nrows, 1, icol, sharex=s[icol-1], sharey=s[icol-1])
    imshow(pmgrid, cmap=cmap, aspect="auto", vmax=vmax, vmin = vmin,
    extent=[gridx[1], gridx[end], gridz[end], gridz[1]])
    plot(gridx, topofine, linewidth=topowidth, "-k")
    title("Percentile 50 conductivity")
    idx == nothing || plotprofile(s[icol], idx, Z, R)
    ylabel("Height m")
    icol += 1
    s[icol] = subplot(nrows, 1, icol, sharex=s[icol-1], sharey=s[icol-1])
    imshow(meangrid, cmap=cmap, aspect="auto", vmax=vmax, vmin = vmin,
    extent=[gridx[1], gridx[end], gridz[end], gridz[1]])
    plot(gridx, topofine, linewidth=topowidth, "-k")
    title("Mean conductivity")
    ylabel("Height m")
    idx == nothing || plotprofile(s[icol], idx, Z, R)
    icol +=1
    s[icol] = subplot(nrows, 1, icol, sharex=s[icol-1], sharey=s[icol-1])
    imlast = imshow(phgrid, cmap=cmap, aspect="auto", vmax=vmax, vmin = vmin,
    extent=[gridx[1], gridx[end], gridz[end], gridz[1]])
    plot(gridx, topofine, linewidth=topowidth, "-k")
    xlabel("Line distance m")
    title("Percentile $(round(Int, 100*qp2)) conductivity")
    ylabel("Height m")
    idx == nothing || plotprofile(s[icol], idx, Z, R)
    xlim(extrema(gridx))
    map(x->x.tick_params(labelbottom=false), s[1:end-1])
    map(x->x.grid(), s)
    nicenup(f, fsize=fontsize)
    plotNEWSlabels(Eislast, Nislast, gridx, gridz, s[(omitconvergence ? 1 : 2 ): end])
    f.subplots_adjust(bottom=0.125)
    cbar_ax = f.add_axes([0.125, 0.05, 0.75, 0.01])
    cb = f.colorbar(imlast, cax=cbar_ax, orientation="horizontal")
    cb.ax.set_xlabel("Log₁₀ S/m")
    (preferNright && !Nislast) && s[end].invert_xaxis()
    (preferEright && !Eislast) && s[end].invert_xaxis()
    saveplot && savefig(lname*".png", dpi=300)
end

function plotsummarygrids2(σmeangrid, ∇zmeangrid, ∇zsdgrid, cigrid, gridx, gridz, topofine, lname;
    qp1=0.05, qp2=0.95, Eislast=true, Nislast=true,
    figsize=(10,10), fontsize=12, cmap="viridis", vmin=-2, vmax=0.5, topowidth=2)

    f = figure(figsize=figsize)
    f.suptitle(lname, fontsize=fontsize)
    s1 = subplot(411)
    imshow(σmeangrid, cmap=cmap, aspect="auto", vmax=vmax, vmin = vmin,
                extent=[gridx[1], gridx[end], gridz[end], gridz[1]])
    title("Mean conductivity")
    ylabel("Height m")
    colorbar()
    s2 = subplot(412, sharex=s1)
    imshow(abs.(cigrid), cmap=cmap, aspect="auto",
                extent=[gridx[1], gridx[end], gridz[end], gridz[1]])
    plot(gridx, topofine, linewidth=topowidth, "-k")
    title("CI: $(round(Int, 100*(qp2-qp1))) conductivity")
    ylabel("Height m")
    colorbar()
    s3 = subplot(413, sharex=s2, sharey=s2)
    imshow(∇zmeangrid, cmap=cmap, aspect="auto",
                extent=[gridx[1], gridx[end], gridz[end], gridz[1]])
    plot(gridx, topofine, linewidth=topowidth, "-k")
    title("Mean conductivity vertical derivative")
    ylabel("Height m")
    colorbar()
    s4 = subplot(414, sharex=s2, sharey=s2)
    imlast = imshow(∇zsdgrid, cmap=cmap, aspect="auto",
                extent=[gridx[1], gridx[end], gridz[end], gridz[1]])
    plot(gridx, topofine, linewidth=topowidth, "-k")
    xlabel("Line distance m")
    title("Std dev of conductivity vertical derivative")
    ylabel("Height m")
    colorbar()
    xlim(extrema(gridx))
    nicenup(f, fsize=fontsize)
end

function plotsummarygrids3(nuhigh, nulow, numid, phgrid, plgrid, pmgrid, gridx, gridz, topofine, R, Z, χ²mean, χ²sd, lname, nunominal, numsize=2; qp1=0.05, qp2=0.95,
    figsize=(10,10), fontsize=12, cmap="viridis", vmin=-2, vmax=0.5, Eislast=true, Nislast=true,
    topowidth=2, idx=nothing, useML=false, preferEright=false, preferNright=false, labelnu=[""],
    saveplot=false)

    # get the finely interpolated nuisances
    # nuhighfine, nulowfine, numidfine = map(X->gridpoints(R, gridx, X), [nuhigh, nulow, numid])

    f = figure(figsize=figsize)
    dr = diff(gridx)[1]
    f.suptitle(lname*" Δx=$dr m, Fids: $(length(R))", fontsize=fontsize)
    nrows = 4 + size(nulow, 1) # add the number of nuisances == no. of rows in nuhigh
    icol = 1
    s = Array{Any, 1}(undef, nrows)
    s[icol] = subplot(nrows, 1, icol)
    if useML
        plot(R, exp.(χ²mean))
        semilogy(R, ones(length(R)), "--k")
        ylabel("variance factor")
        title("Max likelihood noise variance adjustment")
    else
        plot(R, χ²mean)
        plot(R, ones(length(R)), "--k")
        fill_between(R, vec(χ²mean-χ²sd), vec(χ²mean+χ²sd), alpha=0.5)
        ylabel(L"ϕ_d")
        title("Data misfit")
    end
    icol += 1
    icol = plotnuquant(nulow, numid, nuhigh, nunominal, s, R, icol, nrows, numsize, labelnu)
    s[icol] = subplot(nrows, 1, icol, sharex=s[icol-1])
    imshow(plgrid, cmap=cmap, aspect="auto", vmax=vmax, vmin = vmin,
    extent=[gridx[1], gridx[end], gridz[end], gridz[1]])
    plot(gridx, topofine, linewidth=topowidth, "-k")
    idx == nothing || plotprofile(s[icol], idx, Z, R)
    title("Percentile $(round(Int, 100*qp1)) conductivity")
    ylabel("Height m")
    icol += 1
    s[icol] = subplot(nrows, 1, icol, sharex=s[icol-1], sharey=s[icol-1])
    imshow(pmgrid, cmap=cmap, aspect="auto", vmax=vmax, vmin = vmin,
    extent=[gridx[1], gridx[end], gridz[end], gridz[1]])
    plot(gridx, topofine, linewidth=topowidth, "-k")
    title("Percentile 50 conductivity")
    idx == nothing || plotprofile(s[icol], idx, Z, R)
    ylabel("Height m")
    icol += 1
    s[icol] = subplot(nrows, 1, icol, sharex=s[icol-1], sharey=s[icol-1])
    imlast = imshow(phgrid, cmap=cmap, aspect="auto", vmax=vmax, vmin = vmin,
    extent=[gridx[1], gridx[end], gridz[end], gridz[1]])
    plot(gridx, topofine, linewidth=topowidth, "-k")
    xlabel("Line distance m")
    title("Percentile $(round(Int, 100*qp2)) conductivity")
    ylabel("Height m")
    idx == nothing || plotprofile(s[icol], idx, Z, R)
    xlim(extrema(gridx))
    map(x->x.tick_params(labelbottom=false), s[1:end-1])
    map(x->x.grid(), s)
    nicenup(f, fsize=fontsize)
    plotNEWSlabels(Eislast, Nislast, gridx, gridz, s[size(nulow, 1)+1:end])
    f.subplots_adjust(bottom=0.125)
    cbar_ax = f.add_axes([0.125, 0.05, 0.75, 0.01])
    cb = f.colorbar(imlast, cax=cbar_ax, orientation="horizontal")
    cb.ax.set_xlabel("Log₁₀ S/m")
    (preferNright && !Nislast) && s[end].invert_xaxis()
    (preferEright && !Eislast) && s[end].invert_xaxis()
    saveplot && savefig(lname*"_with_nu.png", dpi=300)
end

function plotnuquant(nqlow, nqmid, nqhigh, nunominal, s, gridx, icol, nrows, ms=2, labelnu=[""])
    nnu = size(nqlow, 1)
    labelnu[1] == "" || @assert length(labelnu) == nnu
    for inu = 1:nnu
        s[icol] = subplot(nrows, 1, icol, sharex=s[icol-1])
        s[icol].fill_between(gridx, nqlow[inu,:], nqhigh[inu,:], alpha=0.5)
        s[icol].plot(gridx, nqmid[inu,:])
        s[icol].plot(gridx, nunominal[inu,:], "o", markersize=ms)
        labelnu[1] == "" || s[icol].set_title(labelnu[inu])
        icol += 1
    end
    icol
end

function summaryimages(soundings::Array{TempestSoundingData, 1};
                        qp1=0.05,
                        qp2=0.95,
                        burninfrac=0.5,
                        zstart = 0.0,
                        extendfrac = 1.06,
                        useML = false,
                        dz = 2.,
                        nlayers = 40,
                        nmin = 2,
                        nmax = 40,
                        K = GP.Mat32(),
                        demean = false,
                        sampledc = true,
                        sddc = 0.01,
                        sdpos = 0.05,
                        sdprop = 0.05,
                        fbounds = [-0.5 2.5],
                        λ = [2],
                        δ = 0.1,
                        save_freq = 50,
                        nuisance_sdev   = [0.],
                        nuisance_bounds = [0. 0.],
                        updatenuisances = true,
                        dr = 10,
                        fontsize = 10,
                        vmin = -2,
                        vmax = 0.5,
                        cmap="viridis",
                        figsize=(6,10),
                        topowidth=2,
                        idx = nothing,
                        showderivs = false,
                        omitconvergence = false,
                        preferEright = false,
                        preferNright = false,
                        saveplot = false,
                        numsize=1,
                        labelnu = [""]
                        )
    @assert !(preferNright && preferEright) # can't prefer both labels to the right
    pl, pm, ph, ρmean, vdmean, vddev, χ²mean, χ²sd, zall,
    nulow, numid, nuhigh, nunominal = summarypost(soundings,
                                        qp1 = qp1,
                                        qp2 = qp2,
                                        burninfrac = burninfrac,
                                        zstart = zstart,
                                        extendfrac = extendfrac,
                                        useML = useML,
                                        dz = dz,
                                        nlayers = nlayers,
                                        nmin = nmin,
                                        nmax = nmax,
                                        K = K,
                                        demean = demean,
                                        sampledc = sampledc,
                                        sddc = sddc,
                                        sdpos = sdpos,
                                        sdprop = sdprop,
                                        fbounds = fbounds,
                                        λ = λ,
                                        δ = δ,
                                        save_freq = save_freq,
                                        nuisance_sdev   = nuisance_sdev,
                                        nuisance_bounds = nuisance_bounds,
                                        updatenuisances = updatenuisances)

    phgrid, plgrid, pmgrid, σmeangrid, ∇zmeangrid,
    ∇zsdgrid, gridx, gridz, topofine, R, Z = makesummarygrid(soundings, pl, pm, ph, ρmean,
                                            vdmean, vddev, zall, dz, dr=dr)

    lname = "Line $(soundings[1].linenum)"
    Eislast, Nislast = whichislast(soundings)
    plotsummarygrids1(σmeangrid, phgrid, plgrid, pmgrid, gridx, gridz, topofine, R, Z, χ²mean, χ²sd, lname, qp1=qp1, qp2=qp2,
        figsize=figsize, fontsize=fontsize, cmap=cmap, vmin=vmin, vmax=vmax, Eislast=Eislast,
        Nislast=Nislast, topowidth=topowidth, idx=idx, omitconvergence=omitconvergence, useML=useML,
        preferEright=preferEright, preferNright=preferNright, saveplot=saveplot)

    plotsummarygrids3(nuhigh, nulow, numid, phgrid, plgrid, pmgrid, gridx, gridz, topofine, R, Z, χ²mean, χ²sd, lname, nunominal, numsize, labelnu=labelnu, qp1=qp1, qp2=qp2,
    figsize=figsize, fontsize=fontsize, cmap=cmap, vmin=vmin, vmax=vmax, Eislast=Eislast,
    Nislast=Nislast, topowidth=topowidth, idx=idx, useML=useML,
    preferEright=preferEright, preferNright=preferNright, saveplot=saveplot)

    if showderivs
    cigrid = phgrid - plgrid
    plotsummarygrids2(σmeangrid, ∇zmeangrid, ∇zsdgrid, cigrid, gridx, gridz, topofine, lname, qp1=qp1, qp2=qp2,
        figsize=figsize, fontsize=fontsize, cmap=cmap, vmin=vmin, vmax=vmax, topowidth=topowidth,
        Eislast=Eislast, Nislast=Nislast)
    end
end

function plotindividualsoundings(soundings::Array{TempestSoundingData, 1};
                        burninfrac=0.5,
                        zstart = 0.0,
                        extendfrac = 1.06,
                        dz = 2.,
                        nlayers = 40,
                        nmin = 2,
                        nmax = 40,
                        K = GP.Mat32(),
                        demean = false,
                        sampledc = true,
                        sddc = 0.01,
                        sdpos = 0.05,
                        sdprop = 0.05,
                        fbounds = [-0.5 2.5],
                        λ = [2],
                        δ = 0.1,
                        save_freq = 50,
                        nuisance_sdev   = [0.],
                        nuisance_bounds = [0. 0.],
                        updatenuisances = true,
                        nbins=100,
                        figsize  = (12,6),
                        zfixed   = [-1e5],
                        ρfixed   = [1e12],
                        ntimesperdecade = 10,
                        nfreqsperdecade = 5,
                        computeforwards = false,
                        nforwards = 100,
                        vectorsum = false,
                      idxcompute = [1])
    for idx = 1:length(soundings)
        if in(idx, idxcompute)
            @info "Sounding number: $idx"
            aem, znall = makeoperator(soundings[idx],
                zfixed = zfixed,
                ρfixed = ρfixed,
                zstart = zstart,
                extendfrac = extendfrac,
                dz = dz,
                nlayers = nlayers,
                ntimesperdecade = ntimesperdecade,
                nfreqsperdecade = nfreqsperdecade,
                vectorsum = vectorsum)
            opt, optn = make_tdgp_opt(soundings[idx],
                                        znall = znall,
                                        fileprefix = soundings[idx].sounding_string,
                                        nmin = nmin,
                                        nmax = nmax,
                                        K = K,
                                        demean = demean,
                                        sampledc = sampledc,
                                        sddc = sddc,
                                        sdpos = sdpos,
                                        sdprop = sdprop,
                                        fbounds = fbounds,
                                        save_freq = save_freq,
                                        λ = λ,
                                        δ = δ,
                                        nuisance_bounds = nuisance_bounds,
                                        nuisance_sdev = nuisance_sdev,
                                        updatenuisances = updatenuisances,
                                        dispstatstoscreen = false)
            zall, znall, = setupz(zstart, extendfrac, dz=dz, n=nlayers)
            opt.xall[:] .= zall
            getchi2forall(opt, alpha=0.8) # chi2 errors
            CommonToAll.getstats(opt) # ARs for GP model
            CommonToAll.getstats(optn) # ARs for nuisances
            plot_posterior(aem, opt, burninfrac=burninfrac, nbins=nbins, figsize=figsize) # GP models
            ax = gcf().axes
            ax[1].invert_xaxis()
            plot_posterior(aem, optn, burninfrac=burninfrac, nbins=nbins, figsize=figsize) # nuisances
            if computeforwards
                m = assembleTat1(opt, :fstar, temperaturenum=1, burninfrac=burninfrac)
                mn = CommonToAll.assemblenuisancesatT(optn, temperaturenum=1, burninfrac=burninfrac)
                Random.seed!(10)
                plotmodelfield!(aem, m[randperm(length(m))[1:nforwards]],
                                                          mn[randperm(length(m))[1:nforwards],:],
                                                          dz=dz, extendfrac=extendfrac)
            end
        end
    end
end

# there are definitely nuisances, e.g. TEMPEST
function loopacrosssoundings(soundings::Array{S, 1};
                                nsequentialiters   = -1,
                                nparallelsoundings = -1,
                                zfixed             = [-1e5],
                                ρfixed             = [1e12],
                                zstart             = 0.0,
                                extendfrac         = 1.06,
                                dz                 = 2.,
                                ρbg                = 10,
                                nlayers            = 50,
                                ntimesperdecade    = 10,
                                nfreqsperdecade    = 5,
                                Tmax               = -1,
                                nsamples           = -1,
                                nchainsatone       = -1,
                                nchainspersounding = -1,
                                nmin               = 2,
                                nmax               = 40,
                                K                  = GP.Mat32(),
                                demean             = true,
                                sampledc           = true,
                                sddc               = 0.01,
                                sdpos              = 0.05,
                                sdprop             = 0.05,
                                fbounds            = [-0.5 2.5],
                                λ                  = [2],
                                δ                  = 0.1,
                                pnorm              = 2,
                                save_freq          = 50,
                                nuisance_sdev      = [0.],
                                nuisance_bounds    = [0. 0.],
                                updatenuisances    = true,
                                dispstatstoscreen  = false,
                                useML              = false,
                                restart            = false,
                                C                  = nothing,
                                vectorsum          = false) where S<:Sounding

    @assert nsequentialiters  != -1
    @assert nparallelsoundings != -1
    @assert nchainspersounding != -1
    @assert nsamples != - 1
    @assert nchainsatone != -1
    @assert Tmax != -1

    nsoundings = length(soundings)

    for iter = 1:nsequentialiters
        if iter<nsequentialiters
            ss = (iter-1)*nparallelsoundings+1:iter*nparallelsoundings
        else
            ss = (iter-1)*nparallelsoundings+1:nsoundings
        end
        @info "soundings in loop $iter of $nsequentialiters", ss
        r_nothing = Array{Nothing, 1}(undef, length(ss))
        @sync for (i, s) in Iterators.reverse(enumerate(ss))
            pids = (i-1)*nchainspersounding+i:i*(nchainspersounding+1)
            @info "pids in sounding $s:", pids

            aem, znall = makeoperator(    soundings[s],
                                    zfixed = zfixed,
                                    ρfixed = ρfixed,
                                    zstart = zstart,
                                    extendfrac = extendfrac,
                                    dz = dz,
                                    ρbg = ρbg,
                                    useML = useML,
                                    nlayers = nlayers,
                                    ntimesperdecade = ntimesperdecade,
                                    nfreqsperdecade = nfreqsperdecade,
                                    vectorsum = vectorsum)

            opt, optn = make_tdgp_opt(soundings[s],
                                znall = znall,
                                fileprefix = soundings[s].sounding_string,
                                nmin = nmin,
                                nmax = nmax,
                                K = K,
                                demean = demean,
                                sampledc = sampledc,
                                sddc = sddc,
                                sdpos = sdpos,
                                sdprop = sdprop,
                                fbounds = fbounds,
                                save_freq = save_freq,
                                λ = λ,
                                δ = δ,
                                nuisance_bounds = nuisance_bounds,
                                nuisance_sdev = nuisance_sdev,
                                updatenuisances = updatenuisances,
                                C = C,
                                restart = restart,
                                dispstatstoscreen = dispstatstoscreen
                                )

            @async r_nothing[i] = remotecall_fetch(main, pids[1], opt, optn, aem, collect(pids[2:end]),
                                    Tmax         = Tmax,
                                    nsamples     = nsamples,
                                    nchainsatone = nchainsatone)

        end # @sync
        @info "done $iter out of $nsequentialiters at $(Dates.now())"
    end
end

end

max_stars_repo_name: GeoscienceAustralia/HiQGA

max_issues_repo_name: GeoscienceAustralia/HiQGA

max_forks_repo_name: GeoscienceAustralia/HiQGA


text:
#=
  Second assignment of the CPLINT course:
  https://edu.swi-prolog.org/mod/assign/view.php?id=243
  """
  http://cplint.eu/p/cards.swinb
  Cards

  Suppose you have two decks of poker cards (52 cards, 4 suits, 13 ranks: A, K, Q, J, 10, 9, 8, 7, 6, 5, 4, 3, 2).

  Suppose you draw a card from the first deck and one from the second.

  Write a program to compute the probability that in this way you obtain one pair (two cards of the same rank).

  Can you do it with a single probabilistic clause?

  Add code to compute the probability that you draw at least one ace.
  """

  Note that the two cards are from two different deck of cards.

  Compare woth ~/blog/cards.blog
      ~/psi/cards.blog
      ~/webppl/cards.wppl

    Result:
    Summary Statistics
                   parameters      mean       std   naive_se      mcse         ess      rhat
                       Symbol   Float64   Float64    Float64   Float64     Float64   Float64

  a_king_or_a_queen_in_spades    0.0684    0.2524     0.0009    0.0060   1274.0013    1.0011
              at_least_an_ace    0.1650    0.3712     0.0013    0.0088   1327.6489    1.0020
              exactly_one_ace    0.1603    0.3669     0.0013    0.0087   1311.2268    1.0024
                         pair    0.0736    0.2612     0.0009    0.0061   1442.6291    1.0015
                      rank[1]    6.9292    3.7412     0.0132    0.0890   1256.8658    1.0015
                      rank[2]    7.0677    3.7526     0.0133    0.0885   1412.4389    1.0023
                      suit[1]    2.5018    1.1335     0.0040    0.0270   1184.4944    1.0059
                      suit[2]    2.4984    1.1315     0.0040    0.0275   1208.0383    1.0030
                     two_aces    0.0055    0.0738     0.0003    0.0017   1298.4260    1.0009

    ....

    2.485510 seconds (36.82 M allocations: 1.911 GiB, 20.05% gc time)

    According to the WebPPL model (cards.wppl) the _exact_ results are:
    'pair', 0.07692307692307696
    'at_least_an_ace', 0.14792899408284022
    'exactly_one_ace', 0.1420118343195266
    'two_aces', 0.005917159763313609
    'a_kind_or_a_queen_in_spades', 0.07544378698224857

=#

import Turing #, StatsPlots, DataFrames

@model function cards()
    # TODO: Check if it would work with @enums
    # @enum Suits hearts spades clubs diamonds
    # @enume Ranks ace king queen jack v10 v9 v8 v7 v6 v5 v4 v3 v2

    num_suits = 4
    hearts, spades, clubs, diamonds = 1:num_suits

    num_ranks = 13
    ace, king, queen, jack, v10, v9, v8, v7, v6, v5, v4, v3, v2 = 1:num_ranks

    num_cards = 2

    rank = tzeros(num_cards) # TArray{Int}(undef,num_cards)
    suit = tzeros(num_suits) # TArray{Int}(undef,num_cards)
    for c in 1:num_cards
        suit[c] ~ DiscreteUniform(1,num_suits)
        rank[c] ~ DiscreteUniform(1,num_ranks)
    end

    # Probability of a pair
    # It would be less complex if Turing has an @observe macro
    pair ~ Bernoulli(0.5)
    if pair
       true ~ Dirac(rank[1] == rank[2])
    else
       true ~ Dirac(rank[1] != rank[2])
    end

    # At least an ace
    at_least_an_ace ~ Bernoulli(0.5)
    if at_least_an_ace
        true ~ Dirac(rank[1] == ace || rank[2] == ace)
    else
        true ~ Dirac(rank[1] != ace && rank[2] != ace)
    end

    # Exactly one ace
    exactly_one_ace ~ Bernoulli(0.5)
    if exactly_one_ace
        true ~ Dirac((rank[1] == ace) + (rank[2] == ace) == 1)
    else
        true ~ Dirac((rank[1] == ace) + (rank[2] == ace) != 1)
    end

    # Two aces
    two_aces ~ Bernoulli(0.5)
    if two_aces
        true ~ Dirac((rank[1] == ace) + (rank[2] == ace) == 2)
    else
        true ~ Dirac((rank[1] == ace) + (rank[2] == ace) != 2)
    end

    # A King or a Queen in spades
    a_king_or_a_queen_in_spades ~ Bernoulli(0.5)
    if a_king_or_a_queen_in_spades
        true ~ Dirac(  (
             ((rank[1] == king) || (rank[1] == queen)) && suit[1] == spades
            ) ||
            (
             ((rank[2] == king) || (rank[2] == queen)) && suit[2] == spades
            )
        )
    else
        true ~ Dirac(!((((rank[1] == king) || (rank[1] == queen)) && suit[1] == spades) ||
         (((rank[2] == king) || (rank[2] == queen)) && suit[2] == spades) ))
    end

end

model = cards()

num_chains = 4
# chains = sample(model, Prior(), MCMCThreads(), 1000, num_chains)
chains = sample(model, MH(), MCMCThreads(), 10_000, num_chains) # 1.47s

display(chains)

max_stars_repo_name: Wikunia/hakank

max_issues_repo_name: Wikunia/hakank

max_forks_repo_name: Wikunia/hakank


text:
# syntax: proto2
using Compat
using ProtoBuf
import ProtoBuf.meta

mutable struct TokenProto <: ProtoType
    identifier::Array{UInt8,1}
    password::Array{UInt8,1}
    kind::AbstractString
    service::AbstractString
    TokenProto(; kwargs...) = (o=new(); fillunset(o); isempty(kwargs) || ProtoBuf._protobuild(o, kwargs); o)
end #mutable struct TokenProto
const __req_TokenProto = Symbol[:identifier,:password,:kind,:service]
meta(t::Type{TokenProto}) = meta(t, __req_TokenProto, ProtoBuf.DEF_FNUM, ProtoBuf.DEF_VAL, true, ProtoBuf.DEF_PACK, ProtoBuf.DEF_WTYPES, ProtoBuf.DEF_ONEOFS, ProtoBuf.DEF_ONEOF_NAMES, ProtoBuf.DEF_FIELD_TYPES)

mutable struct GetDelegationTokenRequestProto <: ProtoType
    renewer::AbstractString
    GetDelegationTokenRequestProto(; kwargs...) = (o=new(); fillunset(o); isempty(kwargs) || ProtoBuf._protobuild(o, kwargs); o)
end #mutable struct GetDelegationTokenRequestProto
const __req_GetDelegationTokenRequestProto = Symbol[:renewer]
meta(t::Type{GetDelegationTokenRequestProto}) = meta(t, __req_GetDelegationTokenRequestProto, ProtoBuf.DEF_FNUM, ProtoBuf.DEF_VAL, true, ProtoBuf.DEF_PACK, ProtoBuf.DEF_WTYPES, ProtoBuf.DEF_ONEOFS, ProtoBuf.DEF_ONEOF_NAMES, ProtoBuf.DEF_FIELD_TYPES)

mutable struct GetDelegationTokenResponseProto <: ProtoType
    token::TokenProto
    GetDelegationTokenResponseProto(; kwargs...) = (o=new(); fillunset(o); isempty(kwargs) || ProtoBuf._protobuild(o, kwargs); o)
end #mutable struct GetDelegationTokenResponseProto

mutable struct RenewDelegationTokenRequestProto <: ProtoType
    token::TokenProto
    RenewDelegationTokenRequestProto(; kwargs...) = (o=new(); fillunset(o); isempty(kwargs) || ProtoBuf._protobuild(o, kwargs); o)
end #mutable struct RenewDelegationTokenRequestProto
const __req_RenewDelegationTokenRequestProto = Symbol[:token]
meta(t::Type{RenewDelegationTokenRequestProto}) = meta(t, __req_RenewDelegationTokenRequestProto, ProtoBuf.DEF_FNUM, ProtoBuf.DEF_VAL, true, ProtoBuf.DEF_PACK, ProtoBuf.DEF_WTYPES, ProtoBuf.DEF_ONEOFS, ProtoBuf.DEF_ONEOF_NAMES, ProtoBuf.DEF_FIELD_TYPES)

mutable struct RenewDelegationTokenResponseProto <: ProtoType
    newExpiryTime::UInt64
    RenewDelegationTokenResponseProto(; kwargs...) = (o=new(); fillunset(o); isempty(kwargs) || ProtoBuf._protobuild(o, kwargs); o)
end #mutable struct RenewDelegationTokenResponseProto
const __req_RenewDelegationTokenResponseProto = Symbol[:newExpiryTime]
meta(t::Type{RenewDelegationTokenResponseProto}) = meta(t, __req_RenewDelegationTokenResponseProto, ProtoBuf.DEF_FNUM, ProtoBuf.DEF_VAL, true, ProtoBuf.DEF_PACK, ProtoBuf.DEF_WTYPES, ProtoBuf.DEF_ONEOFS, ProtoBuf.DEF_ONEOF_NAMES, ProtoBuf.DEF_FIELD_TYPES)

mutable struct CancelDelegationTokenRequestProto <: ProtoType
    token::TokenProto
    CancelDelegationTokenRequestProto(; kwargs...) = (o=new(); fillunset(o); isempty(kwargs) || ProtoBuf._protobuild(o, kwargs); o)
end #mutable struct CancelDelegationTokenRequestProto
const __req_CancelDelegationTokenRequestProto = Symbol[:token]
meta(t::Type{CancelDelegationTokenRequestProto}) = meta(t, __req_CancelDelegationTokenRequestProto, ProtoBuf.DEF_FNUM, ProtoBuf.DEF_VAL, true, ProtoBuf.DEF_PACK, ProtoBuf.DEF_WTYPES, ProtoBuf.DEF_ONEOFS, ProtoBuf.DEF_ONEOF_NAMES, ProtoBuf.DEF_FIELD_TYPES)

mutable struct CancelDelegationTokenResponseProto <: ProtoType
    CancelDelegationTokenResponseProto(; kwargs...) = (o=new(); fillunset(o); isempty(kwargs) || ProtoBuf._protobuild(o, kwargs); o)
end #mutable struct CancelDelegationTokenResponseProto

export TokenProto, GetDelegationTokenRequestProto, GetDelegationTokenResponseProto, RenewDelegationTokenRequestProto, RenewDelegationTokenResponseProto, CancelDelegationTokenRequestProto, CancelDelegationTokenResponseProto