using LinearAlgebraicRepresentation
Lar = LinearAlgebraicRepresentation
import Base.show
using ViewerGL,SparseArrays
GL = ViewerGL

function show(filename)
	V, EV = Lar.svg2lar(filename)
	GL.VIEW([ GL.GLLines(V,EV) ])
	return V, EV
end

show("./test/svg/new.svg")
show("./test/svg/curved.svg")
show("./test/svg/twopaths.svg")
show("./test/svg/paths.svg")
show("./test/svg/boundarytest2.svg")
show("./test/svg/tile.svg")
show("./test/svg/interior.svg")
show("./test/svg/holes.svg")
V,EV = show("./test/svg/Lar.svg")

# subdivision of input edges
W = convert(Lar.Points, V')
cop_EV = Lar.coboundary_0(EV::Lar.Cells)
cop_EW = convert(Lar.ChainOp, cop_EV)
V, copEV, copFE = Lar.planar_arrangement(W::Lar.Points, cop_EW::Lar.ChainOp)

# compute containment graph of components
bicon_comps = Lar.Arrangement.biconnected_components(copEV)

# visualization of component graphs
EW = Lar.cop2lar(copEV)
W = convert(Lar.Points, V')
V,EVs = Lar.biconnectedComponent((W,EW::Lar.Cells)) # 2-connected components
VV = [[k] for k=1:size(V,2)]
meshes = [ GL.numbering(.1)((V, (VV,EVs[k])), GL.COLORS[k],1) for k=1:length(EVs) ]
GL.VIEW( cat(meshes) );

# final solid visualizations
FE = [SparseArrays.findnz(copFE[k,:])[1] for k=1:size(copFE,1)]
FV = [collect(Set(cat([EW[e] for e in FE[f]]))) for f=1:length(FE)]
V = convert(Lar.Points, V')
triangulated_faces = Lar.triangulate2D(V, [copEV, copFE])
V = convert(Lar.Points, V')
FVs = convert(Array{Lar.Cells}, triangulated_faces)
GL.VIEW(GL.GLExplode(V,FVs,1,1,1,99,1));  # TODO:  fix inner holes

V,FVs,EVs = Lar.arrange2D(V,EV)
GL.VIEW(GL.GLExplode(V,EVs,1.2,1.2,1.2,99,1));
GL.VIEW(GL.GLExplode(V,FVs,1,1,1,99,1));  # TODO:  fix inner holes

EVs = Lar.FV2EVs(copEV, copFE) # polygonal face boundaries
EVs = convert(Array{Array{Array{Int64,1},1},1}, EVs)
GL.VIEW(GL.GLExplode(V,EVs,1,1,1,3,1));  # TODO:  fix inner holes


"""
    struct Parallel.MultipleDijkstraState{T, U}

An [`AbstractPathState`](@ref) designed for Parallel.dijkstra_shortest_paths calculation.
"""
struct MultipleDijkstraState{T <: Real,U <: Integer} <: AbstractPathState
    dists::Matrix{T}
    parents::Matrix{U}
end

"""
    Parallel.dijkstra_shortest_paths(g, sources=vertices(g), distmx=weights(g))

Compute the shortest paths between all pairs of vertices in graph `g` by running
[`dijkstra_shortest_paths`] for every vertex and using an optional list of source vertex `sources` and
an optional distance matrix `distmx`. Return a [`Parallel.MultipleDijkstraState`](@ref) with relevant
traversal information.
"""
function dijkstra_shortest_paths(g::AbstractGraph{U},
    sources::AbstractVector=vertices(g),
    distmx::AbstractMatrix{T}=weights(g)) where T <: Real where U

    n_v = nv(g)
    r_v = length(sources)

    # TODO: remove `Int` once julialang/#23029 / #23032 are resolved
    dists   = SharedMatrix{T}(Int(r_v), Int(n_v))
    parents = SharedMatrix{U}(Int(r_v), Int(n_v))

    @sync @distributed for i in 1:r_v
        state = LightGraphs.dijkstra_shortest_paths(g, sources[i], distmx)
        dists[i, :] = state.dists
        parents[i, :] = state.parents
    end

    result = MultipleDijkstraState(sdata(dists), sdata(parents))
    return result
end


module SeriesIterators

export iterm, icumsum
export bernoulli

using IterTools
using Base.Iterators

import Base: iterate, IteratorSize, length, size

"""
    iterm(g, ind)

Assume a sequence of terms a[k] can be generated by

    a[i0], b = g(i0)
    a[i], b =  g(i, b) for subsequent indices

The `iterm(g, ind)` iterates the sequence, when `ind` is the
iterator of the indices; first index is `i0`.
"""
iterm(gen, it) = TermIterator(gen, it)
struct TermIterator{F,I}
    generate::F
    indices::I
end
function iterate(t::TermIterator)
    vs = iterate(t.indices)
    k, s = @IterTools.ifsomething vs
    v = t.generate(k)
    v[1], (s, v)
end
function iterate(t::TermIterator, s)
    sx, p = s
    vs = iterate(t.indices, sx)
    k, s = @IterTools.ifsomething vs
    v = t.generate(k, p)
    v[1], (s, v)
end
IteratorSize(::Type{<:TermIterator{<:Any,I}}) where I = IteratorSize(I)
length(t::TermIterator) = length(t.indices)
size(t::TermIterator) = size(t.indices)

function gen_bernoulli(::Type{T}) where T<:Integer
    B = Vector{Rational{T}}(undef, 2)
    B[1] = 1
    B[2] = -1 // 2
    Z = zero(Rational{T})
    function _bernoulli(n::Integer)
        b = B
        n <= 1 && return b[n+1]
        isodd(n) && return Z
        l = length(b)
        if l * 2 - 2 <= n
            ext = min(1_000_000 ÷ l^2, 40)
            m = max(n ÷ 2 + 2, l + ext)
            n0 = (length(b) - 1 ) * 2
            sizehint!(b, ((m + 80) >> 6) << 6)
            resize!(b, m)
            _bernoulli!(b, n0, (m - 1) * 2 - 1, n)
        end
        b[n÷2 + 2]
    end
    _bernoulli
end

function _bernoulli!(B::Vector{Rational{T}}, n0::Integer, n::Integer, nn::Integer) where T
    lastx = 0
    try
        for m = n0:2:n
            sum = -B[1] // T(m + 1)
            sum -= B[2]
            mok = T((m * (m - 1)) ÷ 2)
            for k = 2:2:m-1
                sum -= B[k÷2 + 2] // T(m - k + 1) * mok
                mok *= T(m - k) // T(k + 1)
                mok *= T(m - k - 1) // T(k + 2)
            end
            lastx = m ÷ 2 + 2
            B[lastx] = sum
        end
    catch
        if lastx > 0
            resize!(B, lastx)
        end
        if lastx < nn ÷ 2 + 2
            rethrow()
        end
    end
    nothing
end
"""
    bernoulli(n)

Return Bernoulli number `B_n`. `B_1 = -1/2`.
"""
const bernoulli = gen_bernoulli(BigInt)

struct CumSum{T}
    itr::T
end

"""
    icumsum(itr)

Return iterator, which lazily delivers the same results as `cumsum(itr)`.
"""
icumsum(itr) = CumSum(itr)
function iterate(c::CumSum)
    v, s = @IterTools.ifsomething iterate(c.itr)
    v, (s, v)
end
function iterate(c::CumSum, (s, p))
    v, s = @IterTools.ifsomething iterate(c.itr, s)
    v = p + v
    v, (s, v)
end
IteratorSize(::Type{<:CumSum{T}}) where T = IteratorSize(T)
length(c::CumSum) = length(c.itr)
size(c::CumSum) = size(c.itr)

end


using LinearAlgebra
using Random
using PyPlot
using Interpolations
using SparseArrays
using Optim
using ProgressBars
using StatsBase
using AlphaStableDistributions
using AdvancedHMC
using Dates
using MAT
using MCMCDiagnostics

seed_shift = (@isdefined shift) ? shift : 0

function measurementmatrix(X,MX,kernel)
    s = X[2]-X[1];
    F = zeros(size(MX)[1],size(X)[1])
    for i = 1:size(MX)[1]
        F[i,:] = s.*kernel.(X,MX[i]);
    end
    return F
end

function smatrix(X0,a,b)
    N = length(X0);
    dx = X0[2]-X0[1];

    M = zeros(N,N);
    for i=2:N-1
        M[i,i-1] = -a/dx^2;
        M[i,i+1] = -a/dx^2;
        M[i,i] = 2*a/dx^2 +b;

    end
    #M[1,:] = [b+2*a/dx^2 -a/dx^2 zeros(1,N-3) -a/dx^2];
    #M[N,:] = [-a/dx^2 zeros(1,N-3) -a/dx^2 b+2*a/dx^2];

    M[1,1] = 1;
    #M(N,1) = 1/dx;
    #M(N,2) = -1/dx;

    M[N,N] = 1
    return M

end

function partialsum!(z, s0, sp, spc)
    N = size(spc)[1]
     for i = 1:N
        r = spc[i]
        z[r] = s0[r] + sp[i]
    end
    return nothing
end

function ravel(q)
    return dropdims(q, dims = tuple(findall(size(q) .== 1)...))
end

function difference1(X)
    N = length(X);
    M = zeros(N,N);

    for i  = 2:N
       M[i,i] = -1;
       M[i,i-1] = 1;
    end
    M[1,1] = 1;
    return M
end

function  difference2(X)
    N = length(X);
    M = zeros(N,N);

    for i  = 3:N
       M[i,i] = -1;
       M[i,i-1] = 2;
       M[i,i-2] = -1;
    end
    M[1,1] = 1; M[2,1] = 1; M[2,2] = -1;
    return M

end

function divide(N,NT)
    even = mod(N,NT) == 0
    k = Int64(floor(N/NT))
    t = 0
    if(even)
        q = Vector{Vector{Int64}}(undef,k)
        t = k

    else
        q = Vector{Vector{Int64}}(undef,k+1)
        t = k +1
    end

    for i = 1:t
        q[i] = Vector(NT*(i-1)+1:min(NT*i,N))
    end

    return q

end

@inbounds  function measurementmatrix2d(X,Y,Xt,Yt,kernel;constant=10.0)
    dy = abs(Y[2,1]-Y[1,1])
    dx = abs(X[1,2]-X[1,1])
    d = Float64(dy*dx)
    N = length(X);
    Nt = length(Xt)
    #F = zeros(Nt,N)

    minval = 0.0001*d*kernel(0,0,0,0,constant=constant)
    Nth = Threads.nthreads()
    rows = Vector{Vector{Int64}}(undef, Nth)
    cols = Vector{Vector{Int64}}(undef, Nth)
    vals = Vector{Vector{Float64}}(undef, Nth)
    for p = 1:Nth
        rows[p] = []
        cols[p] = []
        vals[p] = []
    end

    Threads.@threads  for i = 1:N
        for j = 1:Nt
            xi = div(i-1,N)+1;
            yi = rem(i,N)+1
            xj = div(j-1,Nt)+1;
            yj = rem(j,Nt)+1
            val = d*kernel(X[i],Xt[j],Y[i],Yt[j],constant=constant);
            if (val >= minval)
                #F[j,i] = val
                push!(rows[Threads.threadid()],j)
                push!(cols[Threads.threadid()],i)
                push!(vals[Threads.threadid()],val)
            end

            #F[i,j] = d*kernel(X[xi],X[xj],Y[yi],Y[yj],constant=constant);
            #F[j,i] = F[i,j]
        end
    end
    rows = vcat(rows...)
    cols = vcat(cols...)
    vals = vcat(vals...)
    F = sparse(rows,cols,vals,Nt,N)

    return F

end

function  s2dmatrix(x0,a,b)
    N = length(x0);
    dx = x0[2]- x0[1]

   M = zeros(N,N);
   for i=2:N-1
       M[i,i-1] = -a/dx^2;
       M[i,i+1] = -a/dx^2;
       M[i,i] = 2*a/dx^2;

   end
   M[1,:] = [2*a/dx^2 -a/dx^2 zeros(1,N-3) -a/dx^2];
   M[N,:] = [-a/dx^2 zeros(1,N-3) -a/dx^2 2*a/dx^2];

   M = kron(M,I(N)) + kron(I(N),M)
   M = M + I(size(M)[2])*b;

   return M
end

    function linspace(start,stop,length)
        return Vector(range(start,stop=stop,length=length))
    end

    function partialreplace(orig,replace,ix)
        a = copy(orig)
        a[ix] = replace
        return a
    end

   @inline function meshgrid(x,y)
        grid_a = [i for i in x, j in y]
        grid_b = [j for i in x, j in y]

        return grid_a,grid_b
    end

@inbounds @inline function storevector!(target,source)
    N = length(target)
    @assert N == length(source)
    @simd for i = 1:N
        target[i] = source[i]
    end
end

@inbounds @inline function storesubst!(target,a1,a2)
    N = length(target)
    @assert N == length(a1) == length(a2)
    @simd for i = 1:N
        target[i] = a1[i] - a2[i]
    end
end

@inbounds @inline function storeadd!(target,a1,a2)
    N = length(target)
    @assert N == length(a1) == length(a2)
    @simd for i = 1:N
        target[i] = a1[i] + a2[i]
    end
end

@inbounds @inline function slicemadd!(target,cols,vec,ix)
    N = length(ix)
    @assert N == length(vec)
    #mul!(target,cols[ix[1]],vec[1])
    for i = 1:N
        mul!(target,cols[ix[i]],vec[i],1,1)
    end
end

function  logpdiff_prior(x,args,cache)
    D = args.D;
    noisesigma = args.noisesigma
    scale = args.scale; #y  = args.y; F = args.F

    #Fxprop = cache.Fxprop
    Dprop = cache.Dprop
    #res = cache.residual

    #mul!(Fxprop,F,x)
    mul!(Dprop,D,x)

    #res .= Fxprop-y
    #storesubst!(res,Fxprop,y)
    #logp = -0.5/noisesigma^2*dotitself(res);
    logp = sum(log.(scale./((scale^2 .+ Dprop.^2))));

    return logp

end

function  logpdiffpartial(ixl,nth,x,args,cache,current)
    noisesigma = args.noisesigma
    ix = ixl[nth]
    scale = args.scale; y  = args.y; bscale = args.bscale
    # F =   args.F[:,ix]
    # Dx =  args.Dx[:,ix]; Dy =  args.Dy[:,ix]
    # Db =  args.Db[:,ix]
    F =   args.Fs[nth] ; D =   args.Ds[nth]
    Ds = args.Ds; Dnz = args.Dnz;
    Fs = args.Fs; Fnz = args.Fnz

    res = cache.residual
    Dprop = cache.Dprop
    Fxprop = cache.Fxprop
    xprop = cache.xprop

    currx = @view current.xcurr[ix]
    Fxcurr = current.Fxcurr
    Dcurr =  current.Dcurr

    xprop[ix] = x

    # Fnzall = union(vcat(Fnz[ix]...))
    Fnzall = Fnz[nth]
    # Dxnzall = union(vcat(Dxnz[ix]...))
    # Dynzall = union(vcat(Dynz[ix]...))
    Dnzall = Dnz[nth]
    # Dbnzall = union(vcat(Dbnz[ix]...))

    compx = x-currx
    #Fxprop .= Fxcurr
    #Fxprop .+= (F*compx)
    storevector!(view(Fxprop,Fnzall),view(Fxcurr,Fnzall))
    #storevector!(Fxprop,Fxcurr)
    mul!(Fxprop,F,compx,1,1)
    #slicemadd!(Fxprop,Fs,compx,ix)

    #Dxprop .= Dxcurr
    #Dxprop .+= Dx*compx
    storevector!(view(Dprop,Dnzall),view(Dcurr,Dnzall))
    #storevector!(Dprop,Dcurr)
    mul!(Dprop,D,compx,1,1)

    #res .= Fxprop
    storesubst!(res,Fxprop,y)
    #storevector!(res,Fxprop)
    #res .-=  y

    resv = @view res[Fnzall]
    Dv = @view Dprop[Dnzall]

    #return   -0.5/noisesigma^2*res'*res -3/2*sum(log.(scale^2 .+ Dxprop.^2 .+ Dyprop.^2)) - sum(log.(bscale^2 .+ Dbprop.^2))
    return   -0.5/noisesigma^2*dotitself(resv) + sum(log.(scale./((scale^2 .+ Dv.^2))));
end

function  logpdiffgradi_prior(x,arguments,cache;both=true)
    D = arguments.D;
    noisesigma = arguments.noisesigma
    scale = arguments.scale; #y  = arguments.y; F = arguments.F

    logp = 0.0

    #Fxprop = cache.Fxprop
    Dprop = cache.Dprop
    G = cache.gradiprop
    #res = cache.residual

    #mul!(Fxprop,F,x)
    mul!(Dprop,D,x)
    #storesubst!(res,Fxprop,y)

    den = (scale^2 .+ Dprop.^2)
    if both
        #logp = -0.5/noisesigma^2*dotitself(res);
        logp =  sum(log.(scale./(den)));
    end

     #mul!(G,F',(-((res)/noisesigma.^2)))
     #G= F'*(-((res)./noisesigma.^2))

    #Gd =   M'*(-2.0*Dprop./den);
    mul!(G,D',(-2.0*Dprop./den))
    #G = G + Gd;

    return logp,G

end

function  logpisodiff_prior(x,args,cache)
    noisesigma = args.noisesigma
    scale = args.scale; #y  = args.y;# F = args.F
    bscale = args.bscale
    Dx = args.Dx; Dy = args.Dy
    Db = args.Db

    #res = cache.residual
    Dxprop = cache.Dxprop
    Dyprop = cache.Dyprop
    Dbprop = cache.Dbprop
    #Fxprop = cache.Fxprop

    #mul!(Fxprop,F,x)
    #res .= Fxprop - y
    #storesubst!(res,Fxprop,y)
    #Dxprop .= Dx*x
    mul!(Dxprop,Dx,x)
    #Dyprop .= Dy*x
    mul!(Dyprop,Dy,x)
    #Dbprop .= Db*x
    mul!(Dbprop,Db,x)

    return   -sum(log.(bscale^2 .+ Dbprop.^2)) -3/2*sum(log.(scale^2 .+ Dxprop.^2 + Dyprop.^2))

end

function  logpisodiffpartial(ixl,nth,x,args,cache,current)
    noisesigma = args.noisesigma
    ix = ixl[nth]
    scale = args.scale; y  = args.y; bscale = args.bscale
    # F =   args.F[:,ix]
    # Dx =  args.Dx[:,ix]; Dy =  args.Dy[:,ix]
    # Db =  args.Db[:,ix]
    F =   args.Fs[nth]
    Dx =  args.Dxs[nth]; Dy =  args.Dys[nth]
    Db =  args.Dbs[nth]
    Dxs = args.Dxs; Dxnz = args.Dxnz;
    Dys = args.Dys; Dynz = args.Dynz
    Dbs = args.Dbs; Dbnz = args.Dbnz
    Dinz = args.Dinz
    Fs = args.Fs; Fnz = args.Fnz

    #res = cache.residual
    Dxprop = cache.Dxprop
    Dyprop = cache.Dyprop
    Dbprop = cache.Dbprop
    Fxprop = cache.Fxprop
    xprop = cache.xprop

    currx = @view current.xcurr[ix]
    Fxcurr = current.Fxcurr
    Dxcurr =  current.Dxcurr
    Dycurr = current.Dycurr
    Dbcurr = current.Dbcurr

    xprop[ix] = x

    # Fnzall = union(vcat(Fnz[ix]...))
    Fnzall = Fnz[nth]
    # Dxnzall = union(vcat(Dxnz[ix]...))
    Dxnzall = Dxnz[nth]
    # Dynzall = union(vcat(Dynz[ix]...))
    Dynzall = Dynz[nth]
    # Dbnzall = union(vcat(Dbnz[ix]...))
    Dbnzall = Dbnz[nth]
    Diall = Dinz[nth]
    #Diall = union(Dxnzall,Dynzall)

    compx = x-currx
    #Fxprop .= Fxcurr
    #Fxprop .+= (F*compx)
    storevector!(view(Fxprop,Fnzall),view(Fxcurr,Fnzall))
    #storevector!(Fxprop,Fxcurr)
    mul!(Fxprop,F,compx,1,1)
    #slicemadd!(Fxprop,Fs,compx,ix)

    #Dxprop .= Dxcurr
    #Dxprop .+= Dx*compx
    storevector!(view(Dxprop,Dxnzall),view(Dxcurr,Dxnzall))
    #storevector!(Dxprop,Dxcurr)
    mul!(Dxprop,Dx,compx,1,1)

    #Dyprop .= Dycurr
    #Dyprop .+= Dy*compx
    storevector!(view(Dyprop,Dynzall),view(Dycurr,Dynzall))
    #storevector!(Dyprop,Dycurr)
    mul!(Dyprop,Dy,compx,1,1)

    #Dbprop .= Dbcurr
    #Dbprop .+= Db*compx
    storevector!(view(Dbprop,Dbnzall),view(Dbcurr,Dbnzall))
    #storevector!(Dbprop,Dbcurr)
    mul!(Dbprop,Db,compx,1,1)

    #res .= Fxprop
    storesubst!(res,Fxprop,y)
    #storevector!(res,Fxprop)
    #res .-=  y

    resv = @view res[Fnzall]
    Dxv = @view Dxprop[Diall]
    Dyv = @view Dyprop[Diall]
    Dbv = @view Dbprop[Dbnzall]

    #return   -0.5/noisesigma^2*res'*res -3/2*sum(log.(scale^2 .+ Dxprop.^2 .+ Dyprop.^2)) - sum(log.(bscale^2 .+ Dbprop.^2))
    return   -0.5/noisesigma^2*dotitself(resv) - sum(log.(bscale^2 .+ Dbv.^2)) -3/2*sum(log.(scale^2 .+ Dxv.^2 + Dyv.^2))

end

function logpisodiffgradi_prior(x,args,cache;both=true)
    noisesigma = args.noisesigma
    scale = args.scale; #y  = args.y; #F = args.F
    bscale = args.bscale
    Dx = args.Dx; Dy = args.Dy
    Db = args.Db

    logp = 0.0

    #res = cache.residual
    Dxprop = cache.Dxprop
    Dyprop = cache.Dyprop
    Dbprop = cache.Dbprop
    #Fxprop = cache.Fxprop
    G = cache.gradiprop

    #Fxprop .= F*x
    #mul!(Fxprop,F,x)
    #storesubst!(res,Fxprop,y)
    #res .= Fxprop - y
    #Dxprop .= Dx*x
    mul!(Dxprop,Dx,x)
    #Dyprop .= Dy*x
    mul!(Dyprop,Dy,x)
    #Dbprop .= Db*x
    mul!(Dbprop,Db,x)

    #G .= F'*(-((res)./noisesigma.^2))
    #mul!(G,F',-((res)/noisesigma.^2))

    den = (scale^2 .+ Dyprop.^2 + Dxprop.^2)

    if both
        logp =  -3/2*sum(log.(den)) - sum(log.(bscale^2 .+ Dbprop.^2))
    end

    #Gd1 =  Dx'*(-3.0*Lxx./(scale^2 .+ Lxx.^2));
    #G .= G  -3.0*(Dx)'*(Dxprop./(scale^2 .+ Dyprop.^2 + Dxprop.^2))  - 3.0*(Dy)'*(Dyprop./(scale^2 .+ Dyprop.^2 + Dxprop.^2))   -Db'*(2.0*Dbprop./(bscale^2 .+ Dbprop.^2));

    mul!(G,-3.0*(Dx)',(Dxprop./den))
    mul!(G,-3.0*(Dy)',(Dyprop./den),1,1)
    mul!(G,-Db',(2.0*Dbprop./(bscale^2 .+ Dbprop.^2)),1,1)

    return logp, G
end

function regmatrices_first(dim)
    reg1d = spdiagm(Pair(0,-1*ones(dim))) + spdiagm(Pair(1,ones(dim-1))) + spdiagm(Pair(-dim+1,ones(1))) ;reg1d[dim,dim] = 0
    #reg1d = reg1d[1:dim-1,:]
    iden = I(dim)
    regx = kron(reg1d,iden)
    regy = kron(iden,reg1d)

    rmxix = sum(abs.(regx) ,dims=2) .< 2
    rmyix = sum(abs.(regy) ,dims=2) .< 2
    boundary = ((rmxix + rmyix)[:]) .!= 0
    q = findall(boundary .== 1)
    regx = regx[setdiff(1:dim^2,q), :]
    regy = regy[setdiff(1:dim^2,q), :]

    s = length(q)
    bmatrix = sparse(zeros(s,dim*dim))
    for i=1:s
        v = q[i]
        bmatrix[i,v] = 1
    end
    #bmatrix = bmatrix[i,i] .= 1

    return regx,regy,bmatrix
end

function regmatrices_second(dim)
    reg1d = spdiagm(Pair(0,2*ones(dim))) + spdiagm(Pair(1,-1*ones(dim-1))) + spdiagm(Pair(-1,-1*ones(dim-1))) #;reg1d[dim,dim] = 0
    #reg1d = reg1d[1:dim-1,:]
    iden = I(dim)
    regx = kron(reg1d,iden)
    regy = kron(iden,reg1d)

    rmxix = sum(abs.(regx), dims=2) .< 4
    rmyix = sum(abs.(regy), dims=2) .< 4
    boundary = ((rmxix + rmyix)[:]) .!= 0
    q = findall(boundary .== 1)
    regx = regx[setdiff(1:dim^2,q), :]
    regy = regy[setdiff(1:dim^2,q), :]

    #q = findall(boundary .== 1)
    #bmatrix = zeros(dim*dim,dim*dim)
    #for i in q
    #     bmatrix[i,i] = 1
    #end

    h2 = sparse(zeros(2,N))
    h2[1,1] = 1; h2[2,N] = 1
    h1 = hcat([sparse(zeros(dim-2)),I(dim-2),sparse(zeros(dim-2))]...)
    bmatrix = [ kron(h2,I(dim)); kron(h1,h2) ]

    return regx,regy,bmatrix
end

function spdematrix(xs,a,b)
    dx = abs(xs[2]-xs[1])
    N = length(xs)

    M = sparse(zeros(N,N))
    for i=2:N-1
        M[i,i-1] = -a/dx^2;
        M[i,i+1] = -a/dx^2;
        M[i,i] = 2*a/dx^2;# +b;

    end

    M[1,1] = 1 # Option 1
    M[N,N] = 1
    M[N,N-1] = -1

    #M[1,:] = [2*a/dx^2;-a/dx^2; zeros(N-3);-a/dx^2]; # Option 2
    #M[N,:] = [-a/dx^2; zeros(N-3); -a/dx^2;2*a/dx^2];

    M = kron(M,I(N)) + kron(I(N),M)
    M = M + I(size(M,2))*b;

    return M
end

function nzrows(M::SparseVector)
    return unique(M.nzind)
end

function nzrows(M::SparseMatrixCSC)
    return unique(M.rowval)
end

function splitmatrix(M::SparseMatrixCSC,ixlist::Array{Array{Int64,1},1})
    N = length(ixlist)
    zl = Vector{Vector{Int64}}(undef,N)
    Ml = Vector{AbstractSparseArray}(undef,N)
    for i =1:N
        Mi = M[:,ixlist[i]]
        nz = nzrows(Mi)
        zl[i] = nz
        Ml[i] = Mi
        #push!(zl,nz)
        #push!(Ml,Mi)
    end
    return Ml,zl
end

function splitmatrix(M::SparseMatrixCSC)
    N = size(M)[2]
    zl = Vector{Vector{Int64}}(undef,N)
    Ml = Vector{AbstractSparseArray}(undef,N)
    for i =1:N
        Mi = M[:,i]
        nz = nzrows(Mi)
        zl[i] = nz
        Ml[i] = Mi
        #push!(zl,nz)
        #push!(Ml,Mi)
    end
    return Ml,zl
end

mutable struct rram
    n::Int64
    acc::Int64
    Nadapt::Int64
    C::Array{Float64,2}
    Cho::Cholesky{Float64,Array{Float64,2}}
    xm::Vector{Float64}
end

function refreshcurrent!(cache,current;mode=1,arg=nothing,ixl=nothing,nth=nothing)
    if isnothing(arg)
        if (mode ==1)
            current.Mxcurr .= cache.Mxprop
            current.noisecurr .= cache.noiseprop
        elseif (mode==2)
            # current.xcurr .= cache.xprop
            # current.Fxcurr .= cache.Fxprop
            # current.Dcurr .= cache.Dprop
            storevector!(current.xcurr, cache.xprop)
            storevector!(current.Fxcurr, cache.Fxprop)
            storevector!(current.Dcurr, cache.Dprop)

        elseif (mode==3)
            # current.xcurr .= cache.xprop
            # current.Fxcurr .= cache.Fxprop
            # current.Dxcurr .= cache.Dxprop
            # current.Dycurr .= cache.Dyprop
            # current.Dbcurr .= cache.Dbprop
            storevector!(current.xcurr, cache.xprop)
            storevector!(current.Fxcurr, cache.Fxprop)
            storevector!(current.Dxcurr, cache.Dxprop)
            storevector!(current.Dycurr, cache.Dyprop)
            storevector!(current.Dbcurr, cache.Dbprop)
        else
            error("")
        end

    else
        if (mode ==1)
            current.Mxcurr .= cache.Mxprop
            current.noisecurr .= cache.noiseprop
        elseif (mode==2)
            storevector!(view(current.xcurr,ixl[nth]), view(cache.xprop,ixl[nth]))
            storevector!(view(current.Fxcurr,arg.Fnz[nth]), view(cache.Fxprop,arg.Fnz[nth]))
            storevector!(view(current.Dcurr,arg.Dnz[nth]), view(cache.Dprop,arg.Dnz[nth]))
            #current.xcurr .= cache.xprop
            #current.Fxcurr .= cache.Fxprop
            #current.Dcurr .= cache.Dprop

        elseif (mode==3)
            # current.xcurr .= cache.xprop
            # current.Fxcurr .= cache.Fxprop
            # current.Dxcurr .= cache.Dxprop
            # current.Dycurr .= cache.Dyprop
            # current.Dbcurr .= cache.Dbprop
            storevector!(view(current.xcurr,ixl[nth]), view(cache.xprop,ixl[nth]))
            storevector!(view(current.Fxcurr,arg.Fnz[nth]), view(cache.Fxprop,arg.Fnz[nth]))
            storevector!(view(current.Dxcurr,arg.Dxnz[nth]), view(cache.Dxprop,arg.Dxnz[nth]))
            storevector!(view(current.Dycurr,arg.Dynz[nth]), view(cache.Dyprop,arg.Dynz[nth]))
            storevector!(view(current.Dbcurr,arg.Dbnz[nth]), view(cache.Dbprop,arg.Dbnz[nth]))

        else
            error("")
        end
    end
end

function mwgpartial(lpdf,gradilpdf,ixl,nth,x00,args,cache,current,tuning;Nruns=5,Np=3,T=5.0,mode=1,extra=nothing)
    ix = ixl[nth]
    x0 = repeat(x00[ix],Np)
    x = x0# copy(x0)
    Nall = length(x)

    xnew = similar(x)

    logpdf_can(xp) = lpdf(ixl,nth,xp,args,cache,current)
    #logpdf_can(xp) = lpdf(partialreplace(x00,xp,ix),args,cache)

    logpdf(x) = logpdf_can(x)

    ll = logpdf(x)

    for _ = 1:Nruns

        tuning.n = tuning.n + 1

        V=randn(Nall)
        xnew .= x+ tuning.Cho.L*V

        lnew = logpdf(xnew)

        U = rand()
        apr = min(1.0,exp(-ll + lnew))

        if (U < apr)
            ll = lnew
            x .= xnew
            tuning.acc = tuning.acc +1
        end

        if (tuning.n <= tuning.Nadapt)
            xmp = copy(tuning.xm)
            tuning.xm .= tuning.xm + (x-tuning.xm)/tuning.n
            xm = tuning.xm
            tuning.C .= ((tuning.n-1)*tuning.C + (x-xm)*(x-xmp)')/tuning.n
            tuning.Cho = cholesky(Hermitian(2.38^2/Nall*tuning.C + 1e-11*I(Nall)))

            # atarg = 0.3
            # V = V/norm(V)
            # z = sqrt(tuning.n^(-0.7) *abs(tuning.acc/tuning.n - atarg)) *  tuning.Cho.L*V
            # #println(z)
            # if (tuning.acc/tuning.n >= atarg)
            #     lowrankupdate!(tuning.Cho, z)

            # else
            #     lowrankdowndate!(tuning.Cho, z)
            # end

        end

    end

    xret =  x#reweight(x,Np,logpdf_can,T)
    logpdf_can(xret)
    refreshcurrent!(cache,current;mode=mode,arg=args,ixl=ixl,nth=nth)

    return xret

end

function repatpartial(lpdf,placeholder,ixl,nth,x0,args,cache,current,tuning;mode=1,Np=1,Nruns=1,T=-Inf,extra=copy(x0))::Vector{Float64}
    ix = ixl[nth]
    x = copy(x0[ix])
    dim = length(x)

    logpdf(xp) = lpdf(ixl,nth,xp,args,cache,current)
    #logpdf(xp) = lpdf(partialreplace(x0,xp,ixl[nth]),args,cache)
    density = logpdf(x)

    z = copy(extra[ix])
    u = randn(dim)
    densityz = logpdf(z)

    for i = 1:Nruns
        R = tuning.Cho.L/sqrt(2)
		u .= randn(dim)
		xtilde = x + R * u
		densitytilde = logpdf(xtilde)
		while (log(rand()) >=  density - densitytilde)
			u .= randn(dim)
			xtilde .= x + R * u
			densitytilde = logpdf(xtilde)
		end

		u .= randn(dim)
		xstar = xtilde + R * u
		densitystar = logpdf(xstar)

		while (log(rand()) >=  densitystar - densitytilde)
			u .= randn(dim)
			xstar .= x + R * u
			densitystar = logpdf(xstar)
		end

		u .= randn(dim)
		zstar = xstar + R * u
		densityzstar = logpdf(zstar)

		while (log(rand()) >=  densitystar - densityzstar)
			u .= randn(dim)
			zstar .= xstar + R * u
			densityzstar = logpdf(zstar)
		end

        ratio = densitystar + min(0,density-densityz) - density - min(0,densitystar-densityzstar)
        if (log(rand()) <= ratio)
            x .= xstar
            z .= zstar
			extra[ix] = zstar
			density = densitystar
			densityz = densityzstar

        end
        tuning.n = tuning.n + 1

        if (tuning.n <= tuning.Nadapt)
            xmp = copy(tuning.xm)
            tuning.xm .= tuning.xm + (x-tuning.xm)/tuning.n
            xm = tuning.xm
            tuning.C .= ((tuning.n-1)*tuning.C + (x-xm)*(x-xmp)')/tuning.n
            tuning.Cho = cholesky(Hermitian(2.38^2/dim*tuning.C + 1e-11*I(dim)))
        end

    end
    logpdf(x)
    refreshcurrent!(cache,current;mode=mode,arg=args,nth=nth,ixl=ixl)
    return x
end

function totalsample(lpdf,glpdf,x0,meas,arg,ixl;N=1,Nadapt=1000,Nruns=5,thinning=10,algo=repatpartial,mode=2,Np=2,T=2.0)
    x = copy(x0)
    Npar = length(x0)
    chain = zeros(Npar, Int(floor(N / thinning)))
    Nix = length(ixl)

    if (mode==2)
        cache=(xprop=copy(x0),Fxprop=(arg.F*x0),Dprop=(arg.D*x0), residual=similar(arg.y),gradiprop=similar(x0))
        current = (Fxcurr=(arg.F*x0),Dcurr=(arg.D*x0), xcurr = copy(x0))
    elseif (mode==3)
        cache =(xprop=copy(x0),Fxprop=(arg.F*x0),Dxprop=(arg.Dx*x0),Dyprop=(arg.Dy*x0), Dbprop=(arg.Db*x0),  residual=arg.F*x0-arg.y,gradiprop=similar(x0))
        current = (xcurr = copy(x0), Dxcurr = arg.Dx*x0, Dycurr = arg.Dy*x0, Dbcurr = arg.Db*x0, Fxcurr = F*x0)
    else
        error("")
    end

    if algo == repatpartial
        extra = copy(x0)
        tuning = Vector{rram}(undef,Nix)
        for k = 1:Nix
            n = length(ixl[k])
            C = 0.1*Array(I(n))
            Cho = cholesky(C)
            xm = copy(x0[ixl[k]])
            tuning[k] = rram(1,0,Nadapt,C,Cho,xm)
        end

    elseif (algo==mwgpartial)
        extra = []
        tuning = Vector{rram}(undef,Nix)
        #tuning = Vector{dahmc2}(undef,Nix)
        #tuning = Vector{ahmc}(undef,Nix)
        for k = 1:Nix

            n = length(ixl[k])*Np
            C = 0.1*Array(I(n))
            Cho = cholesky(C)
            xm = repeat(copy(x0[ixl[k]]),Np)
            tuning[k] = rram(1,0,Nadapt,C,Cho,xm)

        end
    end

    pb = ProgressBar(1:N)

    for i in pb
        for k = 1:Nix
            x[ixl[k]] = algo(lpdf,glpdf,ixl,k,x,arg,cache,current,tuning[k];mode=mode,Nruns=Nruns,extra=extra,Np=Np,T=T)#lpdf,ix,x0,args,cache,current,tuning
            if (i % thinning == 0)
                chain[ixl[k], Int(i / thinning)] = x[ixl[k]]
            end

        end
    end

    return chain,tuning
end

function savechain(c,prefix)
    p = pwd()
    f = p*"/chains/"*prefix*"_"*Dates.format(Dates.now(), "yyyy_mm_dd_HH_MM_SS")*".mat"
    matwrite(f,Dict("chain"=>c))

end

cw = 150.0
kernel(xi,xj,yi,yj;constant=cw) = constant/pi*exp(-constant*((xi-xj)^2 + (yi-yj)^2) )
tf(x,y) =  1.0*exp.(-20*sqrt.((x .- 0.3).^2 .+ (y .- 0.3).^2)) + 1*((y-x) .< 0.7).*((y-x) .>= -0.7).*((-y-x) .<= 0.8).*((-y-x) .>= 0.4).*(((-x+y) .+ 0.7)/1.39)  + 1*(-x+y .< 1).*(-x+y .>= 0.8).*(abs.(x) .<= 1).*(abs.(y) .<= 1) + ( 4*0.25 .- 4*sqrt.((x.-0.5).^2+(y.+0.6).^2)).*(sqrt.((x.-0.5).^2+(y.+0.6).^2) .<= 0.25);

Random.seed!(1)

noisevar = 0.01^2

Nbig = 300
dimbig = Nbig
N = 180
dim = N
Nmeas = 100
dimmeas = Nmeas

xsbig = -1+1/dimbig:2/dimbig:1-1/(dimbig)
ysbig = -1+1/dimbig:2/dimbig:1-1/(dimbig)

xs = -1+1/dim:2/dim:1-1/(dim)
ys = -1+1/dim:2/dim:1-1/(dim)

xibig = linspace(-1,1,dimbig)
Ybig,Xbig = meshgrid(xibig,xibig)

ximeas = linspace(-1,1,dimmeas)
Ymeas,Xmeas = meshgrid(ximeas,ximeas)
Y,X = meshgrid(xs,xs)

Zbig = tf(Xbig,Ybig)
@time Fbig = measurementmatrix2d(Xbig,Ybig,Xmeas,Ymeas,kernel,constant=cw); # Theory matrix for measurements.
@time F = measurementmatrix2d(X,Y,Xmeas,Ymeas,kernel,constant=cw); # Theory matrix for inversion.

zx = Fbig*Zbig[:]
mebig = reshape(zx,(dimmeas,dimmeas))
y = mebig[:] + randn(Nmeas*Nmeas,)*sqrt(noisevar)

xp = randn(N*N)
ixl = divide(N*N,1)

Dx,Dy,Db = regmatrices_first(N);
Dxs, Dxnz = splitmatrix(Dx,ixl); Dys, Dynz = splitmatrix(Dy,ixl); Dbs, Dbnz = splitmatrix(Db,ixl); #Fs, Fnz = splitmatrix(F,ixl)
S = spdematrix(xs,0.001,0.1)
D = [Dx;Dy;Db]
Ds, Dnz = splitmatrix(D,ixl)
Dinz = [union(Dxnz[i],Dynz[i]) for i = 1:N^2]
argi  = (Dinz = Dinz, Dxs = Dxs, Dxnz = Dxnz , Dys = Dys, Dynz = Dynz, Dbs = Dbs, Dbnz = Dbnz, Dx = Dx, Dy = Dy, D=D,noisesigma = sqrt(noisevar), scale = 0.03, bscale = 0.03, Db = Db, Ds = Ds, Dnz = Dnz )
argis  = ( D=S, F = F,noisesigma = sqrt(noisevar), scale = 0.3, y = y, bscale = 1.0 )

cacheiso =(xprop=zeros(N*N),Dxprop=zeros(size(Dx)[1]),Dyprop=zeros(size(Dy)[1]),Dbprop=zeros(size(Db)[1]),gradiprop=zeros(N*N))
cached =(xprop=zeros(N*N),Dprop=zeros(size(D)[1]),gradiprop=zeros(N*N))
caches =(xprop=zeros(N*N),Fxprop=zeros(Nmeas^2),Dprop=zeros(size(S)[1]), residual=similar(y),gradiprop=zeros(N*N))

currentiso = (xcurr = copy(xp), Dxcurr = Dx*xp, Dycurr = Dy*xp, Dbcurr = Db*xp, Fxcurr = F*xp)
currentaniso = (xcurr = copy(xp), Dcurr = D*xp, Fxcurr = F*xp)

target1iso(x) = -logpisodiff(x,argi,cacheiso)# -logpdiff(w,diff1arg,cached)##
target1isograd(x) = -logpisodiffgradi(x,argi,cacheiso;both=false)[2]#  -logpdiffgradi(w,diff1arg,cached;both=false)[2]#
res = Optim.optimize(target1iso, target1isograd, 0*randn(N*N,),Optim.LBFGS(), Optim.Options(allow_f_increases=true,show_trace=true,iterations=700); inplace = false)
MAP_diff1iso = res.minimizer
imshow(reshape(MAP_diff1iso,(N,N))); colorbar(); clim(-0.1,1.1)

target1(x) = -logpdiff(x,argi,cached)# -logpdiff(w,diff1arg,cached)##
target1grad(x) = -logpdiffgradi(x,argi,cached;both=false)[2]#  -logpdiffgradi(w,diff1arg,cached;both=false)[2]#
res = Optim.optimize(target1, target1grad, 0*randn(N*N,),Optim.LBFGS(), Optim.Options(allow_f_increases=true,show_trace=true,iterations=700); inplace = false)
MAP_diff1 = res.minimizer
imshow(reshape(MAP_diff1,(N,N))); colorbar(); clim(-0.1,1.2)

targets(x) = -logpdiff(x,argis,caches)# -logpdiff(w,diff1arg,cached)##
targetsgrad(x) = -logpdiffgradi(x,argis,caches;both=false)[2]#  -logpdiffgradi(w,diff1arg,cached;both=false)[2]#
res = Optim.optimize(targets, targetsgrad, 0*randn(N*N,),Optim.LBFGS(), Optim.Options(allow_f_increases=true,show_trace=true,iterations=700); inplace = false)
MAP_s = res.minimizer
imshow(reshape(MAP_s,(N,N))); colorbar(); clim(-0.5,12)

function NUTS_sample(Niter,arg,cac)
    n_samples, n_adapts = 2000, 1000
    ss = Vector{Any}(undef,Niter)
    st = Vector{Any}(undef,Niter)
    Np = 1
    T = 1
    for i = (1:Niter) .+ seed_shift

        Ntot = N^2
        Random.seed!(i)
        println(i)
        initial_θ = copy(MAP_diff1)
        global c

        targetdiff(w) = -logpdiff(x,arg,cac)
        targetdiffgrad(w) = logpdiffgradi(w,arg,cac;both=true)

        metric = DiagEuclideanMetric(Ntot)# DenseEuclideanMetric(Ntot)# DiagEuclideanMetric(Ntot)
        hamiltonian = Hamiltonian(metric, targetdiff, targetdiffgrad)

        initial_ϵ = find_good_stepsize(hamiltonian, initial_θ)
        integrator =  Leapfrog(initial_ϵ)#  TemperedLeapfrog(initial_ϵ, 1.05)# Leapfrog(initial_ϵ)
        proposal = NUTS{MultinomialTS, GeneralisedNoUTurn}(integrator,8,1000.0)
        adaptor = StanHMCAdaptor(MassMatrixAdaptor(metric), StepSizeAdaptor(0.8, integrator))
        c, stats = sample(hamiltonian, proposal, initial_θ, n_samples, adaptor, n_adapts; progress=true)
        #ss[i] = hcat(c...)
        c = hcat(c...)
        # ss[i] = c
        #st[i] = stats
        pr="NUTS_"*string("anisodiff1d_prior")
        savechain(c,pr)

    end
    return ss, st
end
ch,b = NUTS_sample(1,argi,cached)

algo =  repatpartial#mwgpartial#
grfunk = nothing#
funk =  logpdiffpartial
par = argi
mode = 2# 1
#x000  = randn(N*N,)

for t=1:0
    Random.seed!(t)
    xz = copy(MAP_diff1)

    Nsamples = 100; Nadapt = 10; thinning = 10
    Np = 1; Nruns = 1; T = 1.0

    c,tun=totalsample(funk,grfunk,xz,y,par,ixl;N=Nsamples,Nadapt=Nadapt,thinning=thinning,algo=algo,mode=mode,Np=Np,T=T,Nruns=Nruns)
    savechain(c,"anisodiff1d_normal_"*string(nameof(algo))*"_"*string(nameof(funk))*"_N"*string(Nsamples)*"_Nadapt"*string(Nadapt)*"_thin"*string(thinning)*"_Np"*string(Np)*"_Nruns"*string(Nruns)*"_T"*string(T))
    println(effective_sample_size(c[1,:]))
    global c
    global tun
end


using NPZ
using Distributions
using LinearAlgebra
using Distributed
using StatsBase
using Distributions
using SpecialFunctions
using LinearAlgebra
using Random

function generate_sph_gaussian_data(N::Int64, D::Int64, K::Int64)
	x = randn(D,N)
	tpi = rand(Dirichlet(ones(K)))
	tzn = rand(Multinomial(N,tpi))
	tz = zeros(N)

	tmean = zeros(D,K)
	tcov = zeros(D,D,K)

	ind = 1
	println(tzn)
	for i=1:length(tzn)
		indices = ind:ind+tzn[i]-1
		tz[indices] .= i
		tmean[:,i] .= rand(MvNormal(zeros(D), 100*Matrix{Float64}(I, D, D)))
		tcov[:,:,i] .= rand(InverseGamma((D+2)/2,1))*Matrix{Float64}(I, D, D)
		# T = chol(slice(tcov,:,:,i))
		# x[:,indices] = broadcast(+, T*x[:,indices], tmean[:,i]);
		d = MvNormal(tmean[:,i], tcov[:,:,i])
		for j=indices
			x[:,j] = rand(d)
		end
		ind += tzn[i]
	end
	x, tz, tmean, tcov
end

function generate_gaussian_data(N::Int64, D::Int64, K::Int64)
	x = randn(D,N)
	tpi = rand(Dirichlet(ones(K)))
	tzn = rand(Multinomial(N,tpi))
	tz = zeros(N)

	tmean = zeros(D,K)
	tcov = zeros(D,D,K)

	ind = 1
	println(tzn)
	for i=1:length(tzn)
		indices = ind:ind+tzn[i]-1
		tz[indices] .= i
		tmean[:,i] .= rand(MvNormal(zeros(D), 100*Matrix{Float64}(I, D, D)))
		tcov[:,:,i] .= rand(InverseWishart(D+2, Matrix{Float64}(I, D, D)))
		# T = chol(slice(tcov,:,:,i))
		# x[:,indices] = broadcast(+, T'*x[:,indices], tmean[:,i]);
		d = MvNormal(tmean[:,i], tcov[:,:,i])
		for j=indices
			x[:,j] = rand(d)
		end
		ind += tzn[i]
	end
	x, tz, tmean, tcov
end

function generate_mnmm_data(N::Int64, D::Int64, K::Int64, trials::Int64)
	clusters = zeros(D,K)
	x = zeros(D,N)
	labels = rand(1:K,(N,))
	for i=1:K
		alphas = rand(1:20,(D,))
		alphas[i] = rand(30:100)
		clusters[:,i] = rand(Dirichlet(alphas))
	end
	for i=1:N
		x[:,i] = rand(Multinomial(trials,clusters[:,labels[i]]))
	end
	return x, labels, clusters
end

function generate_grouped_mnm_data(N::Int64, D_global::Int64, D_local, K_global::Int64, K_local::Int64, groups_count::Int64, rand_local::Bool, trials::Int64)
	global_weights = rand(Dirichlet(ones(K_global)))
	pts_dict = Dict()
	labels_dict = Dict()
	clusters_g = zeros(D_global,K_global)
	for i=1:length(global_weights)
		alphas = ones(D_global)*2
		alphas[rand(1:D_global,Int(floor(D_global/40)))] .= rand(trials/2:trials)
		clusters_g[:,i] = rand(Dirichlet(alphas))
	end
	for j=1:groups_count
		x = zeros(D_global + D_local,N)
		x_labels = zeros(size(x,2),2)
		if rand_local
			local_k = rand(1:K_local*2)
		else
			local_k = K_local
		end
		clusters_tzn = sample(1:length(global_weights),ProbabilityWeights(global_weights),local_k)
		local_weights = rand(Dirichlet(ones(local_k)*100))
		clusters_l = zeros(D_global+D_local,local_k)
		ind = 1
		group_ind = 1
		group_tzn = rand(Multinomial(N,local_weights))
		# group_tzn = rand(Multinomial(N,ones(local_k)))
		for i=1:length(local_weights)
			indices = ind:ind+group_tzn[i]-1
			# println(indices)

			alphas = ones(D_local)*2
			alphas[rand(1:D_local,Int(floor(D_local/1.25)))] .= rand(trials/2:trials)
			# local_part = rand(Dirichlet(alphas))
			# clusters_l[:,i] .= cat(clusters_g[:,clusters_tzn[i]], alphas, dims = [1])

			pvector = rand(Dirichlet(alphas))
			for a=indices
				x[1:D_global,a] = rand(Multinomial(trials,clusters_g[:,clusters_tzn[i]]))
				x[D_global+1:end,a] = rand(Multinomial(trials,pvector))
				x_labels[a,1] = clusters_tzn[i]
				x_labels[a,2] = i
			end
			ind += group_tzn[i]
		end
		pts_dict[j] = x
		labels_dict[j] = x_labels
	end
	return pts_dict, labels_dict
end

function generate_grouped_gaussian_data(N::Int64, D_global::Int64, D_local, K_global::Int64, K_local::Int64, groups_count::Int64, rand_local::Bool, var_size, hdp=false)
	global_weights = rand(Dirichlet(ones(K_global)))
	pts_dict = Dict()
	labels_dict = Dict()
	tmean = zeros(D_global,K_global)
	tcov = zeros(D_global,D_global,K_global)
	for i=1:length(global_weights)
		tmean[:,i] .= rand(MvNormal(zeros(D_global), var_size*Matrix{Float64}(I, D_global, D_global)))
		tcov[:,:,i] .= rand(InverseWishart(D_global+2, Matrix{Float64}(I, D_global, D_global)))
	end
	for j=1:groups_count
		x = randn(D_global + D_local,N)
		x_labels = zeros(size(x,2),2)
		if rand_local
			local_k = rand(K_local-2:K_local+2)
		else
			local_k = K_local
		end
		clusters_tzn = sample(1:length(global_weights),ProbabilityWeights(global_weights),local_k)
		local_weights = rand(Dirichlet(ones(local_k)*100))
		group_mean = zeros(D_global+D_local,local_k)
		group_cov = zeros(D_global+D_local,D_global+D_local,local_k)
		ind = 1
		group_ind = 1
		group_tzn = rand(Multinomial(N,local_weights))
		# group_tzn = rand(Multinomial(N,ones(local_k)))
		for i=1:length(local_weights)
			indices = ind:ind+group_tzn[i]-1
			g_mean = rand(MvNormal(zeros(D_local), var_size*Matrix{Float64}(I, D_local, D_local)))
			g_cov = rand(InverseWishart(D_local+2, Matrix{Float64}(I, D_local, D_local)))
			group_mean[:,i] .= cat(tmean[:,clusters_tzn[i]], g_mean, dims = [1])
			group_cov[:,:,i] .= cat(tcov[:,:,clusters_tzn[i]], g_cov, dims = [1,2])
			d = MvNormal(group_mean[:,i], group_cov[:,:,i])
			for j=indices
				x[:,j] = rand(d)
				x_labels[j,1] = clusters_tzn[i]
				x_labels[j,2] = i
			end
			ind += group_tzn[i]
		end
		if hdp
			x[D_global+1:end,:] .= 0
		end
		pts_dict[j] = x
		labels_dict[j] = x_labels
	end
	return pts_dict, labels_dict
end

function create_mnmm_data()
	 x,labels,clusters = generate_mnmm_data(10^6,100,6,1000)
	 npzwrite("data/mnmm/1milD100K6.npy",x')
	 x,labels,clusters = generate_mnmm_data(10^6,100,60,100)
	 npzwrite("data/mnmm/1milD100K60.npy",x')
 end

function create_gaussian_data()
	x,labels,clusters = generate_gaussian_data(10^5,2,20)
	npzwrite("data/2d-1mil/samples_100k1.npy",x')
	# x,labels,clusters = generate_gaussian_data(10^5,100,60,100)
	# npzwrite("data/30d-1mil/samples.npy",x')
end

function create_grouped_gaussian_data()
	samples_count = 10000
	global_dim = 2
	local_dim = 1
	global_clusters_count = 4
	local_clusters_count = 5
	groups_count = 5
	random_local_cluster_count = true
	var_size = 60
	hdp = true

	path_to_save = "data/3d-gaussians/"
	prefix = "gg"

	pts, labels = generate_grouped_gaussian_data(samples_count,
		global_dim,
		local_dim,
		global_clusters_count,
		local_clusters_count,
		groups_count,
		random_local_cluster_count,
		var_size,
		hdp)
	for (k,v) in pts
		npzwrite(path_to_save * prefix * string(k) * ".npy",v')
	end
	for (k,v) in labels
		npzwrite(path_to_save * prefix*"_labels" * string(k) * ".npy",v)
	end
end

function create_grouped_mnmm_data()
	samples_count = 10000
	global_dim = 100
	local_dim = 100
	global_clusters_count = 10
	local_clusters_count = 20
	groups_count = 4
	random_local_cluster_count = false
	trials = 1000

	path_to_save = "data/200d-mnm/"
	prefix = "g"

	pts, labels = generate_grouped_mnm_data(samples_count,
		global_dim,
		local_dim,
		global_clusters_count,
		local_clusters_count,
		groups_count,
		random_local_cluster_count,
		trials)
	for (k,v) in pts
		npzwrite(path_to_save * prefix * string(k) * ".npy",v')
	end
	for (k,v) in labels
		npzwrite(path_to_save * prefix*"_labels" * string(k) * ".npy",v)
	end
end

function single_crp_draw(tables,α)
	tables = push!(tables,α)
	sum_tables = sum(tables)
	table_probs = [x / sum_tables for x in tables]
	return sample(ProbabilityWeights(table_probs))
end

function hdp_prior_crf_draws(N,J,α,γ)
	groups_tables_counts = Dict()
	for j=1:J
		table_counts = []
		for i=1:N
			point_table = single_crp_draw(table_counts[:],α)
			if point_table == length(table_counts)+1
				push!(table_counts,1)
			else
				table_counts[point_table] += 1
			end
		end
		groups_tables_counts[j] = table_counts
	end
	global_tables_count = []
	groups_tables_assignments = Dict([i=>[] for i=1:J])
	cur_group_table = Dict([i=>1 for i=1:J])
	is_done = [false for i=1:J]
	while any(is_done .== false)
		for j=1:J
			group_assignments = groups_tables_assignments[j]
			i = cur_group_table[j]
			if i > length(groups_tables_counts[j])
				is_done[j] = true
				continue
			end
			table_assignment = single_crp_draw(global_tables_count,γ)
			if table_assignment == length(global_tables_count)+1
				push!(global_tables_count,1)
			else
				global_tables_count[table_assignment] += 1
			end
			push!(group_assignments,table_assignment)
			groups_tables_assignments[j] = group_assignments
			cur_group_table[j]+=1
		end
	end
	global_mixture = [x / sum(global_tables_count) for x in global_tables_count]
	groups_mixtures = Dict()
	for j=1:J
		local_mixture = zeros(length(global_mixture))
		for i = 1:length(global_mixture)
			if i in groups_tables_assignments[j]
				local_mixture[i] = sum(groups_tables_counts[j][groups_tables_assignments[j] .== i])
			end
		end
		groups_mixtures[j] = local_mixture
	end
	return global_mixture,groups_mixtures,groups_tables_counts,groups_tables_assignments
end

function generate_grouped_gaussian_from_hdp_group_counts(group_counts, dim, var_size)
	pts_dict = Dict()
	labels_dict = Dict()
	K = length(group_counts[1])
	tmean = zeros(dim,K)
	tcov = zeros(dim,dim,K)
	J = length(group_counts)
	components = []
	for i=1:length(group_counts[1])
		tmean[:,i] .= rand(MvNormal(zeros(dim), var_size*Matrix{Float64}(I, dim, dim)))
		tcov[:,:,i] .= rand(InverseWishart(dim+2, Matrix{Float64}(I, dim, dim)))
		push!(components,MvNormal(tmean[:,i], tcov[:,:,i]))
	end
	for j=1:J
		points = reduce(hcat,[rand(components[i],Int(group_counts[j][i])) for i=1:length(group_counts[j])])
		labels = reduce(vcat,[Int.(ones(Int(group_counts[j][i])))*i for i=1:length(group_counts[j])])
		pts_dict[j] = points
		labels_dict[j] = labels
	end
	return pts_dict, labels_dict
end


function demoMaxProd()
    #DEMOMAXPROD Max-Product algorithm test :
    # Variable order is arbitary
    variables=1:5
    a, b, c, d, e=variables
    nstates=round(Int64,3*rand(1,5)+2) # random number of states for each variable

    pot=Array(PotArray,5)
    pot[1]=PotArray([a b],rand(nstates[[a b]]...));
    pot[2]=PotArray([b c d],rand(nstates[[b c d]]...));
    pot[3]=PotArray([c],rand(nstates[[c]]...));
    pot[4]=PotArray([e d],rand(nstates[[e d]]...));
    pot[5]=PotArray([d],rand(nstates[[d]]...));

    A = FactorGraph(pot)
    maxstate, maxval, mess=maxprod(pot,A)

    # check if this is correct by brute force max:
    jointpot =prod(pot); V=length(pot)
    maxval2,maxstate2=max(jointpot,variables,ReturnState=true,Ind2Sub=true)
    println("\t\t maxprod \t brute")
    for i=1:V
        println("Variable $i \t $(maxstate[i]...)\t\t $(maxstate2[i])")
    end

end


using SIAN

println("Setting up the problem")

ode = @ODEmodel(
            x1'(t) = -beta * x1(t) * x4(t) - d * x1(t) + s,
            x2'(t) = beta * q1 * x1(t) * x4(t) - k1 * x2(t) - mu1 * x2(t),
            x3'(t) = beta * q2 * x1(t) * x4(t) + k1 * x2(t) - mu2 * x3(t),
            x4'(t) = -c * x4(t) + k2 * x3(t),
            y1(t) = x1(t),
            y2(t) = x4(t)
)

res = identifiability_ode(ode, get_parameters(ode); p=0.99, p_mod=0, nthrds=1)

println(res)


s = open("input") do file
    readlines(file);
end

nums = map(x->parse(Int, x), s);

sort!(nums);

function GetPathCount(nums)
    if nums[1] != 0
        pushfirst!(nums, 0);
    end

    combos = ones(Int128, length(nums));
    for i = 2:length(nums)
        count::Int128 = combos[i-1];
        for j = i-3:i-2
            if j > 0
                if nums[i] - nums[j] <= 3
                    count += combos[j];
                end
            end
        end

        combos[i] = max(1, count);
    end

    return combos[end];
end

println(GetPathCount(nums))


struct Scaler{T} <: Number
    s::T
end

function Scaler(x::Vector{T}; robust::Bool=false) where T
    s = robust ? 1.482602median(abs.(x .- median(x))) : std(x)
    Scaler(s)
end

(s::Scaler)(x::Vector) = x ./ s.s

mutable struct StepScaler{S,F} <: AbstractStep
    selections::S
    params::Union{Nothing, Dict{Symbol,Scaler}}
    robust::Bool
    trained::Bool
    skip::Bool
    prehook::F
end

function transform!(s::StepScaler, df)
    for (k,f) in s.params
        df[k] = f(df[k])
    end
    df
end
function transform(s::StepScaler, df)
    transform!(s, copy(df))
end
function fit!(s::StepScaler, df)
    s.params = Dict(col => Scaler(s.prehook(df[col])) for col in getselectionkeys(df, s.selections))
    s.trained = true
end
function step_scale!(r::Recipe, s...; robust::Bool=false, abs::Bool=false, skip::Bool=false, prehook=identity)
    push!(r.steps,
          StepScaler([s...],
                     nothing,
                     robust,
                     false,
                     skip,
                     prehook))
end


module TestEffectiveNumbers
using Test
using LinearAlgebra
using Statistics

using Diversity
using Diversity: powermean

numbers = [1.0, 2, 4, 8, 16];
numspecies = 100;
fragments = rand(numspecies);
weights = rand(numspecies);
weights /= sum(weights);
Z1 = ones(typeof(weights[1]), (length(weights), length(weights)));
numcommunities = 8;
manyweights = rand(numspecies, numcommunities);
manyweights *= Diagonal(reshape(mapslices(v -> 1. / sum(v), manyweights, dims=1),
                                numcommunities));

# Simple power means - we no longer export these, but we should check
# them anyway as everything relies on them

@testset "powermean" begin
    # Check that an exception is thrown when 'values' and weights are different lengths
    @test_throws DimensionMismatch powermean(numbers, 0, weights)

    # Some simple values
    @test powermean([1.0], 0.0, [1.0]) ≈ 1.0
    @test powermean(numbers, 0.0) ≈ 4.0
    @test powermean(numbers, [-Inf]) ≈ [1]
    @test powermean(numbers, [1.0, -1.0]) ≈ [31.0/5.0, 80.0/31.0]
    @test powermean(numbers, Inf, [1.0, 1.0, 1.0, 1.0, 0.0]) ≈ 8
    @test isnan(powermean(numbers, 0.0, 0.0 * numbers))

    # Power mean with some random numbers
    @test powermean(fragments, 0) ≈ prod(fragments .^ (1. / numspecies))
    @test powermean(fragments, 1) ≈ mean(fragments)
    @test powermean(fragments, Inf) ≈ maximum(fragments)
    @test powermean(fragments, 0, weights) ≈ prod(fragments .^ weights)
    @test powermean(fragments, 1, weights) ≈ sum(fragments .* weights)
    @test powermean(manyweights, -1, manyweights) .^ -1 ≈ numspecies * ones(size(manyweights, 2))
end

@testset "qD" begin
    # Basic qD diversity calculation
    @test qD(weights, 0) ≈ mapreduce((x) -> x ≈ 0 ? 0 : 1, +, weights)
    @test qD(Metacommunity(weights), 0) == qD(weights, 0)
    @test qD(Metacommunity(weights, UniqueTypes(numspecies)), 0) ==
        qD(weights, 0)
    @test_throws ErrorException qD(Metacommunity(weights,
                                   GeneralTypes(rand(numspecies, numspecies))),
                                   0)
    @test qD(weights, 1) ≈ prod(weights .^ -weights)
    @test qD(weights, 2) ≈ 1.0 / sum(weights .^ 2)
    @test qD(weights, Inf) ≈ 1.0 / maximum(weights)

    @test qD(weights, [1, 2]) ≈ [qD(weights, 1), qD(weights, 2)]

    @test typeof(qD(manyweights[:,1], 0)) <: AbstractFloat
    @test typeof(qD(manyweights[:,1], [0])) <: Vector

    for i in 1:size(manyweights, 2)
        @test qD(manyweights[:,i], [0]) ≈ numspecies * ones((1, size(manyweights[:,i], 2)))
    end
end

@testset "qDZ" begin
    # General Leinster-Cobbold diversity calculation
    @test qDZ(weights, [1, 2]) ≈ qD(weights, [1, 2])
    @test qDZ(weights, [0, 1, 2, 3, Inf], Z1) ≈ [1, 1, 1, 1, 1]

    for i in 1:size(manyweights, 2)
        @test qDZ(manyweights[:,i], [0, 1, 2, Inf],
                  ones((size(manyweights[:,i], 1),
                        size(manyweights[:,i], 1)))) ≈ ones((4, size(manyweights[:,i], 2)))
    end
end

end


# Random.seed!(1);
if "data" in keys(conf)
    idx_order = randperm(n_exp)
    pert = DataFrame(CSV.File(string(conf["data"],"/pert.csv"); header=false))
    μ_list = convert(Matrix,pert)[idx_order,:]
else
    μ_list = randomLHC(n_exp, ns) ./ n_exp;
    nμ = Int64(conf["n_mu"])
    for i = 1:n_exp
        nonzeros = findall(μ_list[i, :].>0)
        ind_zero = sample(nonzeros, max(0, length(nonzeros)-nμ), replace=false)
        μ_list[i, ind_zero] .= 0
    end
end

function gen_network(m; weight_params=(-1., 1.), sparsity=0., drop_range=(-1e-1, 1e-1))

    # uniform random for W matrix
    w = rand(Uniform(weight_params[1], weight_params[2]), (m, m))

    # Drop small values
    @inbounds for i in eachindex(w)
        w[i] = ifelse(drop_range[1]<=w[i]<=drop_range[2], 0, w[i])
    end

    # Add sparsity
    p = [sparsity, 1 - sparsity]
    w .*= sample([0, 1], weights(p), (m, m), replace=true)

    # Add α vector
    α = abs.(rand(Uniform(weight_params[1], weight_params[2]), (m))) .+ 0.5

    return hcat(α, w)
end

if "network" in keys(conf)
    df = DataFrame(CSV.File(conf["network"]))
    nodes = names(df)
    w = convert(Matrix, df[:,2:end])
    @assert size(w)[1] == size(w)[2]
    @assert size(w)[1] == ns
    if "randomize_network" in keys(conf)
        w_rand = rand(Normal(1, conf["randomize_network"]), (ns, ns))
        w = w .* w_rand
    end
    if "alpha" in keys(conf)
        α = ones(ns) .* conf["alpha"]
    else
        α = ones(ns) .* 0.2
    end
    p_gold = hcat(α, w)
elseif "data" in keys(conf)
    p_gold = hcat(ones(ns),zeros(ns,ns))
else
    p_gold = gen_network(ns, weight_params=(-1.0, 1.0),
                         sparsity=Float64(conf["sparsity"]),
                         drop_range=(Float64(conf["drop_range"]["lb"]), Float64(conf["drop_range"]["ub"])));
end
p = gen_network(ns; weight_params=(0.0, 0.01), sparsity=0);
# show_network(p)

function show_network(p)
    println("p_gold")
    show(stdout, "text/plain", round.(p_gold, digits=2))
    println("\np_learned")
    show(stdout, "text/plain", round.(p, digits=2))
end

function loss_network(p)
     # distalpha = cosine_dist(p_gold[:,1],p[:,1])
     # distw = cosine_dist(p_gold[:,2:end],p[:,2:end])
     @inbounds coralpha = cor(p_gold[:,1],p[:,1])
     @inbounds corw = cor([p_gold[:,2:end]...],[p[:,2:end]...])
     return coralpha, corw
end

function cellbox!(du, u, p, t)
    @inbounds du .= @view(p[:, 1]) .* tanh.(@view(p[:, 2:end]) * u - μ) .- u
end

tspan = (0, tfinal);
ts = 0:tspan[2]/ntotal:tspan[2];
ts = ts[2:end];
u0 = zeros(ns);
prob = ODEProblem(cellbox!, u0, tspan, saveat=ts);

function max_min(ode_data)
    return maximum(ode_data, dims=2) .- minimum(ode_data, dims=2)
end

if "data" in keys(conf)
    expr = DataFrame(CSV.File(string(conf["data"],"/expr.csv"); header=false))
    ode_data_list = zeros(Float64, (n_exp, ns, ntotal));
    ode_data_list[:,:,1] = convert(Matrix,expr)[idx_order,:]
else
    ode_data_list = zeros(Float64, (n_exp, ns, ntotal));
    yscale_list = [];
    for i = 1:n_exp
        global μ = μ_list[i, 1:ns]
        ode_data = Array(solve(prob, Tsit5(), u0=u0, p=p_gold))

        ode_data += randn(size(ode_data)) .* noise
        ode_data_list[i, :, :] = ode_data

        push!(yscale_list, max_min(ode_data))
    end
    yscale = maximum(hcat(yscale_list...), dims=2);
end

function predict_neuralode(u0, p, i_exp=1, batch=ntotal, saveat=true)
    global μ = μ_list[i_exp, 1:ns]
    if saveat
        @inbounds _prob = remake(prob, p=p, tspan=[0, ts[batch]])
        pred = Array(solve(_prob, Tsit5(), saveat=ts[1:batch],
                     sensealg=InterpolatingAdjoint()))
    else # for full trajectory plotting
        @inbounds _prob = remake(prob, p=p, tspan=[0, ts[end]])
        pred = Array(solve(_prob, Tsit5(), saveat=0:ts[end]/nplot:ts[end]))
    end
    return pred
end
predict_neuralode(u0, p, 1);

function loss_neuralode(p, i_exp=1, batch=ntotal)
    pred = predict_neuralode(u0, p, i_exp, batch)
    @inbounds loss = mae(@views(ode_data_list[i_exp, :, 1:batch]), pred)
    return loss
end
@show loss_neuralode(p, 1);


export TraceTestSystem

"""
    TraceTestSystem(F, p₀, v) <: AbstractSystem

The system `G(x,l) := [F(x, p₀+x[end]v); l ⋅[x[1:end-1];1])`.
"""
struct TraceTestSystem{S<:AbstractSystem, T} <: AbstractSystem
    F::S
    p₀::Vector{T}
    v::Vector{T}
end

function TraceTestSystem(F, p::AbstractVector, v::AbstractVector)
    TraceTestSystem(F, promote(p, v)...)
end

struct TraceTestSystemCache{SC<:AbstractSystemCache, T1, T2} <: AbstractSystemCache
    u::Vector{T1}
    U::Matrix{T1}
    Up::Matrix{T1}
    p::Vector{T2}
    cache::SC
end

function cache(F::TraceTestSystem, x, l)
    c = cache(F.F, x, l)
    u = evaluate(F.F, x, l, c)
    U = similar(u, size(F.F))
    Up = similar(u, size(F.F)[1], length(F.p₀))
    p = F.p₀ + 0.1 * F.v
    TraceTestSystemCache(u, U, Up, p, c)
end

Base.size(F::TraceTestSystem) = size(F.F) .+ 1

function evaluate!(u, F::TraceTestSystem, x, l, c::TraceTestSystemCache)
    λ = x[end]
    for i in eachindex(c.p)
        c.p[i] = F.p₀[i] + λ * F.v[i]
    end
    evaluate!(c.u, F.F, x, c.p, c.cache)
    n = length(x)
    for i in 1:length(c.u)
        u[i] = c.u[i]
    end
    u[end] = l[n]
    for i in 1:n-1
        u[end] += l[i] * x[i]
    end

    u
end

function jacobian!(U, F::TraceTestSystem, x, l, c::TraceTestSystemCache)
    λ = x[end]
    for i in eachindex(c.p)
        c.p[i] = F.p₀[i] + λ * F.v[i]
    end
    jacobian!(c.U, F.F, x, c.p, c.cache)
    differentiate_parameters!(c.Up, F.F, x, c.p, c.cache)
    n, m = length(x), size(U, 1)
    for j in 1:n-1, i in 1:m-1
        U[i,j] = c.U[i,j]
    end
    for i in 1:m
        U[i,n] = zero(eltype(U))
    end
    for k in 1:size(c.Up, 2), i in 1:m-1
        U[i, n] += c.Up[i,k] * F.v[k]
    end

    for j in 1:n-1
        U[m,j] = l[j]
    end
    U
end

function evaluate_and_jacobian!(u, U, F::TraceTestSystem, x, l, c::TraceTestSystemCache)
    λ = x[end]
    for i in eachindex(c.p)
        c.p[i] = F.p₀[i] + λ * F.v[i]
    end
    n, m = length(x), size(U, 1)

    evaluate_and_jacobian!(c.u, c.U, F.F, x, c.p, c.cache)
    for i in 1:length(c.u)
        u[i] = c.u[i]
    end
    u[end] = l[n]
    for i in 1:n-1
        u[end] += l[i] * x[i]
    end

    differentiate_parameters!(c.Up, F.F, x, c.p, c.cache)

    for j in 1:n-1, i in 1:m-1
        U[i,j] = c.U[i,j]
    end
    for i in 1:m
        U[i,n] = zero(eltype(U))
    end
    for k in 1:size(c.Up, 2), i in 1:m-1
        U[i, n] += c.Up[i,k] * F.v[k]
    end

    for j in 1:n-1
        U[m,j] = l[j]
    end

    nothing
end

function differentiate_parameters!(U, F::TraceTestSystem, x, l, c::TraceTestSystemCache)
    U .= zero(eltype(U))
    for i in 1:length(l)-1
        U[end, i] = x[i]
    end
    U[end,length(l)] = one(eltype(U))
    U
end

function differentiate_parameters(F::TraceTestSystem, x, l, c::TraceTestSystemCache)
    differentiate_parameters!(similar(c.U, size(F)[1], length(l)), F, x, l, c)
end

function evaluate(F::TraceTestSystem, x, l, c::TraceTestSystemCache)
    evaluate!(similar(c.u, size(F)[1]), F, x, l, c)
end
function jacobian(F::TraceTestSystem, x, l, c::TraceTestSystemCache)
    jacobian!(similar(c.U, size(F)), F, x, l, c)
end


using Dolo

dolo_dir = Dolo.pkg_path

model = Dolo.yaml_import("examples/models/rbc_iid.yaml")

process = Dolo.MvNormal(0.001)

dp = Dolo.discretize(process)

@time sol = time_iteration(model, dp; verbose=true, maxit=5)

@time dd = Dolo.improved_time_iteration(model, dp, sol.dr; verbose=true)

ivals = [dr(1, [i])[1] for i=kvec ]
nvals = [dr(1, [i])[2] for i=kvec ]

using Gadfly

plot(x=kvec, y=ivals)

plot(x=kvec, y=nvals)

kkvec = [dr(1,[k])[1] for k in kvec]
plot(x=kvec, y=yvec, Geom.line)


export ThinCategory, FreeThinCategory,
  ThinSymmetricMonoidalCategory, FreeThinSymmetricMonoidalCategory,
  Preorder, FreePreorder, Elt, Leq, ≤, lhs, rhs, reflexive, transitive

# Thin category
###############

""" Theory of *thin categories*

Thin categories have at most one morphism between any two objects.
"""
@theory Category(Ob,Hom) => ThinCategory(Ob,Hom) begin
  f == g ⊣ (A::Ob, B::Ob, f::Hom(A,B), g::Hom(A,B))
end

@syntax FreeThinCategory(ObExpr,HomExpr) ThinCategory begin
  compose(f::Hom, g::Hom) = associate(new(f,g; strict=true))
end

@theory SymmetricMonoidalCategory(Ob,Hom) => ThinSymmetricMonoidalCategory(Ob,Hom) begin
  f == g ⊣ (A::Ob, B::Ob, f::Hom(A,B), g::Hom(A,B))
end

@syntax FreeThinSymmetricMonoidalCategory(ObExpr,HomExpr) ThinSymmetricMonoidalCategory begin
  compose(f::Hom, g::Hom) = associate(new(f,g; strict=true))
  otimes(A::Ob, B::Ob) = associate_unit(new(A,B), munit)
  otimes(f::Hom, g::Hom) = associate(new(f,g))
end

# Preorder
##########

""" Theory of *preorders*

Preorders encode the axioms of reflexivity and transitivity as term constructors.
"""
@theory Preorder(Elt,Leq) begin
  Elt::TYPE
  Leq(lhs::Elt, rhs::Elt)::TYPE
  @op (≤) := Leq

  # Preorder axioms are lifted to term constructors in the GAT.
  reflexive(A::Elt)::(A≤A) # ∀ A there is a term reflexive(A) which implies A≤A
  transitive(f::(A≤B), g::(B≤C))::(A≤C) ⊣ (A::Elt, B::Elt, C::Elt)

  # Axioms of the GAT are equivalences on terms or simplification rules in the logic
  f == g ⊣ (A::Elt, B::Elt, f::(A≤B), g::(A≤B))
  # Read as (f⟹ A≤B ∧ g⟹ A≤B) ⟹ f ≡ g
end

@syntax FreePreorder(ObExpr,HomExpr) Preorder begin
  transitive(f::Leq, g::Leq) = associate(new(f,g; strict=true))
end

# TODO: a GAT-homomorphism between the Preorder GAT and the ThinCategory GAT
# this is a morphism is *GAT* the category whose objects are GATs
# and whose morphisms are algebraic structure preserving maps

# @functor F(Preorder(Elt, Leq))::ThinCategory(Ob, Hom) begin
#   Elt ↦ Ob
#   Leq ↦ Hom
#   reflexive ↦ id
#   transitive ↦ compose
# end


export AbstractSystem, AbstractHomotopy, jacobian

"""
    AbstractSystem

An abstract type representing a polynomial system ``F(x, p)`` where ``x`` are variables
and ``p`` are possible parameters.
"""
abstract type AbstractSystem end

# Base.size(F::AbstractSystem)
Base.size(F::AbstractSystem, i::Integer) = size(F)[i]

"""
    parameters(F::AbstractSystem)

Returns the [`Variable`](@ref)s used in the system.
"""
parameters(F::AbstractSystem) = Variable[]

"""
    variable_groups(F::AbstractSystem)

Returns the variable groups of the system.
"""
variable_groups(::AbstractSystem) = nothing

"""
    nvariables(F::AbstractSystem)

Returns the number of variables of a given system `F`.
"""
nvariables(F::AbstractSystem) = size(F, 2)

"""
    nparameters(F::AbstractSystem)

Returns the number of parameters of a given system `F`.
"""
nparameters(F::AbstractSystem) = length(parameters(F))

"""
    System(F::AbstractSystem)

Construct a (symbolic) [`System`](@ref) from `F`.
"""
function System(F::AbstractSystem)
    x, p = variables(F), parameters(F)
    System(F(x, p), x, p, variable_groups(F))
end

(F::AbstractSystem)(x, p = nothing) = evaluate(F, x, p)

"""
    evaluate(F::AbstractSystem, x, p = nothing)

Evaluate the given system.
"""
function evaluate(F::AbstractSystem, x, p = nothing)
    u = Vector{Any}(undef, size(F, 1))
    evaluate!(u, F, x, p)
    ModelKit.to_smallest_eltype(u)
end

"""
    jacobian(F::AbstractSystem, x, p = nothing)

Compute the Jacobian of the given system.
"""
function jacobian(F::AbstractSystem, x, p = nothing)
    u = Vector{Any}(undef, size(F, 1))
    U = Matrix{Any}(undef, size(F))
    evaluate_and_jacobian!(u, U, F, x, p)
    ModelKit.to_smallest_eltype(U)
end

##############
## Homotopy ##
##############

"""
    AbstractHomotopy

An abstract type representing a homotopy ``H(x, t)`` where ``x`` are variables
and ``t`` is space / time.
"""
abstract type AbstractHomotopy end

Base.size(H::AbstractHomotopy, i::Integer) = i == 1 ? first(size(H)) : last(size(H))

(H::AbstractHomotopy)(x, t, p = nothing) = evaluate(H, x, t, p)
function evaluate(H::AbstractHomotopy, x, t, p = nothing)
    U = Vector{Any}(undef, first(size(H)))
    to_smallest_eltype(evaluate!(U, H, x, t, p))
end

function jacobian(H::AbstractHomotopy, x, t, p = nothing)
    n, m = size(H)
    u = Vector{Any}(undef, size(H, 1))
    U = Matrix{Any}(undef, size(H))
    evaluate_and_jacobian!(u, U, H, x, t, p)
    to_smallest_eltype(U)
end


using Flux, Flux.Tracker
using Distributions

include("../ADVI/ADVI.jl")

mutable struct State
  w
  m
  s
  State() = new()
end

function State(K::Integer)
  state = State()

  # NOTE: this works! Notice that dimensions do matter
  #       especially for dirichlet stuff. Needs to be a tuple for size.
  #       So, the minimum must be (1, K-1) as the last dim is the dirichlet
  #       dim. Maybe needs a FIXME?
  state.m = ADVI.ModelParam(K, "real")
  state.w = ADVI.ModelParam((1, K - 1), "simplex",
                            m=fill(0.0, 1, K-1),
                            s=fill(exp(-2), 1, K-1))
  state.s = ADVI.ModelParam(K, "positive",
                            m=fill(0.0, K),
                            s=fill(exp(-2), K))

  return state
end

function rsample(state::State)
  real = State()
  tran = State()

  for key in fieldnames(State)
    f = getfield(state, key)

    r = ADVI.rsample(f)
    t = ADVI.transform(f, r)

    setfield!(real, key, r)
    setfield!(tran, key, t)
  end

  return real, tran
end

function loglike(tran::State, y::Vector{Float64})
  K = length(tran.m)
  N = length(y)

  # NOTE: Dirichlet parameters need special treatment because it changes dims!
  w = reshape(tran.w, 1, K)
  m = reshape(tran.m, 1, K)
  s = reshape(tran.s, 1, K)
  y_rs = reshape(y, N, 1)
  return sum(ADVI.lpdf_gmm(y_rs, m, s, w, dims=2, dropdim=true))
end

function logprior(real::State, tran::State, mp::State)
  lp = 0.0

  lp += sum(ADVI.compute_lpdf(Normal(1.85, 5), tran.m))

  lp += sum(ADVI.compute_lpdf(Gamma(1, 0.1), tran.s))
  lp += sum(ADVI.logabsdetJ(mp.s, real.s, tran.s))

  K = length(tran.w)
  lp += sum(ADVI.compute_lpdf(Dirichlet(ones(K) / K), tran.w))
  lp += sum(ADVI.logabsdetJ(mp.w, real.w, tran.w))

  return lp
end

function logq(real::State, mp::State)
  lq = 0.0
  for key in fieldnames(State)
    vp = getfield(mp, key)
    r = getfield(real, key)
    lq += sum(ADVI.log_q(vp, r))
  end
  return lq
end

function compute_elbo(mp::State, y::Vector{Float64}, N::Integer,
                      metrics::Dict{Symbol, Vector{Float64}})
  real, tran = rsample(mp)

  ll = loglike(tran, y) * N / length(y)
  lp = logprior(real, tran, mp)
  lq = logq(real, mp)

  elbo = ll + lp - lq
  append!(metrics[:elbo], elbo.data)

  return elbo
end


#=

  Euler #39 in Julia.

  """
  If p is the perimeter of a right angle triangle with integral length sides,
  {a,b,c}, there are exactly three solutions for p = 120.

  {20,48,52}, {24,45,51}, {30,40,50}

  For which value of p <= 1000, is the number of solutions maximised?
  """

  This Julia model was created by Hakan Kjellerstrand, hakank@gmail.com
  See also my Julia page: http://www.hakank.org/julia/

=#

include("Euler.jl")

# 0.11346191s
# This is surprinsingly slow!
function euler39a()
    n = 1000;
    # squares = Set((1:n).|>i->i*i) # slightly slower
    squares = Dict((1:n).|>i->i*i=>1)
    valid = []
    for x in keys(squares)
        for y in keys(squares)
            if x < y && haskey(squares,x+y) && (sqrt(x) + sqrt(x) + sqrt(x+y)) < 1000
                push!(valid,[x, y])
            end
        end
    end

    counts = Dict()
    for x in valid
        c = floor(Int,sqrt(x[1]) + sqrt(x[2]) + sqrt(x[1]+x[2]))
        get!(counts,c,0)
        counts[c] += 1
   end

    # find max count
    maxV = maximum(values(counts))
    t = [i for i in keys(counts) if counts[i] == maxV]
    return t[1]
end

run_euler(euler39a)


# TODO
# compile Mod and KMod to elementary gates.
export Mod, KMod

"""
    mod_inverse(x::Int, N::Int) -> Int

Return `y` that `(x*y)%N == 1`, notice the `(x*y)%N` operations in Z* forms a group and this is the definition of inverse.
"""
function mod_inverse(x::Int, N::Int)
    for i=1:N
        (x*i)%N == 1 && return i
    end
    throw(ArgumentError("Can not find the inverse, $x is probably not in Z*($N)!"))
end

"""
    Mod{N} <: PrimitiveBlock{N}

calculates `mod(a*x, L)`, notice `gcd(a, L)` should be 1.
"""
struct Mod{N} <: PrimitiveBlock{N}
    a::Int
    L::Int
    function Mod{N}(a, L) where N
        @assert gcd(a, L) == 1 && L<=1<<N
        new{N}(a, L)
    end
end

function _apply!(reg::ArrayReg{B}, m::Mod{N}) where {B, N}
    nstate = zero(reg.state)
    for i in basis(reg)
        _i = i >= m.L ? i+1 : mod(i*m.a, m.L)+1
        for j in 1:B
            @inbounds nstate[_i,j] = reg.state[i+1,j]
        end
    end
    reg.state = nstate
    reg
end

function Yao.mat(::Type{T}, m::Mod{N}) where {T, N}
    perm = Vector{Int}(undef, 1<<N)
    for i in basis(N)
        @inbounds perm[i >= m.L ? i+1 : mod(i*m.a, m.L)+1] = i+1
    end
    PermMatrix(perm, ones(T, 1<<N))
end

Base.adjoint(m::Mod{N}) where N = Mod{N}(mod_inverse(m.a, m.L), m.L)
Yao.print_block(io::IO, m::Mod{N}) where N = print(io, "Mod{$N}: $(m.a)*x % $(m.L)")

Yao.isunitary(::Mod) = true
# Yao.ishermitian(::Mod) = true  # this is not true for L = 1.
# Yao.isreflexive(::Mod) = true  # this is not true for L = 1.

"""
    KMod{N, K} <: PrimitiveBlock{N}

The first `K` qubits are exponent `k`, and the rest `N-K` are base `a`,
it calculates `mod(a^k*x, L)`, notice `gcd(a, L)` should be 1.
"""
struct KMod{N, K} <: PrimitiveBlock{N}
    a::Int
    L::Int
    function KMod{N, K}(a, L) where {N, K}
        @assert gcd(a, L) == 1 && L<=1<<(N-K)
        new{N, K}(a, L)
    end
end

nqubits_data(m::KMod{N, K}) where {N, K} = N-K
nqubits_control(m::KMod{N, K}) where {N, K} = K

function bint2_reader(T, k::Int)
    mask = bmask(T, 1:k)
    return b -> (b&mask, b>>k)
end

function _apply!(reg::ArrayReg{B}, m::KMod{N, K}) where {B, N, K}
    nstate = zero(reg.state)

    reader = bint2_reader(Int, K)
    for b in basis(reg)
        k, i = reader(b)
        _i = i >= m.L ? i : mod(i*powermod(m.a, k, m.L), m.L)
        _b = k + _i<<K + 1
        for j in 1:B
            @inbounds nstate[_b,j] = reg.state[b+1,j]
        end
    end
    reg.state = nstate
    reg
end

function Yao.mat(::Type{T}, m::KMod{N, K}) where {T, N, K}
    perm = Vector{Int}(undef, 1<<N)
    reader = bint2_reader(Int, K)
    for b in basis(N)
        k, i = reader(b)
        _i = i >= m.L ? i : mod(i*powermod(m.a, k, m.L), m.L)
        _b = k + _i<<K + 1
        @inbounds perm[_b] = b+1
    end
    PermMatrix(perm, ones(T, 1<<N))
end

Base.adjoint(m::KMod{N, K}) where {N, K} = KMod{N, K}(mod_inverse(m.a, m.L), m.L)
Yao.print_block(io::IO, m::KMod{N, K}) where {N, K} = print(io, "Mod{$N, $K}: $(m.a)^k*x % $(m.L)")

Yao.isunitary(::KMod) = true
# Yao.ishermitian(::Mod) = true  # this is not true for L = 1.
# Yao.isreflexive(::Mod) = true  # this is not true for L = 1.


using AutomatonEdification

@testset "mse" begin
    @test AutomatonEdification.LossFunctions.mse([0], [0]) == 0
    @test AutomatonEdification.LossFunctions.mse([1], [-1]) == 4

    @test AutomatonEdification.LossFunctions.mse([-1, -1, -1], [1, 2, 3]) ==
        ((1 - (-1))^2 + (2 - (-1))^2 + (3 - (-1))^2) / 3
end

@testset "cross_entropy" begin
    # TODO: add tests for cross_entropy
end


export facets

"""
    facets(
        vertices :: Vector{Vector{Int64}};
        dir = "./" :: String,
        cleanup=true :: Bool
    ) :: Dict{String, Vector{Vector{Int64}}}

Computes the facet inequalities and equalities which bound the convex polyhedron
defined by the set of `vertices`.
If the optimal representation is used for the vertices, then no equalitities will
be present.
For example, if the "generalized" vertex representation is used the normalization
constraints will be included in the resulting equalities.
The output of this method is structured:

```
Dict(
    "facets" => inequalities, # :: Vector{Vector{Int64}}
    "equalitites" => equalities, # :: Vector{Vector{Int64}}
)
```

Facet inequalities and equalities are represented as single vector `f` where
`f[1:(end-1)]` contains the coefficients of the linear inquality and `f[end]`
is the bound.
Facet inequalities and equalities act upon a vertex `v` where inequalities are
arranged such that there is an implicit `f[1:(end-1)]' * v ≤ f[end]`
and equalities are arranged such that `f[1:(end-1)]' * v == f[end]`

!!! note "Supporting Software"
    The vertex -> facet transformation is performed using the
    [`traf`](https://juliapolyhedra.github.io/XPORTA.jl/dev/exports/#XPORTA.traf)
    method of [XPORTA.jl](https://github.com/JuliaPolyhedra/XPORTA.jl).
    Please refer to the source code for more details.

!!! warning "Performance Limitations"
    Vertex -> facet transformations are notoriously difficult computational problems.
    The performance limits of this method will be met far before the limits the
    vertex enumeration methods.
"""
function facets(
    vertices :: Vector{Vector{Int64}};
    dir = "./" :: String,
    cleanup=true :: Bool
) :: Dict{String, Vector{Vector{Int64}}}
    ieq = traf(POI(vertices = hcat(vertices...)'[:,:]), dir=dir, cleanup=cleanup)

    inequalities = convert.(Int64, ieq.inequalities)
    equalities = convert.(Int64, ieq.equalities)

    Dict(
        "facets" => map(row_id -> inequalities[row_id,:], 1:size(inequalities,1)),
        "equalities" => map(row_id -> equalities[row_id,:], 1:size(equalities,1))
    )
end

"""
    facets(poly::Polyhedron) :: Vector{Vector{Int64}}

Returns the facets of the local polytope `poly`. If the facets have not already
been computed, then they are transformed into the half-space representation from
the vertex representation.
"""
function facets(poly::Polyhedron) :: Vector{Vector{Int64}}
    inequalities = convert.(Int64, hrep(poly).inequalities)
    map(row_id -> inequalities[row_id,:] , 1:size(inequalities,1))
end


"""
    UncertainEmpiricalVectorValue

An empirical value represented by a distribution estimated from actual data.

## Fields
- **`distribution`** The distribution describing the value.
- **`values`**: The values from which `distribution` is estimated.
"""
struct UncertainEmpiricalVectorValue{D <: Distribution, T} <: AbstractEmpiricalValue
    distribution::D
    values::AbstractVector{AbstractVector{T}}
end


using MITgcmTools, MeshArrays
using Test

@testset "MITgcmTools.jl" begin

    #format conversions
    (γ,Γ)=GridOfOnes("CubeSphere",30,30)
    @test isa(convert2gcmfaces(Γ["XC"]),Array)
    @test isa(convert2array(Γ["XC"]),Array)

    #physical oceanography
    (ρP,ρI,ρR) = SeaWaterDensity(3.,35.5,3000.)
    @test isapprox(ρI,1041.83267, rtol=1e-6)

    D=collect(0.0:1.0:500.0); tmp=(1.0.-tanh.(5*(-1 .+ 2/D[end]*D)));
    T=2.0 .+ 8.0*tmp; S=34.0 .+ 0.5*tmp;
    (ρP,ρI,ρR) = SeaWaterDensity(T,S,D);
    mld=MixedLayerDepth(T,S,D,"BM");
    @test isapprox(mld,134.0)

end


module Tools3D

using PyPlot, Random, TransD_GP, GP, Statistics

function geomprogdepth(n, dy, c)
    dy*(1.0-c^n)/(1-c)
end

function getn(z, dy, c)
    log(1 - z/dy*(1-c))/log(c)
end

function nicenup(g::PyPlot.Figure;fsize=16)
    for ax in gcf().axes
        ax.tick_params("both",labelsize=fsize)
        ax.xaxis.label.set_fontsize(fsize)
        ax.yaxis.label.set_fontsize(fsize)
        any(keys(ax) .== :zaxis) && ax.zaxis.label.set_fontsize(fsize)
        ax.title.set_fontsize(fsize)

        if typeof(ax.get_legend_handles_labels()[1]) != Array{Any,1}
            ax.legend(loc="best", fontsize=fsize)
        end
    end
    g.tight_layout()
end

function meshgrid(xs, ys)
    xx, yy = [xs[i] for i in 1:length(xs), j in 1:length(ys)], [ys[j] for i in 1:length(xs), j in 1:length(ys)]
    return xx, yy
end

function makezn(;n          = 20,
                 dz         = 1.6,
                 extendfrac = 1.245)
    @assert typeof(n) == Int
    @assert dz > 0.0
    @assert extendfrac > 1.0
    znrange            = 1.0:n
    zboundaries        = geomprogdepth.(znrange, dz, extendfrac)
    thickness          = [zboundaries[1]; diff(zboundaries)[1:end-1]]
    zall               = [zboundaries[1]/2; 0.5*(zboundaries[1:end-1] + zboundaries[2:end])]
    znall              = getn.(zall, dz, extendfrac)

    figure()
    plot(znall, zall)
    xlabel("depth index")
    ylabel("depth associated km")
    grid()
    nicenup(gcf())

    f, ax = plt.subplots(1, 2, figsize=(10,5))
    ax[1].stem(zboundaries[1:end-1], zboundaries[1:end-1], markerfmt="")
    ax[1].stem(zall, zall, "k--", markerfmt=" ")
    ax[1].set_xlabel("depth km")
    ax[1].set_ylabel("depth km")

    ax[2].stem(znrange[1:end-1], znrange[1:end-1], markerfmt="")
    ax[2].stem(znall, znall, "k--", markerfmt=" ")
    ax[2].set_ylabel("depth index")
    ax[2].set_xlabel("depth index")
    nicenup(gcf())

    f, ax = plt.subplots(1, 2, figsize=(10,5), sharey=true)
    ax[1].stem(zall[1:end-1],thickness, "k--", markerfmt=" ")
    ax[1].set_xlabel("depth km")
    ax[1].set_ylabel("thickness km")
    ax[1].yaxis.grid(which="major")
    ax[2].stem(znall[1:end-1],thickness, "k--", markerfmt=" ")
    ax[2].set_xlabel("depth index")
    ax[2].yaxis.grid(which="major")
    nicenup(f)
    return zall, znall
end

function makeopt(;nmin      = 2,
                  nmax      = 400,
                  λ         = [100.0, 200.0, 1.0],
                  δ         = 0.1,
                  fbounds   = [0.5 4],
                  demean    = true,
                  sdev_prop = [0.4],
                  sdev_pos  = [80.0;80.0;1.0],
                  stat_window = 500,
                  pnorm     = 2.0,
                  λx        = 1600.0,
                  λy        = 1600.0,
                  dxfrac    = 0.025,
                  dyfrac    = 0.025,
                  rseed     = 4,
                  timebirth = true,
                  znall     = nothing
                )
    @assert znall != nothing
    dx = dxfrac*λx
    dy = dyfrac*λy
    x = 0:(dx):λx-dx
    y = 0:(dy):λy-dy
    xall = zeros(3,length(x)*length(y)*length(znall))
    for i in 1:size(xall,2)
        xid, yid, zid = Tuple(CartesianIndices((length(x),length(y),length(znall)))[i])
        xall[:,i] = [x[xid]; y[yid]; znall[zid]]
    end
    xbounds = zeros(Float64,size(xall,1),2)
    for dim in 1:size(xall, 1)
        xbounds[dim,:] = [minimum(xall[dim,:]), maximum(xall[dim,:])]
    end
    ## Initialize a model using these options
    Random.seed!(rseed)
    opt = TransD_GP.OptionsStat(nmin = nmin,
                            nmax = nmax,
                            xbounds = xbounds,
                            fbounds = fbounds,
                            xall = xall,
                            λ = λ,
                            δ = δ,
                            demean = demean,
                            sdev_prop = sdev_prop,
                            sdev_pos = sdev_pos,
                            stat_window = stat_window,
                            pnorm = pnorm,
                            quasimultid = false,
                            updatenonstat = false,
                            needλ²fromlog = false
                            )

    optdummy = TransD_GP.OptionsNonstat(opt,
                            nmin = 2,
                            nmax = 3,
                            fbounds = fbounds,
                            δ = δ,
                            demean = demean,
                            sdev_prop = sdev_prop,
                            sdev_pos = sdev_pos,
                            pnorm = pnorm,
                            K = opt.K
                            )
    return opt, optdummy, x, y
end

function slicemodel(m::TransD_GP.Model,
                    opt::TransD_GP.Options;
                    slicesx    = [],
                    slicesy    = [],
                    slicesz    = [],
                    dz         = nothing,
                    extendfrac = nothing,
                    figsize    = (7,9)
                    )
    @assert !any((dz, extendfrac) .== nothing)
    @assert any([slicesx !=  [], slicesy != [], slicesz !=[]])
    slicemodel(m.fstar, m.n, m.xtrain, m.ftrain, opt::TransD_GP.Options; slicesx=slicesx, slicesy=slicesy, slicesz=slicesz,
                dz=dz, extendfrac=extendfrac, figsize=figsize)
end

function slicemodel(fstar::Array{Float64}, npoints::Int,
                    xtrain::Array{Float64, 2}, ftrain::Array{Float64},
                    opt::TransD_GP.Options;
                    slicesx    = [],
                    slicesy    = [],
                    slicesz    = [],
                    dz         = nothing,
                    extendfrac = nothing,
                    figsize    = (7,9)
                    )
    x = unique(opt.xall[1,:])
    y = unique(opt.xall[2,:])
    z = geomprogdepth.(unique(opt.xall[3,:]), dz, extendfrac)
    @assert !any((dz, extendfrac) .== nothing)
    @assert any([slicesx !=  [], slicesy != [], slicesz !=[]])
    nplots = sum([slicesx != [], slicesy != [], slicesz != []])
    v = reshape(fstar,length(x), length(y), length(z))
    fig = figure(figsize=(nplots*figsize[1],figsize[2]))
    using3D()
    ax = Array{Any,1}(undef, nplots)
    iplot = 0
    if slicesz!=[]
        xx,yy = meshgrid(x,y)
        l = slicesz
        iplot+=1
        ax[iplot]=fig.add_subplot(1, nplots, iplot, projection="3d")
        for i in l
            c=(v[:,:,i].-minimum(v))./(maximum(v)-minimum(v))
            ax[iplot].plot_surface(xx, yy, z[i]*ones(size(yy)), facecolors=plt.cm.jet_r(c), shade=false)
        end
    end
    if slicesx!=[]
        yy,zz = meshgrid(y,z)
        l = slicesx
        iplot+=1
        ax[iplot]=fig.add_subplot(1, nplots, iplot, projection="3d")
        for i in l
            c=(v[i,:,:].-minimum(v))./(maximum(v)-minimum(v))
            ax[iplot].plot_surface(x[i]*ones(size(yy)), yy, zz, facecolors=plt.cm.jet_r(c), shade=false)
        end
    end
   if slicesy!=[]
        xx,zz = meshgrid(x,z)
        l = slicesy
        iplot+=1
        ax[iplot]=fig.add_subplot(1, nplots, iplot, projection="3d")
        for i in l
            c=(v[:,i,:].-minimum(v))./(maximum(v)-minimum(v))
            ax[iplot].plot_surface(xx, y[i]*ones(size(zz)), zz, facecolors=plt.cm.jet_r(c), shade=false)
        end
    end
    for a in ax
        if npoints>0
            a.scatter3D(xtrain[1,1:npoints], xtrain[2,1:npoints], geomprogdepth.(xtrain[3,1:npoints], dz, extendfrac),
                c=ftrain[1:npoints], vmin=minimum(v), vmax=maximum(v), s=50, cmap="jet_r")
        end
        a.set_xlim(extrema(x))
        a.set_ylim(extrema(y))
        a.set_zlim(extrema(z))
        a.set_xlabel("x km")
        a.set_ylabel("y km")
        a.set_zlabel("depth km")
        a.invert_zaxis()
    end
    nicenup(gcf(),fsize=14)
    fig.subplots_adjust(right=0.9)
    cbar_ax = fig.add_axes([0.92, 0.15, 0.025, 0.7])
    cbar = fig.colorbar(plt.cm.ScalarMappable(cmap="jet_r",norm=matplotlib.colors.Normalize(vmin=minimum(v), vmax=maximum(v))), cax=cbar_ax)
    cbar.set_label(L"\log_{10} \rho", fontsize=14)
    cbar.ax.tick_params(labelsize=14)
end

function makecubemodel(opt::TransD_GP.Options;
                        dz         = nothing,
                        extendfrac = nothing,
                        z1         = 10.0,
                        z2         = 400.0,
                        yleft      = 500.0,
                        yright     = 1100.0,
                        xup        = 980.0,
                        xdown      = 580.0,
                        zup        = 70.0,
                        zdown      = 250.0,
                        ρ0         = 2,
                        ρ1         = 4,
                        ρ2         = 0.5,
                        ρanom      = 0.8,
                        ntosmooth  = 1,
                        )
    @assert extendfrac != nothing
    @assert dz != nothing
    @assert z1<zup<zdown<z2
    @assert yleft<yright
    @assert xdown<xup

    ρ = zeros(size(opt.xall, 2))
    zmin, zmax = geomprogdepth.(extrema(opt.xall[3,:]), dz, extendfrac)
    for i = 1:size(opt.xall, 2)
        x, y, zn = opt.xall[:,i]
        z = geomprogdepth(zn, dz, extendfrac)
        if z<z1
            ρ[i] = ρ0
        elseif (zup<z<zdown) & (yleft<y<yright) & (xdown<x<xup)
            ρ[i] = ρanom*(ρ1 + (z-zmin)/(zmax-z1)*(ρ2-ρ1)) #ρanom
        elseif z>z2
            ρ[i] = ρ2
        else
            ρ[i] = ρ1 + (z-zmin)/(zmax-z1)*(ρ2-ρ1)
        end
    end
    nx, ny, nz = length(unique(opt.xall[1,:])), length(unique(opt.xall[2,:])), length(unique(opt.xall[3,:]))
    ρ .= boxcarn(reshape(ρ, nx, ny, nz), ntosmooth)[:]
    opt.fbounds = [0.98*minimum(ρ) 1.02*maximum(ρ)]
    return ρ
end

function get_training_data(fstar::Array{Float64, 1},
                           opt::TransD_GP.Options;
                           dz = nothing,
                           extendfrac = nothing,
                           sdmaxfrac = 0.05,
                           rseed     = 12,
                           zbreak    = 12321.0,
                           takeevery = 4,
                           fractrain = 0.05
                          )
    @assert extendfrac != nothing
    @assert dz != nothing
    @assert zbreak != 12321.0
    @assert sdmaxfrac > 0 && sdmaxfrac < 1
    Random.seed!(rseed)
    x, y =  unique(opt.xall[1,:]), unique(opt.xall[2,:])
    z = geomprogdepth.(unique(opt.xall[3,:]), dz, extendfrac)
    @assert minimum(z) < zbreak < maximum(z)
    noisyd = NaN .+ zeros(Float64, size(fstar))
    ntrain = round(Int, fractrain*length(fstar))
    linidx = randperm(length(fstar))[1:ntrain]
    δtry = sdmaxfrac*maximum(abs.(fstar))
    for l in linidx
        row, col, zid = Tuple(CartesianIndices((length(x),length(y),length(z)))[l])
        if z[zid] > zbreak || (rem(row, takeevery) == 0 && rem(col, takeevery) == 0 && rem(zid, takeevery) == 0)
            noisyd[l] = fstar[l] + δtry*randn()
        end
    end
    f = figure(figsize=(10,10))
    scatter3D(opt.xall[1,:], opt.xall[2,:], geomprogdepth.(opt.xall[3,:], dz, extendfrac),
                          c=noisyd, vmin=minimum(fstar), vmax=maximum(fstar), cmap="jet_r")
    xlim(extrema(x))
    ylim(extrema(y))
    zlim(extrema(z))
    xlabel("x km")
    ylabel("y km")
    zlabel("depth km")
    cbar = f.colorbar(plt.cm.ScalarMappable(cmap="jet_r",norm=matplotlib.colors.Normalize(vmin=minimum(fstar), vmax=maximum(fstar))),ax=gca())
    cbar.set_label(L"\log_{10} \rho")
    gca().zaxis.label.set_fontsize(14); nicenup(gcf())
    gca().invert_zaxis()

    noisyd, δtry
end

function plot_last_target_model(opt_in::TransD_GP.Options, d::Array{Float64, 1}, sd::Float64;
                                slicesx    = [],
                                slicesy    = [],
                                slicesz    = [],
                                dz         = nothing,
                                extendfrac = nothing,
                                fsize      = 14
                               )
    @assert !any((dz, extendfrac) .== nothing)
    @assert any([slicesx !=  [], slicesy != [], slicesz !=[]])
    Tacrosschains = gettargtemps(opt_in)
    last_target_model_idx = findall(abs.(Tacrosschains[end,:] .-1.0) .< 1e-12)
    for idx in last_target_model_idx
        opt_in.fstar_filename = "models_"*opt_in.fdataname*"_$idx.bin"
        m_last = TransD_GP.history(opt_in, stat=:fstar)[end]
        slicemodel(m_last, 0,[0. 0.],[0.], opt_in, slicesx=slicesx, slicesy=slicesy, slicesz=slicesz, dz=dz, extendfrac=extendfrac)
        sd > 0.0 && calc_simple_RMS(d, m_last, sd)
    end

end

function gettargtemps(opt_in::TransD_GP.Options)
    nchains = length(filter( x -> occursin(r"misfits.*bin", x), readdir(pwd()) )) # my terrible regex
    @info "Number of chains is $nchains"
    # now look at any chain to get how many iterations
    costs_filename = "misfits_"*opt_in.fdataname
    opt_in.costs_filename    = costs_filename*"_1.bin"
    iters          = TransD_GP.history(opt_in, stat=:iter)
    niters         = length(iters)
    @info "McMC has run for $(iters[end]) iterations"
    # then create arrays of unsorted by temperature T
    Tacrosschains  = zeros(Float64, niters, nchains)
    # get the values into the arrays
    for ichain in 1:nchains
        opt_in.costs_filename = costs_filename*"_$ichain.bin"
        Tacrosschains[:,ichain] = TransD_GP.history(opt_in, stat=:T)
    end
    Tacrosschains
end

function assembleTat1(opt::TransD_GP.Options, irange::StepRange{Int})
    Tacrosschains = gettargtemps(opt)
    iters = TransD_GP.history(opt, stat=:iter)
    @assert 0<irange.start<=length(iters)
    @assert irange.stop<=length(iters)
    @info "obtaining models $(iters[irange.start]) to $(iters[irange.stop])"
    nat1 = length(findall(abs.(Tacrosschains[end,:] .-1.0) .< 1e-12))
    @info "No. of chains at 1 is $nat1"
    niters = length(irange)
    mat1 = zeros(Float64, size(opt.xall,2), niters, nat1)
    cumT = cumsum(Tacrosschains .== 1,dims=1)
    @info size(mat1)
    for (istep, stp) in enumerate(irange)
        @info "$(istep/niters*100)% done"
        for (itemp, idx) in enumerate(findall(abs.(Tacrosschains[stp,:] .-1.0) .< 1e-12))
            opt.fstar_filename = "models_"*opt.fdataname*"_$idx.bin"
            mat1[:,istep,itemp] = TransD_GP.history(opt, stat=:fstar)[cumT[stp,idx]]
        end
    end
    mat1
end

function assembleTat1(opt::TransD_GP.Options; burninfrac=0.5)
    @assert 0.0<1.0burninfrac<1.0
    Tacrosschains = gettargtemps(opt)
    iters = TransD_GP.history(opt, stat=:iter)
    start = round(Int, length(iters)*burninfrac)
    @info "obtaining models $(iters[start]) to $(iters[end])"
    nmodels = sum((Tacrosschains[start:end,:] .== 1))
    # storing models column wise in nmodels rows
    mat1 = zeros(Float64, size(opt.xall,2), nmodels)
    imodel = 0
    for ichain in 1:size(Tacrosschains, 2)
        @info "chain $ichain"
        at1idx = findall(Tacrosschains[:,ichain].==1).>= start
        ninchain = sum(at1idx)
        ninchain == 0 && continue
        opt.fstar_filename = "models_"*opt.fdataname*"_$ichain.bin"
        mat1[:,imodel+1:imodel+ninchain] .= reduce(hcat, TransD_GP.history(opt, stat=:fstar)[at1idx])
        imodel += ninchain
    end
    mat1
end

function savepngs(opt::TransD_GP.Options, irange::StepRange{Int};
                  slicesx    = [],
                  slicesy    = [],
                  slicesz    = [],
                  dz         = nothing,
                  extendfrac = nothing,
                  fsize      = 14
                 )
    @assert !any((dz, extendfrac) .== nothing)
    @assert any([slicesx !=  [], slicesy != [], slicesz !=[]])
    Tacrosschains = gettargtemps(opt)
    iters = TransD_GP.history(opt, stat=:iter)
    @assert 0<irange.start<=length(iters)
    @assert irange.stop<=length(iters)
    nat1 = length(findall(abs.(Tacrosschains[end,:] .-1.0) .< 1e-12))
    @info "No. of chains at 1 is $nat1"
    @info "obtaining models $(iters[irange.start]) to $(iters[irange.stop])"
    M = assembleTat1(opt,irange)
    for (iridx, ir) in enumerate(irange), ic in 1:nat1
        Tools3D.slicemodel(M[:,iridx,ic],0,[0. 0.],[0.], opt,
                            slicesx=slicesx, slicesy=slicesy, slicesz=slicesz,
                            dz=dz, extendfrac=extendfrac)
        @info "saving target model $(iters[ir]) chain $ic"
        gcf().suptitle("model_"*lpad(iters[ir],7,'0')*"_chain_"*lpad(ic,3,'0'))
        savefig("model_"*lpad(iters[ir],7,'0')*"_chain_"*lpad(ic,3,'0')*".png", dpi=300)
        close()
    end
    nothing
end

function calc_simple_RMS(d::AbstractArray, fstar::Array{Float64, 1}, sd::Float64)

    select = .!isnan.(d)
    r = (d[select] - fstar[select])/sd
    n = sum(select)
    @info "χ^2 error is $(r'*r) for $n points RMS: $(sqrt(r'*r/n))"
    nothing
end

function boxcarn(A::AbstractArray, ntosmooth::Int)
    out = zeros(eltype(A), size(A))
    R = CartesianIndices(A)
    Ifirst, Ilast = first(R), last(R)
    I1 = oneunit(Ifirst)
    for I in R
        n, s = 0, zero(eltype(out))
        for J in max(Ifirst, I-ntosmooth*I1):min(Ilast, I+ntosmooth*I1)
            s += A[J]
            n += 1
        end
        out[I] = s/n
    end
    out
end

end


# example
# numcars = 5;
# x = collect(range(0,10,length = 100));
# ys = randboxcars(x, numcars)
# boxwidthstd = 0.4 unit of x's length
function randboxcars(x::Array{T,1}, numcars::Int; boxwidthstd = 0.4) where {T} #
	halfwidthstd = boxwidthstd*0.5;
	boxcentr = rand(x, numcars);
	halfboxwidth = halfwidthstd .* (x[end]-x[1]) .* abs.(randn(numcars));
	ys = fill(0, length(x));
	for i = 1:numcars
		lb, centr, rb = get_rect(x,boxcentr,halfboxwidth,i);
		ys[lb .< x .< rb] .= ys[lb .< x .< rb] .+ 1;
	end
	return ys
end
# number of rows of ys has to be the same as length(x).
function randboxcars(x::Array{T,1}, ys::Array{N,2}; boxwidthstd = 0.4) where {T,N}
	halfwidthstd = boxwidthstd*0.5;
	numcars = size(ys, 2);
	boxcentr = rand(x, numcars);
	halfboxwidth = halfwidthstd .* (x[end]-x[1]) .* abs.(randn(numcars));
	ys = fill(0, (length(x), numcars));
	for i = 1:numcars
		lb, centr, rb = get_rect(x,boxcentr,halfboxwidth,i);
		ys[lb .< x .< rb,i] .= 1;
	end
	return ys
end

function get_rect(x,boxcentr,halfboxwidth,i)
	centr = findlast(x .<= boxcentr[i]);
	rb = x[centr] + halfboxwidth[i]; # right bound
	lb = x[centr] - halfboxwidth[i]; # left bound
	return lb, centr, rb
end


using Nonconvex, LinearAlgebra, Test

f(x::AbstractVector) = sqrt(x[2])
g(x::AbstractVector, a, b) = (a*x[1] + b)^3 - x[2]

options = NLoptOptions(xtol_rel = 1e-4)

@testset "Simple constraints" begin
    m = Model(f)
    addvar!(m, [0.0, 0.0], [10.0, 10.0])
    add_ineq_constraint!(m, x -> g(x, 2, 0))
    add_ineq_constraint!(m, x -> g(x, -1, 1))

    alg = NLoptAlg(:LD_MMA)
    r = Nonconvex.optimize(m, alg, [1.234, 2.345], options = options)
    @test abs(r.minimum - sqrt(8/27)) < 1e-6
    @test norm(r.minimizer - [1/3, 8/27]) < 1e-6
end

@testset "Block constraints" begin
    m = Model(f)
    addvar!(m, [0.0, 0.0], [10.0, 10.0])
    add_ineq_constraint!(m, FunctionWrapper(x -> [g(x, 2, 0), g(x, -1, 1)], 2))

    alg = NLoptAlg(:LD_MMA)
    r = Nonconvex.optimize(m, alg, [1.234, 2.345], options = options)
    @test abs(r.minimum - sqrt(8/27)) < 1e-6
    @test norm(r.minimizer - [1/3, 8/27]) < 1e-6
end

@testset "Infinite bounds" begin
    @testset "Infinite upper bound" begin
        m = Model(f)
        addvar!(m, [0.0, 0.0], [Inf, Inf])
        add_ineq_constraint!(m, x -> g(x, 2, 0))
        add_ineq_constraint!(m, x -> g(x, -1, 1))

        alg = NLoptAlg(:LD_MMA)
        r = Nonconvex.optimize(m, alg, [1.234, 2.345], options = options)
        @test abs(r.minimum - sqrt(8/27)) < 1e-6
        @test norm(r.minimizer - [1/3, 8/27]) < 1e-6
    end
    @testset "Infinite lower bound" begin
        m = Model(f)
        addvar!(m, [-Inf, -Inf], [10, 10])
        add_ineq_constraint!(m, x -> g(x, 2, 0))
        add_ineq_constraint!(m, x -> g(x, -1, 1))

        alg = NLoptAlg(:LD_MMA)
        r = Nonconvex.optimize(m, alg, [1.234, 2.345], options = options)
        @test abs(r.minimum - sqrt(8/27)) < 1e-6
        @test norm(r.minimizer - [1/3, 8/27]) < 1e-6
    end
    @testset "Infinite upper and lower bound" begin
        m = Model(f)
        addvar!(m, [-Inf, -Inf], [Inf, Inf])
        add_ineq_constraint!(m, x -> g(x, 2, 0))
        add_ineq_constraint!(m, x -> g(x, -1, 1))

        alg = NLoptAlg(:LD_MMA)
        r = Nonconvex.optimize(m, alg, [1.234, 2.345], options = options)
        @test abs(r.minimum - sqrt(8/27)) < 1e-6
        @test norm(r.minimizer - [1/3, 8/27]) < 1e-6
    end
end