export test_jacobian, test_gradients, linedata, lineview, meshdata,
    meshview, gradview, jacview, PCLview, pcolormeshview,
    animate, saveanim, test_hessian, require_optim

function require_pyplot()
    if !isdefined(Main, :PyPlot)
        error("You must load PyPlot to use this function, e.g., `using PyPlot` or `import PyPlot`")
    end
    Main.PyPlot
end

function require_optim()
    if !isdefined(Main, :Optim)
        error("You must load Optim to use this function, e.g., `import Optim` or `using Optim`")
    end
    Main.Optim
end

"""
    test_gradients(f::Function, x0::Array{Float64, 1}; scale::Float64 = 1.0, showfig::Bool = true)

Testing the gradients of a vector function `f`:
`y, J = f(x)` where `y` is a scalar output and `J` is the vector gradient.
"""
function test_gradients(f::Function, x0::Array{Float64, 1};
    scale::Float64 = 1.0, showfig::Bool = true, mpi::Bool = false)

    v0 = rand(Float64,length(x0))
    γs = scale ./10 .^(1:5)
    err2 = Float64[]
    err1 = Float64[]
    f0, J = f(x0)
    for i = 1:5
        f1, _ = f(x0+γs[i]*v0)
        push!(err1, abs(f1-f0))
        push!(err2, abs(f1-f0-γs[i]*sum(J.*v0)))
    end
    if showfig
        if mpi && mpi_rank()>0
            return
        end
        mp = require_pyplot()
        mp.close("all")
        mp.loglog(γs, err1, label="Finite Difference")
        mp.loglog(γs, err2, label="Automatic Differentiation")
        mp.loglog(γs, γs * 0.5*abs(err1[1])/γs[1], "--",label="\$\\mathcal{O}(\\gamma)\$")
        mp.loglog(γs, γs.^2 * 0.5*abs(err2[1])/γs[1]^2, "--",label="\$\\mathcal{O}(\\gamma^2)\$")
        mp.plt.gca().invert_xaxis()
        mp.legend()
        mp.savefig("test.png")
        @info "Results saved to test.png"
        println("Finite difference: $err1")
        println("Automatic differentiation: $err2")
    end
    return err1, err2
end

"""
    test_jacobian(f::Function, x0::Array{Float64, 1}; scale::Float64 = 1.0, showfig::Bool = true)

Testing the gradients of a vector function `f`:
`y, J = f(x)` where `y` is a vector output and `J` is the Jacobian.
"""
function test_jacobian(f::Function, x0::Array{Float64, 1}; scale::Float64 = 1.0, showfig::Bool = true)

    v0 = rand(Float64,size(x0))
    γs = scale ./10 .^(1:5)
    err2 = Float64[]
    err1 = Float64[]
    f0, J = f(x0)
    for i = 1:5
        f1, _ = f(x0+γs[i]*v0)
        push!(err1, norm(f1-f0))
        push!(err2, norm(f1-f0-γs[i]*J*v0))
    end
    if showfig
        mp = require_pyplot()
        mp.close("all")
        mp.loglog(γs, err1, label="Finite Difference")
        mp.loglog(γs, err2, label="Automatic Differentiation")
        mp.loglog(γs, γs * 0.5*abs(err1[1])/γs[1], "--",label="\$\\mathcal{O}(\\gamma)\$")
        mp.loglog(γs, γs.^2 * 0.5*abs(err2[1])/γs[1]^2, "--",label="\$\\mathcal{O}(\\gamma^2)\$")
        mp.plt.gca().invert_xaxis()
        mp.legend()
        mp.savefig("test.png")
        @info "Results saved to test.png"
        println("Finite difference: $err1")
        println("Automatic differentiation: $err2")
    end
    return err1, err2
end

"""
    test_hessian(f::Function, x0::Array{Float64, 1}; scale::Float64 = 1.0)

Testing the Hessian of a scalar function `f`:
`g, H = f(x)` where `y` is a scalar output, `g` is a vector gradient output, and `H` is the Hessian.
"""
test_hessian(f::Function, x0::Array{Float64, 1}; kwargs...) = test_jacobian(f, x0;  kwargs...)

function linedata(θ1, θ2=nothing; n::Integer = 10)
    if θ2 === nothing
        θ2 = θ1 .* (1 .+ randn(size(θ1)...))
    end
    T = []
    for x in LinRange(0.0,1.0,n)
        push!(T, θ1*(1-x)+θ2*x)
    end
    T
end

function lineview(losses::Array{Float64})
    mp = require_pyplot()
    n = length(losses)
    v = collect(LinRange(0.0,1.0,n))
    mp.close("all")
    mp.plot(v, losses)
    mp.xlabel("\$\\alpha\$")
    mp.ylabel("loss")
    mp.grid("on")
end

@doc raw"""
    lineview(sess::PyObject, pl::PyObject, loss::PyObject, θ1, θ2=nothing; n::Integer = 10)

Plots the function

$$h(α) = f((1-α)θ_1 + αθ_2)$$

# Example
```julia
pl = placeholder(Float64, shape=[2])
l = sum(pl^2-pl*0.1)
sess = Session(); init(sess)
lineview(sess, pl, l, rand(2))
```
"""
function lineview(sess::PyObject, pl::PyObject, loss::PyObject, θ1, θ2=nothing; n::Integer = 10)
    mp = require_pyplot()
    dat = linedata(θ1, θ2, n=n)
    V = zeros(length(dat))
    for i = 1:length(dat)
        @info i, n
        V[i] = run(sess, loss, pl=>dat[i])
    end
    lineview(V)
end

function meshdata(θ;
     a::Real=1, b::Real=1, m::Integer=10, n::Integer=10 ,direction::Union{Array{<:Real}, Missing}=missing)
    local δ, γ
    as = LinRange(-a, a, m)
    bs = LinRange(-b, b, n)
    α = zeros(m, n)
    β = zeros(m, n)
    θs = Array{Any}(undef, m, n)
    if !ismissing(direction)
        δ = direction - θ
        γ = randn(size(θ)...)
        γ = γ/norm(γ)*norm(δ)
    else
        δ = randn(size(θ)...)
        γ = randn(size(θ)...)
    end
    for i = 1:m
        for j = 1:n
            α[i,j] = as[i]
            β[i,j] = bs[j]
            θs[i,j] = θ + δ*as[i] + γ*bs[j]
        end
    end
    return θs
end

function meshview(losses::Array{Float64}, a::Real=1, b::Real=1)
    mp = require_pyplot()
    m, n = size(losses)
    α = zeros(m, n)
    β = zeros(m, n)
    as = LinRange(-a, a, m)
    bs = LinRange(-b, b, n)
    for i = 1:m
        for j = 1:n
            α[i,j] = as[i]
            β[i,j] = bs[j]
        end
    end
    mp.close("all")
    mp.mesh(α, β, losses)
    mp.xlabel("\$\\alpha\$")
    mp.ylabel("\$\\beta\$")
    mp.zlabel("loss")
    mp.scatter3D(α[(m+1)÷2, (n+1)÷2], β[(m+1)÷2, (n+1)÷2], losses[(m+1)÷2, (n+1)÷2], color="red", s=100)
    return α, β, losses
end

function pcolormeshview(losses::Array{Float64}, a::Real=1, b::Real=1)
    mp = require_pyplot()
    m, n = size(losses)
    α = zeros(m, n)
    β = zeros(m, n)
    as = LinRange(-a, a, m)
    bs = LinRange(-b, b, n)
    for i = 1:m
        for j = 1:n
            α[i,j] = as[i]
            β[i,j] = bs[j]
        end
    end
    mp.close("all")
    mp.pcolormesh(α, β, losses)
    mp.xlabel("\$\\alpha\$")
    mp.ylabel("\$\\beta\$")
    mp.scatter(0.0,0.0, marker="*", s=100)
    mp.colorbar()
    return α, β, losses
end

function meshview(sess::PyObject, pl::PyObject, loss::PyObject, θ;
        a::Real=1, b::Real=1, m::Integer=9, n::Integer=9,
        direction::Union{Array{<:Real}, Missing}=missing)
    dat = meshdata(θ; a=a, b=b, m=m, n=n, direction=direction)
    m, n = size(dat)
    V = zeros(m, n)
    for i = 1:m
        @info i, m
        for j = 1:n
            V[i,j] = run(sess, loss, pl=>dat[i,j])
        end
    end
    meshview(V, a, b)
end

function pcolormeshview(sess::PyObject, pl::PyObject, loss::PyObject, θ;
    a::Real=1, b::Real=1, m::Integer=9, n::Integer=9,
    direction::Union{Array{<:Real}, Missing}=missing)
    dat = meshdata(θ; a=a, b=b, m=m, n=n, direction=direction)
    m, n = size(dat)
    V = zeros(m, n)
    for i = 1:m
        @info i, m
        for j = 1:n
            V[i,j] = run(sess, loss, pl=>dat[i,j])
        end
    end
    pcolormeshview(V, a, b)
end

function gradview(sess::PyObject, pl::PyObject, loss::PyObject, u0, grad::PyObject;
    scale::Float64 = 1.0, mpi::Bool = false)
    mp = require_pyplot()
    local v
    if length(size(u0))==0
        v = rand()
    else
        v = rand(length(u0))
    end
    γs = scale ./ 10 .^ (1:5)
    v1 = Float64[]
    v2 = Float64[]
    L_ = run(sess, loss, pl=>u0)
    J_ = run(sess, grad, pl=>u0)
    for i = 1:5
        @info i
        L__ = run(sess, loss, pl=>u0+v*γs[i])
        push!(v1, norm(L__-L_))
        if size(J_)==size(v)
            if length(size(J_))==0
                push!(v2, norm(L__-L_-γs[i]*sum(J_*v)))
            else
                push!(v2, norm(L__-L_-γs[i]*sum(J_.*v)))
            end
        else
            push!(v2, norm(L__-L_-γs[i]*J_*v))
        end

    end
    if !(mpi) || (mpi && mpi_rank()==0)
        mp.close("all")
        mp.loglog(γs, abs.(v1), "*-", label="finite difference")
        mp.loglog(γs, abs.(v2), "+-", label="automatic linearization")
        mp.loglog(γs, γs.^2 * 0.5*abs(v2[1])/γs[1]^2, "--",label="\$\\mathcal{O}(\\gamma^2)\$")
        mp.loglog(γs, γs * 0.5*abs(v1[1])/γs[1], "--",label="\$\\mathcal{O}(\\gamma)\$")
        mp.plt.gca().invert_xaxis()
        mp.legend()
        mp.xlabel("\$\\gamma\$")
        mp.ylabel("Error")
        if mpi
            mp.savefig("mpi_gradview.png")
        end
    end

    return v1, v2
end

"""
    gradview(sess::PyObject, pl::PyObject, loss::PyObject, u0; scale::Float64 = 1.0)

Visualizes the automatic differentiation and finite difference convergence converge. For correctly implemented
differentiable codes, the convergence rate for AD should be 2 and for FD should be 1 (if not evaluated at stationary point).

- `scale`: you can control the step size for perturbation.
"""
function gradview(sess::PyObject, pl::PyObject, loss::PyObject, u0; scale::Float64 = 1.0, mpi::Bool = false)
    grad = tf.convert_to_tensor(gradients(loss, pl))
    gradview(sess, pl, loss, u0, grad, scale = scale, mpi = mpi)
end

@doc raw"""
    jacview(sess::PyObject, f::Function, θ::Union{Array{Float64}, PyObject, Missing},
    u0::Array{Float64}, args...)

Performs gradient test for a vector function. `f` has the signature
```
f(θ, u) -> r, J
```
Here `θ` is a nuisance  parameter, `u` is the state variables (w.r.t. which the Jacobian is computed),
`r` is the residual vector, and `J` is the Jacobian matrix (a dense matrix or a [`SparseTensor`](@ref)).

# Example 1
```julia
u0 = rand(10)
function verify_jacobian_f(θ, u)
    r = u^3+u - u0
    r, spdiag(3u^2+1.0)
end
jacview(sess, verify_jacobian_f, missing, u0)
```

# Example 2
```
u0 = rand(10)
rs = rand(10)
function verify_jacobian_f(θ, u)
    r = [u^2;u] - [rs;rs]
    r, [spdiag(2*u); spdiag(10)]
end
jacview(sess, verify_jacobian_f, missing, u0); close("all")
```
"""
function jacview(sess::PyObject, f::Function, θ::Union{Array{Float64}, PyObject, Missing},
                u0::Array{Float64}, args...)
    mp = require_pyplot()
    u = placeholder(Float64, shape=[length(u0)])
    L, J = f(θ, u)
    init(sess)
    L_ = run(sess, L, u=>u0, args...)
    J_ = run(sess, J, u=>u0, args...)
    v = rand(length(u0))
    γs = 1.0 ./ 10 .^ (1:5)
    v1 = Float64[]
    v2 = Float64[]
    for i = 1:5
        L__ = run(sess, L, u=>u0+v*γs[i], args...)
        push!(v1, norm(L__-L_))
        push!(v2, norm(L__-L_-γs[i]*J_*v))
    end
    mp.close("all")
    mp.loglog(γs, abs.(v1), "*-", label="finite difference")
    mp.loglog(γs, abs.(v2), "+-", label="automatic linearization")
    mp.loglog(γs, γs.^2 * 0.5*abs(v2[1])/γs[1]^2, "--",label="\$\\mathcal{O}(\\gamma^2)\$")
    mp.loglog(γs, γs * 0.5*abs(v1[1])/γs[1], "--",label="\$\\mathcal{O}(\\gamma)\$")
    mp.plt.gca().invert_xaxis()
    mp.legend()
    mp.xlabel("\$\\gamma\$")
    mp.ylabel("Error")
end

function PCLview(sess::PyObject, f::Function, L::Function, θ::Union{PyObject,Array{Float64,1}, Float64},
    u0::Union{PyObject, Array{Float64}}, args...; options::Union{Dict{String, T}, Missing}=missing) where T<:Real
    mp = require_pyplot()
    if isa(θ, PyObject)
        θ = run(sess, θ, args...)
    end
    x = placeholder(Float64, shape=[length(θ)])
    l, u, g = NonlinearConstrainedProblem(f, L, x, u0; options=options)
    init(sess)
    L_ = run(sess, l, x=>θ, args...)
    J_ = run(sess, g, x=>θ, args...)
    v = rand(length(x))
    γs = 1.0 ./ 10 .^ (1:5)
    v1 = Float64[]
    v2 = Float64[]
    for i = 1:5
        @info i
        L__ = run(sess, l, x=>θ+v*γs[i], args...)
        # @show L__,L_,J_, v
        push!(v1, L__-L_)
        push!(v2, L__-L_-γs[i]*sum(J_.*v))
    end
    mp.close("all")
    mp.loglog(γs, abs.(v1), "*-", label="finite difference")
    mp.loglog(γs, abs.(v2), "+-", label="automatic linearization")
    mp.loglog(γs, γs.^2 * 0.5*abs(v2[1])/γs[1]^2, "--",label="\$\\mathcal{O}(\\gamma^2)\$")
    mp.loglog(γs, γs * 0.5*abs(v1[1])/γs[1], "--",label="\$\\mathcal{O}(\\gamma)\$")
    mp.plt.gca().invert_xaxis()
    mp.legend()
    mp.xlabel("\$\\gamma\$")
    mp.ylabel("Error")
end

@doc raw"""
    animate(update::Function, frames; kwargs...)

Creates an animation using update function `update`.
# Example
```julia
θ = LinRange(0, 2π, 100)
x = cos.(θ)
y = sin.(θ)
pl, = plot([], [], "o-")
t = title("0")
xlim(-1.2,1.2)
ylim(-1.2,1.2)
function update(i)
    t.set_text("$i")
    pl.set_data([x[1:i] y[1:i]]'|>Array)
end
animate(update, 1:100)
```

!!! info
    `clf()` can be used to clear figures within `update` function.
"""
function animate(update::Function, frames; kwargs...)
    mp = require_pyplot()
    animation_ = pyimport("matplotlib.animation")
    if !isa(frames, Integer)
        frames = Array(frames)
    end
    animation_.FuncAnimation(mp.gcf(), update, frames=frames; kwargs...)
