"""
    JopMix(s::JetAbstractSpace{T,N}, nmix::NTuple{N,Int})

2D spatial mix, and is used to add smoothness to a domain.  JopMix is ported
from the CVX frequency domain FWI tools in SeisSpace.

# Notes:
* This appears to over-lap in funtionality with JopRoughness.

# Examples:
## 2D
```julia
nz, nx = 11,11
spc = JetSpace(Float32, nz, nx)
A = JopMix(spc, (5,5))
y = A * rand(dom)
x = A' * rand(rng)
```
"""
function JopMix(spc::JetAbstractSpace{T,N}, nmix::NTuple{N,Int}) where {T,N}
    for i in nmix
        i < 0 && error("mix lengths must be >= 0 -- $(nmix)")
    end
    JopLn(dom = spc, rng = spc, df! = JopMix_df!, df′! = JopMix_df′!, s= (nmix=nmix,))
end

export JopMix

function JopMix_df!(rngvec::AbstractArray{T,2}, domvec::AbstractArray{T,2}; nmix, kwargs...) where {T}
    rngvec .= 0
    nz,nx = size(domvec)
    lenz = div(nmix[1], 2)
    lenx = div(nmix[2], 2)
    wgt = (2 * lenz + 1) * (2 * lenx + 1)

    for kz = 1:nz
        jz1 = clamp(kz - lenz, 1, nz)
        jz2 = clamp(kz + lenz, 1, nz)
        for kx = 1:nx
            jx1 = clamp(kx - lenx, 1, nx)
            jx2 = clamp(kx + lenx, 1, nx)
            for jx = jx1:jx2
                for jz = jz1:jz2
                    rngvec[kz,kx] += domvec[jz,jx]
                end
            end
        end
    end
    rngvec ./= wgt
    rngvec
end

function JopMix_df′!(domvec::AbstractArray{T,2}, rngvec::AbstractArray{T,2}; nmix, kwargs...) where {T}
    domvec .= 0
    nz,nx = size(domvec)
    lenz = div(nmix[1], 2)
    lenx = div(nmix[2], 2)
    wgt = (2 * lenz + 1) * (2 * lenx + 1)

    for kz = 1:nz
        jz1 = clamp(kz - lenz, 1, nz)
        jz2 = clamp(kz + lenz, 1, nz)
        for kx = 1:nx
            jx1 = clamp(kx - lenx, 1, nx)
            jx2 = clamp(kx + lenx, 1, nx)
            for jx = jx1:jx2
                for jz = jz1:jz2
                    domvec[kz,kx] += rngvec[jz,jx]
                end
            end
        end
    end
    domvec ./= wgt
    domvec
end


import Dolo
import PyPlot
path = Dolo.pkg_path

filename = joinpath(path,"examples","models","rbc_dtcc_iid_ar1.yaml")
model = Dolo.yaml_import(filename)
@time dr = Dolo.time_iteration(model, verbose=true, maxit=10000)

n_steps=100
s0 = model.calibration[:states]
m0 = model.calibration[:exogenous]
index = findfirst(isequal(:z), model.symbols[:states])
bounds = [dr.grid_endo.min[index], dr.grid_endo.max[index]]

df = Dolo.tabulate(model, dr, :z, bounds, s0, m0)
df = Dolo.tabulate(model, dr, :z, s0, m0)

cn =[:i, :n]
cn[1]
Svalues = linspace(bounds[1], bounds[2], n_steps)
PyPlot.plot(Svalues, df[cn[2]])

for j in cn
  fig = PyPlot.figure(j,figsize=(3,3))
  fig
  PyPlot.plot(Svalues, df[j], label=j)
  PyPlot.legend()
  # PyPlot.xlabel('state = {} | mstate = {}'.format(state, i0))
end

plot_dr = Dolo.tabulate(model, dr, :k, bounds, s0, m0, cn)


"""
    brute_force(fla::Formula)::SatSolverResult

Tries all possible variable assignments and stops at the first satisfying assignment.
"""
function brute_force(fla::Formula)::SatSolverResult
  pa = PartialAssignment(fill(-1, fla.variable_count))

  while true
    if all((cl) -> is_clause_satisfied(cl, pa), fla)
      return SatisfiableFormula(pa)
    end

    last_negative = findlast(pa.assignment) do x
      -1 == x
    end

    if nothing == last_negative
      break
    end

    pa[last_negative] = 1
    for i in (last_negative + 1):fla.variable_count
      pa[i] = -1
    end
  end

  UnsatisfiableFormula()
end


using FMMTrees
#using Plots; plotlyjs()

p = rand(40)
p, tp = clustertree(p)

q = rand(40)
q, tq = clustertree(q)

depthfirst(tp) do b,l
    println(l, ": ", b[1].begin_idx:b[1].end_idx-1)
end

depthfirst(tq) do b,l
    println(l, ": ", b[1].begin_idx:b[1].end_idx-1)
end

@assert first(children((tp,tq)))[1] isa FMMTrees.Tree
@assert first(children((tp,tq)))[2] isa FMMTrees.Tree

S = zeros(length(p), length(p))
depthfirst((tp,tp)) do b,l
    I = b[1][1].begin_idx : b[1][1].end_idx-1
    J = b[2][1].begin_idx : b[2][1].end_idx-1
    println(l, ": ", I, "×", J)
    l == 5 && (S[I,J] = rand(1:20))
end

# Visual representation of the block partition at level 3
# heatmap(S)
nmin = 20
η = 1.5
function adm(b)
    I = b[1][1].begin_idx : b[1][1].end_idx-1
    J = b[2][1].begin_idx : b[2][1].end_idx-1
    length(I) < nmin && return true
    length(J) < nmin && return true
    ll1, ur1 = FMMTrees.boundingbox(p[I]); c1 = (ll1+ur1)/2;
    ll2, ur2 = FMMTrees.boundingbox(p[J]); c2 = (ll2+ur2)/2;
    diam1 = norm(ur1-c1)
    diam2 = norm(ur2-c2)
    dist12 = norm(c2-c1)
    return dist12 >= η*max(diam1, diam2)
end

P = admissable_partition((tp,tp), adm)
S = zeros(length(p), length(p))
for b in P
    I = b[1][1].begin_idx : b[1][1].end_idx-1
    J = b[2][1].begin_idx : b[2][1].end_idx-1
    S[I,J] += length(I) + length(J)
end

extrema(S)

# using Plots; plotlyjs();
# heatmap(S)

find(P) do b
    I = b[1][1].begin_idx : b[1][1].end_idx-1
    J = b[2][1].begin_idx : b[2][1].end_idx-1
    length(I) + length(J) == 40
end


import HCubature
export HCubatureAlg

struct HCubatureAlg{K}
    kw::K
end

function HCubatureAlg(;kw...)
    HCubatureAlg(kw)
end

function integral_kernel(f, dom::Domain, alg::HCubatureAlg)
    value, err = HCubature.hcubature(f, dom.lower, dom.upper; alg.kw...)
    (value=value, std=err) # TODO err is not really a standard deviation, what should be done?
end


for N in [Float64, Rational{Int}, Float32]
    # Sum of 2D centered balls in norm 1 and infinity
    b1 = BallInf(N[0, 0], N(2))
    b2 = Ball1(N[0, 0], N(1))
    # Test Construction
    X = MinkowskiSum(b1, b2)
    @test X.X == b1
    @test X.Y == b2

    # swap
    ms2 = swap(X)
    @test X.X == ms2.Y && X.Y == ms2.X

    # Test Dimension
    @test dim(X) == 2
    # Test Support Vector
    d = N[1, 0]
    v = σ(d, X)
    @test v[1] == N(3)
    d = N[-1, 0]
    v = σ(d, X)
    @test v[1] == N(-3)
    d = N[0, 1]
    v = σ(d, X)
    @test v[2] == N(3)
    d = N[0, -1]
    v = σ(d, X)
    @test v[2] == N(-3)

    # Sum of not-centered 2D balls in norm 1 and infinity
    b1 = BallInf(N[-1, 3], N(2))
    b2 = Ball1(N[1, 2], N(1))
    s = b1 + b2
    # Test Support Vector
    d = N[1, 0]
    v = σ(d, s)
    @test v[1] == N(3)
    d = N[-1, 0]
    v = σ(d, s)
    @test v[1] == N(-3)
    d = N[0, 1]
    v = σ(d, s)
    @test v[2] == N(8)
    d = N[0, -1]
    v = σ(d, s)
    @test v[2] == N(2)

    # Sum of array of LazySet
    # 2-elements
    ms = MinkowskiSum(Singleton(N[1]), Singleton(N[2]))
    @test ρ(N[1], ms) == 3
    @test ρ(N[-1], ms) == -3
    # 3-elements
    ms = MinkowskiSum(Singleton(N[1]), MinkowskiSum(Singleton(N[2]), Singleton(N[3])))
    @test ρ(N[1], ms) == N(6)
    @test ρ(N[-1], ms) == N(-6)

    # boundedness
    @test isbounded(ms)
    @test !isbounded(Singleton(N[1]) + HalfSpace(N[1], N(1)))

    # isempty
    @test !isempty(ms)

    # an_element function (falls back to the default implementation
    X = MinkowskiSum(LineSegment(N[0, 0], N[1, 1]),
                     LineSegment(N[1, 0], N[0, 1]))
    v = an_element(X)
    @test v ∈ Ball1(N[1, 1], N(1))

    # constraints list
    ms = MinkowskiSum(BallInf(N[0, 0], N(1)), BallInf(N[0, 0], N(1)))
    clist = constraints_list(ms)
    P = HPolytope(clist)
    H = Hyperrectangle(N[0, 0], N[2, 2])
    @test length(clist) == 4 && P ⊆ H && H ⊆ P

    # =================
    # MinkowskiSumArray
    # =================

    # relation to base type (internal helper functions)
    @test LazySets.array_constructor(MinkowskiSum) == MinkowskiSumArray
    @test LazySets.is_array_constructor(MinkowskiSumArray)

    # array getter
    v = Vector{LazySet{N}}()
    @test array(MinkowskiSumArray(v)) ≡ v

    # constructor with size hint and type
    MinkowskiSumArray(10, N)

    # in-place modification
    msa = MinkowskiSumArray(LazySet{N}[])
    @test MinkowskiSum!(b1, b1) isa MinkowskiSum && length(array(msa)) == 0
    res = MinkowskiSum!(b1, msa)
    @test res isa MinkowskiSumArray && length(array(msa)) == 1
    res = MinkowskiSum!(msa, b1)
    @test res isa MinkowskiSumArray && length(array(msa)) == 2
    res = MinkowskiSum!(msa, msa)
    @test res isa MinkowskiSumArray && length(array(msa)) == 4

    ms = MinkowskiSum(b1, b2)
    msa = MinkowskiSumArray([b1, b2])

    # dimension
    @test dim(msa) == 2

    # support vector
    d = N[1, 1]
    @test σ(d, ms) == σ(d, msa)
    d = N[-1, 1]
    @test σ(d, ms) == σ(d, msa)

    # boundedness
    @test isbounded(msa)
    @test !isbounded(MinkowskiSumArray([Singleton(N[1]), HalfSpace(N[1], N(1))]))

    # isempty
    @test !isempty(msa)

    # =================
    # CacheMinkowskiSum
    # =================

    # caching Minkowski sum
    cms = CacheMinkowskiSum(2, N)
    x1 = BallInf(ones(N, 3), N(3))
    x2 = Ball1(ones(N, 3), N(5))
    d = ones(N, 3)
    a = array(cms)
    push!(a, x1)
    σ(d, cms)
    push!(a, x2)
    svec = σ(d, cms)
    @test σ(d, cms) == svec
    @test dim(cms) == 3
    idx = forget_sets!(cms)
    @test idx == 2 && isempty(array(cms)) && cms.cache[d][1] == 0
    push!(a, x1)
    σ(d, cms)
    push!(a, x2)
    idx = forget_sets!(cms)
    @test idx == 1 && length(array(cms)) == 1 && cms.cache[d][1] == 0
    # test issue #367: inplace modification of the direction modified the cache
    d[1] = zero(N)
    @test haskey(cms.cache, ones(N, 3))
    # getindex
    cp = LazySets.CachedPair(1, N[2])
    @test cp[1] == 1 && cp[2] == N[2]
    @test_throws ErrorException cp[3]

    # boundedness
    @test isbounded(cms)
    @test !isbounded(CacheMinkowskiSum([Singleton(N[1]), HalfSpace(N[1], N(1))]))

    # isempty
    @test !isempty(cms)

    # ================
    # common functions
    # ================

    # neutral element
    z = ZeroSet{N}(2)
    msa = MinkowskiSumArray(LazySet{N}[])
    cms = CacheMinkowskiSum(LazySet{N}[])
    @test neutral(MinkowskiSum) == neutral(MinkowskiSumArray) ==
          neutral(CacheMinkowskiSum) == ZeroSet
    @test b1 + z == z + b1 == b1
    @test msa + z == z + msa == msa
    @test cms + z == z + cms == cms
    # absorbing element
    e = EmptySet{N}()
    @test absorbing(MinkowskiSum) == absorbing(MinkowskiSumArray) ==
          absorbing(CacheMinkowskiSum) == EmptySet
    @test b1 + e == e + b1 == msa + e == e + msa == cms + e == e + cms ==
          e + e == e
    # mix of neutral and absorbing element
    @test z + e == e + z == e
end


abstract type AbstractVecArray{T<:AbstractVec,N} <: AbstractArray{T,N} end

@inline Base.size(u::AbstractVecArray) =
    size(u.x)

@inline Base.length(u::AbstractVecArray) =
    length(u.x)

@inline Base.linearindices(u::AbstractVecArray) =
    linearindices(u.x)

Base.IndexStyle(::Type{<:AbstractVecArray}) =
    IndexLinear()

struct VecArray{T<:Number,N} <: AbstractVecArray{Vec{T},N}
    x::Array{T,N}
    y::Array{T,N}
    z::Array{T,N}

    function VecArray{T,N}(x,y,z) where {T,N}
        (size(x) == size(y) == size(z)) || throw(ArgumentError("Arrays must have the same dimensions"))
        return new{T,N}(x,y,z)
    end

end

VecArray(x,y,z) =
    VecArray{eltype(x),ndims(x)}(x,y,z)

@inline function Base.getindex(u::VecArray{T,N},I) where {T,N}
    x = u.x
    y = u.y
    z = u.z
    @boundscheck checkbounds(x,I)
    @inbounds begin
        xv = x[I]
        yv = y[I]
        zv = z[I]
    end
    return Vec{T}(xv,yv,zv)
end

@inline function Base.setindex!(u::VecArray{T,N},v::Vec,I) where {T,N}
    x = u.x
    y = u.y
    z = u.z
    @boundscheck checkbounds(x,I)
    @inbounds begin
        x[I] = v.x
        y[I] = v.y
        z[I] = v.z
    end
    return u
end

struct Vec2DArray{T<:Number,N} <: AbstractVecArray{Vec2D{T},N}
    x::Array{T,N}
    y::Array{T,N}

    function Vec2DArray{T,N}(x,y) where {T,N}
        (size(x) == size(y)) || throw(ArgumentError("Arrays must have the same dimensions"))
        return new{T,N}(x,y)
    end

end

Vec2DArray(x::AbstractArray,y) =
    Vec2DArray{eltype(x),ndims(x)}(x,y)

Vec2DArray{T}(L::Integer) where T = Vec2DArray{T,1}(zeros(T,L),zeros(T,L))

Base.similar(a::Vec2DArray{T}) where {T} = Vec2DArray{T}(length(a))

@inline function Base.getindex(u::Vec2DArray{T,N},I) where {T,N}
    x = u.x
    y = u.y
    @boundscheck checkbounds(x,I)
    @inbounds begin
        xv = x[I]
        yv = y[I]
    end
    return Vec2D{T}(xv,yv)
end

@inline function Base.setindex!(u::Vec2DArray{T,N},v::Vec2D,I) where {T,N}
    x = u.x
    y = u.y
    @boundscheck checkbounds(x,I)
    @inbounds begin
        x[I] = v.x
        y[I] = v.y
    end
    return u
end

struct Vec1DArray{T<:Number,N} <: AbstractVecArray{Vec1D{T},N}
    x::Array{T,N}

    function Vec1DArray{T,N}(x) where {T,N}
        return new{T,N}(x)
    end

