using FastTransforms
using Base.Test

println("Testing special functions")
n = 0:1000_000
λ = 0.123
@time FastTransforms.Cnλ(n,λ);
@time FastTransforms.Cnλ(n,λ);

x = linspace(0,20,81);
@test norm((FastTransforms.Λ(x)-FastTransforms.Λ(big(x)))./FastTransforms.Λ(big(x)),Inf) < 2eps()

x = 0:0.5:10_000
λ₁,λ₂ = 0.125,0.875
@test norm((FastTransforms.Λ(x,λ₁,λ₂)-FastTransforms.Λ(big(x),big(λ₁),big(λ₂)))./FastTransforms.Λ(big(x),big(λ₁),big(λ₂)),Inf) < 4eps()
λ₁,λ₂ = 1//3,2//3
@test norm((FastTransforms.Λ(x,Float64(λ₁),Float64(λ₂))-FastTransforms.Λ(big(x),big(λ₁),big(λ₂)))./FastTransforms.Λ(big(x),big(λ₁),big(λ₂)),Inf) < 4eps()

n = 0:1000
α = 0.125
β = 0.375

@time FastTransforms.Cnαβ(n,α,β);
@test norm(FastTransforms.Cnαβ(n,α,β)./FastTransforms.Cnαβ(n,big(α),big(β))-1,Inf) < 3eps()

@time FastTransforms.Anαβ(n,α,β);
@test norm(FastTransforms.Anαβ(n,α,β)./FastTransforms.Anαβ(n,big(α),big(β))-1,Inf) < 4eps()

N = 20
f(x) = exp(x)

x,w = FastTransforms.fejer1(N,0.,0.)
@test norm(dot(f(x),w)-2sinh(1)) ≤ 4eps()
x,w = FastTransforms.fejer2(N,0.,0.)
@test norm(dot(f(x),w)-2sinh(1)) ≤ 4eps()
x,w = FastTransforms.clenshawcurtis(N,0.,0.)
@test norm(dot(f(x),w)-2sinh(1)) ≤ 4eps()

#=
x = Fun(identity)
g = f(x)*(1-x)^.25*(1+x)^.35
val = sum(g)
=#

x,w = FastTransforms.fejer1(N,0.25,0.35)
@test norm(dot(f(x),w)-2.0351088204147243) ≤ 4eps()
x,w = FastTransforms.fejer2(N,0.25,0.35)
@test norm(dot(f(x),w)-2.0351088204147243) ≤ 4eps()
x,w = FastTransforms.clenshawcurtis(N,0.25,0.35)
@test norm(dot(f(x),w)-2.0351088204147243) ≤ 4eps()

println("Testing the Chebyshev–Jacobi transform")

Nr = 5
v = zeros(Nr)
Na,Nb = 5,5
V = zeros(Na,Nb)

for N in round(Int,logspace(1,3,3))
    println("")
    println("N = ",N)
    println("")
    for αi=1:Na, βi=1:Nb
        α,β = -0.45+(αi-1)/Na,-0.45+(βi-1)/Nb
        p1,p2 = plan_cjt(rand(N),α,β),plan_icjt(rand(N),α,β)
        for i=1:Nr
            c = rand(N)
            v[i] = log(norm(p2*(p1*c)-c,Inf)/(20N^(1+2max(α,β))*eps()))
        end
        mean(v) > 2 && println("This is the mean log(||Error||_∞/Estimate) that broke the test: ",mean(v)," and the standard deviation: ",std(v))
        V[αi,βi] = mean(v)
    end
    println("This is the mean log(||Error||_∞/Estimate): ",mean(V)," and the standard deviation: ",std(V))
end

println("Testing the special cases length(c) = 0,1,2:")

c = [1.5]
@test cjt(c,0.12,0.34) == c
@test icjt(c,0.12,0.34) == c

pop!(c)

@test cjt(c,0.12,0.34) == c
@test icjt(c,0.12,0.34) == c

c = [1.0;2.0]
@test norm(jjt(c,0.12,0.34,0.12,0.34)-c,Inf) ≤ 2eps()

# test same parameters

c=[0.,1.]
@test jjt(c,1.5,-0.5,1.5,-0.5) == c

println("Testing the special cases (α,β) = (±0.5,±0.5)")

# Chebyshev coefficients of exp
c_cheb = [1.2660658777520084,1.13031820798497,0.27149533953407656,0.044336849848663804,0.005474240442093732,0.0005429263119139438,4.497732295429515e-5,3.1984364624019905e-6,1.9921248066727958e-7,1.1036771725517344e-8,5.505896079673747e-10,2.4979566169849825e-11,1.03915223067857e-12,3.9912633564144015e-14,1.4237580108256572e-15,4.740926102561496e-17,1.4801800572082976e-18,4.3499194949441696e-20,1.2074289272797528e-21,3.175356737059445e-23,7.93367197163804e-25,1.8879484042289162e-26,4.2886738765925875e-28,9.318985281777577e-30,1.9406469749017399e-31,3.8798022492260015e-33,7.45850288739152e-35,1.3807477824110646e-36,2.4648623717710964e-38,4.2485421925059915e-40,7.079001176212895e-42,1.141486778254089e-43,1.783151037543284e-45,2.7011422638738024e-47,3.9714338657415316e-49,5.672351695612112e-51,7.876788130533199e-53,1.0642416282087358e-54,1.4000817885528905e-56,1.794689085390544e-58,2.243019482662972e-60,2.7349926016641967e-62,3.2554929833033215e-64,3.78495689402033e-66,4.300544446926769e-68,4.77780574839433e-70,5.192665133046628e-72,5.523442810988436e-74,5.717985746408847e-76]

# Jacobi coefficients of exp at (±0.5,±0.5)
c_00 = [1.2660658777520084,2.26063641596994,0.723987572090871,0.14187791951572415,0.020020079331085648,0.0022061767595233264,0.00019937999439479753,1.5268992715616027e-5,1.014420290055232e-6,5.9506859054175615e-8,3.124851418974203e-9,1.4852160704995775e-10,6.447143371527455e-12,2.5753240175931934e-13,9.526907336854026e-15,3.281725133829881e-16,1.0576496993856426e-17,3.20238467156949e-19,9.14299140846687e-21,2.4694550834939014e-22,6.3281713790321685e-24,1.542622093731937e-25,3.5857224408995815e-27,7.964666406806272e-29,1.6939043515668796e-30,3.4556190450066486e-32,6.773312266576111e-34,1.2775611026088504e-35,2.3221236341654026e-37,4.072731101981327e-39,6.901079409575277e-41,1.1310395206213411e-42,1.794875997727078e-44,2.760732681473006e-46,4.119631242417262e-48,5.969295899580336e-50,8.405881613439669e-52,1.1512859570474932e-53,1.5347890791678689e-55,1.9929132418968355e-57,2.522290550037395e-59,3.1134869971034086e-61,3.7506700259829663e-63,4.411970068315684e-65,5.0705898146983576e-67,5.696604001638848e-69,6.259278760298182e-71,6.729592140384465e-73,7.039950057933299e-75]
c_01 = [1.8312249817444934,1.4018135475190467,0.4211095858436539,0.07969774446521205,0.011002819207328322,0.0011944708771926122,0.00010677917238679192,8.110005868771403e-6,5.353106062476433e-7,3.1237733951100556e-8,1.6333110219427631e-9,7.735005972383578e-11,3.3473853404557083e-12,1.333595312027857e-13,4.92207038322879e-15,1.6920924742289326e-16,5.443658341283791e-18,1.6456374329092366e-19,4.691719175403533e-21,1.2655773772197326e-22,3.2393803393291905e-24,7.888321906112733e-26,1.8318188278743866e-27,4.065263937281931e-29,8.638847091039724e-31,1.7610248038105726e-32,3.449351261471638e-34,6.50183837007916e-36,1.1810743750838164e-37,2.070295858087742e-39,3.506179552173043e-41,5.743539156119835e-43,9.110325158556408e-45,1.4006615828866467e-46,2.0892357224279913e-48,3.026093616078655e-50,4.2597272086552854e-52,5.832159509538512e-54,7.772313549445791e-56,1.0089104305392274e-57,1.2765228632360985e-59,1.5752735945134057e-61,1.8971383534453687e-63,2.2310498815373167e-65,2.563461449357282e-67,2.8792582164274203e-69,3.1629293838222696e-71,3.399629156416382e-73,3.5199750289666496e-75]
c_10 = [0.7009067737595234,0.8588228684508934,0.30287798624721707,0.062180175050512106,0.009017260123757326,0.0010117058823307142,9.260082200800561e-5,7.1589868468446235e-6,4.791096838075886e-7,2.8269125103075062e-8,1.49154039703144e-9,7.117154732612197e-11,3.099758031071747e-12,1.2417287055653363e-13,4.604836953625235e-15,1.5896326596009483e-16,5.1328386525726346e-18,1.5567472386602532e-19,4.451272233063337e-21,1.2038777062741688e-22,3.088791039702978e-24,7.537899031206638e-26,1.753903613025195e-27,3.899402469524341e-29,8.300196424629072e-31,1.694594241196076e-32,3.3239610051044727e-34,6.273772656009344e-36,1.1410492590815863e-37,2.002435243893585e-39,3.3948998574022335e-41,5.566856050093576e-43,8.838434818714371e-45,1.3600710985863591e-46,2.0303955199892708e-48,2.943202283501681e-50,4.146154404784383e-52,5.68070006093642e-54,7.575577242232898e-56,9.84002811357608e-58,1.2457676868012966e-59,1.5382134025900027e-61,1.8535316725375976e-63,2.1809201867783675e-65,2.5071283653410757e-67,2.8173457852114277e-69,3.0963493764759127e-71,3.3299629839680836e-73,3.5199750289666496e-75]
c_11 = [1.13031820798497,0.7239875720908708,0.21281687927358628,0.040040158662171296,0.005515441898808317,0.0005981399831843923,5.34414745046561e-5,4.057681160220927e-6,2.677808657437903e-7,1.5624257094871016e-8,8.168688387747676e-10,3.8682860229164727e-11,1.6739606114355753e-12,6.668835135797819e-14,2.4612938503724105e-15,8.46119759508514e-17,2.722026970834067e-18,8.228692267620182e-20,2.3459823293192055e-21,6.328171379032168e-23,1.619753198418534e-24,3.944294684989539e-26,9.159366367827211e-28,2.0326852218802552e-29,4.31952380625831e-31,8.805305946548944e-33,1.7247074885219479e-34,3.2509730878315635e-36,5.905460097872925e-38,1.0351619114362916e-39,1.7531112569630786e-41,2.8718015963633246e-43,4.5552089244304594e-45,7.003373112109346e-47,1.044626782426559e-48,1.5130586904191403e-50,2.129879020537863e-52,2.916099250418951e-54,3.8861808216988285e-56,5.04458110007479e-58,6.382648344061989e-60,7.876407054564243e-62,9.485735646882059e-64,1.1155297592778338e-65,1.2817359024131178e-67,1.4396345613848593e-69,1.5814715428162693e-71,1.700107488097128e-73,1.7781317156635653e-75]

@test norm(cjt(c_00,-0.5,-0.5)-c_cheb,Inf) < eps()
@test norm(cjt(c_01,-0.5,0.5)-c_cheb,Inf) < eps()
@test norm(cjt(c_10,0.5,-0.5)-c_cheb,Inf) < eps()
@test norm(cjt(c_11,0.5,0.5)-c_cheb,Inf) < eps()

@test norm(icjt(c_cheb,-0.5,-0.5)-c_00,Inf) < eps()
@test norm(icjt(c_cheb,-0.5,0.5)-c_01,Inf) < eps()
@test norm(icjt(c_cheb,0.5,-0.5)-c_10,Inf) < eps()
@test norm(icjt(c_cheb,0.5,0.5)-c_11,Inf) < eps()

c = exp(-collect(1:1000)./30)

println("Testing increment/decrement operators for α,β ≤ -0.5")

α,β = -0.9,-0.6

@test norm(FastTransforms.fromsquare!(FastTransforms.tosquare!(copy(c),α,β),α,β)-c,Inf) < 10eps()
@test norm(icjt(cjt(c,α,β),α,β)-c,Inf) < 2000eps()
p1,p2 = plan_cjt(c,α,β),plan_icjt(c,α,β)
@test norm(p2*(p1*c)-c,Inf) < 2000eps()

println("Testing increment/decrement operators for α ≤ -0.5, β > -0.5")

α,β = -0.9,3.1

@test norm(FastTransforms.fromsquare!(FastTransforms.tosquare!(copy(c),α,β),α,β)-c,Inf) < 10eps()
@test norm(icjt(cjt(c,α,β),α,β)-c,Inf) < 2000eps()
p1,p2 = plan_cjt(c,α,β),plan_icjt(c,α,β)
@test norm(p2*(p1*c)-c,Inf) < 2000eps()

println("Testing increment/decrement operators for α > -0.5, β ≤ -0.5")

α,β = 2.4,-0.75

@test norm(FastTransforms.fromsquare!(FastTransforms.tosquare!(copy(c),α,β),α,β)-c,Inf) < 1200eps()
@test norm(icjt(cjt(c,α,β),α,β)-c,Inf) < 1e6eps()
p1,p2 = plan_cjt(c,α,β),plan_icjt(c,α,β)
@test norm(p2*(p1*c)-c,Inf) < 1e6eps()

println("Testing increment/decrement operators for α > -0.5, β > -0.5")

α,β = 2.4,3.6

@test norm(FastTransforms.fromsquare!(FastTransforms.tosquare!(copy(c),α,β),α,β)-c,Inf) < 250eps()
@test norm(icjt(cjt(c,α,β),α,β)-c,Inf) < 2e5eps()
p1,p2 = plan_cjt(c,α,β),plan_icjt(c,α,β)
@test norm(p2*(p1*c)-c,Inf) < 2e5eps()

println("Testing for complex coefficients")

α,β = 0.12,0.34
c = complex(rand(100),rand(100))

@test cjt(c,α,β) == complex(cjt(real(c),α,β),cjt(imag(c),α,β))
@test icjt(c,α,β) == complex(icjt(real(c),α,β),icjt(imag(c),α,β))
@test jjt(c,α,β,α,β) == complex(jjt(real(c),α,β,α,β),jjt(imag(c),α,β,α,β))
@test norm(jjt(c,α,β,α,β)-c,Inf) < 200eps()

println("Testing for Vector{Float32}")

c64 = rand(100)
c32 = map(Float32,c64)

cL64 = cjt(c64,0.,0.)
cL32 = cjt(c32,0.f0,0.f0)

@test norm(cL32-cL64,Inf) < 20eps(Float32)

println("Testing for Matrix of coefficients")