end

"""
    saveanim(anim::PyObject, filename::String; kwargs...)

Saves the animation produced by [`animate`](@ref)
"""
function saveanim(anim::PyObject, filename::String; kwargs...)
    if Sys.iswindows()
        anim.save(filename, "ffmpeg"; kwargs...)
    else
        anim.save(filename, "imagemagick"; kwargs...)
    end
end


# # Two Dimensions
#
#md # [![](https://mybinder.org/badge_logo.svg)](@__BINDER_ROOT_URL__/notebooks/random_flow_field.ipynb)
#md # [![](https://img.shields.io/badge/show-nbviewer-579ACA.svg)](@__NBVIEWER_ROOT_URL__/notebooks/random_flow_field.ipynb)
#
# Simulate trajectories of a particle cloud in a two-dimensional flow field.
# A doubly-periodic domain and simple flow fields are used.
# For additional documentation e.g. see :
# [1](https://JuliaClimate.github.io/IndividualDisplacements.jl/dev/),
# [2](https://JuliaClimate.github.io/MeshArrays.jl/dev/),
# [3](https://docs.juliadiffeq.org/latest/solvers/ode_solve.html),
# [4](https://en.wikipedia.org/wiki/Displacement_(vector))
#
# Exercises:
# - change the initial distribution of particles
# - increase the duration of the trajectories simulation
# - treat the non-periodic domain case by padding `u,v` with zeros
# - replace `u,v` with your own two-dimensional flow fields
#
# ![particles in random flow](https://github.com/JuliaClimate/IndividualDisplacements.jl/raw/master/examples/figs/RandomFlow.gif)

#nb # %% {"slideshow": {"slide_type": "slide"}, "cell_type": "markdown"}
# ## 1. Import Software

using IndividualDisplacements, DataFrames
p=dirname(pathof(IndividualDisplacements))
include(joinpath(p,"../examples/jupyter/flow_fields.jl"));

#nb # %% {"slideshow": {"slide_type": "slide"}, "cell_type": "markdown"}
# ## 2. Flow Fields
#
# Let's start with a simple, doubly periodic flow field defined by a streamfunction
# and create the `FlowFields` data structure which will then be drive the
# individual displacement and trajectory computations

nx=16
dx= π/nx
XC = dx*(collect(1:2*nx) .- 0.5)
YC = dx*(collect(1:nx) .- 0.5)

fac=0.1
f(x, y) = sin(x) + cos(y) #streamfunction
ϕ = fac*[f(x, y) for x in XC,y in YC] #streamfunction
uC = -fac*[sin(y) for x in XC,y in YC] #dphi/dy at cell center
vC = -fac*[cos(x) for x in XC,y in YC]; #-dphi/dx at cell center

# When starting as done here with collocated velocities in m/s units (`uC,vC` at
# the grid cell center), one easily obtains the normalized, staggered C-grid
# velocities (`u,v`) as shown below.
# `u` (resp `v`) is staggered by `-0.5` grid point in direction `1` for `u` (`2` for `v`)
# from the grid cell center `uC,vC`. This staggering of variables as well as the
# normalization of velocities by the grid scale are key conventions
# adopted in `IndividualDisplacements.jl`.

u=0.5*(circshift(uC, (1,0))+uC) /dx #staggered u converted to grid point units (m/s -> 1/s)
v=0.5*(circshift(vC, (0,1))+vC) /dx #staggered v converted to grid point units (m/s -> 1/s)

𝐹=FlowFields(u,u,v,v,[0.,10.])

#nb # %% {"slideshow": {"slide_type": "slide"}, "cell_type": "markdown"}
# ## 3. Initialize Individuals
#
# For example, we can initialize 100 particles within a central subdomain as follows.

np,nq=size(u)
x=np*(0.4 .+ 0.2*rand(100))
y=nq*(0.4 .+ 0.2*rand(100));

# The following constructor function wraps everything in the `Individuals` data structure.

𝐼=Individuals(𝐹,x,y)

#nb # %% {"slideshow": {"slide_type": "slide"}, "cell_type": "markdown"}
# ## 4. Compute Trajectories
#
# The time period is `𝐼.𝑃.𝑇` by default, unless `∫!(𝐼,𝑇)` is called instead.

∫!(𝐼)

#nb # %% {"slideshow": {"slide_type": "slide"}, "cell_type": "markdown"}
# ## 5. Plot Results

p=dirname(pathof(IndividualDisplacements))
include(joinpath(p,"../examples/jupyter/recipes_plots.jl"));

🔴_by_t = groupby(𝐼.🔴, :t)
phi_scatter(ϕ,🔴_by_t[end])

# Or to generate a simple animation:
#
# ```
# anim = @animate for t in eachindex(🔴_by_t)
#    phi_scatter(ϕ,🔴_by_t[t])
# end
#
# pth=tempdir()*"/"
# gif(anim, pth*"RandomFlow.gif", fps = 10)
# ```

#nb # %% {"slideshow": {"slide_type": "slide"}, "cell_type": "markdown"}
# ## 6. Alternatives (optional)

# Making the flow field time variable simply involves providing two different
# flow fields. Linear interpolation will then be applied in the specified time
# range. For example :
#
# ```
# 𝐹=FlowFields(-u,u,-v,v,[0.,10.])
# ```

# It is often convenient to set up `FlowFields` using the MeshArrays.jl package, which
# handles such staggered grids in general fashion, instead of plain Arrays. Given a doubly
# periodic flow field defined as before this is readily done via the `convert_to_FlowFields` function.
#
# ```
# u,v,ϕ=random_flow_field()
# #𝐹=FlowFields(u,u,v,v,[0.,10.])
# 𝐹=convert_to_FlowFields(u,v,10.0);
# ```

# We should note that `u,v` derive from streamfunction `ϕ` which is defined at the corner point in `random_flow_field`.
# Doing this ensures that the resulting `u,v` is non-divergent / purely-rotational over the C-grid domain. In brief:
#
# ```
# u=-(circshift(ϕ, (0,-1))-ϕ)
# v=(circshift(ϕ, (-1,0))-ϕ)
# ```

# In contrast uC, vC computed as earlier in this example may contain contain both
# [rotational and divergent](https://en.wikipedia.org/wiki/Helmholtz_decomposition)
# components -- if this were an important consideration, `MeshArrays.jl` provides
# tools that can be used to e.g. ensure non-divergence.

# When using `MeshArray` flow fields instead of plain `Array` flow fields, initial
# positions also include a subdomain array index (`a`;  all ones in our example).
#
# ```
# a=ones(size(x))
# #isa(𝐹,𝐹_MeshArray2D)
# 𝐼=Individuals(𝐹,x,y,a)
# ```


using Test

input = parse.(Int, readlines("01.input"))

star1(input) = sum(input)
@show star1(input)

function star2(input)
    seen = Set{Int}([0])
    freq = 0
    for i in Iterators.cycle(input)
        freq += i
        if freq ∈ seen
            return freq
        else
            push!(seen, freq)
        end
    end
end

@test [+1, -1] |> star2 == 0
@test [+3, +3, +4, -2, -4] |> star2 == 10
@test [-6, +3, +8, +5, -6] |> star2 == 5
@test [+7, +7, -2, -7, -4] |> star2 == 14

@show star2(input)

# can maybe use a BitSet to speed things up (represents a set of integers with a
# bitstring).
function star2bitset(input)
    seen = BitSet([0])
    freq = 0
    for i in Iterators.cycle(input)
        freq += i
        if freq ∈ seen
            return freq
        else
            push!(seen, freq)
        end
    end
end

@test [+1, -1] |> star2bitset == 0
@test [+3, +3, +4, -2, -4] |> star2bitset == 10
@test [-6, +3, +8, +5, -6] |> star2bitset == 5
@test [+7, +7, -2, -7, -4] |> star2bitset == 14

star2bitset(input)

using BenchmarkTools
@btime star2(input)             # 3.8 ms
@btime star2bitset(input)       # 301 μs


using Distributions
using StatsBase
using Base.Test

# Weights

wk = [0.20; 0.50; 0.25; 0.05]

wk = wk / sum(wk)

# Perform resampling

res = Int64[]

for i = 1:100000

    append!(res, resample(wk))

end

# Check results

nr = counts(res, 1:maximum(res))

@test all(wk .== round(nr / sum(nr), 2))


"""
random vector on unit sphere, should be uniformly distributed?
"""
function rand_sphere(T = Float64)
   R = randn(SVector{3, T})
   return R / norm(R)
end

