## General purpose functions

#User is expected to pass functions
#x_k(mhatk, hyperparam)  : provides a solution with data mhatk using hyperparam (e.g., p0, alpha)
#c_ik(x)
#These are often collected into arrays of functions
# xs			: [x_k for k = 1:K]
# c_k  			: [c_ik for i = 1:d]
# cs			: (d, K) Matrix with elements c_ik

#General purpose functions
###
function shrink(phat_k, p0, alpha, Nhat_k; alphaMax = 1e10)
	if Nhat_k == 0 || alpha > alphaMax
		return p0
	end
	@. alpha * p0 / (Nhat_k + alpha) + Nhat_k * phat_k / (Nhat_k + alpha)
end

#takes un-noramlized counts
#assumes out is correctly sized no bounds checking!
function _shrink!(mhat_k, p0, alpha, Nhat_k, out; alphaMax = 1e10)
	if Nhat_k == 0 || alpha > alphaMax
		return p0
	end
	for i = 1:length(p0)
		@inbounds out[i] = alpha * p0[i] / (Nhat_k + alpha) + mhat_k[i] / (Nhat_k + alpha)
	end
	out
end

#need to be careful because some problems may have zero data
function get_GM_anchor(mhats)
	Nhats = vec(sum(mhats, dims=1))
	non_zero_indx = Nhats .> 0
	vec(mean(mhats[:, non_zero_indx] ./ Nhats[non_zero_indx]', dims=2))
end

function sim_path(p_k, N::Int)
    mhats_k = zeros(length(p_k))

    if N == 0 #allow simulating zero data points
    	return mhats_k
    end

    for j in rand(Categorical(p_k), N)
        mhats_k[j] += 1
    end
    mhats_k
end

#ps is a matrix of pk
function sim_path(ps, Nhats)
    K = size(ps, 2)
    out = zeros(size(ps))
    for k = 1:K
        out[:, k] = sim_path(ps[:, k], Nhats[k])
    end
    out
end

#empirical bayes moment-matched estimates.
#returns p0, alpha0
#VG this method does not seem to consistenty yield alpha in [0, 1]
function eb_mm_estimates(mhats)
	Nhats = sum(mhats, dims=1)
	p0 = mean(mhats ./ Nhats, dims=2)
	K = size(mhats, 2)

	C0 = sum(mhats.^2)/K
	sq_norm_p0 = norm(p0)^2
	Nbar = mean(Nhats)
	var_n = mean( @.(Nhats^2 - Nhats) )

	alpha0 = C0 - sq_norm_p0 * var_n - Nbar
	@assert alpha0 > 0 "Moment matching fails"
	alpha0 /= (1 - sq_norm_p0) / var_n
	alpha0 = 1 / alpha0 - 1

	return p0, alpha0
end

#solves an approximation to the Apriori MSE equation for alpha
#for large K this is like AlphaOR for MSE
#for now, super lazy and just search a grid.
#returns AlphaEstimate, minimizingAlphaIndex
function mse_estimates(mhats, supps, p0, alpha_grid)
	K = size(mhats, 2)
	Nhats = sum(mhats, dims=1)
	phats = mhats ./ Nhats
	mu0s = vec(p0' * supps)

	#summary statistics
	muhats = Vector{Float64}(undef, K)
	sigmas = Vector{Float64}(undef, K)
	mse0 = Vector{Float64}(undef, K)
	for k = 1:K
		muhats[k] = dot(phats[:, k], supps[:, k])
		sigmas[k] = dot(phats[:, k], supps[:, k].^2) - muhats[k]^2
		mse0[k] = dot(phats[:, k], @.((supps[:, k] - mu0s[k])^2) )
	end

	#worker function
	function mse(a)
		out = 0.
		for k = 1:K
			out += (Nhats[k] - a^2) / (Nhats[k] + a)^2 * sigmas[k]
			out += ( a / (a + Nhats[k]) )^2 * mse0[k]
		end
		out/K
	end
	out = map(mse, alpha_grid)
	alpha_grid[argmin(out)], argmin(out)
end

#kth element of true perf
function z_k(x_k, c_k, mhat_k, ps_k, lam_k, lamavg, hyper_param)
	x = x_k(mhat_k, hyper_param)
	cs = map(c_ik -> c_ik(x), c_k)
	lam_k/lamavg * dot(ps_k, cs)
end

#actually returns zLOO_k * N * lamavg
function zLOO_k_unsc(x_k, c_k, mhat_k, hyper_param)
	#only compute for terms where mhat_k[i] > 0
	out = 0.
	for i = 1:length(mhat_k)
		if mhat_k[i] > 0
			mhat_k[i] -= 1
		    x = x_k(mhat_k, hyper_param)
		    mhat_k[i] += 1
		    out += mhat_k[i] * c_k[i](x)
		end
	end
	out
end

#average true performance
#hyper_param = p0, alpha
function zbar(xs, cs, mhats, ps, lams, hyper_param)
	lamavg = mean(lams)
	K = size(cs, 2)
	out = 0.
	#VG maybe change this to a more numerically stable way to do avg?
	for k = 1:K
		out += z_k(xs[k], cs[:, k], mhats[:, k], ps[:, k], lams[k], lamavg, hyper_param)
	end
	out/K
end

#full-info for scaling/comparison
#uses the fact that data scale doesn't matter.
#Only works for the regular xs (which is lame)
zstar(xs, cs, ps, lams) = zbar(xs, cs, ps, ps, lams, (ps[:, 1], 0.))

function zLOObar_unsc(xs, cs, mhats, hyper_param)
	K = size(cs, 2)
	out = 0.
	for k = 1:K
		@inbounds @views out += zLOO_k_unsc(xs[k], cs[:, k], mhats[:, k], hyper_param)
	end
	out /K
end

#actually returns N lam_avg * ZCVbar
function zCVbar_unsc(xs, cs, cv_data, hyper_param)
	out = 0.
	numFolds = length(cv_data)
	lams = ones(length(xs))

	for fold = 1:numFolds
		#form the training set and testing distribution
		#be careful about empty problems in testing
		train_data = sum(cv_data[setdiff(1:numFolds, fold)])

		#dispatch to the zbar method, again using scaling not important
		out += zbar(xs, cs, train_data, cv_data[fold], lams, hyper_param)
	end
	return out
end

#returns alphaOR, minimizingAlphaIndex, curveInAlpha
function oracle_alpha(xs, cs, mhats, ps, lams, p0, alpha_grid)
	alphaOR = 0.
	jstar = -1
	best_val = Inf
	out = zeros(length(alpha_grid))
	for (j, alpha) in enumerate(alpha_grid)
		out[j] = zbar(xs, cs, mhats, ps, lams, (p0, alpha))
		if out[j] < best_val
			jstar = j
			alphaOR = alpha
			best_val = out[j]
		end
	end
	return alphaOR, jstar, out
end

#returns alphaLOO, jstar, and cuveInAlpha_unsc
function loo_alpha(xs, cs, mhats, p0, alpha_grid)
	alphaLOO = 0.
	jstar = -1
	best_val = Inf
	out = zeros(length(alpha_grid))

	for (j, alpha) in enumerate(alpha_grid)
		out[j] = zLOObar_unsc(xs, cs, mhats, (p0, alpha))
		if out[j] < best_val
			jstar = j
			alphaLOO = alpha
			best_val = out[j]
		end
	end
	return alphaLOO, jstar, out
end

function cv_alpha(xs, cs, mhats, p0, alpha_grid, numFolds)
	#divy up the data
	cv_data = split_cv(mhats, numFolds)

	alphaCV = 0.
	jstar = -1
	best_val = Inf
	out = zeros(length(alpha_grid))
	for (j, alpha) in enumerate(alpha_grid)
		out[j] = zCVbar_unsc(xs, cs, cv_data, (p0, alpha))
		if out[j] < best_val
			jstar = j
			alphaCV = alpha
			best_val = out[j]
		end
	end
	return alphaCV, jstar, out
end

#########
#Functions to optimize anchor among beta distributions
########
#perform a naive discretization into d+1 points (d bins)
#p_i = Prob((i-1)/(d) <= Beta i <= i /d ) i = 1:d
#assumes p0_out is of size d
function formBetaAnchor!(theta1, theta2, d, p0_out)
	@assert length(p0_out) == d "p0_out must be of length $d"
	dist = Distributions.Beta(theta1, theta2)
	for i = 1:d
		p0_out[i] = cdf(dist, i/d) - cdf(dist, (i-1)/d)
	end
	p0_out
end

#Fits a beta(theta1, theta2) distribution for anchor (discretized)
#returns alphaLOO, p0, best_loo that optimizes loo error
function loo_betaAnchor(xs, cs, mhats, alpha_grid, theta2_grid, mu_grid; info=false)
	alphaLOO = 0.
	theta1LOO = -Inf
	theta2LOO = -Inf
	best_val = Inf
	d = size(mhats, 1)
	p0 = zeros(d)

	for theta2 in theta2_grid
		#theta1 = mu/(1-mu) * theta2
		#since mu in [0, 1] we search a scaled grid
		for mu in mu_grid
			theta1 = mu / (1-mu) * theta2

			#form p0
			formBetaAnchor!(theta1, theta2, d, p0)

			for alpha in alpha_grid
				out = zLOObar_unsc(xs, cs, mhats, (p0, alpha))
				if out < best_val
					alphaLOO = alpha
					theta1LOO = theta1
					theta2LOO = theta2
					best_val = out
				end
			end #over alpha
		end #over mu  (aka theta1)
	end#over theta 2

    #confirm it does better than mm
    theta1, theta2 = mm_beta_param(get_GM_anchor(mhats))
   	formBetaAnchor!(theta1, theta2, d, p0)
   	alpha, jstar, curve = loo_alpha(xs, cs, mhats, p0, alpha_grid)
   	out = curve[jstar]
   	if out < best_val
   		return alpha, p0, out
   	end

	info && println("theta1Loo:\t", theta1LOO, "\t theta2LOO:\t", theta2LOO, "\tMean:\t", theta1LOO/(theta1LOO + theta2LOO))

	#compute the p0 one more time (hopefully fast)
	formBetaAnchor!(theta1LOO, theta2LOO, d, p0)
	return alphaLOO, p0, best_val
end

#Fits a beta(theta1, theta2) distribution for anchor (discretized)
#returns alphaOR, p0, perf that optimizes oracle error
function oracle_betaAnchor(xs, cs, mhats, ps, lams, alpha_grid, theta2_grid, mu_grid; info=false)
	alphaOR = 0.
	theta1OR = -Inf
	theta2OR = -Inf
	best_val = Inf
	d = size(mhats, 1)
	p0 = zeros(d)

	for theta2 in theta2_grid
		#theta1 = mu/(1-mu) * theta2
		#since mu in [0, 1] we search a scaled grid
		for mu in mu_grid
			theta1 = mu / (1-mu) * theta2

			#form p0
			formBetaAnchor!(theta1, theta2, d, p0)

			for alpha in alpha_grid
				out = zbar(xs, cs, mhats, ps, lams, (p0, alpha))
				if out < best_val
					alphaOR = alpha
					theta1OR = theta1
					theta2OR = theta2
					best_val = out

					info && println("alpha\t", alphaOR, "\t theta1Loo:\t", theta1OR, "\t theta2OR:\t", theta2OR, "\tMean:\t", theta1OR/(theta1OR + theta2OR))
				end
			end #over alpha
		end #over mu  (aka theta1)
	end#over theta 2

    #confirm it does better than mm
    theta1, theta2 = mm_beta_param(get_GM_anchor(mhats))
   	formBetaAnchor!(theta1, theta2, d, p0)
   	alpha, jstar, curve = oracle_alpha(xs, cs, mhats, ps, lams, p0, alpha_grid)
   	out = curve[jstar]
   	if out < best_val
   		return alpha, p0, out
   	end

	#compute the p0
	info && println("theta1Loo:\t", theta1OR, "\t theta2OR:\t", theta2OR, "\tMean:\t", theta1OR/(theta1OR + theta2OR))
	formBetaAnchor!(theta1OR, theta2OR, d, p0)
	return alphaOR, p0, best_val
end

#compute the beta dist param by mm to mean and variance
function mm_beta_param(pGM)
	supp = range(0, stop=1, length=length(pGM))
	mu = dot(pGM, supp)
	var = dot(pGM, supp.^2) - mu^2
	if var > mu * (1 - mu)
		#fail silently
		return (1., 1.)
	end
	v = mu * (1 - mu) / var
	return mu * v, (1 - mu) * v
end

#Fits a beta(theta1, theta2) distribution for anchor (discretized)
#returns alphaLOO, p0, best_loo that optimizes loo error
#uses a black-box optimization routine
function loo_betaAnchor(xs, cs, mhats; info=false, time_limit = 45, iterations=1000,
						init_alpha=1., init_theta1= 1., init_theta2=1.)

	d = size(mhats, 1)
	p0 = zeros(d) #used as shared storage

	#ys = [alpha, theta1, theta2]
	function f(ys)
		alpha = ys[1]^2
		theta1 = ys[2]^2
		theta2 = ys[3]^2
		#form p0
		formBetaAnchor!(theta1, theta2, d, p0)
		zLOObar_unsc(xs, cs, mhats, (p0, alpha))
	end
	ys1 = [init_alpha, init_theta1, init_theta2]
	info && println("Initial Value:\t", f(ys1))
    t= @elapsed results = optimize(f, ys1,
                    ParticleSwarm(n_particles=40),
                    Optim.Options(time_limit=time_limit, iterations=iterations))

    z1 = Optim.minimum(results)
    ystar = Optim.minimizer(results)
    info && println("Theta1:\t", ystar[2]^2, "\t Theta2:\t", ystar[3]^2, "\t mu:\t", (ystar[2]^2)/(ystar[2]^2 + ystar[3]^2))
	info && println("Time:\t", t, "\tTime Per Iteration:\t", t/iterations(results))
	formBetaAnchor!(ystar[2]^2, ystar[3]^2, d, p0)

	return ystar[1]^2, p0, z1
end

#### Helper function for binning data
#dat_vec is vector of actual realizations
#computes a grid of size d
#returns
	#d vector of lefthand binpoints
	#d vector of counts
	#raw histogram object for safety
#assumes NA's already filtered from calcualtion
function bin_data(dat_vec, d; TOL = .001)
	#Explicitly make bins to prevent histogram from rounding
	l = minimum(dat_vec)
	u = maximum(dat_vec)
	bin_width = (u - l)/d
    bin_edges = collect(range(l, stop=u, length=d + 1))
    bin_edges[end] += TOL * bin_width  #inflate the last one a little...
    bin_edges[1] -= TOL * bin_width
    dat_hist = fit(Histogram, dat_vec, bin_edges, closed =:left)
    bin_edges[1:d], dat_hist.weights, dat_hist
end

###Helper function for splitting data for k-fold cross-validation
#This is not a memory efficient implementation
#Just lazy split/iterations
function split_cv(mhats, numFolds)
	d, K = size(mhats)
    Ntot = convert(Int, sum(mhats))
    splits = mod.(1:Ntot, numFolds) .+ 1
    shuffle!(splits)
    cv_data = [zero(mhats) for fold = 1:numFolds]

    #iterate through each k in order
    ix_split = 1

    for k = 1:K
        for i = 1:d
            for j = 1:mhats[i, k]
                cv_data[splits[ix_split]][i, k] += 1
                ix_split += 1
            end
        end
    end
    return cv_data
end


using PlotlyJS

function tri_edges(verts_matr::AbstractMatrix, tri_matr::AbstractMatrix;
                   linecolor="rgb(70,70,70)", linewidth=1)
    #def triangulation edges from vertices and triangles, of size (m, 3), resp (n,3)
    if minimum(tri_matr) == 0
        tri_verts = [hcat(verts_matr[t, :]) for t in eachrow(tri_matr .+ 1)]
    else
        tri_verts = [hcat(verts_matr[t, :]) for t in eachrow(tri_matr)]
    end
    Xe = []
    Ye = []
    Ze = []
    for T in tri_verts
        append!(Xe, T[:, 1])
        append!(Xe, [T[1,1], NaN])
        append!(Ye, T[:, 2])
        append!(Ye, [T[1,2], NaN])
        append!(Ze, T[:, 3])
        append!(Ze, [T[1,3], NaN])
    end
    return PlotlyJS.scatter3d(
                   x=Xe,
                   y=Ye,
                   z=Ze,
                   mode="lines",
                   name=" ",
                   line=attr(color=linecolor, width=1.5))
end;

function black_bg(; bg_color = "rgb(20, 20, 20)",
                    grid_color="rgb(150, 150, 150)",
                    zeroline=false)
return  Dict{Any, Any}( #black background
        :showbackground => true,
        :backgroundcolor => bg_color,
        :gridcolor => grid_color,
        :zeroline => zeroline)
end;


module APStatistics

function frequencytable(array::AbstractArray, orderbykeys = false)
    keys = unique(array)
    frequencies = [[i, count(x->x == i, array)] for i in keys]
    if orderbykeys
        sort!(frequencies, by = x->x[1])
    else
        sort!(frequencies, by = x->x[2])
        reverse!(frequencies)
    end

    println("ITEM | COUNT")
    for item in frequencies
        println("$(item[1]) | $(item[2])")
    end
end
function binnedtable(array::AbstractArray, orderbykeys = false, bin = 5)
    keys = unique(array)
    array = [floor(i / bin) * bin for i in array]
    frequencies = [[i, count(x->x == i, array)] for i in keys]
    if orderbykeys
        sort!(frequencies, by = x->x[1])
    else
        sort!(frequencies, by = x->x[2])
        reverse!(frequencies)
    end

    println("BIN | COUNT")
    for item in frequencies
        println("$(item[1]) - $(item[1] + bin) | $(item[2])")
    end
end
end # module


# ranges.jl

(::Colon)(start::QDate, stop::QDate) = StepRange(start, Day(1), stop)

# Given a start and end date, how many steps/periods are in between
Dates.guess(a::QDate, b::QDate, c) = Int64(div(value(b - a), days(c)))

function Dates.len(a::QDate, b::QDate, c)
    lo, hi, st = min(a, b), max(a, b), abs(c)
    i = Dates.guess(a, b, c) - 1
    try
        while lo + st * i <= hi
            i += 1
        end
    catch ex
        if !isa(ex, ArgumentError) && !isa(ex, BoundsError)
            rethrow()
        # else
        #     #pass
        end
    end
    return i - 1
end


# -*- coding: utf-8 -*-
# ---
# jupyter:
#   jupytext:
#     formats: ipynb,jl:hydrogen
#     text_representation:
#       extension: .jl
#       format_name: hydrogen
#       format_version: '1.3'
#       jupytext_version: 1.11.2
#   kernelspec:
#     display_name: Julia 1.8.0-DEV
#     language: julia
#     name: julia-1.8
# ---

# %% [markdown]
# * https://github.com/JuliaMath/Combinatorics.jl/blob/master/src/permutations.jl#L47
# * https://discourse.julialang.org/t/is-there-a-function-behaving-the-same-as-next-permutation-does-in-c/63451/13

# %%
module O

"""
zero allocation version of AquaIndigo's code
"""
function next_perm_old!(v::AbstractVector)
    length(v) ≤ 1 && return false
    k = findlast(isless(v[i], v[i+1]) for i in firstindex(v):lastindex(v)-1)
    isnothing(k) && (reverse!(v); return false)
    i = k + firstindex(v) - 1
    j = findlast(isless(v[i], v[j]) for j in i+1:lastindex(v)) + i
    v[i], v[j] = v[j], v[i]
    reverse!(v, i + 1)
    return true
end

"""
    next_perm!(v::AbstractVector)

changes `v` into the next permutation of `v` under the lexicographic order and returns `false` if it is the last permutation and `true` otherwise.

This is a variant of `Combinatorics.nextpermutation(m, t, state)`.

Examples

```
julia> v = [1, 2, 3]; println(v); while O.next_perm!(v) println(v) end
[1, 2, 3]
[1, 3, 2]
[2, 1, 3]
[2, 3, 1]
[3, 1, 2]
[3, 2, 1]
```

```
julia> v = [1, 2, 2]; println(v); while O.next_perm!(v) println(v) end
[1, 2, 2]
[2, 1, 2]
[2, 2, 1]
```
"""
function next_perm!(v::AbstractVector)
    length(v) ≤ 1 && return false
    i = lastindex(v) - 1
    @inbounds while i ≥ firstindex(v) && !isless(v[i], v[i+1]) i -= 1 end
    i < firstindex(v) && (reverse!(v); return false)
    j = lastindex(v)
    @inbounds while j > i && !isless(v[i], v[j]) j -= 1 end
    @inbounds v[i], v[j] = v[j], v[i]
    reverse!(v, i + 1)
    return true
end

function collect_perm!(v::AbstractVector)
    a = [copy(v)]
    while next_perm!(v) push!(a, copy(v)) end
    a
end

"""
    Perm(v::AbstractVector)

constructs the lexicographic order iterator of the all permutations of `v`.

Examples

```
julia> for s in O.Perm([1, 2, 3]) println(s) end
[1, 2, 3]
[1, 3, 2]
[2, 1, 3]
[2, 3, 1]
[3, 1, 2]
[3, 2, 1]
```

```
julia> for s in O.Perm([1, 2, 2]) println(s) end
[1, 2, 2]
[2, 1, 2]
[2, 2, 1]
```
"""
struct Perm{T<:AbstractVector}
    v::T
    Perm(v::T) where T<: AbstractVector = new{T}(sort(v))
end

function Base.length(p::Perm)
    v = p.v
    N = factorial(length(v))
    i = firstindex(v)
    for j in eachindex(v)
        if isless(v[i], v[j])
            N = N ÷ factorial(j-i)
            i = j
        end
    end
    N = N ÷ factorial(lastindex(v)-i+1)
    N
end

Base.eltype(p::Perm{T}) where T<:AbstractVector = T

function Base.iterate(p::Perm)
    s = copy(p.v)
    t = copy(s)
    next_perm!(t) || return (s, nothing)
    u = copy(t)
    next_perm!(u) || return (s, (t, nothing))
    return (s, (t, u))
end

function Base.iterate(p::Perm, state)
    isnothing(state) && return nothing
    s, t = state
    isnothing(t) && return (s, nothing)
    u = copy(t)
    next_perm!(u) || return (s, (t, nothing))
    return (s, (t, u))
end

#function Base.iterate(p::Perm, s = copy(p.v))
#    isnothing(s) && return nothing
#    next_perm!(s) || return (s, nothing)
#    return (copy(s), s)
#end

end

# %%
@doc O.next_perm!

# %%
@doc O.Perm

# %%
v = [1, 2, 3]; println(v); while O.next_perm!(v) println(v) end

# %%
v = [1, 2, 2]; println(v); while O.next_perm!(v) println(v) end

# %%
for s in O.Perm([1, 2, 3]) println(s) end

# %%
for s in O.Perm([1, 2, 2]) println(s) end

# %%
v = sort([NaN, NaN, 2, 1])
@show v
while O.next_perm!(v)
    @show v
end

# %%
p = O.Perm([NaN, NaN, 2, 1])
for s in p
    @show s
end

# %%
using BenchmarkTools

# %%
p = O.Perm([NaN, NaN, 2, 1])
_, state = iterate(p)

# %%
iterate(p, state)

# %%
p = O.Perm([NaN, NaN, 2, 1])
@btime collect($p)

# %%
v = sort([NaN, NaN, 2, 1])
@btime O.collect_perm!($v)

# %%
test_next_perm!(v) = (c = 1; while O.next_perm!(v) c += 1 end; c)
@show s = [1, 1, 2, 2, 2, 3, 4, 5, 6, 6]
@btime test_next_perm!($s)

# %%
test_next_perm_old!(v) = (c = 1; while O.next_perm_old!(v) c += 1 end; c)
@show s = [1, 1, 2, 2, 2, 3, 4, 5, 6, 6]
@btime test_next_perm_old!($s)

# %%
c = @btime collect($(O.Perm(s)))
length(c)

# %%
@code_warntype O.next_perm!(s)

# %%
using OffsetArrays
@show t = OffsetArray([1, 1, 2, 2, 2, 3, 4, 5, 6, 6], -5:4)
@btime collect($(O.Perm(t)))

# %%
@btime O.collect_perm!($(sort(t)))

# %%
@code_warntype O.next_perm!(t)

# %%


#
#
"""
A piecewise_linear_generator. Used for generating budget constraints used in
conventional microeconomics. See Stark and Duncan 1992 in the biblio.
"""

struct Point2DG{T<:Real}
      x::T
      y::T
end

const Point2D = Point2DG{Float64}

struct Line2DG{T<:Real}
      a::T
      b::T
end

const Line2D = Line2DG{Float64}

const BudgetConstraint = Array{Point2DG,1}
const PointsSet = Set{Point2DG}

const VERTICAL   = 9999999999.9999;
const TOLERANCE  = 0.0001;
const INCREMENT  = 0.0001;
const MAX_DEPTH  = 50;
const MAX_INCOME = 20000.0;
const MIN_INCOME = 0.0;

const ROUND_OUTPUT = false;

struct BCSettings
    mingross :: Float64
    maxgross :: Float64
    increment :: Float64
    tolerance :: Float64
    round_output :: Bool
    maxdepth :: Integer
end

const DEFAULT_SETTINGS = BCSettings( MIN_INCOME, MAX_INCOME, INCREMENT, TOLERANCE, true, MAX_DEPTH )

function makeline( point_1 :: Point2D, point_2 :: Point2D )::Line2D
    a :: Float64 = 0.0
    b :: Float64 = 0.0
    if point_1.x == point_2.x # ?? never taken??
        b = point_1.x;
        a = VERTICAL;
    else
        b = (point_1.y - point_2.y)/(point_1.x - point_2.x );
        b = min( b, VERTICAL );
        a = ( point_1.y - point_1.x*b );
    end
    return Line2D( a, b );
end

function findintersection( line_1::Line2D, line_2 :: Line2D ) :: Point2D
    x :: Float64 = 0.0
    y :: Float64 = 0.0
    if !( line_1.b ≈ line_2.b )
        x = (line_2.a - line_1.a) / (line_1.b - line_2.b);
        y = line_1.a + (x * line_1.b );
    else
        x = 0.0;
        y = line_1.a;
    end
    return Point2D( x, y );
end

function comparepoints( point_1 :: Point2DG, point_2 :: Point2DG ) :: Integer
    if( point_1.x > point_2.x )
        return 1;
    end
    if( point_1.x < point_2.x )
        return -1;
    end
    if( point_1.y > point_2.y )
        return 1;
    end
    if( point_1.y < point_2.y )
        return -1;
    end
    return 0;
end

# sorting for points
import Base.isless

function isless( point_1 :: Point2DG, point_2 :: Point2DG ) :: Bool
    return comparepoints( point_1, point_2 ) < 0
end

function marginalrate( point_1 :: Point2D, point_2 :: Point2D )::Float64
    mr :: Float64 = 0.0
    if !( point_2.x ≈ point_1.x )
        mr = 100.0 * (1-(point_2.y-point_1.y) / (point_2.x - point_1.x));
    else
        mr = VERTICAL;
        if( point_2.y < point_1.y )
            mr *= -1;
        end
    end
    return mr;
end

# approximate equality for points and lines
import Base.≈

function ≈(left :: Point2D, right::Point2D )::Bool
   (left.x ≈ right.x) && ( left.y ≈ right.y )
end

function ≈(left :: Line2D, right::Line2D )::Bool
   # return (left.a ≈ right.a) && ( left.b ≈ right.b )
   #  FIXME this uses the global tolerance even
   # if the user overrides it in settings
   return ((( abs(left.a-right.a)) <= TOLERANCE ) && (( abs(left.b-right.b)) <= TOLERANCE ));
end

# round a float to 2dps
function round2pl( x::Float64 )::Float64
    x *= 100.0
    i = trunc(x)
    x = Float64(i/100)
end

function round2pl!( bc :: BudgetConstraint )
    nbc = size( bc )[1]
    for i in 1:nbc
        p = Point2D( round2pl( bc[i].x ), round2pl( bc[i].y))
        bc[i] = p
    end
end

function toarray( ps :: PointsSet ) :: BudgetConstraint
    bc = BudgetConstraint();
    for p in ps
        push!( bc, p )
    end
    sort!( bc ) # , lt=isless
    bc
end

function censor( ps :: PointsSet, round :: Bool=true ) :: BudgetConstraint
    # println( "censor; ps=$ps")
    bc = toarray( ps )
    sort!( bc )
    nbc = size( bc )[1]
    first = bc[1]
    last = bc[nbc]
    if( nbc < 3 )
        return bc
    end
    i = 2
    while i < (nbc - 1)
        p1 = bc[i-1]
        p2 = bc[i]
        p3 = bc[i+1]
        l1 = makeline( p1, p2 )
        l2 = makeline( p2, p3 )
        if l1 ≈ l2
            deleteat!( bc, i )
            nbc -= 1
        else
            i += 1
        end
    end # while
    i = 1
    if round
        round2pl!( bc )
    end
    # 1 liner which sorts and removes dups
    bc = BudgetConstraint( toarray( PointsSet( bc )))
    if bc[1] != first
        # bc = vcat(first,bc)
    end
    if bc[nbc] != last
        # bc = vcat(bc,last)
    end
    return bc
end

function generate!(
    bc       :: PointsSet,
    getnet,  # TODO give this a signature ( Float64 ) :: Float64
    depth    :: Integer,
    startpos :: Float64,
    endpos   :: Float64,
    settings :: BCSettings ) :: Integer
    diff = abs( startpos - endpos )
    tolerance = settings.tolerance
    if( diff < settings.tolerance )
        return depth
    end
    if depth > settings.maxdepth
        throw( "max depth exceeded $depth"  )
    end
    p1 = Point2D( startpos, getnet(startpos) )
    p2 = Point2D( startpos+settings.increment, getnet(startpos+settings.increment))
    p4 = Point2D( endpos, getnet(endpos) )
    p3 = Point2D( endpos-settings.increment, getnet(endpos-settings.increment) )

    line1 = makeline( p1, p2 )
    line2 = makeline( p3, p4 )

    if line1 ≈ line2
        push!( bc, p1 )
        push!( bc, p4 )

        return depth
    end
    p5 = findintersection( line1, line2 )
    if( p5.x <= startpos ) || ( p5.x >= endpos )
        anchor = startpos + (( endpos - startpos)/2.0)
    else
        anchor = p5.x
    end
    depth += 1
    #
    # expore to the left
    #
    depth = generate!( bc, getnet, depth, startpos, anchor, settings )
    #
    # then the right
    #
    depth = generate!( bc, getnet, depth, anchor, endpos, settings )
    return depth - 1
end

"""
Make a budget constraint using function `getnet` to extract net incomes and `settings` (see above on this struct).
getnet should be a function of the form `net=f(gross)`. See the testcase for an example.
"""
function makebc( getnet, settings :: BCSettings = DEFAULT_SETTINGS ) :: BudgetConstraint
    bc = BudgetConstraint()
    ps = PointsSet()
    depth = 0
    try
        depth = generate!( ps, getnet, depth, settings.mingross, settings.maxgross, settings )
        bc = censor( ps, settings.round_output )
    catch e
        ## FIXME print a fuller stack trace here
        println( "failed! $e")
        println(stacktrace())
    end
    bc;
end

function pointstoarray( bc :: BudgetConstraint ) :: Array{Float64,2}
    sz = size( bc, 1 )
    pts = zeros( Float64, sz, 2 )
    for b in 1:sz
        pts[b,1] = bc[b].x
        pts[b,2] = bc[b].y
    end
    pts
end


### A Pluto.jl notebook ###
# v0.12.17

using Markdown
using InteractiveUtils

# ╔═╡ c9aa684a-43a8-11eb-3400-b383901a7f05
begin
	using Pkg
	Pkg.add(url = "https://github.com/heltonmc/LightPropagation.git")
	using LightPropagation
end

# ╔═╡ 4399d234-43a5-11eb-2389-21beb59e41fd
md"# _When is semi-infinite semi-infinite?_

### Learning outcomes:
- Simulate the time-domain reflectance under the diffusion approximation for **semi-finite, slab, and parralelepiped** geometries
- Utilize the non-linear curve fitting functions to extract optical properties from simulated experimental data
- Gain physical intuition of approximations and how to apply them
"

# ╔═╡ 31c8aab0-43a7-11eb-3fea-417ceb3ba0b5
md"
!!! info \"Disclaimer\"
	The following is not a recommendation on when to apply certain approximations as       these will depend on experimental setup and application. Please use this as an         introduction to the software.
"

# ╔═╡ 117d72ca-43a9-11eb-1394-e584e55e76e9
md"
##### Loading LightPropagation.jl
Let's begin by loading the package into the current environment. This can be done as below or as shown on the main page.
"

# ╔═╡ f7a02712-43a8-11eb-1cff-eb83d99a0989
md"
## Forward Models

The functions we are using to simulate the TPSF for the semi-infinite, slab, and parallelpiped geometries are:
- **TPSF\_DA\_semiinf\_refl**
- **TPSF\_DA\_slab_refl**
- **TPSF\_DA\_paralpip_refl**

"

# ╔═╡ f70fa3ee-43a9-11eb-1599-c9d55ca04787
md"
!!! tip \"Available Functions\"
	To view all the available functions it is often useful to check the exported 		functions in the main file. For example if you go to                                   ~/LightPropagation/src/LightPropagation.jl in the repository you will find             `export TPSF_DA_semiinf_refl` lines that indicate available functions.
"

# ╔═╡ acd4eee0-43ab-11eb-14f7-f95c2291d580
md"
Let's simulate the TPSF for a semi-infinite medium under the following conditions:
- t = 0:0.01:5 (ns)
- β = [0.1, 10.0] -- β[1] and β[2] are the absorption and reduced scattering coefficient (1/cm)
- ρ = 1.0 -- source-detector separation (cm)
- ndet = 1.0 -- index of refraction of detector
- nmed = 1.0 -- index of refraction of medium
"

# ╔═╡ 27dee8a8-43ab-11eb-33a2-81af1f08c66d
begin
	t = 0.01:0.01:5
	TPSF_semiinf = TPSF_DA_semiinf_refl(t, [0.1,10.0], 1.0, 1.0, 1.0)
end

# ╔═╡ a32eb024-43ab-11eb-09fa-9dc34018f56f
begin
	using Plots
	plot(t, TPSF_semiinf, yscale=:log10, lw = 2, ylabel="Counts", 							xlabel="time (ns)", label="TPSF_semiinf")
end

# ╔═╡ 5864c6da-43ad-11eb-117a-b7270071791d
md"
##### Comparison of semi-infinite and slab
Now that we know how to run a function and plot it, let's compare the semi-infinite approximation to a slab (bounded in the z-direction) for two thickness (2cm and 8cm)
"

# ╔═╡ 2a29c586-43ad-11eb-007f-45911e0d0ca5
begin
	TPSF_slab_2cm = TPSF_DA_slab_refl(t, [0.1,10.0], 1.0, 1.0, 1.0, 2.0)
	TPSF_slab_8cm = TPSF_DA_slab_refl(t, [0.1,10.0], 1.0, 1.0, 1.0, 8.0)
end

# ╔═╡ d5282152-43ae-11eb-3381-5bb1db31fece
begin
	plot(t, TPSF_semiinf, yscale=:log10, lw = 6, ylabel="Counts", 			 			xlabel="time (ns)", label="TPSF_semiinf") #make lw thicker for clarity
	plot!(t, TPSF_slab_2cm, lw = 2, label = "TPSF_slab_2cm")
	plot!(t, TPSF_slab_8cm, lw = 2, label = "TPSF_slab_8cm", c="black")
end

# ╔═╡ 5bd4961c-43b0-11eb-0d78-c3850f2f643d
md"
Great. Now change the input parameters for each curve (optical properties and SDS) to investigate different domains. Here, for this set of optical properties and SDS the  semi-infinite approximation matches exactly the slab solution when the slab thickness is 8cm (for t<5ns), though it would be a poor approximation for thin samples (2cm).
"

# ╔═╡ e98e6dae-43b3-11eb-2e8a-01012c7befe2
md"
#### Compare all three geometries
Let's introduce the parralelepiped geometry where it is bounded in the x, y, and z direction. Remember the inputs look like this:

TPSF\_DA\_paralpip\_refl(t, β::Array{Float64,1}, ndet::Float64, nmed::Float64, rd::Array{Float64,1}, rs::Array{Float64,1}, L::Array{Float64,1})

You can also check arguments by typing ? in the REPL and typing in the function name as shown in main page. Here, rd is the x, y location of the source on the top surface, and rs is the location of the detector on the top source. These are measured from the corner. Let's have two cubes one with L = 8cm and L = 4cm with the source located in the center so that rs = [4.0, 4.0] and [2.0, 2.0] respectively. We want the detector to be 1cm from the source to match SDS in previous examples. Below they are defined as rd = [3.0, 4.0] and [1.0, 2.0].
"

# ╔═╡ d40bc0b0-43b0-11eb-1f81-597d8d2a017b
begin
	TPSF_paralpip_8cm = TPSF_DA_paralpip_refl(t, [0.1,10.0], 1.0, 1.0, [3.0,4.0], 			[4.0,4.0], [8.0,8.0,8.0])
	TPSF_paralpip_4cm = TPSF_DA_paralpip_refl(t, [0.1,10.0], 1.0, 1.0, [1.0,2.0], 			[2.0,2.0], [4.0,4.0,4.0])
end

# ╔═╡ 0cb1c518-43b6-11eb-3bac-052d89f819ae
begin
	plot(t, TPSF_semiinf, yscale=:log10, lw = 10, ylabel="Counts", 			 			xlabel="time (ns)", label="TPSF_semiinf") #make lw thicker for clarity
	plot!(t, TPSF_slab_2cm, lw = 2, label = "TPSF_slab_2cm")
	plot!(t, TPSF_slab_8cm, lw = 4, label = "TPSF_slab_8cm", c="black")
	plot!(t, TPSF_paralpip_8cm, lw = 2, label = "TPSF_paralpip_8cm", alpha = 0.9)
	plot!(t, TPSF_paralpip_4cm, lw = 2, label = "TPSF_paralpip_4cm", alpha = 0.9)

end

# ╔═╡ 9a7d26ee-4479-11eb-1eb0-f79e3c07fb28
pwd()

# ╔═╡ 46710c9e-43b8-11eb-0172-65d7f2171f56
md"
#### Key Takeaways
- Good agreement at the earliest times (<1ns) where the boundary effects are minimized.
- For confined geometries, (e.g. TPSF\_paralpip\_4cm and TPSF\_slab\_2cm) we see there are significant differences compared to the semi-infinite approximations.
- If the boundaries are significantly away from the source the semi-infinite approximation appears to be a good approximation.
"

# ╔═╡ 3cb5c3e2-43b9-11eb-2299-7503e415f13c
md"
## Inverse Problem

Now that we have visually observed the differences between the three geometries, let's investigate and try to quantify the error in the extracted optical properties when using different approximations.

The first thing we must do is simulate some experimental data. In practice, samples are not infinite so let's use the parralelepiped geometry as the ground 'truth' and let's fit all three models to that data.
"

# ╔═╡ d2ee9b40-43eb-11eb-1f68-7f94589ba3f4
ydata = TPSF_DA_paralpip_refl(t, [0.1,10.0], 1.0, 1.0, [0.5,1.5], [1.5,1.5], [3.0,3.0,3.0])

# ╔═╡ 2d46a2d4-43b7-11eb-10fa-03d55dba74e8
md"
##### Setup Inverse problem

The first thing we must do is set up the input data and model parameters. These use two separate structures. The input data has the following structure:
```julia
@with_kw struct fitDTOF

    #required inputs
    t::Union{AbstractRange{Float64}, AbstractVector{Float64}}
    DTOF::Array{Float64,1}
    IRF::Array{Float64,1}

    nmed::Float64 = 1.5
	ndet::Float64 = 1.51
	lambda::Float64 = 750.0

    #semi-infinite and slab
    ρ::Union{Float64, Missing} = missing

    #slab
    s::Union{Float64, Missing} = missing

    #parallelepiped
    rd::Union{Array{Float64,1}, Missing} = missing
    rs::Union{Array{Float64,1}, Missing} = missing
    L::Union{Array{Float64,1}, Missing} = missing
end
```
"

# ╔═╡ 27b1a654-43ec-11eb-1b4f-9f165debb017
md"
We can see that we must also describe and IRF. For now, let's assume the IRF is as defined below.
"

# ╔═╡ 6f0f914c-43c3-11eb-2561-f7c8cae30e99
IRF = [1; zeros(length(ydata) - 1)]

# ╔═╡ 369062ea-43c3-11eb-3143-a562cfe3afe5
data = fitDTOF(t = t, DTOF = ydata./maximum(ydata), IRF = IRF, ρ = 1.0, nmed = 1.0, ndet = 1.0, s = 3.0, rd = [0.5,1.5], rs = [1.5,1.5], L = [3.0,3.0,3.0])

# ╔═╡ 495b4f02-43c5-11eb-3ea6-6d35ebc81514
m = DTOF_fitparams(model = TPSF_DA_paralpip_refl, risefactor = 0.01, tailfactor = 1e-10,initparams = [0.05, 20.0])

# ╔═╡ 6508edfe-43c5-11eb-2d67-4fcc34883150
f, tfit, yraw, yfit  = getfit(data, m)

# ╔═╡ 50348fb8-43c8-11eb-2bfc-6102f19bcf0e
begin
	m1 = DTOF_fitparams(model = TPSF_DA_semiinf_refl, risefactor = 0.01, tailfactor = 1e-10, initparams = [0.05, 20.0])
	m2 = DTOF_fitparams(model = TPSF_DA_slab_refl, risefactor = 0.01, tailfactor = 		1e-10, initparams = [0.05, 20.0])
end

# ╔═╡ 6fa2a252-43c8-11eb-2ffd-73dfebb5454e
begin
	f1, tfit1, yraw1, yfit1  = getfit(data, m1)
	f2, tfit2, yraw2, yfit2  = getfit(data, m2)
end

# ╔═╡ 8790e19e-43c8-11eb-11d4-e5b36da13513
begin
	scatter(tfit, yraw, m="*", ms = 8, label = "ydata", c="black")
	plot!(tfit, yfit, lw = 3, alpha = 1, label = "paralpip model", c="blue")
	plot!(tfit1, yfit1, lw = 3, alpha = 1,label = "semiinf model", c="red")
	plot!(tfit2, yfit2,lw = 3, alpha = 1, label = "slab model", c="orange")
	plot!(ylabel="Counts", xlabel="time (ns)")
	savefig("fitcomparison.png")
end

# ╔═╡ a8e788d4-43c8-11eb-3c20-b94cd8884937
f.param

# ╔═╡ af161984-43c8-11eb-0715-079cc75433ec
f1.param

# ╔═╡ b231ab40-43c8-11eb-0819-1d16af7aafb4
f2.param

# ╔═╡ c298b8e6-4488-11eb-2b48-b9a7fcc20740
abs.(f.param .- f2.param)./f.param

# ╔═╡ Cell order:
# ╟─4399d234-43a5-11eb-2389-21beb59e41fd
# ╠═31c8aab0-43a7-11eb-3fea-417ceb3ba0b5
# ╟─117d72ca-43a9-11eb-1394-e584e55e76e9
# ╠═c9aa684a-43a8-11eb-3400-b383901a7f05
# ╟─f7a02712-43a8-11eb-1cff-eb83d99a0989
# ╟─f70fa3ee-43a9-11eb-1599-c9d55ca04787
# ╟─acd4eee0-43ab-11eb-14f7-f95c2291d580
# ╠═27dee8a8-43ab-11eb-33a2-81af1f08c66d
# ╠═a32eb024-43ab-11eb-09fa-9dc34018f56f
# ╟─5864c6da-43ad-11eb-117a-b7270071791d
# ╠═2a29c586-43ad-11eb-007f-45911e0d0ca5
# ╠═d5282152-43ae-11eb-3381-5bb1db31fece
# ╟─5bd4961c-43b0-11eb-0d78-c3850f2f643d
# ╟─e98e6dae-43b3-11eb-2e8a-01012c7befe2
# ╠═d40bc0b0-43b0-11eb-1f81-597d8d2a017b
# ╠═0cb1c518-43b6-11eb-3bac-052d89f819ae
# ╠═9a7d26ee-4479-11eb-1eb0-f79e3c07fb28
# ╟─46710c9e-43b8-11eb-0172-65d7f2171f56
# ╟─3cb5c3e2-43b9-11eb-2299-7503e415f13c
# ╠═d2ee9b40-43eb-11eb-1f68-7f94589ba3f4
# ╟─2d46a2d4-43b7-11eb-10fa-03d55dba74e8
# ╟─27b1a654-43ec-11eb-1b4f-9f165debb017
# ╠═6f0f914c-43c3-11eb-2561-f7c8cae30e99
# ╠═369062ea-43c3-11eb-3143-a562cfe3afe5
# ╠═495b4f02-43c5-11eb-3ea6-6d35ebc81514
# ╠═6508edfe-43c5-11eb-2d67-4fcc34883150
# ╠═50348fb8-43c8-11eb-2bfc-6102f19bcf0e
# ╠═6fa2a252-43c8-11eb-2ffd-73dfebb5454e
# ╠═8790e19e-43c8-11eb-11d4-e5b36da13513
# ╠═a8e788d4-43c8-11eb-3c20-b94cd8884937
# ╠═af161984-43c8-11eb-0715-079cc75433ec
# ╠═b231ab40-43c8-11eb-0819-1d16af7aafb4
# ╠═c298b8e6-4488-11eb-2b48-b9a7fcc20740


using DomainSets, BasisFunctions, FrameFun.FrameFunDomains, Test

function test_extended_domains_3d()
    domains = (atomium(),)
end

function test_extended_domains_2d()
    D1 = characteristic(x->(x[1]-x[2])<0,disk(1.0))
    @test in([0.5;1.0],D1)
    @test ~in([1.0,0.5],D1)
    # basis = ChebyshevT(21)⊗ChebyshevT(21)
    # domain = disk(1.0)
    # f = (x,y)->exp(x+y)
    # F = Fun(f, basis, domain)
    # F2 = F*2
    # D2 = F<2
    # D3 = F<F2
    # @test ~in([0.5;0.5],D2)
    # @test in([0.25;0.25],D2)
    # @test in([0.5;0.5],D3)
    D4 = mandelbrot()
    @test in([0.0;0.0],D4)
    @test ~in([-0.5;0.5],D4)
    D5 = juliaset()
    @test in([0.6;0.0],D5)
    @test ~in([-0.1;-0.3],D5)
    # D6 = polardomain(x->1+0.2*cos(5x),disk(2.0))
    # @test ~in([1.25;0],D6)
    # @test in([0.79;0],D6)
end

domn(point,domain::Domain,normalvector; tol=1e-14) = norm(normal(point,domain)-normalvector)<tol
domb(point,domain::Domain; tol=1e-14) = abs(distance(point,domain))<tol

function test_distances_and_normals()
    D1 = simplex(Val{2})
    @test domb([0.5,0.5],D1)
    @test domn([0.5,0.5],D1,[1/sqrt(2),1/sqrt(2)])
    @test domb([0.5,0],D1)
    @test domn([0.5,0.0],D1,[0,-1])

    # D6 = polardomain(x->1+0.2*cos(5x),disk(2.0))
    # @test domb([-0.8,0],D6)
    # @test domn([-0.8,0],D6,[-1,0],tol=1e-12)

    D2 = disk(sqrt(0.5))
    @test domb([0.5,0.5],D2)
    @test domn([0.5,0.5],D2,[1/sqrt(2),1/sqrt(2)])

    D2b = circle(sqrt(0.5))
    @test domb([0.5,0.5],D2b)
    @test domn([0.5,0.5],D2b,[1/sqrt(2),1/sqrt(2)])

    D3 = -1..1.0
    @test isapprox(distance(0.5,D3),0.5)
    @test isapprox(normal(1.0,D3),1)

    D5 = cube(Val{2})
    @test domb([0.4,0.0],D5)
    @test domn([0.4,0.0],D5,[0,-1])

    # D4 = UnionDomain(D6, disk())\D2
    # @test domb([sqrt(0.5),0.0],D4)
    # @test domb([-1.0,0.0],D4)
    # @test ~domb([-0.8,0.0],D4)
    # @test domn([-1.0,0],D4,[-1,0])
    # @test domn([0.5,0.0],D4,[-1.0,0])

end

@testset "FrameFun Domains" begin

    test_extended_domains_3d()

    test_extended_domains_2d()

    test_distances_and_normals()
end


mutable struct LinearProblem{T} <: AbstractLP{T}
    A::Matrix{T}
    b::Vector{T}
    c::Vector{T}
    status::Status
    solvers::Array{AbstractSolver{T}}
    isMin::Bool
    xstar::Vector{T}
    vstar::T
    function LinearProblem(A::Matrix{T}, b::Vector{T}, c::Vector{T}; isMin::Bool = true) where T
        isaninteger = (T <: Integer)
        isaninteger && @warn "$T cannot be used by solver, $Float64 used instead"
        return new{isaninteger ? Float64 : T}(A, b, c, Unknown, [], isMin, zeros(T, size(A, 2)))
    end
    function LinearProblem{T}(A::Matrix, b::Vector, c::Vector; isMin::Bool = true) where T
        isaninteger = (T <: Integer)
        isaninteger && @warn "$T cannot be used by solver, $Float64 used instead"
        return new{isaninteger ? Float64 : T}(Array{T, 2}(A), Array{T, 1}(b), Array{T, 1}(c), Unknown, [], isMin, zeros(T, size(A, 2)))
    end
end
function LinearProblem(A::Matrix{TA}, b::Vector{Tb}, c::Vector{Tc}; isMin::Bool = true) where {TA, Tb, Tc}
    T = promote_type(TA, Tb, Tc)
    (T <: Integer) && (T = Float64)
    return LinearProblem{T}(A, b, c, isMin = isMin)
end

const LP{T} = LinearProblem{T}

function xstar(lp::LP{T})::Vector{T} where T
    @assert lp.status == Optimal
    return copy(lp.xstar)
end
function vstar(lp::LP{T})::T where T
    @assert lp.status == Optimal
    return lp.vstar
end

function nvar(lp::LP)::Int
    return size(lp.A, 2)
end
function ncons(lp::LP)::Int
    return size(lp.A, 1)
end
function getA(lp::LP{T})::Array{T, 2} where T
    return lp.A
end
function getb(lp::LP{T})::Array{T, 1} where T
    return lp.b
end
function getc(lp::LP{T})::Array{T, 1} where T
    return lp.c
end

function addconstraint!(lp::LinearProblem{T}, v::Vector{T}, b::T) where T
    lp.A = [lp.A; v']
    lp.b = [lp.b; b]
    #lp.c::Vector{T}
    lp.status = Unknown
    addconstraint!.(lp.solvers, v, b)
    lp.xstar[:] .= zero(T)
    lp.vstar = lp.isMin ? abignumber(T) : - abignumber(T)
end
function addvariable!(lp::LinearProblem{T}) where T
    lp.A = [lp.A zeros(T, size(lp.A, 1))]
    lp.c = [lp.c; zero(T)]
    addvariable!.(lp.solvers)
end


# degree matrix
degree_matrix(x::T) where {T<:AbstractHyperGraph} = Diagonal(degrees(x))

# incidence matrix (oriented case)
@traitfn function incidence_matrix(x::T::IsOriented) where {T<:AbstractHyperGraph}
    idx = vertices_to_indices(x)
    I = zeros(Int, nv(x), nhe(x))
    @inbounds for (i, e) in enumerate(hyperedges(x))
        [I[idx[v], i] -= s for (v, s) in zip(objects(e.src), src_multiplicities(e))]
        [I[idx[v], i] += s for (v, s) in zip(objects(e.tgt), tgt_multiplicities(e))]
    end
    I
end

# incidence matrix (unoriented case)
@traitfn function incidence_matrix(x::T::(!IsOriented)) where {T<:AbstractHyperGraph}
    idx = HyperGraphs.vertices_to_indices(x)
    I = zeros(Int, nv(x), nhe(x))
    @inbounds for (i, e) in enumerate(hyperedges(x))
        [I[idx[v], i] += s for (v, s) in zip(vertices(e), fill(1, cardinality(e)))]
    end
    I
end

# cardinality matrix
cardinality_matrix(x::T) where {T<:AbstractHyperGraph} = Diagonal(cardinalities(x))

# weight matrix
weight_matrix(x::T) where {T<:AbstractHyperGraph} = Diagonal(weights(x))


using ADCME
using PyPlot
using Random
using MAT
using DelimitedFiles
matplotlib.use("agg")

include("fwi_data.jl")
reset_default_graph()
tf.set_random_seed(123)
tf.random.set_random_seed(123)

R_ = constant(R)

idx_var = pml+Int(round(0.1/Δx))

inv_sigmoid(y) = @. - log(1.0/y - 1.0)

case = "fwi"

## fwi
if case == "fwi"
if model == "step"
    var= Variable(zeros(N-1-idx_var)) + inv_sigmoid(1.3/2.0)
    var = -2.0*sigmoid(var)
elseif model == "slop"
    var = Variable(zeros(N-1-idx_var))
end
end

## FC
if case == "nnfwi"
x0 = collect( range(0, 1, length=N-1-idx_var) )
init_guess = fc_init([1, 30, 30, 30, 1])
θ = Variable(init_guess)
if model == "step"
    var= fc(x0, [30, 30, 30, 1], θ) * 0.1 + inv_sigmoid(1.3/2.0)
    var = -2.0*sigmoid(var)
elseif model == "slop"
    var= fc(x0, [30, 30, 30, 1], θ) * 0.1
end
end

## CNN
# x = constant(randn(1,4,1))
# x = tf.keras.layers.Conv1D(8, 3, padding="same")(x)
# x = tf.keras.layers.LeakyReLU(alpha=0.1)(x)
# x = tf.keras.layers.UpSampling1D(4)(x)
# x = tf.keras.layers.Conv1D(8, 3, padding="same")(x)
# x = tf.keras.layers.LeakyReLU(alpha=0.1)(x)
# x = tf.keras.layers.UpSampling1D(4)(x)
# x = tf.keras.layers.Conv1D(8, 3, padding="same")(x)
# x = tf.keras.layers.LeakyReLU(alpha=0.1)(x)
# x = tf.keras.layers.UpSampling1D(4)(x)
# x = tf.keras.layers.Conv1D(1, 3, padding="same")(x)
# x = squeeze(x)
# var = tf.slice(x, (length(x) - (N-1-idx_var),).÷2, (N-1-idx_var,)) * 0.1 - 1.0
# var = (3.5 + 0.5) * sigmoid(var) - 0.5

cH = scatter_add(constant(cH_init), idx_var+1:N-1, var)

cE = (cH[1:end-1]+cH[2:end])/2

reg_TV = sum(abs(cH[2:end]-cH[1:end-1]))
reg_smooth = sum((cH[1:end-2] + cH[3:end] - 2*cH[2:end-1])^2)

function condition(i, E_arr, H_arr)
    i<=NT+1
end

function body(i, E_arr, H_arr)
    E = read(E_arr, i-1)
    H = read(H_arr, i-1)

    ΔH = cH * (E[2:end]-E[1:end-1])/Δx - σH*H
    H += ΔH * Δt
    ΔH = 1/(24Δx) * cH[2:N-2] * (-E[4:end] + 3E[3:end-1] - 3E[2:end-2] + E[1:end-3])
    H = scatter_add(H, 2:N-2, ΔH * Δt)

    ΔE = cE * (H[2:end]-H[1:end-1])/Δx - σE[2:end-1]*E[2:end-1] + R_[i] * Z
    E = scatter_add(E, 2:N-1, ΔE * Δt)
    ΔE = 1/(24Δx) * cE[2:end-1] * (-H[4:end] + 3H[3:end-1] - 3H[2:end-2] + H[1:end-3])
    E = scatter_add(E, 3:N-2, ΔE*Δt)

    i+1, write(E_arr, i, E), write(H_arr, i, H)
end

E_arr = TensorArray(NT+1)
H_arr = TensorArray(NT+1)

E_arr = write(E_arr, 1, zeros(N))
H_arr = write(H_arr, 1, zeros(N-1))

i = constant(2, dtype = Int32)

_, E, H = while_loop(condition, body, [i, E_arr, H_arr])

E = stack(E); E = set_shape(E, (NT+1, N))
H = stack(H)

loss = sum((E[:, idx_rcv] - E_true[:, idx_rcv])^2 * [1.0, 1e-3])
loss1 = sum((E[:, idx_rcv[1]] - E_true[:, idx_rcv[1]])^2 * 1.0)
loss2 = sum((E[:, idx_rcv[2]] - E_true[:, idx_rcv[2]])^2 * 1e-3)

if case == "fwi"
    # loss += 1e3*reg_smooth + 0.0*reg_TV
    # case = string(case, "_reg_smooth")
    loss += 1e2*reg_smooth + 1e2*reg_TV
    case = string(case, "_reg1e2_smooth_TV")
end

if !ispath(joinpath("figures", model, case))
    mkpath(joinpath("figures", model, case))
end
if !ispath(joinpath("results", model, case))
    mkpath(joinpath("results", model, case))
end

function cb(vs, iter, loss)
    if (mod(iter, 20) == 1)
        E_est = vs[1]
        cH_est = vs[2]
        @info iter, vs[3], vs[4]
        close("all")
        figure(figsize = (10,6))
        subplot(221)
        plot(xH, cH_true, "C0", label="True")
        plot(xH, cH0, "k--", label="Initial")
        plot(xH, cH_est, "C3--", label="Estimated")
        legend()
        # xlim([0,1])
        autoscale(enable=true, axis="x", tight=true)
        xlabel("Location (km)")
        ylabel("Velocity (km/s)")
        subplot(222)
        plot((0:NT)*Δt, E_true[:,idx_rcv[1]], "C0", label="True #1",  alpha=0.8)
        plot((0:NT)*Δt, E_true[:,idx_rcv[2]], "C2", label="True #2", alpha=0.8)
        plot((0:NT)*Δt, E_est[:,idx_rcv[1]], "C3--",  label="Estimated #1", alpha=0.8)
        plot((0:NT)*Δt, E_est[:,idx_rcv[2]], "C1--",  label="Estimated #2", alpha=0.8)
        legend(loc="upper right")
        gca().ticklabel_format(axis="y", style="sci", scilimits=(0,0))
        autoscale(enable=true, axis="x", tight=true)
        xlabel("Time (s)")
        ylabel("Amplitude")
        subplot(223)
        pcolormesh(xE, (0:NT)*Δt, E_true, cmap="seismic", vmax=maximum(E_true)/3.0, vmin=-maximum(E_true)/3.0, rasterized=true)
        plot(xE[idx_src]*ones(NT+1), (0:NT)*Δt, "k", label="Source", alpha=0.8)
        plot(xE[idx_rcv[1]]*ones(NT+1), (0:NT)*Δt, "C0--", label="Receiver #1", alpha=0.8)
        plot(xE[idx_rcv[2]]*ones(NT+1), (0:NT)*Δt, "C2--", label="Receiver #2", alpha=0.8)
        # xlim([0,1])
        autoscale(enable=true, axis="x", tight=true)
        legend(loc="upper center")
        xlabel("Localtion (km)")
        ylabel("Time (s)")
        title("True")
        subplot(224)
        pcolormesh(xE, (0:NT)*Δt, E_est, cmap="seismic", vmax=maximum(E_true)/3.0, vmin=-maximum(E_true)/3.0, rasterized=true)
        plot(xE[idx_src]*ones(NT+1), (0:NT)*Δt, "k", label="Source", alpha=0.8)
        plot(xE[idx_rcv[1]]*ones(NT+1), (0:NT)*Δt, "C0--", label="Receiver #1", alpha=0.8)
        plot(xE[idx_rcv[2]]*ones(NT+1), (0:NT)*Δt, "C2--", label="Receiver #2", alpha=0.8)
        # xlim([0,1])
        autoscale(enable=true, axis="x", tight=true)
        legend(loc="upper center")
        xlabel("Localtion (km)")
        ylabel("Time (s)")
        title("Estimated")

        # subplot(325)
        # plot((0:NT)*Δt, E_true[:,idx_rcv[1]].-E_est[:,idx_rcv[1]], "C0", label="Diff #1",  alpha=0.8)
        # legend(loc="upper right")
        # gca().ticklabel_format(axis="y", style="sci", scilimits=(0,0))
        # autoscale(enable=true, axis="x", tight=true)
        # xlabel("Time (s)")
        # ylabel("Amplitude")
        # subplot(326)
        # plot((0:NT)*Δt, E_true[:,idx_rcv[2]].-E_est[:,idx_rcv[2]], "C3", label="Diff #2", alpha=0.8)
        # legend(loc="upper right")
        # gca().ticklabel_format(axis="y", style="sci", scilimits=(0,0))
        # autoscale(enable=true, axis="x", tight=true)
        # xlabel("Time (s)")
        # ylabel("Amplitude")

        tight_layout()
        savefig(joinpath("figures", model, case, "fwi$(lpad(iter,5,"0")).png"), bbox_inch="tight")
        if (mod(iter, 200) == 1)
        matwrite(joinpath("results", model, case, "fwi$(lpad(iter,5,"0")).mat"),
                 Dict("idx_rcv"=>idx_rcv, "idx_src"=>idx_src, "xE"=>xE,
                      "cH0"=>cH0, "cH_true"=>cH_true, "cH_est" => cH_est,
                      "E_true" => E_true, "E_est" => E_est))
        # writedlm(joinpath("results", model, case, "fwi$(lpad(iter,5,"0")).txt"), E_est)
        end
    end
end

lr = placeholder(Float64, shape=[])
opt = AdamOptimizer(learning_rate=lr).minimize(loss)
# opt = tf.train.AdagradOptimizer(learning_rate=0.01).minimize(loss)
# opt = tf.train.GradientDescentOptimizer(learning_rate=0.01).minimize(loss)
# opt = tf.train.RMSPropOptimizer(learning_rate=0.01).minimize(loss)

sess = Session(); init(sess)
cH0 = run(sess, cH)
cb(run(sess, [E, cH, loss1, loss2]), -1, run(sess, loss))
BFGS!(sess, loss, vars = [E, cH, loss1, loss2], callback = cb)

# for iter = 0:10000
#     _, l = run(sess, [opt, loss], feed_dict=Dict(lr=>(randexp()*0.001+0.00001)))
#     cb(run(sess, [E, cH]), iter, l)
#     if mod(iter, 20) == 1
#         @info iter, l
#     end
# end


using StanSample, Random, Statistics

stan_chris = "
data {
    int n_rows;
    int<lower=1> n_cols;
    matrix<lower=0>[n_rows,n_cols] x;
}

parameters {
   real mu;
}

model {
    mu ~ normal(0, 1);
}";

seed = 65445
Random.seed!(seed)

n_rows = 10
n_cols = 2
data = [(x = rand(n_rows, n_cols), n_rows = n_rows, n_cols = n_cols) for i in 1:4]

tmpdir = joinpath(@__DIR__, "tmp")
sm = SampleModel("chris", stan_chris, tmpdir);
rc = stan_sample(sm; data, seed);

if success(rc)
	df = read_samples(sm , :dataframe)
    #df |> display
    [mean(df.mu), std(df.mu)] |> display
end

tmpdir = joinpath(@__DIR__, "tmp")
sm = SampleModel("chris", stan_chris, tmpdir);
rc = stan_sample(sm; data=data[1:3], seed);

if success(rc)
    df = read_samples(sm , :dataframe)
    #df |> display
    [mean(df.mu), std(df.mu)] |> display
end


__precompile__()

module Hymod

using DataFrames
using Dates
using Distributions
using Statistics
using Logging
using JSON

export randomparams, hargreaves, simulate, calibrate

"""
    randomparams()

Function to create a random set of parameters as kwargs for Hymod.simulate()
"""
function randomparams()

    param_bounds = Dict{Symbol,Dict}(
        :cmax => Dict{Symbol,Float64}(:lower => 1.0, :upper => 100),
        :bexp => Dict{Symbol,Float64}(:lower => 0.0, :upper => 2.0),
        :alpha => Dict{Symbol,Float64}(:lower => 0.2, :upper => 0.99),
        :ks => Dict{Symbol,Float64}(:lower => 0.01, :upper => 0.5),
        :kq => Dict{Symbol,Float64}(:lower => 0.5, :upper => 1.2),
    )

    out = Dict{Symbol,Float64}()

    for k in keys(param_bounds)
        minV = param_bounds[k][:lower]
        maxV = param_bounds[k][:upper]
        p = rand(Uniform(minV, maxV), 1)
        out[k] = p[1]
    end

    return out
end

function abspower(x::Float64, y::Float64)
    # Needed to capture invalid overflow with netgative values
    x = abs(x)
    return x^y
end

function precipexcess(
    x_loss::Float64,
    cmax::Float64,
    bexp::Float64,
    prc::Float64,
    pet::Float64,
)
    xn_prev = x_loss
    ct_prev =
        cmax *
        (1 - abspower((1 - ((bexp + 1) * (xn_prev) / cmax)), (1 / (bexp + 1))))
    # Calculate Effective rainfall 1
    ER1 = max((prc - cmax + ct_prev), 0.0)
    prc = prc - ER1
    dummy = min(((ct_prev + prc) / cmax), 1)
    xn = (cmax / (bexp + 1)) * (1 - abspower((1 - dummy), (bexp + 1)))

    # Calculate Effective rainfall 2
    ER2 = max(prc - (xn - xn_prev), 0)

    # Alternative approach
    # actual ET is linearly related to the soil moisture state
    evap = (1 - (((cmax / (bexp + 1)) - xn) / (cmax / (bexp + 1)))) * pet
    xn = max(xn - evap, 0)  # update state

    return ER1, ER2, xn
end

"""
    randomparams(xslow::Float64, inflow::Float64, Rs::Float64)

Linear reservoir routing function
"""
function linearreservoir(xslow::Float64, inflow::Float64, Rs::Float64)
    # Linear reservoir routing
    xslow = (1 - Rs) * xslow + (1 - Rs) * inflow
    outflow = (Rs / (1 - Rs)) * xslow
    return xslow, outflow
end

"""
    hargreaves(forcings::DataFrame; tmincol::Symbol = :tmin, tmaxcol::Symbol = :tmax, dtcol::Symbol = :datetime)

Function to calculate PET using Hargreves equation.
"""
function hargreaves(
    forcings::DataFrame;
    tmincol::Symbol = :tmin,
    tmaxcol::Symbol = :tmax,
    dtcol::Symbol = :datetime,
)

    dts = forcings[!, dtcol]
    tmin = forcings[!, tmincol]
    tmax = forcings[!, tmaxcol]

    hargreaves(tmin, tmax, dts)
end

"""
    hargreaves(tmin::AbstractArray, tmax::AbstractArray, dts::AbstractArray)

Function to calculate PET using Hargreves equation.
"""
function hargreaves(
    tmin::AbstractArray,
    tmax::AbstractArray,
    dts::AbstractArray,
)
    hargreaves.(tmin, tmax, dts)
end

"""
    hargreaves(tmin::Real, tmax::Real, dts::Date)

Function to calculate PET using Hargreves equation.
"""
function hargreaves(tmin::Real, tmax::Real, dts::Date)
    Gsc = 367
    lhov = 2.257

    doy = dayofyear(dts)

    tavg = mean((tmin, tmax))

    b = 2 * pi * (Int16(doy) / 365)
    Rav =
        1.00011 +
        0.034221 * cos(b) +
        0.00128 * sin(b) +
        0.000719 * cos(2 * b) +
        0.000077 * sin(2 * b)
    Ho = ((Gsc * Rav) * 86400) / 1e6

    eto = (0.0023 * Ho * (tmax - tmin)^0.5 * (tavg + 17.8))

    return eto
end

"""
    simulate(forcings::DataFrame; precipcol::Symbol = :precip, petcol::Symbol = :pet, kwargs...)

Function to simulate discharge using Hymod rainfall-runoff model.
"""
function simulate(
    forcings::DataFrame;
    precipcol::Symbol = :precip,
    petcol::Symbol = :pet,
    kwargs...,
)
    p = forcings[:, precipcol]
    e = forcings[:, petcol]

    simulate(p, e; kwargs...)

end

"""
    simulate(precip::AbstractArray, pet::AbstractArray; initflow::Bool = true, cmax::Float64 = 1.0, bexp::Float64 = 0.0, alpha::Float64 = 0.2, ks::Float64 = 0.01, kq::Float64 = 0.5)

Function to simulate discharge using Hymod rainfall-runoff model
"""
function simulate(
    precip::AbstractArray,
    pet::AbstractArray;
    initflow::Bool = true,
    cmax::Float64 = 1.0,
    bexp::Float64 = 0.0,
    alpha::Float64 = 0.2,
    ks::Float64 = 0.01,
    kq::Float64 = 0.5,
)

    lt_to_m::Float64 = 0.001

    n_iters = length(precip)

    # HYMOD PROGRAM IS SIMPLE RAINFALL RUNOFF MODEL
    x_loss::Float64 = 0.0
    # Initialize slow tank state
    # value of 0 init flow works ok if calibration data starts with low discharge
    x_slow::Float64 = initflow ? (2.3503 / (ks * 22.5)) : 0.0
    # Initialize state(s) of quick tank(s)
    x_quick = zeros(Float64, 3)
    outflow = zeros(Float64, n_iters)
    output = zeros(Float64, n_iters)
    # START PROGRAMMING LOOP WITH DETERMINING RAINFALL - RUNOFF AMOUNTS

    # prealloc values within loops
    Pval::Float64 = 0.0
    PETval::Float64 = 0.0
    ER1::Float64 = 0.0
    ER2::Float64 = 0.0
    UQ::Float64 = 0.0
    US::Float64 = 0.0
    QS::Float64 = 0.0
    inflow::Float64 = 0.0

    @inbounds for i = 1:n_iters
        Pval = precip[i]
        PETval = pet[i]
        # Compute excess precipitation and evaporation
        ER1, ER2, x_loss = precipexcess(x_loss, cmax, bexp, Pval, PETval)
        # Calculate total effective rainfall
        ET = ER1 + ER2
        #  Now partition ER between quick and slow flow reservoirs
        UQ = alpha * ET
        US = (1 - alpha) * ET
        # Route slow flow component with single linear reservoir
        x_slow, QS = linearreservoir(x_slow, US, ks)
        # Route quick flow component with linear reservoirs
        inflow = UQ

        for j in eachindex(x_quick)
            # Linear reservoir
            x_quick[j], outflow = linearreservoir(x_quick[j], inflow, kq)
            inflow = outflow
        end

        # Compute total flow for timestep
        output[i] = ((QS + outflow) / lt_to_m)
    end

    output[output.<0] .= 0

    return output
end

"""
    calibrate(forcing::DataFrame, paramspace::Dict, nsamples::Int64; precipcol::Symbol = :precip, petcol::Symbol = :pet, obscol::Symbol = :obs, savefinal::Bool = false)

Function to calibrate Hymod rainfall-runoff model
"""
function calibrate(
    forcing::DataFrame,
    paramspace::Dict,
    nsamples::Int64;
    precipcol::Symbol = :precip,
    petcol::Symbol = :pet,
    obscol::Symbol = :obs,
    savefinal::Bool = false,
)
    # implementation of a monte carlo sampling

    obs = forcing[!, obscol]

    params = Dict{Symbol,Array}()

    keyList = collect(keys(paramspace))

    for k in keyList
        p = paramspace[k]
        minV = p[:lower]
        maxV = p[:upper]
        params[k] = collect(rand(Uniform(minV, maxV), nsamples))
    end

    nIter = nsamples
    losses = zeros(nIter)

    @info "Calibrating model with $nIter iterations..."
    for i = 1:nIter
        vals = [params[k][i] for k in keyList]
        pars = Dict(keyList .=> vals)
        q = simulate(forcing, precipcol = :precip, petcol = :pet; pars...)
        # only test after one year of spinup
        losses[i] = nse(q[365:end], obs[365:end])
    end

    # check to make sure that the loss metric is a valid number
    losses = convert(Vector{Union{Missing,Float64}}, losses)
    losses[map(isnan, losses)] .= missing

    finalLoss, idx = findmax(skipmissing(losses))
    vals = [params[k][idx] for k in keyList]
    finalPars = Dict(keyList .=> vals)
    finalQ = simulate(forcing, precipcol = :precip, petcol = :pet; finalPars...)

    if savefinal
        t = now()
        finalPars[:loss] = finalLoss
        jsondata = Dict(String(k) => v for (k, v) in finalPars)
        @debug (jsondata)
        open("hymod_calibration_results_$t.json", "w") do io
            JSON.print(io, jsondata, 4)
        end
    end

    return finalQ, finalPars, finalLoss

end

"""
    nse(y_true::AbstractArray, y_pred::AbstractArray)

Function to calculate Nash-Sutcliffe model efficiency coefficient
"""
function nse(y_true::AbstractArray, y_pred::AbstractArray)
    # sum square difference between observed and simulated
    numerator = sum((y_true - y_pred) .^ 2)
    # sum square difference between observed and observed mean
    denominator = sum((y_true .- mean(y_true)) .^ 2)

    return 1 - (numerator / denominator)
end

end # module


@info "Running tests"
using SpectralDistances# Distributions
using Test, LinearAlgebra, Statistics, Random, ControlSystems, InteractiveUtils, SparseArrays # For subtypes
using DSP, Distances, DoubleFloats
import GLPK, Convex, JuMP, SCS

using SpectralDistances: ngradient, nhessian, njacobian, polyconv, hproots, rev

@testset "SpectralDistances.jl" begin

    @testset "complete distmat" begin
        @info "Testing complete distmat"

        P = randn(2,40)
        D = pairwise(SqEuclidean(), P, dims=2)
        W = rand(size(D)...) .> 0.3 # Create a random mask
        W = (W + W') .> 0           # It makes sense for the mask to be symmetric
        W[diagind(W)] .= true
        D0 = W .* D                 # Remove missing entries

        D2,S = complete_distmat(D0, W)

        @test (norm(D-D2)/norm(D)) < 1e-5
        @test (norm(W .* (D-D2))/norm(D)) < 1e-5

        X  = Diagonal(sqrt.(S.S))*S.Vt

        # Verify that reconstructed `X` is correct up to rotation and translation
        A = [X[1:2,:]' ones(size(D,1))]
        P2 = (A*(A \ P'))'
        @test norm(P-P2)/norm(P) < 1e-3

    end

    @testset "Convolutional" begin
        @info "Testing Convolutional"
        include("test_convolutional.jl")
    end

    @testset "DTW" begin
        @info "Testing DTW"
        include("test_dtw.jl")
    end

    @testset "Model creation" begin
        @info "Testing Model creation"

        @testset "AR" begin
            @info "Testing AR"

            a = randn(3)
            m = AR(a)
            @test_nowarn show(m)
            @test_nowarn show([m, m])
            @test m.a == a
            @test length(m.p) == 2

            m = AR(Continuous(), a)
            @test m.ac == a
            @test length(m.pc) == 2

            r = ContinuousRoots(hproots(rev(a)))
            m = AR(r)
            @test m.pc == r

            r = DiscreteRoots(hproots(rev(a)))
            m = AR(r)
            @test m.p == r
        end

        @testset "ARMA" begin
            @info "Testing ARMA"

            a = randn(3)
            b = [1, 0.1]
            m = ARMA(b,a)
            @test s1(m.a) == s1(a)
            @test length(m.p) == 2
            @test m.b == b
            @test length(m.z) == 1

            m = ARMA(Continuous(), b,a)
            @test s1(m.ac) == s1(a)
            @test length(m.pc) == 2
            @test m.bc == b
            @test length(m.zc) == 1

            r = ContinuousRoots(hproots(rev(a)))
            z = ContinuousRoots(hproots(rev(b)))
            m = ARMA(z,r)
            @test m.pc == r
            @test m.zc == z

            r = DiscreteRoots(hproots(rev(a)))
            z = DiscreteRoots(hproots(rev(b)))
            m = ARMA(z,r)
            @test m.p == r
            @test m.z == z
        end

    end

    @testset "time" begin
        @info "Testing time"
        include("test_time.jl")
    end

    @testset "plot" begin
        @info "Testing plot"
        include("test_plotting.jl")
    end

    @testset "Solvers" begin
        @info "Testing Solvers"

        C = Float64.(Matrix(I(2)))
        a = [1., 0]
        b = [0, 1.]
        Γ,u,v = @inferred sinkhorn(C,a,b,β=0.01)
        @test Γ ≈ [0 1; 0 0]

        a = [0.5, 0.5]
        b = [0, 1.]
        C = Float64.(Matrix(I(2)))
        Γ,u,v = @inferred sinkhorn(C,a,b,β=0.1)
        @test Γ ≈ [0 0.5; 0 0.5]

        a = [1., 0]
        b = [0, 1.]
        C = Float64.(Matrix(I(2)))
        Γ,u,v = @inferred sinkhorn_log(C,a,b,β=0.01)
        @test Γ ≈ [0 1; 0 0]

        a = [0.5, 0.5]
        b = [0, 1.]
        C = Float64.(Matrix(I(2)))
        Γ,u,v = @inferred sinkhorn_log(C,a,b,β=0.01)
        @test Γ ≈ [0 0.5; 0 0.5]

        a = [0.5, 0.5]
        b = [0, 1.]
        C = Float64.(Matrix(I(2)))
        Γ,u,v = @inferred sinkhorn_log!(C,a,b,β=0.01)
        @test Γ ≈ [0 0.5; 0 0.5]

        a = [1., 0]
        b = [0, 1.]
        C = Float64.(Matrix(I(2)))
        Γ,u,v = @inferred IPOT(C,a,b,β=1)
        @test Γ ≈ [0 1; 0 0]

        a = [0.5, 0.5]
        b = [0, 1.]
        C = Float64.(Matrix(I(2)))
        Γ,u,v = @inferred IPOT(C,a,b)
        @test Γ ≈ [0 0.5; 0 0.5]

        a = [1., 0]
        b = [0, 1.]
        C = Float64.(Matrix(I(2)))
        Γ,u,v = @inferred sinkhorn_unbalanced(C,a,b,Balanced(),β=0.01)
        @test Γ ≈ [0 1; 0 0]

        a = [0.5, 0.5]
        b = [0, 1.]
        C = Float64.(Matrix(I(2)))
        Γ,u,v = @inferred sinkhorn_unbalanced(C,a,b,Balanced())
        @test Γ ≈ [0 0.5; 0 0.5]

        a = [1., 0]
        b = [0, 1.]
        C = Float64.(Matrix(I(2)))
        Γ,u,v = ot_jump(C,a,b)
        @test Γ ≈ [0 1; 0 0]

        a = [0.5, 0.5]
        b = [0, 1.]
        C = Float64.(Matrix(I(2)))
        Γ,u,v = ot_jump(C,a,b)
        @test Γ ≈ [0 0.5; 0 0.5]

        a = [1., 0]
        b = [0, 1.]
        C = Float64.(Matrix(I(2)))
        Γ,u,v = ot_convex(C,a,b)
        @test Γ ≈ [0 1; 0 0]

        a = [0.5, 0.5]
        b = [0, 1.]
        C = Float64.(Matrix(I(2)))
        Γ,u,v = ot_convex(C,a,b)
        @test Γ ≈ [0 0.5; 0 0.5]

        a = [0.5, 0.5]
        b = [0, 1.]
        C = Float64.(Matrix(I(2)))
        Γ,u,v = ot_convex(C,a,b)
        @test Γ ≈ [0 0.5; 0 0.5]

        models = examplemodels(2)
        Γ = transport_plan(OptimalTransportRootDistance(domain=Continuous(), weight=unitweight), models...; tol=1e-4)
        @test sum(Γ) ≈ 1 atol=0.001

    end

    @testset "gradients" begin
        @info "Testing gradients"
        include("test_diff.jl")
    end

    @testset "Energy" begin
        @info "testing Energy"
        for σ² = [0.1, 1., 2., 3]
            m = AR(ContinuousRoots([-1.]), σ²)
            e = @inferred spectralenergy(Continuous(),m)
            @test e ≈ σ²
            x = sqrt(σ²)randn(10000)
            m = LS(na=10)(x)
            @test spectralenergy(Continuous(),m) ≈ σ² atol=0.15
            m = TLS(na=10)(x)
            @test spectralenergy(Continuous(),m) ≈ σ² atol=0.15
        end
        # y = filt(numvec(Discrete(),m), denvec(Discrete(),m), x)
        # @test var(y) ≈ var(x)
    end

    @testset "Model estimation" begin
        y = sin.(0:0.1:100)
        fm = @inferred LS(na=2, λ=0)
        m = @inferred fm(y)
        @test fm(m) == m
        @test imag.(m.pc) ≈ [-0.1, 0.1] rtol=1e-4
        mc = @inferred SpectralDistances.change_precision(Float32, m)
        @test m ≈ mc
        @test eltype(mc.a) == Float32

        fm = @inferred TLS(na=2)
        m  = @inferred fm(y)
        @test imag.(m.pc) ≈ [-0.1, 0.1] rtol=1e-4

        fm = IRLS(na=2)
        m = fm(y)
        @test_broken imag.(m.pc) ≈ [-0.1, 0.1] rtol=1e-4

        y = sin.(0:0.1:1000) .+ 0.01 .*randn.()
        fm = PLR(na=2, nc=1)
        m = fm(y)
        @test imag.(log(m.p)) ≈ [-0.1, 0.1] rtol=1e-3

        y = sin.(0:0.1:100) .+ sin.(2 .* (0:0.1:100) .+ 0.3)
        fm = @inferred LS(na=4, λ=0)
        m = @inferred fm(y)
        @test imag.(m.pc) ≈ [-0.2, -0.1, 0.1, 0.2] rtol=1e-1

        fm = TLS(na=4)
        m = fm(y)
        @test spectralenergy(Continuous(), m) ≈ var(y) rtol=1e-3
        @test imag.(m.pc) ≈ [-0.2, -0.1, 0.1, 0.2] rtol=1e-4

        fm = IRLS(na=4)
        m = fm(y)
        @test @inferred(spectralenergy(Continuous(), m)) ≈ var(y) rtol=1e-3
        @test_broken imag.(m.pc) ≈ [-0.2, -0.1, 0.1, 0.2] rtol=1e-4

        y = sin.(0:0.1:1000) .+ sin.(2 .* (0:0.1:1000)) .+ 0.001 .*randn.()
        fm = @inferred PLR(na=4, nc=1, λ=0.0)
        m = fm(y)
        @test imag.(m.pc) ≈ [-0.2, -0.1, 0.1, 0.2] rtol=0.6

    end

@testset "modeldistance" begin
    t = 1:300
    ϵ = 1e-7
    fitmethod = LS
    for fitmethod in [LS, TLS]
        @info "Testing fitmethod $(string(fitmethod))"
        ls_loss = ModelDistance(fitmethod(na=2), EuclideanRootDistance(domain=Discrete()))
        @test @inferred(ls_loss(sin.(t), sin.(t))) < ϵ
        @test ls_loss(sin.(t), -sin.(t)) < ϵ # phase flip invariance
        @test ls_loss(sin.(t), sin.(t .+ 1)) < ϵ # phase shift invariance
        @test ls_loss(sin.(t), sin.(t .+ 0.1)) < ϵ # phase shift invariance
        @test ls_loss(10sin.(t), sin.(t .+ 1)) < ϵ # amplitude invariance
        @test ls_loss(sin.(t), sin.(1.1 .* t)) < 0.2 # small frequency shifts gives small errors
        @test ls_loss(sin.(0.1t), sin.(1.1 .* 0.1t)) < 0.1 # small frequency shifts gives small errors

        @test ls_loss(sin.(t), sin.(1.1 .* t)) ≈ ls_loss(sin.(0.1t), sin.(0.2t)) rtol=1e-2  # frequency shifts of relative size should result in the same error, probably only true for p=1
        ls_loss = ModelDistance(fitmethod(na=10), OptimalTransportRootDistance(domain=Discrete()))
        @test @inferred(ls_loss(filtfilt(ones(10),[10], randn(1000)), filtfilt(ones(10),[10], randn(1000)))) < 0.1 # Filtered through same filter, this test is very non-robust for TLS
        @test ls_loss(filtfilt(ones(10),[10], randn(1000)), filtfilt(ones(10),[10], randn(1000))) < ls_loss(filtfilt(ones(4),[4], randn(1000)), filtfilt(ones(10),[10], randn(1000))) # Filtered through different filters, this test is not robust
    end
    @testset "PLR" begin
        fitmethod = PLR
        @info "Testing fitmethod $(string(fitmethod))"
        t = 1:1000 .+ 0.01 .* randn.()
        ϵ = 0.01
        @info "Creating model distance"
        ls_loss = ModelDistance(fitmethod(na=2,nc=1), EuclideanRootDistance(domain=Discrete()))
        @info "Calc loss"
        @test ls_loss(sin.(t), sin.(t)) < ϵ
        @test ls_loss(sin.(t), -sin.(t)) < ϵ # phase flip invariance
        @test ls_loss(sin.(t), sin.(t .+ 1)) < ϵ # phase shift invariance
        @test ls_loss(sin.(t), sin.(t .+ 0.1)) < ϵ # phase shift invariance
        @info "Halfway"
        @test ls_loss(10sin.(t), sin.(t .+ 1)) < ϵ # amplitude invariance
        @test ls_loss(sin.(t), sin.(1.1 .* t)) < 0.2 # small frequency shifts gives small errors
        @test ls_loss(sin.(0.1t), sin.(1.1 .* 0.1t)) < 0.1 # small frequency shifts gives small errors

        @test ls_loss(sin.(t), sin.(1.1 .* t)) ≈ ls_loss(sin.(0.1t), sin.(0.2t)) rtol=1e-3  # frequency shifts of relative size should result in the same error, probably only true for p=1
        @info "Almost there"
        # ls_loss = ModelDistance(fitmethod(na=10,nc=2), OptimalTransportRootDistance(domain=Discrete())) # This test segfaulted sometimes, couldn't figure out the problem
        # @test ls_loss(randn(1000), randn(1000)) > 0.05
        # @info "Done"
        # @test ls_loss(filtfilt(ones(10),[10], randn(1000)), filtfilt(ones(10),[10], randn(1000))) < 1 # Filtered through same filter, this test is very non-robust for TLS
        # @test ls_loss(filtfilt(ones(10),[10], randn(1000)), filtfilt(ones(10),[10], randn(1000))) < ls_loss(filtfilt(ones(4),[4], randn(1000)), filtfilt(ones(10),[10], randn(1000))) # Filtered through different filters, this test is not robust
    end
end

@testset "discrete_grid_transportplan" begin
    x = [1.,0,0]
    y = [0,0.5,0.5]

    g = @inferred discrete_grid_transportplan(x,y)
    @test sum(g,dims=1)[:] == y
    @test sum(g,dims=2)[:] == x

    # test robustness for long vectors
    x = s1(rand(1000))
    y = s1(rand(1000))
    D = [abs2(abs(i-j)/1000) for i in eachindex(x), j in eachindex(y)]
    g = discrete_grid_transportplan(x,y)
    @test sum(g,dims=1)[:] ≈ y
    @test sum(g,dims=2)[:] ≈ x
    @test dot(g, D) ≈ discrete_grid_transportcost(x,y) rtol=0.01

    g = discrete_grid_transportplan(y,x)
    @test sum(g,dims=1)[:] ≈ x
    @test sum(g,dims=2)[:] ≈ y
    @test dot(g, D) ≈ discrete_grid_transportcost(y,x) rtol=0.01

    # test exception for unequal masses
    x = s1(rand(Float32,1000))
    y = rand(Float32,1000)
    @test_throws ErrorException discrete_grid_transportplan(x,y)
    @test_throws ErrorException discrete_grid_transportplan(y,x)

    @test_throws ErrorException discrete_grid_transportcost(x,y)
    @test_throws ErrorException discrete_grid_transportcost(y,x)

end

# p = randn(20)
# @btime Zygote.gradient(p) do p
#     r = riroots(p)
#     sum([r[1]; r[2]])
# end
#
#
# jacobian(fd, p->begin
#     r = riroots(p)
#     sum([r[1]; r[2]])
# end, p)

@testset "eigenvalue manipulations" begin
    @test @inferred(SpectralDistances.reflectd(2)) ≈ 0.5 + 0.0im
    @test @inferred(SpectralDistances.reflectd(complex(0,2))) ≈ 0 + 0.5im
    @test @inferred(SpectralDistances.reflectd(complex(0,-2))) ≈ 0 - 0.5im

    e = @inferred SpectralDistances.hproots(randn(7))
    ed = @inferred DiscreteRoots(e)
    @test real(e) == real.(e)
    @test real(ed) == real.(ed)
    @test issorted(ed.r, by=angle)
    @test all(<(1) ∘ abs, ed)
    ec = log(ed)
    @test issorted(ec, by=SpectralDistances.imageigsortby)
    @test all(<(0) ∘ real, ec)
    @test @inferred(domain_transform(Continuous(), ed)) isa ContinuousRoots
    @test domain_transform(Continuous(), ed) ≈ SpectralDistances.eigsort(SpectralDistances.reflectc.(log.(ed))) ≈ ContinuousRoots(ed)
    @test @inferred(domain_transform(Discrete(), ed)) == ed
    @test domain(ed) isa Discrete
    @test domain(ContinuousRoots(ed)) isa Continuous
    @test log(ed) isa ContinuousRoots

    @test all(<(1) ∘ abs, reflect(ed))
    @test all(<(0) ∘ real, reflect(ec))

    E = embedding(ec)
    @test length(E) == 12
    @test E == embedding(Vector, ec)
    E = embedding(Matrix, ec)
    @test size(E) == (6,2)
    E = embedding(Matrix, ec, false)
    @test size(E) == (3,2)

    @test_logs (:error, r"real poles") embedding(ContinuousRoots([0,1]), false)

    Random.seed!(0)
    @test SpectralDistances.determine_domain(0.1randn(10)) isa Discrete
    @test SpectralDistances.determine_domain(randn(10).-4) isa Continuous
    @test_throws Exception SpectralDistances.determine_domain(0.1randn(10).-0.3)

    @testset "weight functions" begin
        @info "Testing weight functions"
        r = @inferred ContinuousRoots(randn(ComplexF64, 4))
        m = @inferred AR(r, 2.0)
        @test sum(unitweight(r)) == 1
        @test sum(unitweight(m)) == 1
        @test length(unitweight(r)) == 4
        @test length(unitweight(m)) == 4

        @test m.b^2*sum(residueweight(r)) ≈ sum(residueweight(m)) # Calling residueweight with a model should mutiply energy by b²
        @test length(residueweight(r)) == 4
        @test length(residueweight(m)) == 4
    end

end

@testset "ControlSystems interoperability" begin
    @info "Testing ControlSystems interoperability"
    m = @inferred AR(ContinuousRoots([-1]))
    g = tf(1,[1.,1])
    @test tf(m) == g
    @test_broken tf(m,1) == c2d(g,1)
    @test m*g == g*g
    @test ARMA(g).ac == [1,1]
    @test ARMA(g).bc == [1]
    @test @inferred(denvec(Continuous(), m)) == denvec(g)[1]
    @test numvec(Continuous(), m) == numvec(g)[1]
    @test @inferred(pole(Continuous(), m)) == pole(g)
    @test all(ControlSystems.bode(m) .≈ bode(g))
    @test all(ControlSystems.nyquist(m) .≈ nyquist(g))
    @test ControlSystems.freqresp(m, exp10.(LinRange(-1, 1, 10))) ≈ freqresp(g, exp10.(LinRange(-1, 1, 10)))
    @test all(ControlSystems.step(m, 10) .≈ step(g, 10))

    bodeplot(m)
    pzmap(m)
    nyquistplot(m)

end

@testset "polynomial acrobatics" begin
    a = randn(5)
    b = randn(5)
    @test @inferred(polyconv(a,b)) ≈ DSP.conv(a,b)

    a = randn(5)
    b = randn(10)
    @test polyconv(a,b) ≈ DSP.conv(a,b)
    @test polyconv(b,a) ≈ DSP.conv(b,a)

    a[1] = 1
    @test @inferred(roots2poly(roots(reverse(a)))) ≈ a

end

@testset "preprocess roots with residue weight" begin
    rd  = @inferred EuclideanRootDistance(domain=Continuous(), weight=residueweight)
    m = @inferred AR([1., -0.1])
    @test @inferred(rd(m,m)) < eps()
    @test @inferred(SpectralDistances.preprocess_roots(rd, m))[] ≈ -2.3025850929940455
end

@testset "residues and roots" begin

    # @test inv(1:3) == 1:3
    # @test inv([1,1,2]) == [1,1,3]
    # @test inv([1,2,2]) == [1,1,2]

    a = randn(5); a[1]=1;a
    r = roots(reverse(a))
    @test roots2poly(r) ≈ a

    b,a = 1,0.1randn(5)
    a[1] = 1
    r   = @inferred SpectralDistances.hproots(reverse(a))
    G   = tf(b,a)
    w   = im
    F = evalfr(G,w)[]
    a2 = roots2poly(r)
    @assert a ≈ a2
    roots(reverse(a2))
    @test sum(r) do r
        powers = length(a)-1:-1:1
        ap     = powers.*a[1:end-1]
        G      = tf(b,ap)
        evalfr(G,r)[] *1/(w - r)
    end ≈ F

    res = @inferred residues(a,1)
    @test sum(eachindex(r)) do i
        res[i]/(w-r[i])
    end ≈ F

    @test abs2.(res) ≈ SpectralDistances.abs2residues!(zeros(length(a)-1),a,1)

    @test prod(eachindex(r)) do i
        1/(w-r[i])
    end ≈ F

    n = 4
    a1 = randn(n+1); a1[1] = 1
    r1 = SpectralDistances.reflectc.(hproots(reverse(a1)))
    r1 = complex.(0.01real.(r1), imag.(r1))
    r1 = @inferred SpectralDistances.normalize_energy(ContinuousRoots(r1))

    a2 = randn(n+1); a2[1] = 1
    r2 = SpectralDistances.reflectc.(hproots(reverse(a2)))
    r2 = complex.(0.01real.(r2), imag.(r2))
    r2o = @inferred SpectralDistances.normalize_energy(ContinuousRoots(r2))

    m1,m2 = AR(a1), AR(a2)

    dist2 = @inferred RationalOptimalTransportDistance(domain=Continuous(), p=2, interval=(-20.,20.))
    f    = w -> evalfr(SpectralDistances.domain(dist2), SpectralDistances.magnitude(dist2), w, m1)
    @test SpectralDistances.c∫(f,dist2.interval...)[end] ≈ spectralenergy(Continuous(), m1) rtol=1e-3
    f    = w -> evalfr(SpectralDistances.domain(dist2), SpectralDistances.magnitude(dist2), w, m2)
    @test SpectralDistances.c∫(f,dist2.interval...)[end] ≈ spectralenergy(Continuous(), m2) rtol=1e-3

    @test spectralenergy(Continuous(), AR(ContinuousRoots([-1.]))) ≈ spectralenergy(tf(1,[1,1])) rtol=1e-3

    m1,m2 = ARMA([1], a1), ARMA([1], a2)

    f    = w -> evalfr(SpectralDistances.domain(dist2), SpectralDistances.magnitude(dist2), w, m1)
    @test_broken SpectralDistances.c∫(f,dist2.interval...)[end] ≈ spectralenergy(Continuous(), m1) rtol=1e-3 # The error is in the integration, spectgralenergy produces same result as for AR above
    f    = w -> evalfr(SpectralDistances.domain(dist2), SpectralDistances.magnitude(dist2), w, m2)
    @test_broken SpectralDistances.c∫(f,dist2.interval...)[end] ≈ spectralenergy(Continuous(), m2) rtol=1e-3

end

@testset "d(m,m)" begin
    a = roots2poly([0.9 + 0.1im, 0.9 - 0.1im])
    m = AR(a)
    for D in [  subtypes(SpectralDistances.AbstractDistance);
                subtypes(SpectralDistances.AbstractRootDistance);
                subtypes(SpectralDistances.AbstractCoefficientDistance)]
        (!isempty(methods(D)) && (:domain ∈ fieldnames(D))) || continue
        d = @inferred D(domain=Continuous())
        println(D)
        @test @inferred(d(m,m)) < eps() + 0.001*(d isa OptimalTransportRootDistance)
        d isa Union{RationalOptimalTransportDistance, RationalCramerDistance} && continue
        d = D(domain=Discrete())
        println(D)
        @test d(m,m) < eps() + 0.001*(d isa OptimalTransportRootDistance)
    end
end

@testset "d(m,m̃)" begin
    a1 = [1,-0.1,0.8]
    m1 = @inferred AR(a1)
    a2 = [1,-0.1,0.801]
    m2 = AR(a2)
    for D in [  subtypes(SpectralDistances.AbstractDistance);
                subtypes(SpectralDistances.AbstractRootDistance);
                subtypes(SpectralDistances.AbstractCoefficientDistance)]
        (!isempty(methods(D)) && (:domain ∈ fieldnames(D))) || continue
        d = D(domain=Continuous())
        println(D)
        # @show d(m1,m2)
        @test @inferred(d(m1,m2)) > 1e-10
        d isa Union{RationalOptimalTransportDistance, RationalCramerDistance} && continue
        d = D(domain=Discrete())
        println(D)
        # @show d(m1,m2)
        @test d(m1,m2) > 1e-10
    end
end

@testset "distmat" begin
    e = complex.(randn(3), randn(3))
    D = @inferred SpectralDistances.distmat(SqEuclidean(), e)
    @test issymmetric(D)
    @test tr(D) == 0
    @test @inferred(SpectralDistances.distmat_euclidean(e,e)) ≈ D

    A = randn(10,10)
    A[diagind(A)] .= rand(10);
    A = A + A';
    B = SparseMatrixCSC(A)

    symmetrize!(A)
    symmetrize!(B)
    @test A == B
    @test issymmetric(A)
end

@testset "Welch" begin
    x1 = @inferred SpectralDistances.bp_filter(randn(3000), (0.01,0.1))
    x2 = SpectralDistances.bp_filter(randn(3000), (0.01,0.12))
    x3 = SpectralDistances.bp_filter(randn(3000), (0.01,0.3))
    dist = @inferred WelchOptimalTransportDistance(p=1)
    @test @inferred(dist(x1,x2)) < dist(x1,x3)
    dist = WelchOptimalTransportDistance(p=2)
    @test dist(x1,x2) < dist(x1,x3)

    dist = @inferred WelchLPDistance(p=1)
    @test @inferred(dist(x1,x2)) < dist(x1,x3)
    dist = WelchLPDistance(p=2)
    @test dist(x1,x2) < dist(x1,x3)

    dist = @inferred EnergyDistance()
    @test @inferred(dist(x1,x2)) < dist(x1,x3)
end

@testset "Bures" begin
    fm = LS(na=4)
    x1 = SpectralDistances.bp_filter(randn(3000), (0.01,0.1))  |> fm
    x2 = SpectralDistances.bp_filter(randn(3000), (0.01,0.12)) |> fm
    x3 = SpectralDistances.bp_filter(randn(3000), (0.01,0.3))  |> fm
    dist = BuresDistance()
    @test dist(x1,x2) < dist(x1,x3)
end

@testset "KernelWassersteinRootDistance" begin
    fm = LS(na=4)
    x1 = SpectralDistances.bp_filter(randn(3000), (0.01,0.1))  |> fm
    x2 = SpectralDistances.bp_filter(randn(3000), (0.01,0.12)) |> fm
    x3 = SpectralDistances.bp_filter(randn(3000), (0.01,0.3))  |> fm
    dist = @inferred KernelWassersteinRootDistance(domain=Continuous())
    @test @inferred(dist(x1,x2)) < dist(x1,x3)
end

@testset "ClosedForm" begin
    x1 = SpectralDistances.bp_filter(randn(3000), (0.01,0.1))
    x2 = SpectralDistances.bp_filter(randn(3000), (0.01,0.12))
    x3 = SpectralDistances.bp_filter(randn(3000), (0.01,0.3))
    fm = LS(na=4)
    dist = ModelDistance(fm, RationalOptimalTransportDistance(domain=Continuous(), p=1, interval=(-15., 15)))
    @test dist(x1,x2) < dist(x1,x3)
    dist = RationalOptimalTransportDistance(domain=Continuous(), p=1, interval=(0., 15.))
    @test dist(fm(x1),welch_pgram(x2)) < dist(fm(x1),welch_pgram(x3))
    @test dist(fm(x1),welch_pgram(x2)) < dist(fm(x1),welch_pgram(x3))

    dist = ModelDistance(fm, RationalOptimalTransportDistance(domain=Continuous(), p=2, interval=(-15., 15)))
    @test dist(x1,x2) < dist(x1,x3)
    dist = RationalOptimalTransportDistance(domain=Continuous(), p=2, interval=(0., 15.))
    @test dist(fm(x1),welch_pgram(x2)) < dist(fm(x1),welch_pgram(x3))

    w = LinRange(0, 2pi*15, 3000)
    dist = ModelDistance(fm, RationalOptimalTransportDistance(domain=Continuous(), p=1, interval=(0., 15)))
    ddist = ModelDistance(fm, DiscretizedRationalDistance(w, sqrt.(SpectralDistances.distmat_euclidean(w,w))))
    @test ddist(x1, x1) == 0
    @test_broken ddist(x1, x2) ≈ dist(x1, x2)

    w = LinRange(0, 2pi*15, 3000)
    dist = ModelDistance(fm, RationalOptimalTransportDistance(domain=Continuous(), p=2, interval=(0., 15)))
    ddist = ModelDistance(fm, DiscretizedRationalDistance(w, (SpectralDistances.distmat_euclidean(w,w))))
    @test ddist(x1, x1) == 0
    @test_broken ddist(x1, x2) ≈ dist(x1, x2)

end

@testset "Histogram" begin
    x1 = randn(100)
    x2 = randn(100) .+ 1
    x3 = randn(100) .+ 2
    dist = SpectralDistances.OptimalTransportHistogramDistance(p=1)
    @test dist(x1,x2) < dist(x1,x3)
    dist = SpectralDistances.OptimalTransportHistogramDistance(p=2)
    @test dist(x1,x2) < dist(x1,x3)
end

@testset "Theoretical scaling from paper" begin
    n = 4
    a1 = randn(n+1); a1[1] = 1
    r1 = SpectralDistances.reflectc.(roots(reverse(a1)))
    r1 = ContinuousRoots(complex.(0.01real.(r1), imag.(r1)))
    # r1 = SpectralDistances.normalize_energy(ContinuousRoots(r1))

    a2 = randn(n+1); a2[1] = 1
    r2 = SpectralDistances.reflectc.(roots(reverse(a2)))
    r2 = ContinuousRoots(complex.(0.01real.(r2), imag.(r2)))
    # r2 = SpectralDistances.normalize_energy(ContinuousRoots(r2))

    function scaleddist(alpha,p)
        r1a,r2a = ContinuousRoots(alpha.*r1), ContinuousRoots(alpha.*r2)
        m1,m2 = AR(r1a), AR(r2a)
        # dist = RationalOptimalTransportDistance(domain=Continuous(), p=p, interval=(-20.,20.))
        rdist = EuclideanRootDistance(domain=Continuous(), weight=residueweight, p=p)
        rdist(m1,m2)
    end
    for α ∈ 1:0.5:3, p ∈ 1:3
        @test α^(1. -2n +p) * scaleddist(1,p) ≈ scaleddist(α,p)

        @test α^(1. -n)*residues(ContinuousRoots(r1)) ≈ residues(ContinuousRoots(α*r1))
        @test α^(1. -2n)*residueweight(r1) ≈ residueweight(α*r1)
    end

end

@testset "Interpolations" begin
    # NOTE: Only testing that it runs, not that it does the right thing. Have a look at the figure in the docs to deteermine if it works well
    n = 4
    r1 = complex.(-0.01 .+ 0.001randn(3), 2randn(3))
    r1 = ContinuousRoots([r1; conj.(r1)])
    r2 = complex.(-0.01 .+ 0.001randn(3), 2randn(3))
    r2 = ContinuousRoots([r2; conj.(r2)])
    r1,r2 = normalize_energy.((r1, r2))
    A1 = AR(r1)
    A2 = AR(r2)
    t = 0.1
    dist = RationalOptimalTransportDistance(domain=Continuous(), p=2, interval=(0., exp10(1.01)))
    interp = SpectralDistances.interpolator(dist, A1, A2)
    w = exp10.(LinRange(-1.5, 1, 300))
    for t = LinRange(0, 1, 7)
        Φ = clamp.(interp(w,t), 1e-10, 100)
    end

    rdist = EuclideanRootDistance(domain=Continuous(), p=2)
    interp = SpectralDistances.interpolator(rdist, A1, A2, normalize=false)
    w = exp10.(LinRange(-1.5, 1, 300))
    for t = LinRange(0, 1, 7)
        Φ = interp(w,t)
    end
end

@testset "barycenter" begin
    @info "Testing barycenter"
    include("test_barycenter.jl")
end

@testset "Slerp" begin
    a = n1([1,2,3])
    b = n1([4,4,3])
    c = SpectralDistances.slerp(a,b,0.5)
    @test c ≈ n1((a+b)/2)
end

@testset "Utils" begin
    X = randn(1000,10)
    x1,x2 = twoD(X)
    @test length(x1) == length(x2) == size(X,1)
    x1,x2,x3 = threeD(X)
    @test length(x1) == length(x3) == size(X,1)

    @test norm(n1(x1)) ≈ 1
    @test sum(s1(abs.(x1))) ≈ 1
    @test var(v1(x1)) ≈ 1
    @test mean(v1(x1)) ≈ 0 atol=sqrt(eps())
    @test mean(SpectralDistances.m1(x1)) ≈ 0 atol=0.08
    @test var(SpectralDistances.m1(x1)) ≈ 1 atol=0.2

    # @test norm(n1(X,1)) ≈ 1
    @test  all(sum(s1(abs.(X),1), dims=1) .≈ 1)
    @test  all(var(v1(X,1), dims=1) .≈ 1)
    @test  mean(v1(X,1)) ≈ 0 atol = sqrt(eps())
    @test  mean(SpectralDistances.m1(X,1)) ≈ 0 atol=0.08
    @test  var(SpectralDistances.m1(X,1)) ≈ 1 atol=0.2
end

@testset "Unbalanced transport" begin
    @info "Testing Unbalanced transport"

    fm = LS(na = 10)
    m1 = fm(filtfilt(ones(10), [10], randn(1000)))
    m2 = fm(filtfilt(ones(5), [5], randn(1000)))
    dist = OptimalTransportRootDistance(domain = Continuous(), p=1, divergence=Balanced())
    d1 = evaluate(dist,m1,m2)
    dist = OptimalTransportRootDistance(domain = Continuous(), p=1, divergence=KL(1.0))
    d2 = evaluate(dist,m1,m2)
    dist = OptimalTransportRootDistance(domain = Continuous(), p=1, divergence=KL(10.0))
    d3 = evaluate(dist,m1,m2)
    dist = OptimalTransportRootDistance(domain = Continuous(), p=1, divergence=KL(0.01))
    d4 = evaluate(dist,m1,m2)

    @test d1 > d3 > d2 > d4
end

end

##


push!(LOAD_PATH, expanduser("/home/phelipe/Documentos"))
using SimpleModels
using Trajectory
using Plots
using StaticArrays
using DifferentialEquations
pyplot()

#robot = Dof2([23.902, 1.285], [.45, .45], [0.091, 0.048], [1.266, 0.093])
#model = dinamic(robot)

Ts = 0.08
tend = 4.0
t0 = 0.0
x1,v1,a1,j1 = minimumjerkf(0.,0.,0.,t0,1.05,0.,0.,tend)
x2,v2,a2,j2 = minimumjerkf(0.,0.,0.,t0,1.57,0.,0.,tend)
#kp = SMatrix{2,2}(diagm([3300.,400.]));kv = SMatrix{2,2}(diagm([200.,20.]))
kp = SMatrix{2,2}(diagm([7000.,430.]));kv = SMatrix{2,2}(diagm([450.,30.]))

function myrobot(du, u, p, t)
    m=[23.902, 1.285]
    l= [.45, .45]
    r= [0.091, 0.048]
    I= [1.266, 0.093]
    g = 9.81
    θ = SVector{2}(u[1:2])
    dθ = SVector{2}(u[3:4])
    const cos1 = cos(θ[1])
    const cos2 = cos(θ[2])
    const sin1 = sin(θ[1])
    const sin2 = sin(θ[2])
    const cos12 = cos(θ[1]+θ[2])

    #Matriz de inércia
    m11 = m[1] * r[1]^2 + m[2]* (l[1]^2 + r[2]^2 + 2*l[1] * r[2] * cos2)+ I[1] + I[2]
    m12 = m[2] * (r[2]^2 + l[1] *r[2] * cos2) + I[2]
    m22 = m[2] * r[2]^2 + I[2]
    M = SMatrix{2,2}([m11 m12; m12 m22])

    #Matriz C
    h = -m[2] * l[1] * r[2] * sin2
    c11 = h * dθ[2]
    c12 = h * (dθ[1] + dθ[2])
    c21 = -h * dθ[1]
    c22 = 0
    C = SMatrix{2,2}([c11 c12; c21 c22])

    #Matriz gravidade
    g2 = m[2] * r[2] * g * cos12
    g1 = (m[1] * r[1] + m[2] * l[1]) * g * cos1 + g2
    G = SVector{2}([g1; g2])

    xr = [x1(t), x2(t)]
    vr = [v1(t), v2(t)]
    e = xr - θ
    de = vr - dθ
    tau = kp*e + kv*de

    du[1:2] = dθ
    du[3:4] = inv(M)*(tau - C * dθ - G)
end

tspan = (t0,tend)
prob2 = ODEProblem(myrobot,[0.,0.,0.,0.],tspan)
sol2 = solve(prob2,Tsit5(),saveat = Ts)

x,v,t,a,ta,j,tj = organize(2,sol2)

xd1 = map(x->x1(x),t)
xd2 = map(x->x2(x),t)
vd1 = map(x->v1(x),t)
vd2 = map(x->v2(x),t)
ad1 = map(x->a1(x),ta)
ad2 = map(x->a2(x),ta)
jd1 = map(x->j1(x),tj)
jd2 = map(x->j2(x),tj)

function plotx()
    p1 = plot(t,[xd1,x[1]], xlabel="tempo",ylabel="posição")
    p2 = plot(t,[xd2,x[2]], xlabel="tempo",ylabel="posição")
    plot(p1,p2, tittle="Posições")
end

function plotex()
    p1 = plot(t,xd1 - x[1], xlabel="tempo",ylabel="erro posição")
    p2 = plot(t,xd2 - x[2], xlabel="tempo",ylabel="erro posição")
    plot(p1,p2, tittle="Posições")
end

function plotv()
    p1 = plot(t,[vd1,v[1]], xlabel="tempo",ylabel="velocidade")
    p2 = plot(t,[vd2,v[2]], xlabel="tempo",ylabel="velocidade")
    plot(p1,p2, tittle="Velocidades")
end

function plotev()
    p1 = plot(t,vd1 - v[1], xlabel="tempo",ylabel="erro velocidade")
    p2 = plot(t,vd2 - v[2], xlabel="tempo",ylabel="erro velocidade")
    plot(p1,p2, tittle="Velocidades")
end

function plota()
    p1 = plot(t,[ad1,a[1]], xlabel="tempo",ylabel="aceleração")
    p2 = plot(t,[ad2,a[2]], xlabel="tempo",ylabel="aceleração")
    plot(p1,p2, tittle="Acelerações")
end

function plotea()
    p1 = plot(t,ad1 - a[1], xlabel="tempo",ylabel="erro aceleração")
    p2 = plot(t,ad2 - a[2], xlabel="tempo",ylabel="erro aceleração")
    plot(p1,p2, tittle="Acelerações")
end

function plotj()
    p1 = plot(t,[jd1,j[1]], xlabel="tempo",ylabel="tranco")
    p2 = plot(t,[jd2,j[2]], xlabel="tempo",ylabel="tranco")
    plot(p1,p2, tittle="Tranco")
end

function plotej()
    p1 = plot(t,jd1 - j[1] , xlabel="tempo",ylabel="erro tranco")
    p2 = plot(t,jd2 - j[2], xlabel="tempo",ylabel="erro tranco")
    plot(p1,p2, tittle="Tranco")
end

function plotallx()
    plot(plotx(),plotex())
end

function plotallv()
    plot(plotv(),plotev())
end

function plotalla()
    plot(plota(),plotea())
end

function plotallj()
    plot(plotj(),plotej())
end

#gui()
plotallx()


using BayesianIntegral
using Distributions: Normal, MvNormal, pdf
using LinearAlgebra: diagm, Symmetric
using Statistics
using HCubature: hcubature
using Sobol
using Random
standard_normal = Normal(0.0,1.0)

# Integral of gaussian distribution should be 1.
samples = 29
dims = 1
p(x) = 1.0
X = Array{Float64,2}(undef,samples,1)
X[:,1] =  collect(-2.0:(4.0/(samples-1)):2.0) #
y = p.(X)[:,1]
cov_func_parameters = gaussian_kernel_hyperparameters(1.0, repeat([2.0] , outer = dims))
prob_means = [0.0]
covar = Array{Float64,2}(undef,1,1)
covar[1,1] = 1.0
covar = Symmetric(covar)
noise = 0.2
integ = bayesian_integral_gaussian_exponential(X, y, prob_means, covar, cov_func_parameters, noise)
(integ.expectation - 1) < 0.01

# Integral in two dimensions
samples = 25
dims = 2
p(x) = 1.0
s = SobolSeq(dims)
X = convert( Array{Float64}, hcat([next!(s, repeat([0.5] , outer = dims)     ) for i = 1:samples]...)' )
y = repeat([1.0] , outer = samples)
cov_func_parameters = gaussian_kernel_hyperparameters(1.0, repeat([10.0] , outer = dims))
prob_means = repeat([0.0] , outer = dims)
covar = Symmetric(diagm(0 => ones(dims)))
noise = 0.2
integ = bayesian_integral_gaussian_exponential(X, y, prob_means, covar, cov_func_parameters, noise)
(integ.expectation - 1) < 0.02

# In ten dimensinos
samples = 50
dims = 10
p(x) = 1.0
s = SobolSeq(dims)
X = convert(Array{Float64}, hcat([next!(s, repeat([0.5] , outer = dims)     ) for i = 1:samples]...)')
y = repeat([1.0] , outer = samples)
cov_func_parameters = gaussian_kernel_hyperparameters(1.0, repeat([20.0] , outer = dims))
prob_means = repeat([0.0] , outer = dims)
covar = Symmetric(diagm(0 => ones(dims)))
noise = 0.2
integ = bayesian_integral_gaussian_exponential(X, y, prob_means, covar, cov_func_parameters, noise)
(integ.expectation - 1) < 0.02

##### More complex cases #####
# Now looking at a more complex one.
function f(x::Array{Float64,1})
    ndims = length(x)
    total = 12.0
    for i in 1:ndims
        total = total - x[i]^2 - 2*x[i]
    end
    return total
end
function fp(x::Array{Float64,1})
    ndim = length(x)
    if ndim == 1
        return f(x) * pdf.(Ref(standard_normal), x)
    else
        dist = MvNormal(zeros(ndim),diagm(0 => ones(ndim)))
        return f(x) * pdf(dist, x)
    end
end
function fp(x)
    global counter = counter + 1
    flipped = convert(Array{Float64,1},x)
    return fp(flipped)
end

dims = 1
bayesianAttempts = 10

function compare_for_f(dims::Int, bayesianAttempts::Int = 0, paths::Int = 0;  seed::Int = 1988)
    Random.seed!(1234)
    # Traditional
    global counter = 0
    lims = 3.0
    maxIter = 5000
    numerical_val, numerical_err = hcubature(x -> fp(x), repeat([-lims], outer = dims), repeat([lims], outer = dims); maxevals = maxIter)
    cont = counter

    # Kriging
    if bayesianAttempts < 1
        samples = cont
    else
        samples = bayesianAttempts
    end
    s = SobolSeq(dims)
    X = (hcat([next!(s, repeat([0.5] , outer = dims)     ) for i = 1:samples]...)' .- 0.5) .* (lims*2.0)
    y = Array{Float64,1}(undef, samples)
    for r in 1:samples
        y[r] = f(X[r,:])
    end
    noise = 0.05
    cov_func_parameters = gaussian_kernel_hyperparameters(1.0, repeat([1.0] , outer = dims))
    prob_means = repeat([0.0], outer = dims)
    covar = Symmetric(diagm(0 => ones(dims)))
    bayesian_val, bayesian_err = bayesian_integral_gaussian_exponential(X, y, prob_means, covar, cov_func_parameters, noise)

    # Kriging with new weights
    steps = 2000
    batch_size = convert(Int, floor( samples / 4 ) )
    step_multiple = 0.02
    seed = 1988
    K_mat = K_matrix(X, gaussian_kernel, cov_func_parameters, noise)
    like = log_likelihood(y, K_mat)
    n_cov_func_parameters = calibrate_by_ML_with_SGD(X, y, cov_func_parameters, steps, batch_size, step_multiple, noise, seed)
    K_mat = K_matrix(X, gaussian_kernel, n_cov_func_parameters, noise)
    nlike = log_likelihood(y, K_mat)
    nbayesian_val, nbayesian_err = bayesian_integral_gaussian_exponential(X, y , prob_means, covar, n_cov_func_parameters, noise)

    # Kriging with new weights and RProp
    params = RProp_params(1.01,0.99,0.2,5.0,0.5)
    MaxIter = 2000
    nn_cov_func_parameters = calibrate_by_ML_with_Rprop(X, y, cov_func_parameters, MaxIter, noise, params)
    K_mat = K_matrix(X, gaussian_kernel, n_cov_func_parameters, noise)
    rproplikelihood = log_likelihood(y, K_mat)
    rpropnbayesian_val, rpropnbayesian_err = bayesian_integral_gaussian_exponential(X, y , prob_means, covar, nn_cov_func_parameters, noise)

    # MC Integration
    if paths < 1
        paths = cont
    end
    dist = MvNormal(zeros(dims),diagm(0 => ones(dims)))
    Xs = transpose(convert(Matrix{Float64}, rand(dist, paths)))
    ys = Array{Float64,1}(undef, paths)
    for r in 1:paths
        ys[r] = f(Xs[r,:])
    end
    MC_integral = mean(ys)
    MC_err = std(ys) / sqrt(paths)

    print("For ", dims, " dimensions \n")
    print("    Bayesian    Integral is               ", bayesian_val       , " with error ", bayesian_err       ," and ", samples , " evaluations \n")
    print("    Calibrated Bayesian Integral is       ", nbayesian_val      , " with error ", nbayesian_err      ," and ", samples , " evaluations \n")
    print("    Rprop Calibrated Bayesian Integral is ", rpropnbayesian_val , " with error ", rpropnbayesian_err ," and ", samples , " evaluations \n")
    print("    Traditional Integral is               ", numerical_val[1]   , " with error ", numerical_err[1]   ," and ", cont    , " evaluations \n")
    print("    MC          Integral is               ", MC_integral        , " with error ", MC_err             ," and ", paths   , " evaluations \n")
    print("        Likelihood was ", like   , " and calibrated is ", nlike, " and Rprop calibrated is ", rproplikelihood  ,"\n")
    print("        weights became ", n_cov_func_parameters.w_0   , " and  ", n_cov_func_parameters.w_i, " with SGD \n")
    print("        weights became ", nn_cov_func_parameters.w_0    , " and  ", nn_cov_func_parameters.w_i, " with RProp \n")
    return (bayesian_val = bayesian_val, bayesian_err = bayesian_err, samples = samples, traditional_value = numerical_val[1], traditional_error = numerical_err, traditional_evals = cont, MC_integral = MC_integral, MC_err = MC_err, paths = paths)
end

f1_results = compare_for_f(1, 100)
abs(f1_results.bayesian_val - f1_results.MC_integral) < 0.01
f2_results = compare_for_f(2, 300)
abs(f2_results.bayesian_val - f2_results.MC_integral) < 0.11
#f10_results = compare_for_f(10, 200, 20000)
#f20_results = compare_for_f(20, 2000, 2000)