c = rand(100,100)
@test maxabs(jjt(c,α,β,α,β)-c) < 10000eps()

println("Testing Gaunt coefficients")

include("gaunttest.jl")

println("Testing BigFloat support of FFT and DCT methods")

include("fftBigFloattest.jl")

println("Testing equivalence of CXN and ASY methods")

for k in round(Int,logspace(1,4,20))
    r = randn(k)./√(1:k) # Proven 𝒪(√(log N)) error for ASY method.
    @test_approx_eq leg2cheb(r) cjt(r,0.,0.)
end

@test_approx_eq leg2chebu([1.0,2,3,4,5])  [0.546875,0.5,0.5390625,1.25,1.3671875]

c = randn(1000)./√(1:1000);

@test_approx_eq leg2cheb(cheb2leg(c)) c
@test_approx_eq cheb2leg(leg2cheb(c)) c

@test norm(jac2jac(c,0.,√2/2,-1/4,√2/2)-jjt(c,0.,√2/2,-1/4,√2/2),Inf) < 10length(c)*eps()

@test norm(ultra2ultra(ultra2ultra(c,.5,.75),.75,.5)-c,Inf) < 10length(c)*eps()


@testset "Misc/MPolyFactor" begin

  check_fac(esum, a) = check_fac(esum, esum, a)

  function check_fac(low, hi, a)
      f = factor(a)
      g = f.unit
      esum = 0
      for i in f.fac
          g *= i[1]^i[2]
          esum += i[2]
      end
      @test g == a
      @test low <= esum <= hi
  end

  QA, A = PolynomialRing(Hecke.QQ, "A")
  K, a = number_field(A^3 - 2, "a")
  Kxyz, (x, y, z, t) = PolynomialRing(K, ["x", "y", "z", "t"])

  check_fac(0, Kxyz(0))
  check_fac(0, Kxyz(1))
  check_fac(0, Kxyz(2))
  check_fac(6, 2*x^2*y*(1+y)*(x^3-2))
  check_fac(9, 2*x^2*y*(1+y)*(2*x^3-1)*(a*x+a^2*y+1)*(a^2*x+a*y+1)^2)
  check_fac(2, x^6-2*y^3)
  check_fac(2, (3*a*y*z*x^2+z-a)*(7*x+a*z^3*y))
  check_fac(3, (x+y+a)*(x+a*y+1)*(x+y+2*a))
  check_fac(3, (x^2+y^2+a)*(x^3+a*y^3+1)*(x^2+y+2*a))
  check_fac(8, (1+x+a*y+x*y)^3*(1+x+a^2*y+x*y^2)^5)
  check_fac(4, (x-(a+1)^4*y+z)*(a*x+y+z)*(x+a*y+z)*(x+y+a*z))
  check_fac(4, (1+x+y+x^2*(1+z+a*y)+y^2*(z+a*x^2))*
               (a*x+y+z*y)*(1+x+a*y*z+z)*(a+x*y+y+a*z))
  check_fac(2, x^99-2*y^99*z^33)
  check_fac(4, x^66-y^66*z^33)

  function dodd(n::Int)
    count = 0
    for i in 1:2:n
      if n%i == 0
        count += 1
      end
    end
    return count
  end
  for i in 1:10
    check_fac(1 + dodd(i) + Int(i%3 == 0), ((1+x+y+z+t)^i+1)*((1+x+y+z+t)^i+2))
  end

  QA, A = PolynomialRing(Hecke.QQ, "A")
  K,e3 = number_field(A^2+A+1, "e3")
  R,(a,x,y,z,t,u,v,w) = PolynomialRing(K, ["a", "x", "y", "z", "t", "u", "v", "w"])
  check_fac(24,
    3*(x+y)*(x-y)*(x-e3*y)*(x+e3*y)*(x+(-e3-1)*y)*(x+(e3+1)*y)*
    (a^2*y*z^2+e3*x*t*u*v*w)*(a^2*y*u^2+e3*x*z*t*v*w)*(a^2*y*t^2-e3*x*z*u*v*w)*
    (a^2*y*t^2+e3*x*z*u*v*w)*(a^2*y*z^2-e3*x*t*u*v*w)*(a^2*y*u^2-e3*x*z*t*v*w)*
    (a*x*z*t*v+(e3+1)*y*u^2*w^2)*(a*x*t*u*v+(e3+1)*y*z^2*w^2)*
    (a*x*z*t*v+(-e3-1)*y*u^2*w^2)*(a*x*z*u*w+(-e3-1)*y*t^2*v^2)*
    (a*x*z*u*w+(e3+1)*y*t^2*v^2)*(a*x*z*t*w+(e3+1)*y*u^2*v^2)*
    (a*x*z*t*w+(-e3-1)*y*u^2*v^2)*(a*x*z*u*v+(-e3-1)*y*t^2*w^2)*
    (a*x*t*u*w+(e3+1)*y*z^2*v^2)*(a*x*z*u*v+(e3+1)*y*t^2*w^2)*
    (a*x*t*u*v+(-e3-1)*y*z^2*w^2)*(a*x*t*u*w+(-e3-1)*y*z^2*v^2))

end


using JuMP

m = Model()

# ----- Variables ----- #
@variable(m, objvar)
x_Idx = Any[1, 2]
@variable(m, x[x_Idx])
set_lower_bound(x[1], -10.0)
set_upper_bound(x[1], 10.0)
set_lower_bound(x[2], -10.0)
set_upper_bound(x[2], 10.0)

# ----- Constraints ----- #
@NLconstraint(m, e1, -((-5+sqrt( (x[1])^2))*(-5+sqrt( (x[1])^2))+(-5+sqrt( (x[2])^2))*(-5+sqrt( (x[2])^2)))+objvar == 0.0)

# ----- Objective ----- #
@objective(m, Min, objvar)

m = m 		 # model get returned when including this script.


using Statistics
using Random
using StatsBase
using LinearAlgebra

struct ConstellationHistogram{T1,T2,T3,T4,T5}
    nbins::Int
    dopplers::Array{Float64,1}
    doppler_rates::Array{Float64,1}
    histogram::T1
    weights::T2
    fd_rate::Vector{Float64}
    f_d::Vector{Float64}
    elevations::T3
    elev_hists::T4
    max_elevations::T5
end

function get_histogram(x; nbins)
    return normalize(fit(Histogram, x, nbins=nbins), mode=:probability)
end

function initiate_hist(dopplers, doppler_rates; nbins=1000)
    histogram = normalize(fit(Histogram, (doppler_rates, dopplers),
                              nbins=nbins), mode=:probability)
    weights = collect(Iterators.flatten(histogram.weights))
    weights = Weights(weights)
    fd_rate, f_d = meshgrid(Array(histogram.edges[1][1:end-1]),
                            Array(histogram.edges[2][1:end-1]))
    fd_rate = collect(Iterators.flatten(fd_rate))
    f_d = collect(Iterators.flatten(f_d))
    return ConstellationHistogram(nbins, dopplers, doppler_rates,
                                  histogram, weights, fd_rate, f_d,
                                  missing, missing, missing)
end

function initiate_hist(dopplers, doppler_rates, elevations, max_elevations;
                       nbins=1000, elev_bins=100)
    histogram = normalize(fit(Histogram, (doppler_rates, dopplers),
                              nbins=nbins), mode=:probability)
    histogram3d = normalize(fit(Histogram, (elevations, doppler_rates, dopplers),
                               nbins=(elev_bins, nbins, nbins)), mode=:probability)
    histogram_dopplers =  normalize(fit(Histogram, (elevations, dopplers),
                                        nbins=(elev_bins, nbins)), mode=:probability)
    histogram_doppler_rates =  normalize(fit(Histogram, (elevations, doppler_rates),
                                        nbins=(elev_bins, nbins)), mode=:probability)
    elev_hists = [histogram3d, histogram_dopplers, histogram_doppler_rates]
    weights = collect(Iterators.flatten(histogram.weights))
    weights = Weights(weights)
    fd_rate, f_d = meshgrid(Array(histogram.edges[1][1:end-1]),
                            Array(histogram.edges[2][1:end-1]))
    fd_rate = collect(Iterators.flatten(fd_rate))
    f_d = collect(Iterators.flatten(f_d))
    return ConstellationHistogram(nbins, dopplers, doppler_rates,
                                  histogram, weights, fd_rate, f_d,
                                  elevations, elev_hists, max_elevations)
end

function sample_distribution(histogram::ConstellationHistogram)
    f_d = sample(histogram.f_d, histogram.weights)
    fd_rate = sample(histogram.fd_rate, histogram.weights)
    return (f_d, fd_rate)
end

function max_doppler_curve_per_elev(histogram::ConstellationHistogram, elevs)
    elevations = histogram.elevations

    doppler_bins = histogram.histogram.edges[1]
    dopplers = histogram.dopplers
    doppler_rates = histogram.doppler_rates
    N = size(histogram.elev_hists[2].weights)[1]
    doppler_medians = zeros(N)
    doppler_rate_medians = zeros(N)
    for i in 1:N
        doppler_median_idx = argmax(histogram.elev_hists[2].weights[i,:])
        doppler_rate_median_idx = argmax(histogram.elev_hists[3].weights[i,:])
        doppler_medians[i] = histogram.elev_hists[2].edges[2][doppler_median_idx]
        doppler_rate_medians[i] = histogram.elev_hists[3].edges[2][doppler_rate_median_idx]
    end
    return (doppler_medians, doppler_rate_medians)
end


println("RODE Tests")
using Flux , StochasticDiffEq, DiffEqNoiseProcess , Optim, Test
println("RODE_tests")
using  NeuralPDE

using Random
Random.seed!(100)

println("Test Case 1")
linear = (u,p,t,W) ->   2u*sin(W)
tspan = (0.00f0, 1.00f0)
u0 = 1.0f0
dt = 1/50f0
W = WienerProcess(0.0,0.0,nothing)
prob = RODEProblem(linear, u0 ,tspan, noise = W)
chain = Flux.Chain(Dense(2,8,relu),Dense(8,16 , relu),Dense(16,1))
opt = ADAM(1e-4)
sol = solve(prob, NeuralPDE.NNRODE(chain,W,opt), dt=dt, verbose = true,
            abstol=1e-10, maxiters = 3000)
W2 = NoiseWrapper(sol.W)
prob1 = RODEProblem(linear , u0 , tspan , noise = W2)
sol2 = solve(prob1,RandomEM(),dt=dt)
err = Flux.mse(sol.u , sol2.u)
@test err < 0.3

println("Test Case 2")
linear = (u,p,t,W) -> t^3 + 2*t + (t^2)*((1+3*(t^2))/(1+t+(t^3))) - u*(t + ((1+3*(t^2))/(1+t+t^3))) + 5*W
tspan = (0.00f0, 1.00f0)
u0 = 1.0f0
dt = 1/100f0
W = WienerProcess(0.0,0.0,nothing)
prob = RODEProblem(linear, u0 ,tspan, noise = W)
chain = Flux.Chain(Dense(2,32,sigmoid),  Dense(32,32,sigmoid)  , Dense(32,1))
opt = ADAM(1e-3)
sol = solve(prob, NeuralPDE.NNRODE(chain,W,opt), dt=dt, verbose = true,
            abstol=1e-10, maxiters = 2000)
W2 = NoiseWrapper(sol.W)
prob1 = RODEProblem(linear , u0 , tspan , noise = W2)
sol2 = solve(prob1,RandomEM(),dt=dt)
err = Flux.mse(sol.u , sol2.u)
@test err < 0.4


# Copyright 2017, Iain Dunning, Joey Huchette, Miles Lubin, and contributors    #src
# This Source Code Form is subject to the terms of the Mozilla Public License   #src
# v.2.0. If a copy of the MPL was not distributed with this file, You can       #src
# obtain one at https://mozilla.org/MPL/2.0/.                                   #src

# # The diet problem

# Solve the classic "diet problem", also known as the [Stigler diet](https://en.wikipedia.org/wiki/Stigler_diet).
# The code is based on [an example from Gurobi](https://www.gurobi.com/documentation/9.0/examples/diet_cpp_cpp.html).

using JuMP
import GLPK
import Test

function print_solution(is_optimal, foods, buy)
    println("RESULTS:")
    if is_optimal
        for food in foods
            println("  $(food) = $(value(buy[food]))")
        end
    else
        println("The solver did not find an optimal solution.")
    end
end

function example_diet(; verbose = true)
    ## Nutrition guidelines
    categories = ["calories", "protein", "fat", "sodium"]
    category_data = Containers.DenseAxisArray([
        1800 2200;
        91   Inf;
        0    65;
        0    1779
        ], categories, ["min", "max"]
    )
    Test.@test category_data["protein", "min"] == 91.0
    Test.@test category_data["sodium", "max"] == 1779.0
    ## Foods
    foods = [
        "hamburger", "chicken", "hot dog", "fries", "macaroni", "pizza",
        "salad", "milk", "ice cream",
    ]
    cost = Containers.DenseAxisArray(
        [2.49, 2.89, 1.50, 1.89, 2.09, 1.99, 2.49, 0.89, 1.59],
        foods
    )
    food_data = Containers.DenseAxisArray(
        [
            410 24 26 730;
            420 32 10 1190;
            560 20 32 1800;
            380  4 19 270;
            320 12 10 930;
            320 15 12 820;
            320 31 12 1230;
            100  8 2.5 125;
            330  8 10 180
        ], foods, categories
    )
    Test.@test food_data["hamburger", "calories"] == 410.0
    Test.@test food_data["milk", "fat"] == 2.5
    ## Build model
    model = Model(GLPK.Optimizer)
    @variables(model, begin
        ## Variables for nutrition info
        category_data[c, "min"] <= nutrition[c = categories] <= category_data[c, "max"]
        ## Variables for which foods to buy
        buy[foods] >= 0
    end)
    ## Objective - minimize cost
    @objective(model, Min, sum(cost[f] * buy[f] for f in foods))
    ## Nutrition constraints
    @constraint(model, [c in categories],
        sum(food_data[f, c] * buy[f] for f in foods) == nutrition[c]
    )
    ## Solve
    if verbose
        println("Solving original problem...")
    end
    optimize!(model)
    term_status = termination_status(model)
    is_optimal = term_status == MOI.OPTIMAL
    Test.@test primal_status(model) == MOI.FEASIBLE_POINT
    Test.@test objective_value(model) ≈ 11.8288 atol = 1e-4
    if verbose
        print_solution(is_optimal, foods, buy)
    end
    ## Limit dairy (note that the problem will become infeasible).
    @constraint(model, buy["milk"] + buy["ice cream"] <= 6)
    if verbose
        println("Solving dairy-limited problem...")
    end
    optimize!(model)
    Test.@test termination_status(model) == MOI.INFEASIBLE
    Test.@test primal_status(model) == MOI.NO_SOLUTION
    if verbose
        print_solution(false, foods, buy)
    end
    return