end

Vec1DArray(x::AbstractArray) =
    Vec1DArray{eltype(x),ndims(x)}(x)

@inline function Base.getindex(u::Vec1DArray{T,N},I) where {T,N}
    x = u.x
    @boundscheck checkbounds(x,I)
    @inbounds begin
        xv = x[I]
    end
    return Vec1D{T}(xv)
end

@inline function Base.setindex!(u::Vec1DArray{T,N},v::Vec1D,I) where {T,N}
    x = u.x
    @boundscheck checkbounds(x,I)
    @inbounds begin
        x[I] = v.x
    end
    return u
end


using LinearAlgebraicRepresentation, ViewerGL, SparseArrays
Lar = LinearAlgebraicRepresentation; GL = ViewerGL

using Revise, OhMyREPL

wall = Lar.svg2lar("test/svg/assembly/wall.svg", flag=false)
openings = Lar.svg2lar("test/svg/assembly/openings.svg", flag=false)
rectangle = Lar.svg2lar("test/svg/assembly/rectangle.svg", flag=false)
box = Lar.svg2lar("test/svg/assembly/box.svg", flag=false)

assembly = Lar.Struct([wall , openings, rectangle, box])
V,EV = Lar.struct2lar(assembly)
GL.VIEW([ GL.GLFrame2, GL.GLGrid(V,EV, GL.COLORS[1],1) ]);
W, copEV, copFE, boolmatrix = Lar.bool2d(assembly)
@show Matrix(boolmatrix)