"""
random rotation matrix;
WARN: never checked what the distribution is
"""
rand_rot() = (K = (@SMatrix rand(3,3)) .- 0.5; exp(K - K'))

"""
random reflection, represented as 1, -1 Integer
"""
rand_refl() = rand([-1,1])

"""
random isometry i.e. element of O(3)
"""
rand_O3() = rand_refl() * rand_rot()

# the following two functions are a little hack to make sure
# that the basis spec is read in the same symbol-order as it is written
# (since dicts don't have a specified ordering...)

function _write_dict_1pspec(spec::Vector{NamedTuple{SYMS, NTuple{NSYM, Int}}}) where {SYMS, NSYM}
   inds = Vector{Int}[]
   for b in spec
      vals = [getproperty(b, sym) for sym in SYMS]
      push!(inds, vals)
   end
   return Dict("SYMS" => [ string.(SYMS)... ], "inds" => inds)
end

function _read_dict_1pspec(D::Dict)
   NTPROTO = namedtuple(D["SYMS"]...)
   return [ NTPROTO(binds) for binds in D["inds"] ]
end

#
# This is a useful utility function used e.g. in scal1pbasis.jl
#

"""
returns an `SVector{N}` of the form `x * e_I` where `e_I` is the Ith canonical basis vector.
"""
@generated function __e(::SVector{N}, ::Val{I}, x::T) where {N, I, T}
   code = "SA["
   for i = 1:N
      if i == I
         code *= "x,"
      else
         code *= "0,"
      end
   end
   code *= "]"
   quote
      $( Meta.parse(code) )
   end
end

__e(xx::SVector{N, T}, valI::Val{I}) where {N, T, I} = __e(xx, valI, one(T))

__e(::Number, ::Any, x) = x


issquare(A) = size(A, 1) == size(A, 2) && ndims(A) == 2

perf_hvcat(A, B) =  [A B; B A]

function perf_hvcat_setind(A, B)
    @assert issquare(A) && size(A) == size(B)
    n = size(A, 1)
    C = Matrix{Float64}(undef, 2n, 2n)
    C[1:n, 1:n] = A
    C[1:n, (n+1):end] = B
    C[(n+1):end, 1:n] = B
    C[(n+1):end, (n+1):end] = A
end

perf_hcat(A, B) = [A B B A]

function perf_hcat_setind(A, B)
    @assert issquare(A) && size(A) == size(B)
    n = size(A, 1)
    C = Matrix{Float64}(undef, n, 4n)
    C[:, 1:n] = A
    C[:, (n+1):2n] = B
    C[:, (2n+1):3n] = B
    C[:, (3n+1):end] = A
    return C
end

perf_vcat(A, B) = [A; B; B; A]

function perf_vcat_setind(A, B)
    @assert issquare(A) && size(A) == size(B)
    n = size(A, 1)
    C = Matrix{Float64}(undef, 4n, n)
    C[1:n, :] = A
    C[(n+1):2n, :] = B
    C[(2n+1):3n, :] = B
    C[(3n+1):4n, :] = A
    return C
end

function perf_catnd(n)
    A = samerand(1, n, n, 1)
    B = samerand(1, n, n)
    return cat(A, B, B, A; dims = 3)
end

function perf_catnd_setind(n)
    A = samerand(1, n, n, 1)
    B = samerand(1, n, n)
    C = Array{Float64}(undef, 1, n, 4n, 1)
    C[1, :, 1:n, 1] = A
    C[1, :, (n+1):2n, 1] = B
    C[1, :, (2n+1):3n, 1] = B
    C[1, :, (3n+1):4n, 1] = A
    return C
end

function perf_cat_4467()
    a, b, c, d, e = 1.0, 2.0, 3.0, 4.0, 5
    return [a, b, c, d, e]
end


"""
$(TYPEDEF)
"""
mutable struct Point3D
    # Coordinates
    x::Float64
    y::Float64
    z::Float64

    # Velocity at those coordinates
    ix::Int
    iy::Int
end

"""
$(SIGNATURES)
"""
function dist(p1::Point3D, p2::Point3D)
    dx = p2.x - p1.x
    dy = p2.y - p1.y
    return sqrt(dx * dx + dy * dy)
end

"""
$(SIGNATURES)

A midpoint is the average between two points
"""
midpoint(p1::Point3D, p2::Point3D) =
    Point3D(0.5p1.x + 0.5p2.x, 0.5p1.y + 0.5p2.y, 0.5p1.z + 0.5p2.z, 0, 0)

"""
$(SIGNATURES)

Barycentre of three points.
"""
centroid(p1::Point3D, p2::Point3D, p3::Point3D) = Point3D(
    (p1.x + p2.x + p3.x) / 3.0,
    (p1.y + p2.y + p3.y) / 3.0,
    (p1.z + p2.z + p3.z) / 3.0,
    0,
    0,
)

"""
$(TYPEDEF)
"""
struct Triangle
    pt1::Int64
    pt2::Int64
    pt3::Int64

    Triangle() = new(0, 0, 0)
    Triangle(pt1, pt2, pt3) = new(pt1, pt2, pt3)
end

"""
$(TYPEDEF)
"""
struct Mesh
    nodes::Vector{Point3D}
    nodeArray::Matrix{Point3D}
    triangles::Vector{Triangle}
    width::Int
    height::Int
end

"""
$(SIGNATURES)
"""
function triangle_area(mesh::Mesh, index::Int)
    triangle = mesh.triangles[index]
    pt1 = mesh.nodes[triangle.pt1]
    pt2 = mesh.nodes[triangle.pt2]
    pt3 = mesh.nodes[triangle.pt3]

    return triangle_area(pt1, pt2, pt3)
end

"""
$(SIGNATURES)
"""
function triangle_area(p1::Point3D, p2::Point3D, p3::Point3D)
    a = dist(p1, p2)
    b = dist(p2, p3)
    c = dist(p3, p1)
    s = (a + b + c) / 2.0

    return sqrt(s * (s - a) * (s - b) * (s - c))
end


"""

二重クロス集計表の各マス目の度数を円の面積で表現するのがバルーンプロットである。

以下の関数に freqtable() による集計表を与えれば描画される。注意点は，円の面積が度数に対応するということ。
"""

using Plots
function balloonplot(table; xlabel="", ylabel="",
                     color=:red, unit=100)
    θ = collect(0:360);
    x = cosd.(θ);
    y = sind.(θ);
    gr(linecolor=:black, label="")
    row, col = size(table)
    row_names, col_names = names(table)
    mx = maximum(table)
    p2 = vline(0.5:col+0.5, xlimit=(0.5, col+0.5),
        ylimit=(0.5, row+0.5), showaxis=false,
        xlabel=xlabel, ylabel=ylabel,
        aspect_ratio=1, size=(col*unit, row*unit))
    hline!(0.5:row+0.5)
    for i in 1:row
        for j in 1:col
            r = 0.45sqrt(table[i, j]/mx)
            plot!(r .* x .+ j, r .* y .+ (row - i + 1),
                  seriestype=:shape, fillcolor=color)
            annotate!(j, row - i + 1, text(string(table[i, j]), 8))
        end
    end
    xticks!(1:col, col_names)
    yticks!(1:row, reverse(row_names))
    p2
end;

"""
使用法

julia> using RDatasets

julia> haireyecolor = RDatasets.dataset("datasets", "HairEyeColor");

julia> using DataFrames, FreqTables, Plots

julia> haireyecolor[!, 1:3] = string.(haireyecolor[!, 1:3]); # String7 が使えないパッケージに対する過渡的な対処法

julia> HEC = DataFrame();

julia> for row in eachrow(haireyecolor), j in 1:row.Freq
           push!(HEC, row)
       end

julia> select!(HEC, Not(:Freq)); # Freq 列を除く

julia> EyeHair = freqtable(HEC.Eye, HEC.Hair)
4×4 Named Matrix{Int64}
Dim1 ╲ Dim2 │ Black  Blond  Brown    Red
────────────┼───────────────────────────
Blue        │    20     94     84     17
Brown       │    68      7    119     26
Green       │     5     16     29     14
Hazel       │    15     10     54     14

julia> balloonplot(EyeHair,
           xlabel="Hair color", ylabel="Eye color", color=:blue)
"""


@testset "SGLDAdam" begin
    #-----------------------
    # testing inplace addition
    #-----------------------
    function sample_reference!(s::SGLDAdamState,grad::Function)
        # increment iteration
        s.t += 1
        s.betat *= s.beta

        # improve readability
        t = s.t
        x = s.x
        lambda = s.lambda
        precond = s.preconditioner
        stepsize = s.stepsize(t)
        beta = s.beta
        betat = s.betat

        # update preconditioner
        g = grad(x)
        if t == 1
            precond[:] = g.*g
        else
            precond[:] = beta*precond + (1-beta)*g.*g
        end
        m = lambda + sqrt(precond/(1-betat))
        # gradient step (inplace addition)
        # equivalent to x += stepsize*grad(x)
        x[:] += stepsize*g./m + sqrt(2.0*stepsize./m).*randn(length(x))
        s
    end

    x = 1+ 0.5*randn(2)

    s = SGLDAdamState(x,0.001)
    ref_s = SGLDAdamState(x,0.001)
    sgldadam(x) = x.*x

    srand(1)
    sample!(s,sgldadam)

    srand(1)
    sample_reference!(ref_s,sgldadam)

    @test_approx_eq(s.x,ref_s.x)

    #----------------------
    # Gaussian sanity check.
    #----------------------
    dm = GaussianMixtureModel([0.],[1.],[1.])

    gradient = DataModel.getgrad(dm)
    s = SGLDAdamState(randn(1),0.01)

    x_avg = [0.]
    x2_avg = [0.]
    for i in 1:200000
        sample!(s,gradient)
        x_avg =(i-1)/i*x_avg + s.x/i
        x2_avg = (i-1)/i*x2_avg + s.x.^2/i
    end

    @test abs(x_avg-0)[1] < 0.1
    @test abs(x2_avg-1)[1] < 0.5
end


include("../../../PlotUtils/PlotUtils.jl")
include(joinpath(@__DIR__, "../../imports.jl"))
using PyPlot
const plt = PlotUtils.plt

function tomu(delta)
  mu0 = -cumsum(delta[0])
  mu1 = cumsum(delta[1])
  return [mu0; mu1]
end

function etas_tomat(etas)
  eta0s = [eta[0] for eta in etas]
  eta1s = [eta[1] for eta in etas]
  return (cat(eta0s..., dims=4),
          cat(eta1s..., dims=4))
end

# Results directory
results_dir = "/scratchdata/alui2/cytof/results/repfam/test-sim-6-7-13/"

# Name of output file
OUTPUT_FILE = "output.bson"

# Path to all output files
output_paths = [joinpath(root, OUTPUT_FILE)
                for (root, _, files) in walkdir(results_dir)
                if OUTPUT_FILE in files]

# Load outputs
path_to_output = output_paths[2]
output = BSON.load(path_to_output)
samples = output[:samples][1]
m = output[:d].data.m
yimps = [o[:theta__y_imputed]
         for o in output[:samples][2]]
std(yimps[1][1][Bool.(m[1])])
std(yimps[10][1][Bool.(m[1])])

extract(sym) = [s[sym] for s in samples]

# mu*
mus = [tomu(d) for d in extract(:theta__delta)]
mu_mat = Matrix(hcat(mus...)')

# sig2
sig2s = extract(:theta__sig2)
sig2_mat = Matrix(hcat(sig2s...)')

# W
Ws = extract(:theta__W)
Wmat = cat(Ws..., dims=3)

# Z
Zs = extract(:theta__Z)

# R
Rs = hcat([vec(sum(W .> 0, dims=2)) for W in Ws]...)

# eta
eta0_mat, eta1_mat = etas_tomat(extract(:theta__eta))

# Plot (sig2, mu) vs eta
ell = 1
for i in 1:2, j in 1:20
  println((i, j))

  plt.subplot(2, 1, 1)
  plt.plot(eta0_mat[i, j, ell, :], mu_mat[:, 1], lw=0.5)
  plt.ylabel(L"$\mu^\star_{0, 1}$")
  plt.xlabel(latexstring("\\eta^0_{1, $(j), 1}"))

  plt.subplot(2, 1, 2)
  plt.plot(eta0_mat[i, j, ell, :], sig2_mat[:, i], lw=0.5)
  plt.ylabel(latexstring("\\sigma^2_{$(i)}"))
  plt.xlabel(latexstring("\\eta^0_{1, $(j), 1}"))

  plt.savefig("img/mu-vseta0_$(i)_$(j)_1.pdf", bbox_inches="tight")
  plt.close()
end

# plt.plot(output[:dden][10][1, 13])
# plt.plot(output[:dden][end][1, 13])
# plt.savefig("img/dden11.pdf")
# plt.close()


using Test, Convex, Cliffords, SCS, SchattenNorms, Distributions, QuantumInfo
import Random

import Base.kron

set_default_solver(SCSSolver(verbose=0, eps=1e-6, max_iters=5_000))

function paulichannel(p::Vector{T}) where T <: Real
    n_ = log(4,length(p))
    if !isinteger(n_)
        error("Probability vector must have length 4^n for some integer n")
    end
    n = round(Int,n_)
    allpaulisops = map(m->liou(complex(m)),allpaulis(n))
    return reduce(+,map(*,p,allpaulisops))
end

function randp(n)
    return rand(Dirichlet(ones(4^n)))
end

function dnormp(p1,p2)
    return norm(p1-p2,1)
end

Random.seed!(123456)

p1 = randp(2)
p2 = randp(2)

pc1 = paulichannel(p1)
pc2 = paulichannel(p2)

dnormcptp(pc1,pc2)
SchattenNorms.dnormcptp2(pc1,pc2)
dnorm(pc1-pc2)

times = Vector[]
results = Vector[]
for i in 1:100
    println("Iteration ",i)

    p1 = randp(2)
    p2 = randp(2)

    pc1 = paulichannel(p1)
    pc2 = paulichannel(p2)

    tic()
    r1 = dnormcptp(pc1,pc2)
    te1 = toc()

    tic()
    r2 = SchattenNorms.dnormcptp2(pc1,pc2)
    te2 = toc()

    tic()
    r3 = dnorm(pc1-pc2)
    te3 = toc()

    push!(times, [te1, te2, te3])
    push!(results, [r1, r2, r3, dnormp(p1,p2)])
end
times = reduce(hcat,times)'
results = reduce(hcat,results)'
results = results .- results[:,4]


function num_vecjac!(
    du,
    f,
    x,
    v,
    cache1 = similar(v),
    cache2 = similar(v);
    compute_f0 = true,
)
    if DiffEqBase.numargs(f) != 2
        du .= num_jacvec(f, x, v)
        return du
    end
    compute_f0 && (f(cache1, x))
    T = eltype(x)
    # Should it be min? max? mean?
    ϵ = sqrt(eps(real(T))) * max(one(real(T)), abs(norm(x)))
    vv = reshape(v, size(x))
    for i = 1:length(x)
        x[i] += ϵ
        f(cache2, x)
        x[i] -= ϵ
        du[i] = (((cache2 .- cache1) ./ ϵ)'*vv)[1]
    end
    return du
end

function num_vecjac(f, x, v, f0 = nothing)
    vv = reshape(v, axes(x))
    f0 === nothing ? _f0 = f(x) : _f0 = f0
    T = eltype(x)
    # Should it be min? max? mean?
    ϵ = sqrt(eps(real(T))) * max(one(real(T)), abs(norm(x)))
    du = similar(x)
    for i = 1:length(x)
        x[i] += ϵ
        f0 = f(x)
        x[i] -= ϵ
        du[i] = (((f0 .- _f0) ./ ϵ)'*vv)[1]
    end
    return vec(du)
end


using DiffEqBase, DiffEqJump, DiffEqProblemLibrary, Plots, Statistics
using DiffEqProblemLibrary.JumpProblemLibrary: importjumpproblems; importjumpproblems()
import DiffEqProblemLibrary.JumpProblemLibrary: prob_jump_dnarepressor
gr()
fmt = :png

methods = (Direct(),DirectFW(),FRM(),FRMFW(),SortingDirect(),NRM(),DirectCR(),RSSA())
shortlabels = [string(leg)[12:end-2] for leg in methods]
prob    = prob_jump_dnarepressor.discrete_prob
tf      = prob_jump_dnarepressor.tstop
rn      = prob_jump_dnarepressor.network
ploth   = plot(reuse=false)
for (i,method) in enumerate(methods)
    jump_prob = JumpProblem(prob, method, rn, save_positions=(false,false))
    sol = solve(jump_prob, SSAStepper(), saveat=tf/1000.)
    plot!(ploth,sol.t,sol[3,:],label=shortlabels[i], format=fmt)
end
plot(ploth, title="Protein level", xlabel="time",format=fmt)

p = []
for (i,method) in enumerate(methods)
    jump_prob = JumpProblem(prob, method, rn, save_positions=(false,false))
    sol = solve(jump_prob, SSAStepper(), saveat=tf/1000.)
    push!(p, plot(sol,title=shortlabels[i],leg=false,format=fmt))
end
plot(p...,format=fmt)

function run_benchmark!(t, jump_prob, stepper)
    sol = solve(jump_prob, stepper)
    @inbounds for i in 1:length(t)
        t[i] = @elapsed (sol = solve(jump_prob, stepper))
    end
end

nsims = 500
benchmarks = Vector{Vector{Float64}}()
for method in methods
    jump_prob = JumpProblem(prob, method, rn, save_positions=(false,false))
    stepper = SSAStepper()
    t = Vector{Float64}(undef,nsims)
    run_benchmark!(t, jump_prob, stepper)
    push!(benchmarks, t)
end

medtimes = Vector{Float64}(undef,length(methods))
stdtimes = Vector{Float64}(undef,length(methods))
avgtimes = Vector{Float64}(undef,length(methods))
for i in 1:length(methods)
    medtimes[i] = median(benchmarks[i])
    avgtimes[i] = mean(benchmarks[i])
    stdtimes[i] = std(benchmarks[i])
end
println(medtimes/medtimes[1])

using DataFrames
df = DataFrame(names=shortlabels,medtimes=medtimes,relmedtimes=(medtimes/medtimes[1]),avgtimes=avgtimes, std=stdtimes, cv=stdtimes./avgtimes)
sa = [text(string(round(mt,sigdigits=2),"s"),:center,10) for mt in df.medtimes]
bar(df.names,df.relmedtimes,legend=:false, fmt=fmt)
scatter!(df.names, .05 .+ df.relmedtimes, markeralpha=0, series_annotations=sa, fmt=fmt)
ylabel!("median relative to Direct")
title!("Negative Feedback Gene Expression Model")

using DiffEqBenchmarks
DiffEqBenchmarks.bench_footer(WEAVE_ARGS[:folder],WEAVE_ARGS[:file])


using Revise
using LinearAlgebra, StatsBase, Statistics, Distributions, Random
#Random.seed!(11)

#for displaying matrices
printm(m) = show(IOContext(stdout, :limit=>false), MIME"text/plain"(), m)
function printmln(m)
  show(IOContext(stdout, :limit=>false), MIME"text/plain"(), m)
  println()
end

#generates a well-defined covariance matrix
function randomcovar(K, N=10^4; displaymatrix=false)
  ϵ = rand(Uniform(-1,1), N,K)
  ψ = rand(Uniform(-1,1), N)
  β = (k->k*(-1)^k).(1:K)

  X = β' .* ψ
  @assert size(X) == (N,K)
  X .+= ϵ
  Σ = cov(X)

  displaymatrix && printmln(Σ)
  return Σ
end

function checkmatrixalgebra(dims::NamedTuple; testiter=10^5)
  atol = dims.K/testiter^0.5
  #generate the characteristic matrix
  Ξ = rand(Laplace(), dims.N, dims.K+1)
  #Ξ .-= mean(Ξ, dims=1)
  Ξ[:,end] = rand(dims.N) #last column is the market

  #signal parameters
  δi = rand(dims.K+1)
  Σθ = randomcovar(dims.K+1)
  σ²ϵ = rand(dims.N) .+ 0.0001
  #Σϵ = Diagonal(σ²ϵ)
  Σϵ = randomcovar(dims.N)

  Rf = 1.2
  P = rand(TriangularDist(1.0,11.0), dims.N)
  N1 = ones(dims.N)

  #distributions
  distθ = MvNormal(δi, Σθ)
  distϵ = MvNormal(zeros(dims.N), Σϵ)

  #test covariance of portfolio
  wi = rand(dims.N) .+ 0.01
  wi ./= sum(wi)
  ωi = wi ./ P
  vervec = [wi'* (Ξ * rand(distθ) ./ P .+ rand(distϵ) ./ P .- Rf * N1) for i in 1:testiter]
  vervar = var(vervec)
  testvar = ωi' * (Ξ*Σθ*Ξ' .+ Σϵ) * ωi
  msg  = "covariance of portfolio test with Δ $(abs(vervar-testvar)) and max atol $atol"
  isapprox(testvar, vervar, atol=atol) ? @info("passed $msg") : @warn("failed $msg")

  #test expectation
  vermean = mean(vervec)
  testmean = ωi' * (Ξ * δi .- P * Rf)
  msg  = "expectation of portfolio test with Δ $(abs(vermean-testmean)) and max atol $atol"
  isapprox(testmean, vermean, atol=atol) ? @info("passed $msg") : @warn("failed $msg")

  #test woodbury
  Σθinv = svd(Σθ)\I
  Σϵinv = Σϵ\I
  IN = Diagonal(I,dims.N)

  verω = (Ξ *Σθ * Ξ' + Σϵ)\I * (Ξ * δi .- P .* Rf)
  testw1 = Σϵinv*(I - Ξ*((Σθinv .+ Ξ'*Σϵinv*Ξ)\I)*Ξ'*Σϵinv)*(Ξ*δi .- P .* Rf)
  @assert verω ≈ testw1
  @info "passed test prewoodbury"

  si = (I - ((Σθinv .+ Ξ'*Σϵinv*Ξ)\I) * (Ξ'*Σϵinv * Ξ)) * δi
  κ = (I - Ξ*((Σθinv .+ Ξ'*Σϵinv*Ξ)\I) * (Ξ'*Σϵinv))
  testw2 = Σϵinv*(Ξ*si .- κ * P .* Rf)
  @assert verω ≈ testw2

  testw3 = P .* Σϵinv*(Ξ*si  .- κ * P .* Rf)
  #testw3 = Σϵinv*(Ξ*si .- κ * P .* Rf)
  ver3 = P .* verω
  @assert ver3 ≈ testw3
  @info "passed test woodbury"

  ###Market amalgamation test
  Ξᴷ = Ξ[:,1:(end-1)]
  ξₘ = Ξ[:,end] .* si[end] .- κ * P .* Rf
  siᴷ = si[1:(end-1)]
  testamalg1 = ωi'*Σϵinv*Ξᴷ*siᴷ .+ ωi'*Σϵinv*ξₘ

  Ξᴺ=deepcopy(Ξ)
  Ξᴺ[:,end] .-= κ * P .* Rf/si[end]
  testamalg2 = ωi'*Σϵinv*Ξᴺ*si

  veramalg = ωi'*Σϵinv*(Ξ*si .- κ * P .* Rf)

  @assert testamalg1 ≈ veramalg
  @assert testamalg2 ≈ veramalg
  @info "passed market amalgamation test"

  #Signal test
  verxi = verω .* P#Σϵinv * (Ξ*si - κ*P*Rf) .* P
  testxi = [Σϵinv * Ξ[:,k] .* si[k] .* P for k ∈ 1:(dims.K+1)]
  testxi[end] .= Σϵinv * (Ξ[:,end] .* si[end] .- κ*P .* Rf) .* P
  @assert verxi ≈ sum(testxi)
  @info "passed portfolio weight test. Portfolio weights (net): $(sum.(testxi))"

end

function checkwork()
  dims = (N=1000, K=5)
  checkmatrixalgebra(dims,testiter=10^5)
end

@time checkwork()


"""
    concordance(class, var, tie)

Computes concordant, tied and discordant pairs.\\
`class` can be either a BitVector or a 2 level categorical target variable in which case
`true` is defined by the last value in sorted sequence.\\
`var` is a Vector of predictor, same length as `class`,\\
`tie` (optional) can be a number (default is 1e-6) that defines a tied region,
or it can be a function that when called with a scalar value will return a tuple
of lower bound and upper bound of a tied region,
useful when you want to do percentage tied region for instance.

Pair-wise comparison between class 1 with class 0 values are made as follows:
- class 1 value > class 0 value is Concordant
- class 1 value ≈ class 0 value (within `tie`) is Tied
- class 1 value < class 0 value is Discordant

Returns:
- concordant, number of concordant comparisons
- tied, number of tied comparisons
- discordant, number of discordant comparisons
- auroc, or C, is (Concordant + 0.5Tied) / Total comparisons; same as numeric integration of ROC curve
- gini, 2C-1, also known as Somer's D, is (Concordant - Discordant) / Total comparisons

Concordance calculation is the same as numeric integration of the ROC curve,
but it allows for fuzzy tied regions which can be useful.

Note:
- Goodman-Kruskal Gamma is (Concordant - Discordant) / (Concordant + Discordant)
- Kendall's Tau is (Concordant - Discordant) / (0.5 x Total count x (Total count - 1))
"""
function concordance(class::BitVector, var::Vector, tie = 1e-6)
    n = length(class)
    n == length(var) || throw(ArgumentError("class and var should be the same length"))
    n1 = sum(class)
    n1 == 0 && throw(ArgumentError("there are no class 1"))
    n0 = n - n1
    n0 == 0 && throw(ArgumentError("there are no class 0"))

    c1 = sort!(var[class])              # sorted var of class 1
    c0 = sort!(var[.!class])            # sorted var of class 0

    conc = 0
    tied = 0
    l_ix = 1
    u_ix = 1
    @inbounds for v in c1               # loop over this usually smaller array
        l_v = v - tie
        u_v = v + tie

        while l_ix <= n0                # first index within [l_v, u_v] window
            l_v <= c0[l_ix] && break
            l_ix += 1                   # could exit as n0 + 1
        end

        while u_ix <= n0                # first index beyond [l_v, u_v] window
            u_v < c0[u_ix] && break
            u_ix += 1                   # could exit as n0 + 1
        end

        conc += (l_ix - 1)              # lower than l_ix are Concordant
        tied += (u_ix - l_ix)           # within window are tied
    end

    pairs = n1 * n0
    disc = pairs - conc - tied
    auc = (conc + 0.5tied) / pairs

    (conc = conc, tied = tied, disc = disc, auc = auc, gini = 2auc - 1)
end

function concordance(class::BitVector, var::Vector, tie::Function)
    n = length(class)
    n == length(var) || throw(ArgumentError("class and var should be the same length"))
    n1 = sum(class)
    n1 == 0 && throw(ArgumentError("there are no class 1"))
    n0 = n - n1
    n0 == 0 && throw(ArgumentError("there are no class 0"))

    c1 = sort!(var[class])              # sorted var of class 1
    c0 = sort!(var[.!class])            # sorted var of class 0

    conc = 0
    tied = 0
    l_ix = 1
    u_ix = 1
    @inbounds for v in c1               # loop over this usually smaller array
        l_v, u_v = tie(v)               # tie scalar function returning lower & upper bound

        while l_ix <= n0                # first index within [l_v, u_v] window
            l_v <= c0[l_ix] && break
            l_ix += 1                   # could exit as n0 + 1
        end

        while u_ix <= n0                # first index beyond [l_v, u_v] window
            u_v < c0[u_ix] && break
            u_ix += 1                   # could exit as n0 + 1
        end

        conc += (l_ix - 1)              # lower than l_ix are Concordant
        tied += (u_ix - l_ix)           # within window are tied
    end

    pairs = n1 * n0
    disc = pairs - conc - tied
    auc = (conc + 0.5tied) / pairs

    (conc = conc, tied = tied, disc = disc, auc = auc, gini = 2auc - 1)
end

function concordance(class::Vector, var::Vector, tie = 1e-6)
    uc = sort!(unique(class))
    length(uc) == 2 || throw(ArgumentError("class must be 2 levels"))

    concordance(class .== uc[2], var, tie)
end

roc = concordance                       # synonym


using Diagonalizations, LinearAlgebra, PosDefManifold, Test

####  Create data for testing the case k=1, m>1
# `t` is the number of samples,
# `m` is the number of datasets,
# `n` is the number of variables,
# `noise` must be smaller than 1.0. The smaller the noise,
#  the more data are correlated.
function getData(t, m, n, noise)
    # create m identical data matrices and rotate them by different
    # random orthogonal matrices V_1,...,V_m
    𝐕=[randU(n) for i=1:m] # random orthogonal matrices
    X=randn(n, t)  # data common to all subjects
    # each subject has this common part plus a random part
    𝐗=[𝐕[i]'*((1-noise)*X + noise*randn(n, t)) for i=1:m]
    return 𝐗
end

function getData(::Type{Complex{T}}, t, m, n, noise) where {T<:AbstractFloat}
    # create m identical data matrices and rotate them by different
    # random orthogonal matrices V_1,...,V_m
    𝐕=[randU(ComplexF64, n) for i=1:m] # random orthogonal matrices
    X=randn(ComplexF64, n, t)  # data common to all subjects
    # each subject has this common part plus a random part
    𝐗=[𝐕[i]'*((1-noise)*X + noise*randn(ComplexF64, n, t)) for i=1:m]
    return 𝐗
end

# REAL data: check that for the case m=2 gMCA gives the same result as MCA
t, m, n, noise = 20, 2, 6, 0.1
Xset=getData(t, m, n, noise)
Cx=(Xset[1]*Xset[1]')/t
Cy=(Xset[2]*Xset[2]')/t
Cxy=(Xset[1]*Xset[2]')/t

gm=gmca(Xset; simple=true)
m=mca(Cxy; simple=true)

@test (m.F[1]'*Cxy*m.F[2]) ≈ (gm.F[1]'*Cxy*gm.F[2])
# the following must be the identity matrix out of a possible sign ambiguity
@test abs.(m.F[1]'*gm.F[1]) ≈ I
@test abs.(m.F[2]'*gm.F[2]) ≈ I

# COMPLEX data: check that for the case m=2 gMCA gives the same result as MCA
t, m, n, noise = 20, 2, 6, 0.1
Xcset=getData(ComplexF64, t, m, n, noise)
Ccx=(Xcset[1]*Xcset[1]')/t
Ccy=(Xcset[2]*Xcset[2]')/t
Ccxy=(Xcset[1]*Xcset[2]')/t

gmc=gmca(Xcset; simple=true)
mc=mca(Ccxy; simple=true)

# for complex data just do a sanity check as the order of vectors
# is arbitrary
@test spForm(mc.F[1]'gmc.F[1])<0.01
@test spForm(mc.F[2]'gmc.F[2])<0.01

# REAL data: m>2 case
t, m, n, noise = 20, 4, 6, 0.1
Xset=getData(t, m, n, noise)

# ... selecting subspace dimension allowing an explained variance = 0.9
gm=gmca(Xset, eVar=0.9)

# name of the filter
gm.name

𝒞=Array{Matrix}(undef, 1, m, m)
for i=1:m, j=1:m 𝒞[1, i, j]=(Xset[i]*Xset[j]')/t end

using Plots
# plot regularized accumulated eigenvalues
plot(gm.arev)

# plot the original cross-covariance matrices and the rotated
# cross-covariance matrices

# Get all products 𝐔[i]' * 𝒞[l, i, j] * 𝐔[j]
function _rotate_crossCov(𝐔, 𝒞, m, k)
    𝒮=Array{Matrix}(undef, k, m, m)
    @inbounds for l=1:k, i=1:m, j=1:m 𝒮[l, i, j]=𝐔[i]'*𝒞[l, i, j]*𝐔[j] end
    return 𝒮
end

# Put all cross-covariances in a single matrix of dimension m*n x m*n for visualization
function 𝒞2Mat(𝒞::AbstractArray, m, k)
    n=size(𝒞[1, 1, 1], 1)
    C=Matrix{Float64}(undef, m*n, m*n)
    for i=1:m, j=1:m, x=1:n, y=1:n C[i*n-n+x, j*n-n+y]=𝒞[k, i, j][x, y] end
    return C
end

 C=𝒞2Mat(𝒞, m, 1)
 Cmax=maximum(abs.(C));
 h1 = heatmap(C, clim=(-Cmax, Cmax), yflip=true, c=:bluesreds, title="all cross-covariances")
 𝒮=_rotate_crossCov(gm.F, 𝒞, m, 1)
 S=𝒞2Mat(𝒮, m, 1)
 Smax=maximum(abs.(S));
 h2 = heatmap(S, clim=(0, Smax), yflip=true, c=:amp, title="all rotated cross-covariances")
 📈=plot(h1, h2, size=(700,300))
# savefig(📈, homedir()*"\\Documents\\Code\\julia\\Diagonalizations\\docs\\src\\assets\\FiggMCA.png")

# COMPLEX data: m>2 case
t, m, n, noise = 20, 4, 6, 0.1
Xcset=getData(ComplexF64, t, m, n, noise)

# ... selecting subspace dimension allowing an explained variance = 0.9
gmc=gmca(Xcset, eVar=0.9)


### TRY numero 2

tile_size = 4
board_size = 8 # must be even
mini_board = [zeros(tile_size, tile_size) ones(tile_size, tile_size);
              ones(tile_size, tile_size) zeros(tile_size, tile_size)]

chessboard = repeat(mini_board, outer=(convert(Integer,(board_size/2)), convert(Integer,(board_size/2))))
chessboard = repeat(mini_board, outer=(2, 2))

img = chessboard

### test clean for u

chess_norm = chessboard
chess_norm_f = Float32.(chess_norm)

chess_warp_f = copy(chess_norm_f)
for k in 2:size(chess_warp_f)[1]
    for l in 2:size(chess_warp_f)[2]
        chess_warp_f[k-1, l-1] = chess_warp_f[k, l]
        if l == size(chess_warp_f)[2]
            chess_warp_f[k-1, l] = chess_warp_f[k, l]
        end
        if k == size(chess_warp_f)[1]
            chess_warp_f[k, l-1] = chess_warp_f[k, l]
        end
    end
end

filter_half_size = 2
filter_size = filter_half_size * 2 + 1

u_est, coeffs = LAP_julia.registration.single_lap(chess_norm_f, chess_warp_f, 3, filter_half_size, [filter_size, filter_size])

window_half_size = [1, 1]
u_out = LAP_julia.helpers.smooth_with_gaussian(u_est[3:end-2, 3:end-2], window_half_size)

LAP_julia.visualise.showflow(u_est[3:end-2, 3:end-2])

LAP_julia.visualise.showflow(u_out)

### test gen rand flow
for k in 1:2
    out = gen_tiled_flow((100,100),10)
    showflow(out, 1)
end

out = gen_tiled_flow((100,100),10)
showflow(out)
# --->works

### Tim Holy INTERPOLATION example

using Interpolations, ColorVectorSpace, ImageFiltering, Images

function imWarp(img, dx, dy)
    itp = interpolate(img, BSpline(Linear()))
    inds = indices_spatial(img)
    rng = extrema.(inds)
    imw = similar(img, eltype(itp))
    println(rng[1])
    for I in CartesianIndices(inds)
        dxi, dyi = dx[I], dy[I]
        y, x = I[1]+dyi, I[2]+dxi
        if is_in(y, rng[1]) && is_in(x,rng[2])
            imw[I] = itp(y, x)
        else
            imw[I] = 0
        end
    end
    return imw
end

function is_in(x, (low, high))
    return x >= low && x <= high
end

is_in(1.0, (1, 512))

# Usage demo
using TestImages
img = testimage("mandril");
img

# Let's use a random warp. We blur it with a Gaussian kernel to make it reasonably smooth.
using ImageFiltering

#test1:
kern = KernelFactors.IIRGaussian((10,10))  # IIRGaussian is fast even for very large σ
dx, dy = imfilter(100*randn(size(img)), kern), imfilter(100*randn(size(img)), kern);
imgw = imWarp(img, dx, dy);
imgw

#test2:
flow = gen_tiled_flow(size(img), 50);
dx, dy = real(flow), imag(flow);
imgw = imWarp(img, dx, dy);
imgw

#test3:
n = 20
grid = ones(512, 512);
@views grid[1:n:end, :] .= 0;
@views grid[:, 1:n:end] .= 0;
fig = PyPlot.figure(dpi = 400, figsize = (5, 5))
PyPlot.imshow(grid, cmap = :gray); gcf()
flow = gen_tiled_flow(size(grid), 50);
dx, dy = real(flow), imag(flow);
imgw = imWarp(grid, dx, dy);

fig = PyPlot.figure(dpi = 400, figsize = (5, 5))
PyPlot.imshow(imgw, cmap = :gray); gcf()

fig = PyPlot.figure(dpi = 400, figsize = (5, 5))
PyPlot.close_figs()
showflow(flow,200)

## pyplot disp image !!!! Does not work for TestImages

PyPlot.imshow(chessboard, cmap = :gray)
gcf()

###
using Interpolations

t = 0:.1:1
x = sin.(2π*t)
y = cos.(2π*t)
A = hcat(x,y)

itp = Interpolations.scale(interpolate(A, (BSpline(Cubic(Natural(OnCell()))), NoInterp())), t, 1:2)

tfine = 0:.01:1
xs, ys = [itp(t,1) for t in tfine], [itp(t,2) for t in tfine];

using PyPlot

figure()
scatter(x, y, label="knots")
PyPlot.plot(xs, ys, label="spline")
gcf()

###
using Interpolations
using PyPlot

t = 0:.1:1
x = sin.(2π*t)

itp = Interpolations.interpolate(x, (BSpline(Quadratic(Natural(OnGrid())))))
itp = Interpolations.interpolate(x, (BSpline(Linear())));
sitp = scale(itp, t);
#sitp = itp

tfine = 0:.01:1;
xs = [sitp(t) for t in tfine]

figure()
scatter(t, x, label="knots")
PyPlot.plot(tfine, xs, label="spline")
gcf()

### get rid of unwanted image replication on the right and bottom of images in the interpolation

img = testimage("mandril_gray");
ix = size(img)[1]
iy = size(img)[2]

ixhalf = Int64(ix/2)
iyhalf = Int64(iy/2)

# flow all to centere
flow = Array{ComplexF64}(undef, ix, iy)
flow[1:ixhalf, 1:iyhalf] .= -20 .+ -20 .* im;
flow[1:ixhalf, iyhalf:iy] .= 20 .+ -20 .* im;
flow[ixhalf:ix, 1:iyhalf] .= -20 .+ 20 .* im;
flow[ixhalf:ix, iyhalf:iy] .= 20 .+ 20 .* im;

flow = LAP_julia.helpers.smooth_with_gaussian(flow, 60);

showflow(flow)

dx, dy = real(flow), imag(flow);
imgw = imWarp(img, dx, dy)
imgw

save("img.png", imgw)

###


function psiNext(p::Vector,psi::Vector)
    psi = cat(1,zero(Complex128),psi)
    N = length(psi)
    psik = zero(psi)
    if length(p)==1
        p = cat(1,zero(p[1]),p)
        c = sqrt(1-abs2(p[end]))
        #in this case we assume psi = u(t), rather than the previous psi
        for j=0:(N-1)
            psik[j+1] = c*sum([p[end]^(i)*psi[j-i+1] for i=0:(j-1)])
        end
    else
        c = sqrt((1-abs2(p[end]))/(1-abs2(p[end-1])))
        for j=0:(N-1)
            psik[j+1] = c*sum([p[end]^(i)*(psi[j-i]-conj(p[end-1])*psi[j-i+1]) for i=0:(j-1)])
        end
    end
    return psik[2:end]
end

function psiNextRecursive(p::Vector,psi::Vector)
    psi = cat(1,zero(Complex128),psi)
    N = length(psi)
    psik = zero(psi)
    if length(p)==1
        p = cat(1,zero(p[1]),p)
        c = sqrt(1-abs2(p[end]))
        #in this case we assume psi = u(t), rather than the previous psi
        for j=1:(N-1)
            psik[j+1] = p[end]*psik[j]+c*psi[j+1]
        end
    else
        c = sqrt((1-abs2(p[end]))/(1-abs2(p[end-1])))
        for j=1:(N-1)
            psik[j+1] = p[end]*psik[j] + c*(psi[j] - conj(p[end-1])*psi[j+1])
        end
    end
    return psik[2:end]
end

function genPsiBasis(p::Vector,u::Vector)
    N = length(p)
    B = zeros(Complex128,(length(u),N+1))
    B[:,1] = u
    for k=1:N
        B[:,k+1] = psiNext([p[1:k]],B[:,k])
    end
    return B[:,2:end]
end

function genPsiBasisRecursive(p::Vector,u::Vector)
    N = length(p)
    B = zeros(Complex128,(length(u),N+1))
    B[:,1] = u
    for k=1:N
        B[:,k+1] = psiNextRecursive([p[1:k]],B[:,k])
    end
    return B[:,2:end]
end

function pulseFun(k::Int)
    a = one(Complex128);
    return cat(1,a,zeros(Complex128,k-1))
end

function gaussianPulse(k::Int,fwhm::Float64)
    a = one(Complex128);
    x = [a*exp(-4*log(2)*(i^2/fwhm^2)) for i=1:k]
    return x/norm(x)
end

function unitStep(k::Int)
    return k>=0 ? 1 : 0
end

function dictElem(p::Complex128,K::Int)
    x = [(sqrt(1-abs2(p))*(p^(k-1))*unitStep(k-1))::Complex128 for k=0:(K)]
    return x[2:end]
end

function dictElemRecursive(p::Complex128,u::Vector)
    c = sqrt(1-abs2(p))
    u = cat(1,zero(Complex128),u)
    N = length(u)
    r = zero(u)
    for k=1:(N-1)
        r[k+1] = p*r[k]+c*u[k]
    end
    return r[2:end]
end

function filterR(p::Complex128,R::Vector)
    ai = 1/conj(p)
    R = cat(1,zero(Complex128),R)
    N = length(R)
    Rkp1 = zero(R)
    for k=0:(N-1)
        Rkp1[k+1]  = sum([ai^(i+1)*(p*R[k-i]-R[k-i+1]) for i=0:(k-1)])
    end
    return Rkp1[2:end]
end

function filterRrecursive(p::Complex128,R)
    ai = 1/conj(p)
    Rp = cat(1,zeros(Complex128,(1,size(R,2))),R)
    N = size(Rp,1)
    Rkp1 = zero(R)
    for j=1:size(R,2)
        Rkp1j = zero(Rp[:,j])
        for k=1:(N-1)
            Rkp1j[k+1]  = ai*Rkp1[k]-ai*(Rp[k+1]-p*Rp[k])
        end
        Rkp1[:,j] = Rkp1j[2:end]
    end
    return Rkp1
end

function dictObj(x,g,target)
    p = x[1]*exp(im*2*pi*x[2])
    trial = dictElem(p,size(target,1))
    r = sum(abs2(trial'*target),2)
    return r[1]
end

function dictObjNbest(x,g,target)
    p = [(x[i]*exp(im*2*pi*x[i+1]))::Complex128 for i=1:2:length(x)]
    R = copy(target)
    for k=1:length(p)
        trial = dictElem(p[k],size(target,1))
        #R = filterR(p[k],R - trial*((trial'*R)[1]))
        R = filterRrecursive(p[k],R - trial*((trial'*R)[1]))
    end
    r = norm(R)
    return r
end

function nthRemainder(p,target)
    R = copy(target)
    L = size(target,2);
    for k=1:length(p)
        trial = dictElem(p[k],size(target,1))
        for j=1:L
            R[:,j] = filterRrecursive(p[k],R[:,j] - dot(trial,R[:,j])*trial)
        end
    end
    return R
end

function unitDiskGrid(n::Int,delta::Float64)
    N = nextpow2(n)
    M = N>>1
    v = linspace(delta,one(Float64)-delta,M)
    Zjk = zeros(Complex128,N,N)
    for j=1:(M)
             #bottom, running counter-clockwise
            Zjk[(M+j),((M-(j-1)):(M+1+(j-1)))] = v[j]*exp(im*linspace(5*pi/4,7*pi/4,2*j))
            #top, running clockwise
            Zjk[(M-j+1),((M-(j-1)):(M+1+(j-1)))] = v[j]*exp(im*linspace(3*pi/4,1*pi/4,2*j))
    end
    for j=1:(M)
            #right, running clockwise
            Zjk[((M-(j-1)):(M+1+(j-1))),M+j] = v[j]*exp(im*linspace(pi/4,-pi/4,2*j)) #right
            #left, running counter-clockwise
            Zjk[((M-(j-1)):(M+1+(j-1))),M-j+1] = v[j]*exp(im*linspace(3*pi/4,5*pi/4,2*j)) #left
    end
    return reshape(Zjk,N^2)
end

function afdT(target::Array,Npoles::Int,Nstarts::Int)
    Z = unitDiskGrid(int(sqrt(Nstarts)),0.2)
    opt = Opt(:LN_SBPLX,2)
    ftol_rel!(opt,1e-6) #local stopping criteria
    lb = [0.01; 0]
    ub = [0.999; 1-eps(Float64)]
    lower_bounds!(opt,lb)
    upper_bounds!(opt,ub)
    avec = zeros(Complex128,(length(Z),2))
    ovec = zeros(Float64,length(Z))
    p = zeros(Complex128,Npoles)
    a = zeros(Complex128,(Npoles,size(target,2)))
    gk = copy(target)
    for j=1:Npoles
        max_objective!(opt,(x,g)->dictObj(x,g,gk)[1])
        for k=1:length(Z)
            beta0 = [abs(Z[k]); ((angle(Z[k])+pi)/(2*pi))]
            (ovec[k],avec[k,:],e) = optimize(opt,beta0)
        end
        aopt = avec[indmax(ovec),:]
        p[j] = aopt[1]*exp(im*2*pi*aopt[2])
        ep = dictElem(p[j],size(target,1))
        a[j,:] = ep'*gk
        #for i=1:size(target,2)
        #    gk[:,i] = filterRrecursive(p[j],(gk-a[j,i]*ep))
        #end
        gk = filterRrecursive(p[j],(gk-ep*a[j,:]))
    end
    return (p,a)
end

function afdTCyclic(target::Array,polesIn::Vector,Nstarts::Int,CycleTol::Float64,CycleMax::Int)
    Z = unitDiskGrid(int(sqrt(Nstarts)),0.01)
    Npoles = length(polesIn)
    opt = Opt(:LN_SBPLX,2)
    ftol_rel!(opt,1e-12) #local stopping criteria
    lb = [0.01; 0]
    ub = [0.9999999999; 1-eps(Float64)]
    lower_bounds!(opt,lb)
    upper_bounds!(opt,ub)
    avec = zeros(Complex128,(length(Z),2))
    ovec = zeros(Float64,length(Z))
    p = zero(polesIn)
    prel = CycleTol+1;
    cycle_count = 0;
    while ((prel>CycleTol) && (cycle_count<CycleMax))
        cycle_count += 1
        for j=1:Npoles
            pold = copy(p)
            pprime = pold[(1:Npoles).!=j]
            gk = nthRemainder(pprime,target)
            max_objective!(opt,(x,g)->dictObj(x,g,gk)[1])
            for k=1:length(Z)
                beta0 = [abs(Z[k]); ((angle(Z[k])+pi)/(2*pi))]
                (ovec[k],avec[k,:],errcode) = optimize(opt,beta0)
            end
            aopt = avec[indmax(ovec),:]
            p[j] = aopt[1]*exp(im*2*pi*aopt[2])
            prel = norm(p-pold)/norm(pold)
            #println(prel)
        end
    end
    B = genPsiBasisRecursive(p,pulseFun(size(target,1)))
    a = B\target
    #println(cycle_count)
    return (p,a)
end

function poles2vec(p::Vector{Complex128})
    r = zeros(Float64,length(p)*2)
    r[1:2:end] = abs(p)
    r[2:2:end] = (angle(p)+pi)/(2*pi)
    return r
end

function afdTNbest(target::Array,Npoles::Int,gWallClockTime::Number,beta0::Vector)
    opt = Opt(:GN_MLSL_LDS,Npoles*2)
    lopt = Opt(:LN_SBPLX,Npoles*2)
    ftol_rel!(lopt,1e-6) #local stopping criteria
    maxtime!(opt,gWallClockTime) #global stopping critera
    lb = [0.01; 0]
    ub = [0.99999; 1-eps(Float64)]
    lb = repmat(lb,Npoles)
    ub = repmat(ub,Npoles)
    lower_bounds!(opt,lb)
    upper_bounds!(opt,ub)
    local_optimizer!(opt,lopt)
    p = zeros(Complex128,Npoles)
    a = zeros(Complex128,(Npoles,size(target,2)))
    gk = copy(target)
    min_objective!(opt,(x,g)->dictObjNbest(x,g,gk)[1])
    (r,popt,e) = optimize(opt,beta0)
    p = [(popt[i]*exp(im*2*pi*popt[i+1]))::Complex128 for i=1:2:(Npoles*2)]
    B = genPsiBasis(p,pulseFun(size(target,1)))
    a = B\target
    return (p,a)
end

function akaikeInformationCriterion(RSS,numParams,dataLen)
    #This gives the sample size corrected Akaike Information Criteria
    #commonly known as AICc in the literature
    K = numParams+1
    return dataLen*log(RSS/dataLen)+2*K+2*K*((K+1)/(dataLen-K-1))
end

function processDataStack(Y::Array{Complex128,2},minNpoles::Int,maxNpoles::Int,Nstarts::Int,Textrap::Int)
    pmat = {}
    amat = {}
    Bmat = {}
    for j=1:size(Y,2)
        AICc = {}
        ptemp   = {}
        atemp   = {}
        for k=minNpoles:maxNpoles
            (p,a) = afdT(Y[:,j],k,Nstarts)
            push!(ptemp,p)
            push!(atemp,a)
            B = genPsiBasis(p,pulseFun(size(Y,1)))
            Yhat = squeeze(B*a,2)
            R = (Y[:,j]-Yhat)
            RSS = abs(dot(R,R))
            aicval = akaikeInformationCriterion(RSS,4*k,2*size(Y,1))
            push!(AICc,aicval)
        end
        optInd = indmin(AICc)
        push!(pmat,ptemp[optInd])
        push!(amat,atemp[optInd])
        B = genPsiBasis(pmat[j],pulseFun(Textrap))
        push!(Bmat,B)
        println(optInd+minNpoles-1)
        println(AICc[optInd])
    end
    That = zeros(Complex128,(Textrap,size(Y,2)))
    for k=1:size(Y,2)
        That[:,k] = Bmat[k]*amat[k]
    end
    return (That,pmat,amat)
end

function processDataStackCyclic(Y::Array{Complex128,2},minNpoles::Int,maxNpoles::Int,maxNcycles::Int,Nstarts::Int,CycleTol::Float64,Textrap::Int)
    pmat = {}
    amat = {}
    Bmat = {}
    for j=1:size(Y,2)
        AICc = {}
        ptemp   = {}
        atemp   = {}
        for k=minNpoles:maxNpoles
            (pinit,ainit) = afdT(Y[:,j],k,Nstarts)
            (p,a) = afdTCyclic(Y[:,j],pinit,Nstarts,CycleTol,maxNcycles)
            push!(ptemp,p)
            push!(atemp,a)
            B = genPsiBasis(p,pulseFun(size(Y,1)))
            Yhat = B*a
            R = (Y[:,j]-Yhat)
            RR = R'*R
            RSS = abs(RR)[1]
            aicval = akaikeInformationCriterion(RSS,2*length(p)+2*length(a),2*size(Y,1))
            push!(AICc,aicval)
        end
        optInd = indmin(AICc)
        push!(pmat,ptemp[optInd])
        push!(amat,atemp[optInd])
        B = genPsiBasis(pmat[j],pulseFun(Textrap))
        push!(Bmat,B)
        println(optInd+minNpoles-1)
        println(AICc)
    end
    That = zeros(Complex128,(Textrap,size(Y,2)))
    for k=1:size(Y,2)
        That[:,k] = Bmat[k]*amat[k]
    end
    return (That,pmat,amat)
end

function processDataStackCyclicGlobal(Y::Array{Complex128,2},minNpoles::Int,maxNpoles::Int,maxNcycles::Int,Nstarts::Int,CycleTol::Float64,Textrap::Int)
    Bmat = {}
    AICc = {}
    ptemp   = {}
    atemp   = {}
    for k=minNpoles:maxNpoles
        (pinit,ainit) = afdT(Y,k,Nstarts)
        (p,a) = afdTCyclic(Y,pinit,Nstarts,CycleTol,maxNcycles)
        push!(ptemp,p)
        push!(atemp,a)
        B = genPsiBasis(p,pulseFun(size(Y,1)))
        Yhat = B*a
        R = (Y-Yhat)
        RSS = sum((abs(R'*R)))
        aicval = akaikeInformationCriterion(RSS,2*length(p)+2*length(a),2*length(Y))
        println(aicval)
        push!(AICc,aicval)
    end
    optInd = indmin(AICc)
    popt = ptemp[optInd]
    aopt = atemp[optInd]
    B = genPsiBasis(popt,pulseFun(Textrap))
    #println(optInd+minNpoles-1)
    #println(AICc[optInd])
    That = B*aopt
    return (That,popt,aopt)
end

function processDataStackMLSL(Y::Array{Complex128,2},minNpoles::Int,maxNpoles::Int,Nstarts::Int,SearchTime::Number,Textrap::Int)
    pmat = {}
    amat = {}
    Bmat = {}
    for j=1:size(Y,2)
        AICc = {}
        ptemp   = {}
        atemp   = {}
        for k=minNpoles:maxNpoles
            (pinit,ainit) = afdT(Y[:,j],k,Nstarts)
            (p,a) = afdTNbest(Y[:,j],k,SearchTime,poles2vec(pinit))
            push!(ptemp,p)
            push!(atemp,a)
            B = genPsiBasis(p,pulseFun(size(Y,1)))
            Yhat = B*a
            R = (Y[:,j]-Yhat)
            RSS = abs(dot(R,R))
            aicval = akaikeInformationCriterion(RSS,4*k,2*size(Y,1))
            push!(AICc,aicval)
        end
        optInd = indmin(AICc)
        push!(pmat,ptemp[optInd])
        push!(amat,atemp[optInd])
        B = genPsiBasis(pmat[j],pulseFun(Textrap))
        push!(Bmat,B)
        println(optInd+minNpoles-1)
        println(AICc[optInd])
    end
    That = zeros(Complex128,(Textrap,size(Y,2)))
    for k=1:size(Y,2)
        That[:,k] = Bmat[k]*amat[k]
    end
    return (That,pmat,amat)
end

function processDataStackNpole(Y::Array{Complex128,2},Nsing::Int,Npoles::Int,Nstarts::Int,Textrap::Int)
    pmat = {}
    amat = {}
    Bmat = {}
    for j=1:Nsing
        (p,a) = afdT(Y[:,j],k,Nstarts)
        B = genPsiBasis(p,pulseFun(size(Y,1)))
        push!(pmat,p)
        push!(amat,a)
        B = genPsiBasis(pmat[j],pulseFun(Textrap))
        push!(Bmat,B)
    end
    That = zeros(Complex128,(Textrap,Nsing))
    for k=1:Nsing
        That[:,k] = Bmat[k]*amat[k]
    end
    return (That,pmat,amat)
end

function refineDataStack(Y::Array{Complex128,2},Nsing::Int,pmat,SearchTime::Number,Textrap::Int)
    Bmat = {}
    pref = {}
    aref = {}
    for j=1:Nsing
        pin = poles2vec(pmat[j])
        (p,a) = afdTNbest(Y[:,j],length(pmat[j]),SearchTime,pin)
        push!(pref,p)
        push!(aref,a)
        B = genPsiBasis(pref[j],pulseFun(Textrap))
        push!(Bmat,B)
    end
    That = zeros(Complex128,(Textrap,Nsing))
    for k=1:Nsing
        That[:,k] = Bmat[k]*aref[k]
    end
    return (That,pref,aref)
end

function refineDataStackCyclic(Y::Array{Complex128,2},Nsing::Int,pmat,Nstarts::Int,CycleTol::Float64,CycleMax::Int,Textrap::Int)
    Bmat = {}
    pref = {}
    aref = {}
    for j=1:Nsing
        (p,a) = afdTCyclic(Y[:,j],pmat[j],Nstarts,CycleTol,CycleMax)
        push!(pref,p)
        push!(aref,a)
        B = genPsiBasis(pref[j],pulseFun(Textrap))
        push!(Bmat,B)
    end
    That = zeros(Complex128,(Textrap,Nsing))
    for k=1:Nsing
        That[:,k] = Bmat[k]*aref[k]
    end
    return (That,pref,aref)
end

function unfoldMatrix(M::Array,N::Int)
    uM = M[:,1:N]
    for k=2:int(size(M,2)/N)
        uM = cat(1,uM,M[:,((k-1)*N+1):(k*N)]);
    end
    return uM
end

function cubeMatrix(M::Array,N::Int)
    cM = zeros(typeof(M[1]),(N,size(M,2),int(size(M,1)/N)))
    for k=1:int(size(M,1)/N)
        cM[:,:,k] = M[((k-1)*N+1):(k*N),:]
    end
    return cM
end

function fftaxis(N::Int,TSpac=1.0)
    if abs(mod(N,2))>0
        a = [0:1:int((N-1)/2)]
        b = [-(N-1)/2:1:-1]
        r = cat(1,a,b)/(N*TSpac)
    else
        a = [0:1:int(N/2-1)]
        b = [-int(N/2):1:-1]
        r = cat(1,a,b)/(N*TSpac)
    end
    return r
end

function unfoldCube(M::Array)
    R = zeros(typeof(M[1]),(size(M,1)*size(M,2),size(M,3)))
    N = size(M,1);
    for j=1:size(M,2)
        R[((j-1)*N+1):(j*N),:] = squeeze(M[:,j,:],2);
    end
    return R
end

function optimalSVDTrunc(u,s,v)
    N = size(v,1)
    M = size(v,1)
    if isreal(u[1])
        f = 1
    else
        f = 2;
    end
    AICc = {}
    Y = u*diagm(s)*v';
    for k=1:(N-1)
        Yhat = u[:,1:k]*diagm(s[1:k])*v[:,1:k]'
        R = Y-Yhat
        RSS = abs(R'*R)[1]
        push!(AICc,akaikeInformationCriterion(RSS,k,f*N*M))
    end
    optInd = indmin(AICc)
    return (AICc,optInd)
end

function w2wn(w::Vector)
    c = 299.792458; #nm/fs
    return 1E7*w/(2*pi*c)
end


using Random; Random.seed!(1234)

include("src/data.jl")
using .DataUtil
include("src/forward_stepwise/selective_inference.jl")
using .SIforSFS
include("src/forward_stepwise/data_splitting.jl")
import .DS
using LinearAlgebra:I
using Distributions
using ProgressMeter

const N = 100
const β₀ = 0.0
const 𝛃 = (Float64)[0, 0, 0, 0, 0] # null model
const k = 3

# FPR experiment
function FPR_experiment(β₀, 𝛃, N, k; event=Active(), α=0.05, denominator_upper=100, iter_size=5)
    H0 = filter(i -> iszero(𝛃[i]), 1:length(𝛃)) # {j ∈ {1,...,d} | H_{0,j} is true}
    FPR = Array{Float64}(undef, length(H0), iter_size)
    @showprogress "[$(N), $(event)]: " for l = 1:iter_size
        numerator = zeros(Int, length(H0))
        denominator = zeros(Int, length(H0))
        for _ = 1:denominator_upper
            𝐗, 𝐲_obs = centering(make_dataset(β₀, 𝛃, N)...)
            𝐴_obs, p_values = parametric_SFS_SI(𝐗, 𝐲_obs, k, I; selection_event=event)
            for (j, p_value) in zip(𝐴_obs, p_values)
                for (i, j′) in enumerate(H0)
                    if j == j′ # H_{0,j} is true
                        numerator[i] += p_value < α ? 1 : 0
                        denominator[i] += 1
                        break
                    end
                end
            end
        end
        FPR[:, l] = numerator ./ denominator
    end
    return FPR
end

α = 0.05
zero_num = length(filter(x -> iszero(x), 𝛃))
Ns = [50, 100, 150, 200]
iter_size = 20
FPR_A = fill(NaN, zero_num, iter_size, length(Ns))
FPR_Ao = fill(NaN, zero_num, iter_size, length(Ns))
FPR_As = fill(NaN, zero_num, iter_size, length(Ns))
FPR_Aso = fill(NaN, zero_num, iter_size, length(Ns))
FPR_DS = fill(NaN, zero_num, iter_size, length(Ns))
for (i, N) in enumerate(Ns)
    FPR_A[:,:,i] = Main.FPR_experiment(β₀, 𝛃, N, k; event=Active(), α=α, iter_size=iter_size)
    FPR_As[:,:,i] = Main.FPR_experiment(β₀, 𝛃, N, k; event=ActiveSign(), α=α, iter_size=iter_size)
    FPR_Ao[:,:,i] = Main.FPR_experiment(β₀, 𝛃, N, k; event=ActiveOrder(), α=α, iter_size=iter_size)
    FPR_Aso[:,:,i] = Main.FPR_experiment(β₀, 𝛃, N, k; event=ActiveSignOrder(), α=α, iter_size=iter_size)
    FPR_DS[:,:,i] = DS.FPR_experiment(β₀, 𝛃, N, k; α=α, iter_size=iter_size)
end

using Statistics
FPR_A_mean = reshape(mean(FPR_A, dims=2), zero_num, length(Ns))
FPR_As_mean = reshape(mean(FPR_As, dims=2), zero_num, length(Ns))
FPR_Ao_mean = reshape(mean(FPR_Ao, dims=2), zero_num, length(Ns))
FPR_Aso_mean = reshape(mean(FPR_Aso, dims=2), zero_num, length(Ns))
FPR_DS_mean = reshape(mean(FPR_DS, dims=2), zero_num, length(Ns))

using Plots
Cs = [:red, :orange, :green, :blue, :purple]
p = Array{typeof(plot())}(undef, zero_num)
for j in 1:zero_num
    plot()
    plot!(1:length(Ns), FPR_A_mean[j,:], label="Homotopy", color=Cs[1], markershape=:circle)
    plot!(1:length(Ns), FPR_Ao_mean[j,:], label="Homotopy-H", color=Cs[2], markershape=:circle)
    plot!(1:length(Ns), FPR_As_mean[j,:], label="Homotopy-S", color=Cs[3], markershape=:circle)
    plot!(1:length(Ns), FPR_Aso_mean[j,:], label="Polytope", color=Cs[4], markershape=:circle)
    plot!(1:length(Ns), FPR_DS_mean[j,:], label="DS", color=Cs[5], markershape=:circle)
    # plot setting
    plot!(; ylims=(0, 0.3))
    plot!(; xlabel="Sample size", ylabel="False Positive Rate (FPR)")
    plot!(; legend=:topright)
    plot!(; guidefontsize=16, legendfontsize=12, tickfontsize=10)
    p[j] = plot!(; xticks=(1:length(Ns), Ns))
    savefig(p[j], "img/FPR/FPR$(j).pdf")
end


module MemorylessNonlinearities

using AlphaStableDistributions, DelimitedFiles, Interpolations, QuadGK

export
    Arctangent,
    Blanking,
    CauchyNL,
    Clipping,
    HampelThreePart,
    SαSNL,
    SoftClipping,
    TurkeyBiweight,
    InverseHyperbolicSine,

    filt,
    minmaxrescale,

    NLnames

include("utils.jl")

const DATADIR = joinpath(dirname(pathof(@__MODULE__)), "..", "data")

const _xsas = -200:0.1:200
const _αsas = 1.0:0.001:2.0
const _ysas = readdlm(joinpath(DATADIR, "sas.csv"), ',', Float64)

abstract type AbstractMemorylessNonlinearity end

struct Arctangent{T<:Real} <: AbstractMemorylessNonlinearity
    α::T
end

struct Blanking{T<:Real} <: AbstractMemorylessNonlinearity
    k::T
end

struct CauchyNL{T<:Real} <: AbstractMemorylessNonlinearity
    k::T
end

struct Clipping{T<:Real} <: AbstractMemorylessNonlinearity
    k::T
end

struct HampelThreePart{T<:Real} <: AbstractMemorylessNonlinearity
    a::T
    b::T
    c::T
end

struct SαSNL{T<:Real} <: AbstractMemorylessNonlinearity
    α::T
    scale::T
    location::T
    approx::Bool
end
SαSNL(α, scale, location) = SαSNL(α, scale, location, true)

struct SoftClipping{T<:Real} <: AbstractMemorylessNonlinearity
    k::T
end

struct TurkeyBiweight{T<:Real} <: AbstractMemorylessNonlinearity
    k::T
end

struct InverseHyperbolicSine{T<:Real} <: AbstractMemorylessNonlinearity
    θ::T
end

const NLnames = [:Arctangent,
                 :Blanking,
                 :CauchyNL,
                 :Clipping,
                 :HampelThreePart,
                 :SαSNL,
                 :SoftClipping,
                 :TurkeyBiweight,
                 :InverseHyperbolicSine]

"""
Arctangent nonlinearity.
"""
arctangent(x::S, α::T) where {S<:Real,T<:Real} = (2 / π) * atan(α * x)
function filt(f::Arctangent, x::AbstractVector)
    arctangent.(x, f.α)
end

"""
Blanking nonlinearity.
"""
blanking(x::S, k::T) where {S<:Real,T<:Real} = abs(x) > k ? S(0) : x
function filt(f::Blanking, x::AbstractVector)
    blanking.(x, f.k)
end

"""
Cauchy nonlinearity.
"""
cauchynl(x::S, k::T) where {S<:Real,T<:Real} = (2 * (x / k)) / (1 + (x / k)^2)
function filt(f::CauchyNL, x::AbstractVector)
    cauchynl.(x, f.k)
end

"""
Clipping nonlinearity.
"""
clipping(x::S, k::T) where {S<:Real,T<:Real} = abs(x) > k ? sign(x) * k : x
function filt(f::Clipping, x::AbstractVector)
    clipping.(x, f.k)
end

"""
Hampel's Three Part Redescending nonlinearity.
"""
function hampelthreepart(x::S, a::T, b::T, c::T) where {S<:Real,T<:Real}
    absx = abs(x)
    if absx < a
        return x
    elseif absx >= a && absx < b
        return sign(x) * a
    elseif absx >= b && absx < c
        return (absx - c) * (a / (b - c)) * sign(x)
    else
        return 0.
    end
end
function filt(f::HampelThreePart, x::AbstractVector)
    hampelthreepart.(x, f.a, f.b, f.c)
end

"""
Symmetric Alpha Stable nonlinearity.

`x` is standard symmetric alpha stable distributed.
"""
function sαsnl(x::S, α) where {S<:Real}
    num, _ = quadgk(t -> t * sin(t * x) * exp(-t^α), 0, Inf; atol=1e-12, order=3)
    den, _ = quadgk(t -> cos(t * x) * exp(-t^α), 0, Inf; atol=1e-12, order=3)
    num / den
end
function filt(f::SαSNL, x::AbstractVector)
    (abs(f.α - 2.0) < 0.001) && return (x ./ 2) # Gaussian
    (abs(f.α - 1.0) < 0.001) && return (2 .* x) ./ (1 .+ x.^2) # Cauchy
    xstd = (collect(x) .- f.location) ./ f.scale
    if f.approx
        maxval = maximum(_xsas)
        minval = minimum(_xsas)
        xstd[xstd .> maxval] .= (f.α+1) ./ xstd[xstd .> maxval]
        xstd[xstd .< minval] .= (f.α+1) ./ xstd[xstd .< minval]
        itp = interpolate((_αsas, _xsas), _ysas, Gridded(Linear()))
        xstd[(xstd .<= maxval) .& (xstd .>= minval)] .= itp(f.α, xstd[(xstd .<= maxval) .& (xstd .>= minval)])
        xstd
    else
        sαsnl.(xstd, f.α)
    end
end

"""
Soft clipping nonlinearity.
"""
function softclipping(x::S, k::T) where {S<:Real,T<:Real}
    abs(x) > k ? sign(x) * k : k * (3*(x / k) / 2) * (1- ((x/ k)^2) / 3)
end
function filt(f::SoftClipping, x::AbstractVector)
    softclipping.(x, f.k)
end

"""
Turkey's biweight nonlinearity.
"""
turkeybiweight(x::S, k::T) where {S<:Real,T<:Real} = abs(x) > k ? 0. : x * (1 - (x / k)^2)^2
function filt(f::TurkeyBiweight, x::AbstractVector)
    turkeybiweight.(x, f.k)
end

"""
Inverse hyperbolic sine nonlinearity.
"""
inversehyperbolicsine(x::S, θ::T) where {S<:Real,T<:Real} = asinh(x*θ)
function filt(f::InverseHyperbolicSine, x::AbstractVector)
    inversehyperbolicsine.(x, f.θ)
end

end


module Softmax_CEL
    using LoopVectorization

    function opt_func(values)
        return @avxt exp.(values.-maximum(values))/sum(exp.(values.-maximum(values)))
    end

    function func!(output_matrix::Array{Float32}, value_matrix::Array{Float32})
        @avxt for i in axes(value_matrix, 1), j in axes(value_matrix, 2)
            value_matrix[i,j] = ifelse(value_matrix[i,j]>3.0f38, 3.0f38, value_matrix[i,j])
            value_matrix[i,j] = ifelse(value_matrix[i,j]<-3.0f38, -3.0f38, value_matrix[i,j])
        end
        for b in axes(value_matrix, 2)
            output_matrix[:,b] = opt_func(value_matrix[:,b])
        end
    end

    function pre_diff(inputs::Array{Float32}, position::Int64)
        outputs = opt_func(inputs)
        derivative_vector = zeros(Float32, size(inputs))
        @avxt for i in 1:length(derivative_vector)
            derivative_vector[i] = ifelse(i!=position, outputs[i], outputs[i]-1)
        end
        return derivative_vector
    end

    function get_∇biases!(∇biases::Array{Float32}, input_matrix::Array{Float32}, δ::Array{Float32})
        @avxt for i in axes(∇biases, 1), j in axes(∇biases, 2)
            ∇biases[i,j] = 0.0f0
        end

        for b in axes(input_matrix, 2)
            for l in axes(input_matrix, 1)
                if δ[l,b]!=0
                    ∇biases[:,b] += pre_diff(input_matrix[:,b], l)
                end
            end
        end
    end

    function get_name()
        return "Softmax_CEL"
    end
end

module Softmax
    using LoopVectorization

    function opt_func(values)
        return @avxt exp.(values.-maximum(values))/sum(exp.(values.-maximum(values)))
    end

    function func!(output_matrix::Array{Float32}, value_matrix::Array{Float32})
        @avxt for i in axes(value_matrix, 1), j in axes(value_matrix, 2)
            value_matrix[i,j] = ifelse(value_matrix[i,j]>3.0f38, 3.0f38, value_matrix[i,j])
            value_matrix[i,j] = ifelse(value_matrix[i,j]<-3.0f38, -3.0f38, value_matrix[i,j])
        end
        for b in axes(value_matrix, 2)
            output_matrix[:,b] = opt_func(value_matrix[:,b])
        end
    end

    function pre_diff(inputs::Array{Float32}, position::Int64)
        outputs = opt_func(inputs)
        derivative_vector = zeros(Float32, size(inputs))
        @avxt for i in 1:length(derivative_vector)
            derivative_vector[i] = ifelse(i!=position, -outputs[i]*outputs[position], outputs[i]*(1-outputs[i]))
        end
        return derivative_vector
    end

    function get_∇biases!(∇biases::Array{Float32}, input_matrix::Array{Float32}, δ::Array{Float32})
        @avxt for i in axes(∇biases, 1), j in axes(∇biases, 2)
            ∇biases[i,j] = 0.0f0
        end

        for b in axes(input_matrix, 2)
            for l in axes(input_matrix, 1)
                if δ[l,b]!=0
                    ∇biases[:,b] += pre_diff(input_matrix[:,b], l)
                end
            end
        end

        @avxt for i in axes(∇biases, 1), j in axes(∇biases, 2)
            ∇biases[i,j] *= δ[i,j]
        end
    end

    function get_name()
        return "Softmax"
    end
end


using BenchmarkTools
using Random
using CovarianceEstimation

const SUITE = BenchmarkGroup()

Random.seed!(1234)

const X10_100  = randn(15, 100)
const X100_100  = randn(100, 100)
const X100_1000  = randn(100, 1000)
const X100_10000  = randn(100, 10000)
const X1000_100  = randn(1000, 100)
const X1000_10000  = randn(1000, 10000)

matrices = [X10_100, X100_100, X100_1000, X100_10000, X1000_10000]
estimators = Dict(
    "Simple" => (ce = SimpleCovariance(), maxfeatures = 10000),
    "Linear constant correlation" => (ce = LinearShrinkage(ConstantCorrelation()), maxfeatures = 1000),
    "Linear diagonal unit variance" => (ce = LinearShrinkage(DiagonalUnitVariance()), maxfeatures = 1000),
    "Linear diagonal common variance" => (ce = LinearShrinkage(DiagonalCommonVariance()), maxfeatures = 1000),
    "Linear common covariance" => (ce = LinearShrinkage(CommonCovariance()), maxfeatures = 1000),
    "Linear diagonal unequal variance" => (ce = LinearShrinkage(DiagonalUnequalVariance()), maxfeatures = 1000),
    "Linear perfect positive correlation" => (ce = LinearShrinkage(PerfectPositiveCorrelation()), maxfeatures = 1000),
    "Linear RBLW" => (ce = LinearShrinkage(DiagonalCommonVariance(), :rblw), maxfeatures = 10000),
    "Linear OAS" => (ce = LinearShrinkage(DiagonalCommonVariance(), :oas), maxfeatures = 10000),
    "Analytical nonlinear shrinkage" => (ce = AnalyticalNonlinearShrinkage(), maxfeatures = 1000)
)

for (k, v) in estimators
    SUITE[k] = BenchmarkGroup()
    for m in matrices
        if size(m, 2) ≤ v[:maxfeatures]
            SUITE[k]["Size $(size(m))"] = @benchmarkable cov($v[:ce], $m)
        end
    end
end


using Test
using GeoParams

@testset "DiffusionCreepLaws" begin

# This tests the MaterialParameters structure
CharUnits_GEO   =   GEO_units(viscosity=1Pa*s, length=1m);

# Define a linear viscous creep law ---------------------------------
x1      =   DiffusionCreep()
@test x1.n.val == 1.0
@test x1.p.val == -3.0
@test x1.A.val == 1.5MPa^-1.0*s^-1*m^(3.0)

x2      =   DiffusionCreep(n=3, p=-3.0)
@test x2.A.val == 1.5MPa^-3.0*s^-1*m^(3.0)

# perform a computation with the dislocation creep laws
    # Calculate EpsII, using a set of pre-defined values
CharDim = GEO_units(length=1000km, viscosity=1e19Pa*s, stress=100MPa, temperature=1000C)
EpsII   = GeoUnit(1.0s^-1.0)
        Nondimensionalize!(EpsII,CharDim)
TauII   = GeoUnit(0.3MPa)
        Nondimensionalize!(TauII,CharDim)
P       = GeoUnit(1.0e9Pa)
        Nondimensionalize!(P,CharDim)
T       = GeoUnit(1400C)
        Nondimensionalize!(T,CharDim)
f       = GeoUnit(1000NoUnits)
        Nondimensionalize!(f,CharDim)
d       = GeoUnit(10mm)
        Nondimensionalize!(d,CharDim)
c       = CreepLawVariables(P=P,T=T,f=f,d=d)
Phase   = SetMaterialParams(Name="Viscous Matrix", Phase=1,
                                     Density   = ConstantDensity(),
                                     CreepLaws = DiffusionCreep(n=1.0NoUnits, r=1.0NoUnits, p=-3.0NoUnits, A=1.0e6MPa^(-1.0)*s^(-1.0)*µm^(3.0), E=335000.0J/mol, V=4.0e-6m^(3.0)/mol, Apparatus="Invariant"), CharDim = CharDim)
EpsII.val = ComputeCreepLaw_EpsII(TauII,Phase.CreepLaws[1],c)
# Tested value is Dimensional so need to non-dimensionalize it back to compare
@test EpsII.val  ≈ Nondimensionalize(7.823165072337786e-15s^(-1.0), CharDim) rtol = 1e-12
    # Check that once inverted, we get back the TauII that we used to calculate EpsII
NewTau = ComputeCreepLaw_TauII(EpsII,Phase.CreepLaws[1],c)
#@show NewTau
@test NewTau ≈ TauII.val

# Given stress
#@test ComputeCreepLaw_EpsII(1e6Pa, x1, CreepLawParams())==5e-13/s                # dimensional input

# Given strainrate
#@test ComputeCreepLaw_TauII(1e-13/s, x1, CreepLawParams())==1e18*2*1e-13Pa       # dimensional input
# -------------------------------------------------------------------

end


### A Pluto.jl notebook ###
# v0.12.16

using Markdown
using InteractiveUtils

# This Pluto notebook uses @bind for interactivity. When running this notebook outside of Pluto, the following 'mock version' of @bind gives bound variables a default value (instead of an error).
macro bind(def, element)
    quote
        local el = $(esc(element))
        global $(esc(def)) = Core.applicable(Base.get, el) ? Base.get(el) : missing
        el
    end
end

# ╔═╡ 3f3dbc34-3593-11eb-2250-13fc496e403e
begin
	import Pkg; Pkg.add.(["LaTeXStrings", "PlutoUI", "Distributions", "ColorSchemes", "Measures"])
	using Plots, LaTeXStrings, PlutoUI, Measures
	import Distributions: Beta, pdf
	import ColorSchemes.Paired_12;
	colors = Paired_12[[1,2,7,8,3,4,5,6,9,10]]
	Plots.default()
	gr(legend = nothing, grid = false, color = colors[2], lw = 10, legendfontsize=20,
    		xtickfontsize=20, ytickfontsize=20, xguidefontsize=20, yguidefontsize=20,
		titlefontsize = 20)
	nothing
end

# ╔═╡ 2ea253b6-358a-11eb-10c9-6d257cbca066
md"### Bayesian analysis of the Bernoulli model"

# ╔═╡ 5d6396b0-358a-11eb-2995-eb061a2476aa
md" **Bernoulli model**:

$$x_1,\ldots,x_n |\theta \overset{\mathrm{iid}}{\sim} \mathrm{Bernoulli}(\theta)$$"

# ╔═╡ b7b9303e-358a-11eb-067d-fd0020d579a3
md"**Beta prior**:

$$\theta \sim \mathrm{Beta}(\alpha,\beta)$$"

# ╔═╡ da5f2ce0-358a-11eb-2522-81165870626e
md" **Beta posterior**:

$$\theta | x_1,\ldots,x_n \sim \mathrm{Beta}(\alpha+s,\beta+f),$$

where $s=\sum_{i=1}^n x_i$ is the number of successes and $f = n-s$ is the number of failures."

# ╔═╡ d39a7b72-358b-11eb-1d11-dfdcf238a7f4
md"""
**Input data and prior hyperparameters**:

n =  $(@bind n Slider(1:100; default=10, show_value=true))

s/n =  $(@bind sFrac Slider(0:0.01:1; default=0.5, show_value=true))

\$\alpha = \$ $(@bind α Slider(eps():100; default=3, show_value=true))

\$\beta = \$ $(@bind β Slider(eps():100; default=3, show_value=true))

"""

# ╔═╡ 698e86b0-3590-11eb-185b-df426ce40664
begin
	# Function to compute the Beta posterior and moments
	function BernBeta(s, f, α, β)
		priorMean = α/(α+β)
		priorStd = √(α*β/((α+β)^2*(α+β+1)))
		αPost = α + s
		βPost = β + f
		postMean = αPost/(αPost+βPost)
		postStd = √(αPost*βPost/((αPost+βPost)^2*(αPost+βPost+1)))
		post = Beta(αPost, βPost)
		return post, αPost, βPost, priorMean, priorStd, postMean, postStd
	end
	s = sFrac*n
	post, αPost, βPost, priorMean, priorStd, postMean, postStd = BernBeta(s, n-s, α, β)
end;

# ╔═╡ ca4114a4-3591-11eb-1777-4593b8760a68
begin
	f = n - s
	θGrid = range(eps(),1-eps(), length = 1000)
	binWidth = θGrid[2]-θGrid[1]
	logℓ = s*log.(θGrid) .+ f*log.(1 .- θGrid)
	like = exp.(logℓ .- maximum(logℓ))
	normLike = like/(sum(like)*binWidth)
	plot(size = (1200,800), θGrid, pdf.(Beta(α,β),θGrid), lw = 10, color = colors[6],
		xlabel = L"\theta", label = "Prior", yaxis = false, margin = 5mm)
	plot!(θGrid, normLike, lw = 10, color = colors[2], label = "Likelihood")
	plot!(θGrid, pdf.(Beta(αPost,βPost),θGrid), lw = 10, color = colors[4], label = "Posterior")
end

# ╔═╡ 7cdb3f2c-358d-11eb-03e6-df23873fb7b4
md"""
**Prior mean and standard deviation**:

\$\mathbb{E}(\theta)=\$ $(priorMean)

\$\mathbb{S}(\theta)=\$ $(priorStd)
"""

# ╔═╡ 25b33df2-3598-11eb-2364-a540b649f80d
md"""
**Posterior mean and standard deviation**:

\$\mathbb{E}(\theta|\mathbf{x})=\$  $(postMean)

\$\mathbb{S}(\theta|\mathbf{x})=\$ $(postStd)
"""

# ╔═╡ 3d25b63c-359c-11eb-2aa0-c558e2eb0a6b

# ╔═╡ 1f03b924-359c-11eb-0e7a-41c84b5dc313
md"""
---
Notebook for the book: Villani, M. (2020). [Bayesian Learning]()."""

# ╔═╡ Cell order:
# ╟─2ea253b6-358a-11eb-10c9-6d257cbca066
# ╟─3f3dbc34-3593-11eb-2250-13fc496e403e
# ╟─5d6396b0-358a-11eb-2995-eb061a2476aa
# ╟─b7b9303e-358a-11eb-067d-fd0020d579a3
# ╟─da5f2ce0-358a-11eb-2522-81165870626e
# ╠═d39a7b72-358b-11eb-1d11-dfdcf238a7f4
# ╠═698e86b0-3590-11eb-185b-df426ce40664
# ╟─ca4114a4-3591-11eb-1777-4593b8760a68
# ╟─7cdb3f2c-358d-11eb-03e6-df23873fb7b4
# ╟─25b33df2-3598-11eb-2364-a540b649f80d
# ╟─3d25b63c-359c-11eb-2aa0-c558e2eb0a6b
# ╠═1f03b924-359c-11eb-0e7a-41c84b5dc313


using GLMakie, LaTeXStrings, SpecialFunctions, Random
Random.seed!(13)
GLMakie.activate!()
set_theme!(theme_black())
fig = Figure(resolution = (1200,800))
ax = Axis3(fig; aspect = (1,0.5,0.5), azimuth = 10.42, elevation = 0.027,
    perspectiveness=0.5)
x = 0:0.1:15
y = -1:0.1:7
horizontal = Point3f.(tuple.(15, y, 3exp.(-(y .-3).^2/3)))
for ν in 0:7
    lines!(ax, x, x*0 .+ ν, besselj.(ν, x) .+ ν/3, linewidth = 2,
        color = :grey90, label = latexstring("J_{$(ν)}(x)"))
    band!(ax, Point3f.(tuple.(x, ν, ν/3)),
        Point3f.(tuple.(x, ν, besselj.(ν, x) .+ ν/3)), color = 1:length(x),
        colormap = :plasma)
    text!(ax, latexstring("J_{$(ν)}(x)"), position = Point3f(15.2,ν, ν/3))
end
band!(ax, Point3f.(tuple.(0, y, 0.0)),
    Point3f.(tuple.(0, y, 3exp.(-(y .-3).^2/3))), color = rand(length(y)))
band!(ax, horizontal[1:40], reverse(horizontal[42:end]); color = 1:40,
    colormap = :bone_1)
lines!(ax, horizontal; color= :white, linewidth = 2)
hidedecorations!(ax; grid = false)
fig[1,1] = ax
save("band3d_bessel.png", fig)
display(fig)


@testset "Function interpolation" begin
    rₘₐₓ = 20
    @testset "t₀ = $(t₀), ϕ = $(ϕ)" for (t₀,ϕ) in [(0.0,0.0)# ,
                                                   # (rₘₐₓ/2, π/3)
                                                   ]
        R = FEDVR(range(0,stop=rₘₐₓ,length=11), 10, t₀=t₀, ϕ=ϕ)
        R⁻¹ = pinv(R)
        r = axes(R,1)
        r̃ = range(0,stop=rₘₐₓ,length=1001)
        χ = R[r̃,:]

        fu = r -> r^2*exp(-r)
        fv = r -> r^6*exp(-8r)

        u = R*(R \ fu.(r))
        @test norm(χ * (R⁻¹*u) - fu.(r̃)) < 1e-6
        v = R*(R \ fv.(r))
        @test norm(χ * (R⁻¹*v) - fv.(r̃)) < 5e-5

        @testset "Restricted basis" begin
            R̃ = R[:,2:end-1]
            χ̃ = R̃[r̃,:]

            # # Actually, we want
            # ũ = R̃*(R̃ \ fu.(r))
            # # but R̃*... expands the coefficient vector by two elements
            # # (which are "hidden" by the restriction matrix), which
            # # causes a dimension mismatch below.
            ũ = R̃ \ fu.(r)
            # @test norm(χ̃ * ũ.args[2] - fu.(r̃)) < 5e-6
            @test norm(χ̃ * ũ - fu.(r̃)) < 2e-6
            # ṽ = R̃*(R̃ \ fv.(r))
            ṽ = R̃ \ fv.(r)
            # @test norm(χ̃ * ṽ.args[2] - fv.(r̃)) < 1e-4
            @test norm(χ̃ * ṽ - fv.(r̃)) < 5e-5

            h = r -> (r-rₘₐₓ/2)^2
            c = R̃ \ h.(r)
            # Vandermonde matrix to compare with; its generation is
            # costly, which is why we prefer to compute overlaps with
            # basis functions instead.
            rl = locs(R̃)
            V = R̃[rl,:]
            c̃ = V \ h.(rl)

            @test c ≈ c̃ atol=1e-13
        end
    end
    @testset "Restricted interpolation" begin
        L = 20.0
        n = 10
        t = range(-L/2, stop=L/2, length=n)
        B = FEDVR(t, [4,3,4,2,5,4,2,4,8])
        N = size(B,2)
        xx = axes(B,1)

        Bl = B[:,1:round(Int, 0.6N)+1]
        Bv = B[:,round(Int, 0.4N):end]

        f = sin
        c = B \ f.(xx)
        cl = Bl \ f.(xx)
        cv = Bv \ f.(xx)

        @test length(cl) == size(Bl,2)
        @test c[1:length(cl)-1] == cl[1:end-1]

        @test length(cv) == size(Bv,2)
        @test c[end-length(cv)+1:end] == cv
    end
end


using SVD_Projection
using Test, Random
Random.seed!(1234)
discovery = rand([0., 1., 2.], 30, 20);
target= rand([0., 1., 2.], 30, 15);

V_PC = project_PC(discovery, target, N_PC=3)
@test V_PC[1,:] == [-0.12869065723376472; -0.05776618956235076; -0.12278055691277556]


for N in [Float64, Float32]
    # random ellipsoid
    rand(Ellipsoid)

    # 1D ellipsoid
    E = Ellipsoid(N[0], Diagonal(N[1]))
    # Test Dimension
    @test dim(E) == 1
    # support function
    @test ρ(N[1], E) == ρ(N[-1], E) == N(1)
    # Test Support Vector
    d = N[1]
    @test σ(d, E) == N[1]
    d = N[-1]
    @test σ(d, E) == N[-1]
    # test constructor
    E = Ellipsoid(Diagonal(N[1]))
    @test E.center == N[0]

    # 2D Ellipsoid
    E = Ellipsoid(N[0, 0], Diagonal(N[1, 1]))
    # Test Dimension
    @test dim(E) == 2
    # support function
    @test ρ(N[1, 0], E) == ρ(N[-1, 0], E) == N(1)
    # Test Support Vector
    d = N[1, 0]
    @test σ(d, E) == N[1, 0]
    d = N[-1, 0]
    @test σ(d, E) == N[-1, 0]
    d = N[0, 1]
    @test σ(d, E) == N[0, 1]
    d = N[0, -1]
    @test σ(d, E) == N[0, -1]
    d = N[0, 0]
    @test σ(d, E) ∈ E

    # 2D Ellipsoid not 0-centered
    E = Ellipsoid(N[1, 2], Diagonal(N[1, 1]))
    # Test Dimension
    @test dim(E) == 2
    # Test Support Vector
    d = N[1, 0]
    @test σ(d, E) == N[2, 2]
    d = N[-1, 0]
    @test σ(d, E) == N[0, 2]
    d = N[0, 1]
    @test σ(d, E) == N[1, 3]
    d = N[0, -1]
    @test σ(d, E) == N[1, 1]

    # another shape matrix
    E = Ellipsoid(N[1, 2], Diagonal(N[0.5, 2]))
    # Test Support Vector
    d = N[1, 0]
    @test σ(d, E) ≈ N[1+sqrt(0.5), 2]
    d = N[-1, 0]
    @test σ(d, E) ≈ N[1-sqrt(0.5), 2]
    d = N[0, 1]
    @test σ(d, E) ≈ N[1, 2 + sqrt(2)]
    d = N[0, -1]
    @test σ(d, E) ≈ N[1, 2 - sqrt(2)]

    # boundedness
    @test isbounded(E)

    # isempty
    @test !isempty(E)

    # an_element and set membership functions
    M = Matrix{N}(2I, 2, 2)
    E = Ellipsoid(N[1, 2], M)
    @test an_element(E) ∈ E

    # translation
    @test translate(E, N[1, 2]) == Ellipsoid(N[2, 4], M)
end


using Test
using StaticArrays, IntervalSets
import ClimaCore.DataLayouts: IJFH
import ClimaCore:
    Fields, slab, Domains, Topologies, Meshes, Operators, Spaces, Geometry
using LinearAlgebra: norm

function spectral_space_2D(; n1 = 1, n2 = 1, Nij = 4)
    domain = Domains.RectangleDomain(
        Geometry.XPoint(-3.0)..Geometry.XPoint(5.0),
        Geometry.YPoint(-2.0)..Geometry.YPoint(8.0),
        x1periodic = false,
        x2periodic = false,
        x1boundary = (:east, :west),
        x2boundary = (:south, :north),
    )
    mesh = Meshes.EquispacedRectangleMesh(domain, n1, n2)
    grid_topology = Topologies.GridTopology(mesh)

    quad = Spaces.Quadratures.GLL{Nij}()
    space = Spaces.SpectralElementSpace2D(grid_topology, quad)
    return space
end

@testset "1×1 2D domain space" begin
    Nij = 4
    n1 = n2 = 1
    space = spectral_space_2D(n1 = n1, n2 = n2, Nij = Nij)

    field =
        Fields.Field(IJFH{ComplexF64, Nij}(ones(Nij, Nij, 2, n1 * n2)), space)

    @test sum(field) ≈ Complex(1.0, 1.0) * 8.0 * 10.0 rtol = 10eps()
    @test sum(x -> 3.0, field) ≈ 3 * 8.0 * 10.0 rtol = 10eps()
    @test norm(field) ≈ sqrt(2.0 * 8.0 * 10.0) rtol = 10eps()

    @test Operators.matrix_interpolate(field, 4) ≈
          [Complex(1.0, 1.0) for i in 1:(4 * n1), j in 1:(4 * n2)]

    field_sin = map(x -> sin((x.x) / 2), Fields.coordinate_field(space))
    M = Operators.matrix_interpolate(field_sin, 20)
    @test size(M) == (20, 20)  # 20 x 20 for a 1 element field

    real_field = field.re

    # test broadcasting
    res = field .+ 1
    @test parent(Fields.field_values(res)) == Float64[
        f == 1 ? 2 : 1 for i in 1:Nij, j in 1:Nij, f in 1:2, h in 1:(n1 * n2)
    ]

    res = field.re .+ 1
    @test parent(Fields.field_values(res)) ==
          Float64[2 for i in 1:Nij, j in 1:Nij, f in 1:1, h in 1:(n1 * n2)]
end

@testset "Broadcasting interception for tuple-valued fields" begin
    n1 = n2 = 1
    Nij = 4
    space = spectral_space_2D(n1 = n1, n2 = n2, Nij = Nij)

    nt_field = Fields.Field(
        IJFH{NamedTuple{(:a, :b), Tuple{Float64, Float64}}, Nij}(
            ones(Nij, Nij, 2, n1 * n2),
        ),
        space,
    )
    nt_sum = sum(nt_field)
    @test nt_sum isa NamedTuple{(:a, :b), Tuple{Float64, Float64}}
    @test nt_sum.a ≈ 8.0 * 10.0 rtol = 10eps()
    @test nt_sum.b ≈ 8.0 * 10.0 rtol = 10eps()
    @test norm(nt_field) ≈ sqrt(2.0 * 8.0 * 10.0) rtol = 10eps()
end

@testset "Special case handling for broadcased norm to pass through space local geometry" begin
    space = spectral_space_2D()
    u = Geometry.Covariant12Vector.(ones(space), ones(space))
    @test norm.(u) ≈ hypot(4 / 8 / 2, 4 / 10 / 2) .* ones(space)
end

@testset "FieldVector" begin
    space = spectral_space_2D()
    u = Geometry.Covariant12Vector.(ones(space), ones(space))
    x = Fields.coordinate_field(space)
    Y = Fields.FieldVector(u = u, x = x)

    @test propertynames(Y) == (:u, :x)
    @test Y.u == u

    Y1 = 2 .* Y
    @test parent(Y1.u) == 2 .* parent(u)
    @test parent(Y1.x) == 2 .* parent(x)

    Y1 .= Y1 .+ 2 .* Y
    @test parent(Y1.u) == 4 .* parent(u)
    @test parent(Y1.x) == 4 .* parent(x)
end


@testset "PolygonalAmoeba" begin
    @polyvar x y
    f = x^2*y + y^2 + 3x^2*y^3 + y^4 + x^4*y^4

    A = amoeba(f, accuracy=0.01)
    @test 0 < accuracy(A) < 0.01
end


function checkit(loads::Vector{Float64}, oldloads::Vector{Float64}, tol::Float64)
  neq = length(loads)
  converged = true
  big = 0.0
  for i in 1:neq
    loads[i] > big && (big = abs(loads[i]))
  end
  for i in 1:neq
    abs(loads[i] - oldloads[i])/big > tol && (converged = false)
  end
  converged
end


#=
max_percent_diff.jl
Copyright 2000-9-16, Jeff Fessler, University of Michigan
=#

export max_percent_diff

using Test: @test

"""
    d = max_percent_diff(s1, s2, [options])

compute the "maximum percent difference" between two signals: `s1, s2`
default is to normalize by `maximum(abs.(s1))`

options
- `maxboth::Bool` use max of both arguments to normalize; default `false`
- `normalize::Bool` normalize each before comparing; default `false`

"""
function max_percent_diff(s1, s2 ;
        maxboth::Bool = false, normalize::Bool = false)

    #check to make sure that we have comparable signals
    size(s1) != size(s2) &&
        throw(DomainError("Dimension mismatch $(size(s1)) $(size(s2))"))

    any(isnan.(s1)) && @warn("NaN values in s1 input to max_percent_diff")
    any(isnan.(s2)) && @warn("NaN values in s2 input to max_percent_diff")

    if normalize
        normFactor = sum(abs.(s2)) / sum(abs.(s1))
        print("Normalization factor (matrix 2 / matrix 1): $normFactor \n")
        s1 /= sum(abs.(s1))
        s2 /= sum(abs.(s2))
    end

    if maxboth
        denom = max(maximum(abs.(s1)), maximum(abs.(s2)))
        d = denom == 0 ? 0 : maximum(abs.(s1 - s2)) / denom
    else
        denom = maximum(abs.(s1))
        denom = denom == 0 ? denom : maximum(abs.(s2))
        d = denom == 0 ? 0 : maximum(abs.(s1 - s2)) / denom
    end

    return 100 * d
end

function max_percent_diff_test()
    @test max_percent_diff([200], [202]) ≈ 100/101
    @test max_percent_diff([0 200], [0 202]) ≈ 100/101
    @test max_percent_diff([0 200], [0 198], maxboth=true) ≈ 1
    #≈ does not work with 0 values (it's proportional to the second arg)
    @test max_percent_diff([0 100], [0 200], normalize=true) < .01
    true
end

"""
    max_percent_diff(:test)
self test
"""
function max_percent_diff(test::Symbol)
    test != :test && throw(DomainError("Invalid signal"));
    return max_percent_diff_test()
end

# max_percent_diff(:test)


__precompile__(true)

module VectorField

import Plots: quiver, quiver!
import LinearAlgebra: adjoint
export meshgrid, vectorfield, vectorfield!

function meshgrid(x, y)
    x, y = float.(x), float.(y)
    [repeat(x, inner=length(y))'; repeat(y, outer=length(x))']
end

function vectorfield(points, field::Function; arrowscale=0.1)
    vectors = similar(points)
    d = size(points, 2)
    for i in 1:d
        vectors[:, i] .= collect(field(points[:, i]...))
    end
    vectors .*= arrowscale
    quiver(points[1, :], points[2, :], quiver=(vectors[1, :], vectors[2, :]))
end
# field is a Function that takes d inputs. d = dimension = 2 (3d not supported by Plots - use Makie.jl)
# points is an Array of size 2xd, each column represents a point in space
# vectorfield2d(field::Function, points <: AbstractArray{<:Number,2}, arrowscale::Real=0.1)

function vectorfield!(points, field::Function; arrowscale=0.1)
    vectors = similar(points)
    d = size(points, 2)
    for i in 1:d
        vectors[:, i] .= collect(field(points[:, i]...))
    end
    vectors .*= arrowscale
    quiver!(points[1, :], points[2, :], quiver=(vectors[1, :], vectors[2, :]))
end
# identical, but plots over existing items.
# to keep in line with Plots formatting

function vectorfield(xs, ys, field::Function; arrowscale=0.1)
    points = meshgrid(xs, ys)
    vectors = similar(points)
    d = size(points, 2)
    for i in 1:d
        vectors[:, i] .= collect(field(points[:, i]...))
    end
    vectors .*= arrowscale
    quiver(points[1, :], points[2, :], quiver=(vectors[1, :], vectors[2, :]))
end
# input x and y ranges, grid is interpolated

function vectorfield!(xs, ys, field::Function; arrowscale=0.1)
    points = meshgrid(xs, ys)
    vectors = similar(points)
    d = size(points, 2)
    for i in 1:d
        vectors[:, i] .= collect(field(points[:, i]...))
    end
    vectors .*= arrowscale
    quiver!(points[1, :], points[2, :], quiver=(vectors[1, :], vectors[2, :]))
end

function vectorfield(xs, ys, u::T, v::T; arrowscale=0.1) where {T <: AbstractArray{<:Number,N} where N}
    points = meshgrid(xs, ys)
    d = size(points, 2)
    d == length(u) == length(v) || error("Dimensions do not match. Ensure length(u)==length(v)==length(xs)*length(ys)")
    vectors = [reshape(u', 1, :); reshape(v', 1, :)]
    vectors .*= arrowscale
    quiver(points[1, :], points[2, :], quiver=(vectors[1, :], vectors[2, :]))
end
# u, v are 2 matrices, the value of u[i, j] is the x component of the arrow,
# the value of v[i, j] is the y component of the arrow
# Grid will be interpolated as:
#     ys - >           ys - >
# xs                xs
# |    u            |    v
# v           ,     v

function vectorfield!(xs, ys, u::T, v::T; arrowscale=0.1) where {T <: AbstractArray{<:Number,N} where N}
    points = meshgrid(xs, ys)
    d = size(points, 2)
    d == length(u) == length(v) || error("Dimensions do not match. Ensure length(u)==length(v)==length(xs)*length(ys)")
    vectors = [reshape(u', 1, :); reshape(v', 1, :)]
    vectors .*= arrowscale
    quiver!(points[1, :], points[2, :], quiver=(vectors[1, :], vectors[2, :]))
end

function vectorfield(xs, ys, uv::T; arrowscale=0.1) where {T <: AbstractArray{S,N} where S where N}
    points = meshgrid(xs, ys)
    d = size(points, 2)
    u = [uv[i][1] for i in CartesianIndices(uv)]
    v = [uv[i][2] for i in CartesianIndices(uv)]
    d == length(u) == length(v) || error("Dimensions do not match. Ensure length(u)==length(v)==length(xs)*length(ys)")
    vectors = [reshape(u', 1, :); reshape(v', 1, :)]
    vectors .*= arrowscale
    quiver(points[1, :], points[2, :], quiver=(vectors[1, :], vectors[2, :]))
end
# uv is a matrix of 2-element arrays [x, y]. Arrows will be interpreted componentwise.
# Grid will be interpolated as:
#     ys - >
# xs
# |    uv
# v

function vectorfield!(xs, ys, uv::T; arrowscale=0.1) where {T <: AbstractArray{S,N} where S where N}
    points = meshgrid(xs, ys)
    d = size(points, 2)
    u = [uv[i][1] for i in CartesianIndices(uv)]
    v = [uv[i][2] for i in CartesianIndices(uv)]
    d == length(u) == length(v) || error("Dimensions do not match. Ensure length(u)==length(v)==length(xs)*length(ys)")
    vectors = [reshape(u', 1, :); reshape(v', 1, :)]
    vectors .*= arrowscale
    quiver!(points[1, :], points[2, :], quiver=(vectors[1, :], vectors[2, :]))
end

end