end

example_diet()


using Propagators: Cell, propagator, propagate!, block, constant, switch

# function sqrt_heron(x)
#   g = x/x
#   while !(g^2 ≈ x)
#     g = (x/g+g)/2
#   end
#   return g
# end

function heron_step(x, g, h)
  block(x, g) do
    xdg = Cell()
    gpxdg = Cell()
    two = Cell()
    constant(2)(two)
    propagator(/)(x, g, xdg)
    propagator(+)(g, xdg, gpxdg)
    propagator(/)(gpxdg, two, h)
  end
end

function sqrt_network(x, answer)
  block(x) do
    one = Cell()
    constant(1)(one)
    sqrt_iter(x, one, answer)
  end
end

function sqrt_iter(x, g, answer)
  block(x, g) do
    done = Cell()
    notdone = Cell()
    xifnotdone = Cell()
    gifnotdone = Cell()
    newg = Cell()
    propagator((g, x) -> g^2 ≈ x)(g, x, done)
    propagator(!)(done, notdone)
    switch(done, g, answer)
    switch(notdone, x, xifnotdone)
    switch(notdone, g, gifnotdone)
    heron_step(xifnotdone, gifnotdone, newg)
    sqrt_iter(xifnotdone, newg, answer)
  end
end

# x = Cell()
# y = Cell()
# sqrt_network(x, y)
# x[] = 2
# propagate!()
# y[]


