export
    WignerFunction,
    WignerSurface

#=
    Wigner function by Laguerre Polynominal
=#

abstract type CreateWignerMethod end
struct Load𝐖 <: CreateWignerMethod end
struct Calc𝐖 <: CreateWignerMethod end

function wigner(m::Integer, n::Integer, x::Vector{<:Real}, p::Vector{<:Real})
    w = gaussian_function(x, p) .*
        coefficient_of_wave_function(m, n) .*
        z_to_power(m, n, x, p) .*
        laguerre(m, n, x, p)

    return w
end

function create_wigner(
    m_dim::Integer,
    n_dim::Integer,
    x_range::AbstractRange,
    p_range::AbstractRange,
    ::Type{Calc𝐖}
)
    𝐰 = Array{ComplexF64,4}(undef, m_dim, n_dim, length(x_range), length(p_range))
    @sync for m in 1:m_dim
        Threads.@spawn for n in 1:n_dim
            𝐰[m, n, :, :] = wigner(m, n, collect(x_range), collect(p_range))
        end
    end

    return 𝐰
end

function create_wigner(
    m_dim::Integer,
    n_dim::Integer,
    x_range::AbstractRange,
    p_range::AbstractRange,
    bin_path::String,
    ::Type{Load𝐖}
)
    return load_𝐰(m_dim, n_dim, x_range, p_range, bin_path)
end

function create_wigner(
    m_dim::Integer,
    n_dim::Integer,
    x_range::AbstractRange,
    p_range::AbstractRange,
)
    bin_path = gen_wigner_bin_path(m_dim, n_dim, x_range, p_range)

    if isfile(bin_path)
        return create_wigner(m_dim, n_dim, x_range, p_range, bin_path, Load𝐖)
    end

    𝐰 = create_wigner(m_dim, n_dim, x_range, p_range, Calc𝐖)
    save_𝐰(bin_path, 𝐰)

    return 𝐰
end

mutable struct WignerFunction{T<:Integer, U<:AbstractRange}
    m_dim::T
    n_dim::T
    x_range::U
    p_range::U
    𝐰::Array{ComplexF64,4}

    function WignerFunction(
        m_dim::T,
        n_dim::T,
        x_range::U,
        p_range::U
    ) where {T<:Integer, U<:AbstractRange}
        !check_argv(m_dim, n_dim, x_range, p_range) && throw(ArgumentError)

        𝐰 = create_wigner(m_dim, n_dim, x_range, p_range)
        return new{T, U}(m_dim, n_dim, x_range, p_range, 𝐰)
    end
end

"""
    WignerFunction(x_range::AbstractRange, p_range::AbstractRange; dim=DIM)

Declare a generalized Wigner function with `x` range and `p` range

* `x_range`: range of position.
* `p_range`: range of momentum.
* `dim`: Maximum photon number for truncate, default is $DIM.

# Examples
```julia-repl
julia> state = VacuumState();

julia> wf = WignerFunction(-10:0.1:10, -10:0.1:10);

julia> w = wf(state);
```
"""
function WignerFunction(x_range::AbstractRange, p_range::AbstractRange; dim=DIM)
    return WignerFunction(dim, dim, x_range, p_range)
end

struct WignerSurface{T<:AbstractRange}
    x_range::T
    p_range::T
    𝐰_surface::Matrix{Float64}

    function WignerSurface(
        x_range::T,
        p_range::T,
        𝐰_surface::Matrix{Float64}
    ) where {T<:AbstractRange}
        return new{T}(x_range, p_range, 𝐰_surface)
    end
end

function (wf::WignerFunction)(state::StateMatrix)
    𝛒 = state.𝛒

    𝐰_surface = Matrix{Float64}(undef, length(wf.x_range), length(wf.p_range))
    @sync for i in 1:length(wf.x_range)
        Threads.@spawn for j in 1:length(wf.p_range)
            𝐰_surface[i, j] = real(sum(𝛒 .* wf.𝐰[:, :, i, j]))
        end
    end

    return WignerSurface(wf.x_range, wf.p_range, 𝐰_surface)
end

(wf::WignerFunction)(state::StateVector) = wf(StateMatrix(state))


using Test

ϕ=60
a1 = [1,0].*2
a2 = [cos(π*ϕ/180),sin(π*ϕ/180)].*2
l = Generic2DLattice(a1,a2)
conf = CorrugatedMorse(5,0.4,0.1,l)

x = collect(0:0.1:10)
y = collect(0:0.1:10)
v = Array{Float64,2}(undef, length(x),length(y))

for (i1,xi) in enumerate(x), (i2,yi) in enumerate(y)
    z = 0
    v[i1,i2] = potential(conf,xi,yi,z)
end

using Plots
gr()
#plot(x,y,v,seriestype=:heatmap,aspect_ratio=:equal)
#gui()

start = 0
V0 = []
for s in 5:20:100
push!(V0, real(calc_VG([0,0],l,conf,0.2,s)))
end
plot(V0,seriestype=:line)
gui()

#start = 0
#z = 0.3
#V0 = fill(0.,7,7)
#for gx in -3:3, gy in -3:3
#    V0[gx+4,gy+4] = real(calc_VG([gx,gy],l,conf,z,10))
#end
#@show V0
#plot(V0,seriestype=:heatmap)
#gui()


@testset "Kernels" begin
    import LikelihoodfreeInference: bandwidth, update!, Bandwidth, RosenblattSchedule, Scheduled
    k = Kernel(type = Gaussian,
               bandwidth = Bandwidth(heuristic = nothing, gamma = 2.))
    @test bandwidth(k) == 2.
    update!(k, 0)
    @test bandwidth(k) == 2.
    @test k(1, 4) == exp(-(4 - 1)^2/(2 * 2^2))
    @test k([2, 0, 1], zeros(3)) ≈ exp(-5/(2*2^2))
    k = Kernel(type = Gaussian,
               bandwidth = Bandwidth(heuristic = nothing, gamma = [2., 6.]))
    @test bandwidth(k) == [2., 6.]
    @test k([3, 0], [4, 2]) == exp(-sum(abs2, ([3, 0] - [4, 2]) ./ [2., 6])/2)
    k = Kernel(type = ModifiedGaussian,
               bandwidth = Bandwidth(heuristic = nothing, gamma = .5))
    @test k(1, 0) == exp(-1/(2*.5))
    @test k(.3, 0) == exp(-.3^2/(2*.5^2))
    k = Kernel(bandwidth = Bandwidth(heuristic = MedianHeuristic()))
    update!(k, [2, 2])
    @test bandwidth(k) == 2
    k = Kernel(bandwidth = Bandwidth(heuristic = MedianHeuristic(.1)))
    update!(k, [2, 2])
    @test bandwidth(k) == .1 * 2
    k = Kernel(bandwidth = Bandwidth(heuristic = ScottsHeuristic()))
    update!(k, [2, 2])
    @test bandwidth(k) == 0
    k = Kernel(bandwidth = Scheduled(Bandwidth(heuristic = nothing, gamma = 2),
                                     RosenblattSchedule(2)))
    update!(k, nothing)
    update!(k, nothing)
    @test bandwidth(k) == 2 * 2^(-1/(2 + 4))
end


using StanSample, DiffEqBayes, OrdinaryDiffEq, ParameterizedFunctions,
      ModelingToolkit, RecursiveArrayTools, Distributions, Random, Test

# Uncomment for local testing only, make sure MCMCChains and StatsPlots are available
using MCMCChains, StatsPlots

#Random.seed!(123)
cd(@__DIR__)
isdir("tmp") && rm("tmp", recursive=true)

include("../../stan_string.jl")
include("../../stan_inference.jl")

println("\nFour parameter case\n")
f1 = @ode_def begin
  dx = a*x - b*x*y
  dy = -c*y + d*x*y
end a b c d
u0 = [1.0,1.0]
tspan = (0.0,10.0)
p = [1.5,1.0,3.0,1.0]
prob1 = ODEProblem(f1,u0,tspan,p)
sol = solve(prob1,Tsit5())
t = collect(range(1,stop=10,length=10))
randomized = VectorOfArray([(sol(t[i]) + .5randn(2)) for i in 1:length(t)])
data = convert(Array,randomized)
priors = [truncated(Normal(1.0,1),0.1,2),truncated(Normal(1.5,0.5),0.1,1.5),
          truncated(Normal(2.0,1),0.1,4),truncated(Normal(1.3,0.5),0.1,2)]

@time bayesian_result = stan_inference(prob1,t,data,priors;num_samples=2000, nchains=4,
  num_warmup=1000,vars =(StanODEData(),InverseGamma(4,1)))

if success(bayesian_result.return_code)
  p  = read_samples(bayesian_result.model, output_format=:particles)
  @test mean(p.theta1) ≈ 1.5 atol=5e-1
  @test mean(p.theta2) ≈ 1.0 atol=5e-1
  @test mean(p.theta3) ≈ 3.0 atol=5e-1
  @test mean(p.theta4) ≈ 1.0 atol=5e-1

  # Uncomment for local chain inspection
  p1 = Vector{Plots.Plot{Plots.GRBackend}}(undef, length(priors))
  p2 = Vector{Plots.Plot{Plots.GRBackend}}(undef, length(priors))
  for i in 1:length(priors)
    p1[i] = plot(title="theta.$i", leg=false)
    plot!(p1[i], p[Symbol("theta.$i")].particles)
    p2[i] = plot(title="theta.$i", leg=false)
    density!(p2[i], p[Symbol("theta.$i")].particles)
  end
  plot(p1[1], p2[1], p1[2], p2[2], p1[3], p2[3], p1[4], p2[4], layout=(length(priors),2))
  savefig("$(@__DIR__)/four_par_case.png")

  v1 = [mean(p[Symbol("u_hat.$i.1")]) for i in 1:10]
  v2 = [mean(p[Symbol("u_hat.$i.2")]) for i in 1:10]
  qa_prey = zeros(10, 3);
  qa_pred = zeros(10, 3);

  df = read_samples(bayesian_result.model, output_format=:dataframe)
  achn = Array(df);
  for i in 1:10
    qa_prey[i, :] = quantile(achn[:, 10+i], [0.055, 0.5, 0.945])
    qa_pred[i, :] = quantile(achn[:, 20+i], [0.055, 0.5, 0.945])
  end
  p1 = plot(1:10, v1; ribbon=(qa_prey[:, 1], qa_prey[:, 3]), color=:lightgrey, leg=false)
  title!("Prey u_hat 89% quantiles")
  plot!(v1, lab="prey", xlab="time", ylab="prey", color=:darkred)

  p2 = plot(1:10, v2; ribbon=(qa_pred[:, 1], qa_pred[:, 3]), color=:lightgrey, leg=false)
  title!("Preditors u_hat 89% quantiles")
  plot!(v2, lab="pred", xlab="time", ylab="preditors", color=:darkblue)

  plot(p1, p2, layout=(2,1))
  savefig("$(@__DIR__)/four_par_pred_prey.png")
end


export SpinOneSite,
       spinOneSites

struct SpinOneSite <: Site
    s::Index
    SpinOneSite(I::Index) = new(I)
end
SpinOneSite(n::Int) = SpinOneSite(Index(3, "Site,S=1,n=$n"))

function spinOneSites(N::Int;kwargs...)::SiteSet
  sites = SiteSet(N)
  for n=1:N
    set(sites,n,SpinOneSite(n))
  end
  return sites
end

