__precompile__(false)

module MyLib

using JSXGraph

function example()
	brd = board("board",xlim=[-15,15], ylim=[-15,15],axis=true)
	brd.style="width:600px; height:600px;margin:100px auto;"

	a = slider("a", [[5, 7], [8, 7], [-3, 0.1, 10]])
	b = slider("b", [[5, 6], [8, 6], [-1, 1, 5]])
	c = slider("c", [[5, 5], [8, 5], [-10, -5, 2]])
	[a,b,c] |> brd

	@jsf f(x) = val(a) * x^2 + val(b) * x + val(c)

	brd ++ plot(f)

	return JSXGraph.standalone(brd, full=true)
end

end


Base.@kwdef struct DensityAnalysis{D, K, B}
    datalimits::D=automatic
    npoints::Int=200
    kernel::K=automatic
    bandwidth::B=automatic
end

# Work around lack of length 1 tuple method
# TODO: also add weights support here
_kde(data::NTuple{1, Any}; kwargs...) = kde(data...; kwargs...)
_kde(data::Tuple; kwargs...) = kde(data; kwargs...)

defaultdatalimits(positional) = map(v -> mapreduce(extrema, extend_extrema, v), Tuple(positional))

applydatalimits(f::Function, d) = map(f, d)
applydatalimits(limits::Tuple, _) = limits

function _density(vs::Tuple; datalimits, npoints, kwargs...)
    k = _kde(vs; kwargs...)
    intervals = applydatalimits(datalimits, vs)
    rgs = map(intervals) do (min, max)
        return range(min, max; length=npoints)
    end
    res = pdf(k, rgs...)
    return (rgs..., res)
end

function (d::DensityAnalysis)(input::ProcessedLayer)
    datalimits = d.datalimits === automatic ? defaultdatalimits(input.positional) : d.datalimits
    options = valid_options(; datalimits, d.npoints, d.kernel, d.bandwidth)
    output = map(input) do p, n
        return _density(Tuple(p); pairs(n)..., pairs(options)...), (;)
    end
    N = length(input.positional)
    labels = set(input.labels, N+1 => "pdf")
    plottypes = [LinesFill, Heatmap, Volume]
    default_plottype = plottypes[N]
    plottype = Makie.plottype(input.plottype, default_plottype)
    return ProcessedLayer(output; plottype, labels)
end

"""
    density(; datalimits=automatic, kernel=automatic, bandwidth=automatic, npoints=200)

Fit a kernel density estimation of `data`.
Here, `datalimits` specifies the range for which the density should be calculated,
and `kernel` and `bandwidth` are forwarded to `KernelDensity.kde`.
`npoints` is the number of points used by Makie to draw the line
"""
density(; options...) = transformation(DensityAnalysis(; options...))


# Wishart distribution
#
#   following the Wikipedia parameterization
#
"""
    Wishart(ν, S)
```julia
ν::Real           degrees of freedom (greater than p - 1)
S::AbstractPDMat  p x p scale matrix
```
The [Wishart distribution](http://en.wikipedia.org/wiki/Wishart_distribution)
generalizes the gamma distribution to ``p\\times p`` real, positive definite
matrices ``\\mathbf{H}``. If ``\\mathbf{H}\\sim \\textrm{W}_p(\\nu,\\mathbf{S})``,
then its probability density function is

```math
f(\\mathbf{H};\\nu,\\mathbf{S}) = \\frac{1}{2^{\\nu p/2} \\left|\\mathbf{S}\\right|^{\\nu/2} \\Gamma_p\\left(\\frac {\\nu}{2}\\right ) }{\\left|\\mathbf{H}\\right|}^{(\\nu-p-1)/2} e^{-(1/2)\\operatorname{tr}(\\mathbf{S}^{-1}\\mathbf{H})}.
```

If ``\\nu`` is an integer, then a random matrix ``\\mathbf{H}`` given by

```math
\\mathbf{H} = \\mathbf{X}\\mathbf{X}^{\\rm{T}}, \\quad\\mathbf{X} \\sim \\textrm{MN}_{p,\\nu}(\\mathbf{0}, \\mathbf{S}, \\mathbf{I}_{\\nu})
```

has ``\\mathbf{H}\\sim \\textrm{W}_p(\\nu, \\mathbf{S})``. For non-integer degrees of freedom,
Wishart matrices can be generated via the [Bartlett decomposition](https://en.wikipedia.org/wiki/Wishart_distribution#Bartlett_decomposition).
"""
struct Wishart{T<:Real, ST<:AbstractPDMat} <: ContinuousMatrixDistribution
    df::T     # degree of freedom
    S::ST     # the scale matrix
    logc0::T  # the logarithm of normalizing constant in pdf
end

#  -----------------------------------------------------------------------------
#  Constructors
#  -----------------------------------------------------------------------------

function Wishart(df::T, S::AbstractPDMat{T}) where T<:Real
    p = dim(S)
    df > p - 1 || throw(ArgumentError("df should be greater than dim - 1."))
    logc0 = wishart_logc0(df, S)
    R = Base.promote_eltype(T, logc0)
    prom_S = convert(AbstractArray{T}, S)
    Wishart{R, typeof(prom_S)}(R(df), prom_S, R(logc0))
end

function Wishart(df::Real, S::AbstractPDMat)
    T = Base.promote_eltype(df, S)
    Wishart(T(df), convert(AbstractArray{T}, S))
end

Wishart(df::Real, S::Matrix) = Wishart(df, PDMat(S))

Wishart(df::Real, S::Cholesky) = Wishart(df, PDMat(S))

#  -----------------------------------------------------------------------------
#  REPL display
#  -----------------------------------------------------------------------------

show(io::IO, d::Wishart) = show_multline(io, d, [(:df, d.df), (:S, Matrix(d.S))])

#  -----------------------------------------------------------------------------
#  Conversion
#  -----------------------------------------------------------------------------

function convert(::Type{Wishart{T}}, d::Wishart) where T<:Real
    P = convert(AbstractArray{T}, d.S)
    Wishart{T, typeof(P)}(T(d.df), P, T(d.logc0))
end
function convert(::Type{Wishart{T}}, df, S::AbstractPDMat, logc0) where T<:Real
    P = convert(AbstractArray{T}, S)
    Wishart{T, typeof(P)}(T(df), P, T(logc0))
end

#  -----------------------------------------------------------------------------
#  Properties
#  -----------------------------------------------------------------------------

insupport(::Type{Wishart}, X::Matrix) = isposdef(X)
insupport(d::Wishart, X::Matrix) = size(X) == size(d) && isposdef(X)

dim(d::Wishart) = dim(d.S)
size(d::Wishart) = (p = dim(d); (p, p))
rank(d::Wishart) = dim(d)
params(d::Wishart) = (d.df, d.S)
@inline partype(d::Wishart{T}) where {T<:Real} = T

mean(d::Wishart) = d.df * Matrix(d.S)

function mode(d::Wishart)
    r = d.df - dim(d) - 1.0
    r > 0.0 || throw(ArgumentError("mode is only defined when df > p + 1"))
    return Matrix(d.S) * r
end

function meanlogdet(d::Wishart)
    p = dim(d)
    df = d.df
    v = logdet(d.S) + p * logtwo
    for i = 1:p
        v += digamma(0.5 * (df - (i - 1)))
    end
    return v
end

function entropy(d::Wishart)
    p = dim(d)
    df = d.df
    -d.logc0 - 0.5 * (df - p - 1) * meanlogdet(d) + 0.5 * df * p
end

#  Gupta/Nagar (1999) Theorem 3.3.15.i
function cov(d::Wishart, i::Integer, j::Integer, k::Integer, l::Integer)
    S = Matrix(d.S)
    d.df * (S[i, k] * S[j, l] + S[i, l] * S[j, k])
end

function var(d::Wishart, i::Integer, j::Integer)
    S = Matrix(d.S)
    d.df * (S[i, i] * S[j, j] + S[i, j] ^ 2)
end

#  -----------------------------------------------------------------------------
#  Evaluation
#  -----------------------------------------------------------------------------

function wishart_logc0(df::Real, S::AbstractPDMat)
    h_df = df / 2
    p = dim(S)
    -h_df * (logdet(S) + p * typeof(df)(logtwo)) - logmvgamma(p, h_df)
end

function logkernel(d::Wishart, X::AbstractMatrix)
    df = d.df
    p = dim(d)
    Xcf = cholesky(X)
    0.5 * ((df - (p + 1)) * logdet(Xcf) - tr(d.S \ X))
end

#  -----------------------------------------------------------------------------
#  Sampling
#  -----------------------------------------------------------------------------

function _rand!(rng::AbstractRNG, d::Wishart, A::AbstractMatrix)
    _wishart_genA!(rng, dim(d), d.df, A)
    unwhiten!(d.S, A)
    A .= A * A'
end

function _wishart_genA!(rng::AbstractRNG, p::Int, df::Real, A::AbstractMatrix)
    # Generate the matrix A in the Bartlett decomposition
    #
    #   A is a lower triangular matrix, with
    #
    #       A(i, j) ~ sqrt of Chisq(df - i + 1) when i == j
    #               ~ Normal()                  when i > j
    #
    A .= zero(eltype(A))
    for i = 1:p
        @inbounds A[i,i] = rand(rng, Chi(df - i + 1.0))
    end
    for j in 1:p-1, i in j+1:p
        @inbounds A[i,j] = randn(rng)
    end
end

#  -----------------------------------------------------------------------------
#  Test utils
#  -----------------------------------------------------------------------------

function _univariate(d::Wishart)
    check_univariate(d)
    df, S = params(d)
    α = df / 2
    β = 2Matrix(S)[1]
    return Gamma(α, β)
end