function solver_EPWE(solver::Solver, ky::Number, omega::Number, polarisation::Polarisation; Norm=false)
    epc = solver.epc.*mask(size(solver.epc)[1])
    muc = solver.muc.*mask(size(solver.muc)[1])
    global K = hcat(diag_R(solver.Ky),diag_R(solver.Kx))/2/pi
    M = length(K[:,1])
    LHS11 = zeros(ComplexF64, M, M)
    LHS12 = eye(M)
    Kp = K + ones(M) * [0 ky]

    if polarisation == TE
        LHS21 = epc * (omega^2*muc*eye(M) .- inv(epc).*(Kp*transpose(Kp)))
        LHS22 = -epc*(inv(epc).*(K[:,1]*ones(1,M)+ones(M)*K[:,1]'))
        LHS = [LHS11 LHS12; LHS21 LHS22]
        RHS = epc
        label = L"H_z"
    elseif polarisation == TM
        LHS21 = (omega^2*epc - diagm(K[:,1].^2 + (ky .+ K[:,2]).^2))
        LHS22 = diagm(-2*K[:,1])
        LHS = [LHS11 LHS12; LHS21 LHS22]
        RHS = muc
        label = L"E_z"
    end
    kT, VT = eigen(LHS,sortby=x -> -abs(x))
    VT = VT[1:M, :]
    # Normalization
    if Norm == true
    VT = VT * diagm(vec(1 ./sqrt.(sum(abs.(VT).^2,dims=1))))
    maxVT,idVT=findmax(abs.(VT), dims=1)
    VT=VT*diagm(vec(maxVT./VT[idVT]))
    end
    #direction
    dT=1/omega*diag_R(VT'*inv(RHS)*(diagm(K[:,1])*VT+VT*diagm(kT)))

    # find the solution in 1st BZ
    i0= (abs.(real.(dT)) .> abs.(imag.(dT))).*((real.(dT)+imag.(dT)) .< 0) + (abs.(real.(dT)) .< abs.(imag.(dT))).*(imag.(kT).< 0)
    i0 = Bool.(i0)

    i3 = []
    for i4 = 0:1
        i0 = .!i0
        i1= 1:2*M
        i1 = i1[i0]
        i2 = sortperm(abs.(real.(kT[i0])))
        buf = sort(abs.(real.(kT[i0])))
        i1=i1[i2]
        i2 = (abs.([imag.(kT[i1]); Inf] - [Inf; imag.(kT[i1])] ) .> 1e-10) .| (abs.([imag.(kT[i1]); Inf]) .<= 1e-10) .| (abs.(abs.([real.(kT[i1]); Inf]) .- abs.([Inf; real.(kT[i1])])) .> 1e-10) .| (abs.(abs.([real.(kT[i1]); Inf]) .+ abs.([Inf; real.(kT[i1])]) .- 1) .>= 0.1)
        i1=i1[i2[1:(length(i2)-1)]]
        i3=[i3; i1[1:Int(sqrt(M))]]
    end
    kT=kT[i3]
    dT=dT[i3]
    VT=VT[:,i3]
    return kT, dT, VT
end

function abcd(kT::VecOrMat,dT::VecOrMat,VT::VecOrMat,omega::Number,ky::Number;background=1)
    M = size(VT)[1]
    m = Int((sqrt(M)-1)/2)
    Ky=Int(round(ky))
    ky=ky-Ky
    msum=zeros((2*m+1),(2*m+1)^2)
    for i0=1:(2*m+1)
        msum[i0,((i0-1)*(2*m+1)+1):(i0*(2*m+1))] .= 1
    end

    ep = background
    abcd1=[-eye(2*m+1) msum*VT[:,1:(2*m+1)]; diagm(sqrt.(Complex.(omega^2 .- ((-m:m) .+ky).^2))) msum*(inv(ep)*(diagm(K[:,1])*VT[:,1:(2*m+1)]+VT[:,1:(2*m+1)]*diagm(kT[1:(2*m+1)])))]
    input = zeros(ComplexF64, 4*m+2)
    input[Ky+m+1] = 1
    input[Ky+3*m+2] = sqrt(Complex(omega^2-(ky+Ky)^2))
    output=abcd1\input
    rT=output[1:(2*m+1)]
    tT=output[(2*m+2):(4*m+2)]

    kxi = sqrt(Complex(omega^2 - ky^2))
    if kxi==0
        kxi=eps()
    end
    kxr=sqrt.(Complex.(omega^2 .- (ky - round(ky) .+ (-m:m)).^2))

    i0=abs.(imag.(kxr[1:(2*m+1)])) .<= 1e-15
    #    i0 = Bool.[i0]
    if length(rT[i0])==0
        RT = 0
    else
        RT = rT[i0]'*diagm(kxr[i0])*rT[i0]/kxi;
    end

    i0 = abs.(imag.(dT[1:(2*m+1)])) .<= abs.(real.(dT[1:(2*m+1)]))
    if length(tT[i0])==0
        TT=0
    else
        TT=tT[i0]'*diagm(dT[1:2m+1][i0])*tT[i0]/kxi*omega;
    end

#    RT = abs(RT); TT = abs(TT)
    return RT, TT, rT, tT
end


module AnalyticBSplines

using Polynomials
using BSplinesQuasi
import BSplinesQuasi: AbstractKnotSet

include("intervals.jl")
include("ppoly.jl")

function gen_basis(t::AbstractKnotSet{k,ml,mr,T}, bl=one(T), br=one(T)) where {k,ml,mr,T}
    d(a,b) = b == 0 ? 0 : a/b
    B = [[PPoly([((t[j],t[j+1]) => [one(T)])])
          for j = 1:length(t)-1]]
    for k = 2:order(t)
        Bₖ = [d(Poly([-t[i],1]),(t[i+k-1]-t[i]))*B[end][i] +
              d(Poly([t[i+k],-1]),(t[i+k]-t[i+1]))*B[end][i+1]
              for i = 1:length(t)-k]
        push!(B,Bₖ)
    end
    B[end][1] *= bl
    B[end][end] *= br
    B
end

function scalar_op(Bₖ, t::AbstractKnotSet{k,ml,mr,T}, f::Poly=Poly([one(T)])) where {k,ml,mr,T}
    B = Bₖ[end]
    n = length(B)
    O = Matrix{T}(undef, n,n)
    l(v::Rational) = denominator(v) == 1 ? numerator(v) : 1.0*v
    l(v) = v
    for i = 1:n
        for j = 1:n
            O[i,j] = polyint(B[i]*f*B[j], l(first(t)),l(last(t)))
        end
    end
    O
end

function derop(L, R, t::AbstractKnotSet{k,ml,mr,T}, o::Integer) where {k,ml,mr,T}
    m = length(L[end])
    n = length(R[end])
    O = Matrix{T}(undef, m,n)
    l(v::Rational) = denominator(v) == 1 ? numerator(v) : 1.0*v
    l(v) = v
    d,r = divrem(o,2)
    a,b = d,d+r
    for i = 1:m
        for j = 1:n
            O[i,j] = (iseven(a) ? 1 : -1)*polyint(polyder(L[end][i],a)*polyder(R[end][j],b), l(first(t)),l(last(t)))
        end
    end
    O
end

derop(B, t, o) = derop(B, B, t, o)

export gen_basis, scalar_op, derop

end # module


struct fvflux{ftype,dftype}
    flux::ftype
    Dflux::dftype
end

(F::fvflux)(u) = F.flux(u)

function scalar_flux_function(f, Df)
    if Df == nothing
        Df = x -> ForwardDiff.derivative(f,x)
    end
    fvflux(f,Df)
end

function flux_function(f, Df)
    if Df == nothing
        Df = x -> ForwardDiff.jacobian(f,x)
    end
    fvflux(f,Df)
end

@inline evalJacobian(f::fvflux, uj) = f.Dflux(uj)

function fluxρ(uj,f::fvflux)
        maximum(abs,eigvals(f.Dflux(uj)))
end

function update_dt(alg::AbstractFVAlgorithm,u,Flux,
    CFL,mesh)
  maxρ = zero(eltype(u))
  dx = cell_volume(mesh, 1)
  for i in cell_indices(mesh)
    maxρ = max(maxρ, fluxρ(value_at_cell(u,i,mesh), Flux))
  end
  CFL/(1/dx*maxρ)
end

function maxfluxρ(u, Flux, mesh)
    maxρ = zero(eltype(u))
    for i in cell_indices(mesh)
      maxρ = max(maxρ, fluxρ(value_at_cell(u,i,mesh), Flux))
    end
    maxρ
end


using QuantumOptics
using PyPlot

b = SpinBasis(1//2)

sx = sigmax(b)
sy = sigmay(b)
sz = sigmaz(b)

N = 6

spinchain = b ⊗ b

for ispin in 3:N
    global spinchain
    spinchain = spinchain ⊗ b
end

XX = sx ⊗ sx
YY = sy ⊗ sy
ZZ = sz ⊗ sz

J=1

HHeis = embed(spinchain, [1,2] , J*(XX+YY+ZZ))
for ispin in 2:N-1
    global HHeis
    HHeis += embed(spinchain, [ispin,ispin+1] , J*(XX+YY+ZZ))
end

# make sure H is numerically Hermitian
HHeis = ( HHeis + dagger(HHeis) )/2

s5z = embed(spinchain, [5] , sz)
s6z = embed(spinchain, [6] , sz)

sup = spinup(b)
sdo = spindown(b)

psi0 = sup
for ispin in 2:N
    global psi0
    if ispin == 5
       psi0 = psi0 ⊗ sdo
    else
       psi0 = psi0 ⊗ sup
    end
end

times = 0:.1:1.5

tout, psi_t = timeevolution.schroedinger(times, psi0, HHeis)

m5z = expect(s5z, psi_t)
m6z = expect(s6z, psi_t)

plot(tout, m5z)
plot(tout, m6z)


push!(LOAD_PATH,"../src")

using Test
using TypesParticles
using NuclearRepulsion

"""
	run_nuclearrepulsion()

""" function run_nuclearrepulsion_test()
	 @testset "NuclearRepulsion" begin
	 #H2 structure
	 x0,y0,z0 = 0.00e0,0.00e0,0.00e0;
	 x1,y1,z1 = 0.77e0,0.00e0,0.00e0;
	 Z0,Z1 = 1,1;
	 H1 = Atom(x0,y0,z0,Z0);
	 H2 = Atom(x1,y1,z1,Z1);

	 cell = [1.00e0 0.00e0 0.00e0;
	      	 0.00e0 1.00e0 0.00e0;
		 0.00e0 0.00e0 1.00e0];

             @test begin
	         atomlist = [H1,H2];
	         system = SystemOfAtoms(2,atomlist,cell);
	         calcnuclrepul(system) == 1.2987012987012987e0; #Atomic Units
                 1 == 1
             end
        end #testset
end #run_nuclearrepulsion_test


module FastParzenWindows

    using ProgressMeter, StatsFuns, Random, Printf, SparseArrays, LinearAlgebra, Distributions, MLBase

    include("fastparzenwindows.jl")

    include("softparzen.jl")

    include("partition.jl")

    include("spiraldata.jl")

    include("getmixturemodel.jl")

    include("cvfpw.jl")

    export fpw, spiraldata, cv_fpw

end


# Double Pendulum
T = Float64
model = TrajectoryOptimization.Dynamics.doublependulum
model_d = rk3(model)
n = model.n; m = model.m

# costs
Qf = 100.0*Diagonal(I,n)
Q = (1.0e-2)*Diagonal(I,n)
R = (1.0e-2)*Diagonal(I,m)

x0 = [0.; 0.; 0.; 0.]
xf = [π; 0.; 0.; 0.];
N = 101
dt = 0.01
U0 = [ones(m) for k = 1:N-1]
obj = TrajectoryOptimization.LQRObjective(Q,R,Qf,xf,N)

doublependulum = TrajectoryOptimization.Problem(model_d, obj, x0=x0, xf=xf, N=N, dt=dt)
initial_controls!(doublependulum, U0)


using ChainRulesCore
using ForwardDiff
using FiniteDifferences
using Transits: compute_grad, compute_gn, compute_gn_jac

# @testset "PolynomialLimbDark" begin
#     u = [0.4, 0.26, 0.1, 0.1]
#     # test_frule(PolynomialLimbDark, u)
#     # test_rrule(PolynomialLimbDark, u)

#     ld = PolynomialLimbDark(u)

#     test_frule(compute, ld, 0.0, 0.1; fdm=forward_fdm(10, 1))
#     test_rrule(compute, ld, 0.0, 0.1; fdm=forward_fdm(10, 1))
# end

# @testset "QuadLimbDark" begin
#     u = [0.4, 0.26]
#     test_frule(QuadLimbDark, u)
#     test_rrule(QuadLimbDark, u)

#     ld = QuadLimbDark(u)

#     test_frule(compute, ld, 0.0, 0.1)
#     test_rrule(compute, ld, 0.0, 0.1)
# end

function finite_diff(b_, r_, u_n_, law=PolynomialLimbDark; diff=big(1e-18))
    b = big(b_)
    r = big(r_)
    u_n = big.(u_n_)
    ld = law(u_n)
    f = ld(b, r)

    # modulate r
    f₊ = ld(b, r + diff)
    if r == 1 && length(u_n) > 2
        dfdr = (f₊ - f) / diff
    else
        f₋ = ld(b, r - diff)
        dfdr = (f₊ - f₋) / (2 * diff)
    end

    # modulate b
    f₊ = ld(b + diff, r)
    if b > diff
        f₋ = ld(b - diff, r)
        dfdb = (f₊ - f₋) / (2 * diff)
    elseif iszero(b)
        dfdb = zero(b)
    else
        dfdb = (f₊ - f) / diff
    end

    return Float64.([dfdb, dfdr])
end

function test_compute_grad(b, r, u_n, law=PolynomialLimbDark)
    ld = law(u_n)
    _, _, dfdb, dfdr = @inferred compute_grad(ld, b, r)
    grad_analytical = [dfdb, dfdr]

    atol = r != 1 ? 1e-7 : 3e-6

    @testset "compute_grad - b=$b, r=$r" begin
        # test against BigFloat for numerical accuracy
        _, _, dfdb_big, dfdr_big = compute_grad(ld, big(b), big(r))
        grad_big = [dfdb_big, dfdr_big]
        @test grad_analytical ≈ grad_big atol = atol

        # test against finite diff for correctness
        grad_numerical = finite_diff(b, r, u_n, law)
        @test grad_analytical ≈ grad_numerical atol = atol
    end
end

logspace(start, stop, N) = 10 .^range(log10(start), log10(stop), length=N)

test_coeffs = [[0.0], [1.0], [2.0, -1.0], [3.0, -3.0, 1.0], fill(0.1, 4), fill(0.1, 5), fill(0.1, 10)]
test_names = ["uniform", "linear", "quadratic", "cubic", "quartic", "quintic", "10th order"]

@testset "`compute` grads - $name" for (name, u_n) in zip(test_names, test_coeffs)
    # b + r < 1
    rs = [0.01, 0.1, 0.5, 0.999999]
    nb = 50
    ϵ, δ = 1e-9, 1e-3
    for r in rs
        bs = unique(abs.([
            range(0, ϵ, length=nb);
            range(ϵ, δ, length=nb);
            range(δ, r - δ, length=nb);
            -logspace(δ, ϵ, nb) .+ r;
            range(r - ϵ, r + ϵ, length=nb);
            logspace(ϵ, δ, nb) .+ r;
            range(r + δ, 1 - r - δ, length=nb);
            -logspace(δ, ϵ, nb) .+ (1 - r);
            range(1 - r - ϵ, 1 - r + ϵ, length=nb);
            logspace(ϵ, δ, nb) .+ (1 - r);
            range(1 - r + δ, 1 + r - δ, length=nb);
            -logspace(δ, ϵ, nb) .+ (1 + r);
            range(1 + r - ϵ, 1 + r - 1e-13, length=nb)
        ]))
        test_compute_grad.(bs, r, (u_n,), PolynomialLimbDark)
        if length(u_n) < 3
            test_compute_grad.(bs, r, (u_n,), QuadLimbDark)
        end
    end

    rs = [1, 1.000001, 2, 10, 100]
    for r in rs
        bs = unique(abs.([
            r - 1 + 1e-13;
            logspace(ϵ, δ, nb) .+ (r - 1);
            range(r - 1 + δ, r - δ, length=nb);
            -logspace(δ, ϵ, nb) .+ r;
            range(r - ϵ, r + ϵ, length=nb);
            logspace(ϵ, δ, nb) .+ r;
            range(r + δ, r + 1 - δ, length=nb);
            -logspace(δ, ϵ, nb) .+ (r + 1);
            r + 1 - 1e-13
        ]))
        test_compute_grad.(bs, r, (u_n,), PolynomialLimbDark)
        if length(u_n) < 3
            test_compute_grad.(bs, r, (u_n,), QuadLimbDark)
        end
    end
end

@testset "`compute_gn` jacs - $name" for (name, u_n) in zip(test_names, test_coeffs)
    un_full = [-1, u_n...]
    gn_primal = @inferred compute_gn(un_full)
    gn, jac_analytical = @inferred compute_gn_jac(un_full)

    @test gn_primal == gn

    jac_numerical = FiniteDifferences.jacobian(central_fdm(5, 1), compute_gn, un_full)[1]'

    @test jac_analytical ≈ jac_numerical atol=1e-8
end

@testset "frule, rrule comparison - $name" for (name, u_n) in zip(test_names, test_coeffs)
    @testset "PolynomialLimbDark" begin
        function primal(X)
            ld = PolynomialLimbDark(X[3:end])
            return compute(ld, X[1], X[2])
        end

        function gradr(X)
            ld, ld_pullback = rrule(PolynomialLimbDark, X[3:end])
            f, f_pullback = rrule(compute, ld, X[1], X[2])

            f̄ = one(eltype(X))
            _, l̄d, b̄, r̄ = f_pullback(f̄)
            _, ū_n = ld_pullback(l̄d)
            return [b̄, r̄, ū_n...]
        end

        function gradf(X)
            N = length(X) - 2
            ∂us = zeros(N)
            for i in eachindex(∂us)
                Δus = zeros(N)
                Δus[i] = 1
                ld, ∂ld = frule((0, Δus), PolynomialLimbDark, X[3:end])
                _, ∂us[i] = frule((0, ∂ld, 0, 0), compute, ld, X[1], X[2])
            end
            ld, ∂ld3 = frule((0, zeros(N)), PolynomialLimbDark, X[3:end])
            _, ∂b = frule((0, ∂ld3, 1, 0), compute, ld, X[1], X[2])
            _, ∂r = frule((0, ∂ld3, 0, 1), compute, ld, X[1], X[2])
            return [∂b, ∂r, ∂us...]
        end

        function test_grads(X)
            grad_rev = @inferred gradr(X)
            grad_for = @inferred gradf(X)
            grad_FD = ForwardDiff.gradient(primal, X)
            @test grad_rev ≈ grad_for atol=1e-7
            @test grad_rev ≈ grad_FD atol=1e-7
        end

        test_grads([0.3, 0.2, u_n...])
        test_grads([1.1, 0.3, u_n...])
    end

    if length(u_n) < 3
        @testset "QuadLimbDark" begin
            function primal(X)
                ld = QuadLimbDark(X[3:end])
                return compute(ld, X[1], X[2])
            end

            function gradr(X)
                ld, ld_pullback = rrule(QuadLimbDark, X[3:end])
                f, f_pullback = rrule(compute, ld, X[1], X[2])

                f̄ = one(eltype(X))
                _, l̄d, b̄, r̄ = f_pullback(f̄)
                _, ū_n = ld_pullback(l̄d)
                return [b̄, r̄, ū_n...]
            end

            function gradf(X)
                N = length(X) - 2
                ∂us = zeros(N)
                for i in eachindex(∂us)
                    Δus = zeros(N)
                    Δus[i] = 1
                    ld, ∂ld = frule((0, Δus), QuadLimbDark, X[3:end])
                    _, ∂us[i] = frule((0, ∂ld, 0, 0), compute, ld, X[1], X[2])
                end
                ld, ∂ld3 = frule((0, zeros(N)), QuadLimbDark, X[3:end])
                _, ∂b = frule((0, ∂ld3, 1, 0), compute, ld, X[1], X[2])
                _, ∂r = frule((0, ∂ld3, 0, 1), compute, ld, X[1], X[2])
                return [∂b, ∂r, ∂us...]
            end

            function test_grads(X)
                grad_rev = @inferred gradr(X)
                grad_for = @inferred gradf(X)
                grad_FD = ForwardDiff.gradient(primal, X)
                @test grad_rev ≈ grad_for atol=1e-7
                @test grad_rev ≈ grad_FD atol=1e-7
            end

            test_grads([0.3, 0.2, u_n...])
            test_grads([1.1, 0.3, u_n...])
        end
    end
end


function filter_coeffs(a::Vector{Float64}, b::Vector{Float64}, orders::Vector{Int64}, data::Array{Float64})

    dinv = (b - a)./(orders.-1)
    d = length(a)
    if d == 1
        coeffs = filter_coeffs_1d(dinv,data)
    elseif d == 2
        coeffs = filter_coeffs_2d(dinv,data)
    elseif d == 3
        coeffs = filter_coeffs_3d(dinv,data)
    elseif d == 4
        coeffs = filter_coeffs_4d(dinv,data)
    else
        error()
    end
    return coeffs
end

# function find_coeffs_1d_old(delta_inv, M, data)
#
#     # args: double, int, double[:], double[:
#
#     N = M + 2
#
#     data = data[:]
#
#     if length(data) == N
#         rhs = data
#     else
#         rhs = [0; data; 0]
#     end
#
#     basis = [1.0/6.0, 2.0/3.0, 1.0/6.0]
#
#     vals = repeat( basis, outer=[M])
#     co_x = repeat( collect(2:M+1), inner=[3])
#     co_y = repeat( collect(1:3) , outer=[M]) + co_x .- 2
#
#     db = 4
#     initial = [1.0,-2.0,1.0,0.0]*delta_inv*delta_inv
#     final = [0.0,1.0,-2.0,1.0]*delta_inv*delta_inv
#
#     vals = [initial; vals; final]
#     co_x = [ones(Int,db); co_x; ones(Int,db)*(M+2)]
#     co_y = [1:db; co_y; M+3-db:M+2]
#
#     spmat = sparse(co_x, co_y, vals)
#
#     sol = spmat \ rhs
#
#     return sol
#
# end

function solve_deriv_interp_1d!(bands::Matrix{Float64}, coefs::Vector{Float64})

    M = size(coefs,1) - 2

    # Solve interpolating equations
    # First and last rows are different

    bands[1, 2] /= bands[1, 1]
    bands[1, 3] /= bands[1, 1]
    bands[1, 4] /= bands[1, 1]
    bands[1, 1] = 1.0
    bands[2, 2] -= bands[2, 1] * bands[1, 2]
    bands[2, 3] -= bands[2, 1] * bands[1, 3]
    bands[2, 4] -= bands[2, 1] * bands[1, 4]
    bands[1, 1] = 0.0
    bands[2, 3] /= bands[2, 2]
    bands[2, 4] /= bands[2, 2]
    bands[2, 2] = 1.0

    # Now do rows 2 through M+1
    for row = 3:(M+1)
        bands[row, 2] -= bands[row, 1] * bands[row - 1, 3]
        bands[row, 4] -= bands[row, 1] * bands[row - 1, 4]
        bands[row, 3] /= bands[row, 2]
        bands[row, 4] /= bands[row, 2]
        bands[row, 1] = 0.0
        bands[row, 2] = 1.0
    end
    # Do last row
    bands[M + 2, 2] -= bands[M + 2, 1] * bands[M, 3]
    bands[M + 2, 4] -= bands[M + 2, 1] * bands[M, 4]
    bands[M + 2, 3] -= bands[M + 2, 2] * bands[M + 1, 3]
    bands[M + 2, 4] -= bands[M + 2, 2] * bands[M + 1, 4]
    bands[M + 2, 4] /= bands[M + 2, 3]
    bands[M + 2, 3] = 1.0

    coefs[M + 2] = bands[M + 2, 4]
    # Now back substitute up
    # for row in (M+1):2
    for ii in 1:M
        row = M+2 - ii
        coefs[row] = bands[row, 4] - bands[row, 3] * coefs[row + 1]
    end
    # Finish with first row
    coefs[1] = bands[1, 4] - bands[1, 2] * coefs[2] - bands[1, 3] * coefs[3]

end

function find_coeffs_1d!(delta_inv, M, data, coefs)

    # data is 1d array of length M
    # data is 1d array of length M+2

    basis = [1.0/6.0, 2.0/3.0, 1.0/6.0]
    bands = zeros(M+2, 4)

    # Setup boundary conditions
    abcd_left = zeros(4)
    abcd_right = zeros(4)

    # Left boundary
    abcd_left[1] = 1.0 * delta_inv * delta_inv
    abcd_left[2] = -2.0 * delta_inv * delta_inv
    abcd_left[3] = 1.0 * delta_inv * delta_inv
    abcd_left[4] = 0

    # Right boundary
    abcd_right[1] = 1.0 * delta_inv * delta_inv
    abcd_right[2] = -2.0 * delta_inv * delta_inv
    abcd_right[3] = 1.0 * delta_inv * delta_inv
    abcd_right[4] = 0

    for i in 1:4
        bands[1, i] = abcd_left[i]
        bands[M + 2, i] = abcd_right[i]
    end
    for i in 2:(M+1)
        for j in 1:3
            bands[i, j] = basis[j]
            bands[i, 4] = data[i-1]
        end
    end

    solve_deriv_interp_1d!(bands, coefs)

    return coefs
end

function find_coeffs_1d(delta_inv, M, data)
    coefs = zeros(M+2)
    if length(data) == M
        find_coeffs_1d!(delta_inv, M, data, coefs)
    else
        find_coeffs_1d!(delta_inv, M, data[2:end-1], coefs)
    end
    return coefs
end

function filter_coeffs_1d(dinv, data)

  M = size(data,1)
  N = M+2

  coefs = find_coeffs_1d(dinv[1], M, data)

  return coefs

end

function filter_coeffs_2d(dinv, data)

    Mx = size(data,1)
    My = size(data,2)

    Nx = Mx+2
    Ny = My+2

    coefs = zeros(Nx,Ny)

    # First, solve in the X-direction
    for iy in 1:My
        coefs[:,iy+1] = find_coeffs_1d(dinv[1], Mx, data[:,iy])
    end

    # Now, solve in the Y-direction
    for ix in 1:Nx
        coefs[ix,:] = find_coeffs_1d(dinv[2], My, coefs[ix,:])
    end

    return coefs

end

function filter_coeffs_3d(dinv, data)

    Mx = size(data,1)
    My = size(data,2)
    Mz = size(data,3)

    Nx = Mx+2
    Ny = My+2
    Nz = Mz+2

    coefs = zeros(Nx,Ny,Nz)

    # First, solve in the X-direction
    for iy in 1:My
        for iz in 1:Mz
            coefs[:,iy+1,iz+1] = find_coeffs_1d(dinv[1], Mx, data[:,iy,iz])
        end
    end

    # Now, solve in the Y-direction
    for ix in 1:Nx
        for iz in 1:Mz
            coefs[ix,:,iz+1] = find_coeffs_1d(dinv[2], My, coefs[ix,:,iz+1])
        end
    end

    # Now, solve in the Z-direction
    for ix in 1:Nx
        for iy in 1:Ny
            coefs[ix,iy,:] = find_coeffs_1d(dinv[3], Mz, coefs[ix,iy,:])
        end
    end

    return coefs

end

function filter_coeffs_4d(dinv, data)

    Mx = size(data,1)
    My = size(data,2)
    Mz = size(data,3)
    Mz4 = size(data,4)

    Nx = Mx+2
    Ny = My+2
    Nz = Mz+2
    Nz4 = Mz4+2

    coefs = zeros(Nx,Ny,Nz,Nz4)

    # First, solve in the X-direction
    for iy in 1:My
        for iz in 1:Mz
            for iz4 in 1:Mz4
                coefs[:,iy+1,iz+1,iz4+1] = find_coeffs_1d(dinv[1], Mx, data[:,iy,iz,iz4])
            end
        end
    end

    # Now, solve in the Y-direction
    for ix in 1:Nx
        for iz in 1:Mz
            for iz4 in 1:Mz4
                coefs[ix,:,iz+1,iz4+1] = find_coeffs_1d(dinv[2], My, coefs[ix,:,iz+1,iz4+1])
            end
        end
    end

    # Now, solve in the Z-direction
    for ix in 1:Nx
        for iy in 1:Ny
            for iz4 in 1:Mz4
                coefs[ix,iy,:,iz4+1] = find_coeffs_1d(dinv[3], Mz, coefs[ix,iy,:,iz4+1])
            end
        end
    end
    # Now, solve in the Z4-direction
    for ix in 1:Nx
        for iy in 1:Ny
            for iz in 1:Mz
                coefs[ix,iy,iz,:] = find_coeffs_1d(dinv[4], Mz, coefs[ix,iy,iz,:])
            end
        end
    end

    return coefs

end


#src # Shock capturing with flux differencing and stage limiter

# This tutorial contains a short summary of the idea of shock capturing for DGSEM with flux differencing
# and its implementation in [Trixi.jl](https://github.com/trixi-framework/Trixi.jl).
# In the second part, an implementation of a positivity preserving limiter is added to the simulation.

# # Shock capturing with flux differencing

# The following rough explanation is on a very basic level. More information about an entropy stable
# shock-capturing strategy for DGSEM discretizations of advection dominated problems, such as the
# compressible Euler equations or the compressible Navier-Stokes equations, can be found in
# [Hennemann et al. (2021)](https://doi.org/10.1016/j.jcp.2020.109935). In
# [Rueda-Ramírez et al. (2021)](https://doi.org/10.1016/j.jcp.2021.110580) you find the extension to
# the systems with non-conservative terms, such as the compressible magnetohydrodynamics (MHD) equations.

# The strategy for a shock-capturing method presented by Hennemann et al. is based on a hybrid blending
# of a high-order DG method with a low-order variant. The low-order subcell finite volume (FV) method is created
# directly with the Legendre-Gauss-Lobatto (LGL) nodes already used for the high-order DGSEM.
# Then, the final method is a convex combination with regulating indicator $\alpha$ of these two methods.

# Since the surface integral is equal for both the DG and the subcell FV method, only the volume integral divides
# between the two methods.

# This strategy for the volume integral is implemented in Trixi under the name of
# [`VolumeIntegralShockCapturingHG`](@ref) with the three parameters of the indicator and the volume fluxes for
# the DG and the subcell FV method.

# Note, that the DG method is based on the flux differencing formulation. Hence, you have to use a
# two-point flux, such as [`flux_ranocha`](@ref), [`flux_shima_etal`](@ref), [`flux_chandrashekar`](@ref) or [`flux_kennedy_gruber`](@ref),
# for the DG volume flux. We would recommend to use the entropy conserving flux `flux_ranocha` by
# [Ranocha (2018)](https://cuvillier.de/en/shop/publications/7743) for the compressible Euler equations.
# ````julia
# volume_integral = VolumeIntegralShockCapturingHG(indicator_sc;
#                                                  volume_flux_dg=volume_flux_dg,
#                                                  volume_flux_fv=volume_flux_fv)
# ````

# We now focus on a choice of the shock capturing indicator `indicator_sc`.
# A possible indicator is $\alpha_{HG}$ presented by Hennemann et al. (p.10), which depends on the
# current approximation with modal coefficients $\{m_j\}_{j=0}^N$ of a given `variable`.

# The indicator is calculated for every DG element by itself. First, we calculate a smooth $\alpha$ by
# ```math
# \alpha = \frac{1}{1+\exp(-\frac{-s}{\mathbb{T}}(\mathbb{E}-\mathbb{T}))}
# ```
# with the total energy $\mathbb{E}=\max\big(\frac{m_N^2}{\sum_{j=0}^N m_j^2}, \frac{m_{N-1}^2}{\sum_{j=0}^{N-1} m_j^2}\big)$,
# threshold $\mathbb{T}= 0.5 * 10^{-1.8*(N+1)^{1/4}}$ and parameter $s=ln\big(\frac{1-0.0001}{0.0001}\big)\approx 9.21024$.

# For computational efficiency, $\alpha_{min}$ is introduced und used for
# ```math
# \tilde{\alpha} = \begin{cases}
# 0, & \text{if } \alpha<\alpha_{min}\\
# \alpha, & \text{if } \alpha_{min}\leq \alpha \leq 1- \alpha_{min}\\
# 1, & \text{if } 1-\alpha_{min}<\alpha.
# \end{cases}
# ```

# Moreover, the parameter $\alpha_{max}$ sets a maximal value for $\alpha$ by
# ```math
# \alpha = \min\{\tilde{\alpha}, \alpha_{max}\}.
# ```
# This allows to control the maximal dissipation.

# To remove numerical artifact the final indicator is smoothed with all the neighboring elements'
# indicators. This is activated with `alpha_smooth=true`.
# ```math
# \alpha_{HG} = \max_E \{ \alpha, 0.5 * \alpha_E\},
# ```
# where $E$ are all elements sharing a face with the current element.

# Furthermore, you can specify the variable used for the calculation. For instance you can choose
# `density`, `pressure` or both with `density_pressure` for the compressible Euler equations.
# For every equation there is also the option to use the first conservation variable with `first`.

# This indicator is implemented in Trixi and called [`IndicatorHennemannGassner`](@ref) with the parameters
# `equations`, `basis`, `alpha_max`, `alpha_min`, `alpha_smooth` and `variable`.
# ````julia
# indicator_sc = IndicatorHennemannGassner(equations, basis,
#                                          alpha_max=0.5,
#                                          alpha_min=0.001,
#                                          alpha_smooth=true,
#                                          variable=variable)
# ````

# # Positivity preserving limiter

# Some numerical solutions are physically meaningless, for instance negative values of pressure
# or density for the compressible Euler equations. This often results in crashed simulations since
# the calculation of numerical fluxes or stable time steps uses mathematical operations like roots or
# logarithms. One option to avoid these cases are a-posteriori positivity preserving limiters.
# Trixi provides the fully-discrete positivity-preserving limiter of
# [Zhang, Shu (2011)](https://doi.org/10.1098/rspa.2011.0153).

# It works the following way. For every passed (scalar) variable and for every DG element we calculate
# the minimal value $value_{min}$. If this value falls below the given threshold $\varepsilon$,
# the approximation is slightly adapted such that the minimal value of the relevant variable lies
# now above the threshold.
# ```math
# \underline{u}^{new} = \theta * \underline{u} + (1-\theta) * u_{mean}
# ```
# where $\underline{u}$ are the collected pointwise evaluation coefficients in element $e$ and
# $u_{mean}$ the integral mean of the quantity in $e$. The new coefficients are a convex combination
# of these two values with factor
# ```math
# \theta = \frac{value_{mean} - \varepsilon}{value_{mean} - value_{min}},
# ```
# where $value_{mean}$ is the relevant variable evaluated for the mean value $u_{mean}$.

# The adapted approximation keeps the exact same mean value, but the relevant variable is now greater
# or equal the threshold $\varepsilon$ at every node in every element.

# We specify the variables the way we did before for the shock capturing variables. For the
# compressible Euler equations `density`, `pressure` or the combined variable `density_pressure`
# are a reasonable choice.

# You can implement the limiter in Trixi using [`PositivityPreservingLimiterZhangShu`](@ref) with parameters
# `threshold` and `variables`.
# ````julia
# stage_limiter! = PositivityPreservingLimiterZhangShu(thresholds=thresholds,
#                                                      variables=variables)
# ````
# Then, the limiter is added to the time integration method in the `solve` function. For instance, like
# ````julia
# CarpenterKennedy2N54(stage_limiter!, williamson_condition=false)
# ````
# or
# ````julia
# SSPRK43(stage_limiter!).
# ````

# # Simulation with shock capturing and positivity preserving

# Now, we can run a simulation using the described methods of shock capturing and positivity
# preserving limiters. We want to give an example for the 2D compressible Euler equations.
using OrdinaryDiffEq, Trixi

equations = CompressibleEulerEquations2D(1.4)

# As our initial condition we use the Sedov blast wave setup.
function initial_condition_sedov_blast_wave(x, t, equations::CompressibleEulerEquations2D)
  ## Set up polar coordinates
  inicenter = SVector(0.0, 0.0)
  x_norm = x[1] - inicenter[1]
  y_norm = x[2] - inicenter[2]
  r = sqrt(x_norm^2 + y_norm^2)

  r0 = 0.21875 # = 3.5 * smallest dx (for domain length=4 and max-ref=6)
  ## r0 = 0.5 # = more reasonable setup
  E = 1.0
  p0_inner = 3 * (equations.gamma - 1) * E / (3 * pi * r0^2)
  p0_outer = 1.0e-5 # = true Sedov setup
  ## p0_outer = 1.0e-3 # = more reasonable setup

  ## Calculate primitive variables
  rho = 1.0
  v1  = 0.0
  v2  = 0.0
  p   = r > r0 ? p0_outer : p0_inner

  return prim2cons(SVector(rho, v1, v2, p), equations)
end
initial_condition = initial_condition_sedov_blast_wave
#-
basis = LobattoLegendreBasis(3)

# We set the numerical fluxes and divide between the surface flux and the two volume fluxes for the DG
# and FV method. Here, we are using [`flux_lax_friedrichs`](@ref) and [`flux_ranocha`](@ref).
surface_flux = flux_lax_friedrichs
volume_flux  = flux_ranocha

# Now, we specify the shock capturing indicator $\alpha$.

# We implement the described indicator of Hennemann, Gassner as explained above with parameters
# `equations`, `basis`, `alpha_max`, `alpha_min`, `alpha_smooth` and `variable`.
# Since density and pressure are the critical variables in this example, we use
# `density_pressure = density * pressure = rho * p` as indicator variable.
indicator_sc = IndicatorHennemannGassner(equations, basis,
                                         alpha_max=0.5,
                                         alpha_min=0.001,
                                         alpha_smooth=true,
                                         variable=density_pressure)

# Now, we can use the defined fluxes and the indicator to implement the volume integral using shock
# capturing.
volume_integral = VolumeIntegralShockCapturingHG(indicator_sc;
                                                 volume_flux_dg=volume_flux,
                                                 volume_flux_fv=surface_flux)

# We finalize the discretization by implementing Trixi's `solver`, `mesh`, `semi` and `ode`,
# while `solver` now has the extra parameter `volume_integral`.
solver = DGSEM(basis, surface_flux, volume_integral)

coordinates_min = (-2.0, -2.0)
coordinates_max = ( 2.0,  2.0)
mesh = TreeMesh(coordinates_min, coordinates_max,
                initial_refinement_level=6,
                n_cells_max=10_000)

semi = SemidiscretizationHyperbolic(mesh, equations, initial_condition, solver)

tspan = (0.0, 1.0)
ode = semidiscretize(semi, tspan);

# We add some callbacks to get an solution analysis and use a CFL-based time step size calculation.
analysis_callback = AnalysisCallback(semi, interval=100)

stepsize_callback = StepsizeCallback(cfl=0.8)

callbacks = CallbackSet(analysis_callback, stepsize_callback);

# We now run the simulation using the positivity preserving limiter of Zhang and Shu for the variables
# density and pressure.
stage_limiter! = PositivityPreservingLimiterZhangShu(thresholds=(5.0e-6, 5.0e-6),
                                                     variables=(Trixi.density, pressure))

sol = solve(ode, CarpenterKennedy2N54(stage_limiter!, williamson_condition=false),
            dt=1.0, # solve needs some value here but it will be overwritten by the stepsize_callback
            save_everystep=false, callback=callbacks);

using Plots
plot(sol)


# MIP01 from CPLEX.jl
function int1test(model::MOI.ModelLike, config::TestConfig)
    atol = config.atol
    rtol = config.rtol
    # an example on mixed integer programming
    #
    #   maximize 1.1x + 2 y + 5 z
    #
    #   s.t.  x + y + z <= 10
    #         x + 2 y + z <= 15
    #
    #         x is continuous: 0 <= x <= 5
    #         y is integer: 0 <= y <= 10
    #         z is binary

    @test MOI.supports(model, MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}())
    @test MOI.supports(model, MOI.ObjectiveSense())
    @test MOI.supports_constraint(model, MOI.ScalarAffineFunction{Float64}, MOI.LessThan{Float64})
    @test MOI.supports_constraint(model, MOI.SingleVariable, MOI.GreaterThan{Float64})
    @test MOI.supports_constraint(model, MOI.SingleVariable, MOI.ZeroOne)
    @test MOI.supports_constraint(model, MOI.SingleVariable, MOI.Integer)

    MOI.empty!(model)
    @test MOI.is_empty(model)

    v = MOI.add_variables(model, 3)
    @test MOI.get(model, MOI.NumberOfVariables()) == 3

    cf = MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([1.0,1.0,1.0], v), 0.0)
    c = MOI.add_constraint(model, cf, MOI.LessThan(10.0))
    @test MOI.get(model, MOI.NumberOfConstraints{MOI.ScalarAffineFunction{Float64},MOI.LessThan{Float64}}()) == 1

    cf2 = MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([1.0,2.0,1.0], v), 0.0)
    c2 = MOI.add_constraint(model, cf2, MOI.LessThan(15.0))
    @test MOI.get(model, MOI.NumberOfConstraints{MOI.ScalarAffineFunction{Float64},MOI.LessThan{Float64}}()) == 2

    MOI.add_constraint(model, MOI.SingleVariable(v[1]), MOI.Interval(0.0, 5.0))
    @test MOI.get(model, MOI.NumberOfConstraints{MOI.SingleVariable,MOI.Interval{Float64}}()) == 1

    MOI.add_constraint(model, MOI.SingleVariable(v[2]), MOI.Interval(0.0, 10.0))
    @test MOI.get(model, MOI.NumberOfConstraints{MOI.SingleVariable,MOI.Interval{Float64}}()) == 2
    MOI.add_constraint(model, MOI.SingleVariable(v[2]), MOI.Integer())
    @test MOI.get(model, MOI.NumberOfConstraints{MOI.SingleVariable,MOI.Integer}()) == 1

    MOI.add_constraint(model, MOI.SingleVariable(v[3]), MOI.ZeroOne())
    @test MOI.get(model, MOI.NumberOfConstraints{MOI.SingleVariable,MOI.ZeroOne}()) == 1

    objf = MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([1.1, 2.0, 5.0], v), 0.0)
    MOI.set(model, MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}(), objf)
    MOI.set(model, MOI.ObjectiveSense(), MOI.MaxSense)

    @test MOI.get(model, MOI.ObjectiveSense()) == MOI.MaxSense

    if config.solve
        MOI.optimize!(model)

        @test MOI.get(model, MOI.TerminationStatus()) == MOI.Success

        @test MOI.get(model, MOI.ResultCount()) >= 1

        @test MOI.get(model, MOI.PrimalStatus()) in [ MOI.FeasiblePoint, MOI.NearlyFeasiblePoint ]

        @test MOI.get(model, MOI.ObjectiveValue()) ≈ 19.4 atol=atol rtol=rtol

        @test MOI.get(model, MOI.VariablePrimal(), v) ≈ [4,5,1] atol=atol rtol=rtol

        @test MOI.get(model, MOI.ConstraintPrimal(), c) ≈ 10 atol=atol rtol=rtol

        @test MOI.get(model, MOI.ConstraintPrimal(), c2) ≈ 15 atol=atol rtol=rtol

        @test MOI.get(model, MOI.ObjectiveBound()) >= 19.4
        # FIXME the following are currently not implemented in MockOptimizer