function operator(site::SpinOneSite,
                  opname::AbstractString)::ITensor
    s = site.s
    sP = prime(site.s)
    Up = s(1)
    UpP = sP(1)
    Z0 = s(2)
    Z0P = sP(2)
    Dn = s(3)
    DnP = sP(3)

    Op = ITensor(dag(s), s')

    if opname == "S⁺" || opname == "Splus" || opname == "S+"
        Op[Dn, Z0P] = √2
        Op[Z0, UpP] = √2
    elseif opname == "S⁻" || opname == "Sminus" || opname == "S-"
        Op[Up, Z0P] = √2
        Op[Z0, DnP] = √2
    elseif opname == "Sˣ" || opname == "Sx"
        Op[Up, Z0P] = 1.0/√2
        Op[Z0, UpP] = 1.0/√2
        Op[Z0, DnP] = 1.0/√2
        Op[Dn, Z0P] = 1.0/√2
    elseif opname == "iSʸ" || opname == "iSy"
        Op[Up, Z0P] = -1.0/√2
        Op[Z0, UpP] = +1.0/√2
        Op[Z0, DnP] = -1.0/√2
        Op[Dn, Z0P] = +1.0/√2
    elseif opname == "Sʸ" || opname == "Sy"
        Op = complex(Op)
        Op[Up, Z0P] = +1.0/√2im
        Op[Z0, UpP] = -1.0/√2im
        Op[Z0, DnP] = +1.0/√2im
        Op[Dn, Z0P] = -1.0/√2im
    elseif opname == "Sᶻ" || opname == "Sz"
        Op[Up, UpP] = 1.0
        Op[Dn, DnP] = -1.0
    elseif opname == "Sᶻ²" || opname == "Sz2"
        Op[Up, UpP] = 1.0
        Op[Dn, DnP] = 1.0
    elseif opname == "Sˣ²" || opname == "Sx2"
        Op[Up, UpP] = 0.5
        Op[Up, DnP] = 0.5
        Op[Z0, Z0P] = 1.0
        Op[Dn, UpP] = 0.5
        Op[Dn, DnP] = 0.5
    elseif opname == "Sʸ²" || opname == "Sy2"
        Op[Up, UpP] = 0.5
        Op[Up, DnP] = -0.5
        Op[Z0, Z0P] = 1.0
        Op[Dn, UpP] = -0.5
        Op[Dn, DnP] = 0.5
    elseif opname == "projUp"
        Op[Up, UpP] = 1.
    elseif opname == "projZ0"
        Op[Z0, Z0P] = 1.
    elseif opname == "projDn"
        Op[Dn, DnP] = 1.
    elseif opname == "XUp"
        xup = ITensor(ComplexF64,s)
        xup[Up] = 0.5
        xup[Z0] = im*√2
        xup[Dn] = 0.5
        return xup
    elseif opname == "XZ0"
        xZ0 = ITensor(ComplexF64,s)
        xZ0[Up] = im*√2
        xZ0[Dn] = -im*√2
        return xZ0
    elseif opname == "XDn"
        xdn = ITensor(ComplexF64,s)
        xdn[Up] = 0.5
        xdn[Z0] = -im*√2
        xdn[Dn] = 0.5
        return xdn
    else
        throw(ArgumentError("Operator name '$opname' not recognized for SpinOneSite"))
    end
    return Op
end


@testset "Spaces" begin
  k = FlintQQ

  Qx, x = PolynomialRing(FlintQQ, "x")
  K1, a1 = NumberField(x^2 - 2, "a1")
  K2, a2 = NumberField(x^3 - 2, "a2")

  K1t, t = PolynomialRing(K1, "t")

  L, b = NumberField(t^2 + a1)

  for K in [k, K1, K2, L]
    V = @inferred quadratic_space(K, 2)
    @test_throws ArgumentError quadratic_space(K, -1)

    V = @inferred quadratic_space(K, identity_matrix(FlintZZ, 2))
    @test (@inferred gram_matrix(V)) == identity_matrix(K, 2)
    @test (@inferred dim(V)) == 2
    @test (@inferred rank(V)) == 2
    @test @inferred isregular(V)
    @test (@inferred involution(V)) === identity

    V = @inferred quadratic_space(K, zero_matrix(K, 2, 2))
    @test (@inferred gram_matrix(V)) == zero_matrix(K, 2, 2)
    @test (@inferred gram_matrix(V)) isa dense_matrix_type(K)
    @test (@inferred dim(V)) == 2
    @test (@inferred rank(V)) == 0
    @test @inferred isquadratic(V)
    @test @inferred !ishermitian(V)
    @test (@inferred fixed_field(V)) === K

    @test_throws ArgumentError quadratic_space(K, zero_matrix(K, 1, 2))

    M = identity_matrix(K, 2)
    M[1, 2] = 2
    M[2, 1] = 1
    @test_throws ArgumentError quadratic_space(K, M)
  end

end


"""
Transformation associated with a Generalized Linear Model (GLM). Maps an input
`x` to a random variable with expected value equal to `g(x)` where g(⋅) is a
inverse link function.
"""
immutable GeneralizedLinearTransformation{ST<:StochasticTransformation}
    noise::ST
    invlink::Transformation
end

# constructors
function GeneralizedLinearTransformation{D<:Distribution}(d::D)
    tfm = StochasticTransformation(d)
    invlink = cannonical_inv_link(d)
    GeneralizedLinearTransformation(tfm,invlink)
end
function GeneralizedLinearTransformation{D<:Distribution, T<:Transformation}(d::D, invlink::T)
    tfm = StochasticTransformation(d)
    GeneralizedLinearTransformation(tfm,invlink)
end

# functions
function transform(tfm::GeneralizedLinearTransformation, x)
    transform(tfm.noise, transform(tfm.invlink, x))
end
function rand(tfm::GeneralizedLinearTransformation, x)
    rand(tfm.noise, transform(tfm.invlink, x))
end

# cannonical inverse link functions
cannonical_inv_link(::Normal) = IdentityTransformation()
cannonical_inv_link(::Bernoulli) = LogisticTransformation()
cannonical_inv_link(::Poisson) = ExpTransformation()
cannonical_inv_link(::NegativeBinomial) = ExpTransformation()


# Test dqds algorithms and NDN/LDL factorizations

using Base.Test
using mr3

# @test - macro to test equality
# @test_approx_eq - approximate equality

# test that LDL factorization is accurate
T = SymTridiagonal(randn(5),randn(4))
F = LDL(T)
L, D = mr3.fullLDL(F)
@test_approx_eq (L*D*L')  (full(T))

# test dstqds function
μ = 0.5
G = mr3.dstqds(F,μ)
Lg, Dg = mr3.fullLDL(G)
@test_approx_eq (L*D*L' - μ * I) (Lg*Dg*Lg')

# test dqds function
μ = 0.5
G = mr3.dqds(F,μ)
Dg, Ug = mr3.fullUDU(G)
Lf, Df = mr3.fullLDL(F)
@test_approx_eq (Lf*Df*Lf' - μ*I)  (Ug*Dg*Ug')


#=
two stage stochastic farming problem

sets:
i ∈ 1..n  crops

parameters:
m  ∈ ℝ₊  number of available plots
aᵢ ∈ ℝ₊  fixed cost per farm plot for i
bᵢ ∈ ℝ₊  purchase cost for i
cᵢ ∈ ℝ₊  selling price for i
dᵢ ∈ ℝ₊  contractual demand for i

uncertain parameters:
ξᵢ ∈ ℝ₊  crop yield per farm plot for i
where ξ ~ Ξ

variables:
xᵢ ∈ ℤ₊  number of farm plots for i (1st stage)
yᵢ ∈ P₊  amount of crop i to purchase (2nd stage)
zᵢ ∈ P₊  amount of crop i to sell (2nd stage)
where P₊ are nonnegative polynomials in ξ

objective:
min ∑ᵢ aᵢ xᵢ + 𝔼(ξ ~ Ξ)[∑ᵢ(bᵢ yᵢ - cᵢ zᵢ)]  minimize total expected cost

constraints:
∑ᵢ xᵢ ≤ m  available plots
xᵢ ξᵢ + yᵢ - zᵢ - dᵢ = 0, ∀i  balance inputs and outputs for i

optional linking constraints:
f - ∑ᵢ (yᵢ + zᵢ) ∈ P₊  capacity limit for all crops (f ∈ ℝ₊ is a parameter)

TODO
- CVaR objective or other
- if unlinked, variables can be partitioned when x is fixed, so solve separate subproblems
=#

struct TwoStageStochastic <: ExampleInstance
    n::Int
    halfdeg::Int
    linked::Bool # whether stage two decisions are linked or independent
    use_nat::Bool # use WSOS cone formulation, else SDP formulation
end

function build(inst::TwoStageStochastic)
    n = inst.n
    m = 2n

    # generate parameter values
    a = rand(n)
    b = 1 .+ rand(n)
    c = b - rand(n)
    d = 1 .+ rand(n)
    f = sum(d) + 0.2 * n * rand()

    # setup polynomial interpolation
    n0 = (inst.linked ? n : 1)
    dom = PolyUtils.BoxDomain{Float64}(zeros(n0), ones(n0))
    (U, pts, Ps, _, w) = PolyUtils.interpolate(dom, inst.halfdeg, get_quadr = true)
    w .*= (inst.linked ? 1 : 2^(n - 1))

    # build model
    model = JuMP.Model()
    JuMP.@variable(model, x[1:n], Int)
    JuMP.@variable(model, y[1:n, 1:U])
    JuMP.@variable(model, z[1:n, 1:U])

    JuMP.@objective(
        model,
        Min,
        JuMP.dot(a, x) + JuMP.dot(w, sum(b[i] * y[i, :] - c[i] * z[i, :] for i in 1:n))
    )

    JuMP.@constraint(model, x .>= 0)
    JuMP.@constraint(model, sum(x) <= m)

    add_wsos(aff) = (inst.use_nat ? add_wsos_nat : add_wsos_ext)(Ps, aff, model)

    for i in 1:n
        add_wsos(y[i, :])
        add_wsos(z[i, :])
        pts_i = (inst.linked ? pts[:, i] : vec(pts))
        JuMP.@constraint(model, pts_i * x[i] + y[i, :] - z[i, :] .- d[i] .== 0)
    end

    if inst.linked
        fyz = JuMP.@expression(model, f .- sum(y[i, :] + z[i, :] for i in 1:n))
        add_wsos(fyz)
    end

    # save for use in tests
    model.ext[:x] = x
    model.ext[:y] = y
    model.ext[:z] = z
    if inst.linked
        model.ext[:fyz] = fyz
    end
    model.ext[:Ps] = Ps

    return model
end

function test_extra(inst::TwoStageStochastic, model::JuMP.Model)
    stat = JuMP.termination_status(model)
    @test stat in OPT_OR_LIMIT
    (stat == MOI.OPTIMAL) || return

    tol = eps()^0.2
    x = JuMP.value.(model.ext[:x])
    @test x ≈ round.(Int, x) atol = tol rtol = tol
    # check WSOS feasibility
    Ps = model.ext[:Ps]
    y = JuMP.value.(model.ext[:y])
    z = JuMP.value.(model.ext[:z])
    @test all(check_nonneg(y[i, :] .+ tol, Ps) for i in 1:(inst.n))
    @test all(check_nonneg(z[i, :] .+ tol, Ps) for i in 1:(inst.n))
    if inst.linked
        fyz = JuMP.value.(model.ext[:fyz])
        @test check_nonneg(fyz .+ tol, Ps)
    end
    return
end


vec_simple = @stochastic_model begin
    @stage 1 begin
        @parameters begin
            lb = [40.0, 20.0]
            c  = [100., 150.]
            A  = [1. 1.]
            b  = [120.]
        end
        @decision(model, x[i in 1:2] >= lb[i])
        @objective(model, Min, dot(c, x))
        @constraint(model, A * x - b in MOI.Nonpositives(1))
    end
    @stage 2 begin
        @parameters begin
            lb = [0., 0.]
            T  = [-60. 0.;
                  0.  -80.]
            W  = [6. 10.;
                  8. 5.]
        end
        @uncertain q₁ q₂ d₁ d₂
        ub = [d₁, d₂]
        q  = [q₁, q₂]
        @variable(model, lb[i] <= y[i in 1:2] <= ub[i])
        @objective(model, Max, dot(q, y))
        @constraint(model, T * x + W * y in MOI.Nonpositives(2))
    end
end

ξ₁ = Scenario(q₁ = 24.0, q₂ = 28.0, d₁ = 500.0, d₂ = 100.0, probability = 0.4)
ξ₂ = Scenario(q₁ = 28.0, q₂ = 32.0, d₁ = 300.0, d₂ = 300.0, probability = 0.6)

vec_simple_res = SPResult([46.67, 36.25], -855.83, -1518.75, 662.92, 286.92, -1445.92, -568.92)
push!(problems, (vec_simple, [ξ₁,ξ₂], vec_simple_res, "Simple"))


using Arblib

function sin_naive!(res, x)
    s, t, u = zero(x), zero(x), zero(x)
    tol = one(x)
    Arblib.mul_2exp!(tol, tol, -precision(tol))
    # @show tol
    k = 0
    while true
        Arblib.pow!(t, x, 2k + 1)
        Arblib.fac!(u, UInt(2k + 1))
        Arblib.div!(t, t, u)
        Arblib.abs!(u, t)

        if u ≤ tol
            # @show u
            Arblib.add_error!(s, u)
            break
        end
        if iseven(k)
            Arblib.add!(s, s, t)
        else
            Arblib.sub!(s, s, t)
        end
        k += 1
    end
    Arblib.set!(res, s)
end

global prec = 64
while true
    x = Arb("2016.1"; prec = prec)
    y = zero(x)
    y = sin_naive!(y, x)
    print("Using $(lpad(prec, 5)) bits, sin(x) = ")
    println(Arblib.string_nice(y, 10))
    y < zero(y) && break
    global prec *= 2
end


export IntrinsicSubspaceHomotopy, set_subspaces!

struct LinearSubspaceGeodesicInfo
    Q::Matrix{ComplexF64}
    Q_cos::Matrix{ComplexF64}
    Θ::Vector{Float64}
    U::Matrix{ComplexF64}
    γ1::Matrix{ComplexF64}
end

function LinearSubspaceGeodesicInfo(start::LinearSubspace, target::LinearSubspace)
    Q, Θ, U = geodesic_svd(target, start)
    Q_cos = target.intrinsic.Y * U
    γ1 = similar(Q_cos)
    n, k = size(γ1)
    for j = 1:k
        s, c = sincos(Θ[j])
        for i = 1:n
            γ1[i, j] = Q_cos[i, j] * c + Q[i, j] * s
        end
    end
    LinearSubspaceGeodesicInfo(Q, Q_cos, Θ, U, γ1)
end

const GEODESIC_LRU = LRU{
    Tuple{LinearSubspace{ComplexF64},LinearSubspace{ComplexF64}},
    LinearSubspaceGeodesicInfo,
}(
    maxsize = 128,
)

"""
    IntrinsicSubspaceHomotopy(F::System, V::LinearSubspace, W::LinearSubspace)
    IntrinsicSubspaceHomotopy(F::AbstractSystem, V::LinearSubspace, W::LinearSubspace)

Creates a homotopy ``H(x,t) = F(γ(t)x)`` where ``γ(t)`` is a family of affine subspaces
such that ``γ(1) = V`` and ``γ(0) = W``.
Here ``γ(t)`` is the geodesic between `V` and `W` in the affine Grassmanian, i.e.,
it is the curve of minimal length connecting `V` and `W`.
See also [`LinearSubspace`](@ref) and [`geodesic`](@ref) and the references therein.
"""
Base.@kwdef mutable struct IntrinsicSubspaceHomotopy{S<:AbstractSystem} <: AbstractHomotopy
    system::S

    start::LinearSubspace{ComplexF64}
    target::LinearSubspace{ComplexF64}
    geodesic::LinearSubspaceGeodesicInfo

    J::Matrix{ComplexF64}
    x::Vector{ComplexF64}
    ẋ::Vector{ComplexF64}
    x_high::Vector{ComplexDF64}
    t_cache::Base.RefValue{ComplexF64}
    # For AD
    taylor_t_cache::Base.RefValue{ComplexF64}
    taylor_γ::NTuple{4,Matrix{ComplexF64}}
    v::Vector{ComplexF64}
    tx³::TaylorVector{4,ComplexF64}
    tx²::TaylorVector{3,ComplexF64}
    tx¹::TaylorVector{2,ComplexF64}
end

## Implementation details

# The computation is performed using (implicit) Stiefel coordinates. If `size(F) = m, n`
# and `V` and `W` are affine subspaces of dimension `k, then the computation is performed in
# the affine Grassmanian Graff(k,n) embedded in the Grassmanian(k+1,n+1) using the system
# ``[F(γ(t)v); (γ(t)v)[end] - 1]``. Here the `(γ(t)v)[k+1] - 1` condition ensures that
# we are in the affine Grassmanian and that `(γ(t)v)[1:k]` is the correct value in \C^n.

IntrinsicSubspaceHomotopy(
    F::ModelKit.System,
    start::LinearSubspace,
    target::LinearSubspace;
    compile::Union{Bool,Symbol} = COMPILE_DEFAULT[],
) = IntrinsicSubspaceHomotopy(fixed(F; compile = compile), start, target)

function IntrinsicSubspaceHomotopy(
    system::AbstractSystem,
    start::LinearSubspace,
    target::LinearSubspace;
    compile::Union{Bool,Symbol} = COMPILE_DEFAULT[],
)
    IntrinsicSubspaceHomotopy(
        system,
        convert(LinearSubspace{ComplexF64}, start),
        convert(LinearSubspace{ComplexF64}, target),
    )
end

function IntrinsicSubspaceHomotopy(
    system::AbstractSystem,
    start::LinearSubspace{ComplexF64},
    target::LinearSubspace{ComplexF64},
)
    geodesic = get!(GEODESIC_LRU, (start, target)) do
        LinearSubspaceGeodesicInfo(start, target)
    end
    Q = geodesic.Q
    tx³ = TaylorVector{4}(ComplexF64, size(Q, 1))
    IntrinsicSubspaceHomotopy(
        system = system,
        start = start,
        target = target,
        geodesic = LinearSubspaceGeodesicInfo(start, target),
        J = zeros(ComplexF64, size(system) .+ (1, 1)),
        x = zeros(ComplexF64, size(Q, 1)),
        ẋ = zeros(ComplexF64, size(Q, 1)),
        x_high = zeros(ComplexDF64, size(Q, 1)),
        t_cache = Ref(complex(NaN, NaN)),
        taylor_t_cache = Ref(complex(NaN, NaN)),
        taylor_γ = tuple((similar(Q) for i = 0:3)...),
        v = zeros(ComplexF64, size(Q, 2)),
        tx³ = tx³,
        tx² = TaylorVector{3}(tx³),
        tx¹ = TaylorVector{2}(tx³),
    )
end
Base.size(H::IntrinsicSubspaceHomotopy) = (size(H.system)[1] + 1, dim(H.start) + 1)

"""
    set_subspaces!(H::IntrinsicSubspaceHomotopy, start::LinearSubspace, target::LinearSubspace)

Update the homotopy `H` to track from the affine subspace `start` to `target`.
"""
function set_subspaces!(
    H::IntrinsicSubspaceHomotopy,
    start::LinearSubspace,
    target::LinearSubspace,
)
    H.start = start
    H.target = target
    H.geodesic = get!(GEODESIC_LRU, (start, target)) do
        LinearSubspaceGeodesicInfo(start, target)
    end
    H
end
start_parameters!(H::IntrinsicSubspaceHomotopy, p::LinearSubspace) =
    set_subspaces!(H, p, H.target)
target_parameters!(H::IntrinsicSubspaceHomotopy, q::LinearSubspace) =
    set_subspaces!(H, H.starts, q)
parameters!(H::IntrinsicSubspaceHomotopy, p::LinearSubspace, q::LinearSubspace) =
    set_subspaces!(H, p, q)

function γ!(H::IntrinsicSubspaceHomotopy, t::Number)
    H.t_cache[] != t || return first(H.taylor_γ)
    if isreal(t)
        _γ!(H, real(t))
    else
        _γ!(H, t)
    end
    H.t_cache[] = t

    first(H.taylor_γ)
end
@inline function _γ!(H::IntrinsicSubspaceHomotopy, t::Number)
    @unpack Q, Q_cos, Θ = H.geodesic
    γ = first(H.taylor_γ)
    n, k = size(γ)
    @inbounds for j = 1:k
        Θⱼ = Θ[j]
        s, c = sincos(t * Θⱼ)
        for i = 1:n
            γ[i, j] = Q_cos[i, j] * c + Q[i, j] * s
        end
    end
    γ
end

γ̇!(H::IntrinsicSubspaceHomotopy, t::Number) = isreal(t) ? _γ̇!(H, real(t)) : _γ̇!(H, t)
@inline function _γ̇!(H::IntrinsicSubspaceHomotopy, t::Number)
    @unpack Q, Q_cos, Θ = H.geodesic
    _, γ̇ = H.taylor_γ
    n, k = size(γ̇)
    @inbounds for j = 1:k
        Θⱼ = Θ[j]
        s, c = sincos(t * Θⱼ)
        ċ = -s * Θⱼ
        ṡ = c * Θⱼ
        for i = 1:n
            γ̇[i, j] = Q_cos[i, j] * ċ + Q[i, j] * ṡ
        end
    end
    γ̇
end

function set_solution!(u::Vector, H::IntrinsicSubspaceHomotopy, x::AbstractVector, t)
    (length(x) == length(H.x) - 1) ||
        throw(ArgumentError("Cannot set solution. Expected extrinsic coordinates."))

    set_solution!(view(H.x, 1:length(x)), H.system, x)
    H.x[end] = 1

    if isone(t)
        LA.mul!(u, H.geodesic.γ1', H.x)
    elseif iszero(t)
        LA.mul!(u, H.geodesic.Q_cos', H.x)
    else
        LA.mul!(u, γ!(H, t)', H.x)
    end
end

function get_solution(H::IntrinsicSubspaceHomotopy, u::AbstractVector, t)
    if isone(t)
        (@view H.geodesic.γ1[1:end-1, :]) * u
    elseif iszero(t)
        (@view H.geodesic.Q_cos[1:end-1, :]) * u
    else
        γ = γ!(H, t)
        (@view γ[1:end-1, :]) * u
    end
end

function ModelKit.evaluate!(u, H::IntrinsicSubspaceHomotopy, v::Vector{ComplexDF64}, t)
    γ = γ!(H, t)
    n = first(size(H.system))
    LA.mul!(H.x_high, γ, v)
    evaluate!(u, H.system, H.x_high)
    u[n+1] = H.x_high[end] - 1.0
end

function ModelKit.evaluate!(u, H::IntrinsicSubspaceHomotopy, v::AbstractVector, t)
    γ = γ!(H, t)
    n = first(size(H.system))
    LA.mul!(H.x, γ, v)
    evaluate!(u, H.system, H.x)
    u[n+1] = H.x[end] - 1.0
    u
end

function ModelKit.evaluate_and_jacobian!(
    u,
    U,
    H::IntrinsicSubspaceHomotopy,
    v::AbstractVector,
    t,
)
    γ = γ!(H, t)

    LA.mul!(H.x, γ, v)
    evaluate_and_jacobian!(u, H.J, H.system, H.x)
    LA.mul!(U, H.J, γ)

    n = first(size(H.system))
    u[n+1] = H.x[end] - 1

    m = length(v)
    for j = 1:m
        U[n+1, j] = γ[end, j]
    end

    nothing
end

function ModelKit.taylor!(u, ::Val{1}, H::IntrinsicSubspaceHomotopy, v, t)
    γ = γ!(H, t)
    γ̇ = γ̇!(H, t)

    # apply chain rule
    #    d/dt [F(γ(t)v); (γ(t)v)[end] - 1] = [J_F(γ(t)v)* γ̇(t)*v;  (γ̇(t)v)[end]]
    LA.mul!(H.x, γ, v)
    LA.mul!(H.ẋ, γ̇, v)
    evaluate_and_jacobian!(u, H.J, H.system, H.x)
    LA.mul!(u, H.J, H.ẋ)
    M = size(H, 1)
    u[M] = H.ẋ[end]

    u
end

function _taylor_γ!(H::IntrinsicSubspaceHomotopy, t::Number)
    @unpack geodesic, taylor_γ = H
    @unpack Q, Q_cos, Θ, U = geodesic

    γ, γ¹, γ², γ³ = taylor_γ
    n, k = size(γ)
    @inbounds for j = 1:k
        Θⱼ = Θ[j]
        s, c = sincos(t * Θⱼ)
        c¹ = -s * Θⱼ
        s¹ = c * Θⱼ
        Θⱼ_2 = 0.5 * Θⱼ^2
        s² = -s * Θⱼ_2
        c² = -c * Θⱼ_2
        Θⱼ_3 = Θⱼ_2 * Θⱼ / 3
        s³ = -c * Θⱼ_3
        c³ = s * Θⱼ_3
        for i = 1:n
            γ[i, j] = Q_cos[i, j] * c + Q[i, j] * s
            γ¹[i, j] = Q_cos[i, j] * c¹ + Q[i, j] * s¹
            γ²[i, j] = Q_cos[i, j] * c² + Q[i, j] * s²
            γ³[i, j] = Q_cos[i, j] * c³ + Q[i, j] * s³
        end
    end

end

function taylor_γ!(H::IntrinsicSubspaceHomotopy, t::Number)
    H.taylor_t_cache[] != t || return H.taylor_γ

    if isreal(t)
        _taylor_γ!(H, real(t))
    else
        _taylor_γ!(H, t)
    end
    H.taylor_t_cache[] = t

    H.taylor_γ
end

function ModelKit.taylor!(
    u,
    v::Val{2},
    H::IntrinsicSubspaceHomotopy,
    tv::TaylorVector,
    t,
    incr::Bool,
)
    γ, γ¹, γ², γ³ = taylor_γ!(H, t)
    x, x¹, x² = vectors(H.tx²)
    v, v¹ = vectors(tv)

    if incr
        x .= H.x
        x¹ .= H.ẋ
    else
        LA.mul!(x, γ, v)
        LA.mul!(x¹, γ¹, v)
    end

    H.v .= v¹
    LA.mul!(H.x, γ, v¹)
    x¹ .+= H.x

    LA.mul!(H.x, γ¹, H.v)
    H.v .= v
    LA.mul!(H.x, γ², H.v, true, true)
    x² .= H.x

    taylor!(u, Val(2), H.system, H.tx²)

    n = first(size(H.system))
    u[n+1] = x²[end]

    u
end

function ModelKit.taylor!(
    u,
    v::Val{3},
    H::IntrinsicSubspaceHomotopy,
    tv::TaylorVector,
    t,
    incr::Bool,
)
    γ, γ¹, γ², γ³ = taylor_γ!(H, t)
    x, x¹, x², x³ = vectors(H.tx³)
    v, v¹, v² = vectors(tv)

    if !incr
        LA.mul!(x, γ, v)
        LA.mul!(x¹, γ¹, v)
        LA.mul!(x¹, γ, v¹, true, true)
        LA.mul!(x², γ², v)
        LA.mul!(x², γ¹, v¹, true, true)
    end

    H.v .= v²
    LA.mul!(H.x, γ, H.v)
    x² .+= H.x

    LA.mul!(H.x, γ¹, H.v)
    H.v .= v
    LA.mul!(H.x, γ³, H.v, true, true)
    H.v .= v¹
    LA.mul!(H.x, γ², H.v, true, true)
    x³ .= H.x

    taylor!(u, Val(3), H.system, H.tx³)

    n = first(size(H.system))
    u[n+1] = x³[end]

    u
end


#=
exp_mult.jl
translated from matlab/C file mri_exp_mult_mex.c
originally by Sangwoo Lee, University of Michigan, Jun, 2004
based on the 2003 file int_tim_seg.m by Brad Sutton
2004-6-21 modified by JF
2020-02 by Daniel Wan
=#

export exp_mult

#using MIRT: exp_xform
using Test:@test

"""
    D = exp_mult(A, u, v ; warnboth)

Memory efficient and fast implementation of `D = A' * exp(-u * v^T)`
that is useful for B0-field-corrected MRI image reconstruction.

in:
* `A [N L]` matrix
* `u [N]` vector
* `v [M]` vector
* `warnboth` warn if both `u` and `v` are complex; default: true

out:
* `D [L M]` complex vector: `D = A' * exp(-u * v^T)`
`D_lm = sum_n A_nl^* exp(-u_n v_m)`
"""
function exp_mult(A, u::AbstractVector{<:Number}, v::AbstractVector{<:Number}
        ; warnboth::Bool = true)

    N, L = size(A)
    n = size(u,1)
    M = size(v,1)

    ~isreal(u) && ~isreal(v) && warnboth &&
        @warn "only one of u and v should be complex"

    T = promote_type(eltype(u), eltype(v), ComplexF32)
    D = zeros(T, L, M)

    for m = 1:M
        col = exp.(-u * v[m]) # [N] mth column of N × M matrix B = exp(-u * v^T)
        D[:,m] = A' * col
    end
    return D
end

"""
    exp_mult(:test)
self test
"""
function exp_mult(test::Symbol)
    test != :test && throw("bad test $test")

    L = 10
    N = 30
    M = 20
    A = randn(ComplexF32, N, L)
    ur = randn(N)
    ui = randn(N)
    vr = randn(M)
    vi = randn(M)
    u = ur + im * ui
    v = vr + im * vi

    d1 = A' * exp.(-ur * transpose(vr))
    @test d1 ≈ exp_mult(A, ur, vr)

    d1 = A' * exp.(-ur * transpose(v))
    @test d1 ≈ exp_mult(A, ur, v)

    d1 = A' * exp.(-u * transpose(vr))
    @test d1 ≈ exp_mult(A, u, vr)

    d1 = A' * exp.(-u * transpose(v))
    @test d1 ≈ exp_mult(A, u, v ; warnboth=false)

    d1 = A' * exp.(-ur * transpose(v))
    ur = reshape(ur, 1, :)
    v = reshape(v, 1, :)
    d2 = exp_xform(transpose(A'), ur, v, mode = :matrix)
    d2 = transpose(d2)
    @test d1 ≈ d2

    true
end


# GSA code v0.1.
# Coded by Jesús Mejía.
# Based on MATLAB code of Esmat Rashedi, 2010.
# Adopted from
# https://la.mathworks.com/matlcgsaentral/fileexchange/61116-gsa-+-chaotic-gravitational-constant
# " E. Rashedi, H. Nezamabadi-pour and S. Saryazdi,
# “GSA: A Gravitational Search Algorithm”, Information sciences, vol. 179,
# no. 13, pp. 2232-2248, 2009."
# ------------------------------------------
# Mirjalili, Seyedali, and Amir H. Gandomi.
# "Chaotic gravitational constants for the gravitational search algorithm."
# Applied Soft Computing 53 (2017): 407-419.
# ------------------------------------------

include("physics.jl")
include("chaos.jl")

mutable struct CGSA <: AbstractParameters
    N::Int
    chValueInitial::Real
    chaosIndex::Real
    ElitistCheck::Int
    Rpower::Int
    Rnorm::Int
    wMax::Real
    wMin::Real
    X::Matrix{Float64}
    V::Matrix{Float64}
    fitness::Vector{Float64}
end

"""

    CGSA(;
        N::Int    = 30,
        chValueInitial::Real   = 20,
        chaosIndex::Real   = 9,
        ElitistCheck::Int    = 1,
        Rpower::Int    = 1,
        Rnorm::Int    = 2,
        wMax::Real   = chValueInitial,
        wMin::Real   = 1e-10,
        information = Information(),
        options = Options()
    )

CGSA is an extension of the GSA algorithm but with Chaotic gravitational constants for
the gravitational search algorithm.

Ref. Chaotic gravitational constants for the gravitational search algorithm.
Applied Soft Computing 53 (2017): 407-419.

Parameters:

- N: Population size
- chValueInitial: Initial value for the chaos value
- chaosIndex: Integer 1 ≤ chaosIndex ≤ 10 is the function that model the chaos
- Rpower: power related to the distance norm(x)^Rpower
- Rnorm: is the value as in norm(x, Rnorm)

# Example

```jldoctest
julia> f(x) = sum(x.^2)
f (generic function with 1 method)

julia> optimize(f, [-1 -1 -1; 1 1 1.0], CGSA())
+=========== RESULT ==========+
  iteration: 500
    minimum: 8.63808e-08
  minimizer: [0.0002658098418993323, -1.140808975532608e-5, -0.00012488307670533095]
    f calls: 15000
 total time: 0.1556 s
+============================+

julia> optimize(f, [-1 -1 -1; 1 1 1.0], CGSA(N = 80, chaosIndex = 1))
+=========== RESULT ==========+
  iteration: 500
    minimum: 1.0153e-09
  minimizer: [-8.8507563788141e-6, -1.3050111801923072e-5, 2.7688577445980026e-5]
    f calls: 40000
 total time: 1.0323 s
+============================+
```

"""
function CGSA(;
        N::Int    = 30,
        chValueInitial::Real   = 20,
        chaosIndex::Real   = 8,
        ElitistCheck::Int    = 1,
        Rpower::Int    = 1,
        Rnorm::Int    = 2,
        wMax::Real   = chValueInitial,
        wMin::Real   = 1e-10,
        X::Matrix{Float64} = zeros(0,0),
        V::Matrix{Float64} = zeros(0,0),
        fitness::Vector{Float64} = zeros(0),
        information = Information(),
        options = Options()
    )

    parameters = CGSA(N, chValueInitial, chaosIndex, ElitistCheck, Rpower, Rnorm, wMax, wMin, X, V, fitness)

    Algorithm(
        parameters,
        information = information,
        options = options,
    )
end

function initialize!(
        status,
        parameters::CGSA,
        problem::AbstractProblem,
        information::Information,
        options::Options,
        args...;
        kargs...
       )

    Rnorm = parameters.Rnorm
    N = parameters.N
    D = size(problem.bounds, 2)
    fobj = problem.f

    # bounds vectors
    low, up = problem.bounds[1,:], problem.bounds[2,:]

    max_it = 500
    options.iterations = options.iterations == 0 ? max_it : options.iterations
    options.f_calls_limit = options.f_calls_limit == 0 ? options.iterations * N : options.f_calls_limit

    # random initialization for agents.
    P = generate_population(N, problem)

    # Current best
    theBest = get_best(P)

    status = State(theBest, P)
    status.f_calls = N

    # Velocity
    parameters.V = isempty(parameters.V) ? zeros(N,D) : parameters.V
    # Postions
    parameters.X = isempty(parameters.X) ? positions(status) : parameters.X
    # function values
    if isempty(parameters.fitness)

        parameters.fitness = fval.(status.population)
    end

    return status

end

function update_state!(
        status,
        parameters::CGSA,
        problem::AbstractProblem,
        information::Information,
        options::Options,
        args...;
        kargs...
        )

    wMax = parameters.wMax
    N = parameters.N
    wMin = parameters.wMin
    max_it = options.iterations
    iteration = status.iteration
    searchType = :minimize
    chaosIndex = parameters.chaosIndex
    Rnorm = parameters.Rnorm
    Rpower = parameters.Rpower
    ElitistCheck = parameters.ElitistCheck

    X = parameters.X
    V = parameters.V
    fitness = parameters.fitness

    P = status.population
    theBest = status.best_sol

    # iteration
    chValue = wMax-iteration*((wMax-wMin)/max_it)

    #Calculation of M. eq.14-20
    M = massCalculation(fitness,searchType)

    #Calculation of Gravitational constant. eq.13.
    G = Gconstant(iteration, max_it)

    if 1 <= chaosIndex <= 10
        G += chaos(chaosIndex,iteration,max_it,chValue)
    end

    #Calculation of accelaration in gravitational field. eq.7-10,21.
    a = Gfield(M,X,G,Rnorm,Rpower,ElitistCheck,iteration,max_it)

    #Agent movement. eq.11-12
    X, V = move(X,a,V)

    #
    # Checking allowable range.
    # X = correctPop(X, low, up)
    for i = 1:N
        x = reset_to_violated_bounds!(X[i,:], problem.bounds)
        X[i,:] = x
    end

    status.population = create_solutions(X, problem)

    parameters.X = X
    parameters.fitness = fvals(status.population)
    parameters.V = V

    #Evaluation of agents.
    currentBest = get_best(status.population)

    if is_better(currentBest, theBest)
        status.best_sol = currentBest
    end

    stop_criteria!(status, parameters, problem, information, options)

end

function final_stage!(
        status,
        parameters::CGSA,
        problem::AbstractProblem,
        information::Information,
        options::Options,
        args...;
        kargs...
    )
    status.final_time = time()
end


# Construct a full grid of 100 x 100 and a discrete grid of 10 x 10
# Put reward states of all the same value from (40,40) to (60,60)
# Run discrete VI and localapproxVI with two different resolutions on each grid
# Then randomly sample several points on full grid and compare the average
# difference from the full value function.
# The approximation with higher resolution should have lower avg. error
function test_against_full_grid()

    # Generate reward states and set to reward 10.0
    rewards = Dict{GWPos, Float64}()
    for x = 40:60
        for y = 40:60
            rewards[GWPos(x, y)] = 10
        end
    end

    # Create full MDP - to be used by both!
    mdp = SimpleGridWorld(size=(100, 100), rewards=rewards)

    # Solve with discrete VI
    solver = ValueIterationSolver(max_iterations=1000, verbose=true)
    policy = solve(solver, mdp)

    # Set up two different grids with different step sizes and ensure that the
    # grid with higher resolution has lower error
    STEP_SIZE_LOW = 5
    STEP_SIZE_HI = 10

    grid_low = RectangleGrid(range(1,step=STEP_SIZE_LOW,stop=100), range(1,step=STEP_SIZE_LOW,stop=100))
    grid_hi = RectangleGrid(range(1,step=STEP_SIZE_HI,stop=100), range(1,step=STEP_SIZE_HI,stop=100))

    interp_low = LocalGIFunctionApproximator(grid_low)
    interp_hi = LocalGIFunctionApproximator(grid_hi)

    solver_low = LocalApproximationValueIterationSolver(interp_low, verbose=true, max_iterations = 1000)
    solver_hi = LocalApproximationValueIterationSolver(interp_hi, verbose=true, max_iterations = 1000)

    policy_low = solve(solver_low, mdp)
    policy_hi = solve(solver_hi, mdp)

    total_err_low = 0.0
    total_err_hi = 0.0

    for state in states(mdp)

        full_val = value(policy, state)

        approx_val_low = value(policy_low, state)
        approx_val_hi = value(policy_hi, state)

        total_err_low += abs(full_val-approx_val_low)
        total_err_hi += abs(full_val-approx_val_hi)
    end

    avg_err_low = total_err_low / length(states(mdp))
    avg_err_hi = total_err_hi / length(states(mdp))

    return (avg_err_low < avg_err_hi)
end

@test test_against_full_grid() == true


#########
# AngularMomentum
# Applies the partial derivative with respect to the last angular variable in the coordinate system.
# For example, in polar coordinates (r, θ) in ℝ² or cylindrical coordinates (r, θ, z) in ℝ³, we apply ∂ / ∂θ = (x ∂ / ∂y - y ∂ / ∂x).
# In spherical coordinates (ρ, θ, φ) in ℝ³, we apply ∂ / ∂φ = (x ∂ / ∂y - y ∂ / ∂x).
#########

struct AngularMomentum{T,Ax<:Inclusion} <: LazyQuasiMatrix{T}
    axis::Ax
end

AngularMomentum{T}(axis::Inclusion) where T = AngularMomentum{T,typeof(axis)}(axis)
AngularMomentum{T}(domain) where T = AngularMomentum{T}(Inclusion(domain))
AngularMomentum(axis) = AngularMomentum{eltype(eltype(axis))}(axis)

axes(A::AngularMomentum) = (A.axis, A.axis)
==(a::AngularMomentum, b::AngularMomentum) = a.axis == b.axis
copy(A::AngularMomentum) = AngularMomentum(copy(A.axis))

^(A::AngularMomentum, k::Integer) = ApplyQuasiArray(^, A, k)

@simplify function *(A::AngularMomentum, P::SphericalHarmonic)
    # Spherical harmonics are the eigenfunctions of the angular momentum operator on the unit sphere
    T = real(eltype(P))
    k = mortar(Base.OneTo.(1:2:∞))
    m = T.((-1) .^ (k .- 1) .* k .÷ 2)
    P * Diagonal(im .* m)
end

#=
@simplify function *(A::AngularMomentum, P::RealSphericalHarmonic)
    # The angular momentum operator applied to real spherical harmonics negates orders
    T = real(eltype(P))
    n = mortar(Fill.(oneto(∞),1:2:∞))
    k = mortar(Base.OneTo.(0:2:∞))
    m = T.(iseven.(k) .* (k .÷ 2))
    dat = PseudoBlockArray(Vcat(
        (-m)', # n, k-1
        (0 .* m)', # n, k
        (m)', # n, k+1
        ), (blockedrange(Fill(3, 1)), axes(n,1)))
    P * _BandedBlockBandedMatrix(dat', axes(k,1), (0,0), (1,1))
end
=#


### A Pluto.jl notebook ###
# v0.14.5

using Markdown
using InteractiveUtils

# ╔═╡ 13e5b5ba-9f66-46a8-81cc-44a985163243
using Luxor

# ╔═╡ 546912b4-9af8-47d9-8b13-14abee56f9fd
RULES = Dict('X' => ["F+[[X]-X]-F[-FX]+X"],
			 'F' => ["F[+F]F[−-F][F]", "F[+F]F[--F]F", "FF-[-F+F+F]+[+F-F-F]"])

# ╔═╡ bb2287e2-aad1-4134-8048-cf70c9c1cfe3
function apply_rule(char)
	if char in keys(RULES)
		new_char = RULES[char]
		new_char[rand(1:length(new_char))]
	else
		string(char)
	end
end

# ╔═╡ 8ed3d85d-b975-4ebe-9703-d2febf294733
function rewrite_l(word)
	join([apply_rule(c) for c in word])
end

# ╔═╡ f105753b-5cb6-428b-a888-e28972a57775
function transform_l_tree(axiom, iterations)

	word = axiom

	for i = 1:iterations
		word = rewrite_l(word)
	end

	word
end

# ╔═╡ 0bac16c7-d7ca-44a4-a0c6-c31e1f3d344f
trees = [transform_l_tree("F", rand(3:6)) for i in 1:1]

# ╔═╡ d8d971c8-6d4d-40c9-8d01-7640804f28ba
@png begin
	background("lightgray")

	🐢 = Turtle()
	Pencolor(🐢, "black")
	Penwidth(🐢, 0.4)

	for tree in trees
		Reposition(🐢, rand(-200:200), rand(150:200))
		Orientation(🐢, -90)
		for i in tree
			if i == 'F'
				Forward(🐢, 1)
			elseif i == '-'
				Turn(🐢, -30)
			elseif i == '+'
				Turn(🐢, +30)
			elseif i == '['
				Push(🐢)
			elseif i == ']'
				# Pencolor(🐢, "green")
				# Forward(🐢, 1)
				# Pencolor(🐢, "black")
				Pop(🐢)
			end
		end
	end

	finish()
end 400 400 "turtle.png"

# ╔═╡ Cell order:
# ╠═13e5b5ba-9f66-46a8-81cc-44a985163243
# ╠═546912b4-9af8-47d9-8b13-14abee56f9fd
# ╠═bb2287e2-aad1-4134-8048-cf70c9c1cfe3
# ╠═8ed3d85d-b975-4ebe-9703-d2febf294733
# ╠═f105753b-5cb6-428b-a888-e28972a57775
# ╠═0bac16c7-d7ca-44a4-a0c6-c31e1f3d344f
# ╠═d8d971c8-6d4d-40c9-8d01-7640804f28ba


function create_initial_condition(kx, kz, k2cut, N, parameters::GeneralParams)
    datah, datap, dataq, U = initial_condition(parameters);

    datas1s2r1r2 = zeros(length(datah))'

    # remove high frequency parts by aliasing filter
    datah = datah .* k2cut;
    datap = datap .* k2cut;
    dataq = dataq .* k2cut;

    # reshape data for U vector
    datahSV = reshape(datah', 1, :)
    datapSV = reshape(datap', 1, :)
    dataqSV = reshape(dataq', 1, :)

    U = vcat(datahSV[1:N ÷ 2], dataqSV[1:N ÷ 2], datapSV[1:N ÷ 2])
    U = reshape(U, 1, :)
    return U
end

function initial_condition(parameters::GeneralParams)
    @unpack NN1, NN2 = parameters
    @unpack stSine, spSine, stCosine, spCosine, stNwaves, spNwaves, stNoise, spNoise = parameters
    @unpack k0x, k0z, h0, Lx, Lz = parameters

    Ax = stSine
    Az = spSine
    Bx = stCosine
    Bz = spCosine
    nwx = stNwaves
    nwz = spNwaves

    if (NN1 == 0) k0x = 0; Lx = 0; Ax = 0; Bx = 0 end
    if (NN2 == 0) k0z = 0; Lz = 0; Az = 0; Bz = 0 end

    N1 = 2^NN1; M1 = 2 * N1; N2 = 2^NN2; M2 = 2 * N2; N = 4 * N2 * (N1 + 1);

    datah = zeros((M1, M2))
    dataq = zeros((M1, M2))
    datap = zeros((M1, M2))

    randz = zeros(M2)
    for n = 1:M2
        randz[n] = spNoise * rand()
    end

    for m = 1:M1
        x = (m - 1) * Lx / M1
        for n = 1:M2
            z = (n - 1) * Lz / M2
            datah[m, n] = h0 * (1 + Ax * sin(nwx * k0x * x) + Bx * cos(nwx * k0x * x) + Az * sin(nwz * k0z * z) + Bz * cos(nwz * k0z * z) + stNoise * rand()) + randz[n]
            datap[m, n] = 0
            dataq[m, n] = datah[m, n].^3 / 3
        end
    end

    # calculate mass
    mass = sum(datah)
    mass = mass / (M1 * M2)

    # set mass to h0 for random function
    if spNoise != 0
        datah = datah - mass + h0
    end

    # Fourier transform of the initial condition
    datah = fft(datah)
    datap = fft(datap)
    dataq = fft(dataq)

    datah = conj(datah)
    datap = conj(datap)
    dataq = conj(dataq)

    datahSV = vec(datah')
    datapSV = vec(datap')
    dataqSV = vec(dataq')

    M1M2 = M1 * M2;

    # note ÷ for integer division
    U = vcat(datahSV[1:N ÷ 2], dataqSV[1:N ÷ 2], datapSV[1:N ÷ 2]) ./ M1M2

    datah = datah ./ M1M2;
    dataq = dataq ./ M1M2;
    datap = datap ./ M1M2;

    return datah, datap, dataq, U

end

function aliasing_filter(radius::Float64, parameters::GeneralParams)
    @unpack k0x, k0z, N1, N2, M1, M2 = parameters;

    m = round(radius * N1)
    n = round(radius * N2)
    k2cutCriterion = (m * k0x)^2 + ((-N2 + mod(n + N2, M2)) * k0z)^2
    kx, kz, k2, k2cut = [zeros((N1 + 1, M2)) for i = 1:4]  # initialise
    for m = 1:N1 + 1
        for n = 1:M2
            kx[m, n] = (m - 1) * k0x
            kz[m, n] = (-N2 + mod((n - 1) + N2, M2)) * k0z
            k2[m, n] = kx[m, n] * kx[m, n] + kz[m, n] * kz[m, n]
            if (k2[m, n] > k2cutCriterion)
                k2cut[m, n] = 0
            else
                k2cut[m, n] = 1
            end
        end
    end
    k2cut = [k2cut; [conj(reverse(k2cut[2:N1,1], dims=1)) conj(reverse(reverse(k2cut[2:N1,2:M2], dims=1), dims=2))]]
    return kx, kz, k2cut
end


using Test

using ReadStat, DataValues # for the Petersen data set

using FamaMacBeth

@testset "general" begin
    y = rand(100,1)
    X = rand(100,3)

    # if all in the same group, we should get OLS coefficients
    g = vec(ones(size(y,1),1))
    @test coef(fmb(X,y,g))==X\y

    g2 = repeat(collect(1:10),inner=10)
    resp = fmb(X,y,g2)
    @test all(nobs(resp).==10)
    @test size(coef(resp),1) == 3
    @test size(coef(resp),2) == 1 # because we take the mean!
end

@testset "petersen_dataset" begin
    # download and prcess data
    file = "https://www.kellogg.northwestern.edu/faculty/petersen/htm/papers/se/test_data.dta"
    loc = download(file)
    data = read_dta(loc)
    firmid = Float64.(get.(data.data[1]))
    year = Float64.(get.(data.data[2]))
    x = Float64.(get.(data.data[3]))
    y = Float64.(get.(data.data[4]))

    # fit
    m = fmb([ones(size(x,1),1) x], y, year)

    # compare with https://www.kellogg.northwestern.edu/faculty/petersen/htm/papers/se/test_data.htm
    @test all(round.(coef(m), digits=4) .== [.0313, 1.0356])
    @test all(round.(stderror(m), digits=4) .== [.0234, .0333])
    @test all(round.(tstat(m), digits=2) .== [1.34, 31.06])
    # @test round(adjr2(m),4)==.2078 # this is what Petersen lists, but cannot reproduce
    @test round(adjr2(m);digits=4)==.2076
end


using SymPyCall
using LinearAlgebra
using Test

## SymPyCall only is adding methods for those found in SpecialFunctions
## so not the special polynomials, etc.
using SpecialFunctions

@testset "SpecialFuns" begin
    @syms a b x y
    @syms n m θ ϕ
    @syms ν::integer

    @test sympy.fresnels(Sym(0)) == 0
    @test sympy.fresnels(Sym(oo)) == Sym(1)/2
    @test diff(sympy.fresnels(x), x) == sin(PI*x^2/2)
    #@test evalf(fresnels(Sym(2)), 30) == Sym(parse(BigFloat, "0.343415678363698242195300815958"))

    @test  diff(sympy.fresnelc(x), x) == cos(PI*x^2/2)

    @test sympy.Ei(Sym(-1)) == sympy.Ei(exp(1im*PI))
    @test diff(sympy.Ei(x), x) == exp(x)/x
    @test diff(sympy.Si(x), x) == sin(x)/x
    @test diff(sympy.Ci(x), x) == cos(x)/x

    @test airyai(Sym(0)) == 3^(Sym(1)/3)/(3*gamma(Sym(2)/3))
    @test airybi(Sym(0)) == 3^(Sym(5)/6)/(3*gamma(Sym(2)/3))

    @test sympy.jacobi(0, a, b, x) == 1
    @test sympy.jacobi(1, a, b, x) == a/2 - b/2 + x*(a/2 + b/2 + 1)
    @test sympy.jacobi(n, 0, 0, x) == sympy.legendre(n, x)
    @test sympy.jacobi(n, a, b, -x) == (-1)^n*sympy.jacobi(n, b, a, x)
    @test conj(sympy.jacobi(n, a, b, x)) == sympy.jacobi(n, conj(a), conj(b), conj(x))

    @test sympy.gegenbauer(0, a, x) == 1
    @test sympy.gegenbauer(1, a, x) == 2a*x
    @test sympy.gegenbauer(2, a, x) == -a + x^2*(2a^2 + 2a)
    @test conj(sympy.gegenbauer(n, a, x)) == sympy.gegenbauer(n, conj(a), conj(x))
    @test diff(sympy.gegenbauer(n, a, x), x) == 2a*sympy.gegenbauer(n-1, a+1, x)

    @test sympy.chebyshevt(0, x) == 1
    @test sympy.chebyshevt(1, x) == x
    @test sympy.chebyshevt(2, x) == 2x^2-1

    @test sympy.chebyshevu(0, x) == 1
    @test sympy.chebyshevu(1, x) == 2x
    @test sympy.chebyshevu(2, x) == 4x^2 - 1
    @test sympy.chebyshevu(n, 0) == cos(π*n/2)
    @test sympy.chebyshevu(n, 1) == n + 1

    @test sympy.legendre(0, x) == 1
    @test sympy.legendre(1, x) == x
    @test sympy.legendre(2, x) == 3x^2/2 - 1//2
    @test sympy.legendre(n, x) == sympy.legendre(n, x)
    @test diff(sympy.legendre(n,x), x) == n*(x*sympy.legendre(n, x) - sympy.legendre(n - 1, x))/(x^2 - 1)

    @test sympy.assoc_legendre(0,0, x) == 1
    @test sympy.assoc_legendre(1,0, x) == x
    @test sympy.assoc_legendre(1,1, x) == -sqrt(-x^2 + 1)
    #@test assoc_legendre(n,m,x) == assoc_legendre(n, m, x)

    @test sympy.hermite(0, x) == 1
    @test sympy.hermite(1, x) == 2x
    @test sympy.hermite(2, x) == 4x^2 - 2
    @test sympy.hermite(n, x) == sympy.hermite(n, x)
    @test diff(sympy.hermite(n,x), x) == 2n*sympy.hermite(n - 1, x)
    @test sympy.hermite(n, -x) == (-1)^n*sympy.hermite(n, x)

    @test sympy.laguerre(0, x) == 1
    @test sympy.laguerre(1, x) == -x + 1
    @test sympy.laguerre(2, x) == x^2/2 - 2x + 1
    @test sympy.laguerre(3, x) == -x^3/6 + 3x^2/2 - 3x + 1
    @test diff(sympy.laguerre(n, x), x) == -sympy.assoc_laguerre(n - 1, 1, x)

    @test sympy.assoc_laguerre(0, a, x) == 1
    @test sympy.assoc_laguerre(1, a, x) == a - x + 1
    @test sympy.assoc_laguerre(2, a, x) == a^2/2 + 3a/2 + x^2/2 + x*(-a - 2) + 1
    @test sympy.assoc_laguerre(n, 0, x) == sympy.laguerre(n, x)
    @test diff(sympy.assoc_laguerre(n, a, x), x) == -sympy.assoc_laguerre(n - 1, a + 1, x)

    @test sympy.Ynm(n, m, θ, ϕ) == sympy.Ynm(n, m, θ, ϕ)
    @test sympy.Ynm(n, -m, θ, ϕ) == (-1)^m*exp(-2im*m*ϕ)*sympy.Ynm(n, m, θ, ϕ)
    @test sympy.Ynm(n, m, -θ, ϕ) == sympy.Ynm(n, m, θ, ϕ)
    @test sympy.Ynm(n, m, θ, -ϕ) == exp(-2im*m*ϕ)*sympy.Ynm(n, m, θ, ϕ)
    @test expand(simplify(sympy.Ynm(0, 0, θ, ϕ)), func=true) == 1/(2*sqrt(PI))

    @test sympy.Ynm_c(n, m, θ, ϕ) == conj(sympy.Ynm(n, m, θ, ϕ))

    @test diff(hankelh1(n, x), x) == hankelh1(n - 1, x)/2 - hankelh1(n + 1, x)/2
    @test diff(hankelh2(n, x), x) == hankelh2(n - 1, x)/2 - hankelh2(n + 1, x)/2

    @test sympy.expand_func(sympy.jn(0, x)) == sin(x)/x
    @test sympy.expand_func(sympy.jn(1, x)) == sin(x)/x^2 - cos(x)/x
    @test rewrite(sympy.jn(ν, x), "besselj") == sqrt(2PI/x)*besselj(ν + Sym(1)/2, x)/2
    VERSION < v"0.6.0-dev" && @test rewrite(sympy.jn(ν, x), "bessely") == (-1)^ν*sqrt(2PI/x)*bessely(-ν - Sym(1)/2, x)/2
    u = N(sympy.jn(2, 5.2+0.3im).evalf(20))
    @test norm(real(u) - 0.099419756723640344491) <= 1e-15 && norm(imag(u) + 0.054525080242173562897) <= 1e-15

    @test sympy.expand_func(sympy.yn(0, x)) == -cos(x)/x
    @test sympy.expand_func(sympy.yn(1, x)) == -cos(x)/x^2-sin(x)/x
    VERSION < v"0.6.0-dev" && @test rewrite(sympy.yn(ν, x), "besselj") == (-1)^(ν + 1)*sqrt(2PI/x)*besselj(-ν - Sym(1)/2, x)/2
    @test rewrite(sympy.yn(ν, x), "bessely") == sqrt(2PI/x)*bessely(ν + Sym(1)/2, x)/2
    @test N(sympy.yn(2, 5.2+0.3im)) ≈ 0.18525034196069722536 + 0.014895573969924817587im

    # gamma, beta and related functions
    @test gamma(Sym(1)) == 1
    @test gamma(Sym(3)/2) == sqrt(PI)/2

    @test diff(polygamma(Sym(0), x), x) == polygamma(Sym(1), x)
    @test diff(polygamma(Sym(0), x), x, 2) == polygamma(Sym(2), x)

    @test diff(beta(x, y), x) == (polygamma(Sym(0), x) - polygamma(Sym(0), x + y)) * beta(x, y)
    @test diff(beta(x, y), y) == (polygamma(Sym(0), y) - polygamma(Sym(0), x + y)) * beta(x, y)

    # test numerical consistency with Julia functions
    @test N(gamma(Sym(4.1))) ≈ gamma(4.1)
    @test N(polygamma(Sym(2), Sym(3.2))) ≈ polygamma(2, 3.2)
    VERSION >= v"0.5.0" && @test N(beta(Sym(1)+1im, Sym(1)+1im)) ≈ beta(1.0+1im, 1.0+1im)

    # Elliptic-type functions

    @test sympy.elliptic_k(Sym(0)) == PI/2
    @test N(sympy.elliptic_k(Sym(1.0 + im))) ≈ 1.50923695405127 + 0.625146415202697im

    @test N(sympy.elliptic_f(Sym(3.0 + im/2), Sym(1.0 + im))) ≈ 2.909449841483 + 1.74720545502474im

    @test sympy.elliptic_e(Sym(0)) == PI/2
    @test N(sympy.elliptic_e(Sym(2.0 - im))) ≈ 0.991052601328069 + 0.81879421395609im

    @test sympy.elliptic_pi(Sym(0), Sym(0)) == PI/2
    @test N(sympy.elliptic_pi(Sym(1.0 - im/3), Sym(2.0 + im))) ≈ 3.29136443417283 + 0.32555634906645im

    # Bessel-type functions
    @test diff(besselj(n, x), x) == (besselj(n - 1, x) - besselj(n + 1, x))/2
    @test rewrite(besselj(n, x), "jn") == sqrt(2x/PI)*sympy.jn(n - 1//2, x)

    @test diff(bessely(n, x), x) == (bessely(n - 1, x) - bessely(n + 1, x))/2
    @test rewrite(bessely(n, x), "yn") == sqrt(2x/PI)*sympy.yn(n - 1//2, x)

    @test diff(besseli(n, x), x) == (besseli(n - 1, x) + besseli(n + 1, x))/2

    @test diff(besselk(n, x), x) == -(besselk(n - 1, x) + besselk(n + 1, x))/2

    # test numerical consistency with Julia functions
    @test N(besselj(Sym(3.2), Sym(1.5))) ≈ besselj(3.2, 1.5)
    @test N(bessely(Sym(3.2), Sym(1.5))) ≈ bessely(3.2, 1.5)
    @test N(besseli(Sym(3.2), Sym(1.5))) ≈ besseli(3.2, 1.5)
    @test N(besselk(Sym(3.2), Sym(1.5))) ≈ besselk(3.2, 1.5)

    @test sympy.expand_func(x*sympy.hyper([1, 1], [2], -x)) == log(x + 1)
end


immutable NoncentralT <: ContinuousUnivariateDistribution
    ν::Float64
    λ::Float64

    function NoncentralT(ν::Real, λ::Real)
    	@check_args(NoncentralT, ν > zero(ν))
        @check_args(NoncentralT, λ >= zero(λ))
        new(ν, λ)
    end
end

@distr_support NoncentralT -Inf Inf

### Parameters

params(d::NoncentralT) = (d.ν, d.λ)

### Statistics

function mean(d::NoncentralT)
    if d.ν > 1.0
        isinf(d.ν) ? d.λ :
        sqrt(0.5*d.ν) * d.λ * gamma(0.5*(d.ν-1)) / gamma(0.5*d.ν)
    else
        NaN
    end
end

var(d::NoncentralT) = d.ν > 2.0 ? d.ν*(1+d.λ^2)/(d.ν-2.0) - mean(d)^2 : NaN

### Evaluation & Sampling

@_delegate_statsfuns NoncentralT ntdist ν λ

function rand(d::NoncentralT)
    z = randn()
    v = rand(Chisq(d.ν))
    (z+d.λ)/sqrt(v/d.ν)
end


function momentum_grid(kmin, kmax, nk, nθ; spacing=:momentum)
    kmag = if spacing == :momentum
        range(kmin, stop=kmax, length=nk)
    elseif spacing == :energy
        .√(2range(kmin^2/2, stop=kmax^2/2, length=nk))
    else
        throw(ArgumentError("Unknown spacing $(spacing)"))
    end
    nθ == 1 && return (kmag,kmag,nothing)

    x = range(0, stop=1, length=nθ)
    k = Array{SVector{3,Float64}}(undef, length(kmag), nθ)
    for j = 1:nθ
        s,c = sincospi(x[j])
        k̂ = SVector{3}(s, 0, c)
        for (i,km) in enumerate(kmag)
            k[i,j] = km*k̂
        end
    end
    k,kmag,x*π
end

momentum_grid(Emin::Unitful.Energy, Emax::Unitful.Energy, args...; kwargs...) =
    momentum_grid(√(2austrip(Emin)), √(2austrip(Emax)), args...; kwargs...)

export momentum_grid


using FrankWolfe
using ProgressMeter
using Arpack
using Plots
using DoubleFloats
using ReverseDiff

import LinearAlgebra

# n = Int(1e1)
n = Int(1e4)
k = 5 * Int(1e3)
number_nonzero = 40

xpi = rand(n);
total = sum(xpi);
const xp = xpi # ./ total;

f(x) = LinearAlgebra.norm(x - xp)^2
function grad!(storage, x)
    @. storage = 2 * (x - xp)
end

lmo = FrankWolfe.KSparseLMO(number_nonzero, 1.0);
## alternative lmo
# lmo = FrankWolfe.ProbabilitySimplexOracle(1)
x0 = FrankWolfe.compute_extreme_point(lmo, ones(n));

@time x, v, primal, dual_gap, trajectorylazy, active_set = FrankWolfe.away_frank_wolfe(
    f,
    grad!,
    lmo,
    x0,
    max_iteration=k,
    line_search=FrankWolfe.Adaptive(),
    print_iter=k / 10,
    epsilon=1e-5,
    emphasis=FrankWolfe.memory,
    verbose=true,
    trajectory=true,
    lazy=true,
);

@time x, v, primal, dual_gap, trajectoryAFW, active_set = FrankWolfe.away_frank_wolfe(
    f,
    grad!,
    lmo,
    x0,
    max_iteration=k,
    line_search=FrankWolfe.Adaptive(),
    print_iter=k / 10,
    epsilon=1e-5,
    emphasis=FrankWolfe.memory,
    verbose=true,
    away_steps=true,
    trajectory=true,
);

@time x, v, primal, dual_gap, trajectoryFW = FrankWolfe.away_frank_wolfe(
    f,
    grad!,
    lmo,
    x0,
    max_iteration=k,
    line_search=FrankWolfe.Adaptive(),
    print_iter=k / 10,
    emphasis=FrankWolfe.memory,
    verbose=true,
    epsilon=1e-5,
    trajectory=true,
    away_steps=false,
);

data = [trajectorylazy, trajectoryAFW, trajectoryFW]
label = ["LAFW" "AFW" "FW"]

FrankWolfe.plot_trajectories(data, label, xscalelog=true)


# This file was formerly a part of Julia. License is MIT: https://julialang.org/license

# Benchmark European option
# https://groups.google.com/forum/?hl=en&fromgroups=#!topic/julia-dev/ImhGsqX_IHc

function bench_eu_devec(numPaths)
    steps = 250
    r = 0.05
    sigma = .4
    T = 1
    dt = T/(steps)
    K = 100

    S = fill(100.,numPaths,1)

    t1 = (r-0.5*sigma.^2)*dt
    t2 = sigma*sqrt(dt)
    for i=1:steps
        for j=1:numPaths
            S[j] *= exp(t1 + t2*randn())
        end
    end

    V = mean( exp(-r*T)*max.(K.-S,0) )
end

function bench_eu_vec(numPaths)
    steps = 250
    r = 0.05
    sigma = .4
    T = 1
    dt = T/(steps)
    K = 100

    S = fill(100.,numPaths,1)

    t1 = (r-0.5*sigma.^2)*dt
    t2 = sigma*sqrt(dt)
    for i=1:steps
        S = S .* exp.(t1.+t2*randn(numPaths))
    end

    V = mean( exp(-r*T)*max.(K.-S,0) )
end


#Example demonstrating the use of overlap to solve a long horizon control problem
using Plasmo, Ipopt
using KaHyPar
using SchwarzOpt

T = 100              #number of time points
d = sin.(1:T)        #a disturbance vector
imbalance = 0.1      #partition imbalance
distance = 5         #expand distance
n_parts = 6          #number of partitions

#Create the model (an optigraph)
graph = OptiGraph()
@optinode(graph,state[1:T])
@optinode(graph,control[1:T-1])

for (i,node) in enumerate(state)
    @variable(node,x)
    @constraint(node, x >= 0)
    @objective(node,Min,0.001*x^2) #- 2*x*d[i])
end
for node in control
    @variable(node,u)
    @constraint(node, u >= -1000)
    @objective(node,Min,u^2)
end
n1 = state[1]
@constraint(n1,n1[:x] == 0)

for i = 1:T-1
    @linkconstraint(graph, state[i][:x] + control[i][:u] + d[i] == state[i+1][:x],attach = state[i+1])
end

#Partition the optigraph using recrusive bisection over a hypergraph
hypergraph,hyper_map = hyper_graph(graph) #create hypergraph object based on graph

partition_vector = KaHyPar.partition(hypergraph,n_parts,
configuration = (@__DIR__)*"/cut_rKaHyPar_sea20.ini",
imbalance = imbalance)

partition = Partition(hypergraph,partition_vector,hyper_map)
apply_partition!(graph,partition)

#Inspect the graph structure. It should be RECURSIVE_GRAPH
println(Plasmo.graph_structure(graph))

#calculate subproblems using expansion distance
subgraphs = getsubgraphs(graph)
expanded_subgraphs = Plasmo.expand.(graph,subgraphs,distance)
sub_optimizer = optimizer_with_attributes(Ipopt.Optimizer,"print_level" => 0)

#optimize using schwarz overlapping decomposition
SchwarzOpt.optimize!(graph;
subgraphs = expanded_subgraphs,
sub_optimizer = sub_optimizer,
max_iterations = 50)


using Pkg; Pkg.activate(".."); Pkg.instantiate()

using OceanTurb, Printf

@use_pyplot_utils

# Instantiate a diffusion model with 100 grid points, height 1.0,
# diffusivity 0.01, and a backward Euler timestepper.
model = Diffusion.Model(
    N = 100,
    L = 1.0,
    K = 0.01,
    stepper = :BackwardEuler)

# Set the initial condition to a Gaussian cenetered at z=-0.5.
c₀(z) = exp(-(z + 0.5)^2 / 0.005)
model.solution.c = c₀

# Iterate the model forward
iterate!(model, Δt=0.01, Nt=100)

# Plot some results
fig, axs = subplots()
xlabel(L"c")
ylabel(L"z")

plot(c₀.(model.grid.zc), label=L"t=0")
gcf()

plot(model.solution.c, label=@sprintf("\$ t = %0.2f \$", time(model)))
removespines("top", "right")

legend()
gcf()


function probit_reg(y, X, Xfull, nmcmc)

	N = size(Xfull)[1]
	n = length(y)

	beta_save = fill(0.0, 2, nmcmc)
	p_save = fill(0.0, N, nmcmc)
	y_save = BitArray(undef, N, nmcmc)

	# Priors
	Σ_β = ScalMat(2, 2.25)
	Σ_β_inv = inv(Σ_β)
	μ_β = fill(0.0, 2)

	# Initial values
	β = μ_β

	Σtemp = PDMat(Symmetric(inv(Σ_β_inv + X' * X)))

	# Gibbs loop
	for k in 1:nmcmc

		# Sample z
		z = fill(0.0, length(y))
		μ = X * β

		z[y .== 1] = rand.(Truncated.(Normal.(μ[y .== 1], 1.0), 0.0, Inf))
		z[y .== 0] = rand.(Truncated.(Normal.(μ[y .== 0], 1.0), -Inf, 0.0))

		# Sample beta
		βtemp = Σtemp * (Σ_β_inv * μ_β + X' * z)

		β = rand(MvNormal(βtemp, Σtemp))

		beta_save[:, k] = β
		p_save[:, k] = cdf.(Normal(), Xfull * β)
		y_save[:, k] = rand.(Bernoulli.(p_save[:, k]))

	end

	return (β = beta_save, p = p_save, y = y_save)
end

function probit_pprb(ynew, loc, p_prop, β_prop, nmcmc)

	nprop = size(p_prop)[2]

	β_save = fill(0.0, 2, nmcmc)
	y_save = BitArray(undef, N, nmcmc)
	ynew_save = fill(0, nmcmc)

	init = sample(1:nprop, 1)
	p = p_prop[:, init]

	β = β_prop[:, init]
	β_save[:, 1] = β

	@inbounds for k in 1:nmcmc

		idx = sample(1:nprop, 1)
		pstar = p_prop[:, idx]

		mh1 = logpdf(Bernoulli(pstar[loc]), ynew)
		mh2 = logpdf(Bernoulli(p[loc]), ynew)

		r = exp(mh1 - mh2)
		if r > rand(Uniform())
			p = pstar
			β = β_prop[:, idx]
		end

		β_save[:, k] = β
		# p_save[:, k] = p
		y_save[:, k] = rand.(Bernoulli.(p))
		ynew_save[k] = ynew

	end

	return (y = y_save, ynew = ynew_save, β = β_save)

end


using ManifoldLearning
import ManifoldLearning: BruteForce, knn, swiss_roll
using Test
import Random

Random.seed!(3483743871)

@testset "Nearest Neighbors" begin

    # setup parameters
    k = 12
    X, L = swiss_roll(100)
    DD, EE = knn(X,k)
    @test_throws AssertionError knn(rand(3,10), k)

    NN = fit(BruteForce, X, k)
    D, E = knn(NN, X)
    @test size(X,2) == size(D,2) && size(D, 1) == k
    @test size(X,2) == size(E,2) && size(E, 1) == k
    @test E == EE
    @test D ≈ DD

    D, E = knn(NN, X[:,1:k+1])
    @test k+1 == size(D,2) && size(D, 1) == k
    @test k+1 == size(E,2) && size(E, 1) == k

    D, E = knn(NN, X[:,1:k+1], self=true)
    @test D[1,:] == zeros(k+1)
    @test E[1,:] == collect(1:k+1)

    @test_throws AssertionError knn(NN, X[:,1:k])
end

@testset "Manifold Learning" begin

    # setup parameters
    k = 12
    d = 2
    X, L = swiss_roll()

    # test algorithms
    @testset for algorithm in [Isomap, LEM, LLE, HLLE, LTSA, DiffMap]
        for (k, T) in zip([5, 12], [Float32, Float64])
            # construct KW parameters
            kwargs = [:maxoutdim=>d]
            if algorithm == DiffMap
                push!(kwargs, :t => k)
            else
                push!(kwargs, :k => k)
            end

            # call transformation
            Y = fit(algorithm, convert(Array{T}, X); kwargs...)

            # test results
            @test outdim(Y) == d
            @test size(transform(Y), 2) == size(X, 2)
            @test length(split(sprint(show, Y), '\n')) > 1
            @test length(eigvals(Y)) == d
            if algorithm !== DiffMap
                @test neighbors(Y) == k
                @test length(vertices(Y)) > 1
            end
        end
    end

end

@testset "OOS" begin
    k = 10
    d = 2
    X, _ = swiss_roll()
    M = fit(Isomap, X; k=k, maxoutdim=d)

    @test all(sum(abs2, transform(M) .- transform(M,X), dims=1) .< eps())
end


function norm2_metric(x)
	m, n = size(x)
	D = zeros(eltype(x), n, n)
	@inbounds for i in 1:n
		for j in i+1:n
			D[i,j] = D[j,i] = norm(view(x, :, i) .- view(x, :, j))
		end
	end
	D
end

function norm2_metric_back(x, D, sym_dD)
	dD = 0.5*(sym_dD + sym_dD')  # Cholesky decomposition only use lower triangle, reconstruct the true dD
	m, n = size(x)
	T = eltype(x)
	dx = zeros(T, m, n)
	@inbounds for i in 1:m
		for j in 1:n
			for k in 1:n
				dx[i,j] += 2.0*(dD[j,k]/(D[j,k]+eps(T)))*(x[i,j]-x[i,k])
			end
		end
	end
	dx
end

function norm2_metric(x, xo)
	N1 = size(x, 2)
	N2 = size(xo, 2)
	D = zeros(eltype(x), N1, N2)
	@inbounds for i in 1:N1
		for j in 1:N2
			D[i,j] = norm(view(x, :, i) .- view(xo, :, j))
		end
	end
	D
end


using HighDimPDE
using Random
using Test
using Statistics
using Revise

#relative error l2
function rel_error_l2(u, uanal)
    if abs(uanal) >= 10 * eps(eltype(uanal))
        sqrt((u - uanal)^2 / u^2)
    else # overflow
        abs(u-uanal)
    end
end

# https://en.wikipedia.org/wiki/Heat_equation#Fundamental_solutions
@testset "MLP - heat equation - single thread" begin # solving at one unique point
    for d in [1, 10, 50]
        println("test for d = ", d)
        tspan = (0e0, 5e-1)
        dt = 5e-1  # time step
        μ(x, p, t) = 0e0 # advection coefficients
        σ(x, p, t) = 1e-1 # diffusion coefficients

        u_anal(x, t) = sum(x.^2) + d * σ(0., 0., 0.)^2 * t
        g(x) = sum(x.^2)

        # d = 10
        x0 = fill(2e0,d)

        alg = MLP()

        f(y, z, v_y, v_z, ∇v_y, ∇v_z, p, t) = 0e0 .* v_y #TODO: this fix is not nice

        # defining the problem
        prob = PIDEProblem(g, f, μ, σ, tspan, x = x0)
        # solving
        xs,ts,sol = solve(prob, alg, multithreading = false)
        u1 = sol[end]
        u1_anal = u_anal(x0, tspan[end])
        e_l2 = rel_error_l2(u1, u1_anal)
        println("rel_error_l2 = ", e_l2, "\n")
        @test e_l2 < 0.1
    end
end

# https://en.wikipedia.org/wiki/Heat_equation#Fundamental_solutions
@testset "MLP - heat equation - multi thread" begin # solving at one unique point
    for d in [1, 10, 50]
        println("test for d = ", d)
        tspan = (0e0, 5e-1)

        μ(x, p, t) = 0e0 # advection coefficients
        σ(x, p, t) = 1e-1 # diffusion coefficients

        u_anal(x, t) = sum(x.^2) + d * σ(0., 0., 0.)^2 * t
        g(x) = sum(x.^2)

        # d = 10
        x0 = fill(2e0,d)

        alg = MLP()

        f(y, z, v_y, v_z, ∇v_y, ∇v_z, p, t) = 0e0 .* v_y #TODO: this fix is not nice

        # defining the problem
        prob = PIDEProblem(g, f, μ, σ, tspan, x = x0)
        # solving
        xs,ts,sol = solve(prob, alg, multithreading = true)
        u1 = sol[end]
        u1_anal = u_anal(x0, tspan[end])
        e_l2 = rel_error_l2(u1, u1_anal)
        println("rel_error_l2 = ", e_l2, "\n")
        @test e_l2 < 0.1
    end
end

@testset "MLP - heat equation - single thread - Neumann BC" begin # solving at one unique point
    tspan = (0f0, 5f-1)
    μ(x, p, t) = 0f0 # advection coefficients
    σ(x, p, t) = 1f-1 #1f-1 # diffusion coefficients
    for d in [1,2,5]
        u1s = []
        for _ in 1:2
            neumann_bc = (fill(-5e-1, d), fill(5e-1, d))
            g(x) = sum(x.^2)
            # d = 10
            x0 = fill(3e-1,d)
            alg = MLP()
            f(y, z, v_y, v_z, ∇v_y, ∇v_z, p, t) = 0e0 .* v_y #TODO: this fix is not nice

            # defining the problem
            prob = PIDEProblem(g, f, μ, σ, tspan, x = x0, neumann_bc = neumann_bc)
            # solving
            xs,ts,sol = solve(prob, alg)
            push!(u1s, sol[end])
            println("d = $d, u1 = $(sol[end])")

        end
        e_l2 = mean(rel_error_l2.(u1s[1], u1s[2]))
        println("rel_error_l2 = ", e_l2, "\n")
        @test e_l2 < 0.1
    end
end

@testset "MLP - exponential - interval" begin
    for d in [1, 3, 10]
        println("test for d = ", d)
        tspan = (0f0, 5f-1)

        μ(x, p, t) = 0f0 # advection coefficients
        σ(x, p, t) = 0f0 # diffusion coefficients
        r = 1f-1
        x = fill(0f0, d)

        u_anal(x, t) = (g(x) * exp(r * t))
        g(x) = sum(x.^2) .+ 2f0

        alg = MLP()

        f(y, z, v_y, v_z, ∇v_y, ∇v_z, p, t) = r * v_y #TODO: this fix is not nice

        # defining the problem
        prob = PIDEProblem(g, f, μ, σ, tspan, x=x)
        # solving
        xs,ts,sol = solve(prob, alg)
        u1 = sol[end]
        u1_anal = u_anal.(xs, tspan[end])
        e_l2 = mean(rel_error_l2.(u1, u1_anal))
        println("rel_error_l2 = ", e_l2, "\n")
        @test e_l2 < 0.1
    end
end

@testset "MLP - allen cahn" begin

    tspan = (3f-1, 6f-1)

    μ(x, p, t) = 0f0 # advection coefficients
    σ(x, p, t) = 1f0 # diffusion coefficients

    for d in [10]

        alg = MLP()

        X0 = fill(0f0,d)  # initial point
        g(X) =  1f0 ./ (2f0 .+ 4f-1 * sum(X.^2))   # initial condition
        a(u) = u - u^3
        f(y, z, v_y, v_z, ∇v_y, ∇v_z, p, t) = - a.(v_y)

        # defining the problem
        prob = PIDEProblem(g, f, μ, σ, tspan, x = X0 )
        # solving
        xs,ts,sol = solve(prob, alg,)
        u1 = sol[end]
        # value coming from \cite{Beck2017a}
        e_l2 = rel_error_l2(u1, 0.30879)
        @test e_l2 < 0.5 # this is quite high as a relative error.
        println("d = $d, rel_error_l2 = $e_l2")
    end
end

@testset "MLP - allen cahn local, Neumann BC" begin

    tspan = (0f0, 5f-1)

    μ(x, p, t) = 0f0 # advection coefficients
    σ(x, p, t) = 1f0 # diffusion coefficients

    for d in [1,2,5]
        neumann = (fill(-5f-1, d), fill(5f-1, d))

        alg = MLP()

        X0 = fill(0f0,d)  # initial point
        g(X) = exp.(-0.25f0 * sum(X.^2))   # initial condition
        a(u) = u - u^3
        f(y, z, v_y, v_z, ∇v_y, ∇v_z, p, t) = a.(v_y) # nonlocal nonlinear part of the

        # defining the problem
        prob = PIDEProblem(g, f, μ, σ, tspan, x = X0, neumann = neumann )
        # solving
        xs,ts,sol = solve(prob, alg, )
        u1 = sol[end]
        @test !isnan(u1)
        println("d = $d, u1 = $(u1)")
    end
end

@testset "MLP algorithm - Black-Scholes Equation with Default Risk" begin
    tspan = (0f0, 1f0)

    μ(x, p, t) = 0.02f0 * x # advection coefficients
    σ(x, p, t) = 0.2f0 * x # diffusion coefficients

    d = 20

    alg = MLP()

    X0 = fill(100f0,d)  # initial point
    g(X) =  minimum(X) # initial condition

    δ = 2.0f0/3f0
    R = 0.02f0

    vh = 50.0f0
    vl = 70.0f0
    γh = 0.2f0
    γl = 0.02f0

    Q(u) = (u .< vh) .* γh .+ (u .>= vl) .*  γl .+ ( (u .>= vh) .* (u .< vl)) .* (((γh - γl) / (vh - vl)) * (u .- vh) .+ γh)

    µc = 0.02f0
    σc = 0.2f0

    f(y, z, v_y, v_z, ∇v_y, ∇v_z, p, t) = -(1f0 - δ) * Q.(v_y) .* v_y .- R * v_y

    # defining the problem
    prob = PIDEProblem(g, f, μ, σ, tspan, x = X0 )
    # solving
    xs,ts,sol = solve(prob, alg, )

    u1 = sol[end]

    analytical_ans = 60.781
    error_l2 = rel_error_l2(u1, analytical_ans)

    @test error_l2 < 0.1

    println("d = $d, error_l2 = $(error_l2)")

end

###################################################
########### NON LOCAL #############################
###################################################

@testset "MLP - replicator mutator" begin
    tspan = (0e0,5e-1)
    σ_sampling = 1e0
    μ(x, p, t) = 0e0 # advection coefficients
    σ(x, p, t) = 1e-1 #1f-1 # diffusion coefficients
    ss0 = 1e-2#std g0

    # Analytic sol
    function _SS(x, t, p)
        d = length(x)
        MM = σ(x, p, t) * ones(d)
        SSt = MM .* ((MM .* sinh.(MM *t) .+ ss0 .* cosh.( MM * t)) ./ (MM .* cosh.(MM * t ) .+ ss0 .* sinh.(MM * t)))
        return SSt
    end

    function uanal(x, t, p)
            d = length(x)
            return (2*π)^(-d/2) * prod(_SS(x, t, p) .^(-1/2)) * exp(-0.5 *sum(x .^2 ./ _SS(x, t, p)) )
    end

    sols = []
    xs = []
    for d in [1, 2, 5]

        x = fill(0e0, d)

        alg = MLP(M=4, K=10, L = 4, mc_sample = NormalSampling(σ_sampling) )

        g(x) = (2*π)^(-d/2) * ss0^(- d * 5e-1) * exp.(-5e-1 *sum(x .^2e0 / ss0)) # initial condition
        m(x) = - 5e-1 * sum(x.^2)

        f(y, z, v_y, v_z, ∇v_y, ∇v_z, p, t) = max(0.0, v_y) * (m(y) - max(0.0, v_z) * m(z) * (2.0 * π)^(d/2) * σ_sampling^d * exp(0.5 * sum(z.^2) / σ_sampling^2)) # nonlocal nonlinear part of the

        # defining the problem
        prob = PIDEProblem(g, f, μ, σ, tspan,x=x)
        # solving
        xgrid,ts,sol = solve(prob, alg,)
        u1 = sol[end]
        u1_anal = uanal(x,tspan[2],Dict())
        e_l2 = mean(rel_error_l2.(u1, u1_anal))
        println("rel_error_l2 = ", e_l2, "\n")
        @test e_l2 < 0.1
        push!(sols, sol[end])
    end
end

@testset "MLP- allen cahn non local Neumann BC" begin

    tspan = (0f0, 5f-1)

    μ(x, p, t) = 0f0 # advection coefficients
    σ(x, p, t) = 1f-1 #1f-1 # diffusion coefficients

    for d in [1,2,5]
        u1s = []
        for _ in 1:2
            neumann = (fill(-5f-1, d), fill(5f-1, d))

            alg = MLP(M=4, K=10, L = 4, mc_sample = UniformSampling(neumann...) )

            x = fill(0f0,d)  # initial point
            g(X) = exp.(-0.25f0 * sum(X.^2))   # initial condition
            a(u) = u - u^3
            f(y,z,v_y,v_z,∇v_y,∇v_z, p, t) = a.(v_y) .- a.(v_z) #.* Float32(π^(d/2)) * σ_sampling^d .* exp.(sum(z.^2, dims = 1) / σ_sampling^2) # nonlocal nonlinear part of the

            # defining the problem
            prob = PIDEProblem(g, f, μ, σ, tspan, x = x, neumann = neumann)
            # solving
            xs,ts,sol = solve(prob, alg)
            push!(u1s, sol[end])
            println("d = $d, u1 = $(sol[end])")

        end
        e_l2 = mean(rel_error_l2.(u1s[1], u1s[2]))
        println("rel_error_l2 = ", e_l2, "\n")
        @test e_l2 < 0.1
    end
end


using Interpolations

function getSpline(data; Δx=10)
    x = data[:,1]
    y = data[:,2]
    stepR = x[1]:(x[end] - x[1])/(size(x,1)-1):x[end]
    itp = interpolate(y, BSpline(Cubic(Line(OnGrid()))))
    sitp = scale(itp, stepR)

    ∇_list = collect(x[1]:(x[end] - x[end-1])/Δx:x[end])
    # ∇ = vcat([Interpolations.gradient(sitp, x) for x in ∇_list]...)
    sitp#, ∇_list, ∇
end

function plotSpline(corrType, dict; xmax = 20, mixWindows=true)
    suffix = mixWindows ? "_CorrelationDataTimeWindowsMixed/" : "_CorrelationData/"
    pathout = dict[:PathOut] * split(dict[:title], ".")[1] * suffix
    # files = readdir("/Users/javier/Desktop/Julia/Correlation/Ge00Sb00Te100T823K/Ge00Sb00Te100T823K_CorrelationDataTimeWindowsMixed/")
    files = readdir(pathout)
    hline(0)
    for file in files
        if split(file,"_")[1] == corrType
            dd = readdlm(pathout * file)
            stepR = dd[1,1]:(dd[end,1] - dd[end-1,1])/10:dd[xmax,1]
            sitp = getSpline(dd)
            plot!(stepR, x->sitp(x), label="BSpline $(file)", frame=:box, lw=3)
            plot!(dd[1:xmax,1], dd[1:xmax,2], label="Data $(file)", seriestype=:scatter, markerstrokewidth=0, ms=7)
        end
    end
    plot!(size=(500,700), xlabel="w", ylabel=corrType, margin = 5Plots.mm)
end

function getSplines(corrType, dict; xmax = 20, mixWindows=true)
    suffix = mixWindows ? "_CorrelationDataTimeWindowsMixed/" : "_CorrelationData/"
    pathout = dict[:PathOut] * split(dict[:title], ".")[1] * suffix
    # files = readdir("/Users/javier/Desktop/Julia/Correlation/Ge00Sb00Te100T823K/Ge00Sb00Te100T823K_CorrelationDataTimeWindowsMixed/")
    files = readdir(pathout)
    splineList = []
    # hline(0)
    for file in files
        if split(file,"_")[1] == corrType
            dd = readdlm(pathout * file)
            stepR = dd[1,1]:(dd[end,1] - dd[end-1,1])/10:dd[end,1]
            sitp = getSpline(dd)
            append!(splineList, (stepR, sitp))
            # plot!(stepR, x->sitp(x), label="BSpline $(file)", frame=:box, lw=3)
            # plot!(dd[1:xmax,1], dd[1:xmax,2], label="Data $(file)", seriestype=:scatter, markerstrokewidth=0, ms=7)
        end
    end
    # plot!(size=(500,700), xlabel="w", ylabel=corrType, margin = 5Plots.mm)
    splineList
end


#=
    Least diff problem in Julia + ConstraintSolver.jl

    The model solves the following problem:

    What is the smallest difference between two numbers X - Y
    if you must use all the digits (0..9) exactly once, i.e.
    Minimize the difference
        ABCDE - FGHIJ

    x:[5, 0, 1, 2, 3, 4, 9, 8, 7, 6]
    247 = 50123 - 49876
    0.603865 seconds (833.03 k allocations: 1.107 GiB, 7.89% gc time)

    Note: using diff as a separate variable is much slower.

=#
using ConstraintSolver, JuMP
using Cbc, GLPK
const CS = ConstraintSolver
include("constraints_utils.jl")

#
# 0.603865s
#
function least_diff()

    cbc_optimizer = optimizer_with_attributes(Cbc.Optimizer, "logLevel" => 0)
    glpk_optimizer = optimizer_with_attributes(GLPK.Optimizer)

    model = Model(optimizer_with_attributes(CS.Optimizer,   # "all_solutions"=> true,
                                                            # "all_optimal_solutions"=>true,
                                                            "logging"=>[],

                                                            # "traverse_strategy"=>:BFS,
                                                            "traverse_strategy"=>:DFS,
                                                            # "traverse_strategy"=>:DBFS,

                                                            # "branch_split"=>:Smallest,
                                                            # "branch_split"=>:Biggest,
                                                            "branch_split"=>:InHalf, # Best

                                                            # "simplify"=>false,
                                                            "simplify"=>true,

                                                            "time_limit"=>10,

                                                            # "lp_optimizer" => cbc_optimizer, # much slower (4.7s)
                                                            # "lp_optimizer" => glpk_optimizer, # (0.7s)
                                        ))

    @variable(model, 0 <= x[1:10] <= 9, Int)
    @variable(model, 10000 <= v[1:2] <= 99999, Int)
    # @variable(model, 0 <= diff <= 999, Int) # much slower

    a,b,c,d,e,f,g,h,i,j = x
    @constraint(model, x in CS.AllDifferentSet())
    @constraint(model, v[1] == 10000*a + 1000*b + 100*c + 10*d + e)
    @constraint(model, v[2] == 10000*f + 1000*g + 100*h + 10*i + j)

    # @constraint(model, diff == v[1] - v[2])  # much slower
    # @constraint(model, diff >= 1)
    @constraint(model, v[1] - v[2] >= 1)

    # @constraint(model, diff == 247) # Testing
    # @constraint(model, v[1] - v[2] == 247) # Testing

    # @objective(model, Min, diff)
    # This is much faster than working with a variable `diff`
    @objective(model, Min, v[1]-v[2])

    # Solve the problem
    println("solve")
    optimize!(model)

    status = JuMP.termination_status(model)
    println("status:$status")
    if status == MOI.OPTIMAL
        x = convert.(Integer,JuMP.value.(x))
        abcde = convert.(Integer,JuMP.value.(v[1]))
        fghij = convert.(Integer,JuMP.value.(v[2]))
        # diff = convert.(Integer,JuMP.value.(v[1]-v[2]))
        diff = abcde - fghij
        # diff = convert.(Integer,JuMP.value.(diff))

        println("x:$x")
        println("$diff = $abcde - $fghij")

        num_sols = MOI.get(model, MOI.ResultCount())
        println("\nnum_sols:$num_sols\n")

        # The solutions are in reverse order, i.e. the optimal is the first.
        for sol in num_sols:-1:1
            println("solution #$sol")
            # This works:
            vv = convert.(Integer,JuMP.value.(v,result=sol))
            println("vv:$vv diff: $(vv[1]-vv[2]) $(sol==1 ? """optimal!""" : "")")
        end
    end
end

least_diff()


using SparseArrays
using Mitosis
using LinearAlgebra, Statistics, Random

"""
Graph Laplacian of a `m×n` lattice.
"""
function laplacian(T, m, n)
    S = sparse(T(0.0)I, n*m, n*m)
    linear = LinearIndices((1:m, 1:n))
    for i in 1:m
        for j in 1:n
            for (i2, j2) in ((i + 1, j), (i, j + 1))
                if i2 <= m && j2 <= n
                    S[linear[i, j], linear[i2, j2]] -= 1
                    S[linear[i2, j2], linear[i, j]] -= 1

                    S[linear[i, j], linear[i, j]] += 1
                    S[linear[i2, j2], linear[i2, j2]] += 1
                end
            end
        end
    end
    S
end
# indices
m, n = 20, 20
C = CartesianIndices((m, n))
L = LinearIndices((m, n))

# tuning parameters
κ = 4 # κ-times integrated white noise, so κ is a smoothing parameter
ϵ = 0.1 # ϵ regularises
σ = 0.02 # σ the amount of white noise in the image, in absence of pixel noise regularises
λ = .2 # λ is tradeoff between exploration and exploitation in acquisition (this is just the upper confidence level acquisition function)

# data
W = [exp(-norm(([i,j]-[m÷3,n÷3]))^2/25.0) + σ*randn() for i in 1:m, j in 1:n]

# prior (power of the graph laplacian)
Λ = laplacian(Float32, m, n)
u = Gaussian{(:F,:Γ)}(zeros(m*n), (Λ + ϵ*I)^κ)

# Bayesian optimization
i = CartesianIndex(m÷2, n÷2)
is = [Tuple(i)]
for k in 1:11
    global i, u
    # observation scheme
    H = sparse(zeros(1, m*n))
    H[L[i]] = 1
    k = kernel(Gaussian; μ=AffineMap(H, [0.]), Σ=ConstantMap(Matrix(σ^2*I(1))))
    # update
    _, p = Mitosis.backward(BF(), k, [W[i]])
    # fusion
    _, u = fuse(u, p)
    # aquisition
    if m + n > 50
        _, i_ = findmax(mean(u) + 2λ./sqrt.(diag(u.Γ))) # proxy for inverse
    else
        _, i_ = findmax(mean(u) + λ*sqrt.(diag(inv(Matrix(u.Γ)))))
    end
    i = C[i_]
    push!(is, Tuple(i))
    println(i)
end

imax = C[findmax(mean(u))[2]] # estimate of the max
i = findmax(W)[2] # truth

# plot
using Makie
img(x) = heatmap(reshape(x, m, n))
p1 = img(mean(u))
scatter!(p1, first.(is), last.(is))
scatter!(p1, [i[1]], [i[2]], color=:green)
scatter!(p1, [imax[1]], [imax[2]], color=:red) # red estimate

p2 = img(W)
scatter!(p2, [i[1]], [i[2]], color=:green)
scatter!(p2, [imax[1]], [imax[2]], color=:red) # red estimate
hbox(p1, p2)


using JLD2
using FileIO

# tensor for classical 2-d model
const isingβc = log(1+sqrt(2))/2

"""
    model_tensor(model::Ising, β::Real)

return the  `MT <: HamiltonianModel` bulktensor at inverse temperature `β` for  two-dimensional
square lattice tensor-network.
"""
function model_tensor(model::Ising, β::Real)
    Ni, Nj = model.Ni, model.Nj
    ham = [-1. 1;1 -1]
    w = exp.(- β .* ham)
    wsq = sqrt(w)
    m = ein"ia,ib,ic,id -> abcd"(wsq, wsq, wsq, wsq)
    reshape([m for i=1:Ni*Nj], Ni, Nj)
end

function model_tensor(model::Ising22, β::Real)
    r = model.r
    ham = [-1. 1;1 -1]
    w1 = exp.(- β .* ham)
    w2 = exp.(- r * β .* ham)
    wsq1 = sqrt(w1)
    wsq2 = sqrt(w2)
    M11 = ein"ia,ib,ic,id -> abcd"(wsq1, wsq2, wsq2, wsq1)
    M12 = permutedims(M11, [3,2,1,4])
    M21 = permutedims(M11, [1,4,3,2])
    M22 = permutedims(M11, [3,4,1,2])
    reshape([M11,M21,M12,M22], 2, 2)
end

########## To do: correct Ising33 ##########
function model_tensor(model::Ising33, β::Real)
    r = model.r
    ham = [-1. 1;1 -1]
    w1 = exp.(- β .* ham)
    w2 = exp.(- r * β .* ham)
    wsq1 = sqrt(w1)
    wsq2 = sqrt(w2)
    M11 = ein"ia,ib,ic,id -> abcd"(wsq1, wsq2, wsq2, wsq1)
    M12 = permutedims(M11, [3,2,1,4])
    M21 = permutedims(M11, [1,4,3,2])
    M22 = permutedims(M11, [3,4,1,2])
    reshape([M11,M21,M12,M22], 2, 2)
end

"""
    mag_tensor(::MT, β::Real)

return the  `MT <: HamiltonianModel` the operator for the magnetisation at inverse temperature `β` for a two-dimensional
square lattice tensor-network.
"""
function mag_tensor(model::Ising, β::Real)
    Ni,Nj = model.Ni, model.Nj
    a = reshape(Float64[1 0 0 0; 0 0 0 0; 0 0 0 0; 0 0 0 -1] , 2,2,2,2)
    cβ, sβ = sqrt(cosh(β)), sqrt(sinh(β))
    q = 1/sqrt(2) * [cβ+sβ cβ-sβ; cβ-sβ cβ+sβ]
    m = ein"abcd,ai,bj,ck,dl -> ijkl"(a,q,q,q,q)
    reshape([m for i=1:Ni*Nj], Ni, Nj)
end

function mag_tensor(model::Ising22, β::Real)
    r = model.r
    a = reshape(Float64[1 0 0 0; 0 0 0 0; 0 0 0 0; 0 0 0 -1] , 2,2,2,2)
    ham = [-1. 1;1 -1]
    w1 = exp.(- β .* ham)
    w2 = exp.(- r * β .* ham)
    wsq1 = sqrt(w1)
    wsq2 = sqrt(w2)
    M11 = ein"ijkl,ia,jb,kc,ld -> abcd"(a,wsq1,wsq2,wsq2,wsq1)
    M12 = permutedims(M11, [3,2,1,4])
    M21 = permutedims(M11, [1,4,3,2])
    M22 = permutedims(M11, [3,4,1,2])
    reshape([M11,M21,M12,M22], 2, 2)
end

"""
    energy_tensor(model::Ising, β::Real)

return the  `MT <: HamiltonianModel` the operator for the energy at inverse temperature `β` for a two-dimensional
    square lattice tensor-network.
"""
function energy_tensor(model::Ising, β::Real)
    Ni,Nj = model.Ni, model.Nj
    ham = [-1 1;1 -1]
    w = exp.(-β .* ham)
    we = ham .* w
    wsq = sqrt(w)
    wsqi = wsq^(-1)
    e = (ein"ai,im,bm,cm,dm -> abcd"(wsqi,we,wsq,wsq,wsq) + ein"am,bi,im,cm,dm -> abcd"(wsq,wsqi,we,wsq,wsq) +
        ein"am,bm,ci,im,dm -> abcd"(wsq,wsq,wsqi,we,wsq) + ein"am,bm,cm,di,im -> abcd"(wsq,wsq,wsq,wsqi,we)) ./ 2
    reshape([e for i=1:Ni*Nj], Ni, Nj)
end

function energy_tensor(model::Ising22, β::Real)
    r = model.r
    ham = [-1. 1;1 -1]
    w1 = exp.(- β .* ham)
    w2 = exp.(- r * β .* ham)
    we1 = ham .* w1
    we2 = r .* ham .* w2
    wsq1 = sqrt(w1)
    wsq2 = sqrt(w2)
    wsqi1 = wsq1^(-1)
    wsqi2 = wsq2^(-1)

    M = Array{Array,2}(undef, 2, 2)
    M11 = (ein"ai,im,bm,cm,dm -> abcd"(wsqi1,we1,wsq2,wsq2,wsq1) +
              ein"am,bi,im,cm,dm -> abcd"(wsq1,wsqi2,we2,wsq2,wsq1) +
              ein"am,bm,ci,im,dm -> abcd"(wsq1,wsq2,wsqi2,we2,wsq1) +
              ein"am,bm,cm,di,im -> abcd"(wsq1,wsq2,wsq2,wsqi1,we1)) ./ 2
    M12 = permutedims(M11, [3,2,1,4])
    M21 = permutedims(M11, [1,4,3,2])
    M22 = permutedims(M11, [3,4,1,2])
    reshape([M11,M21,M12,M22], 2, 2)
end


# SPDX-License-Identifier: MIT

# Helper functions for plotting 2D-Normal pdfs.
# Lucas Ondel, 2021

function normalShape(μ, Σ, σ = 1)
    θ = LinRange(0, 2π, 1000)
    λ, U = eigen(Σ)
    B = U * diagm(σ * sqrt.(λ))
    circle = hcat(sin.(θ), cos.(θ))'
    contour = B * circle .+ μ
    contour[1, :], contour[2, :]
end

function plotnormal!(μ::AbstractVector, Σ::AbstractMatrix; σ = 1, kwargs...)
    plot!(normalShape(μ, Σ, σ), seriestype = [:shape];  kwargs...)
end


using Devectorize

if !isdefined(:__EXPRESSION_HASHES__)
    __EXPRESSION_HASHES__ = Set{UInt64}()
end

macro runonce(expr)
    h = hash(expr)
    return esc(quote
        if !in($h, __EXPRESSION_HASHES__)
            push!(__EXPRESSION_HASHES__, $h)
            $expr
        end
    end)
end

"""
  Defines the type `Mat{T}` as a convenience wrapper around the
  `AbstractArray{T,2}` type.
"""
typealias Mat{T} AbstractArray{T, 2}

"""
  Defines the type `Vec{T}` as a convenience wrapper around the
  `AbstractArray{T,1}` type.
"""
typealias Vec{T} AbstractArray{T, 1}

"""
  # Boltzmann.normalize_samples (utils.jl)
  ## Function Calls
    `normalize_samples(X::Mat{Float64})`

  ## Description
    Given a matrix, `X`, assume that each column represents a different
    data sample and that each row represents a different feature. In this
    case, `normalize_samples` will normalize each individual sample to
    the range `[0,1]` according to the minimum and maximum features in the
    sample.

  ## Returns
    1. `::Mat{Float64}`, The normalized dataset.

  ### See also...
    `normalize_samples!`
"""
function normalize_samples(x::Mat{Float64})
    samples = size(x,2)
    y = zeros(x)

    for i=1:samples
      s = x[:,i]
      mins = minimum(s)
      maxs = maximum(s)
      rans = maxs-mins
      y[:,i] = (s-mins)/rans
    end

    return y
end

"""
  # Boltzmann.normalize_samples! (utils.jl)
  ## Function Calls
    `normalize_samples!(x::Mat{Float64})`

  ## Description
    Given a matrix, `x`, assume that each column represents a different
    data sample and that each row represents a different feature. In this
    case, `normalize_samples` will normalize each individual sample to
    the range `[0,1]` according to the minimum and maximum features in the
    sample.

  ## Returns
    Nothing, modifies `x` in place.

  ### See also...
    `normalize_samples`
"""
function normalize_samples!(x::Mat{Float64})
    samples = size(x,2)

    for i=1:samples
      s = x[:,i]
      mins = minimum(s)
      maxs = maximum(s)
      rans = maxs-mins
      x[:,i] = (s-mins)/rans
    end
end

"""
  # Boltzmann.normalize (utils.jl)
  ## Function Calls
    `normalize(x::Mat{Float64})`
    `normalize(x::Vec{Float64})`

  ## Description
    Given an array, `x`, normalize the entire array to the range
    `[0,1]` according to the maximum and miniumum values of the array.

  ## Returns
    1. `::Mat{Float64}` *or* `::Vec{Float64}` depending on the input.

  ### See also...
    `normalize!`
"""
function normalize(x::Mat{Float64})
    y = zeros(x)
    minx = minimum(x)
    maxx = maximum(x)
    ranx = maxx-minx

    # We can index a two-dimension array as if it were a vector.
    @simd for i=1:length(x)
      @inbounds y[i] = (x[i]-minx) / ranx
    end

    return y
end
function normalize(x::Vec{Float64})
    y = zeros(x)
    minx = minimum(x)
    maxx = maximum(x)
    ranx = maxx-minx

    @simd for i=1:length(x)
      @inbounds y[i] = (x[i]-minx) / ranx
    end

    return y
end

"""
  # Boltzmann.normalize! (utils.jl)
  ## Function Calls
    `normalize!(x::Mat{Float64})`
    `normalize!(x::Vec{Float64})`

  ## Description
    Given an array, `x`, normalize the entire array to the range
    `[0,1]` according to the maximum and miniumum values of the array.

  ## Returns
    Nothing. Modifies `x` in place.

  ### See also...
    `normalize`
"""
function normalize!(x::Mat{Float64})
    minx = minimum(x)
    maxx = maximum(x)
    ranx = maxx-minx

    @simd for i=1:length(x)
      @inbounds x[i] = (x[i]-minx) / ranx
    end
end
function normalize!(x::Vec{Float64})
    minx = minimum(x)
    maxx = maximum(x)
    ranx = maxx-minx

    @simd for i=1:length(x)
      @inbounds x[i] = (x[i]-minx) / ranx
    end
end

"""
  # Boltzmann.binarize (utils.jl)
  ## Function Calls
    `binarize(x::Mat{Float64}[,threshold=0.0])`
    `binarize(x::Vec{Float64}[,threshold=0.0])`

  ## Description
    Given an array, `x`, assign each element of the array to
    either `0` or `1` depending on specified value of `threshold`.
    This is done according to the following rule...
    ```
        if element <= threshold: element = 0
        if element >  threshold: element = 1
    ```

  ## Returns
    1. `::Mat{Float64}` *or* `::Vec{Float64}`, depending on input.

  ### See also...
    `binarize!`
"""
function binarize(x::Mat{Float64};threshold=0.0)
  s = copy(x)
  @simd for i=1:length(x)
    @inbounds s[i] = x[i] > threshold ? 1.0 : 0.0
  end
  return s
end
function binarize(x::Vec{Float64};threshold=0.0)
  s = copy(x)
  @simd for i=1:length(x)
    @inbounds s[i] = x[i] > threshold ? 1.0 : 0.0
  end
  return s
end

"""
  # Boltzmann.binarize! (utils.jl)
  ## Function Calls
    `binarize!(x::Mat{Float64}[,threshold=0.0])`
    `binarize!(x::Vec{Float64}[,threshold=0.0])`

  ## Description
    Given an array, `x`, assign each element of the array to
    either `0` or `1` depending on specified value of `threshold`.
    This is done according to the following rule...
    ```
        if element <= threshold: element = 0
        if element >  threshold: element = 1
    ```

  ## Returns
    Nothing. Modifies `x` in place.

  ### See also...
    `binarize`
"""
function binarize!(x::Mat{Float64};threshold=0.0)
  @simd for i=1:length(x)
    @inbounds x[i] = x[i] > threshold ? 1.0 : 0.0
  end
end
function binarize!(x::Vec{Float64};threshold=0.0)
  @simd for i=1:length(x)
    @inbounds x[i] = x[i] > threshold ? 1.0 : 0.0
  end
end

"""
  # Boltzmann.logsig (utils.jl)
  ## Function Calls
    `logsig(x::Mat{Float64})`
    `logsig(x::Vec{Float64})`

  ## Description
    Perform the logistic sigmoid element-by-element for the entries
    of `x`. Here, the logistic sigmoid is defined as
                      `y  =  1 / (1+e^-x)`.

  ## Returns
    1. `::Mat{Float64}` *or* `::Vec{Float64}` (depending on input).

  ### See also...
    `logsig!`
"""
function logsig(x::Mat{Float64})
    @devec s = 1 ./ (1 + exp(-x))
    return s
end
function logsig(x::Vec{Float64})
    @devec s = 1 ./ (1 + exp(-x))
    return s
end

"""
  # Boltzmann.logsig! (utils.jl)
  ## Function Calls
    `logsig!(x::Mat{Float64})`
    `logsig!(x::Vec{Float64})`

  ## Description
    Perform the logistic sigmoid element-by-element for the entries
    of `x`. Here, the logistic sigmoid is defined as
                      `y  =  1 / (1+e^-x)`.

  ## Returns
    Nothing. Modifies `x` in place.

  ### See also...
    `logsig`
"""
function logsig!(x::Mat{Float64})
  # We would like to have used `@devec`, here, but this macro by default
  # makes a new instantiation in memory and thus the result is distinct
  # from the original.
  @simd for i=1:length(x)
    @inbounds x[i] = 1 ./ (1 + exp(-x[i]))
  end
end
function logsig!(x::Vec{Float64})
  @simd for i=1:length(x)
    @inbounds x[i] = 1 ./ (1 + exp(-x[i]))
  end
end

"""
  # Boltzmann.random_samplse (utils.jl)
  ## Function Calls
    `random_columns(x::Mat{Float64})`
  ## Description
    Select a random permutation of the columsn of the specified
    matrix, returning the permuted sub-matrix as well as the
    column indices selected from the original matrix.

  ## Returns
    1. `::Mat{Float64}`, the random sub-matrix.
    2. `::Vec{Int}`, the array of selected column indices
"""
function random_columns(x::Mat{Float64},n_select::Int)
    n_features = size(x,1)
    n_columns = size(x,2)
    perm = shuffle!(collect(1:n_columns))[1:n_select]
    y = x[:,perm]

    return y, perm
end