function _rand_params(::Type{Wishart}, elty, n::Int, p::Int)
    n == p || throw(ArgumentError("dims must be equal for Wishart"))
    ν = elty( n + 1 + abs(10randn()) )
    S = (X = 2rand(elty, n, n) .- 1; X * X')
    return ν, S
end


using ITensors

"""
    trg(T::ITensor; χmax::Int, nsteps::Int) -> κ, T

Perform the TRG algorithm on the partition function composed of the ITensor T.

The indices of T must obey come in pairs `(sₕ => sₕ')` and  `(sᵥ => sᵥ').

χmax is the maximum renormalized bond dimension.

nsteps are the number of renormalization steps performed.

The outputs are κ, the partition function per site, and the final renormalized
ITensor T.
"""
function trg(T::ITensor; χmax::Int, nsteps::Int, cutoff=0.0, svd_alg="divide_and_conquer")
  sₕ, sᵥ = filterinds(T; plev=0)
  @assert hassameinds((sₕ, sₕ', sᵥ, sᵥ'), T)

  # Keep track of the partition function per site
  κ = 1.0
  for n in 1:nsteps
    Fₕ, Fₕ′ = factorize(
      T,
      (sₕ', sᵥ');
      ortho="none",
      maxdim=χmax,
      cutoff=cutoff,
      tags=tags(sₕ),
      svd_alg=svd_alg,
    )

    s̃ₕ = commonind(Fₕ, Fₕ′)
    Fₕ′ *= δ(dag(s̃ₕ), s̃ₕ')

    Fᵥ, Fᵥ′ = factorize(
      T, (sₕ, sᵥ'); ortho="none", maxdim=χmax, cutoff=cutoff, tags=tags(sᵥ), svd_alg=svd_alg
    )

    s̃ᵥ = commonind(Fᵥ, Fᵥ′)
    Fᵥ′ *= δ(dag(s̃ᵥ), s̃ᵥ')

    T =
      (Fₕ * δ(dag(sₕ'), sₕ)) *
      (Fᵥ * δ(dag(sᵥ'), sᵥ)) *
      (Fₕ′ * δ(dag(sₕ), sₕ')) *
      (Fᵥ′ * δ(dag(sᵥ), sᵥ'))

    sₕ, sᵥ = s̃ₕ, s̃ᵥ

    trT = abs((T * δ(sₕ, sₕ') * δ(sᵥ, sᵥ'))[])
    T = T / trT
    κ *= trT^(1 / 2^n)
  end
  return κ, T
end


using CPUTime
using Plots
using DifferentialEquations
using LinearAlgebra
using MatrixEquations

# function ARE(A, B, Q, R)
#     """有本ポッターで解く（未完成）"""
#     ℋ = [
#         A' -B*inv(R)*B'
#         -Q -A
#     ]  #  ハミルトン行列
#     println("ℋ = ", ℋ)

#     λ_ = eigvals(ℋ)  # ハミルトン行列の固有値
#     ω_ = eigvecs(ℋ)  # ハミルトン行列の固有ベクトル
#     println("ℋ の固有値 = ", λ_)
#     println("ℋ の固有ベクトル", ω_)

#     # ハミルトン行列の固有値が負のものを探す
#     index = []
#     for i in 1:size(λ_)[1]
#         print(i)
#         if real(λ_[i]) < 0
#             println("実部が負")
#             push!(index, i)
#         else
#             println("実部が正")
#         end
#     end
#     println("ℋ の固有値の実部が負なのは, ", index)
#     n = size(index)[1]

#     # Y,Zを計算
#     ω = Matrix{ComplexF64}(undef, size(ω_)[1], n)  # 未初期化のMatrix

#     for i in 1:n
#         global ω[:, i] .= ω_[:, index[i]]
#     end

#     Y = ω[1:n, :]
#     Z = ω[n+1:end, :]

#     P = Z * inv(Y)
#     println("P = ", P)
#     return P
# end

function AREbyMatrixEquations(A, B, Q, R)
    """ライブラリ使用
    https://juliapackages.com/p/matrixequations
    """
    S = zero(B)
    P, E, F = arec(A, B, R, Q, S)
    println("P = ", P)
    return P
end

function simu(A, B, C, D, Q, R, x0)
    """解を描写"""
    #P = ARE(A, B, Q, R)  # リカッチ方程式を解く
    P = AREbyMatrixEquations(A, B, Q, R)
    K = inv(R) * B' * P  # 最適フィードバックゲイン
    Ā = A - B * K
    println("Ā = ", Ā)
    println(eigvals(Ā))

    # 解
    x(t) = exp(Ā .* t) * x0  # 状態方程式の解
    u(t) = -K * x(t)  # 最適入力
    y(t) = C * x(t) + D * u(t)  # 出力方程式の解

    # Plot．システムの次元に合わせてplotの数を変更．
    plot(xlims=(0, 10), xaxis="time", )
    plot!(t -> x.(t)[1], label="x1")
    plot!(t -> x.(t)[2], label="x2")
    plot!(t -> x.(t)[3], label="x3")
end

### システムの例1 ###
A = [
    1.1 2.0 3.0
    0 0.95 1.20
    1.2 0.01 10.5
]
B = [
    1.0
    0.0
    0.847
]
C = Matrix(I, 2, 2)
D = [1 1 1]
# 重み行列
Q = diagm(0 => [1,1,1]) .* 1000
R = ones((1,1))
x0 = [1, 2, 3]

# ### システムの例2 ###
# A = [0.0 1.0; -10.0 -1.0]
# B = [0.0, 1.0]
# Q = [300 0; 0 60]
# R = ones((1, 1))
# C = ones((2, 2))
# D = B
# x0 = [1, 2]

@time simu(A, B, C, D, Q, R, x0)  # 実行


module LogDensityTestSuite

export
    # generic
    samples,
    # primitives
    StandardMultivariateNormal,
    # transformations
    linear, shift, elongate,
    # mixtures
    mix, directional_weight,
    # diagnostics
    quantile_boundaries, bin_counts, two_sided_pvalues, print_ascii_plot

using ArgCheck: @argcheck
using DocStringExtensions: FIELDS, FUNCTIONNAME, SIGNATURES, TYPEDEF
using LinearAlgebra: AbstractTriangular, checksquare, diag, Diagonal, I, dot, logabsdet, lu,
    norm
using LogDensityProblems: LogDensityOrder
import LogDensityProblems: capabilities, dimension, logdensity, logdensity_and_gradient
using MCMCDiagnostics: ess_factor_estimate
using Parameters: @unpack
using Printf: @sprintf
using Sobol: next!, SobolSeq
using Statistics: quantile
using StatsFuns: norminvcdf, normcdf, logaddexp, logistic

include("generic.jl")
include("primitives.jl")
include("transformations.jl")
include("mixtures.jl")
include("diagnostics.jl")

end # module


module FloatII
    export interp

    function interp(f, l, u, n)
        # no numerical error caused by adding a lot of small numbers.
        x(idx) = (l*(n-idx)+u*idx)/n
        map( idx -> floor(f(x(idx)),2), 0:(n-1) )
    end

end


using JuMP
using Gurobi
using StatsBase
using DataFrames
using CSV
include("SupportFunction.jl")
include("DataLoad.jl")
println("Leeeeroooy Jenkins")
#Esben's path
path = "$(homedir())/Documents/GitHub/Thesis/Data/MonthlyReturns"

mainData = loadIndexDataNoDurLOGReturn(path)
mainDataArr = Array(mainData)

startRow = 1
endRow = 100
trainingObservations = 100
mainXarr = mainDataArr[startRow:endRow,1:10]
indexes = size(mainXarr)[2]

#### WITH RISK AVERSION

nRows = size(mainDataArr)[1]
nCols = size(mainDataArr)[2]

##LOOP TO TEST IT
portfolioClassTotal = []
portfolioClassTotal = push!(portfolioClassTotal, 1)
portfolioClassIndividual = []
portfolioClassIndividual = push!(portfolioClassIndividual, 0)
wClasstracker = zeros((nRows-trainingObservations-1),10)

wClass = repeat([0.1], outer = 10)
portfolioClassTotal = []
portfolioClassTotal = push!(portfolioClassTotal, 1)
portfolioClassIndividual = []
portfolioClassIndividual = push!(portfolioClassIndividual, 0)

for rStart = 1:(nRows-trainingObservations-1)
    startRow = rStart
    endRow = rStart+trainingObservations
    mainXarr = mainDataArr[startRow:endRow,1:10]
    indexes = size(mainXarr)[2]

    #creating the covariance matrix of the training data
    Sigma = cov(mainXarr)

    #creating the mean of the returns for SAA (forecasts should be used instead)
    LassoOutputHere = [1, 0, 0, 1, 1, 0, 0, 1, 1, 0]
    indicators = LassoOutputHere
    wStar = createWeighting(indicators)

    wMVtracker[rStart,:] = wStar

    #creating the forecastRow in arithmetic returns (not log returns)
    #then we can multiply and sum across the row
    forecastRow = exp(mainDataArr[(endRow+1):(endRow+1),1:10])

    periodReturn = Array(forecastRow*wStar)[1]
    currentLength = length(portfolioClassTotal)
    portfolioClassTotal = push!(portfolioClassTotal, portfolioClassTotal[currentLength]*periodReturn)
    portfolioClassIndividual = push!(portfolioClassIndividual, periodReturn)

    period1NReturn = Array(forecastRow*w1N)[1]
    current1NLength = length(portfolio1NTotal)
    portfolio1NTotal = push!(portfolio1NTotal, portfolio1NTotal[current1NLength]*period1NReturn)
    portfolio1NIndividual = push!(portfolio1NIndividual, period1NReturn)
end

#Checking the total return of the portfolio by checking last index
portfolioClassTotal[length(portfolioClassTotal)]
portfolio1NTotal[length(portfolio1NTotal)]

portfolioClassTotal = convert(Array{Float64,1}, portfolioClassTotal)
portfolioClassIndividual = convert(Array{Float64,1}, portfolioClassIndividual)
(mean(portfolioClassIndividual)-1)*100
std(portfolioClassIndividual)
SharpeRatioMV = (mean(portfolioClassIndividual)-1)/std(portfolioClassIndividual)

portfolio1NTotal = convert(Array{Float64,1}, portfolio1NTotal)
portfolio1NIndividual = convert(Array{Float64,1}, portfolio1NIndividual)
(mean(portfolio1NIndividual)-1)*100
std(portfolio1NIndividual)
SharpeRatio1N = (mean(portfolio1NIndividual)-1)/std(portfolio1NIndividual)

combinedPortfolios = hcat(portfolioClassTotal, portfolio1NTotal, portfolioClassIndividual, portfolio1NIndividual)
writedlm("Class100obsversus1N.csv", combinedPortfolios,",")

writedlm("Class100obsPortfolioWeights.csv", wMVtracker,",")


# Copyright (c) 2021 Idiap Research Institute, http://www.idiap.ch/
#  Niccolò Antonello <nantonel@idiap.ch>

fst = WFST(["a","b","c","d","e","f"])
add_arc!(fst,1=>1,"a"=>"a",2)
add_arc!(fst,1=>2,"a"=>"a",1)
add_arc!(fst,2=>3,"b"=>"b",3)
add_arc!(fst,2=>3,"c"=>"c",4)
add_arc!(fst,3=>3,"d"=>"d",5)
add_arc!(fst,3=>4,"d"=>"d",6)
add_arc!(fst,4=>4,"f"=>"f",2)
initial!(fst,1)
final!(fst,2,3)
final!(fst,4,2)
#println(fst)
sym = get_isym(fst)
W = typeofweight(fst)
D = Int

rfst = reverse(fst)
@test rfst[1][1] == Arc{W,D}(0,0,W(3),3)
@test rfst[1][2] == Arc{W,D}(0,0,W(2),5)
@test rfst[2][1] == Arc{W,D}(sym["a"],sym["a"],W(2),2)
@test rfst[3][1] == Arc{W,D}(sym["a"],sym["a"],W(1),2)
@test rfst[4][1] == Arc{W,D}(sym["b"],sym["b"],W(3),3)
@test rfst[4][2] == Arc{W,D}(sym["c"],sym["c"],W(4),3)
@test rfst[4][3] == Arc{W,D}(sym["d"],sym["d"],W(5),4)
@test rfst[5][1] == Arc{W,D}(sym["d"],sym["d"],W(6),4)
@test rfst[5][2] == Arc{W,D}(sym["f"],sym["f"],W(2),5)
@test size(rfst) == (5,9)

fst = WFST([1,2],["a","b"];W = LeftStringWeight)
add_arc!(fst,1=>2,1=>"a","ciao")
add_arc!(fst,2=>3,2=>"b","bello")
initial!(fst,1)
final!(fst,3)
rfst = reverse(fst)
#println(rfst)

seq = [1,2]
out, w = fst(seq)
outr, wr = rfst(reverse(seq))
@test out == reverse(outr)
@test w == reverse(wr)


import Theta: Lattice, svp, check_size_reduce, size_reduce!, check_lll, lll!

@testset "SVP correctness" begin
# Test data are generated using fplll
    @test norm(svp([116 303;331 963])[1]) == norm([-45, -30])
    @test norm(svp([116 456 99 566; 331 225 649 213; 303 827 395 694; 963 381 975 254])[1]) == norm([110 -12 -133 -127])
    @test norm(svp([116 827 566 597 912 611; 331 381 213 980 451 734; 303 99 694 462 926 668; 963 649 254 533 105 308; 456 395 629 231 162 101; 225 975 303 443 742 863])[1]) == norm([-301 283 -258 203 -61 121])
    @test norm(svp([116 99 629 912 668 438 621 371; 331 649 303 451 308 456 215 554; 303 395 597 926 101 887 6 434; 963 975 980 105 863 466 340 638; 456 566 462 162 279 634 142 939; 225 213 533 742 976 748 411 554; 827 694 231 611 127 307 250 840; 381 254 443 734 549 1002 873 633])[1]) == norm([-17 318 92 12 110 -12 -133 -127])
end

@testset "SVP length" begin
    @test begin
        S = svp(rand(1,1));
        (norm(S[1]))^2 ≈ S[2]
    end
    @test begin
        S = svp(rand(2,2));
        norm(S[1])^2 ≈ S[2]
    end
    @test begin
        S = svp(rand(5,5));
        norm(S[1])^2 ≈ S[2]
    end
    @test begin
        S = svp(rand(4,3));
        norm(S[1])^2 ≈ S[2]
    end
end

@testset "Size reduction" begin
    @test check_size_reduce(size_reduce!(Lattice(rand(1,1))))
    @test check_size_reduce(size_reduce!(Lattice(rand(2,2))))
    @test check_size_reduce(size_reduce!(Lattice(rand(3,3))))
    @test check_size_reduce(size_reduce!(Lattice(rand(4,4))))
    @test check_size_reduce(size_reduce!(Lattice(rand(5,5))))
    @test check_size_reduce(size_reduce!(Lattice(rand(10,10))))
end

@testset "LLL" begin
    @test check_lll(lll!(Lattice(rand(1,1))))
    @test check_lll(lll!(Lattice(rand(2,2))))
    @test check_lll(lll!(Lattice(rand(3,3))))
    @test check_lll(lll!(Lattice(rand(4,4))))
    @test check_lll(lll!(Lattice(rand(5,5))))
    @test check_lll(lll!(Lattice(rand(10,10))))
end


using Stats
# using Logging

function estimateNGFEpsilon(referenceImage::regImage;cutoffPercent = 80)
  m = size(referenceImage.data)
  h = referenceImage.voxelsize
  Rc = referenceImage.data

  shortFiniteDiffX = spdiagm((-ones(m[1]-1,1),ones(m[1]-2,1)),[0, 1],m[1]-1,m[1])/(2*h[1])
  #shortFiniteDiffX[1] = shortFiniteDiffX[2]
  #shortFiniteDiffX[end] = shortFiniteDiffX[end-1]

  shortFiniteDiffY = spdiagm((-ones(m[2]-1,1),ones(m[2]-2,1)),[0, 1],m[2]-1,m[2])/(2*h[2])
  #shortFiniteDiffY[1] = shortFiniteDiffY[2]
  #shortFiniteDiffY[end] = shortFiniteDiffY[end-1]

  averageX = spdiagm((0.5*ones(m[1]-2,1),0.5*ones(m[1]-1,1)),[-1, 0],m[1],m[1]-1)
  averageY = spdiagm((ones(m[2]-2,1)*0.5,ones(m[2]-1,1)*0.5),[-1, 0],m[2],m[2]-1)

  G1 = sparse(kron(speye(m[2]),shortFiniteDiffX));
  G2 = sparse(kron(shortFiniteDiffY,speye(m[1])));

  AvgX = sparse(kron(speye(m[2]),averageX));
  AvgY = sparse(kron(averageY,speye(m[1])));

  gradTx  = G1*referenceImage.data[:]
  gradTy  = G2*referenceImage.data[:]

  quantileX = percentile(AvgX * (gradTx.^2),cutoffPercent)
  quantileY = percentile(AvgY * (gradTy.^2),cutoffPercent)

  estimatedEps = sqrt(0.5*(quantileX + quantileY))

  lengthGT = sqrt(sum(AvgX * (gradTx.^2) + AvgY * (gradTy.^2), 2) + estimatedEps^2) #epsilon-norm

  vectorlength = reshape(lengthGT,(size(referenceImage)))

  try
    figure(51)
    clf()
    showImage(vectorlength)
    s = @sprintf("estimated epsilon: %3.3e", estimatedEps)
    title(s)
    colorbar()
    figure(52)
    clf()
    plt.hist(lengthGT,bins=100)
  catch err
    warn("Error while plotting. This is normal if it happens in the automated test framework.")
    println(err)
  end

  return (estimatedEps,vectorlength)

end


function nterms(t)
    if istree(t)
        return reduce(+, map(nterms, arguments(t)), init=0)
    else
        return 1
    end
end

# Soft pivoted
# Note: we call this function with a matrix of Union{SymbolicUtils.Symbolic, Any}
function sym_lu(A; check=true)
    SINGULAR = typemax(Int)
    m, n = size(A)
    F = map(x->x isa Num ? x : Num(x), A)
    minmn = min(m, n)
    p = Vector{LinearAlgebra.BlasInt}(undef, minmn)
    info = 0
    for k = 1:minmn
        kp = k
        amin = SINGULAR
        for i in k:m
            absi = _iszero(F[i, k]) ? SINGULAR : nterms(F[i,k])
            if absi < amin
                kp = i
                amin = absi
            end
        end

        p[k] = kp

        if amin == SINGULAR && !(amin isa Symbolic) && (amin isa Number) && iszero(info)
            info = k
        end

        # swap
        for j in 1:n
            F[k, j], F[kp, j] = F[kp, j], F[k, j]
        end

        for i in k+1:m
            F[i, k] = F[i, k] / F[k, k]
        end
        for j = k+1:n
            for i in k+1:m
                F[i, j] = F[i, j] - F[i, k] * F[k, j]
            end
        end
    end
    check && LinearAlgebra.checknonsingular(info, Val{true}())
    LU(F, p, convert(LinearAlgebra.BlasInt, info))
end

# Given a vector of equations and a
# list of independent variables,
# return the coefficient matrix `A` and a
# vector of constants (possibly symbolic) `b` such that
# A \ b will solve the equations for the vars
function A_b(eqs::AbstractArray, vars::AbstractArray, check)
    exprs = rhss(eqs) .- lhss(eqs)
    if check
        for ex in exprs
            @assert isaffine(ex, vars)
        end
    end
    A = jacobian(exprs, vars)
    b = A * vars - exprs
    A, b
end

"""
$(TYPEDSIGNATURES)

Solve equation(s) `eqs` for a set of variables `vars`.

Assumes `length(eqs) == length(vars)`

Currently only works if all equations are linear. `check` if the expr is linear
w.r.t `vars`.
"""
function solve_for(eq, var; simplify=false, check=true) # scalar case
    # simplify defaults for `false` as canonicalization should handle most of
    # the cases.
    a, b, islinear = linear_expansion(eq, var)
    check && @assert islinear
    islinear || return nothing
    # a * x + b = 0
    x = a \ -b
    simplify ? SymbolicUtils.simplify(x, expand=true) : x
end

function solve_for(eqs::AbstractArray, vars::AbstractArray; simplify=true, check=true)
    length(eqs) == 1 == length(vars) && return [solve_for(eqs[1], vars[1]; simplify=simplify, check=check)]
    A, b = A_b(eqs, vars, check)
    #TODO: we need to make sure that `solve_for(eqs, vars)` contains no `vars`
    sol = _solve(A, b, simplify)
    map(Num, sol)
end

function _solve(A::AbstractMatrix, b::AbstractArray, do_simplify)
    A = SymbolicUtils.simplify.(Num.(A), expand=true)
    b = SymbolicUtils.simplify.(Num.(b), expand=true)
    sol = value.(sym_lu(A) \ b)
    do_simplify ? SymbolicUtils.simplify.(sol, expand=true) : sol
end

LinearAlgebra.ldiv!(A::UpperTriangular{<:Union{Symbolic,Num}}, b::AbstractVector{<:Union{Symbolic,Num}}, x::AbstractVector{<:Union{Symbolic,Num}} = b) = symsub!(A, b, x)
function symsub!(A::UpperTriangular, b::AbstractVector, x::AbstractVector = b)
    LinearAlgebra.require_one_based_indexing(A, b, x)
    n = size(A, 2)
    if !(n == length(b) == length(x))
        throw(DimensionMismatch("second dimension of left hand side A, $n, length of output x, $(length(x)), and length of right hand side b, $(length(b)), must be equal"))
    end
    @inbounds for j in n:-1:1
        _iszero(A.data[j,j]) && throw(SingularException(j))
        xj = x[j] = b[j] / A.data[j,j]
        for i in j-1:-1:1
            sub = _isone(xj) ? A.data[i,j] : A.data[i,j] * xj
            if !_iszero(sub)
                b[i] -= sub
            end
        end
    end
    x
end

LinearAlgebra.ldiv!(A::UnitLowerTriangular{<:Union{Symbolic,Num}}, b::AbstractVector{<:Union{Symbolic,Num}}, x::AbstractVector{<:Union{Symbolic,Num}} = b) = symsub!(A, b, x)
function symsub!(A::UnitLowerTriangular, b::AbstractVector, x::AbstractVector = b)
    LinearAlgebra.require_one_based_indexing(A, b, x)
    n = size(A, 2)
    if !(n == length(b) == length(x))
        throw(DimensionMismatch("second dimension of left hand side A, $n, length of output x, $(length(x)), and length of right hand side b, $(length(b)), must be equal"))
    end
    @inbounds for j in 1:n
        xj = x[j] = b[j]
        for i in j+1:n
            sub = _isone(xj) ? A.data[i,j] : A.data[i,j] * xj
            if !_iszero(sub)
                b[i] -= sub
            end
        end
    end
    x
end

minor(B, j) = B[2:end, 1:size(B,2) .!= j]
function LinearAlgebra.det(A::AbstractMatrix{<:Num}; laplace=true)
    if laplace
        n = LinearAlgebra.checksquare(A)
        if n == 1
            return A[1, 1]
        elseif n == 2
            return A[1, 1] * A[2, 2] - A[1, 2] * A[2, 1]
        else
            return sum((-1)^(1+j) * A[1,j] * det(minor(A,j), laplace=true)
                       for j in axes(A,2))
        end
    else
        if istriu(A) || istril(A)
            return det(UpperTriangular(A))
        end
        return det(lu(A; check = false))
    end
end

function LinearAlgebra.norm(x::AbstractArray{Num}, p::Real=2)
    p = value(p)
    issym = p isa Symbolic
    if !issym && p == 2
        sqrt(sum(x->abs2(x), x))
    elseif !issym && isone(p)
        sum(abs, x)
    elseif !issym && isinf(p)
        mapreduce(abs, max, x)
    else
        sum(x->abs(x)^p, x)^inv(p)
    end
end

"""
    (a, b, islinear) = linear_expansion(t, x)

When `islinear`, return `a` and `b` such that `a * x + b == t`.
"""
function linear_expansion(t, x)
    a, b, islinear = _linear_expansion(t, x)
    x isa Num ? (Num(a), Num(b), islinear) : (a, b, islinear)
end
# _linear_expansion always returns `Symbolic`
function _linear_expansion(t::Equation, x)
    a₂, b₂, islinear = linear_expansion(t.rhs, x)
    islinear || return (a₂, b₂, false)
    a₁, b₁, islinear = linear_expansion(t.lhs, x)
    # t.rhs - t.lhs = 0
    return (a₂ - a₁, b₂ - b₁, islinear)
end
trival_linear_expansion(t, x) = isequal(t, x) ? (1, 0, true) : (0, t, true) # trival case
function _linear_expansion(t, x)
    t = value(t)
    t isa Symbolic || return (0, t, true)
    x = value(x)
    istree(t) || return trival_linear_expansion(t, x)
    isequal(t, x) && return (1, 0, true)

    op, args = operation(t), arguments(t)
    op isa Differential && trival_linear_expansion(t, x)

    if op === (+)
        a₁ = b₁ = 0
        islinear = true
        # (a₁ x + b₁) + (a₂ x + b₂) = (a₁ + a₂) x + (b₁ + b₂)
        for (i, arg) in enumerate(args)
            a₂, b₂, islinear = linear_expansion(arg, x)
            islinear || return (a₁, b₁, false)
            a₁ += a₂
            b₁ += b₂
        end
        return (a₁, b₁, true)
    elseif op === (-)
        @assert length(args) == 1
        a, b, islinear = linear_expansion(args[1], x)
        return (-a, -b, islinear)
    elseif op === (*)
        # (a₁ x + b₁) (a₂ x + b₂) = a₁ a₂ x² + (a₁ b₂ + a₂ b₁) x + b₁ b₂
        a₁ = 0
        b₁ = 1
        islinear = true
        for (i, arg) in enumerate(args)
            a₂, b₂, islinear = linear_expansion(arg, x)
            (islinear && _iszero(a₁ * a₂)) || return (a₁, b₁, false)
            a₁ = a₁ * b₂ + a₂ * b₁
            b₁ *= b₂
        end
        return (a₁, b₁, true)
    elseif op === (^)
        # (a₁ x + b₁)^(a₂ x + b₂) is linear => a₂ = 0 && (b₂ == 1 || a₁ == 0)
        a₂, b₂, islinear = linear_expansion(args[2], x)
        (islinear && _iszero(a₂)) || return (0, 0, false)
        a₁, b₁, islinear = linear_expansion(args[1], x)
        _isone(b₂) && (a₁, b₁, islinear)
        (islinear && _iszero(a₁)) || return (0, 0, false)
        return (0, b₁^b₂, islinear)
    else
        for (i, arg) in enumerate(args)
            a, b, islinear = linear_expansion(arg, x)
            (_iszero(a) && islinear) || return (0, 0, false)
        end
        return (0, t, true)
    end
end


@testset "Compute Autodiff on CPU" begin
    @testset "Case $case" for case in ["case9.m", "case30.m"]
        datafile = joinpath(dirname(@__FILE__), "..", "..", "data", case)
        tolerance = 1e-8
        pf = PS.PowerNetwork(datafile)
        nbus = pf.nbus
        pv = pf.pv ; npv = length(pv)
        pq = pf.pq ; npq = length(pq)
        ref = pf.ref ; nref = length(ref)

        polar = PolarForm(pf, CPU())
        cache = ExaPF.get(polar, ExaPF.PhysicalState())
        ExaPF.init_buffer!(polar, cache)

        xk = ExaPF.initial(polar, State())
        u = ExaPF.initial(polar, Control())

        jx = AutoDiff.Jacobian(polar, ExaPF.power_balance, State())
        ju = AutoDiff.Jacobian(polar, ExaPF.power_balance, Control())

        # Solve power flow
        conv = powerflow(polar, jx, cache, NewtonRaphson(tol=1e-12))
        # Get solution in complex form.
        V = cache.vmag .* exp.(im .* cache.vang)

        # Test Jacobian w.r.t. State
        @testset "Constraint $(cons)" for cons in [
            ExaPF.voltage_magnitude_constraints,
            ExaPF.power_balance,
            ExaPF.active_power_constraints,
            ExaPF.reactive_power_constraints,
            ExaPF.flow_constraints,
        ]
            m = ExaPF.size_constraint(polar, cons)
            # Allocation
            pbm = AutoDiff.TapeMemory(polar, cons, typeof(u))
            tgt = rand(m)
            c = zeros(m)

            ## STATE JACOBIAN

            xjacobianAD = ExaPF.AutoDiff.Jacobian(polar, cons, State())
            # Evaluate Jacobian with AD
            J = AutoDiff.jacobian!(polar, xjacobianAD, cache)
            # Matpower Jacobian
            Jmat_x = ExaPF.matpower_jacobian(polar, State(), cons, V)
            # Evaluate Jacobian transpose vector product

            # Compare with FiniteDiff
            function jac_fd_x(x)
                cache.vang[pv] .= x[1:npv]
                cache.vang[pq] .= x[npv+1:npv+npq]
                cache.vmag[pq] .= x[npv+npq+1:end]
                c = zeros(m)
                cons(polar, c, cache.vmag, cache.vang, cache.pinj, cache.qinj)
                return c
            end
            x = [cache.vang[pv]; cache.vang[pq]; cache.vmag[pq]]
            Jd = FiniteDiff.finite_difference_jacobian(jac_fd_x, x)
            @test size(J) == (m, length(x))
            # Test that AutoDiff matches MATPOWER
            @test isapprox(Jd, xjacobianAD.J, rtol=1e-5)
            @test isapprox(Jmat_x, xjacobianAD.J, rtol=1e-4)

            ## JACOBIAN VECTOR PRODUCT
            ExaPF.jacobian_transpose_product!(polar, pbm, cache, tgt)
            ∂cons = pbm.stack
            @test isapprox(∂cons.∂x, xjacobianAD.J' * tgt, rtol=1e-6)

            ## CONTROL JACOBIAN
            if !isa(cons, typeof(ExaPF.voltage_magnitude_constraints))
                ujacobianAD = ExaPF.AutoDiff.Jacobian(polar, cons, Control())
                # Evaluate Jacobian with AD
                J = AutoDiff.jacobian!(polar, ujacobianAD, cache)
                # Matpower Jacobian
                Jmat_u = ExaPF.matpower_jacobian(polar, Control(), cons, V)

                # Compare with FiniteDiff
                function jac_fd_u(u)
                    cache.vmag[ref] .= u[1:nref]
                    cache.vmag[pv] .= u[nref+1:npv+nref]
                    cache.pinj[pv] .= u[nref+npv+1:end]
                    c = zeros(m)
                    cons(polar, c, cache.vmag, cache.vang, cache.pinj, cache.qinj)
                    return c
                end
                u = [cache.vmag[ref]; cache.vmag[pv]; cache.pinj[pv]]
                Jd = FiniteDiff.finite_difference_jacobian(jac_fd_u, u)
                @test size(J) == (m, length(u))
                @test isapprox(Jd, ujacobianAD.J, rtol=1e-5)
                @test isapprox(Jmat_u, ujacobianAD.J, rtol=1e-6)
                @test isapprox(∂cons.∂u, ujacobianAD.J' * tgt, rtol=1e-6)
            end

            # ADJOINT
            ExaPF.adjoint!(polar, pbm, tgt, c, cache)
            function test_fd(vvm)
                cache.vmag .= vvm[1:nbus]
                cache.vang .= vvm[1+nbus:2*nbus]
                cons(polar, c, cache.vmag, cache.vang, cache.pinj, cache.qinj)
                return dot(c, tgt)
            end
            vv = [cache.vmag; cache.vang]
            vv_fd = FiniteDiff.finite_difference_gradient(test_fd, vv)
            @test isapprox(vv_fd[1:nbus], ∂cons.∂vm, rtol=1e-6)
            @test isapprox(vv_fd[1+nbus:2*nbus], ∂cons.∂va, rtol=1e-6)
        end
    end
end


using Misfits
using Test
using ForwardDiff
using BenchmarkTools
using LinearAlgebra
using Calculus
using LinearMaps

# =================================================
# generalized least-squares
# =================================================
x=abs2.(randn(10,10))
Q=x*x'
pa=Misfits.P_gls(LinearMap(Q))
x=randn(10);
y=randn(10);
dfdx1=similar(x);
@btime Misfits.func_grad!(dfdx1,x,y, pa)
@btime Misfits.func_grad!(nothing,x,y, pa)
f(x)=Misfits.func_grad!(nothing,x,y, pa)
dfdx2=Calculus.gradient(f,x)

@test dfdx1 ≈ dfdx2

# =================================================
# squared_euclidean!
# =================================================
x=randn(1000,10);
y=randn(1000,10);
w=randn(1000,10);
dfdx1=similar(x);
dfdx2=similar(x);

@btime Misfits.error_squared_euclidean!(dfdx1,x,y,w)
@btime Misfits.error_squared_euclidean!(nothing,x,y,w)
f(x)=Misfits.error_squared_euclidean!(nothing,x,y,w)
ForwardDiff.gradient!(dfdx2,f, x);

@test dfdx1 ≈ dfdx2

# =================================================
# error_after_scaling
# =================================================
# real 1D
x = randn(10,1); α = randn(1); y = α[1] .* x;
J, α1 = Misfits.error_after_scaling(x,y)
@test isapprox(α1.*x, y)
#@test_approx_eq J 0.0

# real 2D
x = randn(10,10); α = randn(1); y = α[1] .* x;
J, α1 = Misfits.error_after_scaling(x,y)
@test isapprox(α1.*x, y)

# complex 2D
x = complex.(randn(10,10), randn(10,10));
α = complex.(randn(1), randn(1)); y = α[1] .* x;
J, α1 = Misfits.error_after_scaling(x,y)
@test isapprox(α1.*x, y)

# =================================================
# weighted norm
# =================================================
x=randn(100,10);
w=randn(100,10);
dfdx1=similar(x);
@btime Misfits.error_weighted_norm!(dfdx1,x,w)
xvec=vec(x)
dfdx2=similar(xvec);
#ForwardDiff.gradient!(dfdx2,x -> Misfits.error_weighted_norm!(nothing, reshape(x,100,10), w), xvec);
dfdx2=Calculus.gradient(x -> Misfits.error_weighted_norm!(nothing, reshape(x,100,10), w), xvec)

@test vec(dfdx1) ≈ vec(dfdx2)

# =================================================
# test derivative_vector_magnitude
# =================================================

# some func
function f(x, z)
    y=x./norm(x)
    J=sum((y-z).^2)
    return J
end

# test derivative_vector_magnitude
function g!(g, x, z)
    xn=norm(x)
    rmul!(x, inv(xn))
    g1=similar(g)
    for i in eachindex(g1)
        g1[i]=2. * (x[i]-z[i])
    end
         rmul!(x, xn)
         nx=length(x)
         X=zeros(nx,nx)
         @time Misfits.derivative_vector_magnitude!(g,g1,x,X)
    return g
end

x=randn(10)
z=randn(10)
g1=zero(x)
f1(x)=f(x,z)
ForwardDiff.gradient!(g1,f1, x)
g2=zero(x)
@time g!(g2,x,z)
@test g1 ≈ g2

# error invariant of translation or global phase
#x=randn(100); y=randn() .* circshift(x,20)
#J, α = Misfits.error_after_autocorr_scaling(x,y)
#@test J < 1e-15

#
#x=randn(100,10);
#dfdx1=similar(x);
#@time Misfits.error_pairwise_corr_dist(dfdx1,x)
#xvec=vec(x)
#dfdx2=similar(xvec);
#Inversion.finite_difference!(x -> Misfits.error_pairwise_corr_dist(nothing, reshape(x,100,10)), xvec, dfdx2, :central)
#
#@test dfdx1 ≈ reshape(dfdx2,100,10)
#
#
#x=randn(10,10);
#dfdx1=similar(x);
#Misfits.error_autocorr_pairwise_corr_dist(dfdx1,x)
#xvec=vec(x)
#dfdx2=similar(xvec);
#Inversion.finite_difference!(x -> Misfits.error_autocorr_pairwise_corr_dist(nothing, reshape(x,10,10)), xvec, dfdx2, :central)
#
#@test dfdx1 ≈ reshape(dfdx2,10,10)


using Revise
using ANOVADDPTest
using Gadfly
using StatsModels
using Statistics
using DataFrames
using Distributions
using Random

# # Simulate sample
# function simulate_sample_normal(rng, N)
#     X = rand(rng, 1:2, N, 2);
#     y = randn(rng, N);
#     for i in 1:N
#         ((X[i, 1] == 2) && (X[i, 2] == 2)) && (y[i] += 1)
#     end
#     return y, X
# end

# Simulate sample
function simulate_sample_poisson(rng, N)
    X = rand(rng, 1:2, N, 2);
    y = rand(rng, Poisson(1.0), N);
    for i in 1:N
        ((X[i, 1] == 2) && (X[i, 2] == 2)) && (y[i] += 1)
    end
    return y, X
end

# Example
N = 1000;
rng = MersenneTwister(1);
y, X = simulate_sample_poisson(rng, N);
fit = anova_bnp_poisson(y, X);
show(group_probs(fit); allrows = true)
show(group_codes(fit); allrows = true)
plot(
    fpost(fit),
    x = :y,
    y = :f,
    color = :group,
    Geom.line,
    Scale.color_discrete()
)


# %% Simulate progression of wealth over many rounds with different fees
# using CairoMakie
using Distributions
using DataFrames
using Gadfly
import Cairo, Fontconfig
using Printf

# %% Functions
g = Geometric(0.5)
toss() = rand(g)+1

function symlog(x)
    return sign(x)*log10(abs(x)+1)
end

function calcWinnings(N::Int,P::Int, F::Float64=log2(N)+2)
    winnings = zeros(Float64, N, P)
    for p in 1:P
        for i in 2:N
            winnings[i,p] = winnings[i-1,p] + 2^toss() - F
        end
    end
    return winnings
end

# %% Figure setup
theme = Theme(major_label_font="PragmataPro Mono",
              minor_label_font="PragmataPro Mono",
              key_title_font="PragmataPro Mono Bold",
              key_position=:none)
Gadfly.push_theme(theme)

nbins = 1000

function plotWinnings(winnings, fee)
    fig = plot(
        winnings,
        x=Row.index,
        y=Col.value,
        Geom.hexbin(xbincount=nbins, ybincount=nbins),
        Guide.xlabel("ROUNDS"),
        Guide.ylabel("CUMULATIVE PAYOUT (SYMLOG)"),
        Guide.xticks(ticks=[0 2500 5000 7500 10000]),
        Guide.yticks(label=false),
        Guide.title(@sprintf("Fee Per Round = %s", fee)),
        Coord.cartesian(xmin=0, xmax=1e4, ymin=-8, ymax=8))
    return fig
end

# %% Plot figures
wN = calcWinnings(Int(1e4), Int(1e3)) .|> symlog
draw(PDF("../../figures/winnings-N.pdf"), plotWinnings(wN, "log2(10000)+2"))

# %%
df2_wide = DataFrame(w2, :auto)
cols = names(df2_wide)
df2_wide.x = 1:(size(df2_wide)[1])
df2_wide.f .= "P=\$2"
df2 = stack(df2_wide, cols)

df100_wide = DataFrame(w100, :auto)
cols = names(df100_wide)
df100_wide.x = 1:(size(df100_wide)[1])
df100_wide.f .= "P=\$100"
df100 = stack(df100_wide, cols)

df = vcat(df2, df100)

# %%
p = plot(df,
         x=:x,
         y=:value,
         ygroup=:f,
         Geom.subplot_grid(Geom.hexbin(xbincount=nbins, ybincount=nbins),
                           Coord.cartesian(xmin=0, xmax=1e4, ymin=-8, ymax=8),
                           # Guide.yticks(ticks=nothing), Does not work :(
                           Guide.xticks(ticks=[0 2500 5000 7500 10000]),
                           ),
         Guide.xlabel("ROUNDS"),
         Guide.ylabel("CUMULATIVE PAYOUT (SYMLOG)"),
         Guide.title("Distribution of Payouts over Time"))

draw(PDF("../../figures/winnings-2-100.pdf"), p)
# %%

Guide.yticks


"""
    fit(template::objective_function_template,
        data::Any,
        theta::Vector{Float64};
        estimation_method::String = "M",
        br_method::String = "implicit_trace",
        regularizer::Function = function regularizer(theta::Vector{Float64}, data::Any) Vector{Float64}() end,
        lower::Vector{Float64} = Vector{Float64}(),
        upper::Vector{Float64} = Vector{Float64}(),
        optim_method = LBFGS(),
        optim_options = Optim.Options(),
        optim_arguments...)

Fit an [`objective_function_template`](@ref) on `data` using M-estimation (`estimation_method = "M"`; default) or RBM-estimation (reduced-bias M estimation; [Kosmidis & Lunardon, 2020](http://arxiv.org/abs/2001.03786); `estimation_method = "RBM"`)

Arguments
===

+ `template`: an [`objective_function_template`](@ref) object.
+ `data`: typically an object of [composite type](https://docs.julialang.org/en/v1/manual/types/#Composite-Types-1) with all the data required to compute the `objective_function`.
+ `theta`: a `Vector{Float64}` of parameter values to use as starting values in `Optim.optimize`.

[Keyword arguments](https://docs.julialang.org/en/v1/manual/functions/#Keyword-Arguments-1)
===

+ `estimation_method`: either "M" (default) or "RBM"; see Details.
+ `br_method`: either "implicit_trace" (default) or "explicit_trace"; see Details.
+ `regularizer`: a function of `theta` and `data` returning a `Float64`, which is added to the (bias-reducing penalized) objective; the default value will result in no regularization.
+ `lower`: a `Vector{Float64}` of dimension equal to `theta` for setting box constraints for the optimization. The default will result in unconstrained optimization. See Details.
+ `upper`: a `Vector{Float64}` of dimension equal to `theta` for setting box constraints for the optimization. The default will result in unconstrained optimization. See Details.
+ `optim_method`: the optimization method to be used; deafult is `Optim.LBFGS()`. See Details.
+ `optim_options`: the result of a call to `Optim.Options` to be passed to `Optim.optimize`. Default is `Optim.Options()`. See details.
+ `optim_arugments...`: extra keyword arguments to be passed to `Optim.optimize`. See Details.

Details
===

Bias reduction is either through the maximization of the bias-reducing penalized objective in Kosmidis & Lunardon (2020) (`br_method = "implicit_trace"`; default) or by subtracting an estimate of the bias from the M-estimates (`br_method = "explicit_trace"`). The bias-reducing penalty is constructed internally using automatic differentiation (using the [**ForwardDiff**](https://github.com/JuliaDiff/ForwardDiff.jl) package), and the bias estimate using a combination of automatic differentiation and numerical differentiation (using the [**FiniteDiff**](https://github.com/JuliaDiff/FiniteDiff.jl) package).

The maximization of the objective or the penalized objective is done using the [**Optim**](https://github.com/JuliaNLSolvers/Optim.jl) package. Optimization methods and options can be supplied directly through the `optim_method` and `optim_options`, respectively. `optim_options` expects an object constructed through `Optim.Options`. Keyword arguments (e.g. `autodiff = :forward`) can be passed directly to `Optim.optimize` through extra keyword arguments. See the [Optim documentation](https://julianlsolvers.github.io/Optim.jl/stable/#user/config/#general-options) for more details on the available options.

An extra additive regularizer to either the objective or the bias-reducing penalized objective can be suplied via the keyword argument `regularizer`, which must be a scalar-valued function of the parameters and the data; the default value will result in no regularization.

`lower` and `upper` can be used to provide box contraints. If valid `lower` and `upper` vectors are supplier, then the internal call to `Optim.optimize` will use `Fminbox(optim_method)` as a method; see the [**Optim** documentation on box minimization](https://julianlsolvers.github.io/Optim.jl/stable/#examples/generated/ipnewton_basics/#box-minimzation) for more details.
"""
function fit(template::objective_function_template,
             data::Any,
             theta::Vector{Float64};
             estimation_method::String = "M",
             br_method::String = "implicit_trace",
             regularizer::Function = function regularizer(theta::Vector{Float64}, data::Any) Vector{Float64}() end,
             lower::Vector{Float64} = Vector{Float64}(),
             upper::Vector{Float64} = Vector{Float64}(),
             optim_method = LBFGS(),
             optim_options = Optim.Options(),
             optim_arguments...)
    if (estimation_method == "M")
        br = false
    elseif (estimation_method == "RBM")
        if (br_method == "implicit_trace")
            br = true
        elseif (br_method == "explicit_trace")
            br = false
        else
            error(br_method, " is not a recognized bias-reduction method")
        end
    else
        error(estimation_method, " is not a recognized estimation method")
    end

    ## Well we only need scalar-valued penalties here but let's leave
    ## it like this
    has_regularizer = length(regularizer(theta, data)) > 0

    ## down the line when/if det is implemented we need to be passing the
    ## bias reduction method to objetive_function
    if has_regularizer
        obj = beta -> - objective_function(beta, data, template, br) - regularizer(beta, data)
    else
        obj = beta -> - objective_function(beta, data, template, br)
    end

    if (length(lower) > 0) & (length(upper) > 0)
        out = optimize(obj, lower, upper, theta, Fminbox(optim_method), optim_options; optim_arguments...)
    else
        out = optimize(obj, theta, optim_method, optim_options; optim_arguments...)
    end
    if (estimation_method == "M")
        theta = out.minimizer
    elseif (estimation_method == "RBM")
        ## FIXME, 12/03/2020: Add stop if br_method is not implicit_trace
        if (br_method == "implicit_trace")
            theta = out.minimizer
        elseif (br_method == "explicit_trace")
            quants = obj_quantities(out.minimizer, data, template, true)
            jmat_inv = quants[2]
            ## We use finite differences to get the adjustment
            adjustment = FiniteDiff.finite_difference_gradient(beta -> obj_quantities(beta, data, template, true)[1], out.minimizer)
            theta = out.minimizer + jmat_inv * adjustment
            ## Reset br
            br = true
        end
    end
    MEstimation_results(out, theta, data, template, regularizer, br, true, has_regularizer, br_method)
end

"""
    fit(template::estimating_function_template,
        data::Any,
        theta::Vector{Float64};
        estimation_method::String = "M",
        br_method::String = "implicit_trace",
        concentrate::Vector{Int64} = Vector{Int64}(),
        regularizer::Function = function regularizer(theta::Vector{Float64}, data::Any) Vector{Float64}() end,
        nlsolve_arguments...)

Fit an [`estimating_function_template`](@ref) on `data` using M-estimation (`estimation_method = "M"`; default) or RBM-estimation (reduced-bias M estimation; [Kosmidis & Lunardon, 2020](http://arxiv.org/abs/2001.03786); `estimation_method = "RBM"`)

Arguments
===
+ `template`: an [`estimating_function_template`](@ref) object.
+ `data`: typically an object of [composite type](https://docs.julialang.org/en/v1/manual/types/#Composite-Types-1) with all the data required to compute the `objective_function`.
+ `theta`: a `Vector{Float64}` of parameter values to use as starting values in `Optim.optimize`.

[Keyword arguments](https://docs.julialang.org/en/v1/manual/functions/#Keyword-Arguments-1)
===
+ `estimation_method`: either "M" (default) or "RBM"; see Details.
+ `br_method`: either "implicit_trace" (default) or "explicit_trace"; see Details.
+ `concentrate`: a `Vector{Int64}`; if specified, empirical bias-reducing adjustments are added only to the subset of estimating functions indexed by `concentrate`. The default is to add empirical bias-reducing adjustments to all estimating functions.
+ `regularizer`: a function of `theta` and `data` returning a `Vector{Float64}` of dimension equal to the number of the estimating functions, which is added to the (bias-reducing) estimating function; the default value will result in no regularization.
+ `nlsolve_arguments...`: extra keyword arguments to be passed to `NLsolve.nlsolve`. See Details.

Details
===

Bias reduction is either through the solution of the empirically adjusted estimating functions in Kosmidis & Lunardon (2020) (`br_method = "implicit_trace"`; default) or by subtracting an estimate of the bias from the M-estimates (`br_method = "explicit_trace"`). The bias-reducing adjustments and the bias estimate are constructed internally using automatic differentiation (using the [**ForwardDiff**](https://github.com/JuliaDiff/ForwardDiff.jl) package).

Bias reduction for only a subset of parameters can be performed by setting `concentrate` to the vector of the indices for those parameters.

The solution of the estimating equations or the adjusted estimating equations is done using the [**NLsolve**](https://github.com/JuliaNLSolvers/NLsolve.jl) package. Keyword arguments can be passed directly to `NLsolve.nlsolve` through extra keyword arguments. See the [NLsolve README](https://github.com/JuliaNLSolvers/NLsolve.jl) for more information on available options.

An extra additive regularizer to either the estimating functions or the bias-reducing adjusted estimating functions can be suplied via the keyword argument `regularizer`, which must be a `length(theta)`-valued function of the parameters and the data; the default value will result in no regularization.
"""
function fit(template::estimating_function_template,
             data::Any,
             theta::Vector{Float64};
             estimation_method::String = "M",
             br_method::String = "implicit_trace",
             concentrate::Vector{Int64} = Vector{Int64}(),
             regularizer::Function = function regularizer(theta::Vector{Float64}, data::Any) Vector{Float64}() end,
             nlsolve_arguments...)
    if (estimation_method == "M")
        br = false
    elseif (estimation_method == "RBM")
        if (br_method == "implicit_trace")
            br = true
        elseif (br_method == "explicit_trace")
            br = false
        else
            error(br_method, " is not a recognized bias-reduction method")
        end
    else
        error(estimation_method, " is not a recognized estimation method")
    end

    ## Well we only need length(theta)-valued penalties here but let's leave
    ## it like this
    has_regularizer = length(regularizer(theta, data)) > 0

    ## down the line when/if det is implemented we need to be passing the
    ## bias reduction method to objetive_function
    if has_regularizer
        ef = get_estimating_function(data, template, br, concentrate, regularizer)
    else
        function non_regularizer(theta::Vector{Float64}, data::Any) fill(0, length(theta)) end
        ef = get_estimating_function(data, template, br, concentrate, non_regularizer)
    end

    out = nlsolve(ef, theta; nlsolve_arguments...)
    if (estimation_method == "M")
        theta = out.zero
    elseif (estimation_method == "RBM")
        if (br_method == "implicit_trace")
            theta = out.zero
        elseif (br_method == "explicit_trace")
            quants = ef_quantities(out.zero, data, template, true, concentrate)
            adjustment = quants[1]
            jmat_inv = quants[2]
            theta = out.zero + jmat_inv * adjustment
            ## Reset br
            br = true
        end
    end
    MEstimation_results(out, theta, data, template, regularizer, br, false, has_regularizer, br_method)
end


using Turing

@model gaussmodel(x) = begin
    N = length(x)
    lam = tzeros(Real,N)
    mu ~ Normal(0, sqrt(1000));
    for i = 1:N
      lam[i] ~ Gamma(1, 1)
      x[i] ~ Normal(mu, sqrt(1 / (sqrt(lam[i]))))
    end
    return(mu, lam)
end


# PREAMBLE

# PKG_SETUP

# ## Setup

using IterativeLQR
using LinearAlgebra
using Plots

# ## horizon
T = 51

# ## car
num_state = 3
num_action = 2
num_parameter = 0

function car(x, u, w)
    [u[1] * cos(x[3]); u[1] * sin(x[3]); u[2]]
end

function midpoint_explicit(x, u, w)
    h = 0.1 # timestep
    x + h * car(x + 0.5 * h * car(x, u, w), u, w)
end

# ## model
dynamics = Dynamics(midpoint_explicit, num_state, num_action, num_parameter)
model = [dynamics for t = 1:T-1]

# ## initialization
x1 = [0.0; 0.0; 0.0]
xT = [1.0; 1.0; 0.0]

# ## rollout
ū = [1.0e-2 * [1.0; 0.1] for t = 1:T-1]
w = [zeros(num_parameter) for t = 1:T]
x̄ = rollout(model, x1, ū, w)

# ## objective
ot = (x, u, w) -> 1.0 * dot(x - xT, x - xT) + 1.0e-2 * dot(u, u)
oT = (x, u, w) -> 1000.0 * dot(x - xT, x - xT)
ct = Cost(ot, num_state, num_action, num_parameter)
cT = Cost(oT, num_state, 0, num_parameter)
objective = [[ct for t = 1:T-1]..., cT]

# ## constraints
ul = -5.0 * ones(num_action)
uu = 5.0 * ones(num_action)

p_obs = [0.5; 0.5]
r_obs = 0.1

function stage_con(x, u, w)
    e = x[1:2] - p_obs
    [
    ul - u; # control limit (lower)
    u - uu; # control limit (upper)
    r_obs^2.0 - dot(e, e); # obstacle
    ]
end

function terminal_con(x, u, w)
    e = x[1:2] - p_obs
    [
    x - xT; # goal
    r_obs^2.0 - dot(e, e); # obstacle
    ]
end

cont = Constraint(stage_con, num_state, num_action, indices_inequality=collect(1:5))
conT = Constraint(terminal_con, num_state, num_action, indices_inequality=collect(3 .+ (1:1)))
constraints = [[cont for t = 1:T-1]..., conT]

# ## problem
prob = Solver(model, objective, constraints)
initialize_controls!(prob, ū)
initialize_states!(prob, x̄)

# ## solve
solve!(prob)

# ## solution
x_sol, u_sol = get_trajectory(prob)

# ## visualize
plot(hcat(x_sol...)')
plot(hcat(u_sol[1:end-1]...)', linetype=:steppost)

# ## benchmark allocations + timing
using BenchmarkTools
info = @benchmark solve!($prob, x̄, ū) setup=(x̄=deepcopy(x̄), ū=deepcopy(ū))


#Get number of detections for each site
function ndetects(y::Array)
  function sum2(row::Array)
    sum(skipmissing(row))
  end
  mapslices(sum2, y, dims=2)
end

#Simulate random covariate frame based on provided formula
function gen_covs(f::FormulaTerm, n::Int)
  covs = varnames(f).rhs
  if isnothing(covs) return DataFrame(_dummy=ones(n)) end
  nc = length(covs)
  out = DataFrame(reshape(rand(Normal(0,1), n*nc), n, nc), :auto)
  DataFrames.rename!(out, covs)
  out
end

#Replicate missing values in a new simulated y matrix
function rep_missing!(ynew::Array, y::Array)
  na_idx = findall(ismissing.(y))
  if length(na_idx) == 0 return nothing end
  ynew[na_idx] = fill(missing, length(na_idx))
end

#Get appropriate value of K
function check_K(K::Int, y::Array{Union{Missing,Int},2})
  if(K < (maximum(skipmissing(y))))
    error("K should be larger than the maximum observed abundance")
  end
  K
end

function check_K(K::Nothing, y::Array{Union{Missing,Int},2})
  maximum(skipmissing(y)) + 20
end

#Get minimum bound for integrating over possible abundance values K
#Minimum bound = maximum abundance ever observed at a site
function get_Kmin(y::Array{Union{Missing,Int64},2})
  N = size(y, 1)
  kmin = zeros(Int64, N)
  for n in 1:N
    if(all(ismissing.(y[n, :]))) continue end
    kmin[n] = maximum(skipmissing(y[n, :]))
  end
  kmin
end


# Regex to validate a Roman numeral
const VALID_ROMAN_PATTERN =
    r"""
    ^\s*                   # Skip leading whitespace
    (
    M*                     # Thousands
    (C{0,9}|CD|DC{0,4}|CM) # Hundreds
    (X{0,9}|XL|LX{0,4}|XC) # Tens
    (I{0,9}|IV|VI{0,4}|IX) # Ones
    )
    \s*$                   # Skip trailing whitespace
    """ix # Be case-insensitive and verbose

const NUMERAL_MAP = [
    (1000, "M")
    (900,  "CM")
    (500,  "D")
    (400,  "CD")
    (100,  "C")
    (90,   "XC")
    (50,   "L")
    (40,   "XL")
    (10,   "X")
    (9,    "IX")
    (5,    "V")
    (4,    "IV")
    (1,    "I")
]

function parseroman(str::ASCIIString)
    m = match(VALID_ROMAN_PATTERN, str)
    m == nothing && throw(InvalidRomanError(str))
    # Strip whitespace
    str = m.captures[1]
    # Make `str` uppercase
    if !isupper(str); str = uppercase(str); end
    i = 1
    val = 0
    strlen = length(str)
    for (num_val, numeral) in NUMERAL_MAP
        numlen = length(numeral)
        while i+numlen-1 <= strlen && str[i:i+numlen-1] == numeral
            val += num_val
            i += numlen
        end
    end
    val
end

function toroman(val::Integer)
    val <= 0 && throw(DomainError())
    val > 5000 && warn("Roman numerals do not handle large numbers well")

    str = ""
    for (num_val, numeral) in NUMERAL_MAP
        i = div(val, num_val)
        # Never concatenate an empty string to `str`
        i == 0 && continue
        str *= (numeral^i)
        val -= i*num_val
        # Stop when ready
        val == 0 && break
    end
    str
end


path = "docs/src/assets/"

using ISA, Plots
𝐶₀ = AMFMtriplet(t->exp(-t^2),t->100,0.0)
ψ₀ = AMFMcomp(𝐶₀)
plot(ψ₀; timeaxis=-1.0:0.0005:1.0)
png(path*"IS_exAM1.png")

using ISA, Plots
𝐶₀ = AMFMtriplet(t->exp(-t^2),t->100,0.0)
ψ₀ = AMFMcomp(𝐶₀)
plot(𝐶₀; timeaxis=-1.0:0.0005:1.0)
png(path*"IS_exAM2.png")

using ISA, Plots
𝐶₀ = AMFMtriplet(t->exp(-t^2),t->100,0.0)
ψ₀ = AMFMcomp(𝐶₀)
plot(𝐶₀; timeaxis=-1.0:0.0005:1.0, view="TF",
     left_margin=15Plots.mm, margin=5Plots.mm)
png(path*"IS_exAM3.png")

# -----------------------------

using ISA, Plots
𝐶₀ = AMFMtriplet(t->sin(11t),t->250,0.0)
ψ₀ = AMFMcomp(𝐶₀)
plot(ψ₀; timeaxis=-1.0:0.0005:1.0)
png(path*"IS_exAM4.png")

using ISA, Plots
𝐶₀ = AMFMtriplet(t->sin(11t),t->250,0.0)
ψ₀ = AMFMcomp(𝐶₀)
plot(𝐶₀; timeaxis=-1.0:0.0005:1.0)
png(path*"IS_exAM5.png")

using ISA, Plots
𝐶₀ = AMFMtriplet(t->sin(11t),t->250,0.0)
ψ₀ = AMFMcomp(𝐶₀)
plot(𝐶₀; timeaxis=-1.0:0.0005:1.0, view="TF",
     left_margin=15Plots.mm, margin=5Plots.mm)
png(path*"IS_exAM6.png")

#------------------------------------

using ISA, Plots
𝐶₀ = AMFMtriplet(t->cos(t)+exp(-t),t->100,1.0)
ψ₀ = AMFMcomp(𝐶₀)
plot(ψ₀; timeaxis=-1.0:0.0005:1.0)
png(path*"IS_exAM7.png")

using ISA, Plots
𝐶₀ = AMFMtriplet(t->cos(t)+exp(-t),t->100,1.0)
ψ₀ = AMFMcomp(𝐶₀)
plot(𝐶₀; timeaxis=-1.0:0.0005:1.0)
png(path*"IS_exAM8.png")

using ISA, Plots
𝐶₀ = AMFMtriplet(t->cos(t)+exp(-t),t->100,1.0)
ψ₀ = AMFMcomp(𝐶₀)
plot(𝐶₀; timeaxis=-1.0:0.0005:1.0, view="TF",
     left_margin=15Plots.mm, margin=5Plots.mm)
png(path*"IS_exAM9.png")


function (SL::SingleLayerKernel{T,S})(x,y)::T  where {T,S<:Maxwell}
    N = ambient_dimension(S)
    x==y && return zero(T)
    k = SL.op.k
    r = x - y
    d = norm(r)
    if N==2
        return error("Maxwell operator not implemented in 2d")
    elseif N==3
        g   = 1/(4π)/d * exp(im*k*d)
        gp  = im*k*g - g/d
        gpp = im*k*gp - gp/d + g/d^2
        ID    = SMatrix{3,3,Float64,9}(1,0,0,0,1,0,0,0,1)
        RRT   = r*transpose(r) # r ⊗ rᵗ
        return  g*ID + 1/k^2*(gp/d*ID + (gpp/d^2 - gp/d^3)*RRT)
    end
end

# Double Layer Kernel
# n × ∇ × G = γ₁ G
function (DL::DoubleLayerKernel{T,S})(x,y,ny)::T where {T,S<:Maxwell}
    N = ambient_dimension(S)
    x==y && return zero(T)
    k = DL.op.k
    r = x - y
    d = norm(r)
    g   = 1/(4π)/d * exp(im*k*d)
    gp  = im*k*g - g/d
    if N==2
        return error("Maxwell operator not yet defined in 2d")
    elseif N==3
        ID    = SMatrix{3,3,Float64,9}(1,0,0,0,1,0,0,0,1)
        ncross = transpose(SMatrix{3,3,Float64,9}(0,-ny[3],ny[2],
                                         ny[3],0,-ny[1],
                                         -ny[2],ny[1],0))
        rcross = transpose(SMatrix{3,3,Float64,9}(0,-r[3],r[2],
                                         r[3],0,-r[1],
                                         -r[2],r[1],0))
        # return -gp/d*ncross*rcross
        return gp/d*rcross*ncross
    end
end


module WeightTest

using TestSetup, OnlineStats, Base.Test
O = OnlineStats

@testset "Weighting" begin
    @testset "EqualWeight" begin
        w = EqualWeight()
        O.updatecounter!(w)
        @test O.weight(w) == 1.0
        O.updatecounter!(w, 5)
        @test O.weight(w, 5) == 5 / 6
    end
    @testset "ExponentialWeight" begin
        w = ExponentialWeight(.5)
        O.updatecounter!(w)
        @test O.weight(w) == 0.5
        O.updatecounter!(w, 5)
        @test O.weight(w) == 0.5
        @test ExponentialWeight(5).λ == ExponentialWeight(1/3).λ
    end
    @testset "BoundedEqualWeight" begin
        w = BoundedEqualWeight(.1)
        O.updatecounter!(w)
        @test O.weight(w, 1) == 1.0
        O.updatecounter!(w, 100)
        @test O.weight(w) == 0.1
        @test BoundedEqualWeight(5).λ == BoundedEqualWeight(1/3).λ
    end
    @testset "LearningRate" begin
        w = LearningRate(.6)
        O.updatecounter!(w)
        @test O.weight(w) == 1.0
        O.updatecounter!(w, 5)
        @test O.weight(w, 5) == 2 ^ -.6
        @test O.nups(w) == 2
    end
    @testset "LearningRate2" begin
        w = LearningRate2(0.5)
        O.updatecounter!(w)
        @test O.weight(w) == 1.0
        O.updatecounter!(w, 5)
        @test O.weight(w, 5) == 1 / (1 + .5)
    end
end  # facts
end  # module


function get_ub_lb_lanczos(pw::PWGrid, Potentials::PotentialsT, nlancz::Int64)
    #
    Ngwx = pw.gvecw.Ngwx
    V = zeros(ComplexF64,Ngwx,nlancz)
    HV = zeros(ComplexF64,Ngwx,nlancz)
    T = zeros(Float64,nlancz,nlancz)
    f = zeros(ComplexF64,Ngwx)
    s = zeros(ComplexF64,nlancz)
    h = zeros(ComplexF64,nlancz)
    #
    V[:,1] = randn(Ngwx) + im*randn(Ngwx)
    beta = norm(V[:,1])
    V[:,1] = V[:,1] ./ beta
    #
    HV[:,1] = op_H( pw, Potentials, V[:,1] )
    h[1] = real( V[:,1]' * HV[:,1] )
    #
    T[1,1] = h[1]
    # One-step of reorthogonalization
    f[:] = HV[:,1] - V[:,1]*h[1]
    s[1] = V[:,1]' * f[:]
    h[1] = h[1] + s[1]
    f[:] = f[:] - V[:,1]*s[1]
    # MAIN LOOP
    for j = 2:nlancz
        #@printf("iter lanczos = %d\n", j)
        beta = norm(f)
        T[j,j-1] = beta
        V[:,j] = f[:]/beta
        HV[:,j] = op_H( pw, Potentials, V[:,j] )
        #
        for jj = 1:j
            h[jj] = V[:,jj]' * HV[:,j]
        end
        f[:] = HV[:,j] - V[:,1:j]*h[1:j]

        # One-step of reorthogonalization
        #s = V' * f
        #h = h + s
        for jj=1:j
            s[jj] = V[:,jj]' * f[:]
            h[jj] = h[jj] + s[jj]
        end
        f[:] = f[:] - V[:,1:j]*s[1:j]
        #
        T[1:j,j] = real(h[1:j])
    end
    #
    evalsT = eigvals(T)
    #lb = evalsT[Nstates+2]
    #ub = evalsT[2*Nstates]
    lb = evalsT[Int(nlancz/2)]
    ub = norm_matrix_induced(T) + norm(f)
    #
    return lb, ub
end


# Implements the method of PPT relaxations of MX operations

using Convex
using SCS

export pptRelax, pptRelaxCopies;

""" `(problem, F, p_succ) = pptRelax(rho,n,k,delta, verbose, eps, max_iters)`

Implements the PPT relaxation computing an upper bound on the fidelity achievable using measure and exchange operations for a given input state and fixed succes probability.

Inputs:
- *rho* quantum state to be distilled
- *n* number of qubits on one side, assuming dimensions nA=nB =2^n in the input.
- *k* desired output dimension of the maximally entangled state
- *delta* desired success probability
- *verbose* flag indicating whether the SCS solver should be verbose (default false)
- *eps* desired precision (default 1e-4) in SCS solver
- *max_iters* maximum number of iterations in SCS solver (default 2e4)

Outputs:
- *problem* SDP object from Convex.jl allowing further processing.
- *F* Fidelity achieved
- *p_succ* success probability attained in SDP (should equal *delta*)
"""

function pptRelax(rho::AbstractMatrix, n::Int, k::Int, delta::Real; verbose::Bool = false, eps::Number = 1e-4, max_iters = 2e4)
	# Presume the input dimensions are equally divided and we're using qubits
	nA = 2^n;
	nB = nA;

	return pptRelax(rho, nA, nB, k, delta, verbose = verbose, eps = eps, max_iters = max_iters)
end

""" `(problem, F, p_succ) = pptRelax(rho,nA, nB, k,delta, verbose, eps, max_iters)`

Implements the PPT relaxation computing an upper bound on the fidelity achievable using measure and exchange operations for a given input state and fixed succes probability.

Inputs:
- *rho* quantum state to be distilled on A and B
- *nA* dimension of the A system
- *nB* dimension of the B system
- *k* desired output dimension of the maximally entangled state
- *delta* desired success probability
- *verbose* flag indicating whether the SCS solver should be verbose (default true)
- *eps* desired precision (default 1e-4) in SCS solver
- *max_iters* maximum number of iterations in SCS solver (default 2e4)

Outputs:
- *problem* SDP object from Convex.jl allowing further processing.
- *F* Fidelity achieved
- *p_succ* success probability attained in SDP (should equal *delta*)
"""

function pptRelax(rho::AbstractArray, nA::Int, nB::Int, k::Int, delta::Number; verbose::Bool = true, eps::Number = 1e-4, max_iters::Number = 2e4)

	# Check whether rho is a quantum state
	@assert isQuantumState(rho) "The input state must be a valid quantum state."

	# Check whether dimensions match
	(d, db) = size(rho);
	@assert d == nA * nB "Input state doesn't match given dimensions." , d , " != ", nA*nB

	# Check whether delta is a valid probability
	@assert delta <= 1 "Success probability must be less than one."
	@assert 0 <= delta "Success probability must be positive."

	# define the identity matrix on the whole space
	id = eye(d);

	# define the variables - if the input is real it is sufficient to optimize over reals
	if real(rho) == rho
		D = Semidefinite(d);
		E = Semidefinite(d);
	else
		# complex variable support advertises HermitianSemidefinite
		# but it's buggy
		D = ComplexVariable(d,d);
		E = ComplexVariable(d,d);
	end

	# define the objective
	tv = Variable(1);
	problem = maximize(tv);
	problem.constraints += tv == trace(nA * nB * D * rho.');

	problem.constraints += D in :SDP;
	problem.constraints += E in :SDP;
	problem.constraints += (id/d - (D+E)) in :SDP;

	# Constraints relating to the PPT Condition
	EPT = ptranspose(E);
	DPT = ptranspose(D);
	problem.constraints += (DPT + EPT/(k+1)) in :SDP;
	problem.constraints += (- DPT + EPT/(k-1)) in :SDP;
	problem.constraints += (id/d - (DPT+EPT)) in :SDP;

	# Constraint coming from the success probability
	problem.constraints += nA * nB * trace(rho.'*(D+E)) == delta;

	solve!(problem, SCSSolver(verbose = verbose, eps = eps, max_iters = max_iters),verbose=verbose);
	p_succ = nA * nB * trace(rho.'*(D.value + E.value));
	F = problem.optval / p_succ;

	return (problem, F, p_succ)
end

""" `(problem, F, p) = pptRelaxCopies(rho, n, nA, nB, k, delta, verbose, eps, max_iters)`

Calls pptRelax for an input state of the form *rho&*^(tensor *n*), reordered appropriately.

"""

function pptRelaxCopies(rho::AbstractMatrix,n::Int,nA::Int,nB::Int,k::Int,delta::Real; verbose::Bool = false, eps::Number = 1e-4, max_iters::Number = 2e4)

	@assert n > 1 "Need at least two copies."

	# copy rho n times
	rhoNew = copies(rho, n);

  	# sort all the entries, such that all the systems on A are first and all
	# the systems on B are second
	rhoSorted = sortAB(rhoNew, nA, n)

	return pptRelax(rhoSorted,nA^n,nB^n, k, delta, verbose=verbose, eps=eps,max_iters=max_iters);
end


# Scalar product and norm: only valid for EuclideanSpace
Base.dot{S<:EuclideanSpace}(t1::Tensor{S},t2::Tensor{S})= (space(t1)==space(t2) ? dot(vec(t1),vec(t2)) : throw(SpaceError()))
Base.vecnorm{S<:EuclideanSpace}(t::Tensor{S})=vecnorm(t.data) # frobenius norm

# # Factorizations:
# #-----------------
# for (S,TT) in ((CartesianSpace,CartesianTensor),(ComplexSpace,ComplexTensor))
#     @eval function svd!(t::$TT,n::Int,::NoTruncation)
#         # Perform singular value decomposition corresponding to bipartion of the
#         # tensor indices into the left indices 1:n and remaining right indices,
#         # thereby destroying the original tensor.
#         N=numind(t)
#         spacet=space(t)
#         leftspace=spacet[1:n]
#         rightspace=spacet[n+1:N]
#         leftdim=dim(leftspace)
#         rightdim=dim(rightspace)
#         data=reshape(t.data,(leftdim,rightdim))
#         F=svdfact!(data)
#         newdim=length(F[:S])
#         newspace=$S(newdim)
#         U=tensor(F[:U],leftspace*newspace')
#         S=tensor(diagm(F[:S]),newspace*newspace')
#         V=tensor(F[:Vt],newspace*rightspace)
#         return U,S,V,abs(zero(eltype(t)))
#     end
#
#     @eval function svd!(t::$TT,n::Int,trunc::Union(TruncationDimension,TruncationSpace))
#         # Truncate rank corresponding to bipartition into left indices 1:n
#         # and remain right indices, based on singular value decomposition,
#         # thereby destroying the original tensor.
#         # Truncation parameters are given as keyword arguments: trunctol should
#         # always be one of the possible arguments for specifying truncation, but
#         # truncdim can be replaced with different parameters for other types of tensors.
#
#         N=numind(t)
#         spacet=space(t)
#         leftspace=spacet[1:n]
#         rightspace=spacet[n+1:N]
#         leftdim=dim(leftspace)
#         rightdim=dim(rightspace)
#         dim(trunc) >= min(leftdim,rightdim) && return svd!(t,n)
#         data=reshape(t.data,(leftdim,rightdim))
#         F=svdfact!(data)
#         sing=F[:S]
#
#         # compute truncation dimension
#         if eps(trunc)!=0
#             sing=F[:S]
#             normsing=vecnorm(sing)
#             truncdim=0
#             while truncdim<length(sing) && vecnorm(sing[(truncdim+1):end])>eps(trunc)*normsing
#                 truncdim+=1
#             end
#             truncdim=min(truncdim,dim(trunc))
#         else
#             truncdim=dim(trunc)
#         end
#         newspace=$S(truncdim)
#
#         # truncate
#         truncerr=vecnorm(sing[(truncdim+1):end])/vecnorm(sing)
#         U=tensor(F[:U][:,1:truncdim],leftspace*newspace')
#         S=tensor(diagm(sing[1:truncdim]),newspace*newspace')
#         V=tensor(F[:Vt][1:truncdim,:],newspace*rightspace)
#         return U,S,V,truncerr
#     end
#
#     @eval function svd!(t::$TT,n::Int,trunc::TruncationError)
#         # Truncate rank corresponding to bipartition into left indices 1:n
#         # and remain right indices, based on singular value decomposition,
#         # thereby destroying the original tensor.
#         # Truncation parameters are given as keyword arguments: trunctol should
#         # always be one of the possible arguments for specifying truncation, but
#         # truncdim can be replaced with different parameters for other types of tensors.
#
#         N=numind(t)
#         spacet=space(t)
#         leftspace=spacet[1:n]
#         rightspace=spacet[n+1:N]
#         leftdim=dim(leftspace)
#         rightdim=dim(rightspace)
#         data=reshape(t.data,(leftdim,rightdim))
#         F=svdfact!(data)
#
#         # find truncdim based on trunctolinfo
#         sing=F[:S]
#         normsing=vecnorm(sing)
#         truncdim=0
#         while truncdim<length(sing) && vecnorm(sing[(truncdim+1):end])>eps(trunc)*normsing
#             truncdim+=1
#         end
#         newspace=$S(truncdim)
#
#         # truncate
#         truncerr=vecnorm(sing[(truncdim+1):end])/normsing
#         U=tensor(F[:U][:,1:truncdim],leftspace*newspace')
#         S=tensor(diagm(sing[1:truncdim]),newspace*newspace')
#         V=tensor(F[:Vt][1:truncdim,:],newspace*rightspace)
#         return U,S,V,truncerr
#     end
#
#     @eval function leftorth!(t::$TT,n::Int,::NoTruncation)
#         # Create orthogonal basis U for indices 1:n, and remainder C for right
#         # indices, thereby destroying the original tensor.
#         # Decomposition should be unique, such that it always returns the same
#         # result for the same input tensor t. UC = QR is fastest but only unique
#         # after correcting for phases.
#         N=numind(t)
#         spacet=space(t)
#         leftspace=spacet[1:n]
#         rightspace=spacet[n+1:N]
#         leftdim=dim(leftspace)
#         rightdim=dim(rightspace)
#         data=reshape(t.data,(leftdim,rightdim))
#         if leftdim>rightdim
#             newdim=rightdim
#             tau=zeros(eltype(data),(newdim,))
#             Base.LinAlg.LAPACK.geqrf!(data,tau)
#
#             C=zeros(eltype(t),(newdim,newdim))
#             for j in 1:newdim
#                 for i in 1:j
#                     @inbounds C[i,j]=data[i,j]
#                 end
#             end
#             Base.LinAlg.LAPACK.orgqr!(data,tau)
#             U=data
#             for i=1:newdim
#                 tau[i]=sign(C[i,i])
#             end
#             scale!(U,tau)
#             scale!(conj!(tau),C)
#         else
#             newdim=leftdim
#             C=data
#             U=eye(eltype(data),newdim)
#         end
#
#         newspace=$S(newdim)
#         return tensor(U,leftspace*newspace'), tensor(C,newspace*rightspace), abs(zero(eltype(t)))
#     end
#
#     @eval function leftorth!(t::$TT,n::Int,trunc::Union(TruncationDimension,TruncationSpace))
#         # Truncate rank corresponding to bipartition into left indices 1:n
#         # and remain right indices, based on singular value decomposition,
#         # thereby destroying the original tensor.
#         # Truncation parameters are given as keyword arguments: trunctol should
#         # always be one of the possible arguments for specifying truncation, but
#         # truncdim can be replaced with different parameters for other types of tensors.
#
#         N=numind(t)
#         spacet=space(t)
#         leftspace=spacet[1:n]
#         rightspace=spacet[n+1:N]
#         leftdim=dim(leftspace)
#         rightdim=dim(rightspace)
#         dim(trunc) >= min(leftdim,rightdim) && return leftorth!(t,n)
#         data=reshape(t.data,(leftdim,rightdim))
#         F=svdfact!(data)
#         sing=F[:S]
#
#         # compute truncation dimension
#         if eps(trunc)!=0
#             sing=F[:S]
#             normsing=vecnorm(sing)
#             truncdim=0
#             while truncdim<length(sing) && vecnorm(sing[(truncdim+1):end])>eps(trunc)*normsing
#                 truncdim+=1
#             end
#             truncdim=min(truncdim,dim(trunc))
#         else
#             truncdim=dim(trunc)
#         end
#         newspace=$S(truncdim)
#
#         # truncate
#         truncerr=vecnorm(sing[(truncdim+1):end])/vecnorm(sing)
#         U=tensor(F[:U][:,1:truncdim],leftspace*newspace')
#         C=tensor(scale!(sing[1:truncdim],F[:Vt][1:truncdim,:]),newspace*rightspace)
#         return U,C,truncerr
#     end
#
#     @eval function leftorth!(t::$TT,n::Int,trunc::TruncationError)
#         # Truncate rank corresponding to bipartition into left indices 1:n
#         # and remain right indices, based on singular value decomposition,
#         # thereby destroying the original tensor.
#         # Truncation parameters are given as keyword arguments: trunctol should
#         # always be one of the possible arguments for specifying truncation, but
#         # truncdim can be replaced with different parameters for other types of tensors.
#
#         N=numind(t)
#         spacet=space(t)
#         leftspace=spacet[1:n]
#         rightspace=spacet[n+1:N]
#         leftdim=dim(leftspace)
#         rightdim=dim(rightspace)
#         data=reshape(t.data,(leftdim,rightdim))
#         F=svdfact!(data)
#
#         # compute truncation dimension
#         sing=F[:S]
#         normsing=vecnorm(sing)
#         truncdim=0
#         while truncdim<length(sing) && vecnorm(sing[(truncdim+1):end])>eps(trunc)*normsing
#             truncdim+=1
#         end
#         newspace=$S(truncdim)
#
#         # truncate
#         truncerr=vecnorm(sing[(truncdim+1):end])/vecnorm(sing)
#         U=tensor(F[:U][:,1:truncdim],leftspace*newspace')
#         C=tensor(scale!(sing[1:truncdim],F[:Vt][1:truncdim,:]),newspace*rightspace)
#
#         return U,C,truncerr
#     end
#
#     @eval function rightorth!(t::$TT,n::Int,::NoTruncation)
#         # Create orthogonal basis U for right indices, and remainder C for left
#         # indices. Decomposition should be unique, such that it always returns the
#         # same result for the same input tensor t. CU = LQ is typically fastest but only
#         # unique after correcting for phases.
#         N=numind(t)
#         spacet=space(t)
#         leftspace=spacet[1:n]
#         rightspace=spacet[n+1:N]
#         leftdim=dim(leftspace)
#         rightdim=dim(rightspace)
#         data=reshape(t.data,(leftdim,rightdim))
#         if leftdim<rightdim
#             newdim=leftdim
#             tau=zeros(eltype(data),(newdim,))
#             datat=transpose(data)
#             Base.LinAlg.LAPACK.geqrf!(datat,tau)
#
#             C=zeros(eltype(t),(newdim,newdim))
#             for j in 1:newdim
#                 for i in 1:j
#                     @inbounds C[j,i]=datat[i,j]
#                 end
#             end
#             Base.LinAlg.LAPACK.orgqr!(datat,tau)
#             Base.transpose!(data,datat)
#             U=data
#
#             for i=1:newdim
#                 tau[i]=sign(C[i,i])
#             end
#             scale!(tau,U)
#             scale!(C,conj!(tau))
#         else
#             newdim=rightdim
#             C=data
#             U=eye(eltype(data),newdim)
#         end
#
#         newspace=$S(newdim)
#         return tensor(C,leftspace ⊗ dual(newspace)), tensor(U,newspace ⊗ rightspace), abs(zero(eltype(t)))
#     end
#
#     @eval function rightorth!(t::$TT,n::Int,trunc::Union(TruncationDimension,TruncationSpace))
#         # Truncate rank corresponding to bipartition into left indices 1:n
#         # and remain right indices, based on singular value decomposition,
#         # thereby destroying the original tensor.
#         # Truncation parameters are given as keyword arguments: trunctol should
#         # always be one of the possible arguments for specifying truncation, but
#         # truncdim can be replaced with different parameters for other types of tensors.
#
#         N=numind(t)
#         spacet=space(t)
#         leftspace=spacet[1:n]
#         rightspace=spacet[n+1:N]
#         leftdim=dim(leftspace)
#         rightdim=dim(rightspace)
#         dim(trunc) >= min(leftdim,rightdim) && return rightorth!(t,n)
#         data=reshape(t.data,(leftdim,rightdim))
#         F=svdfact!(data)
#         sing=F[:S]
#
#         # compute truncation dimension
#         if eps(trunc)!=0
#             sing=F[:S]
#             normsing=vecnorm(sing)
#             truncdim=0
#             while truncdim<length(sing) && vecnorm(sing[(truncdim+1):end])>eps(trunc)*normsing
#                 truncdim+=1
#             end
#             truncdim=min(truncdim,dim(trunc))
#         else
#             truncdim=dim(trunc)
#         end
#         newspace=$S(truncdim)
#
#         # truncate
#         truncerr=vecnorm(sing[(truncdim+1):end])/vecnorm(sing)
#         C=tensor(scale!(F[:U][:,1:truncdim],sing[1:truncdim]),leftspace*newspace')
#         U=tensor(F[:Vt][1:truncdim,:],newspace*rightspace)
#         return C,U,truncerr
#     end
#
#     @eval function rightorth!(t::$TT,n::Int,trunc::TruncationError)
#         # Truncate rank corresponding to bipartition into left indices 1:n
#         # and remain right indices, based on singular value decomposition,
#         # thereby destroying the original tensor.
#         # Truncation parameters are given as keyword arguments: trunctol should
#         # always be one of the possible arguments for specifying truncation, but
#         # truncdim can be replaced with different parameters for other types of tensors.
#
#         N=numind(t)
#         spacet=space(t)
#         leftspace=spacet[1:n]
#         rightspace=spacet[n+1:N]
#         leftdim=dim(leftspace)
#         rightdim=dim(rightspace)
#         data=reshape(t.data,(leftdim,rightdim))
#         F=svdfact!(data)
#
#         # compute truncation dimension
#         sing=F[:S]
#         normsing=vecnorm(sing)
#         truncdim=0
#         while truncdim<length(sing) && vecnorm(sing[(truncdim+1):end])>eps(trunc)*normsing
#             truncdim+=1
#         end
#         newspace=$S(truncdim)
#
#         # truncate
#         truncerr=vecnorm(sing[(truncdim+1):end])/vecnorm(sing)
#         C=tensor(scale!(F[:U][:,1:truncdim],sing[1:truncdim]),leftspace*newspace')
#         U=tensor(F[:Vt][1:truncdim,:],newspace*rightspace)
#         return C,U,truncerr
#     end
# end
#
# # Methods below are only implemented for CartesianMatrix or ComplexMatrix:
# #++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
# typealias ComplexMatrix{T} ComplexTensor{T,2}
# typealias CartesianMatrix{T} CartesianTensor{T,2}
#
# function Base.pinv(t::Union(ComplexMatrix,CartesianMatrix))
#     # Compute pseudo-inverse
#     spacet=space(t)
#     data=copy(t.data)
#     leftdim=dim(spacet[1])
#     rightdim=dim(spacet[2])
#
#     F=svdfact!(data)
#     Sinv=F[:S]
#     for k=1:length(Sinv)
#         if Sinv[k]>eps(Sinv[1])*max(leftdim,rightdim)
#             Sinv[k]=one(Sinv[k])/Sinv[k]
#         end
#     end
#     data=F[:V]*scale(F[:S],F[:U]')
#     return tensor(data,spacet')
# end
#
# function Base.eig(t::Union(ComplexMatrix,CartesianMatrix))
#     # Compute eigenvalue decomposition.
#     spacet=space(t)
#     spacet[1] == spacet[2]' || throw(SpaceError("eigenvalue factorization only exists if left and right index space are dual"))
#     data=copy(t.data)
#
#     F=eigfact!(data)
#
#     Lambda=tensor(diagm(F[:values]),spacet)
#     V=tensor(F[:vectors],spacet)
#     return Lambda, V
# end
#
# function Base.inv(t::Union(ComplexMatrix,CartesianMatrix))
#     # Compute inverse.
#     spacet=space(t)
#     spacet[1] == spacet[2]' || throw(SpaceError("inverse only exists if left and right index space are dual"))
#
#     return tensor(inv(t.data),spacet)
# end


function mandelbrot_set(cplane::Array{ComplexF64,1}, rmin,rmax,imin,imax,width,height,maxiter, epsilon)

    cartesian =  Array{Float64, 3}(undef, 3, height, width)
    bs = 1
    total = width * height

    println("start computing pmap")
    if length(cplane) == width*height
	println("using mmap")
	zset =pmap(c->mandel(c,maxiter, total), distributed=true, batch_size=bs, cplane)
    else
	println("computing complex comprehension")
    	zset =pmap(c->mandel(c,maxiter), distributed=true, batch_size=bs, [Complex(r,i) for r in LinRange(rmin ,rmax , width) for i in LinRange(imin, imax, height)])
    end

    println("color zset")
    for y = 1:height, x = 1:width

        index = y + (((x-1) % width))*height
	point = zset[index]
	if point == maxiter
		point = 0
	else
		point /= convert(Float64, maxiter)

	end

 	#red_color, blue_color, green_color = 1.0, 1.0, 1.0
    	if point < epsilon
		red_color, blue_color, green_color = (1.0-point), (1.0 - point * 2), (1.0 - point * 3)
    	else
            red_color = point * 0.60
            blue_color = point * 0.80
            green_color = point * 0.35
    	end

	cartesian[1,y,x] = red_color
	cartesian[2,y,x] = blue_color
	cartesian[3,y,x] = green_color
    end

    return cartesian
end

# the mandelbrot iteration function
@everywhere count = 0
@everywhere start = time()
@everywhere delta = 0
@everywhere function mandel(c, maxiter::Int64, total::Int64)
    global count += 1
    if (count % convert(Int, total * 0.01) == 0)
	    global delta = time() - start
	    println("processed $(count*1.0/total*100.0) in $(delta) seconds")
    end
    z=0+0im
    for n = 1:maxiter
        if abs(z) > 2
            return n-1
        end
        z = z^2 + c
    end
    return maxiter
end


using RCall
using Distributions
using StatsBase
using Iterators
using Gurobi
using JLD
using Test
using EmpiricalBayes

using KernelDensity
nreps =20
k_max = 12 #in total 25*16 replications

true_dists = [MixtureModel([ Normal(-0.3,.2), Normal(0,0.9)],[0.8, 0.2]), EmpiricalBayes.ash_flattop ]
xs = collect(-3:0.5:3)
minimax_σs = [0.2, 0.5]
ms = [20_000, 100_000]

combs = []

for (true_dist, x, σ, m) in product(true_dists, xs, minimax_σs, ms)
    push!(combs, Dict(:dist=>true_dist, :x=>x, :σ=>σ, :m=>m) )
end

#i = parse(Int64, ARGS[1])
i=1
comb = combs[i]

m = comb[:m]
σ = comb[:σ]
true_dist = comb[:dist]
x = comb[:x]

marginal_grid = collect(range(-7.0, stop=7.0, length=1001));
prior_grid = collect(range(-3, stop=3, length=121));
marginal_h = marginal_grid[2]-marginal_grid[1];

d_true = NormalConvolutionProblem(true_dist, marginal_grid)
ds = MixingNormalConvolutionProblem(Normal, σ, prior_grid, marginal_grid);

target = PosteriorTarget(LFSRNumerator(x))

Xs_train = rand(d_true,Int(m/2))

Xs_test = rand(d_true, Int(m/2))

dry_run = CEB_ci(Xs_train, Xs_test, ds, target)

m_train = length(Xs_train)
m_test = length(Xs_test)

marginal_grid = ds.marginal_grid
marginal_h = ds.marginal_h

M_bd = marginal_h/sqrt(2*pi)

f_const = BinnedMarginalDensity(M_bd, marginal_grid, marginal_h)
# Train
# Estimate the numerator
M_max_num = MinimaxCalibrator(ds, f_const, m_train, target.num;
             C=Inf);

# Estimate the denominator
M_max_denom = MinimaxCalibrator(ds, f_const, m_train, target.denom;
            C=Inf);

f_nb = fit(BinnedMarginalDensityNeighborhood, Xs_train, ds)

C_bias = f_nb.C_bias

f_nb.f_kde

CEB_ci(Xs_train, Xs_test, ds, target,
            M_max_num,
            M_max_denom,
            C_bias;
            C=:auto, conf=0.1)

tmp1 = sinc_kde(Xs_train, marginal_grid)

marginal_grid2 = collect(range(-6.0, stop=6.0, length=1001));
tmp2= sinc_kde(Xs_train, marginal_grid2)

marginal_h
tmp3 = kde(Xs_train, -7:(marginal_h-1e-10):7)

marginal_h2 = marginal_grid2[2]-marginal_grid2[1]
tmp4 = kde(Xs_train, -6:marginal_h2:6)


using MOCNeutronTransport
@testset "AABB_2D" begin
    @testset "$F" for F in [Float32, Float64, BigFloat]
        @testset "Methods" begin
            # getproperty
            aabb = AABB_2D(Point_2D{F}(1, 0), Point_2D{F}(3, 2))
            @test aabb.xmin ≈ 1
            @test aabb.ymin ≈ 0
            @test aabb.xmax ≈ 3
            @test aabb.ymax ≈ 2

            # width, height, area
            @test width(aabb) ≈ 2
            @test height(aabb) ≈ 2
            @test area(aabb) ≈ 4

            # in
            @test Point_2D{F}(2, 1) ∈ aabb
            @test Point_2D{F}(3, 1) ∈ aabb
            @test !(Point_2D{F}(4, 1) ∈ aabb)
            @test !(Point_2D{F}(2, 5) ∈ aabb)

            # intersect
            # Horizontal
            hit, points = LineSegment_2D(Point_2D{F}(-1, 1), Point_2D{F}(4, 1)) ∩ aabb
            @test hit
            @test points[1] ≈ Point_2D{F}(1, 1)
            @test points[2] ≈ Point_2D{F}(3, 1)

            # Horizontal miss
            hit, points = LineSegment_2D(Point_2D{F}(-1, 5), Point_2D{F}(4, 5)) ∩ aabb
            @test !hit

            # Vertical
            hit, points = LineSegment_2D(Point_2D{F}(2, -1), Point_2D{F}(2, 5)) ∩ aabb
            @test hit
            @test points[1] ≈ Point_2D{F}(2, 0)
            @test points[2] ≈ Point_2D{F}(2, 2)

            # Vertical miss
            hit, points = LineSegment_2D(Point_2D{F}(5, -1), Point_2D{F}(5, 5)) ∩ aabb
            @test !hit

            # Angled
            hit, points = LineSegment_2D(Point_2D{F}(0, 0), Point_2D{F}(5, 2)) ∩ aabb
            @test hit
            @test points[1] ≈ Point_2D{F}(1, 0.4)
            @test points[2] ≈ Point_2D{F}(3, 1.2)

            # Angled miss
            hit, points = LineSegment_2D(Point_2D{F}(0, 5), Point_2D{F}(5, 7)) ∩ aabb
            @test !hit

            # union
            r1 = AABB_2D(Point_2D{F}(0, 0), Point_2D{F}(2, 2))
            r2 = AABB_2D(Point_2D{F}(1, 1), Point_2D{F}(4, 4))
            r3 = r1 ∪ r2
            @test r3.xmin ≈ 0
            @test r3.ymin ≈ 0
            @test r3.xmax ≈ 4
            @test r3.ymax ≈ 4
        end
    end
end


mtcars = dataset("datasets", "mtcars")

priors = [
    Prior_Ridge(),
    Prior_Laplace(),
    Prior_Cauchy(),
    Prior_TDist(),
    Prior_Uniform(),
]

model = @fitmodel((MPG ~ HP + WT + Gear), mtcars, LinearRegression())
@test sizeof(model) > 0

for prior in priors
    model = @fitmodel((MPG ~ HP + WT + Gear), mtcars, LinearRegression(), prior)
    @test sizeof(model) > 0
end


using BayesianNonparametrics

(X, Y) = bloobs(randomize = false)

μ0 = vec(mean(X, 1))
κ0 = 5.0
ν0 = 9.0
Σ0 = cov(X)

H = WishartGaussian(μ0, κ0, ν0, Σ0)

model = DPM(H)
initialisation = KMeansInitialisation(k = 10)

modelBuffer = init(X, model, initialisation)
models = train(modelBuffer, DPMHyperparam(), Gibbs(maxiter = 500))

densities = Float64[m.energy for m in models]
activeComponents = Int[sum(m.weights .> 0) for m in models]
assignments = [m.assignments for m in models]

A = reduce(hcat, assignments)
(N, D) = size(X)

PSM = zeros(N, N)
M = size(A, 2)
for i in 1:N
  for j in 1:i
    PSM[i, j] = sum(A[i,:] .== A[j,:]) / M
    PSM[j, i] = sum(A[i,:] .== A[j,:]) / M
  end
end

mink = minimum([length(m.weights) for m in models])
maxk = maximum([length(m.weights) for m in models])

(peassignments, _) = pointestimate(PSM, method = :average, mink = mink, maxk = maxk)


using DoubleFloats
using LinearAlgebra

using SemiLagrangian:
    Advection,
    standardsplit,
    strangsplit,
    triplejumpsplit,
    sizeall,
    AdvectionData,
    getdata,
    advection!,
    UniformMesh,
    gettranslationvar,
    AbstractInterpolation,
    interpolate!,
    Lagrange,
    B_SplineLU,
    B_SplineFFT,
    Hermite
"""

   exact(tf, mesh1, mesh2)

   Julia function to compute exact solution "

```math
\frac{d f}{dt} +  (y \frac{df}{delta1} - x \frac{df}{delta2}) = 0
```

"""
function exact!(f, decall::NTuple{Nsum,T}, t::T) where {T,Nsum}
    for ind in CartesianIndices(f)
        f[ind] = exp(
            -sum([
                sin(2T(pi) * (ind.I[i] + decall[i] * t) / size(f, i) + T(big"0.56") * i^2) for i = 1:Nsum
            ]),
        )
    end
end

function test_translation(
    sz::NTuple{2,Int},
    interp_sp::AbstractInterpolation{T},
    interp_v::AbstractInterpolation{T},
    dt::T;
    tc::Vector{T}=strangsplit(dt),
) where {T}
@time begin
    spmin, spmax, nsp = T(-5), T(5), sz[1]
    vmin, vmax, nv = -T(5.55), T(5), sz[2]

    mesh_sp = UniformMesh(spmin, spmax, nsp)
    mesh_v = UniformMesh(vmin, vmax, nv)
end
println("trace1")
@time begin
    nbdt = Int(round(T(1)/dt))
    adv = Advection(
        (mesh_sp, mesh_v),
        [interp_sp, interp_v],
        dt,
        [([1, 2], 1, 1, true), ([2, 1], 1, 2, true)],
        tab_coef= tc
    )

    tabref = zeros(T, sz)

    v1 = T(big"0.83545655467782872872782870029282982828737872878776717190927267611111")
    v2 = T(big"0.945678101929276765616767176761671771766717828781828998101092981877817176")
end
println("trace2")
@time begin

    exact!(tabref, (v1, v2), T(0))

    pvar = gettranslationvar((v1, v2))

    advdata = AdvectionData(adv, tabref, pvar)

    diffmax = 0
    data = getdata(advdata)
end
println("trace3")
    for ind = 1:nbdt
        data = getdata(advdata)
        while advection!(advdata)
        end
        #        GC.gc()
        exact!(tabref, (v1, v2), T(ind)*dt)
        data = getdata(advdata)
        diff = norm(data .- tabref, Inf)
        @show v1, v2, ind, diff
        diffmax = max(diffmax, diff)
    end
    println(
        "length(tab_coef)=$(length(tc)) test_translation sz=$sz interp=$interp_sp, $interp_v nbdt=$nbdt diffmax=$diffmax",
    )
    return diffmax

end
function test_translation(
    sz::NTuple{N,Int},
    interp_t::Vector{I},
    nbdt::Int,
) where {T,N,I<:AbstractInterpolation{T}}

    dt = T(1)

    vall = ntuple(x -> T(10rand(BigFloat) - 5), N)

    vall = T(1) / 10 .* vall

    @show vall

    tabref = zeros(T, sz)
    exact!(tabref, vall, T(0))

    data = copy(tabref)

    buf = zeros(T, sz)

    # dec1 = fill(v1, sz)
    # dec2 = fill(v2, sz)

    diffmax = 0
    for ind = 1:nbdt
        interpolate!(buf, data, (i, _) -> vall[i], interp_t)
        #        interpolate!(buf, data, (dec1, dec2), interp)
        copyto!(data, buf)
        exact!(tabref, vall, T(ind))
        diff = norm(data .- tabref, Inf)
        diffmax = max(diffmax, diff)
        @show ind, diff
    end
    println("test_translation sz=$sz interp=$interp_t nbdt=$nbdt diffmax=$diffmax")
    return diffmax

end

@testset "test translation" begin
    T = Float64
    dt = T(1)/11
    @time @test test_translation(
        (200, 300),
        Lagrange(5, T),
        Lagrange(5, T),
        dt,
    ) < 1e-3
    @time @test test_translation(
        (200, 300),
        Lagrange(5, T),
        Lagrange(5, T),
        dt, tc=triplejumpsplit(dt)
    ) < 1e-3
    @time @test test_translation(
        (200, 300),
        Lagrange(5, T),
        Lagrange(5, T),
        dt, tc=standardsplit(dt)
    ) < 1e-3
    @time @test test_translation(
        (200, 300),
        Hermite(5, T),
        Hermite(5, T),
        dt,
    ) < 1e-3
    @time @test test_translation(
        (200, 300),
        Hermite(5, T),
        Hermite(5, T),
        dt, tc=triplejumpsplit(dt)
    ) < 1e-3
    @time @test test_translation(
        (200, 300),
        Hermite(5, T),
        Hermite(5, T),
        dt,tc=standardsplit(dt)
    ) < 1e-3
    @time @test test_translation(
        (128, 64),
        B_SplineLU(5, 128, T),
        B_SplineLU(5, 64, T),
        dt,
    ) < 1e-6
    @time @test test_translation(
        (128, 64),
        B_SplineLU(5, 128, T),
        B_SplineLU(5, 64, T),
        dt, tc=triplejumpsplit(dt)
    ) < 1e-6
    @time @test test_translation(
        (128, 64),
        B_SplineFFT(5, 128, T),
        B_SplineFFT(5, 64, T),
        dt,
    ) < 1e-6
    T = Double64
    dt = T(1)/11
    @time @test test_translation((100, 120), Lagrange(15, T), Lagrange(15, T), dt) < 1e-14
    @time @test test_translation(
        (128, 64),
        B_SplineLU(15, 128, T),
        B_SplineLU(15, 64, T),
        dt,
    ) < 1e-18
    @time @test test_translation(
        (128, 64),
        B_SplineFFT(15, 128, T),
        B_SplineFFT(15, 64, T),
        dt,
    ) < 1e-17

    T = BigFloat
    dt = T(1)/11
@time @test test_translation((100, 90), Lagrange(25, T), Lagrange(25, T), dt) < 1e-20
@time @test test_translation(
        (128, 64),
        B_SplineLU(25, 128, T),
        B_SplineLU(25, 64, T),
        dt,
    ) < 1e-22
   @time @test test_translation(
        (128, 64),
        B_SplineFFT(25, 128, T),
        B_SplineFFT(25, 64, T),
        dt,
    ) < 1e-22

end