#        @test MOI.get(model, MOI.RelativeGap()) >= 0.0
#        @test MOI.get(model, MOI.SolveTime()) >= 0.0
#        @test MOI.get(model, MOI.SimplexIterations()) >= 0
#        @test MOI.get(model, MOI.BarrierIterations()) >= 0
#        @test MOI.get(model, MOI.NodeCount()) >= 0
    end
end

Base.isapprox(a::T, b::T; kwargs...) where T <: Union{MOI.SOS1, MOI.SOS2} = isapprox(a.weights, b.weights; kwargs...)

# sos from CPLEX.jl" begin
function int2test(model::MOI.ModelLike, config::TestConfig)
    atol = config.atol
    rtol = config.rtol
    @testset "SOSI" begin
        @test MOI.supports(model, MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}())
        @test MOI.supports(model, MOI.ObjectiveSense())
        @test MOI.supports_constraint(model, MOI.VectorOfVariables, MOI.SOS1{Float64})
        @test MOI.supports_constraint(model, MOI.SingleVariable, MOI.LessThan{Float64})

        MOI.empty!(model)
        @test MOI.is_empty(model)

        v = MOI.add_variables(model, 3)
        @test MOI.get(model, MOI.NumberOfVariables()) == 3
        MOI.add_constraint(model, MOI.SingleVariable(v[1]), MOI.LessThan(1.0))
        MOI.add_constraint(model, MOI.SingleVariable(v[2]), MOI.LessThan(1.0))
        MOI.add_constraint(model, MOI.SingleVariable(v[3]), MOI.LessThan(2.0))

        c1 = MOI.add_constraint(model, MOI.VectorOfVariables([v[1], v[2]]), MOI.SOS1([1.0, 2.0]))
        c2 = MOI.add_constraint(model, MOI.VectorOfVariables([v[1], v[3]]), MOI.SOS1([1.0, 2.0]))
        @test MOI.get(model, MOI.NumberOfConstraints{MOI.VectorOfVariables,MOI.SOS1{Float64}}()) == 2

        #=
            To allow for permutations in the sets and variable vectors
            we're going to sort according to the weights
        =#
        cs_sos = MOI.get(model, MOI.ConstraintSet(), c2)
        cf_sos = MOI.get(model, MOI.ConstraintFunction(), c2)
        p = sortperm(cs_sos.weights)
        @test cs_sos.weights[p] ≈ [1.0, 2.0] atol=atol rtol=rtol
        @test cf_sos.variables[p] == v[[1,3]]

        objf = MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([2.0, 1.0, 1.0], v), 0.0)
        MOI.set(model, MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}(), objf)
        MOI.set(model, MOI.ObjectiveSense(), MOI.MaxSense)
        @test MOI.get(model, MOI.ObjectiveSense()) == MOI.MaxSense

        if config.solve
            MOI.optimize!(model)

            @test MOI.get(model, MOI.TerminationStatus()) == MOI.Success

            @test MOI.get(model, MOI.ResultCount()) >= 1

            @test MOI.get(model, MOI.PrimalStatus()) == MOI.FeasiblePoint

            @test MOI.get(model, MOI.ObjectiveValue()) ≈ 3 atol=atol rtol=rtol

            @test MOI.get(model, MOI.VariablePrimal(), v) ≈ [0,1,2] atol=atol rtol=rtol
        end

        MOI.delete(model, c1)
        MOI.delete(model, c2)

        if config.solve
            MOI.optimize!(model)

            @test MOI.get(model, MOI.TerminationStatus()) == MOI.Success

            @test MOI.get(model, MOI.ResultCount()) >= 1

            @test MOI.get(model, MOI.PrimalStatus()) == MOI.FeasiblePoint

            @test MOI.get(model, MOI.ObjectiveValue()) ≈ 5 atol=atol rtol=rtol

            @test MOI.get(model, MOI.VariablePrimal(), v) ≈ [1,1,2] atol=atol rtol=rtol
        end
    end
    @testset "SOSII" begin
        @test MOI.supports(model, MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}())
        @test MOI.supports(model, MOI.ObjectiveSense())
        @test MOI.supports_constraint(model, MOI.VectorOfVariables, MOI.SOS1{Float64})
        @test MOI.supports_constraint(model, MOI.VectorOfVariables, MOI.SOS2{Float64})
        @test MOI.supports_constraint(model, MOI.SingleVariable, MOI.ZeroOne)
        @test MOI.supports_constraint(model, MOI.ScalarAffineFunction{Float64}, MOI.EqualTo{Float64})

        MOI.empty!(model)
        @test MOI.is_empty(model)

        v = MOI.add_variables(model, 10)
        @test MOI.get(model, MOI.NumberOfVariables()) == 10

        bin_constraints = []
        for i in 1:8
            MOI.add_constraint(model, MOI.SingleVariable(v[i]), MOI.Interval(0.0, 2.0))
            push!(bin_constraints, MOI.add_constraint(model, MOI.SingleVariable(v[i]), MOI.ZeroOne()))
        end

        MOI.add_constraint(model,
            MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([1.0,2.0,3.0,-1.0], v[[1,2,3,9]]), 0.0),
            MOI.EqualTo(0.0)
        )

        MOI.add_constraint(model,
            MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([5.0,4.0,7.0,2.0,1.0,-1.0], v[[4,5,6,7,8,10]]), 0.0),
            MOI.EqualTo(0.0)
        )

        MOI.add_constraint(model,
            MOI.VectorOfVariables(v[[1, 2, 3]]),
            MOI.SOS1([1.0, 2.0, 3.0])
        )

        vv   = MOI.VectorOfVariables(v[[4,5,6,7,8]])
        sos2 = MOI.SOS2([5.0, 4.0, 7.0, 2.0, 1.0])
        c = MOI.add_constraint(model, vv, sos2)

        #=
            To allow for permutations in the sets and variable vectors
            we're going to sort according to the weights
        =#
        cs_sos = MOI.get(model, MOI.ConstraintSet(), c)
        cf_sos = MOI.get(model, MOI.ConstraintFunction(), c)
        p = sortperm(cs_sos.weights)
        @test cs_sos.weights[p] ≈ [1.0, 2.0, 4.0, 5.0, 7.0] atol=atol rtol=rtol
        @test cf_sos.variables[p] == v[[8,7,5,4,6]]

        objf = MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([1.0, 1.0], [v[9], v[10]]), 0.0)
        MOI.set(model, MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}(), objf)
        MOI.set(model, MOI.ObjectiveSense(), MOI.MaxSense)
        @test MOI.get(model, MOI.ObjectiveSense()) == MOI.MaxSense

        if config.solve
            MOI.optimize!(model)

            @test MOI.get(model, MOI.TerminationStatus()) == MOI.Success

            @test MOI.get(model, MOI.ResultCount()) >= 1

            @test MOI.get(model, MOI.PrimalStatus()) == MOI.FeasiblePoint

            @test MOI.get(model, MOI.ObjectiveValue()) ≈ 15.0 atol=atol rtol=rtol

            @test MOI.get(model, MOI.VariablePrimal(), v) ≈ [0.0, 0.0, 1.0, 1.0, 0.0, 1.0, 0.0, 0.0, 3.0, 12.0] atol=atol rtol=rtol
        end

        for cref in bin_constraints
            MOI.delete(model, cref)
        end

        if config.solve
            MOI.optimize!(model)

            @test MOI.get(model, MOI.TerminationStatus()) == MOI.Success

            @test MOI.get(model, MOI.ResultCount()) >= 1

            @test MOI.get(model, MOI.PrimalStatus()) == MOI.FeasiblePoint

            @test MOI.get(model, MOI.ObjectiveValue()) ≈ 30.0 atol=atol rtol=rtol

            @test MOI.get(model, MOI.VariablePrimal(), v) ≈ [0.0, 0.0, 2.0, 2.0, 0.0, 2.0, 0.0, 0.0, 6.0, 24.0] atol=atol rtol=rtol
        end
    end
