export AbstractApproximationMethod
export approximation_name, approximation_short_name

abstract type AbstractApproximationMethod end

function approximation_name end
function approximation_short_name end

function approximate_meancov(method::AbstractApproximationMethod, g::Function, distribution)
    return approximate_meancov(method, g, mean(distribution), cov(distribution))
end

function approximate_meancov(method::AbstractApproximationMethod, g::Function, m::T, v::T) where { T <: Real }
    weights = getweights(method, m, v)
    points  = getpoints(method, m, v)

    cs   = Vector{eltype(m)}(undef, length(weights))
    norm = 0.0
    mean = 0.0

    for (index, (weight, point)) in enumerate(zip(weights, points))
        gv = g(point)
        cv = weight * gv

        mean += point * cv
        norm += cv

        @inbounds cs[index] = cv
    end

    mean /= norm

    var = 0.0
    for (index, (point, c)) in enumerate(zip(points, cs))
        point -= mean
        var += c * point ^ 2
    end

    var /= norm

    return mean, var
end

function approximate_meancov(method::AbstractApproximationMethod, g::Function, m::AbstractVector{T}, P::AbstractMatrix{T}) where { T <: Real }
    ndims = length(m)

    weights = getweights(method, m, P)
    points  = getpoints(method, m, P)

    cs = similar(m, eltype(m), length(weights))
    norm = zero(T)
    mean = zeros(T, ndims)

    for (index, (weight, point)) in enumerate(zip(weights, points))
        gv = g(point)
        cv = weight * gv

        # mean = mean + point * weight * g(point)
        broadcast!(*, point, point, cv)  # point *= cv
        broadcast!(+, mean, mean, point) # mean += point
        norm += cv

        @inbounds cs[index] = cv
    end

    broadcast!(/, mean, mean, norm)

    cov = zeros(T, ndims, ndims)
    foreach(enumerate(zip(points, cs))) do (index, (point, c))
        broadcast!(-, point, point, mean)                # point -= mean
        mul!(cov, point, reshape(point, (1, ndims)), c, 1.0) # cov = cov + c * (point)⋅(point)' where c = weight * g(point)
    end

    broadcast!(/, cov, cov, norm)

    return mean, cov
end

function approximate_kernel_expectation(method::AbstractApproximationMethod, g::Function, distribution)
    return approximate_kernel_expectation(method, g, mean(distribution), cov(distribution))
end

function approximate_kernel_expectation(method::AbstractApproximationMethod, g::Function, m::AbstractVector{T}, P::AbstractMatrix{T}) where { T <: Real }
    ndims = length(m)

    weights = getweights(method, m, P)
    points  = getpoints(method, m, P)

    gbar = zeros(ndims, ndims)
    foreach(zip(weights, points)) do (weight, point)
        axpy!(weight, g(point), gbar) # gbar = gbar + weight * g(point)
    end

    return gbar
end


# n = mult(3, 4) #=> ErrorException("mult not defined"))
function mult(x, y)
    println("x is $x and y is $y")
    x * y
end

n = mult(3, 4) #=> 12
# one-line assignment syntax:
mult(x, y) = x * y
f(x,y) = x^3-y+x*y
println(f(3,2) ) # 31

function mult(x, y)
    println("x is $x and y is $y")
    if x == 1 return y end
    x * y
end

function mult(x::Float64, y::Float64)
    println("x is $x and y is $y")
    x * y
end
# following gives error:
# mult(5, 6)
# ERROR: `mult` has no method matching mult(::Int64, ::Int64)

function multi(n,m)
    n*m, div(n,m), n%m
end
x, y, z = multi(8, 2)

function varargs(n, m, args...)
    println("arguments: $n $m $args")
end
varargs(1, 2, 3, 4) # n is 1, m is 2, and args has value (3, 4)