innerpoints = Lar.internalpoints2d(W,copEV,copFE[1:end,:])
points = convert(Array{Float64,2}, hcat(innerpoints...)')
GL.VIEW([ GL.GLFrame2, GL.GLGrid(V,EV, GL.COLORS[1],1), GL.GLPoints(points) ]);

A = boolmatrix[:,1]
B = boolmatrix[:,2]
C = boolmatrix[:,3]
D = boolmatrix[:,4]
expr = ( D .& .!(.&(A, .!B, .!C)) )# A - B - C
formula = Matrix(copFE)' * Int.(expr) # coord vector of Faces
V = convert(Lar.Points,W')
EV = Lar.cop2lar(copEV)
EVmin = [ev for (k,ev) in enumerate(EV) if abs(formula[k])==1 ]
V = Lar.normalize(V)
V = Lar.normalize(V)
GL.VIEW([ GL.GLFrame2, GL.GLGrid(V,EVmin, GL.COLORS[1],1) ]);

VV = [[k] for k in 1:size(V,2)];
model = (V,[VV, EV])
GL.VIEW( GL.numbering(.1)( model,GL.COLORS[1],0.1 ) );


using LinearAlgebra, OrdinaryDiffEq, Test
f = (du,u,p,t) -> du .= u ./ t
jac = (J,u,p,t) -> (J[1,1] = 1/t; J[2,2] = 1/t; J)

jp_diag = Diagonal(zeros(2))
fun = ODEFunction(f; jac=jac, jac_prototype=jp_diag)
prob = ODEProblem(fun,ones(2),(1.0,10.0))
sol = solve(prob,Rosenbrock23(autodiff=false))
@test sol[end] ≈ [10.0,10.0]
@test length(sol) < 60

jp = Tridiagonal(jp_diag)
fun = ODEFunction(f; jac=jac, jac_prototype=jp)
prob = ODEProblem(fun,ones(2),(1.0,10.0))
sol = solve(prob,Rosenbrock23(autodiff=false))
@test sol[end] ≈ [10.0,10.0]
@test length(sol) < 60

jp = SymTridiagonal(jp_diag)
fun = ODEFunction(f; jac=jac, jac_prototype=jp)
prob = ODEProblem(fun,ones(2),(1.0,10.0))
sol = solve(prob,Rosenbrock23(autodiff=false))
@test sol[end] ≈ [10.0,10.0]
@test length(sol) < 60

@test_broken begin
    jp = Hermitian(jp_diag)
    fun = ODEFunction(f; jac=jac, jac_prototype=jp)
    prob = ODEProblem(fun,ones(2),(1.0,10.0))
    sol = solve(prob,Rosenbrock23(autodiff=false))
    @test sol[end] ≈ [10.0,10.0]
    @test length(sol) < 60
end

@test_broken begin
    jp = Symmetric(jp_diag)
    fun = ODEFunction(f; jac=jac, jac_prototype=jp)
    prob = ODEProblem(fun,ones(2),(1.0,10.0))
    sol = solve(prob,Rosenbrock23(autodiff=false))
    @test sol[end] ≈ [10.0,10.0]
    @test length(sol) < 60
end


using DelaySSAToolkit
using Catalyst
# using Plots
################################################################################
# the following example illustrates how to deal with a oscillatory system
# d: X-> 0; k1/(1+Y^2): 0-> X; [trigger X-> Y after τ time;] k2*Y/(1+Y):  Y-> 0;

rn = @reaction_network begin
    1/(1+Y^2), 0 --> X
    1/(1+Y),   Y --> 0
end
jumpsys = convert(JumpSystem, rn, combinatoric_ratelaws = false)

u0 = [0,0]
de_chan0 = [[]]
tf = 400.
tspan = (0,tf)
τ = 20.
dprob = DiscreteProblem(jumpsys, u0, tspan)

delay_trigger_affect! = function (integrator, rng)
  append!(integrator.de_chan[1], τ)
end
delay_trigger = Dict(1=>delay_trigger_affect!)
delay_complete = Dict(1=>[2=>1, 1=>-1])
delay_interrupt = Dict()
delayjumpset = DelayJumpSet(delay_trigger, delay_complete, delay_interrupt)

djprob = DelayJumpProblem(jumpsys, dprob, DelayRejection(), delayjumpset, de_chan0, save_positions=(true,true))
sol = solve(djprob, SSAStepper(),seed=1234)

djprob2 = DelayJumpProblem(jumpsys, dprob, DelayDirect(), delayjumpset, de_chan0, save_positions = (true,true))
sol2 = solve(djprob2, SSAStepper(),seed=1234)

plot(sol.t,[sol.u[i][1] for i in 1:length(sol.u)])
plot!(sol.t,[sol.u[i][2] for i in 1:length(sol.u)])
scatter(sol2.t,[sol2.u[i][1] for i in 1:length(sol2.u)],legend = false)
scatter(sol2.t,[sol2.u[i][2] for i in 1:length(sol2.u)],legend = false)


function get_affine_map(ds::ArchGDAL.Dataset)
    # ArchGDAL fails hard on datasets without
    # an affinemap. GDAL documents that on fail
    # a default affinemap should be returned.
    try
        global gt = ArchGDAL.getgeotransform(ds)
    catch y
        @warn y.msg
        global gt = [0.,1.,0.,0.,0.,1.]
    end
    geotransform_to_affine(SVector{6, Float64}(gt))
end

function geotransform_to_affine(gt::SVector{6, Float64})
    # See https://lists.osgeo.org/pipermail/gdal-dev/2011-July/029449.html
    # for an explanation of the geotransform format
    AffineMap(SMatrix{2,2}([gt[2] gt[3]; gt[5] gt[6]]), SVector{2}([gt[1], gt[4]]))
end
geotransform_to_affine(A::Vector{Float64}) = geotransform_to_affine(SVector{6}(A))

function affine_to_geotransform(am::AffineMap{SArray{Tuple{2,2},Float64,2,4},SArray{Tuple{2},Float64,1,2}})
    l = am.linear
    t = am.translation
    [t[1], l[1], l[3], t[2], l[2], l[4]]
end

function affine_to_geotransform(am::AffineMap{Array{Float64,2},Array{Float64,1}})
    l = am.linear
    t = am.translation
    (length(l) != 4 || length(t) != 2) || error("AffineMap has wrong dimensions.")
    [t[1], l[1], l[3], t[2], l[2], l[4]]
end

"""Check wether the AffineMap of a GeoArray contains rotations."""
function is_rotated(ga::GeoArray)
    ga.f.linear[2] != 0. || ga.f.linear[3] != 0.
end

function bbox(ga::GeoArray)
    ax, ay = ga.f(SVector(0,0))
    bx, by = ga.f(SVector(size(ga)[1:2]))
    (min_x=min(ax, bx), min_y=min(ay, by), max_x=max(ax, bx), max_y=max(ay, by))
end

function unitrange_to_affine(x::StepRangeLen, y::StepRangeLen)
    δx, δy = step(x), step(y)
    AffineMap(
        SMatrix{2,2}(δx, 0, 0, δy),
        SVector(x[1] - δx/2, y[1] - δy/2)
    )
end

function bbox_to_affine(size::Tuple{Integer, Integer}, bbox::NamedTuple{(:min_x, :min_y, :max_x, :max_y),Tuple{Float64, Float64, Float64, Float64}})
    AffineMap(
        SMatrix{2,2}((bbox.max_x - bbox.min_x) / size[1], 0, 0, (bbox.max_y - bbox.min_y)/size[2]),
        SVector(bbox.min_x, bbox.min_y)
        )
end

"""Set geotransform of `GeoArray` by specifying a bounding box.
Note that this only can result in a non-rotated or skewed `GeoArray`."""
function bbox!(ga::GeoArray, bbox::NamedTuple{(:min_x, :min_y, :max_x, :max_y),Tuple{Float64, Float64, Float64, Float64}})
    ga.f = bbox_to_affine(size(ga)[1:2], bbox)
    ga
end

"""Generate bounding boxes for GeoArray cells."""
function bboxes(ga::GeoArray)
    c = coords(ga)::Array{StaticArrays.SArray{Tuple{2},Float64,1,2},2}
    m, n = size(c)
    cellbounds = Matrix{NamedTuple}(undef, (m-1, n-1))
    for j = 1:n-1, i = 1:m-1
        v = c[i:i+1, j:j+1]
        minx, maxx = extrema(first.(v))::Tuple{Float64, Float64}
        miny, maxy = extrema(last.(v))::Tuple{Float64, Float64}
        cellbounds[i, j] = (min_x=minx, max_x=maxx, min_y=miny, max_y=maxy)
    end
    cellbounds
end

# Extend CoordinateTransformations
CoordinateTransformations.compose(ga::GeoArray, t2::AffineMap) = CoordinateTransformations.compose(ga.f, t2)
CoordinateTransformations.compose(ga::GeoArray, t2::LinearMap) = CoordinateTransformations.compose(ga.f, t2)
CoordinateTransformations.compose(ga::GeoArray, t2::Translation) = CoordinateTransformations.compose(ga.f, t2)

"""Transform an GeoArray by applying a Transformation."""
function compose!(ga::GeoArray, t2::AffineMap)
    ga.f = compose(ga, t2)
end
function compose!(ga::GeoArray, t2::LinearMap)
    ga.f = compose(ga, t2)
end
function compose!(ga::GeoArray, t2::Translation)
    ga.f = compose(ga, t2)
end

"""Function to flip GeoArray upside down to adjust to GDAL ecosystem."""
function flipud!(ga::GeoArray)
    # Flip data upside down
    ga.A = reverse(ga.A, dims=2)

    # Find new corner coordinates
    ux, uy = ga.f(SVector{2}([0, size(ga)[2]]))

    # Define y mirror and compose
    lm = LinearMap(SMatrix{2,2}([1.0 0.0; 0.0 -1.0]))
    am = compose(ga.f, lm)  # AffineMap
    translate = SVector{2}([ux, uy])
    f = AffineMap(am.linear, translate)

    ga.f = f
    ga
end


module Math

export Vector3, Point, len, unit, dot, randominunitsphere, squaredlength, cross, lenhost, unithost, Transform
export identitytransform, translation, rotation, permute, scale

using CUDAnative
using CthulhuVision.Random

struct Vector3
    x::Float32
    y::Float32
    z::Float32
end

@inline Base.:+(a::Vector3, b::Vector3) :: Vector3 = Vector3(a.x + b.x, a.y + b.y, a.z + b.z)
@inline Base.:-(a::Vector3, b::Vector3) :: Vector3 = Vector3(a.x - b.x, a.y - b.y, a.z - b.z)
@inline Base.:-(a::Vector3) :: Vector3 = Vector3(-a.x, -a.y, -a.z)
@inline Base.:*(t::Float32, v::Vector3) :: Vector3 = Vector3(t*v.x, t*v.y, t*v.z)
@inline Base.:/(a::Vector3, t::Float32) :: Vector3 = Vector3(a.x / t, a.y / t, a.z / t)
@inline len(a::Vector3) = CUDAnative.sqrt(a.x*a.x + a.y*a.y + a.z*a.z)
@inline lenhost(a::Vector3) = sqrt(a.x*a.x + a.y*a.y + a.z*a.z)
@inline squaredlength(a::Vector3) :: Float32 = a.x*a.x + a.y*a.y + a.z*a.z

@inline dot(a::Vector3, b::Vector3) :: Float32 = a.x * b.x + a.y * b.y + a.z * b.z

@inline function cross(a::Vector3, b::Vector3) :: Vector3
    Vector3(
          a.y * b.z - a.z * b.y,
        -(a.x * b.z - a.z * b.x),
          a.x * b.y - a.y * b.x
    )
end

@inline function unit(a::Vector3) :: Vector3
    l = len(a)
    a / l
end

@inline function unithost(a::Vector3) :: Vector3
    l = lenhost(a)
    a / l
end

@inline function byindex(v::Vector3, index) :: Float32
    if index == 0
        v.x
    elseif index == 1
        v.y
    else
        v.z
    end
end

@inline permute(v::Vector3, x, y, z) :: Vector3 = Vector3(byindex(v, x), byindex(v, y), byindex(v, z))

@inline function randominunitsphere(rng::UniformRNG) :: Vector3
    p = Vector3(1.0f0, 1.0f0, 1.0f0)

    while squaredlength(p) >= 1.0f0
        p = 2.0f0 * Vector3(next(rng), next(rng), next(rng)) - Vector3(1f0, 1f0, 1f0)
    end

    p
end

#
# Point
#

struct Point
    x::Float32
    y::Float32
    z::Float32
end

@inline function byindex(p::Point, index) :: Float32
    if index == 0
        p.x
    elseif index == 1
        p.y
    else
        p.z
    end
end

@inline Base.:-(p::Point, v::Vector3) :: Point = Point(p.x - v.x, p.y - v.y, p.z - v.z)
@inline Base.:-(a::Point, b::Point) :: Vector3 = Vector3(a.x - b.x, a.y - b.y, a.z - b.z)
@inline Base.:+(a::Point, b::Vector3) :: Point = Point(a.x + b.x, a.y + b.y, a.z + b.z)
@inline Base.abs(p::Point) :: Point = Point(abs(p.x), abs(p.y), abs(p.z))
@inline permute(p::Point, x, y, z) :: Point = Point(byindex(p, x), byindex(p, y), byindex(p, z))
@inline dot(a::Vector3, b::Point) :: Float32 = a.x * b.x + a.y * b.y + a.z * b.z

#
# Transformation matrix 4x4
#

struct Transform
    e11::Float32
    e12::Float32
    e13::Float32
    e14::Float32
    e21::Float32
    e22::Float32
    e23::Float32
    e24::Float32
    e31::Float32
    e32::Float32
    e33::Float32
    e34::Float32
    e41::Float32
    e42::Float32
    e43::Float32
    e44::Float32
end

@inline function Base.:*(a::Transform, b::Transform) :: Transform
    e11 = a.e11 * b.e11 + a.e12 * b.e21 + a.e13 * b.e31 + a.e14 * b.e41
    e12 = a.e11 * b.e12 + a.e12 * b.e22 + a.e13 * b.e32 + a.e14 * b.e42
    e13 = a.e11 * b.e13 + a.e12 * b.e23 + a.e13 * b.e33 + a.e14 * b.e43
    e14 = a.e11 * b.e14 + a.e12 * b.e24 + a.e13 * b.e34 + a.e14 * b.e44

    e21 = a.e21 * b.e11 + a.e22 * b.e21 + a.e23 * b.e31 + a.e24 * b.e41
    e22 = a.e21 * b.e12 + a.e22 * b.e22 + a.e23 * b.e32 + a.e24 * b.e42
    e23 = a.e21 * b.e13 + a.e22 * b.e23 + a.e23 * b.e33 + a.e24 * b.e43
    e24 = a.e21 * b.e14 + a.e22 * b.e24 + a.e23 * b.e34 + a.e24 * b.e44

    e31 = a.e31 * b.e11 + a.e32 * b.e21 + a.e33 * b.e31 + a.e34 * b.e41
    e32 = a.e31 * b.e12 + a.e32 * b.e22 + a.e33 * b.e32 + a.e34 * b.e42
    e33 = a.e31 * b.e13 + a.e32 * b.e23 + a.e33 * b.e33 + a.e34 * b.e43
    e34 = a.e31 * b.e14 + a.e32 * b.e24 + a.e33 * b.e34 + a.e34 * b.e44

    e41 = a.e41 * b.e11 + a.e42 * b.e21 + a.e43 * b.e31 + a.e44 * b.e41
    e42 = a.e41 * b.e12 + a.e42 * b.e22 + a.e43 * b.e32 + a.e44 * b.e42
    e43 = a.e41 * b.e13 + a.e42 * b.e23 + a.e43 * b.e33 + a.e44 * b.e43
    e44 = a.e41 * b.e14 + a.e42 * b.e24 + a.e43 * b.e34 + a.e44 * b.e44

    Transform(
        e11, e12, e13, e14,
        e21, e22, e23, e24,
        e31, e32, e33, e34,
        e41, e42, e43, e44,
    )
end

@inline function Base.:*(a::Transform, b::Vector3) :: Vector3
    w = 1f0
    x = a.e11 * b.x + a.e12 * b.y + a.e13 * b.z + a.e14 * w
    y = a.e21 * b.x + a.e22 * b.y + a.e23 * b.z + a.e24 * w
    z = a.e31 * b.x + a.e32 * b.y + a.e33 * b.z + a.e34 * w

    Vector3(x, y, z)
end

@inline function Base.:*(a::Transform, b::Point) :: Point
    w = 1f0
    x = a.e11 * b.x + a.e12 * b.y + a.e13 * b.z + a.e14 * w
    y = a.e21 * b.x + a.e22 * b.y + a.e23 * b.z + a.e24 * w
    z = a.e31 * b.x + a.e32 * b.y + a.e33 * b.z + a.e34 * w

    Point(x, y, z)
end

########
# Host #
########

function identitytransform() :: Transform
    Transform(
        1f0, 0f0, 0f0, 0f0,
        0f0, 1f0, 0f0, 0f0,
        0f0, 0f0, 1f0, 0f0,
        0f0, 0f0, 0f0, 1f0,
    )
end

function translation(v::Vector3) :: Transform
    Transform(
        1f0, 0f0, 0f0, v.x,
        0f0, 1f0, 0f0, v.y,
        0f0, 0f0, 1f0, v.z,
        0f0, 0f0, 0f0, 0f0,
    )
end

function rotation(θ::Float32, axis::Vector3) :: Transform
    u = unithost(axis)

    e11 = cos(θ) + u.x*u.x*(1f0 - cos(θ))
    e12 = u.x * u.y * (1f0 - cos(θ)) - u.z * sin(θ)
    e13 = u.x * u.z * (1f0 - cos(θ)) + u.y * sin(θ)
    e14 = 0f0

    e21 = u.y * u.x * (1f0 - cos(θ)) + u.z * sin(θ)
    e22 = cos(θ) + u.y * u.y * (1f0 - cos(θ))
    e23 = u.y * u.z * (1f0 - cos(θ)) - u.x * sin(θ)
    e24 = 0f0

    e31 = u.z * u.x * (1f0 - cos(θ)) - u.y * sin(θ)
    e32 = u.z * u.y * (1f0 - cos(θ)) + u.x * sin(θ)
    e33 = cos(θ) + u.z * u.z * (1f0 - cos(θ))
    e34 = 0f0

    e41 = 0f0
    e42 = 0f0
    e43 = 0f0
    e44 = 1f0

    Transform(
        e11, e12, e13, e14,
        e21, e22, e23, e24,
        e31, e32, e33, e34,
        e41, e42, e43, e44,
    )
end

function scale(v::Vector3) :: Transform
    Transform(
        v.x, 0f0, 0f0, 0f0,
        0f0, v.y, 0f0, 0f0,
        0f0, 0f0, v.z, 0f0,
        0f0, 0f0, 0f0, 1f0,
    )
end

end


get_total_fitness_variance(history::History) = var(get_fitness.(get_detailed_history(history)))

function get_inherent_variance(history::History)
    organisms = get_detailed_history(history)
    return mean(var(mean(organisms, [EnvironmentType(), PopulationType(), NumOffspringsType()]), [EnvironmentType(), PopulationType()]))
end

function get_mutational_variance(history::History)
    organisms = get_detailed_history(history)
    return mean(var(organisms, [EnvironmentType(), PopulationType(), NumOffspringsType()]))
end

function get_abiotic_variance(history::History)
    organisms = get_detailed_history(history)
    return var(mean(organisms, [EnvironmentType()]))
end

function get_biotic_variance(history::History)
    organisms = get_detailed_history(history)
    return mean(var(mean(organisms, [EnvironmentType(), PopulationType()]), [EnvironmentType()]))
end


function column_model_1D_kpp(constants, BCs, ICs, times, parameters=OceanTurb.KPP.Parameters())
    f, α, g, Nz, H = constants.f, constants.α, constants.g, constants.Nz, constants.H
    ρ₀ = 1027.0
    cₚ = 4000.0
    β  = 0.0
    kpp_constants = OceanTurb.Constants(Float64, ρ₀=ρ₀, cP=cₚ, f=f, α=α, β=β, g=g)

    model = OceanTurb.KPP.Model(N=Nz, H=H, stepper=:BackwardEuler, constants=kpp_constants, parameters=parameters)

    U₀ = ICs.u
    V₀ = ICs.v
    T₀ = ICs.T

    model.solution.U.data[1:Nz] .= U₀
    model.solution.V.data[1:Nz] .= V₀
    model.solution.T.data[1:Nz] .= T₀

    UW_flux = BCs.uw.top
    VW_flux = BCs.vw.top
    WT_flux = BCs.wT.top

    Δz = H / Nz

    ∂U₀∂z = (U₀[2] - U₀[1]) / Δz
    ∂V₀∂z = (V₀[2] - V₀[1]) / Δz
    ∂T₀∂z = (T₀[2] - T₀[1]) / Δz

    model.bcs.U.top = OceanTurb.FluxBoundaryCondition(UW_flux)
    model.bcs.V.top = OceanTurb.FluxBoundaryCondition(VW_flux)

    if WT_flux isa Number
        model.bcs.T.top = OceanTurb.FluxBoundaryCondition(WT_flux)
    else
        model.bcs.T.top = OceanTurb.FluxBoundaryCondition(WT_flux(0))
    end

    model.bcs.U.bottom = OceanTurb.GradientBoundaryCondition(∂U₀∂z)
    model.bcs.V.bottom = OceanTurb.GradientBoundaryCondition(∂V₀∂z)
    model.bcs.T.bottom = OceanTurb.GradientBoundaryCondition(∂T₀∂z)

    Nt = length(times)
    U = zeros(Nz, Nt)
    V = zeros(Nz, Nt)
    T = zeros(Nz, Nt)

    UW = zeros(Nz+1, Nt)
    VW = zeros(Nz+1, Nt)
    WT = zeros(Nz+1, Nt)

    # loop the model
    Δt = times[2] - times[1]
    for n in 1:Nt
        OceanTurb.run_until!(model, Δt, times[n])

        if !isa(WT_flux, Number)
            model.bcs.T.top.condition = WT_flux(model.clock.time)
        end

        U[:, n] .= model.solution.U[1:Nz]
        V[:, n] .= model.solution.V[1:Nz]
        T[:, n] .= model.solution.T[1:Nz]

        UW[:, n] .= OceanTurb.diffusive_flux(:U, model)[1:Nz+1]
        VW[:, n] .= OceanTurb.diffusive_flux(:V, model)[1:Nz+1]
        WT[:, n] .= OceanTurb.diffusive_flux(:T, model)[1:Nz+1] .+ OceanTurb.KPP.nonlocal_temperature_flux(model)[1:Nz+1]

        UW[Nz+1, n] = UW_flux
        VW[Nz+1, n] = VW_flux

        if WT_flux isa Number
            WT[Nz+1, n] = WT_flux
        else
            WT[Nz+1, n] = WT_flux(model.clock.time)
        end
    end

    return (; U, V, T, UW, VW, WT)
end


#=
Very similar to examples/plasticityFFtest.jl
Here I consider triplet rule, two high frequency neurons and two low
frequency ones (all Poisson).  Desired outcome: weight increases between the
high firing neurons, decreases in all other cases.
=#
push!(LOAD_PATH, abspath(@__DIR__,".."))

using Test,ProgressMeter
using LinearAlgebra,Statistics,StatsBase,Distributions
using Plots,NamedColors ; theme(:dark)
using SparseArrays
using SpikingRNNs; const global S = SpikingRNNs
using BenchmarkTools

using Random ; Random.seed!(1)

##
const dt = 0.1E-3
const Ttot = 20.0
const highrate = 80.0
const lowrate = 5.0
const nneus = 4

nt,ps = let spkgen = S.SGPoissonMulti([highrate,highrate,lowrate,lowrate])
  # spkgen = S.SGPoissonMulti([highrate,highrate])
  sker = S.SKExp(-1E6) # this will never be used
  vrev = -3000.0 # same
  nt = S.NTInputConductance(spkgen,sker,vrev)
  ps = S.PSInputConductance(nt,nneus)
  (nt,ps)
end

# plasticity and initial weights

wstart = let w = 100.0
  wm = fill(w,(nneus,nneus))
  wm[diagind(wm)] .= 0.0
  sparse(wm)
end

triplets_plasticity = let (n_post,n_pre) = size(wstart),
  τplus = 17E-3 # 16.8
  τminus = 34E-3 # 33.7
  τx = 100E-3 # 101
  τy = 120E-3 # 125
  plast_eps = 1E-3
  A2plus = 7.5E-7*plast_eps
  A3plus = 9.3*plast_eps
  A2minus = 7.0*plast_eps
  A3minus = 2.3E-1*plast_eps
  S.PlasticityTriplets(τplus,τminus,τx,τy,A2plus,A3plus,
    A2minus,A3minus,n_post,n_pre)
end

conn = S.ConnectionPlasticityTest(wstart,triplets_plasticity)

pop = S.Population(ps,(conn,ps))

# that's it, let's make the network
ntw = S.RecurrentNetwork(dt,pop)

## recorders : spikes
krec = 1
rec_spikes = S.RecSpikes(ps,highrate,Ttot)
rec_weights = S.RecWeights(conn,krec,dt,Ttot)

## Run

times = (0.0:ntw.dt:Ttot)
nt = length(times)
# clean up
S.reset!.([ps,rec_spikes,rec_weights,conn.plasticities[1]])
# initial conditions

@time @showprogress for (k,t) in enumerate(times)
  rec_spikes(t,k,ntw)
  rec_weights(t,k,ntw)
  S.dynamics_step!(t,ntw)
end

spk_out = S.get_spiketimes_dictionary(rec_spikes)
tw,weights_all = S.rebuild_weights(rec_weights)
##
# check spikes
_ = let plt=plot(leg=false)
  for i in 1:nneus
   scatter!(plt, _->i, spk_out[i];marker=:vline)
  end
  plt
 end

##

_ = let plt=plot(legend=:topleft,xlabel="time",ylabel="weight")
  x = rec_weights.times
  n = size(weights_all[1],1)
  pre = 1
  for post in 1:n
    if post != pre
      y = getindex.(weights_all,post,pre)
      plot!(plt,x,y;linewidth=2,label="$(post) <- $(pre)")
    end
  end
  plot(plt)
end

##

##############################
# Let's make a function to test it

function test_plasticity(lowrate::R,highrate::R,
    myplasticity::S.PlasticityRule;(Ttot::R)=50.0,(dt::R)=0.1E-3) where R
  nneus = 2
  nt,ps = let spkgen = S.SGPoissonMulti([highrate,lowrate])
    sker = S.SKExp(-1E6) # this will never be used
    vrev = -3000.0 # same
    nt = S.NTInputConductance(spkgen,sker,vrev)
    ps = S.PSInputConductance(nt,nneus)
    (nt,ps)
  end
  wstart = let w = 100.0
    wm = fill(w,(nneus,nneus))
    wm[diagind(wm)] .= 0.0
    sparse(wm)
  end
  conn = S.ConnectionPlasticityTest(wstart,myplasticity)
  pop = S.Population(ps,(conn,ps))
  ntw = S.RecurrentNetwork(dt,pop)

  ## recorders : spikes
  krec = 1
  rec_spikes = S.RecSpikes(ps,highrate,Ttot)
  rec_weights = S.RecWeights(conn,krec,dt,Ttot)

  ## Run
  times = (0.0:ntw.dt:Ttot)
  nt = length(times)
  # clean up
  S.reset!.([ps,rec_spikes,rec_weights,conn.plasticities[1]])
  for (k,t) in enumerate(times)
    rec_spikes(t,k,ntw)
    rec_weights(t,k,ntw)
    S.dynamics_step!(t,ntw)
  end
  # FROM HIGH TO LOW
  idx_low_high = findfirst(idx->Tuple(idx)==(2,1),rec_weights.idx_save)
  w_low_high = rec_weights.weights_now[idx_low_high,:]
  # FROM LOW TO HIGH
  idx_high_low = findfirst(idx->Tuple(idx)==(1,2),rec_weights.idx_save)
  w_high_low = rec_weights.weights_now[idx_high_low,:]
  _dostuff = function(v::Vector{R})
    vd=diff(v)
    filter!(!=(0.0),vd)
    sum(vd)/Ttot
  end
  return _dostuff.((w_low_high,w_high_low))
end
##
triplets_plasticity = let (n_post,n_pre) = size(wstart),
  τplus = 17E-3 # 16.8
  τminus = 34E-3 # 33.7
  τx = 100E-3 # 101
  τy = 120E-3 # 125
  plast_eps = 1E-3
  A2plus = 7.5E-7*plast_eps
  A3plus = 9.3*plast_eps
  A2minus = 7.0*plast_eps
  A3minus = 2.3E-1*plast_eps
  S.PlasticityTriplets(τplus,τminus,τx,τy,A2plus,A3plus,
    A2minus,A3minus,n_post,n_pre)
end

booh = test_plasticity(10.0,80.0,triplets_plasticity)

##
nrats = 25
ratestest = range(0.5,60;length=nrats)
testret = fill(NaN,nrats,nrats)

for i in 1:nrats, j in i:nrats
  # j is high, i is low or equal
  @info "now processing $i,$j (out of $nrats)"
  (low_high,high_low) = test_plasticity(ratestest[i],
    ratestest[j], triplets_plasticity ; Ttot=25.0)
  testret[i,j] = low_high
  testret[j,i] = high_low
end
##

heatmap(ratestest,ratestest,testret;color=:vik)

function plus_minus_rescale(mat::Array{R}) where R
  down,up = extrema(mat)
  @assert (down < 0.0) && (up > 0.0)
  pluspart = mat ./ up
  minuspart = mat ./ abs(down)
  ret = similar(mat)
  for i in eachindex(mat)
    if mat[i] > 0
      ret[i] = pluspart[i]
    else
      ret[i] = minuspart[i]
    end
  end
  return ret
end

heatmap(ratestest,ratestest,testret;
  color=:seismic,ratio=1)

heatmap(ratestest,ratestest,plus_minus_rescale(testret);
  color=:seismic,ratio=1)


# Advent of Code 2020 - Day 22

test_str = "Player 1:
9
2
6
3
1

Player 2:
5
8
4
7
10"

test_list = split(test_str, '\n')

list = readlines("day22_input.txt")

function play_combat(list)
    num_cards = div(length(list)+1, 2) - 2
    decks = fill(0, 2, 2*num_cards)
    decks[1, 1:num_cards] = parse.(Int, list[2:num_cards+1])
    decks[2, 1:num_cards] = parse.(Int, list[num_cards+4:end])
    N = num_cards
    while 0 < N < 2*num_cards
        cards = decks[:, 1]
        decks[:, 1:2*num_cards-1] = decks[:, 2:end]
        winner = argmax(cards)
        slot = winner == 1 ? N : 2*num_cards-N
        N += winner == 1 ? 1 : -1
        decks[winner, slot:slot+1] = [cards[winner] cards[3-winner]]
    end

    winner = N == 0 ? 2 : 1
    points = sum([decks[winner, j]*(2*num_cards-j+1) for j in 1:2*num_cards])

    return points
end

@show play_combat(test_list) == 306
@show play_combat(list) == 32495

function play_recursive_combat(list)
    num_cards = div(length(list)+1, 2) - 2
    decks = fill(0, 2, 2*num_cards)
    decks[1, 1:num_cards] = parse.(Int, list[2:num_cards+1])
    decks[2, 1:num_cards] = parse.(Int, list[num_cards+4:end])

    function playgame(decks)
        num_cards = div(size(decks)[2], 2)
        history1 = Set()
        history2 = Set()
        gamerepeated = false
        while minimum(sum(decks, dims=2)) > 0 && !gamerepeated
            if  decks[1,:] ∈ history1 || decks[2,:] ∈ history2
                gamerepeated = true
                winner = 1
            else
                push!(history1, decks[1,:])
                push!(history2, decks[2,:])
                cards = decks[:, 1]
                decks[:, 1:2*num_cards-1] = decks[:, 2:end]
                if sum(sum(decks.>0, dims=2) .≥ cards) == 2
                    new_decks = copy(decks)
                    new_decks[1,cards[1]+1:end] .= 0
                    new_decks[2,cards[2]+1:end] .= 0
                    winner = playgame(new_decks)
                else
                    winner = argmax(cards)
                end
                slot = findfirst(==(0), decks[winner,:])
                decks[winner, slot:slot+1] = [cards[winner] cards[3-winner]]
            end
        end
        return winner
    end

    winner = playgame(decks)

    points = sum([decks[winner, j]*(2*num_cards-j+1) for j in 1:2*num_cards])
    return points
end

@show play_recursive_combat(test_list) == 291
@show play_recursive_combat(list) == 32665

nothing


module TestGMM

using Statistics
using Test
import TreeParzen: GMM

N_SAMPLES = 100_000

# see here for an expression of the variance of a gaussian mixture
# https://stats.stackexchange.com/a/16609
mixture_variance(weights, sigmas, means) = sum(weights .* (sigmas .^ 2)) + sum(weights .* (means .^ 2)) - sum(weights .* means) ^ 2

@testset "GMM1" begin

    # Check that mu is specified correctly
    mu_test_std = 1e-7
    mu_test_mu = 1e1
    mu_test_weight = 1e0
    # for whatever standard deviation we use we don't expect more than 1-2 orders of magnitude deviation from this
    # if sigma was larger the rtol might need adjusting upwards a little but not a lot.
    @test all(isapprox.(mu_test_mu, GMM.GMM1([mu_test_weight], [mu_test_mu], [mu_test_std], N_SAMPLES); rtol = mu_test_std * 1e1))

    # Check that sigma is specified correctly
    sigma_test_std = 1e1
    sigma_test_mu = 0e0
    sigma_test_weight = 1e0
    @test isapprox(std(GMM.GMM1([sigma_test_weight], [sigma_test_mu], [sigma_test_std], N_SAMPLES)), sigma_test_std; rtol=1e-1)

    # Check a mixture of gaussians, equally weighted, and contributes to mixture variance
    mixture_test_stds = [1e-6, 1e-6]
    mixture_test_mus = [0e0, 1e0]
    mixture_test_weights = [0.5e0, 0.5e0]
    samples = GMM.GMM1(mixture_test_weights, mixture_test_mus, mixture_test_stds, N_SAMPLES)
    expected_variance = mixture_variance(mixture_test_weights, mixture_test_stds, mixture_test_mus)
    expected_mean = sum(mixture_test_weights .* mixture_test_mus)
    @test isapprox(expected_mean, mean(samples); rtol=3e-2)
    @test isapprox(expected_variance, var(samples); rtol=1e-1)

    # Check an uneven weighted mixture has expected variance (this is quite a bit more variable)
    uneven_mixture_test_stds = [1e-6, 1e-6]
    uneven_mixture_test_mus = [0e0, 1e0]
    uneven_mixture_test_weights = [1 - 1e-4, 1e-4]
    samples = GMM.GMM1(uneven_mixture_test_weights, uneven_mixture_test_mus, uneven_mixture_test_stds, N_SAMPLES)
    expected_variance = mixture_variance(uneven_mixture_test_weights, uneven_mixture_test_stds, uneven_mixture_test_mus)
    expected_mean = sum(uneven_mixture_test_weights .* uneven_mixture_test_mus)
    @test size(samples) == (N_SAMPLES,)
    # these are whopper tolerances because although the numbers come out close
    # the highly unbalanced nature of the mixture makes drawn samples from run to run extremely variable
    # even with 1e5 samples drawn we don't expect most of the time to draw from the 2nd
    # distribution, but that can differ by maybe an ofer of magnitude from run to run
    @test isapprox(expected_mean, mean(samples); rtol=1e0)
    @test isapprox(expected_variance, var(samples); rtol=1e0)

    # lpdf scalar one component
    llval = GMM.GMM1_lpdf([1.0], [1.], [1.0], [2.0])
    @test size(llval) == (1,) # Shape should match first parameter above
    @test isapprox(llval, [log(1.0 / sqrt(2pi * 2.0 ^ 2))])

    # lpdf vector N components
    llval = GMM.GMM1_lpdf([1.0, 0.0], [0.25, 0.25, .5], [0.0, 1.0, 2.0], [1.0, 2.0, 5.0])

    a = .25 / sqrt(2pi * 1^2) * exp(-.5 * 1^2)
    a += .25 / sqrt(2pi * 2^2)
    a += .5 / sqrt(2pi * 5^2) * exp(-.5 * (1 / 5)^2)
    @test size(llval) == (2,)
    @test isapprox(llval[1], log(a))

    a = .25 / sqrt(2pi * 1^2)
    a += .25 / sqrt(2pi * 2^2) * exp(-.5 * (1 / 2)^2)
    a += .5 / sqrt(2pi * 5^2) * exp(-.5 * (2 / 5)^2)
    @test isapprox(llval[2], log(a))

    # weights, mus and sigmas have different lengths
    @test_throws DimensionMismatch GMM.GMM1([1.0, 0.0], [0.0, 1.0, 2.0], [10.0], 1000)

    # non-1 sum of weights throws
    @test_throws ArgumentError GMM.GMM1([1., 2.], [1., 2.], [1., 2.], 1)
    @test_throws ArgumentError GMM.GMM1([0.5, 0.6], [1., 2.], [1., 2.], 1)
    @test_throws ArgumentError GMM.GMM1([0.2, 0.1], [1., 2.], [1., 2.], 1)
    @test_throws ArgumentError GMM.GMM1([-1., 2.], [1., 2.], [1., 2.], 1)
    @test_throws ArgumentError GMM.GMM1([-0.5, -0.5], [1., 2.], [1., 2.], 1)

end

end
true


using MPI
using Test
using ClimateMachine
using ClimateMachine.Mesh.Topologies
using ClimateMachine.Mesh.Grids
using ClimateMachine.DGMethods
using ClimateMachine.DGMethods.NumericalFluxes
using ClimateMachine.MPIStateArrays
using ClimateMachine.ODESolvers
using ClimateMachine.GenericCallbacks
using ClimateMachine.VariableTemplates: flattenednames
using ClimateMachine.BalanceLaws
using ClimateMachine.Ocean.ShallowWater
using ClimateMachine.Ocean.ShallowWater:
    TurbulenceClosure, LinearDrag, ConstantViscosity
using ClimateMachine.Ocean
using ClimateMachine.Ocean.OceanProblems

using LinearAlgebra
using StaticArrays
using Logging, Printf, Dates
using ClimateMachine.VTK

using CLIMAParameters
using CLIMAParameters.Planet: grav
struct EarthParameterSet <: AbstractEarthParameterSet end
const param_set = EarthParameterSet()

if !isempty(ARGS)
    stommel = Bool(parse(Int, ARGS[1]))
    linear = Bool(parse(Int, ARGS[2]))
    test = parse(Int, ARGS[3])
else
    stommel = true
    linear = true
    test = 1
end

###################
# PARAM SELECTION #
###################
const FT = Float64

const τₒ = 2e-4 # value includes τₒ, g, and ρ
const fₒ = 1e-4
const β = 1e-11
const λ = 1e-6
const ν = 5e2

const Lˣ = 1e6
const Lʸ = 1e6
const timeend = 300 * 24 * 60 * 60
const H = 1000
const c = sqrt(grav(param_set) * H)

if stommel
    gyre = "stommel"
else
    gyre = "munk"
end

if linear
    advec = "linear"
else
    advec = "nonlinear"
end

outname = "vtk_new_dt_" * gyre * "_" * advec

function setup_model(FT, stommel, linear, τₒ, fₒ, β, γ, ν, Lˣ, Lʸ, H)
    problem = HomogeneousBox{FT}(
        Lˣ,
        Lʸ,
        H,
        τₒ = τₒ,
        BC = OceanBC(Impenetrable(NoSlip()), Insulating()),
    )

    if stommel
        turbulence = LinearDrag{FT}(λ)
    else
        turbulence = ConstantViscosity{FT}(ν)
    end

    if linear
        advection = nothing
    else
        advection = NonLinearAdvectionTerm()
    end

    model = ShallowWaterModel{FT}(
        param_set,
        problem,
        turbulence,
        advection,
        c = c,
        fₒ = fₒ,
        β = β,
    )
end

#########################
# Timestepping function #
#########################

function test_run(mpicomm, topl, ArrayType, N, dt, FT, model, test)
    grid = DiscontinuousSpectralElementGrid(
        topl,
        FloatType = FT,
        DeviceArray = ArrayType,
        polynomialorder = N,
    )

    dg = DGModel(
        model,
        grid,
        RusanovNumericalFlux(),
        CentralNumericalFluxSecondOrder(),
        CentralNumericalFluxGradient(),
    )

    Q = init_ode_state(dg, FT(0))
    Qe = init_ode_state(dg, FT(timeend))

    lsrk = LSRK144NiegemannDiehlBusch(dg, Q; dt = dt, t0 = 0)

    nt_freq = floor(Int, 1 // 10 * timeend / dt)

    cbcs_dg = ClimateMachine.StateCheck.sccreate(
        [(Q, "2D state")],
        nt_freq;
        prec = 12,
    )

    cb = (cbcs_dg,)

    ##CNH if test > 2
        outprefix = @sprintf("ic_mpirank%04d_ic", MPI.Comm_rank(mpicomm))
        statenames = flattenednames(vars_state(model, Prognostic(), eltype(Q)))
        auxnames = flattenednames(vars_state(model, Auxiliary(), eltype(Q)))
        writevtk(outprefix, Q, dg, statenames, dg.state_auxiliary, auxnames)

        outprefix = @sprintf("exact_mpirank%04d", MPI.Comm_rank(mpicomm))
        statenames = flattenednames(vars_state(model, Prognostic(), eltype(Qe)))
        auxnames = flattenednames(vars_state(model, Auxiliary(), eltype(Qe)))
        writevtk(outprefix, Qe, dg, statenames, dg.state_auxiliary, auxnames)

        vtkstep = [0]
        vtkpath = outname
        mkpath(vtkpath)
        cbvtk = GenericCallbacks.EveryXSimulationSteps(1000) do
            outprefix = @sprintf(
                "%s/mpirank%04d_step%04d",
                vtkpath,
                MPI.Comm_rank(mpicomm),
                vtkstep[1]
            )
            @debug "doing VTK output" outprefix
            statenames =
                flattenednames(vars_state(model, Prognostic(), eltype(Q)))
            auxiliarynames =
                flattenednames(vars_state(model, Auxiliary(), eltype(Q)))
            writevtk(
                outprefix,
                Q,
                dg,
                statenames,
                dg.state_auxiliary,
                auxiliarynames,
            )
            vtkstep[1] += 1
            nothing
        end
        cb = (cb..., cbvtk)
    ##CNH end

    solve!(Q, lsrk; timeend = timeend, callbacks = cb)

    error = euclidean_distance(Q, Qe) / norm(Qe)
    @info @sprintf(
        """Finished
        error = %.16e
        """,
        error
    )

    return error
end

################
# Start Driver #
################

let
    ClimateMachine.init()
    ArrayType = ClimateMachine.array_type()
    mpicomm = MPI.COMM_WORLD

    model = setup_model(FT, stommel, linear, τₒ, fₒ, β, λ, ν, Lˣ, Lʸ, H)

    if test == 1
        cellsrange = 10:10
        orderrange = 4:4
        testval = 1.6068814534535144e-03
    elseif test == 2
        cellsrange = 5:5:10
        orderrange = 3:4
    elseif test > 2
        cellsrange = 5:5:25
        orderrange = 6:10
    end

    errors = zeros(FT, length(cellsrange), length(orderrange))
    for (i, Ne) in enumerate(cellsrange)
        brickrange = (
            range(FT(0); length = Ne + 1, stop = Lˣ),
            range(FT(0); length = Ne + 1, stop = Lʸ),
        )
        topl = BrickTopology(
            mpicomm,
            brickrange,
            periodicity = (false, false),
            boundary = ((1, 1), (1, 1)),
        )

        for (j, N) in enumerate(orderrange)
            @info "running Ne $Ne and N $N with"
            dt = (Lˣ / c) / Ne / N^2
            @info @sprintf("\n dt = %f", dt)
            errors[i, j] =
                test_run(mpicomm, topl, ArrayType, N, dt, FT, model, test)
        end
    end

    @test errors[end, end] ≈ testval

    #=
    msg = ""
    for i in length(cellsrange)-1
      rate = log2(errors[i, end] - log2(errors[i+1, end]))
      msg *= @sprintf("\n rate for Ne %d = %e", cellsrange[i], rate)
    end
    @info msg

    msg = ""
    for j in length(orderrange)-1
      rate = log2(errors[end, j] - log2(errors[end, j+1]))
      msg *= @sprintf("\n rate for N  %d = %e", orderrange[j], rate)
    end
    @info msg
    =#

end


using ApproxFun, SingularIntegralEquations, DualNumbers, RiemannHilbert, LinearAlgebra, FastTransforms, SpecialFunctions, Test
import ApproxFunBase: ArraySpace, pieces, dotu, interlace
import RiemannHilbert: RiemannDual, LogNumber, fpstieltjesmatrix!, fpstieltjesmatrix, orientedleftendpoint, orientedrightendpoint, finitepart, fpcauchymatrix, collocationvalues, collocationpoints
import SingularIntegralEquations: stieltjesmoment, stieltjesmoment!, undirected, Directed, ⁺, ⁻, istieltjes
import SingularIntegralEquations.HypergeometricFunctions: speciallog

@testset "RiemannDual" begin
    for h in (0.1,0.01), a in (2exp(0.1im),1.1)
        @test log(RiemannDual(0,a))(h) ≈ log(h*a)
        @test log(RiemannDual(Inf,a))(h) ≈ log(a/h)
    end

    for h in (0.1,0.01), a in (2exp(0.1im),1.1)
        @test log1p(RiemannDual(-1,a))(h) ≈ log(h*a)
        @test log1p(RiemannDual(Inf,a))(h) ≈ log(a/h)
    end

    h = 0.0000001
    for z in (RiemannDual(-1,-1), RiemannDual(1,1), RiemannDual(-1,2exp(0.1im)), RiemannDual(1,2exp(0.1im))),
            k = 0:1
        l = stieltjesjacobimoment(0,0,k,z)
        @test l(h) ≈ stieltjesjacobimoment(0,0,k,realpart(z)+epsilon(z)h) atol=1E-5
    end

    h=0.0001
    for z in (RiemannDual(1,3exp(0.2im)), RiemannDual(1,0.5exp(-1.3im)),
                RiemannDual(-1,3exp(0.2im)), RiemannDual(-1,0.5exp(-1.3im)),
                RiemannDual(-1,1), RiemannDual(1,-1))
        @test atanh(z)(h) ≈  atanh(realpart(z)+epsilon(z)h) atol = 1E-4
    end

    z = RiemannDual(1,-0.25)
    h = 0.0000001
    @test speciallog(z)(h) ≈ speciallog(realpart(z)+epsilon(z)h) atol=1E-4

    h = 0.00001
    for z in (RiemannDual(-1,-1), RiemannDual(-1,exp(0.1im)), RiemannDual(-1,exp(-0.1im)))
        @test stieltjesjacobimoment(0.5,0,0,z)(h) ≈ stieltjesjacobimoment(0.5,0,0,realpart(z)+epsilon(z)h) atol=1E-4
    end
end

@testset "Legendre Cauchy" begin
    f = Fun(exp,Legendre())

    h = 0.00001
    for z in  (RiemannDual(-1,-1), RiemannDual(-1,1+im), RiemannDual(-1,1-im))
        @test cauchy(f, z)(h) ≈ cauchy(f, realpart(z) + epsilon(z)h) atol=1E-4
    end
end

@testset "Directed and RiemannDual" begin
    @test undirected(Directed{false}(RiemannDual(0,-1))) == 0

    @test real(LogNumber(2im,im+1)) == LogNumber(0,1)
    @test imag(LogNumber(2im,im+1)) == LogNumber(2,1)
    @test conj(LogNumber(2im,im+1)) == LogNumber(-2im,1-im)

    @test log(Directed{false}(RiemannDual(0,-1))) == LogNumber(1,π*im)
    @test log(Directed{true}(RiemannDual(0,-1))) == LogNumber(1,-π*im)

    @test log(Directed{false}(RiemannDual(0,-1-eps()*im))) == LogNumber(1,π*im)
    @test log(Directed{false}(RiemannDual(0,-1+eps()*im))) == LogNumber(1,π*im)

    @test log(Directed{true}(RiemannDual(0,-1-eps()*im))) == LogNumber(1,-π*im)
    @test log(Directed{true}(RiemannDual(0,-1+eps()*im))) == LogNumber(1,-π*im)

    z = Directed{false}(RiemannDual(1,-2))

    for k=0:1, s=(false,true)
        z = Directed{s}(RiemannDual(-1,2))
        l = stieltjesmoment(Legendre(),k,z)
        h = 0.00000001
        @test l(h) ≈ stieltjesmoment(Legendre(),k,-1 + epsilon(z.x)h + (s ? 1 : -1)*eps()*im) atol=1E-5
    end

    @test RiemannHilbert.orientedleftendpoint(ChebyshevInterval()) == RiemannDual(-1.0,1)
    @test RiemannHilbert.orientedrightendpoint(ChebyshevInterval()) == RiemannDual(1.0,-1)
end

@testset "Interval FPStieltjes" begin
    Γ = ChebyshevInterval()
    f = Fun(x->exp(-40(x-0.1)^2), Legendre())
    C = Array{ComplexF64}(undef, ncoefficients(f), ncoefficients(f))
    d = Segment(im,2im)

    fpstieltjesmatrix!(C, space(f), d)
    c = Fun(d, chebyshevtransform(C*coefficients(f); kind=2))
    @test c(1.5im) ≈ stieltjes(f,1.5im)

    d = Segment(-1,-1+im)
    fpstieltjesmatrix!(C, space(f), d)
    @test norm(C) ≤ 100
    c = Fun(d, chebyshevtransform(C*coefficients(f); kind=2))
    @test c(-1+0.5im) ≈ stieltjes(f,-1+0.5im)

    d = Segment(1,1+im)
    fpstieltjesmatrix!(C, space(f), d)
    @test norm(C) ≤ 200
    c = Fun(d, chebyshevtransform(C*coefficients(f); kind=2))
    @test c(1+0.5im) ≈ stieltjes(f,1+0.5im)

    d = ChebyshevInterval()
    fpstieltjesmatrix!(C, space(f), d)
    @test norm(C) ≤ 200
    c = Fun(d, chebyshevtransform(C*coefficients(f); kind=2))
    @test c(0.5) ≈ stieltjes(f,0.5⁻)

    d = Segment(0,1)
    f = Fun(x->exp(-200(x-0.6)^2), Legendre(d))
    C = fpstieltjesmatrix(space(f), ncoefficients(f), ncoefficients(f))
    @test norm(C) ≤ 200
    c = Fun(d, chebyshevtransform(C*coefficients(f); kind=2))
    @test c(0.5) ≈ stieltjes(f,0.5⁻)
end

@testset "finitepart stieltjes" begin
    f = Fun(exp,Legendre())
    f1 = Fun(exp,Legendre(-1..0))
    f2 = Fun(exp,Legendre(0..1))
    fp = f1+f2

    @test stieltjes(f,0.0⁻) ≈ finitepart(stieltjes(f1,RiemannDual(0.0,-im)) + stieltjes(f2,RiemannDual(0.0,-im)))
    @test stieltjes(f,0.0⁻) ≈ finitepart(stieltjes(f1,RiemannDual(0.0,exp(-0.1im))) + stieltjes(f2,RiemannDual(0.0,exp(-0.1im))))
    @test stieltjes(f,0.0⁻) ≈ finitepart(stieltjes(f1,Directed{false}(RiemannDual(0.0,-1.0))) + stieltjes(f2,RiemannDual(0.0,-1.0)))

    @test stieltjes(fp,RiemannDual(0.0,-im)) ≈ stieltjes(f1,RiemannDual(0.0,-im)) + stieltjes(f2,RiemannDual(0.0,-im))
    @test stieltjes(f,0.0⁻) ≈ finitepart(stieltjes(fp,RiemannDual(0.0,-im)))
end

@testset "Two interval" begin
    @testset "-1..0 and 0..1" begin
        sp = Legendre(-1 .. 0) ⊕ Legendre(0 .. 1)
        f = Fun(x->exp(-40(x-0.1)^2), sp)
        v = components(f)
        ns = ncoefficients.(v)
        C11 = fpstieltjesmatrix(space(v[1]), ncoefficients(v[1]), ncoefficients(v[1]))
        c_vals11 = C11*v[1].coefficients
        h = 0.00000001; @test stieltjes(v[1],h*im) ≈ stieltjes(v[1],RiemannDual(0.0,im))(h) atol=1E-6
        @test finitepart(stieltjes(v[1],Directed{false}(RiemannDual(0.0,-1)))) ≈ c_vals11[1]

        C22 = fpstieltjesmatrix(space(v[2]), ncoefficients(v[2]), ncoefficients(v[2]))
        c_vals22 = C22*v[2].coefficients
        h = 0.00000001; @test stieltjes(v[2],h*im) ≈ stieltjes(v[2],RiemannDual(0.0,im))(h) atol=1E-6
        @test finitepart(stieltjes(v[2],Directed{false}(RiemannDual(0.0,1)))) ≈ c_vals22[end]

        C12 = fpstieltjesmatrix(space(v[2]), domain(v[1]), ncoefficients(v[1]), ncoefficients(v[2]))
        c_vals12 = C12*v[2].coefficients
        h = 0.00000001; @test stieltjes(v[2],-h) ≈ stieltjes(v[2],RiemannDual(0.0,-1))(h) atol=1E-6
        @test finitepart(stieltjes(v[2],RiemannDual(0.0,-1))) ≈ c_vals12[1]

        C = fpstieltjesmatrix(space(f), ns, ns)
        @test norm(C) ≤ 200

        @test C[1:44,1:2:end] == C11

        c_vals = C*coefficients(f)
        pts = RiemannHilbert.collocationpoints(space(f), ns)

        f_ex = Fun(x->exp(-40(x-0.1)^2), Legendre())
        @test stieltjes(f_ex, Directed{false}(0.0)) ≈ finitepart(stieltjes(f,Directed{false}(RiemannDual(0.0,-im))))
        @test finitepart(stieltjes(v[1],RiemannDual(0.0,-im))+stieltjes(v[2],RiemannDual(0.0,-im))) ≈ stieltjes(f_ex,Directed{false}(0.0))

        @test c_vals[1] ≈ finitepart(stieltjes(f,Directed{false}(RiemannDual(0.0,-im))))
        @test c_vals[1] ≈ finitepart(stieltjes(f,Directed{false}(RiemannDual(0.0,exp(-0.1im)))))
        @test c_vals[1] ≈ finitepart(stieltjes(f,RiemannDual(0.0,-im)))
        @test c_vals[2:ns[1]-1] ≈  stieltjes.(f,pts[2:ns[1]-1]⁻)
        @test c_vals[ns[1]] ≈ finitepart(stieltjes(f,RiemannDual(-1.0,-1.0)))

        @test c_vals[ns[1]+1] ≈ finitepart(stieltjes(f,RiemannDual(1.0,1.0)))
        @test c_vals[ns[1]+2:end-1] ≈  stieltjes.(f,pts[ns[1]+2:end-1]⁻)
        @test c_vals[end] ≈ finitepart(stieltjes(f,RiemannDual(0.0,-im)))

        h =0.00001
        @test stieltjes(v[1], Directed{false}(RiemannDual(0.0,-1.0))) ≈ stieltjes(v[1], RiemannDual(0.0,-1.0-eps()*im))

        @test finitepart(stieltjes(v[1], Directed{false}(RiemannDual(0.0,-1.0)))+ stieltjes(v[2], RiemannDual(0.0,-1.0))) ≈
            stieltjes(v[1], -0.00000000001im)+stieltjes(v[2], -0.00000000001im)

        @test finitepart(stieltjes(v[1], Directed{false}(RiemannDual(0.0,-1.0)))+ stieltjes(v[2], RiemannDual(0.0,-1.0))) ≈
            stieltjes(f, -0.00000000001im)

        @test c_vals[1] ≈ stieltjes(f, -0.0000000001im)
        @test c_vals[end] ≈ stieltjes(f, -0.0000000001im)

        @test finitepart(stieltjes(f, RiemannDual(0.0,-im))) ≈
            finitepart(stieltjes(v[1], RiemannDual(0.0,-im))+stieltjes(v[2], RiemannDual(0.0,-im)))

        C11 = fpstieltjesmatrix(space(v[1]), ncoefficients(v[1]), ncoefficients(v[1]))
        C12 = fpstieltjesmatrix(space(v[2]), domain(v[1]), ncoefficients(v[1]), ncoefficients(v[2]))

        @test finitepart(stieltjes(v[1], Directed{false}(RiemannDual(0.0,-1.0)))) ≈ (C11*coefficients(v[1]))[1]
        @test finitepart(stieltjes(v[2], RiemannDual(0.0,-1.0))) ≈ dotu(stieltjesmoment!(Array{ComplexF64}(undef,ncoefficients(v[2])), space(v[2]), orientedrightendpoint(domain(v[1])), finitepart),
                    coefficients(v[2]))

        @test C12[1,:] ≈ stieltjesmoment!(Array{ComplexF64}(undef,ncoefficients(v[2])), space(v[2]), orientedrightendpoint(domain(v[1])), finitepart)

        @test finitepart(stieltjes(v[2], RiemannDual(0.0,-1.0))) ≈ (C12*coefficients(v[2]))[1]
        @test C[1,:] ≈ interlace(C11[1,:], C12[1,:])
    end

    @testset "0..1 and 0..1" begin
        sp = Legendre(Segment(0 , -1)) ⊕ Legendre(0 .. 1)

        f = Fun(x->sign(x)*exp(-40(x-0.1)^2), sp)
        v = components(f)
        ns = ncoefficients.(v)
        C = fpstieltjesmatrix(space(f), ns, ns)
        @test norm(C) ≤ 200
        c_vals = C*coefficients(f)
        pts = RiemannHilbert.collocationpoints(space(f), ns)

        @test stieltjes(f,-1-eps()) ≈ stieltjes(f,RiemannDual(-1.0,-1.0)).c
        @test c_vals[1] ≈ finitepart(stieltjes(f,RiemannDual(-1.0,-1.0)))
        @test c_vals[2:ns[1]-1] ≈ stieltjes.(f,pts[2:ns[1]-1]⁻)
        @test c_vals[ns[1]] ≈ finitepart(stieltjes(f,RiemannDual(0.0,+im)))
        @test c_vals[ns[1]+1] ≈ finitepart(stieltjes(f,RiemannDual(1.0,1.0)))
        @test c_vals[ns[1]+2:end-1] ≈  stieltjes.(f,pts[ns[1]+2:end-1]⁻)
        @test c_vals[end] ≈ finitepart(stieltjes(f,RiemannDual(0.0,-im)))

        h =0.00001
        @test stieltjes(v[1], Directed{false}(RiemannDual(0.0,-1.0))) ≈ stieltjes(v[1], RiemannDual(0.0,-1.0+eps()*im))
        @test stieltjes(v[1], Directed{true}(RiemannDual(0.0,-1.0))) ≈ stieltjes(v[1], RiemannDual(0.0,-1.0-eps()*im))

        @test stieltjes(v[2], RiemannDual(0.0,-1.0))(h) ≈ stieltjes(v[2], -h) atol=1E-3

        @test stieltjes(v[1], Directed{true}(RiemannDual(0.0,-1.0)))(h) ≈ stieltjes(v[1], -h-eps()*im) atol=1E-3

        @test finitepart(stieltjes(v[1], Directed{true}(RiemannDual(0.0,-1.0)))+ stieltjes(v[2], RiemannDual(0.0,-1.0))) ≈
            stieltjes(v[1], -0.00000000001im)+stieltjes(v[2], -0.00000000001im)

        @test finitepart(stieltjes(v[1], Directed{true}(RiemannDual(0.0,-1.0)))+ stieltjes(v[2], RiemannDual(0.0,-1.0))) ≈
            stieltjes(f, -0.00000000001im)

        @test c_vals[ns[1]] ≈ stieltjes(f, +0.0000000001im)
        @test c_vals[end] ≈ stieltjes(f, -0.0000000001im)
    end

    @testset "-1..0 and 0..1" begin
        sp = Legendre(-1 .. 0) ⊕ Legendre(0 .. 1)
        f = Fun(x->exp(-40(x-0.1)^2), sp)
        v = components(f)
        ns = ncoefficients.(v)
        C = fpstieltjesmatrix(space(f), ns, ns)
        @test norm(C) ≤ 200

        c_vals = C*coefficients(f)
        pts = RiemannHilbert.collocationpoints(space(f), ns)

        @test c_vals[1] ≈ finitepart(stieltjes(f,RiemannDual(0.0,-im)))
        @test c_vals[2:ns[1]-1] ≈  stieltjes.(f,pts[2:ns[1]-1]⁻)
        @test c_vals[ns[1]] ≈ finitepart(stieltjes(f,RiemannDual(-1.0,-1.0)))
        @test c_vals[ns[1]+1] ≈ finitepart(stieltjes(f,RiemannDual(1.0,1.0)))
        @test c_vals[ns[1]+2:end-1] ≈  stieltjes.(f,pts[ns[1]+2:end-1]⁻)
        @test c_vals[end] ≈ finitepart(stieltjes(f,RiemannDual(0.0,-im)))
    end
end

@testset "ArraySpace" begin
    sp = ArraySpace(Legendre() ,2)
    f = Fun(x->[exp(-40(x-0.1)^2); cos(x-0.1)*exp(-40(x-0.1)^2)], sp)
    ns = ncoefficients.(Array(f))
    C = fpstieltjesmatrix(sp, ns, ns)
    C1 = fpstieltjesmatrix(sp[1], ns[1], ns[1])
    C2 = fpstieltjesmatrix(sp[2], ns[2], ns[2])

    @test C*coefficients(f) ≈ [C1*coefficients(f[1]) ; C2*coefficients(f[2])]

    sp = ArraySpace(Legendre(), 2, 2)
    f = Fun(x->[exp(-40(x-0.1)^2) sin(x-0.1)exp(-40(x-0.1)^2); cos(x-0.1)*exp(-40(x-0.1)^2) airyai(x-0.1)*exp(-40(x-0.1)^2)], sp)
    ns = ncoefficients.(Array(f))
    C = fpstieltjesmatrix(sp, ns, ns)
    C11 = fpstieltjesmatrix(sp[1,1], ns[1,1], ns[1,1])
    C21 = fpstieltjesmatrix(sp[2,1], ns[2,1], ns[2,1])
    C12 = fpstieltjesmatrix(sp[1,2], ns[1,2], ns[1,2])
    C22 = fpstieltjesmatrix(sp[2,2], ns[2,2], ns[2,2])

    @test C[1:ns[1], 1:4:end] ≈ C11
    @test norm(C[1:ns[1], 2:4:end]) ≤ 100eps()
    @test norm(C[1:ns[1], 3:4:end]) ≤ 100eps()
    @test norm(C[1:ns[1], 4:4:end]) ≤ 100eps()

    @test C[ns[1]+1:ns[1]+ns[2], 2:4:end] ≈ C21

    @test C*coefficients(f) ≈ [C11*coefficients(f[1,1]) ; C21*coefficients(f[2,1]) ; C12*coefficients(f[1,2]) ; C22*coefficients(f[2,2])]

    sp = ArraySpace(Legendre(-1 .. 0) ∪ Legendre(0 .. 1), 2)
    f = Fun(x->[exp(-40(x-0.1)^2); cos(x-0.1)*exp(-40(x-0.1)^2)], sp)
    ns = ncoefficients.(Array(f))
    C = fpstieltjesmatrix(sp, ns, ns)

    C1 = fpstieltjesmatrix(sp[1], ns[1], ns[1])
    C2 = fpstieltjesmatrix(sp[2], ns[2], ns[2])

    @test C*coefficients(f) ≈ [C1*coefficients(f[1]) ; C2*coefficients(f[2])]

    sp = ArraySpace(Legendre(-1 .. 0) ∪ Legendre(0 .. 1), 2, 2)
    f = Fun(x->[exp(-40(x-0.1)^2) sin(x-0.1)exp(-40(x-0.1)^2); cos(x-0.1)*exp(-40(x-0.1)^2) airyai(x-0.1)*exp(-40(x-0.1)^2)], sp)
    ns = ncoefficients.(Array(f))
    C = fpstieltjesmatrix(sp, ns, ns)

    C11 = fpstieltjesmatrix(sp[1,1], ns[1,1], ns[1,1])
    C21 = fpstieltjesmatrix(sp[2,1], ns[2,1], ns[2,1])
    C12 = fpstieltjesmatrix(sp[1,2], ns[1,2], ns[1,2])
    C22 = fpstieltjesmatrix(sp[2,2], ns[2,2], ns[2,2])

    @test C*coefficients(f) ≈ [C11*coefficients(f[1,1]) ; C21*coefficients(f[2,1]) ; C12*coefficients(f[1,2]) ; C22*coefficients(f[2,2])]
end

@testset "rhsolve" begin
    sp = Legendre()
    g = 1-0.3Fun(x->exp(-40x^2), sp)
    n = 2ncoefficients(g)
    g_v = RiemannHilbert.collocationvalues(g-1, n)
    u = Fun(sp, rhmatrix(g,n) \ g_v)

    @testset "-1 .. 1" begin
        n = 2ncoefficients(g)
        g = pad(g,n)
        C₋ = fpcauchymatrix(sp, n, n)
        pts = RiemannHilbert.collocationpoints(sp, n)
        @test C₋[2:end-1,:]*coefficients(g) ≈ cauchy.(g, pts[2:end-1]⁻)
        g_v = RiemannHilbert.collocationvalues(g-1, n)
        @test g_v ≈ g.(pts).-1
        G = Diagonal(g_v)

        @test G*g_v ≈ (g.(pts) .- 1).^2

        E = RiemannHilbert.evaluationmatrix(sp, pts, length(pts))
        @test E*coefficients(g) ≈ g.(pts)

        @test (g.(pts).-1).*(C₋*coefficients(g)) ≈ (g.(pts).-1).*cauchy.(g, pts.-0.000000001im)
        @test G*(C₋*coefficients(g)) ≈ (g.(pts).-1).*cauchy.(g, pts.-0.000000001im)

        L = E - G*C₋
        @test L*coefficients(g) ≈ g.(pts) - (g.(pts).-1).*cauchy.(g, pts.-0.000000001im)

        @test L == rhmatrix(g,n)

        φ = z -> 1 + cauchy(u,z)
        @test φ(0.1⁺)  ≈ g(0.1)φ(0.1⁻)

        @test 1+cauchy(u)(0.1+0.2im) ≈ φ(0.1+0.2im)
        @test (1+cauchy(u) )(0.1+0.2im) ≈ φ(0.1+0.2im)
        @test (1+cauchy(u) )(0.1⁻) ≈ φ(0.1⁻)
    end

    @testset "-1 .. 0 and 0 .. 1" begin
        sp = Legendre(-1 .. 0) ∪ Legendre(0 .. 1)
        u_1 = u
        u_ex = Fun(x->u_1(x), sp)
        g_1 = 1-0.3Fun(x->exp(-40x^2), Legendre())

        g = 1-0.3Fun(x->exp(-40x^2), sp)

        n = 2ncoefficients(g)
        g = pad(g,n)
        u_ex = pad(u_ex,n)
        @test (1+cauchy(u_ex,0.1⁺)) ≈ g(0.1)*(1+cauchy(u_ex,0.1⁻))

        C₋ = fpcauchymatrix(sp, n, n)
        pts = RiemannHilbert.collocationpoints(sp, n)

        @test transpose(C₋[1,:])*coefficients(u_ex) ≈ cauchy(u_ex,pts[1]-eps()im)
        @test transpose(C₋[5,:])*coefficients(u_ex) ≈ cauchy(u_ex,pts[5]-eps()im)
        @test transpose(C₋[end,:])*coefficients(u_ex) ≈ cauchy(u_ex, 0.0-eps()im)
        @test transpose(C₋[end-1,:])*coefficients(u_ex) ≈ cauchy(u_ex,pts[end-1]-eps()im)
        @test transpose(C₋[n÷2,:])*coefficients(u_ex) ≈ cauchy(u_ex,-1.0-eps())
        @test transpose(C₋[(n÷2)+1,:])*coefficients(u_ex) ≈ cauchy(u_ex,1.0+eps())

        @test C₋*coefficients(u_ex) ≈ cauchy.(u_ex, pts.-eps()im)
        g_v = RiemannHilbert.collocationvalues(g-1, n)
        @test g_v ≈ g.(pts).-1
        G = Diagonal(g_v)

        @test G*g_v ≈ (g.(pts) .- 1).^2

        g1 = component(g,2)
        pts1 = RiemannHilbert.collocationpoints(component(sp,2), n÷2)
        E1=RiemannHilbert.evaluationmatrix(component(sp,1), length(pts1))

        @test E1*g1.coefficients ≈ g1.(pts1)

        E = RiemannHilbert.evaluationmatrix(sp, n)
        @test E*coefficients(g) ≈ g.(pts)

        @test (g.(pts).-1).*(C₋*coefficients(g)) ≈ (g.(pts).-1).*cauchy.(g, pts.-0.000000001im)
        @test G*(C₋*coefficients(g)) ≈ (g.(pts).-1).*cauchy.(g, pts.-0.000000001im)

        L = E - G*C₋
        @test L*coefficients(g) ≈ g.(pts) - (g.(pts).-1).*cauchy.(g, pts.-0.000000001im)

        @test L == rhmatrix(g,n)

        u = Fun(sp, rhmatrix(g,n) \ g_v)
        φ = z -> 1 + cauchy(u,z)
        @test φ(0.1⁺)  ≈ g(0.1)φ(0.1⁻)

        @test 1+cauchy(u)(0.1+0.2im) ≈ φ(0.1+0.2im)
        @test (1+cauchy(u) )(0.1+0.2im) ≈ φ(0.1+0.2im)
        @test (1+cauchy(u) )(0.1⁻) ≈ φ(0.1⁻)

        @time φ = rhsolve(g, 2ncoefficients(g))
        @test φ(0.1⁺)  ≈ g(0.1)φ(0.1⁻)
    end

    @testset "0..-1 and 0..1" begin
        sp = Legendre(Segment(0 , -1)) ∪ Legendre(0 .. 1)
        g = Fun(x->x ≥ 0 ? 1-0.3exp(-40x^2) : inv(1-0.3exp(-40x^2)), sp)
        u_1 = u
        u_ex = Fun(x->sign(x)u_1(x), sp)

        n = 2ncoefficients(g)
        g = pad(g,n)
        u_ex = pad(u_ex,n)
        @test (1+cauchy(u_ex,0.1⁺)) ≈ g(0.1)*(1+cauchy(u_ex,0.1⁻))

        C₋ = fpcauchymatrix(sp, n, n)
        pts = RiemannHilbert.collocationpoints(sp, n)

        @test transpose(C₋[1,:])*coefficients(u_ex) ≈ cauchy(u_ex,pts[1]+eps()im)
        @test transpose(C₋[5,:])*coefficients(u_ex) ≈ cauchy(u_ex,pts[5]+eps()im)
        @test transpose(C₋[n÷2,:])*coefficients(u_ex) ≈ cauchy(u_ex,0.0+eps()im)
        @test transpose(C₋[(n÷2)+1,:])*coefficients(u_ex) ≈ cauchy(u_ex,1.0+eps())
        @test transpose(C₋[end,:])*coefficients(u_ex) ≈ cauchy(u_ex, pts[end]-eps()im)
        @test transpose(C₋[end,:])*coefficients(u_ex) ≈ cauchy(u_ex, 0.0-eps()im)
        @test transpose(C₋[end-1,:])*coefficients(u_ex) ≈ cauchy(u_ex,pts[end-1]-eps()im)

        @test C₋*coefficients(u_ex) ≈ [cauchy.(u_ex, pts[1:n÷2].+eps()im); cauchy.(u_ex, pts[(n÷2)+1:end].-eps()im)]
        g_v = RiemannHilbert.collocationvalues(g-1, n)

        g1 = component(g,1)
        pts1 = RiemannHilbert.collocationpoints(component(sp,1), n÷2)
        E1=RiemannHilbert.evaluationmatrix(component(sp,1), length(pts1))

        @test E1*g1.coefficients ≈ g1.(pts1)

        g2 = component(g,2)
        pts2 = RiemannHilbert.collocationpoints(component(sp,2), n÷2)

        @test g_v ≈ [g1.(pts1).-1; g2.(pts2).-1]
        G = Diagonal(g_v)

        @test g1.(pts1) ≈ g.(pts1)
        @test_broken g2.(pts2) ≈ g.(pts2) # evaluation at 0
        @test_broken G*g_v ≈ (g.(pts) .- 1).^2

        E = RiemannHilbert.evaluationmatrix(sp, n)
        @test E[1:end-1,:]*coefficients(g) ≈ g.(pts[1:end-1])
        @test transpose(E[end,:])*coefficients(g) ≈ component(g,2)(0.0)

        @test (g2.(pts2).-1).*(C₋[(n÷2)+1:end,:]*coefficients(u_ex)) ≈
            (g2.(pts2).-1).*cauchy.(u_ex, pts2.-eps()im)
        @test (g.(pts1).-1).*(C₋[1:n÷2,:]*coefficients(u_ex)) ≈
            (g.(pts1).-1).*cauchy.(u_ex, pts1.+0.000000001im)
        @test (g.(pts2).-1).*(C₋[(n÷2)+1:end,:]*coefficients(u_ex)) ≈
            (g.(pts2).-1).*cauchy.(u_ex, pts2.-0.000000001im)

        u_ex1,u_ex2 = components(u_ex)
        @test (C₋*coefficients(u_ex))[ncoefficients(u_ex1)+1:end] ≈ cauchy.(u_ex, pts2.-eps()im)
        @test G*C₋*coefficients(u_ex) ≈ [(g1.(pts1).-1).*cauchy.(u_ex, pts1.+eps()im);
                                         (g2.(pts2).-1).*cauchy.(u_ex, pts2.-eps()im)]

        L = E - G*C₋
        @test L*coefficients(u_ex) ≈ [u_ex1.(pts1) - (g1.(pts1).-1).*cauchy.(u_ex, pts1.+eps()im);
                                        u_ex2.(pts2) - (g2.(pts2).-1).*cauchy.(u_ex, pts2.-eps()im)]

        @test L == rhmatrix(g,n)

        u = Fun(sp, rhmatrix(g,n) \ g_v)
        φ1 = z -> 1 + cauchy(u,z)
        @test φ1(0.1⁺)  ≈ g(0.1)φ1(0.1⁻)
        @test φ1(-0.1⁺)  ≈ g(-0.1)φ1(-0.1⁻)

        @test 1+cauchy(u)(0.1+0.2im) ≈ φ1(0.1+0.2im)
        @test (1+cauchy(u) )(0.1+0.2im) ≈ φ1(0.1+0.2im)
        @test (1+cauchy(u) )(0.1⁻) ≈ φ1(0.1⁻)
        @test (1+cauchy(u) )(0.1⁺) ≈ φ1(0.1⁺)
        @test (1+cauchy(u) )(-0.1⁻) ≈ φ1(-0.1⁻)
        @test (1+cauchy(u) )(-0.1⁺) ≈ φ1(-0.1⁺)

        @time φ = rhsolve(g, 2ncoefficients(g))
        @test φ(0.1⁺)  ≈ g(0.1)φ(0.1⁻)
    end

    @testset "Chebyshev g" begin
        G = Fun(x -> 1 + exp(-30x^2), -1..1)
        Φ = rhsolve(G, 200)
        @test Φ(0.1⁺) ≈ G(0.1)Φ(0.1⁻)
    end
end

@testset "Matrix rhsolve" begin
    sp = ArraySpace(Legendre(), 2)
    f = Fun(Fun(x->[cos(x);sin(x)], Chebyshev()), ArraySpace(Legendre(), 2))
    G = Fun(Fun(x->[1 exp(-40x^2); 0.1exp(-40x^2) 1], Chebyshev()), ArraySpace(Legendre(), 2, 2))

    n = 2ncoefficients(G)
    E = RiemannHilbert.evaluationmatrix(sp, n)
    pts = RiemannHilbert.collocationpoints(sp, n÷2)

    @test E*coefficients(pad(f,n)) ≈ [f[1].(pts); f[2].(pts)]

    M = RiemannHilbert.multiplicationmatrix(G-I, n)
    @test M*(E*coefficients(pad(f,n))) ≈ mapreduce(f -> f.(pts), vcat, (G-I)*f)

    L = E - M*fpcauchymatrix(sp, n,n)
    u1 = L \ mapreduce(f -> f.(pts), vcat, (G-I)[:,1])
    u2 = L \ mapreduce(f -> f.(pts), vcat, (G-I)[:,2])
    z=2+I
        Φ = z -> I + [cauchy(Fun(sp, u1),z) cauchy(Fun(sp, u2),z)]

    @test Φ(0.1+eps()im) ≈ G(0.1)*Φ(0.1-eps()im)

    @test RiemannHilbert.rhmatrix(G, n) == L

    @test RiemannHilbert.collocationvalues(G-I, n) ≈
        [RiemannHilbert.collocationvalues((G-I)[:,1], n) RiemannHilbert.collocationvalues((G-I)[:,2], n)]

    Φ = rhsolve(G,n)
    @test G(0.1)*Φ(0.1⁻) ≈ Φ(0.1⁺)
end

@testset "Matrix rhsolve two interval" begin
    sp = ArraySpace(Legendre(-1 .. 0) ∪ Legendre(0 .. 1), 2)
    G = Fun(x->[1 exp(-40x^2); 0.1exp(-40x^2) 1], ArraySpace(sp[1], 2, 2))

    n=2ncoefficients(G)
    Φ = rhsolve(G,n)

    @test G(0.1)*Φ(0.1⁻) ≈ Φ(0.1⁺)
    U = RiemannHilbert.rh_sie_solve(G,n)

    U1 = U[:,1]
    g = G
    sp = space(g)[:,1]
    C₋ = fpcauchymatrix(sp, n, n)
    G = RiemannHilbert.multiplicationmatrix(g-I, n)
    E = RiemannHilbert.evaluationmatrix(sp, n)

    pts = RiemannHilbert.collocationpoints(sp, n÷2)
    @test last(E*coefficients(U1)) ≈ U1(0.00000000001)[2]
    @test last(C₋*coefficients(U1)) ≈ cauchy(U1[2], 0.00000000001-eps()*im)
    @test cauchy(U1[1], 0.00000000001-eps()*im) ≈ (C₋*coefficients(U1))[n÷2]
    @test last(G*C₋*coefficients(U1)) ≈ ((g(0.00000000001)-I)*cauchy(U1, 0.00000000001-eps()*im))[2]

    @test RiemannHilbert.collocationvalues((g-I)[:,1], n)[end] ≈ 0.1

    sp = ArraySpace(Legendre(Segment(0 , -1)) ∪ Legendre(0 .. 1), 2)
    G = g = Fun(x-> x ≥ 0 ? [1 exp(-40x^2); 0.1exp(-40x^2) 1] : inv([1 exp(-40x^2); 0.1exp(-40x^2) 1]), ArraySpace(sp[1], 2, 2))
    n=2ncoefficients(G)
    Ũ1 = Fun(x-> x ≥ 0 ? U1(x) : -U1(x), sp, n÷2)
    @test cauchy(Ũ1, 0.1⁻) ≈ cauchy(Ũ1, 0.1-eps()im)
    @test cauchy(Ũ1, 0.1+eps()im) ≈ cauchy(U1, 0.1+eps()im)
    @test cauchy(Ũ1, -0.1+eps()im) ≈ cauchy(U1, -0.1+eps()im)

    L = rhmatrix(G, n)
    rhs = RiemannHilbert.collocationvalues((g-I)[:,1], n)
    @test rhs[end] ≈ 0.1

    Φ = rhsolve(G,n)

    g = G
    sp = space(g)[:,1]
    C₋ = fpcauchymatrix(sp, n, n)
    G = RiemannHilbert.multiplicationmatrix(g-I, n)
    E = RiemannHilbert.evaluationmatrix(sp, n)

    pts = RiemannHilbert.collocationpoints(sp, n÷2)

    @test last(E*coefficients(Ũ1)) ≈ Ũ1(0.00000000001)[2]
    @test last(C₋*coefficients(Ũ1)) ≈ cauchy(Ũ1[2], 0.00000000001-eps()*im)
    @test cauchy(Ũ1[1], 0.00000000001-eps()*im) ≈ (C₋*coefficients(Ũ1))[n÷2]
    @test G[[n÷2,n],[n÷2,n]] ≈ (g(0.00000000001)-I)

    @test last(G*C₋*coefficients(Ũ1)) ≈ transpose((g(0.00000000001)-I)[2,:])*cauchy(Ũ1, 0.00000000001-eps()*im)

    @test transpose(L[end,:])*coefficients(Ũ1) ≈ Ũ1(0.00000000001)[2] - transpose((g(0.00000000001)-I)[2,:])*cauchy(Ũ1, 0.00000000001-eps()*im)

    @test U1(0.00000000001)[2] - transpose((g(0.00000000001)-I)[2,:])*cauchy(U1, 0.00000000001-eps()*im) ≈ 0.1
    @test g(0.1)*Φ(0.1⁻) ≈ Φ(0.1⁺)
end

@testset "4 rays" begin
    Γ = Segment(0, 2.5exp(im*π/6)) ∪ Segment(0, 2.5exp(5im*π/6)) ∪
            Segment(0, 2.5exp(-5im*π/6)) ∪ Segment(0, 2.5exp(-im*π/6))
    sp = ArraySpace(PiecewiseSpace(Legendre.(components(Γ))), 2,2)

    s₁ = im
    s₃ = -im

    G = Fun( z -> if angle(z) ≈ π/6
                        [1 0; s₁*exp(8im/3*z^3) 1]
                    elseif angle(z) ≈ 5π/6
                        [1 0; s₃*exp(8im/3*z^3) 1]
                    elseif angle(z) ≈ -π/6
                        [1 -s₃*exp(-8im/3*z^3); 0 1]
                    elseif angle(z) ≈ -5π/6
                        [1 -s₁*exp(-8im/3*z^3); 0 1]
                    end
                        , sp)

    Φ = transpose(rhsolve(transpose(G), 2*4*100))

    s=exp(im*π/6)
    @test Φ((s)⁻)*G(s) ≈ Φ((s)⁺)
    @test map(g->first(components(g)[1]), G)*map(g->first(components(g)[2]), G)*
        map(g->first(components(g)[3]), G)*map(g->first(components(g)[4]), G) ≈ Matrix(I,2,2)

    for _=1:5
        @test RiemannHilbert.evaluationmatrix(sp[:,1], 32) ≈ RiemannHilbert.evaluationmatrix(sp[:,1], 32)
        @test RiemannHilbert.rhmatrix(G, 32) ≈ RiemannHilbert.rhmatrix(G, 32)
    end

    @test RiemannHilbert.rhmatrix(transpose(G), 900) ≈ RiemannHilbert.rhmatrix(transpose(G), 900)
    U = RiemannHilbert.rh_sie_solve(transpose(G), 2*4*100)
    @test -0.36706155154807807 ≈ sum(U[1,2])/(-π*im)

    G̃ = Fun( z -> if angle(z) ≈ π/6
                        [1 0; s₁*exp(8im/3*z^3) 1]
                    elseif angle(z) ≈ 5π/6
                        [1 0; s₃*exp(8im/3*z^3) 1]
                    elseif angle(z) ≈ -π/6
                        [1 -s₃*exp(-8im/3*z^3); 0 1]
                    elseif angle(z) ≈ -5π/6
                        [1 -s₁*exp(-8im/3*z^3); 0 1]
                    end
                        , Γ)

    Ũ = RiemannHilbert.rh_sie_solve(transpose(G̃), 2*4*100)

    @test rhmatrix(transpose(G), 2*4*100) ≈ rhmatrix(transpose(G̃), 2*4*100)
    @test -0.36706155154807807 ≈ sum(Ũ[1,2])/(-π*im)

    V = istieltjes(Φ)
    x = Fun(domain(V))

    @test 10.0I + sum.(Array(V)) + stieltjes(x*V, 10.0) ≈ 10.0Φ(10.0)

    # check rhmatrix relationship for below
    U = V*(-2π*im)
    U1 = U[1,:]
    n= ncoefficients(U1)
    L = rhmatrix(transpose(G),n)
    vals = collocationvalues(transpose(G)-I, n)
    @test L*coefficients(U1) ≈ vals[:,1]
end

@testset "6 rays" begin
    @testset "HM on 6 rays" begin
        s₁,s₂,s₃ = -im,0,im
        @assert s₁ - s₂ + s₃ + s₁*s₂*s₃ ≈ 0

        # construct true solution using 4 rays
        Γ = Segment(0, 2.5exp(im*π/6)) ∪ Segment(0, 2.5exp(5im*π/6)) ∪
        Segment(0, 2.5exp(-5im*π/6)) ∪ Segment(0, 2.5exp(-im*π/6))
        sp = ArraySpace(PiecewiseSpace(Legendre.(components(Γ))), 2,2)

        G = Fun( z -> if angle(z) ≈ π/6
                            [1 0; s₁*exp(8im/3*z^3) 1]
                        elseif angle(z) ≈ 5π/6
                            [1 0; s₃*exp(8im/3*z^3) 1]
                        elseif angle(z) ≈ -π/6
                            [1 -s₃*exp(-8im/3*z^3); 0 1]
                        elseif angle(z) ≈ -5π/6
                            [1 -s₁*exp(-8im/3*z^3); 0 1]
                        end
                            , sp)

        Φ = transpose(rhsolve(transpose(G), 2*6*100))

        # debug
        V4 = istieltjes(Φ)
        x = 0.0

        @test stieltjes(V4,1+im)+I == Φ(1+im)

        Γ = Segment(0, 2.5exp(im*π/6))   ∪
        Segment(0, 2.5exp(im*π/2))       ∪
        Segment(0, 2.5exp(5im*π/6))      ∪
        Segment(0, 2.5exp(-5im*π/6))     ∪
        Segment(0, 2.5exp(-im*π/2))      ∪
        Segment(0, 2.5exp(-im*π/6));

        G = Fun( z -> if angle(z) ≈ π/6
                        [1                             0;
                        s₁*exp(8im/3*z^3+2im*x*z)     1]
                    elseif angle(z) ≈ π/2
                        [1                 s₂*exp(-8im/3*z^3-2im*x*z);
                        0                 1]
                    elseif angle(z) ≈ 5π/6
                        [1                             0;
                        s₃*exp(8im/3*z^3+2im*x*z)     1]
                    elseif angle(z) ≈ -π/6
                        [1                -s₃*exp(-8im/3*z^3-2im*x*z);
                        0                 1]
                    elseif angle(z) ≈ -π/2
                        [1                             0;
                        -s₂*exp(8im/3*z^3+2im*x*z)    1]
                    elseif angle(z) ≈ -5π/6
                        [1                -s₁*exp(-8im/3*z^3-2im*x*z);
                        0                 1]
                    end
                        , Γ);
        sp = ArraySpace(PiecewiseSpace(Legendre.(components(Γ))), 2,2)
        V = Fun(V4, sp)

        @test stieltjes(V,1+im)+I ≈ Φ(1+im)

        U = V*(-2π*im)

        U1 = U[1,:]
        @test cauchy(U1,1+im)+[1,0] ≈ Φ(1+im)[1,:]
        @test abs(sum(first.(components(U1[1])))) ≤ 100eps()
        @test abs(sum(first.(components(U1[2])))) ≤ 100eps()

        U11 = U1[1]
        n = ncoefficients(U11)
        C₋ = fpcauchymatrix(space(U11), n, n)
        pts = collocationpoints(space(U11), n)
        c = C₋*coefficients(U11)

        @test c[1] ≈ finitepart(cauchy(U11,orientedrightendpoint(component(Γ,1))))
        @test c[2] ≈ cauchy(U11,pts[2])
        @test c[150] ≈ finitepart(cauchy(U11,orientedleftendpoint(component(Γ,1))⁻))
        @test c[151] ≈ finitepart(cauchy(U11,orientedrightendpoint(component(Γ,2))))

        n = ncoefficients(U1)
        L = rhmatrix(transpose(G),n)
        vals = collocationvalues(transpose(G)-I, n)
        pts = collocationpoints(space(U1),n)

        c = L*coefficients(U1)
        @test c ≈ vals[:,1]

        @test coefficients(U1) ≈ L \ vals[:,1]
    end
end

include("test_nls.jl")


using SeisIO, JLD2, Wavelets, Statistics, DSP, FFTW, Plots, GLM, DataFrames

include("wavelet_methods.jl")

function test_WXS(finame::String, foname::String, InputDict::Dict, type::String; plot::Bool=true)
    valData = jldopen(finame)
    dvVlist = valData["info/dvVlist"]
    noiselist = valData["info/noiselist"]

    fmin    = InputDict["fmin"]
    fmax    = InputDict["fmax"]
    tmin    = InputDict["tmin"]
    tmax    = InputDict["tmax"]
    dt      = InputDict["dt"]
    dj      = InputDict["dj"]
    s0      = InputDict["s0"]
    J       = InputDict["J"]
    wvn     = InputDict["wvn"]
    sig     = InputDict["sig"]
    fbands  = InputDict["fbands"]
    unwrap  = InputDict["unwrap"]
    mintime = InputDict["mintime"]
    winflag = InputDict["winflag"]
    nwindow = InputDict["nwindow"]
    siglvl  = InputDict["siglvl"]
    twindow = [tmin, tmax]
    fwindow = [fmin, fmax]

    # output file
    f_out = jldopen(foname, "w")

    for dvV in dvVlist
        println("dv/v: $dvV")
        for noiselvl in noiselist
            signals = valData["$type/$dvV.$noiselvl"]
            ref = signals["ref"]
            cur = signals["cur"]

            npts = length(cur)
            tvec = (collect(0:npts-1) .* dt) .+ mintime

            # Wavelet cross-spectrum on discrete frequency bands
            (freqbands, dvv, err) = wxs_freqbands(cur, ref, tvec, twindow, fbands, dj, s0, J, wvn=wvn, unwrapflag=unwrap, sig=sig, siglvl=siglvl, windowflag=winflag, nwindow=nwindow)
            f_out["bands/$dvV/$noiselvl"] = [freqbands, dvv, err]
            p=PlotlyJS.plot(dvv, mode="lines+markers")
            display(p)
            readline()

            # Wavelet cross-spectrum for each frequency
            (freqs, dvv, err) = wxs_allfreq(cur, ref, tvec, twindow, fwindow, dj, s0, J, wvn=wvn, unwrapflag=unwrap, sig=sig, siglvl=siglvl, windowflag=winflag, nwindow=nwindow)
            f_out["full/$dvV/$noiselvl"] = [freqbands, dvv, err]
            p=PlotlyJS.plot(freqs, dvv, mode="lines+markers")
            display(p)
            readline()
        end
    end
    close(f_out)
end

function test_WTS(finame::String, foname::String, InputDict::Dict, type::String, plot::Bool=true)
    valData = jldopen(finame)
    dvVlist = valData["info/dvVlist"]
    noiselist = valData["info/noiselist"]

    fmin    = InputDict["fmin"]
    fmax    = InputDict["fmax"]
    tmin    = InputDict["tmin"]
    tmax    = InputDict["tmax"]
    dt      = InputDict["dt"]
    dj      = InputDict["dj"]
    s0      = InputDict["s0"]
    J       = InputDict["J"]
    wvn     = InputDict["wvn"]
    ndv     = InputDict["ndv"]
    norm    = InputDict["norm"]
    dvmax   = InputDict["dvmax"]
    fbands  = InputDict["fbands"]
    mintime = InputDict["mintime"]
    winflag = InputDict["winflag"]
    nwindow = InputDict["nwindow"]

    twindow = [tmin, tmax]
    fwindow = [fmin, fmax]

    # output file
    f_out = jldopen(foname, "w")

    for dvV in dvVlist
        println("dv/v: $dvV")
        for noiselvl in noiselist
            signals = valData["$type/$dvV.$noiselvl"]
            ref = signals["ref"]
            cur = signals["cur"]

            npts = length(cur)
            tvec = (collect(0:npts-1) .* dt) .+ mintime

            # Wavelet trace stretching on discrete frequency bands
            (freqbands, dvv, cc, cdp, err) = wts_freqbands(cur, ref, tvec, twindow, fbands, dj, s0, J, dvmax=dvmax, ndv=ndv, wvn=wvn, normalize=norm, windowflag=winflag, nwindow=nwindow)
            f_out["bands/$dvV/$noiselvl"] = [freqbands, dvv, err, cc, cdp]

            # Wavelet trace stretching for each frequency
            (freqs, dvv, cc, cdp, err) = wts_allfreqs(cur, ref, tvec, twindow, fbands, dj, s0, J, dvmax=dvmax, ndv=ndv, wvn=wvn, normalize=norm, windowflag=winflag, nwindow=nwindow)
            f_out["full/$dvV/$noiselvl"] = [freqs, dvv, err, cc, cdp]
        end
    end
    close(f_out)
end

function test_WTDTW(finame::String, foname::String, InputDict::Dict, type::String; plot::Bool=true)
    valData = jldopen(finame)
    dvVlist = valData["info/dvVlist"]
    noiselist = valData["info/noiselist"]

    fmin    = InputDict["fmin"]
    fmax    = InputDict["fmax"]
    tmin    = InputDict["tmin"]
    tmax    = InputDict["tmax"]
    dt      = InputDict["dt"]
    dj      = InputDict["dj"]
    s0      = InputDict["s0"]
    J       = InputDict["J"]
    b       = InputDict["b"]
    wvn     = InputDict["wvn"]
    ndv     = InputDict["ndv"]
    dir     = InputDict["dir"]
    norm    = InputDict["norm"]
    dvmax   = InputDict["dvmax"]
    fbands  = InputDict["fbands"]
    maxLag  = InputDict["maxLag"]
    winflag = InputDict["winflag"]
    mintime = InputDict["mintime"]
    nwindow = InputDict["nwindow"]

    twindow = [tmin, tmax]
    fwindow = [fmin, fmax]

    # output file
    f_out = jldopen(foname, "w")

    for dvV in dvVlist
        println("dv/v: $dvV")
        for noiselvl in noiselist
            signals = valData["$type/$dvV.$noiselvl"]
            ref = signals["ref"]
            cur = signals["cur"]

            npts = length(cur)
            tvec = (collect(0:npts-1) .* dt) .+ mintime

            # Wavelet dynamic time warping on discrete frequency bands
            (freqbands, dvv, err) = wtdtw_freqbands(cur, ref, tvec, twindow, fbands, maxLag, b, dir, dj, s0, J, wvn=wvn, normalize=norm, windowflag=winflag, nwindow=nwindow)
            f_out["bands/$dvV/$noiselvl"] = [freqbands, dvv, err]

            # Wavelet dynamic time warping for each frequency
            (freqs, dvv, err) = wtdtw_allfreqs(cur, ref, tvec, twindow, fwindow, maxLag, b, dir, dj, s0, J, wvn=wvn, normalize=norm, windowflag=winflag, nwindow=nwindow)
            f_out["full/$dvV/$noiselvl"] = [freqs, dvv, err]
        end
    end
    close(f_out)
end

finame = "/Users/jared/SCECintern2019/data/validation/verificationData.jld2"
foname_WXS = "verificationData_WXS.jld2"
foname_WTS = "verificationData_WTS.jld2"
foname_DTW = "verificationData_DTW.jld2"

InputDict_WXS = Dict( "fmin"    => 0.0,
                      "fmax"    => 10.0,
                      "tmin"    => 0.0,
                      "tmax"    => 200.0,
                      "dt"      => 0.05,
                      "dj"      => 1/12,
                      "s0"      => -1,
                      "J"       => -1,
                      "wvn"     => "morlet",
                      "sig"     => false,
                      "unwrap"  => false,
                      "mintime" => 0,
                      "siglvl"  => 0.95,
                      "fbands"  => [0.1 0.2; 0.2 0.4; 0.4 0.8; 0.8 1.6; 1.6 3.2],
                      "winflag" => true,
                      "nwindow" => 1.5)

InputDict_WTS = Dict( "fmin"    => 0.0,
                      "fmax"    => 10.0,
                      "tmin"    => 0.0,
                      "tmax"    => 200.0,
                      "dt"      => 0.05,
                      "dj"      => 1/12,
                      "s0"      => -1,
                      "J"       => -1,
                      "wvn"     => "morlet",
                      "ndv"     => 100,
                      "dvmax"   => 0.1,
                      "mintime" => 0,
                      "norm"    => true,
                      "fbands"  => [0.1 0.2; 0.2 0.4; 0.4 0.8; 0.8 1.6; 1.6 3.2],
                      "winflag" => true,
                      "nwindow" => 1.5)

InputDict_DTW = Dict( "fmin"    => 0.0,
                      "fmax"    => 10.0,
                      "tmin"    => 0.0,
                      "tmax"    => 200.0,
                      "dt"      => 0.05,
                      "dj"      => 1/12,
                      "s0"      => -1,
                      "b"       => 1,
                      "J"       => -1,
                      "wvn"     => "morlet",
                      "dir"     => 0,
                      "ndv"     => 100,
                      "dvmax"   => 0.1,
                      "mintime" => 0,
                      "norm"    => true,
                      "maxLag"  => 150,
                      "fbands"  => [0.1 0.2; 0.2 0.4; 0.4 0.8; 0.8 1.6; 1.6 3.2],
                      "winflag" => true,
                      "nwindow" => 1.5)

for type in ["spectSynth", "rickerConv", "dampedSin", "realData"]
    test_WXS(finame, foname_WXS, InputDict_WXS, type)
    test_WTS(finame, foname_WTS, InputDict_WTS, type)
    test_WTDTW(finame, foname_DTW, InputDict_DTW, type)
end


#---------------------------------------------------------
# # [ImageGeoms overview](@id 1-overview)
#---------------------------------------------------------

# This page explains the Julia package
# [`ImageGeoms`](https://github.com/JuliaImageRecon/ImageGeoms.jl).

# ### Setup

# Packages needed here.

using ImageGeoms
using MIRTjim: jim, prompt
using Plots: scatter, plot!, default; default(markerstrokecolor=:auto)

# The following line is helpful when running this example.jl file as a script;
# this way it will prompt user to hit a key after each figure is displayed.

isinteractive() ? jim(:prompt, true) : prompt(:draw);

# ### Overview

# When performing tomographic image reconstruction,
# one must specify the geometry of the grid of image pixels.
# (In contrast, for image denoising and image deblurring problems,
# one works with the given discrete image
# and no physical coordinates are needed.)

# The key parameters of a grid of image pixels are
# * the size (dimensions) of the grid, e.g., `128 × 128`,
# * the spacing of the pixels, e.g., `1mm × 1mm`,
# * the offset of the pixels relative to the origin, e.g., `(0,0)`

# The data type `ImageGeom` describes such a geometry,
# for arbitrary dimensions (2D, 3D, etc.).

# There are several ways to construct this structure.
# The default is a `128 × 128` grid
# with pixel size ``\Delta_X = \Delta_Y = 1`` (unitless) and zero offset:

ig = ImageGeom()

# Here is a 3D example with non-cubic voxel size:

ig = ImageGeom( (512,512,128), (1,1,2), (0,0,0) )

# To avoid remembering the order of the arguments,
# named keyword pairs are also supported:

ig = ImageGeom( dims=(512,512,128), deltas=(1,1,2), offsets=(0,0,0) )

# ### Units

# The pixel dimensions `deltas` can (and should!) be values with units.

# Here is an example for a video (2D+time) with 12 frames per second:
using UnitfulRecipes
using Unitful: mm, s

ig = ImageGeom( dims=(640,480,1000), deltas=(1mm,1mm,(1//12)s) )

# ### Methods

# An ImageGeom object has quite a few methods;
# `axes` and `axis` are especially useful:

ig = ImageGeom( dims=(7,8), deltas=(3,2), offsets=(0,0.5) )
axis(ig, 2)

# For an axis of length `n` with spacing `Δ` (possibly with units)
# and (always unitless but possibly non-integer) `offset` the axis
# is a subtype of `AbstractRange` of the form
# `( (0:n-1) .- ((n - 1)/2 + offset) ) * Δ`

# These axes are useful for plotting:
ig = ImageGeom( dims=(12,10), deltas=(1mm,1mm), offsets=(0.5,0.5) )

#
showgrid = (ig) -> begin # x,y grid locations of pixel centers
    x = axis(ig, 1)
    y = axis(ig, 2)
    (xg, yg) = grids(ig)
    scatter(xg, yg, label="", xlabel="x", ylabel="y",
        xlims = maximum(abs, x) * 1.2 .* (-1,1),
        xticks = [x[1], zero(eltype(x)), x[end]],
        ylims = maximum(abs, y) * 1.2 .* (-1,1),
        yticks = [y[1], zero(eltype(y)), y[end]],
        aspect_ratio=1, title="offsets $(ig.offsets)")
end
showgrid(ig)

# Unit labels on the axes due to `UnitfulRecipes.jl`
prompt();

# ### Offsets (unitless translation of grid)

# The default `offsets` are zeros,
# corresponding to symmetric sampling around origin:

ig = ImageGeom( dims=(12,10), deltas=(1mm,1mm) )
p = showgrid(ig)

#
prompt();

# That default for `offsets` is natural for tomography
# when considering finite pixel size:

square = (x,y,Δ) -> plot!(p, label="", color=:black,
    x .+ Δ[1] * ([0,1,1,0,0] .- 0.5),
    y .+ Δ[2] * ([0,0,1,1,0] .- 0.5),
)
square2 = (x,y) -> square(x, y, ig.deltas)
square2.(grids(ig)...)
plot!(p)

#
prompt();

# In that default geometry, the center `(0,0)` of the image
# is at a corner of the middle 4 pixels (for even image sizes).
# That default is typical for tomographic imaging (e.g., CT, PET, SPECT).
# One must be careful when using operations like `imrotate` or `fft`.

# ### AxisArrays

# There is a natural connection between `ImageGeom` and `AxisArrays`.
# Note the automatic labeling of units (when relevant) on all axes by
# [MIRTjim.jim](https://github.com/JeffFessler/MIRTjim.jl).

using AxisArrays
using Unitful: mm
ig = ImageGeom( dims=(60,48), deltas=(1.5mm,1mm) )
za = AxisArray( ellipse(ig) * 10/mm ; x=axis(ig,1), y=axis(ig,2) )
jim(za, "AxisArray example")

#
prompt();

# ### Resizing

# Often we have a target grid in mind but want coarser sampling for debugging.
# The `downsample` method is useful for this.

ig = ImageGeom( dims = (512,512), deltas = (500mm,500mm) ./ 512 )
ig_down = downsample(ig, 4)

# Other times we want to avoid an "inverse crime" by using finer sampling
# to simulate data; use `oversample` for this.

ig_over = oversample(ig, (2,2))

# Just for debugging:
using Pkg
Pkg.status("ImageGeoms")


################################################################################
### For generic methods
################################################################################

# Heuristic for selecting the level with the (presumably) most reliable
# standard error estimate:
# Take the highest lvl with at least 32 bins.
# (Chose 32 based on https://doi.org/10.1119/1.3247985)
function _reliable_level(B::LogBinner{T,N})::Int64 where {T,N}
    isempty(B) && (return 1)                # results in NaN in std_error
    i = findlast(x -> x.count >= 32, B.accumulators)
    return something(i, 1)
end

_eachlevel(B::LogBinner) = 1:findlast(acc -> acc.count > 1, B.accumulators)

################################################################################
### Statistics
################################################################################

"""
    mean(B::LogBinner[, lvl])

Calculates the mean for a given level in the Binning Analysis.
"""
mean(B::LogBinner, lvl = 1) = mean(B.accumulators[lvl])

"""
    var(B::LogBinner[, lvl])

Calculates the variance of a given level in the Binning Analysis.
"""
var(B::LogBinner, lvl = _reliable_level(B)) = var(B.accumulators[lvl])

"""
    varN(B::LogBinner[, lvl])

Calculates the variance/N of a given level in the Binning Analysis.
"""
varN(B::LogBinner, lvl = _reliable_level(B)) = varN(B.accumulators[lvl])

################################################################################
### Experimental
################################################################################

"""
    convergence(B::LogBinner, lvl)

Computes the difference between the variance of this lvl and the last,
normalized to the last lvl. If this value tends to 0, the Binning Analysis has
converged.
"""
function convergence(B::LogBinner) end

function convergence(B::LogBinner{T,N}, lvl = _reliable_level(B)) where {N, T <: Number}
    return abs((varN(B, lvl+1) - varN(B, lvl)) / varN(B, lvl))
end
function convergence(B::LogBinner{T,N}, lvl = _reliable_level(B)) where {N, T <: AbstractArray}
    return mean(abs.((varN(B, lvl+1) .- varN(B, lvl)) ./ varN(B, lvl)))
end

"""
    has_converged(B::LogBinner, lvl[, threshhold = 0.05])

Returns true if the Binning Analysis has converged for a given lvl.
"""
function has_converged(B::LogBinner, lvl = _reliable_level(B), threshhold::Float64 = 0.05)
    return convergence(B, lvl) <= threshhold
end


using NeutralLandscapes, Test

algorithms = (
   DiamondSquare(),
   DiamondSquare(; H=0.2),
   DiamondSquare(0.4),
   DiscreteVoronoi(),
   DiscreteVoronoi(; n=2),
   DiscreteVoronoi(3),
   DistanceGradient(),
   DistanceGradient(; sources=[2]),
   DistanceGradient([1, 2]),
   EdgeGradient(),
   EdgeGradient(; direction=120.0),
   EdgeGradient(90.0),
   MidpointDisplacement(),
   MidpointDisplacement(;),
   NearestNeighborElement(),
   NearestNeighborElement(; k=1),
   NearestNeighborElement(2),
   NearestNeighborCluster(),
   # NearestNeighborCluster(; n=:queen),
   NearestNeighborCluster(; n=:diagonal),
   NearestNeighborCluster(0.2, :diagonal),
   NoGradient(),
   PerlinNoise(),
   PerlinNoise(; octaves=3, lacunarity=3, persistance=0.3),
   PerlinNoise((1, 2), 3, 2, 0.2),
   PlanarGradient(),
   PlanarGradient(; direction=120.0),
   PlanarGradient(90),
   RectangularCluster(),
   RectangularCluster(; maximum=5),
   RectangularCluster(2),
   WaveSurface(),
   WaveSurface(90.0),
   WaveSurface(; periods=2),
)

sizes = (50, 200), (100, 100), (301, 87)

for alg in algorithms, sze in sizes
    A = rand(alg, sze)
    @test size(A) == sze
end


using LinearAlgebra, OrdinaryDiffEq, Test, PreallocationTools, LabelledArrays, RecursiveArrayTools

#Base array
function foo(du, u, (A, tmp), t)
    tmp = get_tmp(tmp, u)
    mul!(tmp, A, u)
    @. du = u + tmp
    nothing
end
#with defined chunk_size
chunk_size = 5
u0 = ones(5, 5)
A = ones(5,5)
cache = dualcache(zeros(5,5), chunk_size)
prob = ODEProblem(foo, u0, (0., 1.0), (A, cache))
sol = solve(prob, TRBDF2(chunk_size=chunk_size))
@test sol.retcode == :Success

#with auto-detected chunk_size
prob = ODEProblem(foo, ones(5, 5), (0., 1.0), (ones(5,5), dualcache(zeros(5,5))))
sol = solve(prob, TRBDF2())
@test sol.retcode == :Success

#Base array with LBC
function foo(du, u, (A, lbc), t)
    tmp = lbc[u]
    mul!(tmp, A, u)
    @. du = u + tmp
nothing
end
prob = ODEProblem(foo, ones(5, 5), (0., 1.0), (ones(5,5), LazyBufferCache()))
sol = solve(prob, TRBDF2())
@test sol.retcode == :Success

#LArray
A = LArray((2,2); a=1.0, b=1.0, c=1.0, d=1.0)
c = LArray((2,2); a=0.0, b=0.0, c=0.0, d=0.0)
u0 = LArray((2,2); a=1.0, b=1.0, c=1.0, d=1.0)
function foo(du, u, (A, tmp), t)
    tmp = get_tmp(tmp, u)
    mul!(tmp, A, u)
    @. du = u + tmp
    nothing
end
#with specified chunk_size
chunk_size = 4
prob = ODEProblem(foo, u0, (0., 1.0), (A, dualcache(c, chunk_size)))
sol = solve(prob, TRBDF2(chunk_size = chunk_size))
@test sol.retcode == :Success
#with auto-detected chunk_size
prob = ODEProblem(foo, u0, (0., 1.0), (A, dualcache(c)))
sol = solve(prob, TRBDF2())
@test sol.retcode == :Success


include("DSGEmoments.jl")
include("DSGEmodel.jl") # defines prior and log-likelihood
using LinearAlgebra, Statistics, DelimitedFiles
function main()
    data = readdlm("dsgedata.txt")
    tuning = [0.001, 0.15, 0.1, 0.006, 0.15, 0.004, 0.005] # fix this somehow
    include("parameters.jl")
    lb = lb_param_ub[:,1]
    ub = lb_param_ub[:,3]
    θtrue = lb_param_ub[:,2]
    θinit = (ub + lb) / 2.0 # prior mean to start
    lnL = θ -> logL(θ, data)
    Prior = θ -> prior(θ) # uniform, doesn't matter
    # define things for MCMC
    verbosity = true
    burnin = 10000
    ChainLength = 50000
    # initial proposal moves one at a time
    Proposal = θ -> proposal1(θ, tuning)
    tuning = [0.0001, 0.1, 0.01, 0.001, 0.05, 0.001, 0.001]
    chain = mcmc(θinit, ChainLength, burnin, Prior, lnL, Proposal, verbosity)
    # keep every 10th
    i = 1:size(chain,1)
    keep = mod.(i,10.0).==0
    chain = chain[keep,:]
    # now use a MVN random walk proposal
    Σ = cov(chain[:,1:7])
    tuning = 0.1
    for j = 1:10
        P = (cholesky(Σ)).U
        Proposal = θ -> proposal2(θ,tuning*P)
        θinit = vec(mean(chain[:,1:7],dims=1))
        if j == 10
            ChainLength = 200000
        end
        chain = mcmc(θinit, ChainLength, 0, Prior, lnL, Proposal, verbosity)
        accept = mean(chain[:,end])
        #println("tuning: ", tuning, "  acceptance rate: ", accept)
        if accept > 0.35
            tuning *= 1.5
        elseif accept < 0.25
            tuning *= 0.5
        end
        # keep every 10th
        i = 1:size(chain,1)
        keep = mod.(i,10.0).==0
        θinit = vec(mean(chain[:,1:7],dims=1))
        Σ = 0.5*Σ + 0.5*cov(chain[:,1:7])
    end
    # keep every 10th to reduce autocorrelation
    i = 1:size(chain,1)
    j = mod.(i,10.0).==0
    chain = chain[j,:]
    # plain MCMC fit
    posmean = vec(mean(chain[:,1:7],dims=1))
    inci = zeros(7)
    for i = 1:7
        lower = quantile(chain[:,i],0.05)
        upper = quantile(chain[:,i],0.95)
        inci[i] = θtrue[i] >= lower && θtrue[i] <= upper
    end
    p = npdensity(chain[:,1]) # example of posterior plot
    display(p)
    prettyprint([posmean inci])
    return chain
end
main();