end

# CPLEX #76
function int3test(model::MOI.ModelLike, config::TestConfig)
    atol = config.atol
    rtol = config.rtol
    # integer knapsack problem
    # max   z - 0.5 ( b1 + b2 + b3) / 40
    # s.t.  0 <= z - 0.5 eᵀ b / 40 <= 0.999
    #       b1, b2, ... b10 ∈ {0, 1}
    #       z in {0, 1, 2, ..., 100}

    MOI.empty!(model)
    @test MOI.is_empty(model)

    @test MOI.supports(model, MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}())
    @test MOI.supports(model, MOI.ObjectiveSense())
    @test MOI.supports_constraint(model, MOI.SingleVariable, MOI.ZeroOne)
    @test MOI.supports_constraint(model, MOI.SingleVariable, MOI.Integer)
    @test MOI.supports_constraint(model, MOI.SingleVariable, MOI.Interval{Float64})
    @test MOI.supports_constraint(model, MOI.ScalarAffineFunction{Float64}, MOI.Interval{Float64})

    z = MOI.add_variable(model)
    MOI.add_constraint(model, MOI.SingleVariable(z), MOI.Integer())
    MOI.add_constraint(model, MOI.SingleVariable(z), MOI.Interval(0.0, 100.0))

    b = MOI.add_variables(model, 10)

    for bi in b
        MOI.add_constraint(model, MOI.SingleVariable(bi), MOI.ZeroOne())
    end

    c = MOI.add_constraint(model, MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.(vcat(1.0, fill(-0.5 / 40, 10)), vcat(z, b)), 0.0), MOI.Interval(0.0, 0.999))

    MOI.set(model, MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}(), MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.(vcat(1.0, fill(-0.5 / 40, 3)), vcat(z, b[1:3])), 0.0))
    MOI.set(model, MOI.ObjectiveSense(), MOI.MaxSense)

    if config.solve
        MOI.optimize!(model)

        @test MOI.get(model, MOI.TerminationStatus()) == MOI.Success

        @test MOI.get(model, MOI.PrimalStatus()) == MOI.FeasiblePoint

        @test MOI.get(model, MOI.ObjectiveValue()) ≈ 1 atol=atol rtol=rtol

        # test for CPLEX.jl #76
        MOI.optimize!(model)

        @test MOI.get(model, MOI.TerminationStatus()) == MOI.Success

        @test MOI.get(model, MOI.PrimalStatus()) == MOI.FeasiblePoint

        @test MOI.get(model, MOI.ObjectiveValue()) ≈ 1 atol=atol rtol=rtol
    end
end

# Mixed-integer linear problems

function knapsacktest(model::MOI.ModelLike, config::TestConfig)
    atol = config.atol
    rtol = config.rtol
    # integer knapsack problem
    # max 5a + 3b + 2c + 7d + 4e
    # st  2a + 8b + 4c + 2d + 5e <= 10
    #                  a,b,c,d,e ∈ binary

    MOI.empty!(model)
    @test MOI.is_empty(model)

    @test MOI.supports(model, MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}())
    @test MOI.supports(model, MOI.ObjectiveSense())
    @test MOI.supports_constraint(model, MOI.SingleVariable, MOI.ZeroOne)
    @test MOI.supports_constraint(model, MOI.ScalarAffineFunction{Float64}, MOI.LessThan{Float64})

    v = MOI.add_variables(model, 5)
    @test MOI.get(model, MOI.NumberOfVariables()) == 5

    for vi in v
        MOI.add_constraint(model, MOI.SingleVariable(vi), MOI.ZeroOne())
    end
    @test MOI.get(model, MOI.NumberOfConstraints{MOI.SingleVariable,MOI.ZeroOne}()) == 5
    c = MOI.add_constraint(model, MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([2.0, 8.0, 4.0, 2.0, 5.0], v), 0.0), MOI.LessThan(10.0))
    @test MOI.get(model, MOI.NumberOfConstraints{MOI.ScalarAffineFunction{Float64},MOI.LessThan{Float64}}()) == 1

    MOI.set(model, MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}(), MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([5.0, 3.0, 2.0, 7.0, 4.0], v), 0.0))
    MOI.set(model, MOI.ObjectiveSense(), MOI.MaxSense)

    if MOI.supports(model, MOI.VariablePrimalStart(), MOI.VariableIndex)
        MOI.set(model, MOI.VariablePrimalStart(), v, [0.0, 0.0, 0.0, 0.0, 0.0])
    end

    if config.solve
        MOI.optimize!(model)

        @test MOI.get(model, MOI.TerminationStatus()) == MOI.Success

        @test MOI.get(model, MOI.PrimalStatus()) in [ MOI.FeasiblePoint, MOI.NearlyFeasiblePoint ]

        @test MOI.get(model, MOI.ObjectiveValue()) ≈ 16 atol=atol rtol=rtol

        @test MOI.get(model, MOI.VariablePrimal(), v) ≈ [1, 0, 0, 1, 1] atol=atol rtol=rtol
    end
end

const intlineartests = Dict("knapsack" => knapsacktest,
                            "int1"     => int1test,
                            "int2"     => int2test,
                            "int3"     => int3test)

@moitestset intlinear


#!/usr/bin/env julia

const DATA_DIR   = "../18_data"
const INPUT_FILE = "input.txt"

const LPAREN = '('
const RPAREN = ')'

const OP_PLUS  = '+'
const OP_TIMES = '*'