x = (3, 4)
varargs2(1, 2, x...) # args is (1,2,3,4)
x = [10, 11, 12]
varargs2(1, 2, x...) # args is (1,2,10,11,12)
varargs2(k1="name1", k2="name2, k3=7)

apply(mult, 3, 4) #=> 12

∑(x,y) = x + y
∑(3, 4) #=> 7

# anonymous functions:
function (x)
    x + 2
end

g = (x, y) -> x^3-y+x*y
println(g(3,2) ) # 31

x -> x + 2
() -> x * 5

3+4 #=> 7
+(3,4) #=> 7

m = mult
m(6, 6) #=> 36


# Frequent Directions
# from: Liberty, "Simple and Deterministic Matrix Sketching"
#
# Danny Perry (dperry@cs.utah.edu)
# Oct 2014

using Debug

function reduce(S,V,p)
	delta = S[end-p+1]^2
	SS = (S .^2) - delta
	SS[find(SS .< 0)] = 0
	diagm( sqrt(SS) ) * V'
end

# Frequent Directions
# @param A - original matrix
# @param l - size of sketch
# @param p - the number of components to zero out at a time
#
# @return B - the sketch
function FD(A, l, p)
  n = size(A,1)
  d = size(A,2)
  if l < p
    error("l must be >= p")
  end
  if p > d
    error ("p must be <= d")
  end
	B = A[1:l,:]
	for i=l+1:p:n
    numz = p
    if n-i < p
      numz = n-i+1
    end
		U,S,V = svd(B)
		corankB = length(find(S .<= eps(typeof(S[1]))))
		if corankB < numz # d-rank(B) < numz
    	B = reduce(S,V,numz)
		else
			B = diagm(S)*V'
		end
    B[end-numz+1:end,:] = A[i:i+numz-1,:]
	end
	B
end

# Frequent Directions - bit by bit
# @param A - next p or less rows
# @param B - the current B matrix
# @param l - size of sketch
# @param p - the number of components to zero out at a time
#
# @return B - the sketch
function FD(A, B, l, p)
	numz = size(A,1)
  d = size(A,2)
  if l < p
    error("l must be >= p")
  end
  if p > d
    error ("p must be <= d")
  end
	U,S,V = svd(B)
	if length(find(S .<= eps(typeof(S[1])))) < numz # d-rank(B) < numz
		B = reduce(S,V,numz)
	else
		B = diagm(S)*V'
	end
	B[end-numz+1:end,:] = A
	B
end


# unit test he SVRG solver
include("../src/DMD_Util.jl")
include("../src/DMD_Type.jl")
include("../src/DMD_Func.jl")
include("../src/DMD_SVRG.jl")
include("../src/DMD_BFGS.jl")
include("../src/DMD_Loss.jl")

# Generate DMD Synthetic Data
#------------------------------------------------------------------------------
# dimensions
m = 500;    # temporal dimension
n = 100;    # spatial dimension
k = 3;      # number of modes
T = Float64;
# generate data
X, t, at, Bt = simDMD(m, n, k, T; seed=123456);

# additive noise
sigma = 1.0e-4 # background noise level
lds = 0.1 # larger deviation size
ldf = 0.05 # frequency of larger deviations
X1 = X .+ randn(real(eltype(X)),size(X))*sigma .+ (rand(real(eltype(X)),size(X)) .< ldf)*lds
X2 = copy(X1)
# specify the loss
kappa = lds/5.0
lossFunc1(r) = huber_func(r,kappa)
lossGrad1(r) = huber_grad!(r,kappa)
lossFunc2(r) = l2_func(r)
lossGrad2(r) = l2_grad!(r)

#
# Create object
#------------------------------------------------------------------------------
params1 = DMDParams(k, X1, t, lossFunc1, lossGrad1);
params2 = DMDParams(k, X2, t, lossFunc2, lossGrad2);
params2B = DMDParams(k, X2, t, lossFunc2, lossGrad2);
Random.randn!(params1.ar);
copyto!(params2.ar,params1.ar);

#
# Apply Solver
#------------------------------------------------------------------------------
tau = 10;
eta = 1e1;
itm = 2000;
tol = 1e-6;
ptf = 1000;

function prox_stab(ar)
    # project for numerical stability
    k = length(ar) >> 1;
    @inbounds for i = 1:k
        ar[i<<1-1] = min(0.0, ar[i<<1-1]);
    end
end

opts1 = DMD_SVRG_Options(tau, eta, params1, itm=itm, tol=tol, ptf=ptf,
                         prox=prox_stab);
opts2 = DMD_SVRG_Options(tau, eta, params2, itm=itm, tol=tol, ptf=ptf,
                         prox=prox_stab);

sigma   = 1e-3;
itm = 1000;
tol = 1e-6;
ptf = 10;
opts2B = DMD_BFGS_Options(sigma, params2B, itm=itm, tol=tol, ptf=ptf);

@time obj_his1, err_his1 = solveDMD_withSVRG(params1, opts1);
@time obj_his2, err_his2 = solveDMD_withSVRG(params2, opts2);

@time obj_his2B, err_his2B = solveDMD_withBFGS(params2B, opts2B);

@printf "%7.4e abs err evals w/ Huber\n" besterrperm(params1.a,at)
@printf "%7.4e abs err evals w/ L2\n" besterrperm(params2.a,at)
@printf "%7.4e abs err evals w/ L2 (BFGS)\n" besterrperm(params2.a,at)


@testset "Karger Minimum Cut" begin

    gx = path_graph(5)

    #Assumes cut[1] = 1
    for g in testgraphs(gx)
        cut = @inferred(karger_min_cut(g))
        @test findfirst(isequal(2), cut) ==  findlast(isequal(1), cut)+1
        @test karger_cut_cost(g, cut) == 1
        @test karger_cut_edges(g, cut) == [Edge(findlast(isequal(1), cut), findfirst(isequal(2), cut))]
    end

    add_vertex!(gx)

    for g in testgraphs(gx)
        cut = @inferred(karger_min_cut(g))
        @test cut == [1, 1, 1, 1, 1, 2]
        @test karger_cut_cost(g, cut) == 0
        @test karger_cut_edges(g, cut) == Vector{Edge}()

    end

    gx = star_graph(5)
    for g in testgraphs(gx)
        cut = @inferred(karger_min_cut(g))
        @test count(isequal(2), cut) == 1
        @test karger_cut_cost(g, cut) == 1
        @test karger_cut_edges(g, cut) == [Edge(1, findfirst(isequal(2), cut))]
    end

    gx = SimpleGraph(1)
    for g in testgraphs(gx)
        cut = @inferred(karger_min_cut(g))
        @test cut == zeros(Int, 1)
    end

end


struct SimpleOrbit{T,S,V,R} <: AbstractOrbit
    period::T
    t0::T
    b::S
    duration::T
    r_star::V

    b_norm::V
    speed::R
    half_period::T
    ref_time::T
end

"""
    SimpleOrbit(; period, duration, t0=0, b=0, r_star=1)

Circular orbit parameterized by the basic observables of a transiting system.

# Parameters
* `period` - The orbital period of the planets, nominally in days
* `duration` The duration of the transit, same units as `period`
* `t0` - The midpoint time of the reference transit, same units as `period`
* `b` - The impact parameter of the orbit
* `r_star` - The radius of the star, nominally in solar radii
"""
function SimpleOrbit(;period, duration, t0=zero(period), b=0.0, r_star=1.0)
    half_period = 0.5 * period
    duration > half_period && error("duration cannot be longer than half the period")
    b_norm = b * r_star
    speed = 2 * sqrt(r_star^2 - b_norm^2) / duration
    ref_time =  t0 - half_period
    # normalize time types
    period, t0, duration, half_period, ref_time = promote(period, t0, duration, half_period, ref_time)
    # normalize unitless types
    r_star, b_norm = promote(r_star, b_norm)
    SimpleOrbit(period, t0, b, duration, r_star, b_norm, speed, half_period, ref_time)
end

period(orbit::SimpleOrbit) = orbit.period
duration(orbit::SimpleOrbit) = orbit.duration

planet_position(orbit::SimpleOrbit, t) = relative_position(orbit, t)

function relative_position(orbit::SimpleOrbit, t)
    dt = (t - orbit.ref_time) % period(orbit) - orbit.half_period
    x = orbit.speed * dt
    y = orbit.b_norm
    z = abs(dt) < 0.5 * orbit.duration ? one(x) : -one(x)
    return SA[x, y, z]
end

function in_transit(orbit::SimpleOrbit, t, r; texp=0)
    dt = (t - orbit.ref_time) % period(orbit) - orbit.half_period
    tol = sqrt((r + orbit.r_star)^2 - orbit.b_norm^2) / orbit.speed
    tol += 0.5 * texp
    return abs(dt) < tol
end

function in_transit(orbit::SimpleOrbit, t; texp=0)
    dt = (t - orbit.ref_time) % period(orbit) - orbit.half_period
    tol = 0.5 * (orbit.duration + texp)
    return abs(dt) < tol
end

function flip(orbit::SimpleOrbit, r_planet)
    t0 = orbit.t0 + orbit.half_period
    b = orbit.b_norm / r_planet
    duration = 2 * sqrt(r_planet^2 + orbit.b_norm^2) / orbit.speed
    return SimpleOrbit(orbit.period,
                       t0,
                       b,
                       duration,
                       r_planet,
                       orbit.b_norm,
                       orbit.speed,
                       orbit.half_period,
                       orbit.t0)
end


include("SortSparse.jl")
include("KoLesky.jl")
include("Cholesky.jl")
include("CovFuncs.jl")
using LinearAlgebra
using RecursiveArrayTools
using Distributions
using Distances
using Random
using Profile
using ProfileView
using JLD
using MAT
using IterativeSolvers
using BlockArrays
Random.seed!(123)

# Creating uniform random data
NTrain = 1000000
d =  2
NTest = 50
N = NTest + NTrain
x = rand(d, N)
setSplit = randperm(N)
trainInds = setSplit[1:NTrain]
testInds = setSplit[(NTrain + 1) : end]
xTrain = x[ :, trainInds ]
xTest = x[:, testInds ]
x = hcat(xTrain, xTest)
y = randn(N)
y = PseudoBlockArray(y, [NTest, NTrain])

ρ = 3.0
λ = 1.5

# reordering points and forming skeletons
# @profview colptr, rowval, P, revP, distances = sortSparseRev(xTrain, xTest, ρ, 1)
@time colptr, rowval, P, revP, distances = sortSparseRev(xTrain, xTest, ρ, 1)
x = x[:, P]

LTest = SparseMatrixCSC(N, N, colptr, rowval, ones(length(rowval)))
LTest22 = LTest[(NTest + 1) : end, (NTest + 1) : end]

xTest = x[:, 1:NTest]
xTrain = x[:, (NTest + 1) : end]
skeletons = construct_skeletons(colptr, rowval, distances, λ)

l = 0.1
# σ = vcat(zeros(NTest), fill(10.0, NTrain))
σ = vcat(zeros(NTest), 1.0 * exp.(randn(NTrain)) .+ 1e-4)
σ = PseudoBlockArray(σ, [NTest, NTrain])

function covfunc!(r)
  matern32!(r,l)
end

LinearAlgebra.BLAS.set_num_threads(1)
# @profview L = assembleL(skeletons, x, covfunc!)
@time L = assembleL(skeletons, x, covfunc!)
# @time LNaive = assembleL(skeletons, x, covfunc!, σ)
# @time L22 = L[(NTest + 1) : end, (NTest + 1) : end]
# @time A22 = L22 * L22' .+ spdiagm(0 => 1 ./ max.(σ[Block(2)], eps(eltype(σ))))
# # @time U22 = sparse(triu(A22))
# @time U22 = squareSparse(L22)+ spdiagm(0 => 1 ./ max.(σ[Block(2)], 1e-30))
# @time icholU_high_level!(U22)
#
#
# # The vecchia posterior mean
# z = (L22 * (L22' * y[Block(2)]))
# z = cg(A22 , z, Pl = IChol(U22',U22), verbose=true, maxiter=5, tol=1e-16)
# z = (Diagonal(σ[Block(2)]))\ z
# μVechia = (L' \ (L \ (vcat(zero(y[Block(1)]), z))))[1:NTest]
# # display(μVechia)
# @show mean(abs.(μVechia))
# @show median(abs.(μVechia))
#
#
# # The naive application of vecchia
# μNaive = -LNaive[1:NTest, 1:NTest]' \ (LNaive[(NTest + 1) : end, 1 : NTest]' * y[Block(2)])
# # display(μNaive)
# @show mean(abs.(μNaive))
# @show median(abs.(μNaive))
#


@testset "Fitting" begin
  @testset "Block" begin
    gev_seed = StableRNG(12345)
    dist1    = GeneralizedExtremeValue(0.0, 1.0, 0.0)
    dist2    = GeneralizedExtremeValue(0.0, 1.0, 0.5)
    dist3    = GeneralizedExtremeValue(1.0, 1.0, 0.5) # The problem distribution...
    r_gev_1  = readdlm(joinpath(datadir,"R_gev_dist1_fits.csv"), ',')
    r_gev_2  = readdlm(joinpath(datadir,"R_gev_dist2_fits.csv"), ',')
    r_gev_3  = readdlm(joinpath(datadir,"R_gev_dist3_fits.csv"), ',')
    r_result = zip(eachrow(r_gev_1), eachrow(r_gev_2), eachrow(r_gev_3))
    for (test_ix, r_gev_ix) in enumerate(r_result)
      for (r_mle, dist) in zip(r_gev_ix, (dist1, dist2, dist3))
        sample = rand(gev_seed, dist, 1000)
        try
          julia_fit = fit(GeneralizedExtremeValue, BlockMaxima(sample, 1))
          j_mle = [julia_fit.μ, julia_fit.σ, julia_fit.ξ]
          @test isapprox(j_mle, r_mle, atol=1e-3)
        catch er
          @info "Optimization failed for $dist on trial $test_ix with error $er."
        end
      end
    end
  end

  @testset "PoT" begin
    gpd_seed = StableRNG(12345)
    dist1    = GeneralizedPareto(0.0, 1.0, 0.0)
    dist2    = GeneralizedPareto(0.0, 1.0, 0.5)
    dist3    = GeneralizedPareto(1.0, 1.0, 0.5)
    r_gpd_1  = readdlm(joinpath(datadir,"R_gpd_dist1_fits.csv"), ',')
    r_gpd_2  = readdlm(joinpath(datadir,"R_gpd_dist2_fits.csv"), ',')
    r_gpd_3  = readdlm(joinpath(datadir,"R_gpd_dist3_fits.csv"), ',')
    r_result = zip(eachrow(r_gpd_1), eachrow(r_gpd_2), eachrow(r_gpd_3))
    for (test_ix, r_gpd_ix) in enumerate(r_result)
      for (r_mle, dist) in zip(r_gpd_ix, (dist1, dist2, dist3))
        sample = rand(gpd_seed, dist, 1000)
        try
          julia_fit = fit(GeneralizedPareto, PeakOverThreshold(sample, dist.μ))
          j_mle     = [julia_fit.σ, julia_fit.ξ]
          @test isapprox(j_mle, r_mle, atol=1e-3)
        catch er
          @info "Optimization failed for $dist on trial $test_ix with error $er."
        end
      end
    end
  end
end

#= Code to generate the R test files:
using RCall, StableRNGs, Distributions
R"library(ismev)"
gev_R_results = NTuple{3, Vector{Float64}}[]
gev_seed = StableRNG(12345)
dist1    = GeneralizedExtremeValue(0.0, 1.0, 0.0)
dist2    = GeneralizedExtremeValue(0.0, 1.0, 0.5)
dist3    = GeneralizedExtremeValue(1.0, 1.0, 0.5)
for test_ix in 1:10
  mles_trialj = map((dist1, dist2, dist3)) do dist
    sample = rand(gev_seed, dist, 1000)
    R"""
      r_fit = gev.fit($sample, show=FALSE)
      r_mle = r_fit$mle
    """
    @rget r_mle
  end
  push!(gev_R_results, mles_trialj)
end

gpd_R_results = NTuple{3, Vector{Float64}}[]
gpd_seed = StableRNG(12345)
dist1    = GeneralizedPareto(0.0, 1.0, 0.0)
dist2    = GeneralizedPareto(0.0, 1.0, 0.5)
dist3    = GeneralizedPareto(1.0, 1.0, 0.5)
for test_ix in 1:10
  mles_trialj = map((dist1, dist2, dist3)) do dist
    sample = rand(gpd_seed, dist, 1000)
    R"""
      r_fit = gpd.fit($sample, threshold=$(dist.μ), show=FALSE)
      r_mle = r_fit$mle
    """
    @rget r_mle
  end
  push!(gpd_R_results, mles_trialj)
end

using DelimitedFiles
writedlm("data/R_gev_dist1_fits.csv", getindex.(gev_R_results, 1), ',')
writedlm("data/R_gev_dist2_fits.csv", getindex.(gev_R_results, 2), ',')
writedlm("data/R_gev_dist3_fits.csv", getindex.(gev_R_results, 3), ',')
writedlm("data/R_gpd_dist1_fits.csv", getindex.(gpd_R_results, 1), ',')
writedlm("data/R_gpd_dist2_fits.csv", getindex.(gpd_R_results, 2), ',')
writedlm("data/R_gpd_dist3_fits.csv", getindex.(gpd_R_results, 3), ',')
=#


# dependencies
using LFAToolkit
using LinearAlgebra
using Plots
# setup
p = 2
dimension = 1
mesh = Mesh1D(1.0)

# operator
diffusion = GalleryOperator("diffusion", p+1, p+1, mesh)

# NoPreco smoother
noPreco = NoPreco(diffusion)
# full operator symbols
numbersteps = 250
maxeigenvalue = 0
θ_min = -π/2
θ_max = 3π/2
θ_step = 2π/(numbersteps-1)
θ_range = θ_min:θ_step:θ_max

# compute and plot smoothing factor
# setup
ω = [1.00]
eigenvalues = zeros(numbersteps, p)

# compute
for i in 1:numbersteps
    θ = [θ_range[i]]
    if abs(θ[1]) >  π/512
        A = computesymbols(noPreco, ω, θ)
        currenteigenvalues = [real(val) for val in eigvals(I-A)]
        eigenvalues[i, :] = currenteigenvalues
    end
end

# plot
xrange = θ_range/π
plot(
    xrange,
    xlabel="θ/π",
    xtickfont=font(12, "Courier"),
    eigenvalues,
    ytickfont=font(12, "Courier"),
    ylabel="λ",
    linewidth=3,
    legend=:none,
    title="Spectrum of NoPreco Symbol",
    palette=palette(:tab10)
)
ylims!(min(0.0, eigenvalues...) * 1.1, max(eigenvalues...) * 1.1)
savefig("noPreco_spectrum_2")

## p = 4

# setup
p = 4
dimension = 1
mesh = Mesh1D(1.0)

# operator
diffusion = GalleryOperator("diffusion", p+1, p+1, mesh)

# NoPreco smoother
noPreco = NoPreco(diffusion)
# full operator symbols
numbersteps = 250
maxeigenvalue = 0
θ_min = -π/2
θ_max = 3π/2
θ_step = 2π/(numbersteps-1)
θ_range = θ_min:θ_step:θ_max

# compute and plot smoothing factor
# setup
ω = [1.00]
eigenvalues = zeros(numbersteps, p)

# compute
for i in 1:numbersteps
    θ = [θ_range[i]]
    if abs(θ[1]) >  π/512
        A = computesymbols(noPreco, ω, θ)
        currenteigenvalues = [real(val) for val in eigvals(I-A)]
        eigenvalues[i, :] = currenteigenvalues
    end
end

# plot
xrange = θ_range/π
plot(
    xrange,
    xlabel="θ/π",
    xtickfont=font(12, "Courier"),
    eigenvalues,
    ytickfont=font(12, "Courier"),
    ylabel="λ",
    linewidth=3,
    legend=:none,
    title="Spectrum of NoPreco Symbol",
    palette=palette(:tab10)
)
ylims!(min(0.0, eigenvalues...) * 1.1, max(eigenvalues...) * 1.1)
savefig("noPreco_spectrum_4")

## p = 7

# setup
p = 7
dimension = 1
mesh = Mesh1D(1.0)

# operator
diffusion = GalleryOperator("diffusion", p+1, p+1, mesh)

# NoPreco smoother
noPreco = NoPreco(diffusion)
# full operator symbols
numbersteps = 250
maxeigenvalue = 0
θ_min = -π/2
θ_max = 3π/2
θ_step = 2π/(numbersteps-1)
θ_range = θ_min:θ_step:θ_max

# compute and plot smoothing factor
# setup
ω = [1.00]
eigenvalues = zeros(numbersteps, p)

# compute
for i in 1:numbersteps
    θ = [θ_range[i]]
    if abs(θ[1]) >  π/512
        A = computesymbols(noPreco, ω, θ)
        currenteigenvalues = [real(val) for val in eigvals(I-A)]
        eigenvalues[i, :] = currenteigenvalues
    end
end

# plot
xrange = θ_range/π
plot(
    xrange,
    xlabel="θ/π",
    xtickfont=font(12, "Courier"),
    eigenvalues,
    ytickfont=font(12, "Courier"),
    ylabel="λ",
    linewidth=3,
    legend=:none,
    title="Spectrum of NoPreco Symbol",
    palette=palette(:tab10)
)
ylims!(min(0.0, eigenvalues...) * 1.1, max(eigenvalues...) * 1.1)
savefig("noPreco_spectrum_7")


using CMPSKit
using Test
using LinearAlgebra
using Random
using OptimKit: GradientDescent, ConjugateGradient, LBFGS
using KrylovKit: GMRES, Arnoldi
Random.seed!(111134)

include("functionseries.jl")
include("uniformcmps.jl")
include("periodiccmps.jl")


DIR = joinpath("tmp","BSplineManifold")
rm(DIR, force=true, recursive=true)
mkpath(DIR)

@testset "BSplineManifold" begin
    Random.seed!(42)

    @testset "1d2d" begin
        p = 3 # degree of polynomial
        k = KnotVector(1:12) # knot vector
        P = BSplineSpace{p}(k) # B-spline space
        rand_a = [randn(2) for i in 1:dim(P)] / 2
        a = [SVector(2i-6, 0) for i in 1:dim(P)] + rand_a # random generated control points
        M = BSplineManifold(a, (P,)) # Define B-spline manifold

        @testset "luxor-svg" begin
            path = joinpath(DIR, "1d2d.svg")
            save_svg(path,M)
            @test isfile(path)
        end
        @testset "luxor-png" begin
            path = joinpath(DIR, "1d2d.png")
            save_png(path,M)
            @test isfile(path)
        end
    end

    @testset "2d2d" begin
        p = 3 # degree of polynomial
        k = KnotVector(1:12) # knot vector
        P = BSplineSpace{p}(k) # B-spline space
        rand_a = [randn(2) for i in 1:dim(P), j in 1:dim(P)] / 2
        a = [SVector(2i-6, 2j-6) for i in 1:dim(P), j in 1:dim(P)] + rand_a # random generated control points
        M = BSplineManifold(a, (P, P)) # Define B-spline manifold

        @testset "luxor-svg" begin
            path = joinpath(DIR, "2d2d.svg")
            save_svg(path,M)
            @test isfile(path)
        end
        @testset "luxor-png" begin
            path = joinpath(DIR, "2d2d.png")
            save_png(path,M)
            @test isfile(path)
        end
        @testset "luxor-color-a-png" begin
            color(u1,u2) = rand(RGB)
            path = joinpath(DIR, "2d2d_color-a.png")
            save_png(path, M, color)
            @test isfile(path)
        end
        @testset "luxor-color-b-png" begin
            n1, n2 = dim.(bsplinespaces(M))
            colors = rand(RGB,n1,n2)
            path = joinpath(DIR, "2d2d_color-b.png")
            save_png(path, M, colors)
            @test isfile(path)
        end
    end

    @testset "1d3d" begin
        p = 4 # degree of polynomial
        k = KnotVector(rand(12)) # knot vector
        P = BSplineSpace{p}(k) # B-spline space
        rand_a = [randn(3) for i in 1:dim(P)]
        a = [SVector(2i-6, 2i-6, 0) for i in 1:dim(P)] + rand_a # random generated control points
        M = BSplineManifold(a, (P,)) # Define B-spline manifold

        @testset "povray" begin
            path = joinpath(DIR, "1d3d.inc")
            save_pov(path,M)
            @test isfile(path)
        end
    end

    @testset "2d3d" begin
        p = 3 # degree of polynomial
        k = KnotVector(1:12) # knot vector
        P = BSplineSpace{p}(k) # B-spline space
        rand_a = [randn(3) for i in 1:dim(P), j in 1:dim(P)] / 2
        a = [SVector(2i-6, 2j-6, 0) for i in 1:dim(P), j in 1:dim(P)] + rand_a # random generated control points
        M = BSplineManifold(a, (P, P)) # Define B-spline manifold

        @testset "povray" begin
            path = joinpath(DIR, "2d3d.inc")
            save_pov(path, M , maincolor=RGBA(0,1,1,0.95))
            @test isfile(path)
        end
    end

    @testset "3d3d" begin
        p = 3 # degree of polynomial
        k = KnotVector(1:12) # knot vector
        P = BSplineSpace{p}(k) # B-spline space
        rand_a = [randn(3) for i in 1:dim(P), j in 1:dim(P), k in 1:dim(P)] / 4
        a = [SVector(2i-6, 2j-6, 2k-6) for i in 1:dim(P), j in 1:dim(P), k in 1:dim(P)] + rand_a # random generated control points
        M = BSplineManifold(a, (P, P, P)) # Define B-spline manifold

        @testset "povray" begin
            path = joinpath(DIR, "3d3d.inc")
            save_pov(path, M, maincolor=RGBA(0.5,0.1,1,0.8))
            @test isfile(path)
        end
    end
end


using StochasticDiffEq, Plots, DiffEqDevTools, DiffEqProblemLibrary
using ParameterizedFunctions, DiffEqMonteCarlo, Base.Test, OrdinaryDiffEq
using Plots; pyplot()
using BenchmarkTools

################################################################################

## Additive

prob = prob_sde_additive

dts = 1./2.^(10:-1:2) #14->7 good plot
sim1 = test_convergence(dts,prob,SOSRA(),numMonte=Int(1e3))
sim2 = test_convergence(dts,prob,SOSRA2(),numMonte=Int(1e3))
sim3 = test_convergence(dts,prob,SKenCarp(),numMonte=Int(1e3))

p1 = plot(dts,[sim1.errors[:l2],sim2.errors[:l2],sim3.errors[:l2]],
          title="Additive Convergence Tests",
          label=["SOSRA" "SOSRA2" "SKenCarp"],
          xtickfont = font(16, "Latin Modern Roman"),
          titlefont = font(20, "Latin Modern Roman"),
          ytickfont = font(16, "Latin Modern Roman"),
          guidefont = font(18, "Latin Modern Roman"),
          legendfont = font(14, "Latin Modern Roman"),
          ylabel = "Error", xlabel="dt",
          yscale=:log10,xscale=:log10,lw=3)

dts = 1./2.^(10:-1:2) #14->7 good plot
ref = dts.^2 * 1e-4
plot!(p1,dts,ref,linestyle=:dash,lw=3,label="2nd Order Reference")

prob = prob_sde_2Dlinear

dts = 1./2.^(7:-1:4) #14->7 good plot
srand(100)
sim1 = test_convergence(dts,prob,SOSRI(),numMonte=1000)
sim2 = test_convergence(dts,prob,SOSRI2(),numMonte=1000)

p2 = plot(dts,[sim1.errors[:l2],sim2.errors[:l2]],
          title="Linear Convergence Tests",label=["SOSRI" "SOSRI2"],
          ylabel = "Error", xlabel="dt",
          xtickfont = font(16, "Latin Modern Roman"),titlefont = font(20, "Latin Modern Roman"),
          ytickfont = font(16, "Latin Modern Roman"),guidefont = font(18, "Latin Modern Roman"),
          legendfont = font(14, "Latin Modern Roman"),
          yscale=:log10,xscale=:log10,lw=3)

dts = 1./2.^(7:-1:4) #14->7 good plot
ref = dts.^1.5 * 5e-1
plot!(p2,dts,ref,linestyle=:dash,lw=3,label="1.5 Order Reference")

α = 0.1
β = 0.5
ff1 = (u,p,t) -> β./sqrt.(1+t)
ff2 = (u,p,t) -> - u./(2*(1+t))
σ2 = (u,p,t) -> α*β./sqrt.(1+t)
prob = SplitSDEProblem(ff1,ff2,σ2,1.,(0.0,1.0))
(::typeof(prob.f))(::Type{Val{:analytic}},u0,p,t,W) = u0./sqrt.(1+t) + β*(t+α*W)./sqrt.(1+t)

dts = 1./2.^(10:-1:2) #14->7 good plot
sim = test_convergence(dts,prob,SKenCarp(),numMonte=Int(1e3))

p3 = plot(dts,sim.errors[:l2],
          title="IMEX Convergence Tests",label="IMEX SKenCarp",
          ylabel = "Error", xlabel="dt",
          xtickfont = font(16, "Latin Modern Roman"),titlefont = font(20, "Latin Modern Roman"),
          ytickfont = font(16, "Latin Modern Roman"),guidefont = font(18, "Latin Modern Roman"),
          legendfont = font(14, "Latin Modern Roman"),
          yscale=:log10,xscale=:log10,lw=3)

ref = dts.^2 * 1e-4
plot!(p3,dts,ref,linestyle=:dash,lw=3,label="2nd Order Reference")

plot(p1,p2,p3,layout=(3,1),size=(600,800))

savefig("convergence.png")
savefig("convergence.pdf")


struct PdeBSM{B <: AbstractBlackScholesProcess}
  process::B
end

diffusion(pde::PdeBSM, t::Float64, x::Float64) = diffusion(pde.process, t, x)
drift(pde::PdeBSM, t::Float64, x::Float64) = drift(pde.process, t, x)
function discount(pde::PdeBSM, t::Float64, ::Float64)
  if abs(t) < 1e-8
    t = 0
  end

  return forward_rate(pde.process.riskFreeRate, t, t, ContinuousCompounding(), NoFrequency()).rate
end

mutable struct PdeConstantCoeff
  diffusion::Float64
  drift::Float64
  discount::Float64
end

function PdeConstantCoeff(process::AbstractBlackScholesProcess, t::Float64, x::Float64)
  pde = PdeBSM(process)
  diffusion_ = diffusion(pde, t, x)
  drift_ = drift(pde, t, x)
  discount_ = discount(pde, t, x)

  return PdeConstantCoeff(diffusion_, drift_, discount_)
end

# diffusion(pdecc::PdeConstantCoeff, ::Float64, ::Float64) = pdecc.diffusion
# drift(pdecc::PdeConstantCoeff, ::Float64, ::Float64) = pdecc.drift
# discount

function generate_operator!(pdecc::PdeConstantCoeff, t::Float64, tg::TransformedGrid, L::TridiagonalOperator)
  @inbounds @simd for i = 2:length(tg.grid) - 1
    sigma = pdecc.diffusion
    nu = pdecc.drift
    r = pdecc.discount

    sigma2 = sigma * sigma

    pd = -(sigma2 / tg.dxm[i] - nu) / tg.dx[i]
    pu = -(sigma2 / tg.dxp[i] + nu) / tg.dx[i]
    pm = sigma2 / (tg.dxm[i] * tg.dxp[i]) + r
    set_mid_row!(L, i, pd, pm, pu)
  end

  return L
end


function setupTravelTimeTomography(filenamePrefix::String, Minv,Iact,sback,modFun)
###########################################################
### Read receivers and sources files
RCVfile = string(filenamePrefix,"_rcvMap.dat");
SRCfile = string(filenamePrefix,"_srcMap.dat");

srcNodeMap = readSrcRcvLocationFile(SRCfile,Minv);
rcvNodeMap = readSrcRcvLocationFile(RCVfile,Minv);

Q = generateSrcRcvProjOperators(Minv.n+1,srcNodeMap);
Q = Q.*1/(norm(Minv.h)^2);
P = generateSrcRcvProjOperators(Minv.n+1,rcvNodeMap);

println("Travel time tomography: ",size(Q,2)," sources.");
#############################################################################################

HO = false;

println("Reading data:");

(DobsEik,WdEik) =  readDataFileToDataMat(string(filenamePrefix,"_travelTime.dat"),srcNodeMap,rcvNodeMap);

N = prod(Minv.n+1);

########################################################################################################
##### Set up remote workers ############################################################################
########################################################################################################

EikMPIWorkers = nworkers(); # this just set the maximal MPI workers. To activate parallelism, run addprocs()

(pFor,contDiv,SourcesSubInd) = getEikonalInvParam(Minv,Q,P,HO,EikMPIWorkers);

misfun = SSDFun

Wd 		= Array{Array{Float64}}(length(pFor));
dobs 	= Array{Array{Float64}}(length(pFor));
for i=1:length(pFor)
	I_i = SourcesSubInd[i];
	Wd[i]   = 0.5*WdEik[:,I_i];
	dobs[i] = DobsEik[:,I_i];
end
WdEik = 0;
DobsEik = 0;

pMisRFs = getMisfitParam(pFor, Wd, dobs, misfun, Iact,sback,ones(length(pFor)),modFun);

return (Q,P,pMisRFs,SourcesSubInd,contDiv);
end


"""
    invariantset(embedding, tolerance)

Determine whether an embedding forms an invariant set.
"""
function invariantset(embedding)
    lastpoint = embedding[end, :]
    dim = size(embedding, 2)

    # Triangulate the embedding using all points but the last
    triangulation = Simplices.Delaunay.delaunayn(embedding[1:end-1, :])

    points = embedding[1:end-1, :]
    simplex_indices = triangulation

    # Centroids and radii of simplices in the triangulation
    centroids, radii = centroids_radii2(points, simplex_indices)

    lastpoint_matrix = repmat(lastpoint', size(centroids, 1), 1)

    # Find simplices that can contain the last point (not all can)
    dists_lastpoint_and_centroids = sum((lastpoint_matrix - centroids).^2, 2)
    distdifferences = radii.^2 - dists_lastpoint_and_centroids

    # Find the row indices of the simplices that possibly contain the last point (meaning that
    # dist(simplex_i, lastpoint) <= radius(simplex), so the circumsphere of the simplex
    # contains the last point.
    valid_simplex_indices = find(heaviside0(distdifferences) .* collect(1:size(triangulation[2], 1)))

    n_validsimplices = size(valid_simplex_indices, 1)

    # Loop over valid simplices and check whether the corresponding simplex actually
    # contains the last point.
    i = 1
    lastpoint_contained = false

    # Continue checking if simplices contain the last point until some simplex contains it.
    # Then the point must necessarily be inside the convex hull of the triangulation formed
    # by those simplices.
    while i <= n_validsimplices && !lastpoint_contained
        # Subsample the valid simplex and compute its orientation
        simplex = embedding[simplex_indices[valid_simplex_indices[i], :], :]
        orientation_simplex = det([ones(dim + 1, 1) simplex])

        beta = 1 # Convex expansion coefficient

        j = 1
        while j <= dim + 1 && beta >= 0
            tmp = copy(simplex)
            tmp[j, :] = copy(lastpoint)
            beta = det([ones(dim + 1, 1) tmp]) * sign(orientation_simplex)
            j = j + 1
        end

        # If the last convex expansion coefficient is positive, the last point is contained
        # in the triangulation (because all the previous coefficients must have been nonnegative)
        if beta >= 0
            lastpoint_contained = true
        end

        i = i + 1
    end
    # If the last point is contained in the triangulation, the set is invariant.
    return lastpoint_contained
end


using MultilinearOpt
using JuMP
using Gurobi
using Test
using Random
using LinearAlgebra
using MathOptInterface

const MOI = MathOptInterface

function test_constraint_violation(f, xbounds::NTuple{2}, ybounds::NTuple{2};
        method::Symbol, disc_level::Integer, rng::AbstractRNG, atol::Real, num_tests::Integer,
        sense::MOI.OptimizationSense)
    m = Model(with_optimizer(Gurobi.Optimizer, OutputFlag=0))
    @variable m x lower_bound=xbounds[1] upper_bound=xbounds[2]
    @variable m y lower_bound=ybounds[1] upper_bound=ybounds[2]
    @variable m z
    @constraint m z == f(x, y)
    relaxbilinear!(m, method=method, disc_level=disc_level)
    for i = 1 : num_tests
        if sense != MOI.FEASIBILITY_SENSE
            x̄ = xbounds[1] + rand(rng) * (xbounds[2] - xbounds[1])
            ȳ = ybounds[1] + rand(rng) * (ybounds[2] - ybounds[1])
            z̄ = f(x̄, ȳ) + rand(rng) - 0.5
            min_objective = sum(x -> x^2, [x, y, z] - [x̄, ȳ, z̄])
            if sense == MOI.MIN_SENSE
                @objective m Min min_objective
            else
                @objective m Max -min_objective
            end
        end
        optimize!(m)
        bounds_tol = 1e-8
        @test xbounds[1] - bounds_tol <= value(x) <= xbounds[2] + bounds_tol
        @test ybounds[1] - bounds_tol <= value(y) <= ybounds[2] + bounds_tol
        violation = f(value(x), value(y)) - value(z)
        # @show violation
        # @show value(x) value(y) value(z)
        # println()
        @test violation ≈ 0 atol=atol
    end
end

@testset "constraint violation" begin
    rng = MersenneTwister(1)
    atol = 1.1e-2
    num_tests = 10
    xbounds = (-0.5, 0.5)
    ybounds = (-0.4, 0.6)
    for (method, disc_level) in  [
            (:Logarithmic1D, 31),
            (:Logarithmic2D, 6),
            (:ZigZag1D, 31),
            (:ZigZag2D, 6),
            # (:Unary, 31) # broken
            (:MisenerLinear, 31),
            (:MisenerLog1, 31),
            (:MisenerLog2, 31)]

        for sense in [MOI.MIN_SENSE, MOI.MAX_SENSE, MOI.FEASIBILITY_SENSE]
            test_constraint_violation(*, xbounds, ybounds;
                method=method, disc_level=disc_level,
                atol=atol, rng=rng, num_tests=num_tests, sense=sense)

            test_constraint_violation((x, y) -> x * y + 3, xbounds, ybounds;
                method=method, disc_level=disc_level,
                atol=atol, rng=rng, num_tests=num_tests, sense=sense)

            test_constraint_violation((x, y) -> x^2, xbounds, ybounds;
                method=method, disc_level=disc_level,
                atol=atol, rng=rng, num_tests=num_tests, sense=sense)

            test_constraint_violation((x, y) -> x^2 + x * y - y, xbounds, ybounds;
                method=method, disc_level=disc_level,
                atol=2 * atol, rng=rng, num_tests=num_tests, sense=sense)
        end
    end
end


# Imports
using Random
using LinearAlgebra
using SparseArrays
using Arpack: eigs

# ESN struct
mutable struct ESN
    inSize::Integer # Tamanho da entrada
    outSize::Integer # Tamanho da saída
    resSize::Integer # Tamanho do reservatório
    leaking_rate::AbstractFloat # leaking rate
    sparsity::AbstractFloat
    spectral_radius::AbstractFloat
    W
    Win
    Wout
    states
    x
    Yt

    function ESN(; inSize::Integer=1,
        outSize::Integer=1,
        resSize::Integer=500,
        leaking_rate::AbstractFloat=0.5,
        sparsity::AbstractFloat=0.6,
        spectral_radius::AbstractFloat=0.9)

        Random.seed!(42)
        # pesos da entrada
        Win = rand(Float64, (resSize, 1+inSize)) .- 0.5

        # pesos do reservatório
        W = sprand(Float64,resSize,resSize, 1-sparsity)
        W = W - 0.5*(W .!= 0)

        rhoW = 0

        try
            rhoW = maximum(abs.(eigs(W)[1])) # mais rápido
        catch e
            rhoW = maximum(abs.(eigvals(Matrix(W)))) #funciona sempre
        end

        W = W * spectral_radius / rhoW

        esn = new()
        esn.inSize = inSize
        esn.outSize = outSize
        esn.resSize = resSize
        esn.leaking_rate = leaking_rate
        esn.sparsity = sparsity
        esn.spectral_radius = spectral_radius
        esn.Win = Win
        esn.W = W

        return esn
    end
end

function fit!(esn::ESN, data)
    # matriz de estados
    states = zeros(Float64, (1+esn.inSize+esn.resSize, length(data)-1))

    # matriz "target"
    Yt = (data[(2):(length(data))])'

    # rodando o reservatório e armazenando valores na matriz de estados
    x = zeros(Float64, (esn.resSize, 1))

    for t = 1:length(data)-1
        u = data[t]
        x = (1-esn.leaking_rate)*x + esn.leaking_rate*tanh.( esn.Win*[1;u] + esn.W*x )
        states[:, t] = [1;u;x]
    end

    reg = 1e-8  # regularization

    esn.Wout = Yt*states' * pinv(states*states' + reg*I(1+esn.inSize+esn.resSize))

    esn.states = states
    esn.x = x
    esn.Yt = Yt

end

function predict!(esn::ESN, steps, v1)
    Y = zeros(esn.outSize, steps)
    u = esn.Yt[length(esn.Yt)]
    x = esn.x
    for t = 1:steps
        x = (1-esn.leaking_rate)*x + esn.leaking_rate*tanh.( esn.Win*[1;u] + esn.W*x )
        y = esn.Wout*[1;u;x]
        Y[:,t] = y
        # generative mode:
        u = y
    end

    return Y
end


"""
Implementation of Differential Evolution: DE/**selection**/n/**recombination**

The constructor takes following keyword arguments:

- `populationSize`: The size of the population
- `F`: the differentiation (mutation) scale factor (default: 0.9). It's usually defined in range ``F \\in (0, 1+]``
- `n`: the number of differences used in the perturbation (default: 1)
- `selection`: the selection strategy function (default: [`random`](@ref))
- `recombination`: the recombination functions (default: [`uniformbin(0.5)`](@ref))
- `K`: the recombination scale factor (default: 0.5*(F+1))
"""
@kwdef struct DE{F1,F2} <: AbstractOptimizer
    populationSize::Integer = 50
    F::Real = 0.9
    n::Integer = 1
    K::Real = 0.5 * (F + 1)
    selection::F1 = random
    recombination::F2 = uniformbin(0.5)
end
population_size(method::DE) = method.populationSize
default_options(method::DE) = (abstol = 1e-10,)
summary(m::DE) = "DE/$(m.selection)/$(m.n)/$(m.recombination)"

mutable struct DEState{T,IT} <: AbstractOptimizerState
    N::Int
    fitness::Vector{T}
    fittest::IT
end
value(s::DEState) = minimum(s.fitness)
minimizer(s::DEState) = s.fittest

"""Initialization of ES algorithm state"""
function initial_state(method::DE, options, objfun, population)
    T = typeof(value(objfun))
    individual = first(population)
    N = length(individual)
    fitness = fill(maxintfloat(T), method.populationSize)

    # setup initial state
    return DEState(N, fitness, copy(individual))
end

function update_state!(objfun, constraints, state, population::AbstractVector{IT}, method::DE, itr) where {IT}

    # setup
    Np = method.populationSize
    n = method.n
    F = method.F

    offspring = Array{IT}(undef, Np)

    # select base vectors
    bases = method.selection(state.fitness, Np)

    # select target vectors
    for (i,b) in enumerate(bases)
        # mutation
        base = population[b]
        offspring[i] = copy(base)
        # println("$i => base:", offspring[i])

        targets = randexcl(1:Np, [i], 2 * n)
        offspring[i] = differentiation(offspring[i], @view population[targets]; F=F)
        # println("$i => mutated:", offspring[i], ", targets:", targets)

        # recombination
        offspring[i], _ = method.recombination(offspring[i], base)
        # println("$i => recombined:", offspring[i])
    end

    # Create new generation
    fitidx = 0
    minfit = Inf
    for i in 1:Np
        o = apply!(constraints, offspring[i])
        v = value(objfun, o) + penalty(constraints, o)
        if (v <= state.fitness[i])
            population[i] = o
            state.fitness[i] = v
            if v < minfit
                minfit = v
                fitidx = i
            end
        end
    end

    # set best individual
    if fitidx > 0
        state.fittest = population[fitidx]
    end

    return false
end


using Distributed
addprocs(Sys.CPU_THREADS)

import PyPlot; const plt = PyPlot
import DSP: conv
import Random: seed!, randperm
import Statistics: std
import LinearAlgebra: norm, dot
import Arpack: eigs
using Printf

import KirchMig

dz, dx = 15, 15
x = 0:dx:1000; nx = length(x)
z = 0:dz:1000; nz = length(z)

rho = 1000ones(nz, nx)
rho[div(330,dx):end,:] .+= 1000
rho[div(670,dz):end,:] .-= 500

vel = 2000 .+ zero(rho)

imp = 1e-3rho.*vel;

refl = [zeros(nx)'
    (imp[2:end,:] .- imp[1:end-1,:])./(imp[2:end,:] .+ imp[1:end-1,:])]

refl[div(nz, 2)-1, div(nx, 2)] += 0.2
refl[div(nz, 2)+1, div(nx, 2)] += 0.2;

nr = 46

rec_z = zeros(nr)
rec_x = range(x[1], stop=x[end], length=nr)

ns = 10
seed!(12)
src_z = zeros(ns)
src_x = sort(rec_x[randperm(nr)][1:ns])

trav_r = KirchMig.eikonal_const_vel([rec_z rec_x], z, x, vel[1]);
trav_s = KirchMig.eikonal_const_vel([src_z src_x], z, x, vel[1]);

t = 0:0.008:1; nt = length(t)
G = KirchMig.KirchMap(t, trav_r, trav_s)

ricker(t0, f) = @. (1 - 2pi^2 * f^2 * t0^2) * exp(-pi^2 * f^2 * t0^2)
rick_dtt = ricker(t .- t[div(nt,5)], 15);
@views rick_dtt[2:end-1] = (rick_dtt[1:end-2] - 2.0*rick_dtt[2:end-1] + rick_dtt[3:end])/(t[2] - t[1])^2;
rick_dtt /= maximum(abs.(rick_dtt));

W = KirchMig.ConvMap(rick_dtt, nr, ns, nt);

L = W*G;

@time d = L*refl[:];

seed!(1)
n = randn(size(d))
d += n*std(d)/std(n);

@time m_mig = L'd;

xran = x[end]-x[1]
mod_mig = reshape(m_mig, nz, nx)
vmin, vmax = maximum(abs.(mod_mig))*[-1,1]
fig, ax = plt.subplots(figsize=(6,4))
cax = ax.imshow(mod_mig, extent=[x[1], x[end], z[end], z[1]],
    vmin=vmin, vmax=vmax, aspect="equal", cmap="gray", interpolation=nothing)
cbar = fig.colorbar(cax, ax=ax)
cbar.ax.set(ylabel="Amplitude");
ax.plot(xran*mod_mig[:,div(end,2)]./(4vmax) .+ xran./2, z, color="#d62728");
ax.set(xlim=(x[1], x[end]), ylim=(z[end], z[1]), xlabel="Position [m]", ylabel="Depth [m]",
    title="Migrated image")
fig.tight_layout()

J0(m) = norm(L*m - d)^2

function ∇J0!(storage, m)
    storage[1:end] = 2L'L*m - 2m_mig
end

seed!(123);

nL = size(L, 2)
u = randn(nL)
v = randn(nL)
h = 0.001*maximum(abs.(u))

storage = zero(u)
∇J0!(storage, u)

g1 = (J0(u + h*v) - J0(u - h*v))/2h
g2 = dot(storage, v)
err = 100abs(g1 - g2)/((g1+g2)/2)
println("$(@sprintf("%.2f", err))% error")

import Optim

@time res = Optim.optimize(J0, ∇J0!, zeros(size(L, 2)),
                           Optim.ConjugateGradient(),
                           Optim.Options(iterations = 20, show_trace=true))

pclip=0.5
xran = x[end]-x[1]
mod_lsm_0 = reshape(res.minimizer, nz, nx)
vmin, vmax = pclip*maximum(abs.(mod_lsm_0))*[-1,1]
fig, ax = plt.subplots(figsize=(6,4))
cax = ax.imshow(mod_lsm_0, extent=[x[1], x[end], z[end], z[1]],
    vmin=vmin, vmax=vmax, aspect="equal", cmap="gray", interpolation=nothing)
cbar = fig.colorbar(cax, ax=ax)
cbar.ax.set(ylabel="Amplitude");
ax.plot(xran*mod_lsm_0[:,div(end,2)]./(4vmax) .+ xran./2, z, color="#d62728");
ax.set(xlim=(x[1], x[end]), ylim=(z[end], z[1]), xlabel="Position [m]", ylabel="Depth [m]",
    title="LSM image: no reg.")
fig.tight_layout()

@time λ, ϕ = eigs(L'L; nev=1, maxiter=2); λ = λ[1]

function TV(x)
    grad = KirchMig.gradient(x, nz, nx) # Returns a nz × nx × 2 array in 2D
    norm_grad = sqrt.(sum(grad.^2, dims=3))
    e = maximum([100eps(), 0.01maximum(norm_grad)])
    return 0.002*λ*sum(sqrt.(norm_grad.^2 .+ e^2))
end
function ∇TV!(storage, x)
    grad = KirchMig.gradient(x, nz, nx)
    norm_grad = sqrt.(sum(grad.^2, dims=3))
    e = maximum([100eps(), 0.01maximum(norm_grad)])
    grad ./= sqrt.(norm_grad.^2 .+ e^2)

    storage[1:end] = -0.002*λ*KirchMig.divergence(grad[:], nz, nx)[:]
end

seed!(123);

nL = size(L, 2)
u = randn(nL)
v = randn(nL)
h = 0.001*maximum(abs.(u))

storage = zero(u)
∇TV!(storage, u)

g1 = (TV(u + h*v) - TV(u - h*v))/2h
g2 = dot(storage, v)
err = 100abs(g1 - g2)/((g1+g2)/2)
println("$(@sprintf("%.2f", err))% error")

JTV(x) = J0(x) + TV(x)

function ∇JTV!(storage, x)
    ∇J0!(storage, x)
    s = storage[:]
    ∇TV!(storage, x)
    storage[1:end] = s + storage[1:end]
end

# Check gradient
seed!(123);

nL = size(L, 2)
u = randn(nL)
v = randn(nL)
h = 0.001*maximum(abs.(u))

storage = zero(u)
∇JTV!(storage, u)

g1 = (JTV(u + h*v) - JTV(u - h*v))/2h
g2 = dot(storage, v)
err = 100abs(g1 - g2)/((g1+g2)/2)
println("$(@sprintf("%.2f", err))% error")

@time res_tv = Optim.optimize(JTV, ∇JTV!, zeros(size(L, 2)),
                              Optim.ConjugateGradient(),
                              Optim.Options(iterations = 20, show_trace=true))

pclip=0.5
xran = x[end]-x[1]
mod_lsm_tv = reshape(res_tv.minimizer, nz, nx)
vmin, vmax = pclip*maximum(abs.(mod_lsm_tv))*[-1,1]
fig, ax = plt.subplots(figsize=(6,4))
cax = ax.imshow(mod_lsm_tv, extent=[x[1], x[end], z[end], z[1]],
    vmin=vmin, vmax=vmax, aspect="equal", cmap="gray", interpolation=nothing)
cbar = fig.colorbar(cax, ax=ax)
cbar.ax.set(ylabel="Amplitude");
ax.plot(xran*mod_lsm_tv[:,div(end,2)]./(4vmax) .+ xran./2, z, color="#d62728");
ax.set(xlim=(x[1], x[end]), ylim=(z[end], z[1]), xlabel="Position [m]", ylabel="Depth [m]",
    title="LSM image: TV reg.")
fig.tight_layout()


if basename(pwd()) == "aoc"
    cd("2017/13")
end

using OffsetArrays

function loadfirewall(filename::AbstractString)
    data = collect(parse.(Int, split(line, ": ")) for line in eachline(filename))
    firewall = OffsetArray(zeros(Int, first(last(data)) + 1), -1)
    for (i, v) in data
        firewall[i] = v
    end
    firewall
end

function part1(firewall, start = 0)
    total = 0
    for (depth, range) in zip(Iterators.countfrom(start), firewall)
        if !iszero(range) && iszero(scannerposition(range, depth))
            total += depth * range
        end
    end
    total
end

function scannerposition(range, step)
    if range < 2
        return 0
    end
    roundrange = 2 * range - 2
    roundstep = step % roundrange
    if roundstep < range
        roundstep
    else
        roundrange - roundstep
    end
end

@assert collect(scannerposition(5, x) for x in 0:10) == [0:4..., 3:-1:0..., 1:2...]
@assert part1(loadfirewall("example.txt")) == 24
part1(loadfirewall("input.txt"))

function part2(filename::AbstractString)
    firewall = loadfirewall(filename)
    for start in Iterators.countfrom(1)
        if all(
            iszero(range) || !iszero(scannerposition(range, depth))
            for (depth, range) in zip(Iterators.countfrom(start), firewall)
        )
            return start
        end
    end
end

@assert part2("example.txt") == 10
part2("input.txt")


using ADCME
using PyCall
using LinearAlgebra
using PyPlot
using Random
Random.seed!(233)

function compute_vel(a,v0,psi,sigma,tau,eta)
    compute_vel_ = load_op_and_grad("./build/libComputeVel","compute_vel")
a,v0,psi,sigma,tau,eta = convert_to_tensor([a,v0,psi,sigma,tau,eta], [Float64,Float64,Float64,Float64,Float64,Float64])
    compute_vel_(a,v0,psi,sigma,tau,eta)
end

# verify
n = 10
a = 2.0*ones(n)|>constant
x = 4.0*ones(n)|>constant
v0 = 0.001|>constant
Ψ = 2000.0*ones(n)|>constant
σ = 3.0*ones(n)|>constant
η = 2.0|>constant
s = x/2v0
inv_o = exp(-Ψ/a)
τ = a * (Ψ/a + log(s + sqrt(s*s+inv_o*inv_o)))*σ + η*x
# τ = a * asinh(*exp(Ψ/a))*σ + η*(x-u)/Δt

# TODO: specify your input parameters
u = compute_vel(a,v0,Ψ,σ,τ,η)
sess = Session(); init(sess)
@show run(sess, u)
# error()

# TODO: change your test parameter to `m`
#       in the case of `multiple=true`, you also need to specify which component you are testings
# gradient check -- v
function scalar_function(m)
    # return sum(compute_vel(a,v0,Ψ,σ,τ,η)^2)
    return sum(compute_vel(a,v0,Ψ,σ,τ,m)^2)
end

# TODO: change `m_` and `v_` to appropriate values
# m_ = constant(rand(n))
# v_ = rand(n)

# m_ = τ
# v_ = rand(n)

m_ = constant(rand())
v_ = rand()

y_ = scalar_function(m_)
dy_ = gradients(y_, m_)
ms_ = Array{Any}(undef, 5)
ys_ = Array{Any}(undef, 5)
s_ = Array{Any}(undef, 5)
w_ = Array{Any}(undef, 5)
gs_ =  @. 1 / 10^(1:5)

for i = 1:5
    g_ = gs_[i]
    ms_[i] = m_ + g_*v_
    ys_[i] = scalar_function(ms_[i])
    s_[i] = ys_[i] - y_
    w_[i] = s_[i] - g_*sum(v_*dy_)
end

sess = Session(); init(sess)
sval_ = run(sess, s_)
wval_ = run(sess, w_)
close("all")
loglog(gs_, abs.(sval_), "*-", label="finite difference")
loglog(gs_, abs.(wval_), "+-", label="automatic differentiation")
loglog(gs_, gs_.^2 * 0.5*abs(wval_[1])/gs_[1]^2, "--",label="\$\\mathcal{O}(\\gamma^2)\$")
loglog(gs_, gs_ * 0.5*abs(sval_[1])/gs_[1], "--",label="\$\\mathcal{O}(\\gamma)\$")

plt.gca().invert_xaxis()
legend()
xlabel("\$\\gamma\$")
ylabel("Error")


#Structure containg design matrix and related info for each submodel"
mutable struct UmDesign
  name::Symbol
  formula::FormulaTerm
  link::Link
  data::DataFrame
  coefnames::Array{String}
  mat::Array{Float64,2}
  idx::Union{UnitRange{Int64}, Nothing}
end

#Check formulas for validity
function check_formula(formula::FormulaTerm, data::DataFrame)
  vars = varnames(formula).rhs
  if isnothing(vars) return nothing end
  vars = String.(vars)
  covs = names(data)
  for x in vars
    if x ∉ covs error(string(x," not found in data")) end
  end
  return nothing
end

#Add dummy column for response variable to data
add_resp = function(x::DataFrame, f::FormulaTerm)
  lhs = f.lhs.sym
  if lhs in names(x) return x end
  dc = deepcopy(x)
  dc[!,lhs] = zeros(nrow(dc))
  return dc
end

#Workaround for predicting from newdata with missing cat variable levels

#Get number of unique levels for a categorical variable
function n_levels(x::AbstractArray)
  typeof(x)<:CategoricalArray ? length(levels(x)) : 0
end

#Repeat levels of cat variable to certain length
function rep_levels(x::AbstractArray, n::Int)
  if typeof(x)<:CategoricalArray
    reps = Int(ceil(n / length(levels(x))))
    out = categorical(repeat(levels(x), reps))
    levels!(out, levels(x))
    return out[1:n]
  end
  return repeat([0], n)
end

#Create data frame containing all categorical levels to bind to original data
function add_levels(df::DataFrame)
  nlev = [n_levels(col) for col = eachcol(df)]
  maxlev = max(nlev...)
  ndf = DataFrame([rep_levels(col, maxlev) for col = eachcol(df)], :auto)
  rename!(ndf, names(df))
  return ndf
end

#Build and apply schema
function get_schema(formula::FormulaTerm, data::DataFrame)
  check_formula(formula, data)
  dat = add_resp(data, formula)
  dat_aug = [dat; add_levels(dat)]
  sch = schema(formula, dat_aug)
  apply_schema(formula, sch, StatisticalModel)
end

#Outer constructor for UmDesign objects
function UmDesign(name::Symbol, formula::FormulaTerm, link::Link,
                  data::DataFrame)

  asch = get_schema(formula, data)
  cn = coefnames(asch.rhs)
  cn = cn isa String ? [cn] : cn
  mat = modelcols(asch.rhs, data)
  UmDesign(name, formula, link, data, cn, mat, nothing)
end

#Add indexes to map combined coefs back to submodels
function add_idx!(dm::Array{UmDesign})
  idx = 1
  for x in dm
    np = length(x.coefnames)
    x.idx = idx:(idx+np-1)
    idx += np
  end
  return nothing
  #return idx - 1
end

#Get total parameters
function get_np(dm::Array{UmDesign})
  sum(map(x -> length(x.coefnames), dm))
end

#Transform linear predictor back to response scale
function transform(dm::UmDesign, coefs::Array)
  lp = dm.mat * coefs[dm.idx]
  invlink(lp, dm.link)
end


@testset "Overrides" begin
    g3 = SimpleWeightedGraph(path_graph(5))
    g3_d = [1. 0. 0. 0. 0.;
            0. 2. 0. 0. 0.;
            0. 0. 2. 0. 0.;
            0. 0. 0. 2. 0.;
            0. 0. 0. 0. 1.]
    g3_l = [1. -1. 0. 0. 0.;
           -1. 2. -1. 0. 0.;
            0. -1. 2. -1. 0.;
            0. 0. -1. 2. -1.;
            0. 0. 0. -1. 1.]
    g5 = SimpleWeightedDiGraph(4)
    g5_din = [0. 0. 0. 0.;
              0. 2. 0. 0.;
              0. 0. 3. 0.;
              0. 0. 0. 5.]
    g5_dout = [3. 0. 0. 0.;
               0. 2. 0. 0.;
               0. 0. 5. 0.;
               0. 0. 0. 0.]
    g5_dboth = [3. 0. 0. 0.;
                0. 4. 0. 0.;
                0. 0. 8. 0.;
                0. 0. 0. 5.]
    g5_l = [3. -2. -1. 0.;
            0. 2. -2. 0.;
            0. 0. 5. -5.;
            0. 0. 0. 0.]

    for g in testgraphs(g3)
        @test @inferred(adjacency_matrix(g, Bool)) == adjacency_matrix(g, Bool; dir=:out)
        @test @inferred(adjacency_matrix(g))[3, 2] == 1
        @test degree_matrix(g, Float64, dir=:out) == g3_d
        @test degree_matrix(g, Float64, dir=:out) == degree_matrix(g, Float64, dir=:in)
        @test adjacency_matrix(g)[2, 4] == 0
        @test adjacency_matrix(g; dir=:out) == adjacency_matrix(g; dir=:in)'
        @test issymmetric(laplacian_matrix(g))
        @test laplacian_matrix(g, Float64) ≈ g3_l
        @test g[1:3] == SimpleWeightedGraph{eltype(g), weighttype(g)}(path_graph(3))
        gx = copy(g)
        add_edge!(gx, 2, 3, 99)
        gi = gx[2:4]
        @test weights(gi)[1, 2] == 99

        h = @inferred(cartesian_product(g, g))
        @test nv(h) == 25
        @test ne(h) == 40
        gz = g3[1:4]
        add_edge!(gz, 3, 4, 87)
        @test weights(cartesian_product(g3,gz))[11,12]==weights(gz)[3,4]
    end

    add_edge!(g5, 1, 2, 2); add_edge!(g5, 2, 3, 2); add_edge!(g5, 1, 3, 1); add_edge!(g5, 3, 4, 5)
    for g in testdigraphs(g5)
        @test degree_matrix(g, Float64, dir=:out) == g5_dout
        @test degree_matrix(g, Float64, dir=:in) == g5_din
        @test degree_matrix(g, Float64, dir=:both) == g5_dboth
        @test_throws DomainError degree_matrix(g, dir=:other)
        @test @inferred(adjacency_matrix(g, Int64)) == adjacency_matrix(g, Int64; dir=:out)
        @test adjacency_matrix(g; dir=:out) == adjacency_matrix(g; dir=:in)'
        @test !issymmetric(laplacian_matrix(g))
        @test laplacian_matrix(g, Float64) ≈ g5_l
        @test @inferred(pagerank(g))[3] ≈ 0.2266 atol=0.001
        @test length(@inferred(pagerank(g))) == nv(g)
        @test_throws ErrorException pagerank(g, 2)
        @test_throws ErrorException pagerank(g, 0.85, 2)

        gc = SimpleWeightedDiGraph(path_digraph(2), 2)
        @test g[2:3] == SimpleWeightedDiGraph{eltype(g5), weighttype(g5)}(gc)
        @test weights(g[2:3])[1, 2] == 2
    end
end


using QuantumCumulants
using SymbolicUtils
using Test

@testset "nlevel" begin

# Symbolic levels
ha = NLevelSpace(:atom, (:g,:e))
@test ha.GS == :g

σ = Transition(ha, :σ, :g,:e)
@test σ' == Transition(ha, :σ, :e, :g)
σee = Transition(ha, :σ, :e,:e)
@test σee'==σee

@test σ'*σ == σee

ex = σ*σ'
@test isequal(simplify(σ*σ'), simplify(1 - σee))

sz = σ'*σ - σ*σ'
@test isequal(simplify(sz),simplify(2*σee - 1))

# Integer levels
ha = NLevelSpace(:atom, 2)
@test ha.GS == 1

@test_throws AssertionError Transition(ha, :σ, :g,:e)
σ = Transition(ha, :σ, 1, 2)
@test σ' == Transition(ha, :σ, 2,1)
σee = Transition(ha, :σ, :2,2)
@test σee'==σee

@test σ'*σ == σee

ex = σ*σ'
σgg = Transition(ha, :σ,1,1)
@test isequal(ex, σgg)

sz = σ'*σ - σ*σ'
@test isequal(simplify(sz),simplify(2*σee - one(σgg)))

# Product space
ha1 = NLevelSpace(:atom, (:g,:e))
ha2 = NLevelSpace(:atom, 2)
@test ha1 != ha2

hprod = ha1⊗ha2
σ1 = Transition(hprod,:σ,:g,:e,1)
σ2 = Transition(hprod,:σ,1,2,2)
@test σ1'*σ1==Transition(hprod,:σ,:e,:e,1)
@test isequal(simplify(σ2*σ2'), simplify(1 - Transition(hprod,:σ,2,2,2)))
@test isequal(simplify(σ1*σ2), σ1*σ2)

@test_throws ErrorException Transition(hprod,:σ,:g,:e)
@test Transition(hprod,:σ,:g,:e,1)==σ1
@test Transition(hprod,:σ,1,2,2)==σ2
@test_throws AssertionError Transition(hprod,:σ,1,2,1)

# Callable constructor
@test_throws ErrorException Transition(hprod,:σ)
s = Transition(hprod,:σ,1)
@test s isa QuantumCumulants.CallableTransition
@test s(:g,:e) == σ1
@test acts_on(Transition(ha2,:s))==1

end # testset


"""
    xnew,unew,cnew,sigmanew = forward_pass(traj_new, x0,u,x,α,f,costfun,lims,diff)
# Arguments
- α: step size (αk is applied to old trajectory)
- diff: function to determine difference `diff(xnew[:,i], x[:,i])`
- f: forward dynamics `x(k+1)  = f(x(k), u(k), k)`
- `cnew = costfun(xnew, unew)`
"""
@views function forward_pass(traj_new, x0,u,x,α,f,costfun,lims,diff)
    n         = size(x0,1)
    m,N       = size(u)
    xnew      = Array{eltype(x0)}(undef,n,N)
    xnew[:,1] = x0
    unew      = copy(u)
    cnew      = zeros(N)
    for i = 1:N
        if !isempty(traj_new)
            unew[:,i] .+= traj_new.k[:,i]*α
            dx = diff(xnew[:,i], x[:,i])
            unew[:,i] .+= traj_new.K[:,:,i]*dx
        end
        if !isempty(lims)
            unew[:,i] = clamp.(unew[:,i],lims[:,1], lims[:,2])
        end
        xnewi  = f(xnew[:,i], unew[:,i], i)
        if i < N
            xnew[:,i+1] = xnewi
        end
    end
    cnew = costfun(xnew, unew)

    return xnew,unew,cnew
end

@views function forward_covariance(model, x, u, traj)
    fx,fu,fxx,fxu,fuu = df(model, x, u)
    n        = size(fx,1)
    m        = size(fu,2)
    N        = size(fx,3)
    R1       = covariance(model,x,u) # Simple empirical prediction covariance
    Σ0       = R1 # TODO: I was lazy here
    ix = 1:n
    iu = n+1:n+m
    sigmanew = Array{eltype(x)}(undef,n+m,n+m,N)
    sigmanew[ix,ix,1] = Σ0
    for i = 1:N-1
        K,Σ = traj.K[:,:,i], traj.Σ[:,:,i]
        sigmanew[ix,ix,i+1] = fx[:,:,i]*sigmanew[ix,ix,i]*fx[:,:,i]' + R1 # Iterate dLyap forward
        sigmanew[iu,ix,i] = K*sigmanew[ix,ix,i]
        sigmanew[ix,iu,i] = sigmanew[ix,ix,i]*K'
        sigmanew[iu,iu,i] = K*sigmanew[ix,ix,i]*K' + Σ
    end
    sigmanew
end


using Pkg
Pkg.activate(joinpath(@__DIR__, ".."))
Pkg.instantiate()

# ## setup
using Dojo
using IterativeLQR
using LinearAlgebra

# ## system
gravity=-9.81
timestep = 0.1
env = get_environment(:cartpole,
    representation=:minimal,
    timestep=timestep,
    gravity=gravity);

## mujoco_inertia!(env.mechanism)

# ## visualizer
open(env.vis)

# ## dimensions
n = env.num_states
m = env.num_inputs

# ## states
z1 = maximal_to_minimal(env.mechanism, Dojo.cartpole_nominal_max())
zT = maximal_to_minimal(env.mechanism, Dojo.cartpole_goal_max())

# ## horizon
T = 26

# ## model
dyn = IterativeLQR.Dynamics(
    (y, x, u, w) -> dynamics(y, env, x, u, w),
    (dx, x, u, w) -> dynamics_jacobian_state(dx, env, x, u, w),
    (du, x, u, w) -> dynamics_jacobian_input(du, env, x, u, w),
    n, n, m)
model = [dyn for t = 1:T-1]

# ## rollout
ū = [t < 5 ? 1.0 * rand(m) : (t < 10 ? -1.0 * rand(m) : zeros(m)) for t = 1:T-1]
x̄ = IterativeLQR.rollout(model, z1, ū)
visualize(env, x̄)

# ## objective
ot = (x, u, w) -> transpose(x - zT) * Diagonal(1.0e-1 * ones(n)) * (x - zT) + transpose(u) * Diagonal(1.0e-3 * ones(m)) * u
oT = (x, u, w) -> transpose(x - zT) * Diagonal(1.0e-1 * ones(n)) * (x - zT)

ct = IterativeLQR.Cost(ot, n, m)
cT = IterativeLQR.Cost(oT, n, 0)
obj = [[ct for t = 1:T-1]..., cT]

# ## constraints
function goal(x, u, w)
    x - zT
end

cont = IterativeLQR.Constraint()
conT = IterativeLQR.Constraint(goal, n, 0)
cons = [[cont for t = 1:T-1]..., conT]

# ## solver
s = IterativeLQR.solver(model, obj, cons,
    opts=IterativeLQR.Options(
        max_al_iter=10,
        verbose=true))
IterativeLQR.initialize_controls!(s, ū)
IterativeLQR.initialize_states!(s, x̄)

# ## solve
@time IterativeLQR.solve!(s)

# ## solution
z_sol, u_sol = IterativeLQR.get_trajectory(s)
@show IterativeLQR.eval_obj(s.m_data.obj.costs, s.m_data.x, s.m_data.u, s.m_data.w)
@show s.s_data.iter[1]
@show norm(goal(s.m_data.x[T], zeros(0), zeros(0)), Inf)

# ## visualize
visualize(env, z_sol)


using ADCME
using PyPlot

db = Database("simulation.db")
c = execute(db, "select size, forward, backward from acoustic where level = 'cpu_tensor' and size <= 600 order by size")
c = collect(c)
s = [x[1] for x in c]
t1 = [x[2] for x in c]
t2 = [x[3] for x in c]
close("all")
loglog(s, t1, "*-", label = "Forward Computation")
loglog(s, t2, "*-", label = "Gradient Backpropagation")
savefig("result.png")


"""
    brim(N::NT, L::Dict{E,Int64}) where {NT<:AbstractEcologicalNetwork,E}

Uses BRIM to optimize the modularity of an ecological network. The `L` argument
is a dictionary mapping every species in the network to its module. This
function returns a tuple of the network and its module assignment.

#### References

- Barber, M.J., 2007. Modularity and community detection in bipartite networks.
  Phys. Rev. E 76, 066102. https://doi.org/10.1103/PhysRevE.76.066102

- Newman, M.E., 2006. Modularity and community structure in networks.
  Proceedings of the national academy of sciences 103, 8577–8582.

- Thébault, E., 2013. Identifying compartments in presence–absence matrices and
  bipartite networks: insights into modularity measures. Journal of Biogeography
  40, 759–768. https://doi.org/10.1111/jbi.12015
"""
function brim(N::NT, L::Dict{E,Int64}) where {NT<:AbstractEcologicalNetwork,E}
  for s in species(N)
    @assert haskey(L, s)
  end
  EcologicalNetworks.tidy_modules!(L)

  old_Q = Q(N, L)
  new_Q = old_Q+0.00001

  m = links(N)

  c = length(unique(collect(values(L))))

  R = zeros(Int64, (richness(N; dims=1),c))
  for si in eachindex(species(N; dims=1))
    R[si,L[species(N; dims=1)[si]]] = 1
  end

  T = zeros(Int64, (richness(N; dims=2),c))
  for sj in eachindex(species(N; dims=2))
    T[sj,L[species(N; dims=2)[sj]]] = 1
  end

  d_in, d_out = degree(N; dims=2), degree(N; dims=1)
  s_d_in = [d_in[s] for s in species(N; dims=2)]
  s_d_out = [d_out[s] for s in species(N; dims=1)]
  P = kron(s_d_out, s_d_in')./m
  B = N.edges .- P

  while old_Q < new_Q

    t_tilde = B*T
    R = map(Int64, t_tilde .== maximum(t_tilde; dims=2))
    r_tilde = B'*R
    T = map(Int64, r_tilde .== maximum(r_tilde; dims=2))
    S = vcat(R, T)
    L = Dict(zip(species(N), vec(mapslices(r -> StatsBase.sample(findall(x->x!=0, r)), S; dims=2))))
    EcologicalNetworks.tidy_modules!(L)
    old_Q = new_Q
    new_Q = Q(N,L)

  end

  EcologicalNetworks.tidy_modules!(L)
  return (N, L)

end


function fva(S, b, lb, ub, ridxs = eachindex(lb);
        check_obj::Union{Nothing, Int} = nothing,
        check_obj_atol::Real = 1e-4,
        verbose = true,
        batchlen::Int = 50,
        on_non_optimal_sol::Function = (idx, lp_model) -> error("FBA failed, non OPTIMAL solution returned!!!"),
        nths = nthreads(),
        solver = Clp.Optimizer
    )

    ridxs = collect(ridxs)

    # bounds
    fvalb = lb[ridxs]
    fvaub = ub[ridxs]

    verbose && (prog = Progress(length(ridxs); desc = "Doing FVA (-t$nths)  "))
    ii_ch = Channel{Int}(nths) do ch_
        for ii in eachindex(ridxs)
            put!(ch_, ii)
            verbose && next!(prog)
        end
    end

    @threads for _ in 1:nths

        lp_model = build_lp_model(S, b, lb, ub, solver)
        x = _get_vars(lp_model)

        for ii in ii_ch
            ridx = ridxs[ii]

            # optimize max
            JuMP.@objective(lp_model, MIN_SENSE, x[ridx])
            JuMP.optimize!(lp_model)
            status = JuMP.termination_status(lp_model)
            if status == JuMP.MOI.OPTIMAL
                fvalb[ii] = JuMP.objective_value(lp_model)
                else; on_non_optimal_sol(ridx, lp_model)
            end

            # optimize min
            JuMP.@objective(lp_model, MAX_SENSE, x[ridx])
            JuMP.optimize!(lp_model)
            status = JuMP.termination_status(lp_model)
            if status == JuMP.MOI.OPTIMAL
                fvaub[ii] = max(JuMP.objective_value(lp_model), fvalb[ii])
                else; on_non_optimal_sol(ridx, lp_model)
            end
        end
    end
    verbose && finish!(prog)

    # Check bounds
    if !isnothing(check_obj)
        return MetLP.check_newbounds(S, b, lb, ub, fvalb, fvaub,
            check_obj, ridxs; check_obj_atol, verbose, batchlen
        )
    end

    return fvalb, fvaub
end

function fva(model::MetNet, iders = eachindex(model.lb);
        check_obj = nothing, kwargs...
    )
    obj_idx = isnothing(check_obj) ? nothing : rxnindex(model, check_obj)
    idxs = [rxnindex(model, idx) for idx in iders]
    model_fields = _extract_dense(model, [:S, :b, :lb, :ub])
    return fva(model_fields..., idxs; check_obj = obj_idx, kwargs...);
end

fva(model::MetNet, ider::IDER_TYPE; kwargs...) = fva(model, [ider]; kwargs...)


function ψtcNewtonSolve!(F, ∇ₓF, x, nrm, τstop)
  # TODO x_hist = x
    @printf "Solving F(x) = 0 (Using Pseudotransient Continuation)\n"
    nit_NS = 200
    i = 0
    s = copy(x)
    FNorm = nrm(F(x))
    FNorm₀ = copy(FNorm)
    xNorm = nrm(x)
    τ = xNorm / FNorm
    sNorm = nrm(s)
  # evaluate F at the initial iterate
  # compute the stop tolerance
  # main iteration loop
    @printf "  │ i = %d: |F(x)| = %.2g\n" i FNorm
    while (i == 0) || ((τ < τstop) && (i < nit_NS))
        i += 1
        @printf "  │ i = %d: " i
    # compute the ψtc step
        τ = xNorm / FNorm
        s .= (I / τ - ∇ₓF(x)) \ F(x)
        x .= x + s # ψtc step
        FNorm = nrm(F(x))
        xNorm = nrm(x)
        sNorm = nrm(s)
        @printf "|F(x)| = %.2g" FNorm
        @printf "; |x| = %.2g" xNorm
        @printf ", |dx|/|x| = %.2g\n" sNorm / xNorm
    end
  # on failure, set the error flag
    if (nrm(x) / nrm(F(x)) < τstop)
        @printf "  └────> Newton has reached max iterations, "
        @printf "|x|/|F(x)| = %.2g years\n" ((xNorm / FNorm) / (365 * 24 * 60 * 60))
    else
        @printf "  └────> Newton has converged, "
        @printf "|x|/|F(x)| = %.2g years\n" ((xNorm / FNorm) / (365 * 24 * 60 * 60))
    end
  # TODO return sol, x_hist
    return x
end

export ψtcNewtonSolve!


##################################################
# Useful constants
##################################################

using SpecialFunctions

# Numerical Constants
# Calculation done with with Julia

"""
    PI

`π` constant, set equal to `3.1415926535897`.
Computed by `Julia`.
"""
const PI = 3.1415926535897

##################################################
# Constants of the problem
##################################################

"""
    nbGlobularCluster

Number of stars within the Plummer cluster.
"""
const nbGlobularCluster = 10^5

"""
    parameterCoulomb

Coulomb logarithm of the cluster.
Set as `ln Λ = ln(0.11 N)` where `N` is the number of stars in the cluster.

See `Heggie et Hut (2003)`.
"""
const logCoulomb = log(0.11*nbGlobularCluster)

"""
    _G

Newton's gravitational constant.
"""
const _G = 1.0

"""
    _M

Total mass of the globular cluster.
"""
const _M = 1.0

"""
    _b

Plummer radius of the globular cluster.
"""
const _b = 1.0

##################################################
# Dimensional units (Plummer potential)
##################################################

"""
    _v0

Velocity unit of the Plummer cluster.
"""
const _v0 = sqrt(_G*_M/_b)

"""
    _E0

Energy unit of the Plummer cluster.
"""
const _E0 = -_G*_M/_b

"""
    _L0

Action unit of the Plummer cluster.
"""
const _L0 = sqrt(_G*_M*_b)

"""
    _F0

Distribution function unit of the Plummer cluster.
"""
const _F0 = (_G*_M*_b)^(-3/2)

"""
    _Omega0

Frequency unit of the Plummer cluster.
"""
const _Omega0 = sqrt(_G*_M/_b^3)

"""
    _rho0

Density unit of the Plummer cluster.
"""
const _rho0 = _M/(4*PI*_b^3/3)

##################################################
# Constant used to compute H(x,q) and the DF
##################################################

"""
    GAMMA_ca

`gamma(c-a) = gamma(9/2 - qCalc)`.
"""
const GAMMA_ca = gamma(9/2 - qCalc)

"""
    GAMMA_db

`gamma(d-b) = gamma(1 - qCalc/2)`.
"""
const GAMMA_db = gamma(1 - qCalc/2)

"""
    GAMMA_bc

`gamma(b-c) = gamma(9/2 - qCalc/2)`.
"""
const GAMMA_bc = gamma(9/2 - qCalc/2)

"""
    GAMMA_6q

`gamma(6-q) = gamma(6-qCalc)`.
"""
const GAMMA_6q = gamma(6-qCalc)


#   Copyright (c) 2020-2022 Velexi Corporation
#
#   Licensed under the Apache License, Version 2.0 (the "License");
#   you may not use this file except in compliance with the License.
#   You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#   Unless required by applicable law or agreed to in writing, software
#   distributed under the License is distributed on an "AS IS" BASIS,
#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#   See the License for the specific language governing permissions and
#   limitations under the License.

"""
Unit tests for isapprox(x, y) and !isapprox(x, y)
"""

# --- Imports

# Standard library
import InteractiveUtils.subtypes
using Test

# GeometricAlgebra.jl
using GeometricAlgebra

# --- File inclusions

# Test utilities
include("test_utils.jl")

# --- Tests

# ------ B::Blade

@testset "isapprox(B::Blade, C::Blade)" begin
    # Preparations
    vectors = [3 3; 4 4; 0 1]
    B = Blade(vectors)

    # dim(B) == dim(C), grade(B) == grade(C), volume(B) ≈ volume(C)
    # basis(B) ≈ basis(C)
    for precision_type1 in subtypes(AbstractFloat)
        for precision_type2 in subtypes(AbstractFloat)
            B = Blade(convert(Array{precision_type1}, vectors))
            C = Blade(convert(Array{precision_type2}, vectors))
            @test B ≈ C
        end
    end

    # dim(B) != dim(C)
    for precision_type1 in subtypes(AbstractFloat)
        vectors2 = [3 3; 4 4; 0 1; 0 1]
        for precision_type2 in subtypes(AbstractFloat)
            B = Blade(convert(Array{precision_type1}, vectors))
            C = Blade(convert(Array{precision_type2}, vectors2))
            @test B ≉ C
        end
    end

    # grade(B) != grade(C)
    for precision_type1 in subtypes(AbstractFloat)
        vectors2 = [3 3 3; 4 4 4; 0 1 0; 0 0 1]
        for precision_type2 in subtypes(AbstractFloat)
            B = Blade(convert(Array{precision_type1}, vectors))
            C = Blade(convert(Array{precision_type2}, vectors2))
            @test B ≉ C
        end
    end

    # volume(B) ≉ volume(C)
    for precision_type1 in subtypes(AbstractFloat)
        for precision_type2 in subtypes(AbstractFloat)
            B = Blade(convert(Array{precision_type1}, vectors))
            C = Blade(convert(Array{precision_type2}, vectors),
                      volume=2*volume(B))
            @test B ≉ C
        end
    end

    # basis(B) ≉ basis(C)
    test_volume = 5.0
    for precision_type1 in subtypes(AbstractFloat)
        vectors2 = [3 4; 4 10; 0 1]
        for precision_type2 in subtypes(AbstractFloat)
            B = Blade(convert(Array{precision_type1}, vectors),
                      volume=test_volume)
            C = Blade(convert(Array{precision_type2}, vectors2),
                      volume=test_volume)
            @test B ≉ C
        end
    end

    # B and C have opposite orientations
    test_volume = 5.0
    for precision_type1 in subtypes(AbstractFloat)
        for precision_type2 in subtypes(AbstractFloat)
            B = Blade(convert(Array{precision_type1}, vectors),
                      volume=test_volume)
            C = Blade(convert(Array{precision_type2}, vectors),
                      volume=-test_volume)
            @test B ≉ C
        end
    end
end

@testset "!isapprox(B::Blade, C::Pseudoscalar)" begin
    test_vectors = [3 3; 4 4; 0 1]

    test_dim = 3
    test_value = 5

    for precision_type in subtypes(AbstractFloat)
        B = Blade{precision_type}(test_vectors)
        C = Pseudoscalar{precision_type}(test_dim, test_value)
        @test B ≉ C
    end
end

@testset "!isapprox(B::Blade, C::Scalar)" begin
    test_vectors = [3 3; 4 4; 0 1]
    test_value = 5

    for precision_type in subtypes(AbstractFloat)
        B = Blade{precision_type}(test_vectors)
        C = Scalar{precision_type}(test_value)
        @test B ≉ C
    end
end

@testset "!isapprox(B::Blade, C::One)" begin
    test_vectors = [3 3; 4 4; 0 1]
    for precision_type in subtypes(AbstractFloat)
        B = Blade{precision_type}(test_vectors)
        C = One{precision_type}()
        @test B ≉ C
    end
end

@testset "!isapprox(B::Blade, C::Zero)" begin
    test_vectors = [3 3; 4 4; 0 1]
    for precision_type in subtypes(AbstractFloat)
        B = Blade{precision_type}(test_vectors)
        C = Zero{precision_type}()
        @test B ≉ C
    end
end

@testset "!isapprox(B::Blade, C::Real)" begin
    test_vectors = [3 3; 4 4; 0 1]
    test_value = 5
    for precision_type in subtypes(AbstractFloat)
        B = Blade{precision_type}(test_vectors)
        C = precision_type(test_value)
        @test B ≉ C
    end
end

@testset "isapprox(B::Blade, C::Vector)" begin
    # --- Preparations

    test_vector = [3; 4; 1]

    # B ≉ C
    test_vectors = [3 3; 4 4; 0 1]
    for precision_type in subtypes(AbstractFloat)
        B = Blade{precision_type}(test_vectors)
        C = Vector{precision_type}(test_vector)
        @test B ≉ C
    end

    test_vectors = [3; 4; 10]
    for precision_type in subtypes(AbstractFloat)
        B = Blade{precision_type}(test_vectors)
        C = Vector{precision_type}(test_vector)
        @test B ≉ C
    end

    # B ≈ C
    for precision_type1 in subtypes(AbstractFloat)
        for precision_type2 in subtypes(AbstractFloat)
            B = Blade{precision_type1}(test_vector)
            C = Vector{precision_type2}(test_vector)
            @test B ≈ C
        end
    end
end

# ------ B::Pseudoscalar

@testset "!isapprox(B::Pseudoscalar, C::Blade)" begin
    test_dim = 3
    test_value = 5

    test_vectors = [3 3; 4 4; 0 1]

    for precision_type in subtypes(AbstractFloat)
        B = Pseudoscalar{precision_type}(test_dim, test_value)
        C = Blade{precision_type}(test_vectors)
        @test B ≉ C
    end
end

@testset "isapprox(B::Pseudoscalar, C::Pseudoscalar)" begin
    # Preparations
    dim = 10

    test_value = get_random_value(1)  # add 1 to keep value away from 0

    # dim(B) == dim(C), value(B) == value(C)
    for precision_type1 in subtypes(AbstractFloat)
        for precision_type2 in subtypes(AbstractFloat)
            B = Pseudoscalar(dim, precision_type1(test_value))
            C = Pseudoscalar(dim, precision_type2(test_value))
            @test B ≈ C
        end
    end

    # dim(B) != dim(C), value(B) == value(C)
    for precision_type1 in subtypes(AbstractFloat)
        for precision_type2 in subtypes(AbstractFloat)
            B = Pseudoscalar(dim, precision_type1(test_value))
            C = Pseudoscalar(dim + 1, precision_type2(test_value))
            @test B ≉ C
        end
    end

    # dim(B) == dim(C), value(B) != value(C)
    for precision_type1 in subtypes(AbstractFloat)
        for precision_type2 in subtypes(AbstractFloat)
            B = Pseudoscalar(dim, precision_type1(test_value))
            C = Pseudoscalar(dim, precision_type2(test_value) + 1)
            @test B ≉ C
        end
    end
end

@testset "isapprox(B::Pseudoscalar, C::Scalar)" begin
    test_dim = 10
    test_value = 5

    for precision_type in subtypes(AbstractFloat)
        B = Pseudoscalar{precision_type}(test_dim, test_value)
        C = Scalar{precision_type}(test_value)
        @test B ≉ C
    end
end

@testset "isapprox(B::Pseudoscalar, C::One)" begin
    test_dim = 10
    test_value = 5

    for precision_type in subtypes(AbstractFloat)
        B = Pseudoscalar{precision_type}(test_dim, test_value)
        C = One{precision_type}()
        @test B ≉ C
    end
end

@testset "isapprox(B::Pseudoscalar, C::Zero)" begin
    test_dim = 10
    test_value = 5

    for precision_type in subtypes(AbstractFloat)
        B = Pseudoscalar{precision_type}(test_dim, test_value)
        C = Zero{precision_type}()
        @test B ≉ C
    end
end

@testset "isapprox(B::Pseudoscalar, C::Real)" begin
    test_dim = 10
    test_value = 5

    for precision_type in subtypes(AbstractFloat)
        B = Pseudoscalar{precision_type}(test_dim, test_value)
        C = precision_type(test_value)
        @test B ≉ C
    end
end

@testset "isapprox(B::Pseudoscalar, C::Vector)" begin
    test_value = 5

    test_dim_1 = 3
    test_vector_1 = [3; 4; 1]

    test_dim_2 = 1
    test_vector_2 = [test_value]

    for precision_type1 in subtypes(AbstractFloat)
        # B ≉ C
        B = Pseudoscalar{precision_type1}(test_dim_1, test_value)
        C = Vector{precision_type1}(test_vector_1)
        @test B ≉ C

        # B ≈ C
        B = Pseudoscalar{precision_type1}(test_dim_2, test_value)
        for precision_type2 in subtypes(AbstractFloat)
            C = Vector{precision_type2}(test_vector_2)
            @test B ≈ C
        end
    end
end

# ------ B::Scalar

@testset "!isapprox(B::Scalar, C::Blade)" begin
    test_value = 5
    test_vectors = [3 3; 4 4; 0 1]
    for precision_type in subtypes(AbstractFloat)
        B = Scalar{precision_type}(test_value)
        C = Blade{precision_type}(test_vectors)
        @test B ≉ C
    end
end

@testset "isapprox(B::Scalar, C::Pseudoscalar)" begin
    test_value = 5
    test_dim = 10
    for precision_type in subtypes(AbstractFloat)
        B = Scalar{precision_type}(test_value)
        C = Pseudoscalar{precision_type}(test_dim, test_value)
        @test B ≉ C
    end
end

@testset "isapprox(B::Scalar, C::Scalar)" begin
    # Preparations

    test_value = get_random_value(2)  # add 2 to keep value away from 0 and 1

    # B::Scalar, C::Scalar
    for precision_type1 in subtypes(AbstractFloat)
        for precision_type2 in subtypes(AbstractFloat)
            # B ≈ C
            B = Scalar(precision_type1(test_value))
            C = Scalar(precision_type2(test_value))
            @test B ≈ C

            # B ≉ C
            B = Scalar(precision_type1(test_value))
            C = Scalar(precision_type2(test_value + 1))
            @test B ≉ C
        end
    end
end

@testset "isapprox(B::Scalar, C::One)" begin
    for precision_type in subtypes(AbstractFloat)
        C = One{precision_type}()

        # B ≈ C
        test_value = 1 - (√eps(precision_type)) / 2
        B = Scalar{precision_type}(test_value)
        @test B ≈ C

        # B ≉ C
        test_value = 5
        B = Scalar{precision_type}(test_value)
        @test B ≉ C
    end
end

@testset "isapprox(B::Scalar, C::Zero)" begin
    test_value = 5
    for precision_type in subtypes(AbstractFloat)
        B = Scalar{precision_type}(test_value)
        C = Zero{precision_type}()
        @test B ≉ C

        # B ≈ C is not being tested as isapprox between a Scalar and a
        # Zero can never return true with the default atol and rtol value
    end
end

@testset "isapprox(B::Scalar, C::Real)" begin
    test_value = 5
    for precision_type1 in subtypes(AbstractFloat)
        B = Scalar{precision_type1}(test_value)

        # B ≉ C
        C = precision_type1(test_value + 1)
        @test B ≉ C

        # B ≈ C
        for precision_type2 in subtypes(AbstractFloat)
            C = precision_type2(test_value)
            @test B ≈ C
        end
    end
end

@testset "isapprox(B::Scalar, C::Vector)" begin
    test_value = 5
    test_vector = [3; 4; 1]
    for precision_type in subtypes(AbstractFloat)
        B = Scalar{precision_type}(test_value)
        C = Vector{precision_type}(test_vector)
        @test B ≉ C
    end
end

# ------ B::One

@testset "!isapprox(B::One, C::Blade)" begin
    test_vectors = [3 3; 4 4; 0 1]
    for precision_type in subtypes(AbstractFloat)
        B = One{precision_type}()
        C = Blade{precision_type}(test_vectors)
        @test B ≉ C
    end
end

@testset "isapprox(B::One, C::Pseudoscalar)" begin
    test_dim = 10
    test_value = 5
    for precision_type in subtypes(AbstractFloat)
        B = One{precision_type}()
        C = Pseudoscalar{precision_type}(test_dim, test_value)
        @test B ≉ C
    end
end

@testset "isapprox(B::One, C::Scalar)" begin
    for precision_type in subtypes(AbstractFloat)
        B = One{precision_type}()

        # B ≈ C
        test_value = 1 - (√eps(precision_type)) / 2
        C = Scalar{precision_type}(test_value)
        @test B ≈ C

        # B ≉ C
        test_value = 5
        C = Scalar{precision_type}(test_value)
        @test B ≉ C
    end
end

@testset "isapprox(B::One, C::One)" begin
    for precision_type1 in subtypes(AbstractFloat)
        for precision_type2 in subtypes(AbstractFloat)
            B = One{precision_type1}()
            C = One{precision_type2}()
            @test B ≈ C
        end
    end
end

@testset "isapprox(B::One, C::Zero)" begin
    for precision_type in subtypes(AbstractFloat)
        B = One{precision_type}()
        C = Zero{precision_type}()
        @test B ≉ C
    end
end

@testset "isapprox(B::One, C::Real)" begin
    test_value = 5
    for precision_type1 in subtypes(AbstractFloat)
        B = One{precision_type1}()

        # B ≉ C
        C = precision_type1(test_value)
        @test B ≉ C

        # B ≈ C
        for precision_type2 in subtypes(AbstractFloat)
            C = precision_type2(1)
            @test B ≈ C
        end
    end
end

@testset "isapprox(B::One, C::Vector)" begin
    test_vector = [3; 4; 1]
    for precision_type in subtypes(AbstractFloat)
        B = One{precision_type}()
        C = Vector{precision_type}(test_vector)
        @test B ≉ C
    end
end

# ------ B::Zero

@testset "!isapprox(B::Zero, C::Blade)" begin
    test_vectors = [3 3; 4 4; 0 1]
    for precision_type in subtypes(AbstractFloat)
        B = Zero{precision_type}()
        C = Blade{precision_type}(test_vectors)
        @test B ≉ C
    end
end

@testset "isapprox(B::Zero, C::Pseudoscalar)" begin
    test_dim = 10
    test_value = 5
    for precision_type in subtypes(AbstractFloat)
        B = Zero{precision_type}()
        C = Pseudoscalar{precision_type}(test_dim, test_value)
        @test B ≉ C
    end
end

@testset "isapprox(B::Zero, C::Scalar)" begin
    test_value = 5
    for precision_type in subtypes(AbstractFloat)
        B = Zero{precision_type}()
        C = Scalar{precision_type}(test_value)
        @test B ≉ C

        # B ≈ C is not being tested as isapprox between a Zero and a
        # Scalar can never return true with the default atol and rtol value
    end
end

@testset "isapprox(B::Zero, C::One)" begin
    for precision_type in subtypes(AbstractFloat)
        B = Zero{precision_type}()
        C = One{precision_type}()
        @test B ≉ C
    end
end

@testset "isapprox(B::Zero, C::Zero)" begin
    for precision_type1 in subtypes(AbstractFloat)
        for precision_type2 in subtypes(AbstractFloat)
            B = Zero{precision_type1}()
            C = Zero{precision_type2}()
            @test B ≈ C
        end
    end
end

@testset "isapprox(B::Zero, C::Real)" begin
    test_value = 5
    for precision_type1 in subtypes(AbstractFloat)
        B = Zero{precision_type1}()

        # B ≉ C
        C = precision_type1(test_value)
        @test B ≉ C

        # B ≈ C
        for precision_type2 in subtypes(AbstractFloat)
            C = precision_type2(0)
            @test B ≈ C
        end
    end
end

@testset "isapprox(B::Zero, C::Vector)" begin
    # B ≉ C
    test_vector = [3; 4; 1]
    for precision_type in subtypes(AbstractFloat)
        B = Zero{precision_type}()
        C = Vector{precision_type}(test_vector)
        @test B ≉ C
    end

    test_vector = [0.001; 0.002; 0.003]
    for precision_type in subtypes(AbstractFloat)
        B = Zero{precision_type}()
        C = Vector{precision_type}(test_vector)
        @test !isapprox(B, C, atol=0.001)
    end

    # B ≈ C
    test_vector = [0; 0; 0]
    for precision_type in subtypes(AbstractFloat)
        B = Zero{precision_type}()
        C = Vector{precision_type}(test_vector)
        @test B ≈ C
    end

    test_vector = [0.001; 0.002; 0.003]
    for precision_type in subtypes(AbstractFloat)
        B = Zero{precision_type}()
        C = Vector{precision_type}(test_vector)
        @test isapprox(B, C, atol=0.01)
    end
end

# ------ B::Real

@testset "!isapprox(B::Real, C::Blade)" begin
    test_value = 5
    test_vectors = [3 3; 4 4; 0 1]
    for precision_type in subtypes(AbstractFloat)
        B = precision_type(test_value)
        C = Blade{precision_type}(test_vectors)
        @test B ≉ C
    end
end

@testset "isapprox(B::Real, C::Pseudoscalar)" begin
    test_value = 5
    test_dim = 10
    for precision_type in subtypes(AbstractFloat)
        B = precision_type(test_value)
        C = Pseudoscalar{precision_type}(test_dim, test_value)
        @test B ≉ C
    end
end

@testset "isapprox(B::Real, C::Scalar)" begin
    test_value = 5
    for precision_type1 in subtypes(AbstractFloat)
        B = precision_type1(test_value)

        # B ≉ C
        C = Scalar{precision_type1}(test_value + 1)
        @test B ≉ C

        # B ≈ C
        for precision_type2 in subtypes(AbstractFloat)
            C = Scalar{precision_type2}(test_value)
            @test B ≈ C
        end
    end
end

@testset "isapprox(B::Real, C::One)" begin
    test_value = 5
    for precision_type1 in subtypes(AbstractFloat)
        # B ≉ C
        B = precision_type1(test_value)
        C = One{precision_type1}()
        @test B ≉ C

        # B ≈ C
        B = precision_type1(1)
        for precision_type2 in subtypes(AbstractFloat)
            C = One{precision_type2}()
            @test B ≈ C
        end
    end
end

@testset "isapprox(B::Real, C::Zero)" begin
    test_value = 5
    for precision_type1 in subtypes(AbstractFloat)
        # B ≉ C
        B = precision_type1(test_value)
        C = Zero{precision_type1}()
        @test B ≉ C

        # B ≈ C
        B = precision_type1(0)
        for precision_type2 in subtypes(AbstractFloat)
            C = Zero{precision_type2}()
            @test B ≈ C
        end
    end
end

# ------ B::Vector

@testset "isapprox(B::Vector, C::Blade)" begin
    # --- Preparations

    test_vector = [3; 4; 1]

    # B ≉ C
    test_vectors = [3 3; 4 4; 0 1]
    for precision_type in subtypes(AbstractFloat)
        B = Vector{precision_type}(test_vector)
        C = Blade{precision_type}(test_vectors)
        @test B ≉ C
    end

    test_vectors = [3; 4; 10]
    for precision_type in subtypes(AbstractFloat)
        B = Vector{precision_type}(test_vector)
        C = Blade{precision_type}(test_vectors)
        @test B ≉ C
    end

    # B ≈ C
    for precision_type1 in subtypes(AbstractFloat)
        for precision_type2 in subtypes(AbstractFloat)
            B = Vector{precision_type1}(test_vector)
            C = Blade{precision_type2}(test_vector)
            @test B ≈ C
        end
    end
end

@testset "isapprox(B::Vector, C::Pseudoscalar)" begin
    test_value = 5

    test_dim_1 = 3
    test_vector_1 = [3; 4; 1]

    test_dim_2 = 1
    test_vector_2 = [test_value]

    for precision_type1 in subtypes(AbstractFloat)
        # B ≉ C
        B = Vector{precision_type1}(test_vector_1)
        C = Pseudoscalar{precision_type1}(test_dim_1, test_value)
        @test B ≉ C

        # B ≈ C
        B = Vector{precision_type1}(test_vector_2)
        for precision_type2 in subtypes(AbstractFloat)
            C = Pseudoscalar{precision_type2}(test_dim_2, test_value)
            @test B ≈ C
        end
    end
end

@testset "isapprox(B::Vector, C::Scalar)" begin
    test_vector = [3; 4; 1]
    test_value = 5
    for precision_type in subtypes(AbstractFloat)
        B = Vector{precision_type}(test_vector)
        C = Scalar{precision_type}(test_value)
        @test B ≉ C
    end
end

@testset "isapprox(B::Vector, C::One)" begin
    test_vector = [3; 4; 1]
    for precision_type in subtypes(AbstractFloat)
        B = Vector{precision_type}(test_vector)
        C = One{precision_type}()
        @test B ≉ C
    end
end

@testset "isapprox(B::Vector, C::Zero)" begin
    # B ≉ C
    test_vector = [3; 4; 1]
    for precision_type in subtypes(AbstractFloat)
        B = Vector{precision_type}(test_vector)
        C = Zero{precision_type}()
        @test B ≉ C
    end

    test_vector = [0.001; 0.002; 0.003]
    for precision_type in subtypes(AbstractFloat)
        B = Vector{precision_type}(test_vector)
        C = Zero{precision_type}()
        @test !isapprox(B, C, atol=0.001)
    end

    # B ≈ C
    test_vector = [0; 0; 0]
    for precision_type in subtypes(AbstractFloat)
        B = Vector{precision_type}(test_vector)
        C = Zero{precision_type}()
        @test B ≈ C
    end

    test_vector = [0.001; 0.002; 0.003]
    for precision_type in subtypes(AbstractFloat)
        B = Vector{precision_type}(test_vector)
        C = Zero{precision_type}()
        @test isapprox(B, C, atol=0.01)
    end
end


export ess, actime

# Effective sample size (ESS)
actypes = (:bm, :imse, :ipse)

function ess(c::MCMCChain, pars::Ranges=1:ncol(c.samples); vtype::Symbol=:imse, args...)
  @assert in(vtype, actypes) "Unknown ESS type $vtype"

  return nrow(c.samples)*var(c, pars; vtype=:iid)./var(c, pars; vtype=vtype, args...)
end

ess(c::MCMCChain, par::Real; vtype::Symbol=:imse, args...) = ess(c, par:par; vtype=vtype, args...)

# Integrated autocorrelation time
function actime(c::MCMCChain, pars::Ranges=1:ncol(c.samples); vtype::Symbol=:imse, args...)
  @assert in(vtype, actypes) "Unknown integrated autocorrelation time type $vtype"

  return var(c, pars; vtype=vtype, args...)./var(c, pars; vtype=:iid)
end

actime(c::MCMCChain, par::Real; vtype::Symbol=:imse, args...) = actime(c, par:par; vtype=vtype, args...)


# Functions for dynamic systems
# Increment of x[ix] through RK4-integration
# dxdh is the array of derivatives, dt is the time increment
function dx_rk4(dxdt, x, ix, dt)
    x = collect(x)
    incr = zeros(length(x))
    k1 = dxdt[ix](x...)
    incr[ix] = dt*k1/2
    k2 = dxdt[ix]((x+incr)...)
    incr[ix] = dt*k2/2
    k3 = dxdt[ix]((x+incr)...)
    incr[ix] = dt*k3
    k4 = dxdt[ix]((x+incr)...)
    dt/6*(k1 + 2k2 + 2k3 + k4)
end
# Integrate dynamical system through n intervals of length dt
function dynamical_system(x0, dxdt, dt, n)
    m = length(x0)
    x = zeros(n,m)
    x[1,:] = collect(x0)
    for t=1:n-1
        dx = zeros(m)
        for ix=1:m
            dx[ix] = dx_rk4(dxdt, x[t,:], ix, dt)
        end
        x[t+1,:] = x[t,:][:] + dx
    end
    x
end
# Rössler system
rossler_eq(a, b, c) = (
    (x,y,z) -> -(y+z),
    (x,y,z) -> x + a*y,
    (x,y,z) -> b + z*(x-c))

a=.25; b=.25; c=4
x0 = zeros(3)
m = zeros(3000,24)
for i=0:11
    x = randn(3)
    rossler_data = dynamical_system(x0, rossler_eq(a,b,c), .05, 4000)
    m[:,2i.+(1:2)] .= rossler_data[1001:end,1:2]
end
writedlm("rossler.txt",m,'\t')


using LinearAlgebra, ForwardDiff, Distributions, StaticArrays
using Convex, SCS, ECOS
using IterativeSolvers
using Plots

# noc
function κ_no(z)
    max.(0.0, z)
end

function Jκ_no(z)
    p = zero(z)
    for (i, pp) in enumerate(z)
        println(pp)
        println(i)
        if pp >= 0.0
            p[i] = 1.0
        end
    end
    return Diagonal(p)
end

# soc
function κ_soc(z)
    z1 = z[1:end-1]
    z2 = z[end]

    z_proj = zero(z)

    if norm(z1) <= z2
        z_proj = copy(z)
    elseif norm(z1) <= -z2
        z_proj = zero(z)
    else
        a = 0.5 * (1.0 + z2 / norm(z1))
        z_proj[1:end-1] = a * z1
        z_proj[end] = a * norm(z1)
    end
    return z_proj
end

function Jκ_soc(z)
    z1 = z[1:end-1]
    z2 = z[end]
    m = length(z)

    if norm(z1) <= z2
        return Diagonal(ones(m))
    elseif norm(z1) <= -z2
        return Diagonal(zeros(m))
    else
        D = zeros(m, m)
        for i = 1:m
            if i < m
                D[i, i] = 0.5 + 0.5 * z2 / norm(z1) - 0.5 * z2 * ((z1[i])^2.0) / norm(z1)^3.0
            else
                D[i, i] = 0.5
            end
            for j = 1:m
                if j > i
                    if j < m
                        D[i, j] = -0.5 * z2 * z1[i] * z1[j] / norm(z1)^3.0
                        D[j, i] = -0.5 * z2 * z1[i] * z1[j] / norm(z1)^3.0
                    elseif j == m
                        D[i, j] = 0.5 * z1[i] / norm(z1)
                        D[j, i] = 0.5 * z1[i] / norm(z1)
                    end
                end
            end
        end
        return D
    end
end

# friction problem
v = 1.0 * ones(2)
y = 1.0 #1.0e-3

n = 3
m = 1
c = [v; 0.0]
A = [0.0; 0.0; 1.0]'
b = copy(y)

θ_friction = [vec(A); b; c]

function r_friction(z)
    x = z[1:3]
    μ = z[3 + 1]
    λ = z[3 + 1 .+ (1:3)]

    [[v; 0.0] - [0.0; 0.0; 1.0] * μ - λ;
     y - [0.0; 0.0; 1.0]' * x;
     x - κ_soc(x - λ)]
end

function r_friction(z, θ)
    x = z[1:3]
    μ = z[3 + 1]
    λ = z[3 + 1 .+ (1:3)]

    A = reshape(θ[1:1 * 3], 1, 3)
    b = θ[1 * 3 .+ (1:1)]
    c = θ[1 * 3 + 1 .+ (1:3)]

    [c - A' * μ - λ;
     b - A * x;
     x - κ_soc(x - λ)]
end

function R_friction(z)
    x = z[1:3]
    μ = z[3 + 1]
    λ = z[3 + 1 .+ (1:3)]
    J = Jκ_soc(x - λ)

    [zeros(3, 3) -[0.0; 0.0; 1.0]'' -Diagonal(ones(3));
     -[0.0; 0.0; 1.0]' zeros(1, 1) zeros(1, 3);
     (Diagonal(ones(3)) - J) zeros(3, 1) J]
end

function solve()
    z = 0.001 * rand(3 + 1 + 3)

    extra_iters = 0

    for i = 1:10
        _F = r_friction(z)
        _J = R_friction(z) #R_friction(z)
        Δ = gmres(_J, 1.0 * _F, abstol = 1.0e-12, maxiter = i + extra_iters)
        # Δ = (_J' * _J + 1.0e-5 * I) \ (_J' * _F)
        iter = 0
        α = 1.0
        while norm(r_friction(z - α * Δ))^2.0 >= (1.0 - 0.001 * α) * norm(_F)^2.0 && α > 1.0e-4
            α = 0.5 * α
            # println("   α = $α")
            iter += 1
            if iter > 100
                @error "line search fail"

                return z
            end
        end

        if α <= 1.0e-4
            extra_iters += 1
        end

        println("iter ($i) - norm: $(norm(r_friction(z)))")

        z .-= α * Δ
    end

    return z
end

z_sol = solve()

@show b_sol = z_sol[1:2]

θ = [vec(A); b; c]
r_friction(z_sol, θ)

_rz(w) = r_friction(w, θ)
_rθ(w) = r_friction(z_sol, w)

Rz = ForwardDiff.jacobian(_rz, z_sol)
Rθ = ForwardDiff.jacobian(_rθ, θ)

dθ = ((Rz' * Rz + 1.0e-5 * I) \ (Rz' * Rθ))[1:2, m * n .+ (1:m+(n-1))]

# impact problem
# rod
struct Particle
      m # mass
      g # gravity

      nq # configuration dimension
end

# jacobian
function jacobian(model, q)
    Diagonal(@SVector ones(model.nq))
end

# mass matrix
function mass_matrix(model)
      m = model.m

      Diagonal(@SVector [m, m, m])
end

# gravity
function gravity(model, q)
      m = model.m
      g = model.g

      @SVector [0.0, 0.0, m * g]
end

# dynamics
function dynamics(model, q1, q2, q3, λ, h)
      nq = model.nq
      SVector{nq}(mass_matrix(model) * (2.0 * q2 - q1 - q3) / h
            - h * gravity(model, q2)
            + h * jacobian(model, q3)' * λ)
end

# create model
model = Particle(1.0, 9.81, 3)

q1 = [0.0; 0.0; 1.0]
q2 = [1.0; 1.0; 1.0]

function r_action(z)
    q3 = z[1:3]
    q3z = z[3]
    λ = z[4]

    [dynamics(model, q1, q2, q3, [0.0; 0.0; λ], 0.1);
     q3z - κ_no(q3z - λ)]
end

function solve()
    z = [q2; 0.001]

    extra_iters = 0

    for i = 1:10
        _F = r_action(z)
        _J = ForwardDiff.jacobian(r_action, z)
        Δ = gmres(_J, 1.0 * _F, abstol = 1.0e-12, maxiter = i + extra_iters)
        # Δ = (_J' * _J + 1.0e-5 * I) \ (_J' * _F)
        iter = 0
        α = 1.0
        while norm(r_action(z - α * Δ))^2.0 >= (1.0 - 0.001 * α) * norm(_F)^2.0 && α > 1.0e-4
            α = 0.5 * α
            # println("   α = $α")
            iter += 1
            if iter > 100
                @error "line search fail"

                return z
            end
        end

        if α <= 1.0e-4
            extra_iters += 1
        end

        println("iter ($i) - norm: $(norm(r_action(z)))")

        z .-= α * Δ
    end

    return z
end

z_sol = solve()

function r_friction(z)
    x = z[1:3]
    μ = z[3 + 1]
    λ = z[3 + 1 .+ (1:3)]

    [[v; 0.0] - [0.0; 0.0; 1.0] * μ - λ;
     y - [0.0; 0.0; 1.0]' * x;
     x - κ_soc(x - λ)]
end

# impact and friction

# function r_impact_friction(z)
#     q3 = z[1:3]
#     q3z = z[3]
#     v = (q3[1:2] - q2[1:2]) ./ 0.1
#     λ = z[4]
#     b = z[5:6]
#     b̄ = z[5:7]
#     μ = z[8]
#     η = z[9:11]
#
#     [dynamics(model, q1, q2, q3, [b; λ], 0.1);
#      q3z - κ_no(q3z - λ);
#
#      [v; 0.0] - [0.0; 0.0; 1.0] * μ - η;
#       λ - [0.0; 0.0; 1.0]' * b̄;
#       b̄ - κ_soc(b̄ - η)]
# end
#
# function R_impact_friction(z)
#     R = ForwardDiff.jacobian(r_impact_friction, z)
#
#     # fix projection
#     b̄ = z[5:7]
#     μ = z[8]
#     η = z[9:11]
#     J = Jκ_soc(b̄ - η)
#
#     R[(end-2):end, 5:11] = [(Diagonal(ones(3)) - J) zeros(3, 1) J]
#
#     return R
# end

function solve(q1, q2, h)
    function r(z)
        q3 = z[1:3]
        ϕ = z[3]
        λ = z[4]
        b = z[5:6]
        b̄ = z[5:7]
        μ = z[8]
        η = z[9:11]

        [dynamics(model, q1, q2, q3, [b; λ], h);
         ϕ - κ_no(ϕ - λ);

         [(q3[1:2] - q2[1:2]) ./ h; 0.0] - [0.0; 0.0; 1.0] * μ - η;
          λ - [0.0; 0.0; 1.0]' * b̄;
          b̄ - κ_soc(b̄ - η)]
    end

    function R(z)
        _R = ForwardDiff.jacobian(r, z)

        # fix projection
        b̄ = z[5:7]
        μ = z[8]
        η = z[9:11]
        J = Jκ_soc(b̄ - η)

        _R[(end-2):end, 5:11] = [(Diagonal(ones(3)) - J) zeros(3, 1) J]

        return _R
    end

    z = 1.0e-6 * rand(11)
    # z = zeros(11)
    z[1:3] = copy(q2)

    extra_iters = 0

    for i = 1:100
        _F = r(z)
        _J = R(z)
        Δ = gmres(_J, 1.0 * _F, abstol = 1.0e-12, maxiter = i + extra_iters)
        # Δ = (_J' * _J + 1.0e-6 * I) \ (_J' * _F)
        iter = 0
        α = 1.0
        while norm(r(z - α * Δ))^2.0 >= (1.0 - 0.001 * α) * norm(_F)^2.0 && α > 1.0e-4
            α = 0.5 * α
            # println("   α = $α")
            iter += 1
            if iter > 100
                @error "line search fail"

                return z
            end
        end

        if α <= 1.0e-4
            extra_iters += 1
        end

        println("iter ($i) - norm: $(norm(r(z)))")

        z .-= α * Δ
    end

    if norm(r(z)) < 1.0e-5
        status = true
    else
        status = false
    end

    return z, status
end

h = 0.01
v1 = [1.0; 1.0; 1.0]
q2 = [1.0; -2.0; 1.0]
q1 = q2 - h * v1

function simulate(q1, q2, T, h)
    println("simulation")
    q = [q1, q2]
    y = [0.0]
    b = [zeros(2)]
    for t = 1:T
        println("   t = $t")
        z_sol, status = solve(q[end-1], q[end], h)

        if !status
            @error "failed step (t = $t)"
            return q, y, b
        else
            push!(q, z_sol[1:3])
            push!(y, z_sol[4])
            push!(b, z_sol[5:6])
        end
    end

    return q, y, b
end

q_sol, y_sol, b_sol = simulate(q1, q2, 100, h)
# q_sol[end]
# q_sol[end-1]
plot(hcat(q_sol...)[3:3, :]', xlabel = "")
plot!(h .* hcat(y_sol...)', xlabel = "", linetype = :steppost)

plot(hcat(q_sol...)[1:2, :]', xlabel = "")
plot!(h * hcat(b_sol...)', linetype = :steppost)