# ----------------------------------------
macro assert(bool_expr)
    message = string("Assertion: ", bool_expr, " failed")
    :($(esc(bool_expr)) || error($message))
end
# ----------------------------------------
@enum Op op_plus op_times

mutable struct PartialExpr
    is_subexpr::Bool
    lval::Union{Int, Nothing}
    op::Union{Op, Nothing}
    # --------------------
    function PartialExpr(is_subexpr=false)
        return new(is_subexpr, nothing, nothing)
    end
end
# ----------------------------------------

function pexpr_addval!(pexpr, val)
    if isnothing(pexpr.lval)
        pexpr.lval = val
    else
        if isnothing(pexpr.op)
            error("Encountered digit where operator expected")
        else
            new_lval = eval_op(pexpr.lval, pexpr.op, val)
            pexpr.lval = new_lval
            pexpr.op = nothing
        end
    end
end

function pexpr_print(pexpr::PartialExpr)
    op_str = isnothing(pexpr.op) ? "" : (pexpr.op == op_plus ? "+" : "*")
    subexpr_str = pexpr.is_subexpr ? "(subexpr)" : ""
    output = "PExpr: $(pexpr.lval) $(op_str) $(subexpr_str)"
    println(strip(output))
end

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

function eval_op(arg_left::Int, op::Op, arg_right::Int)
    if op == op_plus
        return arg_left + arg_right
    elseif op == op_times
        return arg_left * arg_right
    elseif isnothing(op)
        error("operator is not set")
    else
        error("Unrecognized operator value")
    end
end

function eval_expr(expr::String)
    partials = Array{PartialExpr,1}()
    push!(partials, PartialExpr())
    for c in expr
        if isspace(c)
            continue
        end
        top = last(partials)
        if isdigit(c)
            val = parse(Int, c)
            pexpr_addval!(top, val)
        elseif c == LPAREN
            push!(partials, PartialExpr(true))
        elseif c == RPAREN
            @assert top.is_subexpr
            @assert !isnothing(top.lval)
            @assert isnothing(top.op)
            val = top.lval
            _ = pop!(partials)
            top = last(partials)
            pexpr_addval!(top, val)
        elseif c == OP_PLUS
            @assert !isnothing(top.lval)
            @assert isnothing(top.op)
            top.op = op_plus
        elseif c == OP_TIMES
            @assert !isnothing(top.lval)
            @assert isnothing(top.op)
            top.op = op_times
        else
            error("Unrecognized character: $(c)")
        end
    end
    @assert length(partials) == 1
    only = partials[1]
    @assert !isnothing(only.lval)
    @assert isnothing(only.op)
    return only.lval
end

function main(filename)
    lines::Array{String} = readlines(filename)
    result = 0
    for line in lines
        # println("line=$(line)")
        val = eval_expr(line)
        # println("\tvalue=$(val)")
        result += val
    end
    println("Total of all exprs=$(result)")
end

main("$DATA_DIR/$INPUT_FILE")


using DataFrames
#using StatsBase
using Statistics
using Distributions
using CSV
using StatsBase
using Printf
export spearman_cor, consolidate_dataframe, read_dataframe, write_dataframe_with_comments, write_dataframe
export goal_lookup_from_df, add_edf_col_from_df, frequencies_from_counts, variance1, variance2

# Returns a dataframe by reading a CSV file with comments that start with "#"
function read_dataframe( df_filename::String )
  df = CSV.read(df_filename,DataFrame,header=true,comment="#")
  df
end

function scor_df( filename::String )
  df = read_dataframe( filename )
  rnames = vcat( [:logsteps], names(df)[12:20])  # result column names
  cdf = DataFrame()
  cdf.names = String.(rnames)
  for i = 1:length(rnames)
    cdf[!,rnames[i]] = zeros(Float64,length(rnames))
    for j = (i+1):length(rnames)
      cdf[j,rnames[i]] = StatsBase.corspearman( df[!,rnames[i]], df[!,rnames[j]] )
    end
  end
  cdf
end

function t_value( v1::Vector{Float64}, v2::Vector{Float64} )
  r = corspearman( v1, v2 )
  r*sqrt((length(v1)-2)/(1-r^2))
end

# Returns a pair of the Spearman correlation between two columns of df and a p-value for this correlation
function spearman_cor( df::DataFrame, name1::Symbol, name2::Symbol )
  v1 = Vector{Float64}(df[!,name1])
  v2 = Vector{Float64}(df[!,name2])
  r = StatsBase.corspearman( v1, v2 )
  t_value = r*sqrt((length(v1)-2)/(1-r^2))
  # p_value is always positive
  p_value = r>=0.0 ? Distributions.ccdf( TDist(length(v1)-2), t_value ) : Distributions.cdf( TDist(length(v1)-2), t_value )
  (r, p_value)
end

function spearman_cor( v1::AbstractVector, v2::AbstractVector )
  r = StatsBase.corspearman( v1, v2 )
  t_value = r*sqrt((length(v1)-2)/(1-r^2))
  # p_value is always positive
  p_value = r>=0.0 ? Distributions.ccdf( TDist(length(v1)-2), t_value ) : Distributions.cdf( TDist(length(v1)-2), t_value )
  (r, p_value)
end

# Consolidates a dataframe with the columns specified below by averaging columns with the same parameter values
# Revised 10/27/20
# Similar to R aggregate function.
function consolidate_dataframe( in_filename::String, out_filename::String )
  df = read_dataframe( in_filename )
  cdf = consolidate_dataframe( df )
  open( out_filename, "w" ) do outfile
    open( in_filename, "r" ) do infile
      line = readline(infile)
      while line[1] == '#'
        #println("line: ",line)
        #if line[1:6] != "# cor("
          write(outfile,string(line,"\n"))
        #end
        line = readline(infile)
      end
    end
   #write(outfile,"# consolidate increement: $increment \n")
   CSV.write( outfile, cdf, append=true, writeheader=true )
  end
end

# Consolidates a dataframe with the columns specified below by averaging rows with the same parameter values over all columns except goals
# Assumes that the first column is the "goal" column which may either be Vectors of MyInts or strings.
# The :evo_count column is not averaged since these are cummulative over the rows corresponding to a goal.
# Thus, for the :evo_count column, the last row corresponding to the goal is the consolidated value.
# Revised 10/18/21   Previous values for evo_count were not correct
# Similar to R aggregate function.
function consolidate_dataframe( df::DataFrame )
  new_df = DataFrame()
  #df.n_combined = map(Float64,df.n_combined)  # Only for neutral walk files, remove otherwise
  for n in names(df)
    if n!= "ntries" && n != "evo_count"
      new_df[!,n] = typeof(df[1,n])[]
    else
      new_df[!,n] = Float64[]
    end
  end
  #println("size(new_df): ",size(new_df))
  #println("names(new_df): ",names(new_df))
  new_names = vcat(["goal"],names(df)[2:end])
  #println("new_names: ",new_names)
  select!(df,new_names)
  df = sort(df,[:goal])
  ff = findfirst(x->x!=df.goal[1],df.goal)
  increment = ff != nothing ? ff - 1 : size(df)[1]
  #println("increment: ",increment)
  if Int(floor(size(df)[1]/increment)) != Int(ceil(size(df)[1]/increment))
    error("The number of rows in the dataframe must be a multiple of increment")
  end
  for r = 1:increment:Int(floor(size(df)[1]))
    ndf_row = Any[df.goal[r]]   # first entry of row is the current goal
    isum = zeros(Int64,size(df)[2])
    fsum = zeros(Float64,size(df)[2])
    for row = r:(r+increment-1)
      #println("r: ",r,"  row: ",row)
      for i = 2:size(df)[2]
        if names(df)[i] == "evo_count"
          #println("evo_count skip value: ",df[row,:evo_count])
          continue
        elseif typeof(df[row,i]) <: Int64
          isum[i] += df[row,i]
        elseif typeof(df[row,i]) <: Float64
          fsum[i] += df[row,i]
        else
          error("Illegal type of dataframe element")
        end
      end
    end
    for i = 2:size(df)[2]
      if names(df)[i] == "evo_count"
        #println("evo_count push value: ",df[r+increment-1,:evo_count])
        push!(ndf_row,df[r+increment-1,:evo_count])  # last row for goal has correct evo_count
      elseif typeof(df[r,i]) <: Int64
        #print(" ",names(df)[i]," ",isum[i])
        push!(ndf_row,isum[i]/increment)
      elseif typeof(df[r,i]) <: Float64
        #print(" ",names(df)[i]," ",fsum[i])
        push!(ndf_row,fsum[i]/increment)
      end
      #println()
    end
    #println("size(new_df): ",new_df)
    #println("length(ndf_row): ",length(ndf_row))
    #println("r:",r," ",ndf_row)
    push!(new_df,ndf_row)
  end
  new_df
end

function write_dataframe( df::DataFrame, out_filename::String )
  CSV.write( out_filename, df, append=true, writeheader=true )
end

# Writes the datarame df to the file out_filename with comments taken from file comments_filename.
# No dataframe is read from comments_filename.
function write_dataframe_with_comments( df::DataFrame, comments_filename::String, out_filename::String )
  open( out_filename, "w" ) do outfile
    open( comments_filename, "r" ) do infile
      line = readline(infile)
      while line[1] == '#'
        #println("line: ",line)
        if line[1:6] != "# cor("
          write(outfile,string(line,"\n"))
        end
        line = readline(infile)
      end
    end
   #write(outfile,"# increement: $increment \n")
   CSV.write( outfile, df, append=true, writeheader=true )
  end
end

function write_comments( in_filename::String, out_filename::String )
  open( out_filename, "w" ) do outfile
    open( in_filename, "r" ) do infile
      line = readline(infile)
      while line[1] == '#'
        println("line: ",line)
        write(outfile,string(line,"\n"))
        line = readline(infile)
      end
    end
  end
end

# Given a goal represented as a string, extract the df values of the columns fields for that goal.
# Example:  goal_lookup_from_df(df,"0x4208",[:counts10ints]) returns 1192.
function goal_lookup_from_df( df::DataFrame, goal::String, fields::Vector{Symbol} )
  [df[findfirst(x->x==goal,df.goal),f] for f in fields]
end

# Adds a column to dataframe edf whose values are obtained by goal_lookup_from_df( df, goal_string, [field] )
function add_edf_col_from_df( df::DataFrame, edf::DataFrame, field::Symbol )
  edf[!,field] =  zeros(Float64,size(edf)[1])
  for i = 1:size(edf)[1]
    # Extract MyInt goal from edf.goal[i] by parsing.  See Julia documentation for Metaprogramming.
    goal = Meta.parse(edf[i,:goal]).args[2]
    goal_string = @sprintf("0x%x",goal)   # Convert to string
    edf[i,field] = goal_lookup_from_df(df,goal_string,[field])[1]
  end
  edf
end

# Return the frequencies of the unique elements of V
function frequencies_from_counts( V::Vector{Int64} )
  sort!(V)
  #println("V: ",V)
  freq = Int64[]
  i = 1
  while i <= length(V)
    push!(freq,1)
    while i < length(V) && V[i+1] == V[i]
      freq[end] += 1
      i += 1
    end
    i += 1
  end
  #println("sum(freq): ",sum(freq),"  freq: ",freq)
  # Two assertions to check correctness
  @assert sum(freq) == length(V)
  @assert length(freq) == length(unique(V))
  freq
end

function pad_missing( V::Vector, new_length::Int64 )
  result = Union{Missing,Int64}[]
    for v in V push!(result,v)
  end
  for i = length(V)+1:new_length
    push!(result,missing)
  end
  result
end

# Simple sample variance
function variance1( X::Vector )
  N = length(X)
  avg = sum(X)/N
  sum( (x - avg)^2 for x in X )/(N-1)
end

# Simple sample variance
function variance2( X::Vector )
  N = length(X)
  ( sum( x^2 for x in X) - sum(X)^2/N )/(N-1)
end


### A Pluto.jl notebook ###
# v0.12.11

using Markdown
using InteractiveUtils

# ╔═╡ 3c57ceb0-f934-11ea-1bef-c16b272d0c6b
using Pkg, DrWatson

# ╔═╡ 3c580a7e-f934-11ea-3b0c-4d5d7d5a7382
begin
	@quickactivate "StatisticalRethinkingTuring"
	using Turing
	using StatisticalRethinking
end

# ╔═╡ 83ad95c0-f933-11ea-1552-23ac7681803c
md"## Pluto-helpers-prettytable.jl"

# ╔═╡ b4656386-2db1-11eb-089f-993d79325858
begin
	# Define the experiment
	niter = 4000
	nparams = 3
	nchains = 2

	# some sample experiment results
	val = randn(niter, nparams, nchains) .+ [1, 2, 3]'
	val = hcat(val, rand(1:2, niter, 1, nchains))

	# construct a Chains object
	chns = Chains(val, start = 1, thin = 2)
end;

# ╔═╡ d1670910-2c0d-11eb-12d0-616c518e22da
Text(sprint(show, "text/plain", chns))

# ╔═╡ 5e493ac4-2db2-11eb-0567-7d635483d547
begin
	bounds = hpd(chns)
	Text(sprint(show, "text/plain", bounds))
end

# ╔═╡ 36733c28-2da0-11eb-3b0b-d5c9b85e16ba
CHNS(chns)

# ╔═╡ eff124b4-2dda-11eb-2717-dde3f4d416e6
HPD(chns)

# ╔═╡ 3ce617ae-f934-11ea-0e4a-ff1ec9ee2610
md"## End pluto-helpers-prettytable.jl"

# ╔═╡ Cell order:
# ╟─83ad95c0-f933-11ea-1552-23ac7681803c
# ╠═3c57ceb0-f934-11ea-1bef-c16b272d0c6b
# ╠═3c580a7e-f934-11ea-3b0c-4d5d7d5a7382
# ╠═b4656386-2db1-11eb-089f-993d79325858
# ╠═d1670910-2c0d-11eb-12d0-616c518e22da
# ╠═5e493ac4-2db2-11eb-0567-7d635483d547
# ╠═36733c28-2da0-11eb-3b0b-d5c9b85e16ba
# ╠═eff124b4-2dda-11eb-2717-dde3f4d416e6
# ╟─3ce617ae-f934-11ea-0e4a-ff1ec9ee2610


# ------------------------------------------------------------------
# Licensed under the ISC License. See LICENSE in the project root.
# ------------------------------------------------------------------

"""
    NearestNeighborSearcher(object, k; locations=all, metric=Euclidean())

A method for searching `k` nearest neighbors in spatial `object`
`locations` according to `metric`.
"""
struct NearestNeighborSearcher{O,K} <: AbstractBoundedNeighborSearcher
  # input fields
  object::O
  k::Int
  locations::Vector{Int}

  # state fields
  kdtree::K
end

function NearestNeighborSearcher(object::O, k::Int; locations=1:npoints(object), metric=Euclidean()) where {O}
  @assert 1 ≤ k ≤ length(locations) "number of neighbors must be smaller than number of locations"
  @assert locations ⊆ 1:npoints(object) "invalid locations for spatial object"
  kdtree = KDTree(coordinates(object, locations), metric)
  NearestNeighborSearcher{O,typeof(kdtree)}(object, k, locations, kdtree)
end

maxneighbors(searcher::NearestNeighborSearcher) = searcher.k

function search!(neighbors::AbstractVector{Int}, xₒ::AbstractVector,
                 searcher::NearestNeighborSearcher; mask=nothing)
  k       = searcher.k
  inds, _ = knn(searcher.kdtree, xₒ, k, true)
  locs    = view(searcher.locations, inds)

  if mask ≠ nothing
    nneigh = 0
    @inbounds for i in 1:k
      if mask[locs[i]]
        nneigh += 1
        neighbors[nneigh] = locs[i]
      end
    end
  else
    nneigh = k
    @inbounds for i in 1:k
      neighbors[i] = locs[i]
    end
  end

  nneigh
end


using StochasticDiffEq, Plots, DiffEqDevTools, DiffEqProblemLibrary
using DiffEqProblemLibrary.SDEProblemLibrary: importsdeproblems; importsdeproblems()
import DiffEqProblemLibrary.SDEProblemLibrary: prob_sde_additivesystem,
            prob_sde_additive, prob_sde_2Dlinear, prob_sde_linear, prob_sde_wave
gr()
const N = 1000

prob = prob_sde_additivesystem
prob = remake(prob,tspan=(0.0,1.0))

reltols = 1.0 ./ 10.0 .^ (1:5)
abstols = reltols#[0.0 for i in eachindex(reltols)]
setups = [Dict(:alg=>SRIW1())
          Dict(:alg=>EM(),:dts=>1.0./5.0.^((1:length(reltols)) .+ 1))
          Dict(:alg=>RKMil(),:dts=>1.0./5.0.^((1:length(reltols)) .+ 1),:adaptive=>false)
          Dict(:alg=>SRIW1(),:dts=>1.0./5.0.^((1:length(reltols)) .+ 1),:adaptive=>false)
          Dict(:alg=>SRA1(),:dts=>1.0./5.0.^((1:length(reltols)) .+ 1),:adaptive=>false)
          Dict(:alg=>SRA1())
          ]
names = ["SRIW1","EM","RKMil","SRIW1 Fixed","SRA1 Fixed","SRA1"]
wp = WorkPrecisionSet(prob,abstols,reltols,setups;numruns=N,names=names,maxiters=1e7,error_estimate=:l2)
plot(wp)

prob = prob_sde_additivesystem
prob = remake(prob,tspan=(0.0,1.0))

reltols = 1.0 ./ 10.0 .^ (1:5)
abstols = reltols#[0.0 for i in eachindex(reltols)]
setups = [
          Dict(:alg=>SRA1())
          Dict(:alg=>SRA2())
          Dict(:alg=>SRA3())
          Dict(:alg=>SOSRA())
          Dict(:alg=>SOSRA2())
          ]
wp = WorkPrecisionSet(prob,abstols,reltols,setups;numruns=N,maxiters=1e7,error_estimate=:l2)
plot(wp)

prob = prob_sde_additive
prob = remake(prob,tspan=(0.0,1.0))

reltols = 1.0 ./ 10.0 .^ (1:5)
abstols = reltols#[0.0 for i in eachindex(reltols)]

setups = [Dict(:alg=>SRIW1())
          Dict(:alg=>EM(),:dts=>1.0./5.0.^((1:length(reltols)) .+ 1))
          Dict(:alg=>RKMil(),:dts=>1.0./5.0.^((1:length(reltols)) .+ 1),:adaptive=>false)
          Dict(:alg=>SRIW1(),:dts=>1.0./5.0.^((1:length(reltols)) .+ 1),:adaptive=>false)
          Dict(:alg=>SRA1(),:dts=>1.0./5.0.^((1:length(reltols)) .+ 1),:adaptive=>false)
          Dict(:alg=>SRA1())
          ]
names = ["SRIW1","EM","RKMil","SRIW1 Fixed","SRA1 Fixed","SRA1"]
wp = WorkPrecisionSet(prob,abstols,reltols,setups;numruns=N,names=names,maxiters=1e7,error_estimate=:l2)
plot(wp)

prob = prob_sde_additive
prob = remake(prob,tspan=(0.0,1.0))

reltols = 1.0 ./ 10.0 .^ (1:5)
abstols = reltols#[0.0 for i in eachindex(reltols)]
setups = [
          Dict(:alg=>SRA1())
          Dict(:alg=>SRA2())
          Dict(:alg=>SRA3())
          Dict(:alg=>SOSRA())
          Dict(:alg=>SOSRA2())
          ]
wp = WorkPrecisionSet(prob,abstols,reltols,setups;numruns=N,error_estimate=:l2)
plot(wp)

prob = prob_sde_2Dlinear
prob = remake(prob,tspan=(0.0,1.0))

reltols = 1.0 ./ 10.0 .^ (1:5)
abstols = reltols#[0.0 for i in eachindex(reltols)]

setups = [Dict(:alg=>SRIW1())
          Dict(:alg=>EM(),:dts=>1.0./5.0.^((1:length(reltols)) .+ 1))
          Dict(:alg=>RKMil(),:dts=>1.0./5.0.^((1:length(reltols)) .+ 1),:adaptive=>false)
          Dict(:alg=>SRIW1(),:dts=>1.0./5.0.^((1:length(reltols)) .+ 1),:adaptive=>false)
          ]
names = ["SRIW1","EM","RKMil","SRIW1 Fixed"]
wp = WorkPrecisionSet(prob,abstols,reltols,setups;numruns=N,names=names,maxiters=1e7,error_estimate=:l2)
plot(wp)

prob = prob_sde_2Dlinear
prob = remake(prob,tspan=(0.0,1.0))

reltols = 1.0 ./ 10.0 .^ (1:5)
abstols = reltols#[0.0 for i in eachindex(reltols)]

setups = [Dict(:alg=>EM(),:dts=>1.0./5.0.^((1:length(reltols)) .+ 2))
          Dict(:alg=>RKMil(),:dts=>1.0./5.0.^((1:length(reltols)) .+ 2),:adaptive=>false)
          Dict(:alg=>SRI())
          Dict(:alg=>SRIW1())
          Dict(:alg=>SRIW2())
          Dict(:alg=>SOSRI())
          Dict(:alg=>SOSRI2())
          ]
wp = WorkPrecisionSet(prob,abstols,reltols,setups;numruns=N,maxiters=1e7,error_estimate=:l2)
plot(wp)

prob = prob_sde_linear
prob = remake(prob,tspan=(0.0,1.0))

reltols = 1.0 ./ 10.0 .^ (1:5)
abstols = reltols#[0.0 for i in eachindex(reltols)]

setups = [Dict(:alg=>SRIW1())
          Dict(:alg=>EM(),:dts=>1.0./5.0.^((1:length(reltols)) .+ 1))
          Dict(:alg=>RKMil(),:dts=>1.0./5.0.^((1:length(reltols)) .+ 1),:adaptive=>false)
          Dict(:alg=>SRIW1(),:dts=>1.0./5.0.^((1:length(reltols)) .+ 1),:adaptive=>false)
          ]
names = ["SRIW1","EM","RKMil","SRIW1 Fixed"]
wp = WorkPrecisionSet(prob,abstols,reltols,setups;numruns=N,names=names,maxiters=1e7,error_estimate=:l2)
plot(wp)

setups = [Dict(:alg=>EM(),:dts=>1.0./5.0.^((1:length(reltols)) .+ 2))
          Dict(:alg=>RKMil(),:dts=>1.0./5.0.^((1:length(reltols)) .+ 2),:adaptive=>false)
          Dict(:alg=>SRI())
          Dict(:alg=>SRIW1())
          Dict(:alg=>SRIW2())
          Dict(:alg=>SOSRI())
          Dict(:alg=>SOSRI2())
          ]
wp = WorkPrecisionSet(prob,abstols,reltols,setups;numruns=N,maxiters=1e7,error_estimate=:l2)
plot(wp)

prob = prob_sde_wave
prob = remake(prob,tspan=(0.0,1.0))

reltols = 1.0 ./ 10.0 .^ (1:5)
abstols = reltols#[0.0 for i in eachindex(reltols)]

setups = [Dict(:alg=>SRIW1())
          Dict(:alg=>EM(),:dts=>1.0./5.0.^((1:length(reltols)) .+ 1))
          Dict(:alg=>RKMil(),:dts=>1.0./5.0.^((1:length(reltols)) .+ 1),:adaptive=>false)
          Dict(:alg=>SRIW1(),:dts=>1.0./5.0.^((1:length(reltols)) .+ 1),:adaptive=>false)
          ]
names = ["SRIW1","EM","RKMil","SRIW1 Fixed"]
wp = WorkPrecisionSet(prob,abstols,reltols,setups;numruns=N,names=names,maxiters=1e7,error_estimate=:l2)
plot(wp)

setups = [Dict(:alg=>EM(),:dts=>1.0./5.0.^((1:length(reltols)) .+ 2))
          Dict(:alg=>RKMil(),:dts=>1.0./5.0.^((1:length(reltols)) .+ 2),:adaptive=>false)
          Dict(:alg=>SRI())
          Dict(:alg=>SRIW1())
          Dict(:alg=>SRIW2())
          Dict(:alg=>SOSRI())
          Dict(:alg=>SOSRI2())
          ]
wp = WorkPrecisionSet(prob,abstols,reltols,setups;numruns=N,maxiters=1e7,error_estimate=:l2)
plot(wp)

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


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

using Test
using BenchmarkTools
using OffsetArrays

@testset "transfinite_interpolate" begin
    include(normpath(joinpath(@__DIR__, "../../src/misc-util.jl")))
    using .__CFD2021__misc_util__: tuplejoin

    include("../src/transfinite-interpolate.jl")
    import .transfinite_interpolate: transfinite_interpolate_2d!, transfinite_interpolate_2d
    @testset "transfinite_interpolate_2d" begin
        Mx, My = 3, 3

        @testset "on functions" begin
            interp_func = transfinite_interpolate_2d( (x->2x, y->3y), (x->4x+3, y->2+5y) )
            @test interp_func((0.5,0.5)) == 3.0
        end

        @testset "on arrays" begin
            x_lo = ([LinRange(0, 1, Mx)...], [LinRange(0, 0, My)...]); x_hi = ([LinRange(0, 1, Mx)...], [LinRange(1, 1, My)...]);
            y_lo = ([LinRange(0, 0, Mx)...], [LinRange(0, 1, My)...]); y_hi = ([LinRange(1, 1, Mx)...], [LinRange(0, 1, My)...]);
            in_situ_interp_x = Array{Float64}(undef, (length(v) for v in x_lo)...);
            in_situ_interp_y = Array{Float64}(undef, (length(v) for v in y_lo)...);
            transfinite_interpolate_2d!(in_situ_interp_x, x_lo, x_hi)
            transfinite_interpolate_2d!(in_situ_interp_y, y_lo, y_hi)
            ab_situ_interp_x = transfinite_interpolate_2d(x_lo, x_hi)
            ab_situ_interp_y = transfinite_interpolate_2d(y_lo, y_hi)

            @testset "on ordinary arrays" begin
                @test in_situ_interp_x == ab_situ_interp_x == [0.0 0.0 0.0; 0.5 0.5 0.5; 1.0 1.0 1.0]
                @test in_situ_interp_y == ab_situ_interp_y == [0.0 0.5 1.0; 0.0 0.5 1.0; 0.0 0.5 1.0]
            end

            x_lo = map(x->OffsetArray(x, OffsetArrays.Origin(0)), x_lo)
            y_lo = map(x->OffsetArray(x, OffsetArrays.Origin(0)), y_lo)
            x_hi = map(x->OffsetArray(x, OffsetArrays.Origin(0)), x_hi)
            y_hi = map(x->OffsetArray(x, OffsetArrays.Origin(0)), y_hi)
            in_situ_interp_x = OffsetArray(in_situ_interp_x, tuplejoin((axes(v) for v in x_lo)...))
            in_situ_interp_y = OffsetArray(in_situ_interp_y, tuplejoin((axes(v) for v in y_lo)...))
            transfinite_interpolate_2d!(in_situ_interp_x, x_lo, x_hi)
            transfinite_interpolate_2d!(in_situ_interp_y, y_lo, y_hi)
            ab_situ_interp_x = transfinite_interpolate_2d(x_lo, x_hi)
            ab_situ_interp_y = transfinite_interpolate_2d(y_lo, y_hi)

            @testset "on OffsetArrays" begin
                @test in_situ_interp_x == ab_situ_interp_x != [0.0 0.0 0.0; 0.5 0.5 0.5; 1.0 1.0 1.0]
                @test in_situ_interp_y == ab_situ_interp_y != [0.0 0.5 1.0; 0.0 0.5 1.0; 0.0 0.5 1.0]
                @test in_situ_interp_x == ab_situ_interp_x == OffsetArray([0.0 0.0 0.0; 0.5 0.5 0.5; 1.0 1.0 1.0], OffsetArrays.Origin(0))
                @test in_situ_interp_y == ab_situ_interp_y == OffsetArray([0.0 0.5 1.0; 0.0 0.5 1.0; 0.0 0.5 1.0], OffsetArrays.Origin(0))
            end

            @testset "exception handling" begin
                x_lo = ([LinRange(0, 1, Mx)...], [LinRange(0, 0, My)...]); x_hi = ([LinRange(0, 1, 2Mx+1)...], [LinRange(1, 1, 2My+1)...]);
                @test_throws DimensionMismatch transfinite_interpolate_2d(x_lo, x_hi)
            end

        end # @testset "on arrays"
    end # @testset "transfinite_interpolate_2d"
end # @testset "transfinite_interpolate"

@testset "elliptic_grid_gen" begin
    include("../src/elliptic-grid-gen.jl")
    import .elliptic_grid_gen

    @testset "inverted_poisson_jacobi" begin
        @test_skip elliptic_grid_gen.inverted_poisson_jacobi_step
        @test_skip elliptic_grid_gen.inverted_poisson_jacobi_step!
        @test_skip elliptic_grid_gen.inverted_poisson_jacobi_iterate
        @test_skip elliptic_grid_gen.inverted_poisson_jacobi_iterate!
    end # @testset "inverted_poisson_jacobi"
end # @testset "elliptic_grid_gen"